datachain 0.12.0__tar.gz → 0.13.0__tar.gz

{datachain-0.12.0 → datachain-0.13.0}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.9.10'
+    rev: 'v0.11.0'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.12.0 → datachain-0.13.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: datachain
-Version: 0.12.0
+Version: 0.13.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/catalog/catalog.py

@@ -777,6 +777,8 @@ class Catalog:
         validate_version: Optional[bool] = True,
         listing: Optional[bool] = False,
         uuid: Optional[str] = None,
+        description: Optional[str] = None,
+        labels: Optional[list[str]] = None,
     ) -> "DatasetRecord":
         """
         Creates new dataset of a specific version.
@@ -803,6 +805,8 @@ class Catalog:
                 query_script=query_script,
                 schema=schema,
                 ignore_if_exists=True,
+                description=description,
+                labels=labels,
             )
 
         version = version or default_version
@@ -1608,7 +1612,7 @@ class Catalog:
             except TerminationSignal as exc:
                 signal.signal(signal.SIGTERM, orig_sigterm_handler)
                 signal.signal(signal.SIGINT, orig_sigint_handler)
-                logging.info("Shutting down process %s, received %r", proc.pid, exc)
+                logger.info("Shutting down process %s, received %r", proc.pid, exc)
                 # Rather than forwarding the signal to the child, we try to shut it down
                 # gracefully. This is because we consider the script to be interactive
                 # and special, so we give it time to cleanup before exiting.
@@ -1623,7 +1627,7 @@ class Catalog:
                 if thread:
                     thread.join()  # wait for the reader thread
 
-        logging.info("Process %s exited with return code %s", proc.pid, proc.returncode)
+        logger.info("Process %s exited with return code %s", proc.pid, proc.returncode)
         if proc.returncode == QUERY_SCRIPT_CANCELED_EXIT_CODE:
             raise QueryScriptCancelError(
                 "Query script was canceled by user",

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/cli/commands/ls.py

@@ -38,11 +38,12 @@ def ls_local(
 ):
     from datachain import DataChain
 
-    if catalog is None:
-        from datachain.catalog import get_catalog
-
-        catalog = get_catalog(client_config=client_config)
     if sources:
+        if catalog is None:
+            from datachain.catalog import get_catalog
+
+            catalog = get_catalog(client_config=client_config)
+
         actual_sources = list(ls_urls(sources, catalog=catalog, long=long, **kwargs))
         if len(actual_sources) == 1:
             for _, entries in actual_sources:
@@ -61,8 +62,9 @@ def ls_local(
                 for entry in entries:
                     print(format_ls_entry(entry))
     else:
-        chain = DataChain.listings()
-        for ls in chain.collect("listing"):
+        # Collect results in a list here to prevent interference from `tqdm` and `print`
+        listing = list(DataChain.listings().collect("listing"))
+        for ls in listing:
             print(format_ls_entry(f"{ls.uri}@v{ls.version}"))  # type: ignore[union-attr]
 
 

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/cli/commands/show.py

@@ -40,6 +40,13 @@ def show(
         .offset(offset)
     )
     records = query.to_db_records()
+    print("Name: ", name)
+    if dataset.description:
+        print("Description: ", dataset.description)
+    if dataset.labels:
+        print("Labels: ", ",".join(dataset.labels))
+    print("\n")
+
     show_records(records, collapse_columns=not no_collapse, hidden_fields=hidden_fields)
 
     if schema and dataset_version.feature_schema:

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/client/gcs.py

@@ -30,7 +30,7 @@ class GCSClient(Client):
         if kwargs.pop("anon", False):
             kwargs["token"] = "anon"  # noqa: S105
 
-        return cast(GCSFileSystem, super().create_fs(**kwargs))
+        return cast("GCSFileSystem", super().create_fs(**kwargs))
 
     def url(self, path: str, expires: int = 3600, **kwargs) -> str:
         """

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/client/s3.py

@@ -55,7 +55,7 @@ class ClientS3(Client):
             except NotImplementedError:
                 pass
 
-        return cast(S3FileSystem, super().create_fs(**kwargs))
+        return cast("S3FileSystem", super().create_fs(**kwargs))
 
     def url(self, path: str, expires: int = 3600, **kwargs) -> str:
         """

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/data_storage/metastore.py

@@ -119,6 +119,8 @@ class AbstractMetastore(ABC, Serializable):
         query_script: str = "",
         schema: Optional[dict[str, Any]] = None,
         ignore_if_exists: bool = False,
+        description: Optional[str] = None,
+        labels: Optional[list[str]] = None,
     ) -> DatasetRecord:
         """Creates new dataset."""
 
@@ -518,6 +520,8 @@ class AbstractDBMetastore(AbstractMetastore):
         query_script: str = "",
         schema: Optional[dict[str, Any]] = None,
         ignore_if_exists: bool = False,
+        description: Optional[str] = None,
+        labels: Optional[list[str]] = None,
         **kwargs,  # TODO registered = True / False
     ) -> DatasetRecord:
         """Creates new dataset."""
@@ -533,6 +537,8 @@ class AbstractDBMetastore(AbstractMetastore):
             sources="\n".join(sources) if sources else "",
             query_script=query_script,
             schema=json.dumps(schema or {}),
+            description=description,
+            labels=json.dumps(labels or []),
         )
         if ignore_if_exists and hasattr(query, "on_conflict_do_nothing"):
             # SQLite and PostgreSQL both support 'on_conflict_do_nothing',

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/lib/dc.py

@@ -6,6 +6,7 @@ import sys
 from collections.abc import Iterator, Sequence
 from functools import wraps
 from typing import (
+    IO,
     TYPE_CHECKING,
     Any,
     BinaryIO,
@@ -270,6 +271,18 @@ class DataChain:
         self._setup: dict = setup or {}
         self._sys = _sys
 
+    def __repr__(self) -> str:
+        """Return a string representation of the chain."""
+        classname = self.__class__.__name__
+        if not self._effective_signals_schema.values:
+            return f"Empty {classname}"
+
+        import io
+
+        file = io.StringIO()
+        self.print_schema(file=file)
+        return file.getvalue()
+
     @property
     def schema(self) -> dict[str, DataType]:
         """Get schema of the chain."""
@@ -323,9 +336,9 @@ class DataChain:
         """Return `self.union(other)`."""
         return self.union(other)
 
-    def print_schema(self) -> None:
+    def print_schema(self, file: Optional[IO] = None) -> None:
         """Print schema of the chain."""
-        self._effective_signals_schema.print_tree()
+        self._effective_signals_schema.print_tree(file=file)
 
     def clone(self) -> "Self":
         """Make a copy of the chain in a new table."""
@@ -629,7 +642,8 @@ class DataChain:
                 model_name=model_name,
                 jmespath=jmespath,
                 nrows=nrows,
-            )
+            ),
+            "params": {"file": File},
         }
         # disable prefetch if nrows is set
         settings = {"prefetch": 0} if nrows else {}
@@ -773,7 +787,12 @@ class DataChain:
         )
 
     def save(  # type: ignore[override]
-        self, name: Optional[str] = None, version: Optional[int] = None, **kwargs
+        self,
+        name: Optional[str] = None,
+        version: Optional[int] = None,
+        description: Optional[str] = None,
+        labels: Optional[list[str]] = None,
+        **kwargs,
     ) -> "Self":
         """Save to a Dataset. It returns the chain itself.
 
@@ -781,11 +800,18 @@ class DataChain:
             name : dataset name. Empty name saves to a temporary dataset that will be
                 removed after process ends. Temp dataset are useful for optimization.
             version : version of a dataset. Default - the last version that exist.
+            description : description of a dataset.
+            labels : labels of a dataset.
         """
         schema = self.signals_schema.clone_without_sys_signals().serialize()
         return self._evolve(
             query=self._query.save(
-                name=name, version=version, feature_schema=schema, **kwargs
+                name=name,
+                version=version,
+                description=description,
+                labels=labels,
+                feature_schema=schema,
+                **kwargs,
             )
         )
 
@@ -1003,8 +1029,9 @@ class DataChain:
         func: Optional[Union[Callable, UDFObjT]],
         params: Union[None, str, Sequence[str]],
         output: OutputType,
-        signal_map,
+        signal_map: dict[str, Callable],
     ) -> UDFObjT:
+        is_batch = target_class.is_input_batched
         is_generator = target_class.is_output_batched
         name = self.name or ""
 
@@ -1015,7 +1042,9 @@ class DataChain:
         if self._sys:
             signals_schema = SignalSchema({"sys": Sys}) | signals_schema
 
-        params_schema = signals_schema.slice(sign.params, self._setup)
+        params_schema = signals_schema.slice(
+            sign.params, self._setup, is_batch=is_batch
+        )
 
         return target_class._create(sign, params_schema)
 

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/lib/file.py

@@ -193,7 +193,14 @@ class File(DataModel):
         "last_modified": DateTime,
         "location": JSON,
     }
-    _hidden_fields: ClassVar[list[str]] = ["version", "source"]
+    _hidden_fields: ClassVar[list[str]] = [
+        "source",
+        "version",
+        "etag",
+        "is_latest",
+        "last_modified",
+        "location",
+    ]
 
     _unique_id_keys: ClassVar[list[str]] = [
         "source",

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/lib/meta_formats.py

@@ -10,7 +10,7 @@ import jmespath as jsp
 from pydantic import BaseModel, ConfigDict, Field, ValidationError  # noqa: F401
 
 from datachain.lib.data_model import DataModel  # noqa: F401
-from datachain.lib.file import File
+from datachain.lib.file import TextFile
 
 
 class UserModel(BaseModel):
@@ -130,7 +130,7 @@ def read_meta(  # noqa: C901
     #
 
     def parse_data(
-        file: File,
+        file: TextFile,
         data_model=spec,
         format=format,
         jmespath=jmespath,

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/lib/signal_schema.py

@@ -5,6 +5,7 @@ from dataclasses import dataclass
 from datetime import datetime
 from inspect import isclass
 from typing import (  # noqa: UP035
+    IO,
     TYPE_CHECKING,
     Annotated,
     Any,
@@ -154,9 +155,9 @@ class SignalSchema:
             if not callable(func):
                 raise SetupError(key, "value must be function or callable class")
 
-    def _init_setup_values(self):
+    def _init_setup_values(self) -> None:
         if self.setup_values is not None:
-            return self.setup_values
+            return
 
         res = {}
         for key, func in self.setup_func.items():
@@ -398,7 +399,7 @@ class SignalSchema:
         return SignalSchema(signals)
 
     @staticmethod
-    def get_flatten_hidden_fields(schema):
+    def get_flatten_hidden_fields(schema: dict):
         custom_types = schema.get("_custom_types", {})
         if not custom_types:
             return []
@@ -464,19 +465,61 @@ class SignalSchema:
         return False
 
     def slice(
-        self, keys: Sequence[str], setup: Optional[dict[str, Callable]] = None
+        self,
+        params: dict[str, Union[DataType, Any]],
+        setup: Optional[dict[str, Callable]] = None,
+        is_batch: bool = False,
     ) -> "SignalSchema":
-        # Make new schema that combines current schema and setup signals
-        setup = setup or {}
-        setup_no_types = dict.fromkeys(setup.keys(), str)
-        union = SignalSchema(self.values | setup_no_types)
-        # Slice combined schema by keys
-        schema = {}
-        for k in keys:
-            try:
-                schema[k] = union._find_in_tree(k.split("."))
-            except SignalResolvingError:
-                pass
+        """
+        Returns new schema that combines current schema and setup signals.
+        """
+        setup_params = setup.keys() if setup else []
+        schema: dict[str, DataType] = {}
+
+        for param, param_type in params.items():
+            # This is special case for setup params, they are always treated as strings
+            if param in setup_params:
+                schema[param] = str
+                continue
+
+            schema_type = self._find_in_tree(param.split("."))
+
+            if param_type is Any:
+                schema[param] = schema_type
+                continue
+
+            schema_origin = get_origin(schema_type)
+            param_origin = get_origin(param_type)
+
+            if schema_origin is Union and type(None) in get_args(schema_type):
+                schema_type = get_args(schema_type)[0]
+                if param_origin is Union and type(None) in get_args(param_type):
+                    param_type = get_args(param_type)[0]
+
+            if is_batch:
+                if param_type is list:
+                    schema[param] = schema_type
+                    continue
+
+                if param_origin is not list:
+                    raise SignalResolvingError(param.split("."), "is not a list")
+
+                param_type = get_args(param_type)[0]
+
+            if param_type == schema_type or (
+                isclass(param_type)
+                and isclass(schema_type)
+                and issubclass(param_type, File)
+                and issubclass(schema_type, File)
+            ):
+                schema[param] = schema_type
+                continue
+
+            raise SignalResolvingError(
+                param.split("."),
+                f"types mismatch: {param_type} != {schema_type}",
+            )
+
         return SignalSchema(schema, setup)
 
     def row_to_features(
@@ -696,16 +739,20 @@ class SignalSchema:
                     substree, new_prefix, depth + 1, include_hidden
                 )
 
-    def print_tree(self, indent: int = 4, start_at: int = 0):
+    def print_tree(self, indent: int = 2, start_at: int = 0, file: Optional[IO] = None):
         for path, type_, _, depth in self.get_flat_tree():
             total_indent = start_at + depth * indent
-            print(" " * total_indent, f"{path[-1]}:", SignalSchema._type_to_str(type_))
+            col_name = " " * total_indent + path[-1]
+            col_type = SignalSchema._type_to_str(type_)
+            print(col_name, col_type, sep=": ", file=file)
 
             if get_origin(type_) is list:
                 args = get_args(type_)
                 if len(args) > 0 and ModelStore.is_pydantic(args[0]):
                     sub_schema = SignalSchema({"* list of": args[0]})
-                    sub_schema.print_tree(indent=indent, start_at=total_indent + indent)
+                    sub_schema.print_tree(
+                        indent=indent, start_at=total_indent + indent, file=file
+                    )
 
     def get_headers_with_length(self, include_hidden: bool = True):
         paths = [

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/lib/udf.py

@@ -159,6 +159,7 @@ class UDFBase(AbstractUDF):
         ```
     """
 
+    is_input_batched = False
     is_output_batched = False
     prefetch: int = 0
 
@@ -395,6 +396,7 @@ class Mapper(UDFBase):
 class BatchMapper(UDFBase):
     """Inherit from this class to pass to `DataChain.batch_map()`."""
 
+    is_input_batched = True
     is_output_batched = True
 
     def run(
@@ -481,6 +483,7 @@ class Generator(UDFBase):
 class Aggregator(UDFBase):
     """Inherit from this class to pass to `DataChain.agg()`."""
 
+    is_input_batched = True
     is_output_batched = True
 
     def run(

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/lib/udf_signature.py

@@ -1,7 +1,7 @@
 import inspect
 from collections.abc import Generator, Iterator, Sequence
 from dataclasses import dataclass
-from typing import Callable, Union, get_args, get_origin
+from typing import Any, Callable, Union, get_args, get_origin
 
 from datachain.lib.data_model import DataType, DataTypeNames, is_chain_type
 from datachain.lib.signal_schema import SignalSchema
@@ -18,7 +18,7 @@ class UdfSignatureError(DataChainParamsError):
 @dataclass
 class UdfSignature:
     func: Union[Callable, UDFBase]
-    params: Sequence[str]
+    params: dict[str, Union[DataType, Any]]
     output_schema: SignalSchema
 
     DEFAULT_RETURN_TYPE = str
@@ -58,15 +58,23 @@ class UdfSignature:
         if not isinstance(udf_func, UDFBase) and not callable(udf_func):
             raise UdfSignatureError(chain, f"UDF '{udf_func}' is not callable")
 
-        func_params_map_sign, func_outs_sign, is_iterator = (
-            UdfSignature._func_signature(chain, udf_func)
+        func_params_map_sign, func_outs_sign, is_iterator = cls._func_signature(
+            chain, udf_func
         )
+
+        udf_params: dict[str, Union[DataType, Any]] = {}
         if params:
-            udf_params = [params] if isinstance(params, str) else params
-        elif not func_params_map_sign:
-            udf_params = []
-        else:
-            udf_params = list(func_params_map_sign.keys())
+            udf_params = (
+                {params: Any} if isinstance(params, str) else dict.fromkeys(params, Any)
+            )
+        elif func_params_map_sign:
+            udf_params = {
+                param: (
+                    param_type if param_type is not inspect.Parameter.empty else Any
+                )
+                for param, param_type in func_params_map_sign.items()
+            }
+
         if output:
             udf_output_map = UdfSignature._validate_output(
                 chain, signal_name, func, func_outs_sign, output

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/query/dataset.py

@@ -1646,6 +1646,8 @@ class DatasetQuery:
         name: Optional[str] = None,
         version: Optional[int] = None,
         feature_schema: Optional[dict] = None,
+        description: Optional[str] = None,
+        labels: Optional[list[str]] = None,
         **kwargs,
     ) -> "Self":
         """Save the query as a dataset."""
@@ -1678,6 +1680,8 @@ class DatasetQuery:
                 version=version,
                 feature_schema=feature_schema,
                 columns=columns,
+                description=description,
+                labels=labels,
                 **kwargs,
             )
             version = version or dataset.latest_version

{datachain-0.12.0 → datachain-0.13.0}/src/datachain/sql/sqlite/base.py

@@ -290,9 +290,9 @@ def adapt_datetime(val: datetime) -> str:
             val = val.astimezone(timezone.utc)
         except (OverflowError, ValueError, OSError):
             if val.year == MAXYEAR:
-                val = datetime.max
+                val = datetime.max.replace(tzinfo=timezone.utc)
             elif val.year == MINYEAR:
-                val = datetime.min
+                val = datetime.min.replace(tzinfo=timezone.utc)
             else:
                 raise
     return val.replace(tzinfo=None).isoformat(" ")

{datachain-0.12.0 → datachain-0.13.0}/src/datachain.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: datachain
-Version: 0.12.0
+Version: 0.13.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0

{datachain-0.12.0 → datachain-0.13.0}/tests/func/test_datachain.py

@@ -447,6 +447,21 @@ def test_show(capsys, test_session):
         assert f"{i} {first_name[i]}" in normalized_output
 
 
+def test_save(test_session):
+    chain = DataChain.from_values(key=["a", "b", "c"])
+    chain.save(
+        name="new_name",
+        version=1,
+        description="new description",
+        labels=["new_label", "old_label"],
+    )
+
+    ds = test_session.catalog.get_dataset("new_name")
+    assert ds.name == "new_name"
+    assert ds.description == "new description"
+    assert ds.labels == ["new_label", "old_label"]
+
+
 def test_show_nested_empty(capsys, test_session):
     files = [File(size=s, path=p) for p, s in zip(list("abcde"), range(5))]
     DataChain.from_values(file=files, session=test_session).limit(0).show()
@@ -707,7 +722,7 @@ def test_udf_parallel_boostrap(test_session_tmpfile):
             self.value = MyMapper.DEFAULT_VALUE
             self._had_teardown = False
 
-        def process(self, *args) -> int:
+        def process(self, key) -> int:
             return self.value
 
         def setup(self):

{datachain-0.12.0 → datachain-0.13.0}/tests/unit/lib/test_datachain.py

@@ -79,6 +79,55 @@ def sort_files(files):
     return sorted(files, key=lambda f: (f.path, f.size))
 
 
+def test_repr(test_session):
+    dc = DataChain.from_values(
+        sign1=features_nested, col1=["a", "b", "c"], session=test_session
+    )
+    assert (
+        repr(dc)
+        == """\
+sign1: MyNested
+  label: str
+  fr: MyFr
+    nnn: str
+    count: int
+col1: str
+"""
+    )
+
+    # datachain without any columns
+    assert repr(dc.select_except("col1", "sign1")) == "Empty DataChain"
+
+    dc = dc.map(col2=lambda col1: col1 * 2)
+    assert (
+        repr(dc)
+        == """\
+sign1: MyNested
+  label: str
+  fr: MyFr
+    nnn: str
+    count: int
+col1: str
+col2: str
+"""
+    )
+
+    dc = dc.mutate(countplusone=dc.column("sign1.fr.count") + 1)
+    assert (
+        repr(dc)
+        == """\
+sign1: MyNested
+  label: str
+  fr: MyFr
+    nnn: str
+    count: int
+col1: str
+col2: str
+countplusone: int
+"""
+    )
+
+
 def test_pandas_conversion(test_session):
     df = pd.DataFrame(DF_DATA)
     df1 = DataChain.from_pandas(df, session=test_session)

{datachain-0.12.0 → datachain-0.13.0}/tests/unit/lib/test_datachain_bootstrap.py

@@ -15,7 +15,7 @@ def test_udf():
             self.value = MyMapper.DEFAULT_VALUE
             self._had_teardown = False
 
-        def process(self, *args) -> int:
+        def process(self, key) -> int:
             return self.value
 
         def setup(self):
@@ -40,7 +40,7 @@ def test_no_bootstrap_for_callable():
             self._had_bootstrap = False
             self._had_teardown = False
 
-        def __call__(self, *args):
+        def __call__(self, key):
             return None
 
         def bootstrap(self):