datachain 0.1.10__tar.gz → 0.1.11__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.1.10
+Version: 0.1.11
 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.11}/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.11/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.11}/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.11}/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.11}/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.11'
+__version_tuple__ = version_tuple = (0, 1, 11)

{datachain-0.1.10 → datachain-0.1.11}/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.11}/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.11}/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.11}/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,7 +492,7 @@ class DataChain(DatasetQuery):
         chain.signals_schema = new_schema
         return chain
 
-    def get_values(self) -> Iterator[Sequence]:
+    def get_values(self) -> Iterator[list]:
         """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]
@@ -607,3 +691,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.11}/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,

datachain-0.1.11/src/datachain/lib/meta_formats.py

@@ -0,0 +1,164 @@
+# pip install datamodel-code-generator
+# pip install jmespath
+#
+import csv
+import io
+import json
+import subprocess
+import sys
+import uuid
+from collections.abc import Iterator
+from typing import Any, Callable
+
+import jmespath as jsp
+
+from datachain.lib.feature_utils import pydantic_to_feature  # noqa: F401
+from datachain.lib.file import File
+
+# from datachain.lib.dc import C, DataChain
+
+
+def generate_uuid():
+    return uuid.uuid4()  # Generates a random UUID.
+
+
+# JSON decoder
+def load_json_from_string(json_string):
+    try:
+        data = json.loads(json_string)
+        print("Successfully parsed JSON", file=sys.stderr)
+        return data
+    except json.JSONDecodeError:
+        print("Failed to decode JSON: The string is not formatted correctly.")
+    return None
+
+
+# Read valid JSON and return a data object sample
+def process_json(data_string, jmespath):
+    json_dict = load_json_from_string(data_string)
+    if jmespath:
+        json_dict = jsp.search(jmespath, json_dict)
+        # we allow non-list JSONs here to print the root schema
+        # but if jmespath expression is given, we assume a list
+        if not isinstance(json_dict, list):
+            raise ValueError("JMESPATH expression must resolve to a list")
+            return None
+        json_dict = json_dict[0]  # sample the first object
+    return json.dumps(json_dict)
+
+
+# Print a dynamic datamodel-codegen output from JSON or CSV on stdout
+def read_schema(source_file, data_type="csv", expr=None):
+    data_string = ""
+    uid_str = str(generate_uuid()).replace("-", "")  # comply with Python class names
+    # using uiid to get around issue #1617
+    model_name = f"Model{uid_str}"
+    try:
+        with source_file.open() as fd:  # CSV can be larger than memory
+            if data_type == "csv":
+                data_string += fd.readline().decode("utf-8", "ignore").replace("\r", "")
+                data_string += fd.readline().decode("utf-8", "ignore").replace("\r", "")
+            else:
+                data_string = fd.read()  # other meta must fit into RAM
+    except OSError as e:
+        print(f"An unexpected file error occurred: {e}")
+        return
+    if data_type == "json":
+        data_string = process_json(data_string, expr)
+    command = [
+        "datamodel-codegen",
+        "--input-file-type",
+        data_type,
+        "--class-name",
+        model_name,
+    ]
+    try:
+        result = subprocess.run(
+            command,  # noqa: S603
+            input=data_string,
+            text=True,
+            capture_output=True,
+            check=True,
+        )
+        model_output = (
+            result.stdout
+        )  # This will contain the output from datamodel-codegen
+    except subprocess.CalledProcessError as e:
+        model_output = f"An error occurred in datamodel-codegen: {e.stderr}"
+    print(f"{model_output}")
+    print("\n" + f"spec=pydantic_to_feature({model_name})" + "\n")
+
+
+#
+# UDF mapper which calls chain in the setup to infer the dynamic schema
+#
+def read_meta(
+    spec=None, schema_from=None, meta_type="json", jmespath=None, show_schema=False
+) -> Callable:
+    from datachain.lib.dc import DataChain
+
+    # ugly hack: datachain is run redirecting printed outputs to a variable
+    if schema_from:
+        captured_output = io.StringIO()
+        current_stdout = sys.stdout
+        sys.stdout = captured_output
+        try:
+            chain = (
+                DataChain.from_storage(schema_from)
+                .limit(1)
+                .map(  # dummy column created (#1615)
+                    meta_schema=lambda file: read_schema(
+                        file, data_type=meta_type, expr=jmespath
+                    ),
+                    output=str,
+                )
+            )
+            # dummy executor (#1616)
+            chain.save()
+        finally:
+            sys.stdout = current_stdout
+        model_output = captured_output.getvalue()
+        captured_output.close()
+        if show_schema:
+            print(f"{model_output}")
+        # Below 'spec' should be a dynamically converted Feature from Pydantic datamodel
+        if not spec:
+            local_vars: dict[str, Any] = {}
+            exec(model_output, globals(), local_vars)  # noqa: S102
+            spec = local_vars["spec"]
+
+    if not (spec) and not (schema_from):
+        raise ValueError(
+            "Must provide a static schema in spec: or metadata sample in schema_from:"
+        )
+
+    #
+    # UDF mapper parsing a JSON or CSV file using schema spec
+    #
+    def parse_data(
+        file: File, data_model=spec, meta_type=meta_type, jmespath=jmespath
+    ) -> Iterator[spec]:
+        if meta_type == "csv":
+            with (
+                file.open() as fd
+            ):  # TODO: if schema is statically given, should allow CSV without headers
+                reader = csv.DictReader(fd)
+                for row in reader:  # CSV can be larger than memory
+                    json_string = json.dumps(row)
+                    yield data_model.model_validate_json(json_string)
+        if meta_type == "json":
+            try:
+                with file.open() as fd:  # JSON must fit into RAM
+                    data_string = fd.read()
+            except OSError as e:
+                print(f"An unexpected file error occurred: {e}")
+            json_object = load_json_from_string(data_string)
+            if jmespath:
+                json_object = jsp.search(jmespath, json_object)
+            if not isinstance(json_object, list):
+                raise ValueError("JSON expression must resolve in a list of objects")
+            for json_dict in json_object:
+                json_string = json.dumps(json_dict)
+                yield data_model.model_validate_json(json_string)
+
+    return parse_data