datachain 0.1.10__tar.gz → 0.1.12__tar.gz

{datachain-0.1.10/src/datachain.egg-info → datachain-0.1.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.1.10
+Version: 0.1.12
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -36,6 +36,8 @@ Requires-Dist: multiprocess==0.70.16
 Requires-Dist: dill==0.3.8
 Requires-Dist: ujson>=5.9.0
 Requires-Dist: pydantic<3,>=2
+Requires-Dist: jmespath>=1.0
+Requires-Dist: datamodel-code-generator>=0.25
 Provides-Extra: cv
 Requires-Dist: Pillow<11,>=10.0.0; extra == "cv"
 Requires-Dist: torch>=2.1.0; extra == "cv"

{datachain-0.1.10 → datachain-0.1.12}/examples/clip.py

@@ -4,26 +4,22 @@ from torch.nn.functional import cosine_similarity
 from torch.utils.data import DataLoader
 
 from datachain.lib.dc import C, DataChain
-from datachain.lib.image import ImageReader
-from datachain.lib.text import TextReader
-from datachain.sql.functions import path
 
 source = "gs://dvcx-50k-laion-files/000000/00000000*"
 
 
 def create_dataset():
     imgs = (
-        DataChain(source)
+        DataChain.from_storage(source, type="image")
         .filter(C.name.glob("*.jpg"))
-        .mutate(stem=path.file_stem(C.name))
+        .map(stem=lambda name: name.split(".")[0], output=str)
     )
     captions = (
-        DataChain.from_storage(source, is_text=True)
+        DataChain.from_storage(source, type="text")
         .filter(C.name.glob("*.txt"))
-        .mutate(stem=path.file_stem(C.name))
-        .map(lambda file: file.get_value(), output={"caption": str})
+        .map(stem=lambda name: name.split(".")[0], output=str)
     )
-    return imgs.join(captions.select("stem", "text"), "stem").save("laion-50k")
+    return imgs.merge(captions, on="stem")
 
 
 if __name__ == "__main__":
@@ -34,9 +30,9 @@ if __name__ == "__main__":
     )
     tokenizer = open_clip.get_tokenizer("ViT-B-32")
 
-    ds = q.to_pytorch(
-        ImageReader(transform=preprocess),
-        TextReader("text", tokenizer=tokenizer),
+    ds = q.select("file", "right_file").to_pytorch(
+        transform=preprocess,
+        tokenizer=tokenizer,
     )
     loader = DataLoader(ds, batch_size=16)
 

datachain-0.1.12/examples/json-csv-reader.py

@@ -0,0 +1,87 @@
+#
+# TODO:
+# refactor lib/meta_formats/read_scema into a Datachain method
+#
+# ER: add support for Optional fields in read_schema()
+# ER: add support for headless CSV within static schema only
+# ER: fix the bug in datamodel-codegen failing to recognize csv float and int columns
+#
+# Open issues:
+# 1. A single filename cannot be passed as schema source (#1563)
+# 2. Need syntax like "file.open(encoding='utf-8')" to avoid "type=text" (#1614)
+# 3. Need syntax like "datachain.collate(func -> Any)" (#1615)
+# 4. "Feature" does not tolerate creating a class twice (#1617)
+# 5. Unsure how to deal with 'folder' pseudo-files in cloud systems(#1618)
+# 6. There should be exec() method to force-run the existing chain (#1616)
+# 7. data-model-codegenerator: datamodel-codegen reports all CSV fields as 'str'.
+# 8. from_json and from_csv methods do not filter empty files from AWS
+# dependencies:
+# pip install datamodel-code-generator
+# pip install jmespath
+
+from typing import Optional
+
+from pydantic import BaseModel
+
+from datachain.lib.dc import C, DataChain
+from datachain.lib.feature_utils import pydantic_to_feature
+from datachain.lib.meta_formats import read_schema
+
+
+# Sample model for static JSON model
+class LicenseModel(BaseModel):
+    url: str
+    id: int
+    name: str
+
+
+LicenseFeature = pydantic_to_feature(LicenseModel)
+
+
+# Sample model for static CSV model
+class ChatDialog(BaseModel):
+    id: Optional[int] = None
+    count: Optional[int] = None
+    sender: Optional[str] = None
+    text: Optional[str] = None
+
+
+ChatFeature = pydantic_to_feature(ChatDialog)
+
+
+def main():
+    uri = "gs://datachain-demo/coco2017/annotations_captions/"
+
+    print("Reading schema from the root COCO annotation")
+    chain = (
+        DataChain.from_storage(uri)
+        .filter(C.name.glob("*.json"))
+        .limit(1)
+        .map(  # dummy column created (#1615)
+            meta_schema=lambda file: read_schema(file, data_type="json"), output=str
+        )
+    )
+    # dummy executor (#1616)
+    chain.save()
+
+    print("static JSON schema test parsing 7 objects")
+    static_json_ds = DataChain.from_json(uri, jmespath="licenses", spec=LicenseFeature)
+    print(static_json_ds.to_pandas())
+
+    print("dynamic JSON schema test parsing 5K objects")
+    dynamic_json_ds = DataChain.from_json(uri, jmespath="images", show_schema=True)
+    print(dynamic_json_ds.to_pandas())
+
+    uri = "gs://datachain-demo/chatbot-csv/"
+    print("static CSV with header schema test parsing 3.5K objects")
+    static_csv_ds = DataChain.from_csv(uri, spec=ChatFeature)
+    print(static_csv_ds.to_pandas())
+
+    uri = "gs://datachain-demo/laion-aesthetics-csv"
+    print("dynamic CSV with header schema test parsing 3M objects")
+    dynamic_csv_ds = DataChain.from_csv(uri, show_schema=True)
+    print(dynamic_csv_ds.to_pandas())
+
+
+if __name__ == "__main__":
+    main()

{datachain-0.1.10 → datachain-0.1.12}/examples/torch-loader.py

@@ -6,8 +6,7 @@ from torch.utils.data import DataLoader
 from torchvision.transforms import v2
 
 from datachain.lib.dc import C, DataChain
-from datachain.lib.image import ImageReader
-from datachain.lib.reader import LabelReader
+from datachain.lib.pytorch import label_to_int
 
 STORAGE = "gs://dvcx-datalakes/dogs-and-cats/"
 
@@ -45,17 +44,13 @@ class CNN(nn.Module):
 
 if __name__ == "__main__":
     ds = (
-        DataChain(STORAGE)
+        DataChain.from_storage(STORAGE, type="image")
         .filter(C.name.glob("*.jpg"))
-        .map(lambda name: (name[:3],), output={"label": str})
+        .map(label=lambda name: label_to_int(name[:3], CLASSES), output=int)
     )
 
     train_loader = DataLoader(
-        ds.to_pytorch(
-            ImageReader(),
-            LabelReader("label", classes=CLASSES),
-            transform=transform,
-        ),
+        ds.to_pytorch(transform=transform),
         batch_size=16,
         num_workers=2,
     )

{datachain-0.1.10 → datachain-0.1.12}/pyproject.toml

@@ -39,7 +39,9 @@ dependencies = [
   "multiprocess==0.70.16",
   "dill==0.3.8",
   "ujson>=5.9.0",
-  "pydantic>=2,<3"
+  "pydantic>=2,<3",
+  "jmespath>=1.0",
+  "datamodel-code-generator>=0.25"
 ]
 
 [project.optional-dependencies]

{datachain-0.1.10 → datachain-0.1.12}/src/datachain/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.10'
-__version_tuple__ = version_tuple = (0, 1, 10)
+__version__ = version = '0.1.12'
+__version_tuple__ = version_tuple = (0, 1, 12)

{datachain-0.1.10 → datachain-0.1.12}/src/datachain/catalog/catalog.py

@@ -1580,10 +1580,54 @@ class Catalog:
 
         return dst
 
-    def open_object(self, row: RowDict, use_cache: bool = True, **config: Any):
+    def get_file_signals(
+        self, dataset_name: str, dataset_version: int, row: RowDict
+    ) -> Optional[dict]:
+        """
+        Function that returns file signals from dataset row.
+        Note that signal names are without prefix, so if there was 'laion__file__source'
+        in original row, result will have just 'source'
+        Example output:
+            {
+                "source": "s3://ldb-public",
+                "parent": "animals/dogs",
+                "name": "dog.jpg",
+                ...
+            }
+        """
+        from datachain.lib.signal_schema import SignalSchema
+
+        version = self.get_dataset(dataset_name).get_version(dataset_version)
+
+        file_signals_values = SignalSchema.deserialize(
+            version.feature_schema
+        ).get_file_signals_values(row)
+        if not file_signals_values:
+            return None
+
+        # there can be multiple file signals in a schema, but taking the first
+        # one for now. In future we might add ability to choose from which one
+        # to open object
+        return next(iter(file_signals_values.values()))
+
+    def open_object(
+        self,
+        dataset_name: str,
+        dataset_version: int,
+        row: RowDict,
+        use_cache: bool = True,
+        **config: Any,
+    ):
+        file_signals = self.get_file_signals(dataset_name, dataset_version, row)
+        if not file_signals:
+            raise RuntimeError("Cannot open object without file signals")
+
         config = config or self.client_config
-        client = self.get_client(row["source"], **config)
-        return client.open_object(self._get_row_uid(row), use_cache=use_cache)
+        client = self.get_client(file_signals["source"], **config)
+        return client.open_object(
+            self._get_row_uid(file_signals),  # type: ignore [arg-type]
+            use_cache=use_cache,
+        )
 
     def _get_row_uid(self, row: RowDict) -> UniqueId:
         return UniqueId(

{datachain-0.1.10 → datachain-0.1.12}/src/datachain/data_storage/metastore.py

@@ -1142,6 +1142,8 @@ class AbstractDBMetastore(AbstractMetastore):
                 if field == "schema":
                     dataset_version.update(**{field: DatasetRecord.parse_schema(value)})
                     values[field] = json.dumps(value) if value else None
+                elif field == "feature_schema":
+                    values[field] = json.dumps(value) if value else None
                 elif field == "preview" and isinstance(value, list):
                     values[field] = json.dumps(value, cls=JSONSerialize)
                 else:

{datachain-0.1.10 → datachain-0.1.12}/src/datachain/dataset.py

@@ -157,7 +157,7 @@ class DatasetVersion:
     dataset_id: int
     version: int
     status: int
-    feature_schema: Optional[str]
+    feature_schema: dict
     created_at: datetime
     finished_at: Optional[datetime]
     error_message: str
@@ -199,7 +199,7 @@ class DatasetVersion:
             dataset_id,
             version,
             status,
-            feature_schema,
+            json.loads(feature_schema) if feature_schema else {},
             created_at,
             finished_at,
             error_message,
@@ -263,9 +263,9 @@ class DatasetRecord:
     labels: list[str]
     shadow: bool
     schema: dict[str, Union[SQLType, type[SQLType]]]
+    feature_schema: dict
     versions: list[DatasetVersion]
     status: int = DatasetStatus.CREATED
-    feature_schema: Optional[dict] = None
     created_at: Optional[datetime] = None
     finished_at: Optional[datetime] = None
     error_message: str = ""
@@ -320,8 +320,6 @@ class DatasetRecord:
         version_job_id: Optional[str] = None,
         version_is_job_result: bool = False,
     ) -> "DatasetRecord":
-        fr_schema = json.loads(feature_schema) if feature_schema else {}
-
         labels_lst: list[str] = json.loads(labels) if labels else []
         schema_dct: dict[str, Any] = json.loads(schema) if schema else {}
         version_schema_dct: dict[str, str] = (
@@ -333,7 +331,7 @@ class DatasetRecord:
             version_dataset_id,
             version,
             version_status,
-            fr_schema,
+            version_feature_schema,
             version_created_at,
             version_finished_at,
             version_error_message,
@@ -356,9 +354,9 @@ class DatasetRecord:
             labels_lst,
             bool(shadow),
             cls.parse_schema(schema_dct),  # type: ignore[arg-type]
+            json.loads(feature_schema) if feature_schema else {},
             [dataset_version],
             status,
-            fr_schema,
             created_at,
             finished_at,
             error_message,

{datachain-0.1.10 → datachain-0.1.12}/src/datachain/lib/dc.py

@@ -6,6 +6,7 @@ import sqlalchemy
 from datachain.lib.feature import Feature, FeatureType
 from datachain.lib.feature_utils import features_to_tuples
 from datachain.lib.file import File, get_file
+from datachain.lib.meta_formats import read_meta
 from datachain.lib.settings import Settings
 from datachain.lib.signal_schema import SignalSchema
 from datachain.lib.udf import (
@@ -219,6 +220,89 @@ class DataChain(DatasetQuery):
         """
         return DataChain(name=name, version=version)
 
+    @classmethod
+    def from_csv(
+        cls,
+        path,
+        type: Literal["binary", "text", "image"] = "text",
+        anon: bool = False,
+        spec: Optional[FeatureType] = None,
+        schema_from: Optional[str] = "auto",
+        show_schema: Optional[bool] = False,
+    ) -> "DataChain":
+        """Get data from CSV. It returns the chain itself.
+
+        Parameters
+        ----------
+        path : storage URI with directory. URI must start with storage prefix such
+               as `s3://`, `gs://`, `az://` or "file:///"
+        type : read file as "binary", "text", or "image" data. Default is "binary".
+        anon : use anonymous mode to access the storage.
+        spec : optional Data Model
+        schema_from : path to sample to infer spec from
+        show_schema : print auto-generated schema
+
+        Examples
+        --------
+
+        >>> chain = DataChain.from_csv("gs://csv")
+        """
+        if schema_from == "auto":
+            schema_from = path
+
+        chain = DataChain.from_storage(path=path, type=type, anon=anon)
+        return chain.gen(
+            csv=read_meta(
+                schema_from=schema_from,
+                meta_type="csv",
+                spec=spec,
+                show_schema=show_schema,
+            )
+        )
+
+    @classmethod
+    def from_json(
+        cls,
+        path,
+        type: Literal["binary", "text", "image"] = "text",
+        anon: bool = False,
+        spec: Optional[FeatureType] = None,
+        schema_from: Optional[str] = "auto",
+        jmespath: Optional[str] = None,
+        show_schema: Optional[bool] = False,
+    ) -> "DataChain":
+        """Get data from CSV. It returns the chain itself.
+
+        Parameters
+        ----------
+        path : storage URI with directory. URI must start with storage prefix such
+               as `s3://`, `gs://`, `az://` or "file:///"
+        type : read file as "binary", "text", or "image" data. Default is "binary".
+        anon : use anonymous mode to access the storage.
+        spec : optional Data Model
+        schema_from : path to sample to infer spec from
+        show_schema : print auto-generated schema
+        jmespath : JMESPATH expression to reduce JSON
+        name : return object name
+        Examples
+        --------
+
+        >>> chain = DataChain.from_json("gs://json")
+        """
+        if schema_from == "auto":
+            schema_from = path
+
+        chain = DataChain.from_storage(path=path, type=type, anon=anon)
+        return chain.gen(
+            json=read_meta(
+                schema_from=schema_from,
+                meta_type="json",
+                spec=spec,
+                show_schema=show_schema,
+                jmespath=jmespath,
+            )
+        )
+
     def save(  # type: ignore[override]
         self, name: Optional[str] = None, version: Optional[int] = None
     ) -> "DataChain":
@@ -408,16 +492,43 @@ class DataChain(DatasetQuery):
         chain.signals_schema = new_schema
         return chain
 
-    def get_values(self) -> Iterator[Sequence]:
-        """Iterate over rows, getting feature values and applying reader calls."""
-        for features in self.iterate():
-            yield [fr.get_value() if isinstance(fr, Feature) else fr for fr in features]
+    def get_values(self, *cols: str) -> Iterator[list]:
+        """Iterate over rows, getting feature values and applying reader calls.
+        If columns are specified - limit them to specified columns.
+        """
+        for features in self.iterate(*cols):
+            yield [fr.get_value() if isinstance(fr, Feature) else fr for fr in features]  # type: ignore[union-attr,call-arg]
+
+    def get_one_value(self, col: str) -> Iterator:
+        for item in self.get_values(col):
+            yield item[0]
 
-    def iterate(self) -> Iterator[Sequence[Feature]]:
-        db_signals = self.signals_schema.db_signals()
+    def iterate(self, *cols: str) -> Iterator[list[FeatureType]]:
+        """Iterate over rows. If columns are specified - limit them to specified
+        columns.
+        """
+        chain = self.select(*cols) if cols else self
+
+        db_signals = chain.signals_schema.db_signals()
         with super().select(*db_signals).as_iterable() as rows_iter:
             for row in rows_iter:
-                yield self.signals_schema.row_to_features(row, self.session.catalog)
+                yield chain.signals_schema.row_to_features(row, chain.session.catalog)
+
+    def iterate_one(self, col: str) -> Iterator[FeatureType]:
+        for item in self.iterate(col):
+            yield item[0]
+
+    def collect(self, *cols: str) -> list[list[FeatureType]]:
+        return list(self.iterate(*cols))
+
+    def collect_one(self, col: str) -> list[FeatureType]:
+        return list(self.iterate_one(col))
+
+    def collect_values(self, *cols: str) -> list[list]:
+        return list(self.get_values(*cols))
+
+    def collect_one_value(self, col: str) -> list:
+        return list(self.get_one_value(col))
 
     def to_pytorch(self, **kwargs):
         """Convert to pytorch dataset format."""
@@ -607,3 +718,35 @@ class DataChain(DatasetQuery):
 
     def max(self, fr: FeatureType):  # type: ignore[override]
         return self._extend_features("max", fr)
+
+    @detach
+    def gen_random(self) -> "DataChain":
+        from random import getrandbits
+
+        from datachain.data_storage.warehouse import RANDOM_BITS
+
+        if "random" not in self.signals_schema.values:
+            chain = self.map(random=lambda: getrandbits(RANDOM_BITS), output=int).save()
+            return chain.select_except("random")
+
+        return self
+
+    @detach
+    def shuffle(self) -> "DataChain":
+        """Return results in deterministic random order."""
+        chain = self.gen_random()
+        return DatasetQuery.shuffle(chain)
+
+    @detach
+    def chunk(self, index: int, total: int) -> "DataChain":
+        """Split a query into smaller chunks for e.g. parallelization.
+        Example:
+            >>> dc = DataChain(...)
+            >>> chunk_1 = dc._chunk(0, 2)
+            >>> chunk_2 = dc._chunk(1, 2)
+        Note:
+            Bear in mind that `index` is 0-indexed but `total` isn't.
+            Use 0/3, 1/3 and 2/3, not 1/3, 2/3 and 3/3.
+        """
+        chain = self.gen_random()
+        return DatasetQuery.chunk(chain, index, total)

{datachain-0.1.10 → datachain-0.1.12}/src/datachain/lib/feature.py

@@ -78,16 +78,6 @@ DATACHAIN_TO_TYPE = {
     JSON: dict,
 }
 
-NAMES_TO_TYPES = {
-    "int": int,
-    "str": str,
-    "float": float,
-    "bool": bool,
-    "list": list,
-    "dict": dict,
-    "bytes": bytes,
-    "datetime": datetime,
-}
 
 NUMPY_TO_DATACHAIN = {
     np.dtype("int8"): Int,