datachain 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

datachain/catalog/catalog.py

@@ -676,7 +676,7 @@ class Catalog:
 
     def parse_url(self, uri: str, **config: Any) -> tuple[Client, str]:
         config = config or self.client_config
-        return Client.parse_url(uri, self.metastore, self.cache, **config)
+        return Client.parse_url(uri, self.cache, **config)
 
     def get_client(self, uri: StorageURI, **config: Any) -> Client:
         """
@@ -1627,8 +1627,17 @@ class Catalog:
         version = self.get_dataset(dataset_name).get_version(dataset_version)
 
         file_signals_values = {}
+        file_schemas = {}
+        # TODO: To remove after we properly fix deserialization
+        for signal, type_name in version.feature_schema.items():
+            from datachain.lib.model_store import ModelStore
 
-        schema = SignalSchema.deserialize(version.feature_schema)
+            type_name_parsed, v = ModelStore.parse_name_version(type_name)
+            fr = ModelStore.get(type_name_parsed, v)
+            if fr and issubclass(fr, File):
+                file_schemas[signal] = type_name
+
+        schema = SignalSchema.deserialize(file_schemas)
         for file_signals in schema.get_signals(File):
             prefix = file_signals.replace(".", DEFAULT_DELIMITER) + DEFAULT_DELIMITER
             file_signals_values[file_signals] = {

datachain/client/fsspec.py

@@ -37,7 +37,6 @@ from datachain.storage import StorageURI
 if TYPE_CHECKING:
     from fsspec.spec import AbstractFileSystem
 
-    from datachain.data_storage import AbstractMetastore
 
 logger = logging.getLogger("datachain")
 
@@ -116,13 +115,12 @@ class Client(ABC):
     @staticmethod
     def parse_url(
         source: str,
-        metastore: "AbstractMetastore",
         cache: DataChainCache,
         **kwargs,
     ) -> tuple["Client", str]:
         cls = Client.get_implementation(source)
         storage_url, rel_path = cls.split_url(source)
-        client = cls.from_name(storage_url, metastore, cache, kwargs)
+        client = cls.from_name(storage_url, cache, kwargs)
         return client, rel_path
 
     @classmethod
@@ -136,7 +134,6 @@ class Client(ABC):
     def from_name(
         cls,
         name: str,
-        metastore: "AbstractMetastore",
         cache: DataChainCache,
         kwargs: dict[str, Any],
     ) -> "Client":

datachain/client/local.py

@@ -2,7 +2,7 @@ import os
 import posixpath
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import Any
 from urllib.parse import urlparse
 
 from fsspec.implementations.local import LocalFileSystem
@@ -12,9 +12,6 @@ from datachain.storage import StorageURI
 
 from .fsspec import Client
 
-if TYPE_CHECKING:
-    from datachain.data_storage import AbstractMetastore
-
 
 class FileClient(Client):
     FS_CLASS = LocalFileSystem
@@ -97,9 +94,7 @@ class FileClient(Client):
         return cls.root_dir(), uri.removeprefix(cls.root_path().as_uri())
 
     @classmethod
-    def from_name(
-        cls, name: str, metastore: "AbstractMetastore", cache, kwargs
-    ) -> "FileClient":
+    def from_name(cls, name: str, cache, kwargs) -> "FileClient":
         use_symlinks = kwargs.pop("use_symlinks", False)
         return cls(name, kwargs, cache, use_symlinks=use_symlinks)
 

datachain/data_storage/schema.py

@@ -67,7 +67,11 @@ def convert_rows_custom_column_types(
     for row in rows:
         row_list = list(row)
         for idx, t in custom_columns_types:
-            row_list[idx] = t.on_read_convert(row_list[idx], dialect)
+            row_list[idx] = (
+                t.default_value(dialect)
+                if row_list[idx] is None
+                else t.on_read_convert(row_list[idx], dialect)
+            )
 
         yield tuple(row_list)
 
@@ -136,7 +140,15 @@ class DataTable:
         self.column_types: dict[str, SQLType] = column_types or {}
 
     @staticmethod
-    def copy_column(column: sa.Column):
+    def copy_column(
+        column: sa.Column,
+        primary_key: Optional[bool] = None,
+        index: Optional[bool] = None,
+        nullable: Optional[bool] = None,
+        default: Optional[Any] = None,
+        server_default: Optional[Any] = None,
+        unique: Optional[bool] = None,
+    ) -> sa.Column:
         """
         Copy a sqlalchemy Column object intended for use as a signal column.
 
@@ -150,12 +162,14 @@ class DataTable:
         return sa.Column(
             column.name,
             column.type,
-            primary_key=column.primary_key,
-            index=column.index,
-            nullable=column.nullable,
-            default=column.default,
-            server_default=column.server_default,
-            unique=column.unique,
+            primary_key=primary_key if primary_key is not None else column.primary_key,
+            index=index if index is not None else column.index,
+            nullable=nullable if nullable is not None else column.nullable,
+            default=default if default is not None else column.default,
+            server_default=(
+                server_default if server_default is not None else column.server_default
+            ),
+            unique=unique if unique is not None else column.unique,
         )
 
     @classmethod

datachain/data_storage/sqlite.py

@@ -122,6 +122,11 @@ class SQLiteDatabaseEngine(DatabaseEngine):
             engine = sqlalchemy.create_engine(
                 "sqlite+pysqlite:///", creator=lambda: db, future=True
             )
+            # ensure we run SA on_connect init (e.g it registers regexp function),
+            # also makes sure that it's consistent. Otherwise in some cases it
+            # seems we are getting different results if engine object is used in a
+            # different thread first and enine is not used in the Main thread.
+            engine.connect().close()
 
             db.isolation_level = None  # Use autocommit mode
             db.execute("PRAGMA foreign_keys = ON")

datachain/data_storage/warehouse.py

@@ -17,7 +17,7 @@ from sqlalchemy.sql.expression import true
 
 from datachain.client import Client
 from datachain.data_storage.serializer import Serializable
-from datachain.dataset import DatasetRecord, RowDict
+from datachain.dataset import DatasetRecord
 from datachain.node import DirType, DirTypeGroup, Entry, Node, NodeWithPath, get_path
 from datachain.sql.functions import path as pathfunc
 from datachain.sql.types import Int, SQLType
@@ -201,23 +201,17 @@ class AbstractWarehouse(ABC, Serializable):
     def dataset_select_paginated(
         self,
         query,
-        limit: Optional[int] = None,
-        order_by: tuple["ColumnElement[Any]", ...] = (),
         page_size: int = SELECT_BATCH_SIZE,
-    ) -> Generator[RowDict, None, None]:
+    ) -> Generator[Sequence, None, None]:
         """
         This is equivalent to `db.execute`, but for selecting rows in batches
         """
-        cols = query.selected_columns
-        cols_names = [c.name for c in cols]
+        limit = query._limit
+        paginated_query = query.limit(page_size)
 
-        if not order_by:
-            ordering = [cols.sys__id]
-        else:
-            ordering = order_by  # type: ignore[assignment]
-
-        # reset query order by and apply new order by id
-        paginated_query = query.order_by(None).order_by(*ordering).limit(page_size)
+        if not paginated_query._order_by_clauses:
+            # default order by is order by `sys__id`
+            paginated_query = paginated_query.order_by(query.selected_columns.sys__id)
 
         results = None
         offset = 0
@@ -236,7 +230,7 @@ class AbstractWarehouse(ABC, Serializable):
                 processed = False
                 for row in results:
                     processed = True
-                    yield RowDict(zip(cols_names, row))
+                    yield row
                     num_yielded += 1
 
                 if not processed:

datachain/lib/dc.py

@@ -508,7 +508,7 @@ class DataChain(DatasetQuery):
 
     def print_json_schema(  # type: ignore[override]
         self, jmespath: Optional[str] = None, model_name: Optional[str] = None
-    ) -> "DataChain":
+    ) -> "Self":
         """Print JSON data model and save it. It returns the chain itself.
 
         Parameters:
@@ -533,7 +533,7 @@ class DataChain(DatasetQuery):
 
     def print_jsonl_schema(  # type: ignore[override]
         self, jmespath: Optional[str] = None, model_name: Optional[str] = None
-    ) -> "DataChain":
+    ) -> "Self":
         """Print JSON data model and save it. It returns the chain itself.
 
         Parameters:
@@ -549,7 +549,7 @@ class DataChain(DatasetQuery):
 
     def save(  # type: ignore[override]
         self, name: Optional[str] = None, version: Optional[int] = None
-    ) -> "DataChain":
+    ) -> "Self":
         """Save to a Dataset. It returns the chain itself.
 
         Parameters:
@@ -785,7 +785,7 @@ class DataChain(DatasetQuery):
             descending (bool): Whether to sort in descending order or not.
         """
         if descending:
-            args = tuple([sqlalchemy.desc(a) for a in args])
+            args = tuple(sqlalchemy.desc(a) for a in args)
 
         return super().order_by(*args)
 
@@ -1206,14 +1206,14 @@ class DataChain(DatasetQuery):
         """
         headers, max_length = self._effective_signals_schema.get_headers_with_length()
         if flatten or max_length < 2:
-            df = pd.DataFrame.from_records(self.to_records())
+            columns = []
             if headers:
-                df.columns = [".".join(filter(None, header)) for header in headers]
-            return df
+                columns = [".".join(filter(None, header)) for header in headers]
+            return pd.DataFrame.from_records(self.to_records(), columns=columns)
 
-        transposed_result = list(map(list, zip(*self.results())))
-        data = {tuple(n): val for n, val in zip(headers, transposed_result)}
-        return pd.DataFrame(data)
+        return pd.DataFrame(
+            self.results(), columns=pd.MultiIndex.from_tuples(map(tuple, headers))
+        )
 
     def show(
         self,
@@ -1232,6 +1232,12 @@ class DataChain(DatasetQuery):
         """
         dc = self.limit(limit) if limit > 0 else self
         df = dc.to_pandas(flatten)
+
+        if df.empty:
+            print("Empty result")
+            print(f"Columns: {list(df.columns)}")
+            return
+
         if transpose:
             df = df.T
 
@@ -1270,7 +1276,7 @@ class DataChain(DatasetQuery):
         source: bool = True,
         nrows: Optional[int] = None,
         **kwargs,
-    ) -> "DataChain":
+    ) -> "Self":
         """Generate chain from list of tabular files.
 
         Parameters:
@@ -1390,7 +1396,8 @@ class DataChain(DatasetQuery):
             dc = DataChain.from_csv("s3://mybucket/dir")
             ```
         """
-        from pyarrow.csv import ParseOptions, ReadOptions
+        from pandas.io.parsers.readers import STR_NA_VALUES
+        from pyarrow.csv import ConvertOptions, ParseOptions, ReadOptions
         from pyarrow.dataset import CsvFileFormat
 
         chain = DataChain.from_storage(path, **kwargs)
@@ -1414,7 +1421,14 @@ class DataChain(DatasetQuery):
 
         parse_options = ParseOptions(delimiter=delimiter)
         read_options = ReadOptions(column_names=column_names)
-        format = CsvFileFormat(parse_options=parse_options, read_options=read_options)
+        convert_options = ConvertOptions(
+            strings_can_be_null=True, null_values=STR_NA_VALUES
+        )
+        format = CsvFileFormat(
+            parse_options=parse_options,
+            read_options=read_options,
+            convert_options=convert_options,
+        )
         return chain.parse_tabular(
             output=output,
             object_name=object_name,
@@ -1623,7 +1637,7 @@ class DataChain(DatasetQuery):
 
             Using glob to match patterns
             ```py
-            dc.filter(C("file.name").glob("*.jpg))
+            dc.filter(C("file.name").glob("*.jpg"))
             ```
 
             Using `datachain.sql.functions`

datachain/lib/meta_formats.py

@@ -11,12 +11,16 @@ from collections.abc import Iterator
 from typing import Any, Callable
 
 import jmespath as jsp
-from pydantic import Field, ValidationError  # noqa: F401
+from pydantic import BaseModel, ConfigDict, Field, ValidationError  # noqa: F401
 
 from datachain.lib.data_model import DataModel  # noqa: F401
 from datachain.lib.file import File
 
 
+class UserModel(BaseModel):
+    model_config = ConfigDict(populate_by_name=True)
+
+
 def generate_uuid():
     return uuid.uuid4()  # Generates a random UUID.
 
@@ -72,6 +76,8 @@ def read_schema(source_file, data_type="csv", expr=None, model_name=None):
         data_type,
         "--class-name",
         model_name,
+        "--base-class",
+        "datachain.lib.meta_formats.UserModel",
     ]
     try:
         result = subprocess.run(  # noqa: S603
@@ -87,7 +93,7 @@ def read_schema(source_file, data_type="csv", expr=None, model_name=None):
     except subprocess.CalledProcessError as e:
         model_output = f"An error occurred in datamodel-codegen: {e.stderr}"
     print(f"{model_output}")
-    print("\n" + "from datachain.lib.data_model import DataModel" + "\n")
+    print("from datachain.lib.data_model import DataModel")
     print("\n" + f"DataModel.register({model_name})" + "\n")
     print("\n" + f"spec={model_name}" + "\n")
     return model_output

datachain/lib/udf.py

@@ -1,6 +1,5 @@
 import sys
 import traceback
-from collections.abc import Iterable, Iterator
 from typing import TYPE_CHECKING, Callable, Optional
 
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
@@ -14,16 +13,19 @@ from datachain.lib.model_store import ModelStore
 from datachain.lib.signal_schema import SignalSchema
 from datachain.lib.udf_signature import UdfSignature
 from datachain.lib.utils import AbstractUDF, DataChainError, DataChainParamsError
-from datachain.query.batch import RowBatch
+from datachain.query.batch import UDFInputBatch
 from datachain.query.schema import ColumnParameter
 from datachain.query.udf import UDFBase as _UDFBase
-from datachain.query.udf import UDFProperties, UDFResult
+from datachain.query.udf import UDFProperties
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator, Sequence
+
     from typing_extensions import Self
 
     from datachain.catalog import Catalog
-    from datachain.query.batch import BatchingResult
+    from datachain.query.batch import RowsOutput, UDFInput
+    from datachain.query.udf import UDFResult
 
 
 class UdfError(DataChainParamsError):
@@ -42,22 +44,27 @@ class UDFAdapter(_UDFBase):
 
     def run(
         self,
-        udf_inputs: "Iterable[BatchingResult]",
+        udf_fields: "Sequence[str]",
+        udf_inputs: "Iterable[RowsOutput]",
         catalog: "Catalog",
         is_generator: bool,
         cache: bool,
         download_cb: Callback = DEFAULT_CALLBACK,
         processed_cb: Callback = DEFAULT_CALLBACK,
-    ) -> Iterator[Iterable["UDFResult"]]:
+    ) -> "Iterator[Iterable[UDFResult]]":
         self.inner._catalog = catalog
         if hasattr(self.inner, "setup") and callable(self.inner.setup):
             self.inner.setup()
 
-        for batch in udf_inputs:
-            n_rows = len(batch.rows) if isinstance(batch, RowBatch) else 1
-            output = self.run_once(catalog, batch, is_generator, cache, cb=download_cb)
-            processed_cb.relative_update(n_rows)
-            yield output
+        yield from super().run(
+            udf_fields,
+            udf_inputs,
+            catalog,
+            is_generator,
+            cache,
+            download_cb,
+            processed_cb,
+        )
 
         if hasattr(self.inner, "teardown") and callable(self.inner.teardown):
             self.inner.teardown()
@@ -65,12 +72,12 @@ class UDFAdapter(_UDFBase):
     def run_once(
         self,
         catalog: "Catalog",
-        arg: "BatchingResult",
+        arg: "UDFInput",
         is_generator: bool = False,
         cache: bool = False,
         cb: Callback = DEFAULT_CALLBACK,
-    ) -> Iterable[UDFResult]:
-        if isinstance(arg, RowBatch):
+    ) -> "Iterable[UDFResult]":
+        if isinstance(arg, UDFInputBatch):
             udf_inputs = [
                 self.bind_parameters(catalog, row, cache=cache, cb=cb)
                 for row in arg.rows

datachain/node.py

@@ -47,7 +47,7 @@ class DirTypeGroup:
 @attrs.define
 class Node:
     sys__id: int = 0
-    sys__rand: int = -1
+    sys__rand: int = 0
     vtype: str = ""
     dir_type: Optional[int] = None
     path: str = ""

datachain/query/batch.py

@@ -5,21 +5,29 @@ from collections.abc import Generator, Sequence
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Callable, Optional, Union
 
-import sqlalchemy as sa
-
 from datachain.data_storage.schema import PARTITION_COLUMN_ID
 from datachain.data_storage.warehouse import SELECT_BATCH_SIZE
 
 if TYPE_CHECKING:
+    from sqlalchemy import Select
+
     from datachain.dataset import RowDict
 
 
 @dataclass
-class RowBatch:
+class RowsOutputBatch:
+    rows: Sequence[Sequence]
+
+
+RowsOutput = Union[Sequence, RowsOutputBatch]
+
+
+@dataclass
+class UDFInputBatch:
     rows: Sequence["RowDict"]
 
 
-BatchingResult = Union["RowDict", RowBatch]
+UDFInput = Union["RowDict", UDFInputBatch]
 
 
 class BatchingStrategy(ABC):
@@ -28,9 +36,9 @@ class BatchingStrategy(ABC):
     @abstractmethod
     def __call__(
         self,
-        execute: Callable,
-        query: sa.sql.selectable.Select,
-    ) -> Generator[BatchingResult, None, None]:
+        execute: Callable[..., Generator[Sequence, None, None]],
+        query: "Select",
+    ) -> Generator[RowsOutput, None, None]:
         """Apply the provided parameters to the UDF."""
 
 
@@ -42,10 +50,10 @@ class NoBatching(BatchingStrategy):
 
     def __call__(
         self,
-        execute: Callable,
-        query: sa.sql.selectable.Select,
-    ) -> Generator["RowDict", None, None]:
-        return execute(query, limit=query._limit, order_by=query._order_by_clauses)
+        execute: Callable[..., Generator[Sequence, None, None]],
+        query: "Select",
+    ) -> Generator[Sequence, None, None]:
+        return execute(query)
 
 
 class Batch(BatchingStrategy):
@@ -59,31 +67,24 @@ class Batch(BatchingStrategy):
 
     def __call__(
         self,
-        execute: Callable,
-        query: sa.sql.selectable.Select,
-    ) -> Generator[RowBatch, None, None]:
+        execute: Callable[..., Generator[Sequence, None, None]],
+        query: "Select",
+    ) -> Generator[RowsOutputBatch, None, None]:
         # choose page size that is a multiple of the batch size
         page_size = math.ceil(SELECT_BATCH_SIZE / self.count) * self.count
 
         # select rows in batches
-        results: list[RowDict] = []
-
-        with contextlib.closing(
-            execute(
-                query,
-                page_size=page_size,
-                limit=query._limit,
-                order_by=query._order_by_clauses,
-            )
-        ) as rows:
+        results: list[Sequence] = []
+
+        with contextlib.closing(execute(query, page_size=page_size)) as rows:
             for row in rows:
                 results.append(row)
                 if len(results) >= self.count:
                     batch, results = results[: self.count], results[self.count :]
-                    yield RowBatch(batch)
+                    yield RowsOutputBatch(batch)
 
             if len(results) > 0:
-                yield RowBatch(results)
+                yield RowsOutputBatch(results)
 
 
 class Partition(BatchingStrategy):
@@ -95,27 +96,30 @@ class Partition(BatchingStrategy):
 
     def __call__(
         self,
-        execute: Callable,
-        query: sa.sql.selectable.Select,
-    ) -> Generator[RowBatch, None, None]:
+        execute: Callable[..., Generator[Sequence, None, None]],
+        query: "Select",
+    ) -> Generator[RowsOutputBatch, None, None]:
         current_partition: Optional[int] = None
-        batch: list[RowDict] = []
-
-        with contextlib.closing(
-            execute(
-                query,
-                order_by=(PARTITION_COLUMN_ID, "sys__id", *query._order_by_clauses),
-                limit=query._limit,
-            )
-        ) as rows:
+        batch: list[Sequence] = []
+
+        query_fields = [str(c.name) for c in query.selected_columns]
+        partition_column_idx = query_fields.index(PARTITION_COLUMN_ID)
+
+        ordered_query = query.order_by(None).order_by(
+            PARTITION_COLUMN_ID,
+            "sys__id",
+            *query._order_by_clauses,
+        )
+
+        with contextlib.closing(execute(ordered_query)) as rows:
             for row in rows:
-                partition = row[PARTITION_COLUMN_ID]
+                partition = row[partition_column_idx]
                 if current_partition != partition:
                     current_partition = partition
                     if len(batch) > 0:
-                        yield RowBatch(batch)
+                        yield RowsOutputBatch(batch)
                         batch = []
                 batch.append(row)
 
             if len(batch) > 0:
-                yield RowBatch(batch)
+                yield RowsOutputBatch(batch)

datachain/query/dataset.py

@@ -461,6 +461,8 @@ class UDFStep(Step, ABC):
 
         processes = determine_processes(self.parallel)
 
+        udf_fields = [str(c.name) for c in query.selected_columns]
+
         try:
             if workers:
                 from datachain.catalog.loader import get_distributed_class
@@ -473,6 +475,7 @@ class UDFStep(Step, ABC):
                     query,
                     workers,
                     processes,
+                    udf_fields=udf_fields,
                     is_generator=self.is_generator,
                     use_partitioning=use_partitioning,
                     cache=self.cache,
@@ -489,6 +492,7 @@ class UDFStep(Step, ABC):
                     "warehouse_clone_params": self.catalog.warehouse.clone_params(),
                     "table": udf_table,
                     "query": query,
+                    "udf_fields": udf_fields,
                     "batching": batching,
                     "processes": processes,
                     "is_generator": self.is_generator,
@@ -528,6 +532,7 @@ class UDFStep(Step, ABC):
                     generated_cb = get_generated_callback(self.is_generator)
                     try:
                         udf_results = udf.run(
+                            udf_fields,
                             udf_inputs,
                             self.catalog,
                             self.is_generator,
@@ -1244,21 +1249,23 @@ class DatasetQuery:
         actual_params = [normalize_param(p) for p in params]
         try:
             query = self.apply_steps().select()
+            query_fields = [str(c.name) for c in query.selected_columns]
 
-            def row_iter() -> Generator[RowDict, None, None]:
+            def row_iter() -> Generator[Sequence, None, None]:
                 # warehouse isn't threadsafe, we need to clone() it
                 # in the thread that uses the results
                 with self.catalog.warehouse.clone() as warehouse:
-                    gen = warehouse.dataset_select_paginated(
-                        query, limit=query._limit, order_by=query._order_by_clauses
-                    )
+                    gen = warehouse.dataset_select_paginated(query)
                     with contextlib.closing(gen) as rows:
                         yield from rows
 
-            async def get_params(row: RowDict) -> tuple:
+            async def get_params(row: Sequence) -> tuple:
+                row_dict = RowDict(zip(query_fields, row))
                 return tuple(
                     [
-                        await p.get_value_async(self.catalog, row, mapper, **kwargs)
+                        await p.get_value_async(
+                            self.catalog, row_dict, mapper, **kwargs
+                        )
                         for p in actual_params
                     ]
                 )