datachain 0.8.11__tar.gz → 0.8.13__tar.gz

{datachain-0.8.11 → datachain-0.8.13}/PKG-INFO

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

datachain-0.8.13/docs/references/func.md

@@ -0,0 +1,5 @@
+# Functions
+
+Use built-in functions for data manipulation and analysis to operate on the underlying database storing the chain data. These functions are useful for operations like [`DataChain.filter`](datachain.md#datachain.lib.dc.DataChain.filter) and [`DataChain.mutate`](datachain.md#datachain.lib.dc.DataChain.mutate). Import these functions from `datachain.func`.
+
+::: datachain.func

{datachain-0.8.11 → datachain-0.8.13}/docs/references/index.md

@@ -10,5 +10,5 @@ DataChain's API is organized into several modules:
 - [DataType](./datatype.md) - Type system and schema definitions
 - [File](./file.md) - File handling and storage operations
 - [UDF](./udf.md) - User-defined functions and transformations
-- [SQL](./sql.md) - SQL query integration
+- [Functions](./func.md) - Built-in functions for data manipulation and analysis
 - [Torch](./torch.md) - PyTorch data loading utilities

{datachain-0.8.11 → datachain-0.8.13}/examples/get_started/common_sql_functions.py

@@ -9,7 +9,7 @@ def num_chars_udf(file):
     return ([],)
 
 
-dc = DataChain.from_storage("gs://datachain-demo/dogs-and-cats/")
+dc = DataChain.from_storage("gs://datachain-demo/dogs-and-cats/", anon=True)
 dc.map(num_chars_udf, params=["file"], output={"num_chars": list[str]}).select(
     "file.path", "num_chars"
 ).show(5)
@@ -32,6 +32,12 @@ dc.map(num_chars_udf, params=["file"], output={"num_chars": list[str]}).select(
     .show(5)
 )
 
+parts = string.split(path.name(C("file.path")), ".")
+chain = dc.mutate(
+    isdog=array.contains(parts, "dog"),
+    iscat=array.contains(parts, "cat"),
+)
+chain.select("file.path", "isdog", "iscat").show(5)
 
 chain = dc.mutate(
     a=array.length(string.split("file.path", "/")),
@@ -79,6 +85,15 @@ Processed: 400 rows [00:00, 16364.66 rows/s]
 3  dogs-and-cats/cat.10.json   cat.10  json
 4  dogs-and-cats/cat.100.jpg  cat.100   jpg
 
+[Limited by 5 rows]
+                        file isdog iscat
+                        path
+0    dogs-and-cats/cat.1.jpg     0     1
+1   dogs-and-cats/cat.1.json     0     1
+2   dogs-and-cats/cat.10.jpg     0     1
+3  dogs-and-cats/cat.10.json     0     1
+4  dogs-and-cats/cat.100.jpg     0     1
+
 [Limited by 5 rows]
 Processed: 400 rows [00:00, 16496.93 rows/s]
    a  b  greatest  least

{datachain-0.8.11 → datachain-0.8.13}/mkdocs.yml

@@ -73,7 +73,7 @@ nav:
           - File: references/file.md
           - UDF: references/udf.md
           - Torch: references/torch.md
-          - SQL: references/sql.md
+          - Functions: references/func.md
       - 🤝 Contributing: contributing.md
 
   - DataChain Website ↗: https://datachain.ai" target="_blank"

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/catalog/catalog.py

@@ -38,7 +38,6 @@ from datachain.dataset import (
     DatasetDependency,
     DatasetListRecord,
     DatasetRecord,
-    DatasetStats,
     DatasetStatus,
     StorageURI,
     create_dataset_uri,
@@ -1235,17 +1234,6 @@ class Catalog:
         dataset = self.get_dataset(name)
         return self.warehouse.dataset_table_export_file_names(dataset, version)
 
-    def dataset_stats(self, name: str, version: Optional[int]) -> DatasetStats:
-        """
-        Returns tuple with dataset stats: total number of rows and total dataset size.
-        """
-        dataset = self.get_dataset(name)
-        dataset_version = dataset.get_version(version or dataset.latest_version)
-        return DatasetStats(
-            num_objects=dataset_version.num_objects,
-            size=dataset_version.size,
-        )
-
     def remove_dataset(
         self,
         name: str,
@@ -1391,19 +1379,12 @@ class Catalog:
         except DatasetNotFoundError:
             pass
 
-        stats_response = studio_client.dataset_stats(
-            remote_ds_name, remote_ds_version.version
-        )
-        if not stats_response.ok:
-            raise_remote_error(stats_response.message)
-        ds_stats = stats_response.data
-
         dataset_save_progress_bar = tqdm(
             desc=f"Saving dataset {remote_ds_uri} locally: ",
             unit=" rows",
             unit_scale=True,
             unit_divisor=1000,
-            total=ds_stats.num_objects,  # type: ignore [union-attr]
+            total=remote_ds_version.num_objects,  # type: ignore [union-attr]
             leave=False,
         )
 

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/cli/__init__.py

@@ -11,7 +11,6 @@ from datachain.telemetry import telemetry
 from .commands import (
     clear_cache,
     completion,
-    dataset_stats,
     du,
     edit_dataset,
     garbage_collect,
@@ -182,13 +181,6 @@ def handle_dataset_command(args, catalog):
             all=args.all,
             team=args.team,
         ),
-        "stats": lambda: dataset_stats(
-            catalog,
-            args.name,
-            args.version,
-            show_bytes=args.bytes,
-            si=args.si,
-        ),
     }
 
     handler = dataset_commands.get(args.datasets_cmd)

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/cli/commands/__init__.py

@@ -1,5 +1,4 @@
 from .datasets import (
-    dataset_stats,
     edit_dataset,
     list_datasets,
     list_datasets_local,
@@ -15,7 +14,6 @@ from .show import show
 __all__ = [
     "clear_cache",
     "completion",
-    "dataset_stats",
     "du",
     "edit_dataset",
     "garbage_collect",

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/cli/commands/datasets.py

@@ -3,8 +3,6 @@ from typing import TYPE_CHECKING, Optional
 
 from tabulate import tabulate
 
-from datachain import utils
-
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
 
@@ -109,20 +107,3 @@ def edit_dataset(
 
     if (all or studio) and token:
         edit_studio_dataset(team, name, new_name, description, labels)
-
-
-def dataset_stats(
-    catalog: "Catalog",
-    name: str,
-    version: int,
-    show_bytes=False,
-    si=False,
-):
-    stats = catalog.dataset_stats(name, version)
-
-    if stats:
-        print(f"Number of objects: {stats.num_objects}")
-        if show_bytes:
-            print(f"Total objects size: {stats.size}")
-        else:
-            print(f"Total objects size: {utils.sizeof_fmt(stats.size, si=si): >7}")

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/cli/parser/__init__.py

@@ -307,31 +307,6 @@ def get_parser() -> ArgumentParser:  # noqa: PLR0915
         help="The team to delete a dataset. By default, it will use team from config",
     )
 
-    dataset_stats_parser = datasets_subparser.add_parser(
-        "stats", parents=[parent_parser], description="Show basic dataset statistics."
-    )
-    dataset_stats_parser.add_argument("name", type=str, help="Dataset name")
-    dataset_stats_parser.add_argument(
-        "--version",
-        action="store",
-        default=None,
-        type=int,
-        help="Dataset version",
-    )
-    dataset_stats_parser.add_argument(
-        "-b",
-        "--bytes",
-        default=False,
-        action="store_true",
-        help="Display size in bytes instead of human-readable size",
-    )
-    dataset_stats_parser.add_argument(
-        "--si",
-        default=False,
-        action="store_true",
-        help="Display size using powers of 1000 not 1024",
-    )
-
     parse_ls = subp.add_parser(
         "ls", parents=[parent_parser], description="List storage contents."
     )

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/dataset.py

@@ -150,12 +150,6 @@ class DatasetDependency:
         return hash(f"{self.type}_{self.name}_{self.version}")
 
 
-@dataclass
-class DatasetStats:
-    num_objects: Optional[int]  # None if table is missing
-    size: Optional[int]  # in bytes None if table is missing or empty
-
-
 class DatasetStatus:
     CREATED = 1
     PENDING = 2

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/func/__init__.py

@@ -15,7 +15,7 @@ from .aggregate import (
     row_number,
     sum,
 )
-from .array import cosine_distance, euclidean_distance, length, sip_hash_64
+from .array import contains, cosine_distance, euclidean_distance, length, sip_hash_64
 from .conditional import case, greatest, ifelse, isnone, least
 from .numeric import bit_and, bit_hamming_distance, bit_or, bit_xor, int_hash_64
 from .random import rand
@@ -34,6 +34,7 @@ __all__ = [
     "case",
     "collect",
     "concat",
+    "contains",
     "cosine_distance",
     "count",
     "dense_rank",

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/func/array.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from typing import Union
+from typing import Any, Union
 
 from datachain.sql.functions import array
 
@@ -140,6 +140,44 @@ def length(arg: Union[str, Sequence, Func]) -> Func:
     return Func("length", inner=array.length, cols=cols, args=args, result_type=int)
 
 
+def contains(arr: Union[str, Sequence, Func], elem: Any) -> Func:
+    """
+    Checks whether the `arr` array has the `elem` element.
+
+    Args:
+        arr (str | Sequence | Func): Array to check for the element.
+            If a string is provided, it is assumed to be the name of the array column.
+            If a sequence is provided, it is assumed to be an array of values.
+            If a Func is provided, it is assumed to be a function returning an array.
+        elem (Any): Element to check for in the array.
+
+    Returns:
+        Func: A Func object that represents the contains function. Result of the
+            function will be 1 if the element is present in the array, and 0 otherwise.
+
+    Example:
+        ```py
+        dc.mutate(
+            contains1=func.array.contains("signal.values", 3),
+            contains2=func.array.contains([1, 2, 3, 4, 5], 7),
+        )
+        ```
+    """
+
+    def inner(arg):
+        is_json = type(elem) in [list, dict]
+        return array.contains(arg, elem, is_json)
+
+    if isinstance(arr, (str, Func)):
+        cols = [arr]
+        args = None
+    else:
+        cols = None
+        args = [arr]
+
+    return Func("contains", inner=inner, cols=cols, args=args, result_type=int)
+
+
 def sip_hash_64(arg: Union[str, Sequence]) -> Func:
     """
     Computes the SipHash-64 hash of the array.

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/func/conditional.py

@@ -9,7 +9,7 @@ from datachain.sql.functions import conditional
 
 from .func import ColT, Func
 
-CaseT = Union[int, float, complex, bool, str, Func]
+CaseT = Union[int, float, complex, bool, str, Func, ColumnElement]
 
 
 def greatest(*args: Union[ColT, float]) -> Func:
@@ -94,11 +94,12 @@ def case(
     """
     Returns the case function that produces case expression which has a list of
     conditions and corresponding results. Results can be python primitives like string,
-    numbers or booleans but can also be other nested function (including case function).
+    numbers or booleans but can also be other nested functions (including case function)
+    or columns.
     Result type is inferred from condition results.
 
     Args:
-        args (tuple((ColumnElement, Func), (str | int | float | complex | bool, Func))):
+        args tuple((ColumnElement | Func),(str | int | float | complex | bool, Func, ColumnElement)):
             Tuple of condition and values pair.
         else_ (str | int | float | complex | bool, Func): optional else value in case
             expression. If omitted, and no case conditions are satisfied, the result
@@ -113,13 +114,16 @@ def case(
             res=func.case((C("num") > 0, "P"), (C("num") < 0, "N"), else_="Z"),
         )
         ```
-    """
+    """  # noqa: E501
     supported_types = [int, float, complex, str, bool]
 
     def _get_type(val):
         if isinstance(val, Func):
             # nested functions
             return val.result_type
+        if isinstance(val, Column):
+            # at this point we cannot know what is the type of a column
+            return None
         return type(val)
 
     if not args:
@@ -129,13 +133,16 @@ def case(
 
     for arg in args:
         arg_type = _get_type(arg[1])
+        if arg_type is None:
+            # we couldn't figure out the type of case value
+            continue
         if type_ and arg_type != type_:
             raise DataChainParamsError(
                 f"Statement values must be of the same type, got {type_} and {arg_type}"
             )
         type_ = arg_type
 
-    if type_ not in supported_types:
+    if type_ is not None and type_ not in supported_types:
         raise DataChainParamsError(
             f"Only python literals ({supported_types}) are supported for values"
         )
@@ -151,15 +158,15 @@ def ifelse(
     """
     Returns the ifelse function that produces if expression which has a condition
     and values for true and false outcome. Results can be one of python primitives
-    like string, numbers or booleans, but can also be nested functions.
+    like string, numbers or booleans, but can also be nested functions or columns.
     Result type is inferred from the values.
 
     Args:
         condition (ColumnElement, Func):  Condition which is evaluated.
-        if_val (str | int | float | complex | bool, Func): Value for true
+        if_val (str | int | float | complex | bool, Func, ColumnElement): Value for true
             condition outcome.
-        else_val (str | int | float | complex | bool, Func): Value for false condition
-            outcome.
+        else_val (str | int | float | complex | bool, Func, ColumnElement): Value for
+            false condition outcome.
 
     Returns:
         Func: A Func object that represents the ifelse function.

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/func/func.py

@@ -424,10 +424,9 @@ class Func(Function):
 
 def get_db_col_type(signals_schema: "SignalSchema", col: ColT) -> "DataType":
     if isinstance(col, tuple):
-        raise DataChainParamsError(
-            "Cannot get type from tuple, please provide type hint to the function"
-        )
-
+        # we can only get tuple from case statement where the first tuple item
+        # is condition, and second one is value which type is important
+        col = col[1]
     if isinstance(col, Func):
         return col.get_result_type(signals_schema)
 
@@ -435,7 +434,7 @@ def get_db_col_type(signals_schema: "SignalSchema", col: ColT) -> "DataType":
         return sql_to_python(col)
 
     return signals_schema.get_column_type(
-        col.name if isinstance(col, ColumnElement) else col
+        col.name if isinstance(col, ColumnElement) else col  # type: ignore[arg-type]
     )
 
 

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/remote/studio.py

@@ -16,14 +16,12 @@ from urllib.parse import urlparse, urlunparse
 import websockets
 
 from datachain.config import Config
-from datachain.dataset import DatasetStats
 from datachain.error import DataChainError
 from datachain.utils import STUDIO_URL, retry_with_backoff
 
 T = TypeVar("T")
 LsData = Optional[list[dict[str, Any]]]
 DatasetInfoData = Optional[dict[str, Any]]
-DatasetStatsData = Optional[DatasetStats]
 DatasetRowsData = Optional[Iterable[dict[str, Any]]]
 DatasetJobVersionsData = Optional[dict[str, Any]]
 DatasetExportStatus = Optional[dict[str, Any]]
@@ -309,7 +307,7 @@ class StudioClient:
             "datachain/datasets",
             {
                 "dataset_name": name,
-                "version": version,
+                "dataset_version": version,
                 "force": force,
             },
             method="DELETE",
@@ -347,16 +345,6 @@ class StudioClient:
             method="GET",
         )
 
-    def dataset_stats(self, name: str, version: int) -> Response[DatasetStatsData]:
-        response = self._send_request(
-            "datachain/datasets/stats",
-            {"dataset_name": name, "dataset_version": version},
-            method="GET",
-        )
-        if response.ok:
-            response.data = DatasetStats(**response.data)
-        return response
-
     def export_dataset_table(
         self, name: str, version: int
     ) -> Response[DatasetExportSignedUrls]:

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/sql/functions/array.py

@@ -1,6 +1,6 @@
 from sqlalchemy.sql.functions import GenericFunction
 
-from datachain.sql.types import Float, Int64
+from datachain.sql.types import Boolean, Float, Int64
 from datachain.sql.utils import compiler_not_implemented
 
 
@@ -37,6 +37,17 @@ class length(GenericFunction):  # noqa: N801
     inherit_cache = True
 
 
+class contains(GenericFunction):  # noqa: N801
+    """
+    Checks if element is in the array.
+    """
+
+    type = Boolean()
+    package = "array"
+    name = "contains"
+    inherit_cache = True
+
+
 class sip_hash_64(GenericFunction):  # noqa: N801
     """
     Computes the SipHash-64 hash of the array.
@@ -51,4 +62,5 @@ class sip_hash_64(GenericFunction):  # noqa: N801
 compiler_not_implemented(cosine_distance)
 compiler_not_implemented(euclidean_distance)
 compiler_not_implemented(length)
+compiler_not_implemented(contains)
 compiler_not_implemented(sip_hash_64)

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/sql/sqlite/base.py

@@ -87,6 +87,7 @@ def setup():
     compiles(sql_path.file_stem, "sqlite")(compile_path_file_stem)
     compiles(sql_path.file_ext, "sqlite")(compile_path_file_ext)
     compiles(array.length, "sqlite")(compile_array_length)
+    compiles(array.contains, "sqlite")(compile_array_contains)
     compiles(string.length, "sqlite")(compile_string_length)
     compiles(string.split, "sqlite")(compile_string_split)
     compiles(string.regexp_replace, "sqlite")(compile_string_regexp_replace)
@@ -269,13 +270,16 @@ def register_user_defined_sql_functions() -> None:
 
     _registered_function_creators["string_functions"] = create_string_functions
 
-    has_json_extension = functions_exist(["json_array_length"])
+    has_json_extension = functions_exist(["json_array_length", "json_array_contains"])
     if not has_json_extension:
 
         def create_json_functions(conn):
             conn.create_function(
                 "json_array_length", 1, py_json_array_length, deterministic=True
             )
+            conn.create_function(
+                "json_array_contains", 3, py_json_array_contains, deterministic=True
+            )
 
         _registered_function_creators["json_functions"] = create_json_functions
 
@@ -428,10 +432,22 @@ def py_json_array_length(arr):
     return len(orjson.loads(arr))
 
 
+def py_json_array_contains(arr, value, is_json):
+    if is_json:
+        value = orjson.loads(value)
+    return value in orjson.loads(arr)
+
+
 def compile_array_length(element, compiler, **kwargs):
     return compiler.process(func.json_array_length(*element.clauses.clauses), **kwargs)
 
 
+def compile_array_contains(element, compiler, **kwargs):
+    return compiler.process(
+        func.json_array_contains(*element.clauses.clauses), **kwargs
+    )
+
+
 def compile_string_length(element, compiler, **kwargs):
     return compiler.process(func.length(*element.clauses.clauses), **kwargs)
 

{datachain-0.8.11 → datachain-0.8.13}/src/datachain/sql/sqlite/types.py

@@ -31,6 +31,10 @@ def adapt_array(arr):
     return orjson.dumps(arr).decode("utf-8")
 
 
+def adapt_dict(dct):
+    return orjson.dumps(dct).decode("utf-8")
+
+
 def convert_array(arr):
     return orjson.loads(arr)
 
@@ -52,6 +56,7 @@ def adapt_np_generic(val):
 
 def register_type_converters():
     sqlite3.register_adapter(list, adapt_array)
+    sqlite3.register_adapter(dict, adapt_dict)
     sqlite3.register_converter("ARRAY", convert_array)
     if numpy_imported:
         sqlite3.register_adapter(np.ndarray, adapt_np_array)

{datachain-0.8.11 → datachain-0.8.13}/src/datachain.egg-info/PKG-INFO

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

{datachain-0.8.11 → datachain-0.8.13}/src/datachain.egg-info/SOURCES.txt

@@ -31,8 +31,8 @@ docs/overrides/main.html
 docs/references/datachain.md
 docs/references/datatype.md
 docs/references/file.md
+docs/references/func.md
 docs/references/index.md
-docs/references/sql.md
 docs/references/torch.md
 docs/references/udf.md
 examples/computer_vision/iptc_exif_xmp_lib.py

{datachain-0.8.11 → datachain-0.8.13}/tests/func/test_catalog.py

@@ -17,7 +17,8 @@ from tests.utils import DEFAULT_TREE, skip_if_not_sqlite, tree_from_path
 def listing_stats(uri, catalog):
     list_dataset_name, _, _ = parse_listing_uri(uri, catalog.client_config)
     dataset = catalog.get_dataset(list_dataset_name)
-    return catalog.dataset_stats(dataset.name, dataset.latest_version)
+    dataset_version = dataset.get_version(dataset.latest_version)
+    return dataset_version.num_objects, dataset_version.size
 
 
 @pytest.fixture
@@ -582,23 +583,23 @@ def test_listing_stats(cloud_test_catalog):
         listing_stats(src_uri, catalog)
 
     catalog.enlist_source(src_uri)
-    stats = listing_stats(src_uri, catalog)
-    assert stats.num_objects == 7
-    assert stats.size == 36
+    num_objects, size = listing_stats(src_uri, catalog)
+    assert num_objects == 7
+    assert size == 36
 
     catalog.enlist_source(f"{src_uri}/dogs/", update=True)
-    stats = listing_stats(src_uri, catalog)
-    assert stats.num_objects == 7
-    assert stats.size == 36
+    num_objects, size = listing_stats(src_uri, catalog)
+    assert num_objects == 7
+    assert size == 36
 
-    stats = listing_stats(f"{src_uri}/dogs/", catalog)
-    assert stats.num_objects == 4
-    assert stats.size == 15
+    num_objects, size = listing_stats(f"{src_uri}/dogs/", catalog)
+    assert num_objects == 4
+    assert size == 15
 
     catalog.enlist_source(f"{src_uri}/dogs/")
-    stats = listing_stats(src_uri, catalog)
-    assert stats.num_objects == 7
-    assert stats.size == 36
+    num_objects, size = listing_stats(src_uri, catalog)
+    assert num_objects == 7
+    assert size == 36
 
 
 @pytest.mark.parametrize("cloud_type", ["s3", "azure", "gs"], indirect=True)
@@ -608,15 +609,15 @@ def test_enlist_source_handles_slash(cloud_test_catalog):
     src_path = f"{src_uri}/dogs"
 
     catalog.enlist_source(src_path)
-    stats = listing_stats(src_path, catalog)
-    assert stats.num_objects == len(DEFAULT_TREE["dogs"])
-    assert stats.size == 15
+    num_objects, size = listing_stats(src_path, catalog)
+    assert num_objects == len(DEFAULT_TREE["dogs"])
+    assert size == 15
 
     src_path = f"{src_uri}/dogs"
     catalog.enlist_source(src_path, update=True)
-    stats = listing_stats(src_path, catalog)
-    assert stats.num_objects == len(DEFAULT_TREE["dogs"])
-    assert stats.size == 15
+    num_objects, size = listing_stats(src_path, catalog)
+    assert num_objects == len(DEFAULT_TREE["dogs"])
+    assert size == 15
 
 
 @pytest.mark.parametrize("cloud_type", ["s3", "azure", "gs"], indirect=True)
@@ -626,10 +627,10 @@ def test_enlist_source_handles_glob(cloud_test_catalog):
     src_path = f"{src_uri}/dogs/*.jpg"
 
     catalog.enlist_source(src_path)
-    stats = listing_stats(src_path, catalog)
+    num_objects, size = listing_stats(src_path, catalog)
 
-    assert stats.num_objects == len(DEFAULT_TREE["dogs"])
-    assert stats.size == 15
+    assert num_objects == len(DEFAULT_TREE["dogs"])
+    assert size == 15
 
 
 @pytest.mark.parametrize("cloud_type", ["s3", "azure", "gs"], indirect=True)

{datachain-0.8.11 → datachain-0.8.13}/tests/func/test_datachain.py

@@ -20,7 +20,7 @@ from sqlalchemy import Column
 from datachain import DataModel, func
 from datachain.catalog.catalog import QUERY_SCRIPT_CANCELED_EXIT_CODE
 from datachain.data_storage.sqlite import SQLiteWarehouse
-from datachain.dataset import DatasetDependencyType, DatasetStats
+from datachain.dataset import DatasetDependencyType
 from datachain.func import path as pathfunc
 from datachain.lib.dc import C, DataChain
 from datachain.lib.file import File, ImageFile
@@ -515,8 +515,9 @@ def test_from_storage_dataset_stats(tmp_dir, test_session):
     dc = DataChain.from_storage(tmp_dir.as_uri(), session=test_session).save(
         "test-data"
     )
-    stats = test_session.catalog.dataset_stats(dc.name, dc.version)
-    assert stats == DatasetStats(num_objects=4, size=20)
+    version = test_session.catalog.get_dataset(dc.name).get_version(dc.version)
+    assert version.num_objects == 4
+    assert version.size == 20
 
 
 def test_from_storage_check_rows(tmp_dir, test_session):