datachain 0.5.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

datachain/catalog/catalog.py

@@ -988,6 +988,14 @@ class Catalog:
         schema = {
             c.name: c.type.to_dict() for c in columns if isinstance(c.type, SQLType)
         }
+
+        job_id = job_id or os.getenv("DATACHAIN_JOB_ID")
+        if not job_id:
+            from datachain.query.session import Session
+
+            session = Session.get(catalog=self)
+            job_id = session.job_id
+
         dataset = self.metastore.create_dataset_version(
             dataset,
             version,

datachain/data_storage/metastore.py

@@ -50,7 +50,6 @@ if TYPE_CHECKING:
     from datachain.data_storage import AbstractIDGenerator, schema
     from datachain.data_storage.db_engine import DatabaseEngine
 
-
 logger = logging.getLogger("datachain")
 
 
@@ -384,6 +383,11 @@ class AbstractMetastore(ABC, Serializable):
     ) -> None:
         """Set the status of the given job and dataset."""
 
+    @abstractmethod
+    def get_job_dataset_versions(self, job_id: str) -> list[tuple[str, int]]:
+        """Returns dataset names and versions for the job."""
+        raise NotImplementedError
+
 
 class AbstractDBMetastore(AbstractMetastore):
     """
@@ -1519,3 +1523,18 @@ class AbstractDBMetastore(AbstractMetastore):
                 .values(status=dataset_status)
             )
             self.db.execute(query, conn=conn)  # type: ignore[attr-defined]
+
+    def get_job_dataset_versions(self, job_id: str) -> list[tuple[str, int]]:
+        """Returns dataset names and versions for the job."""
+        dv = self._datasets_versions
+        ds = self._datasets
+
+        join_condition = dv.c.dataset_id == ds.c.id
+
+        query = (
+            self._datasets_versions_select(ds.c.name, dv.c.version)
+            .select_from(dv.join(ds, join_condition))
+            .where(dv.c.job_id == job_id)
+        )
+
+        return list(self.db.execute(query))

datachain/data_storage/sqlite.py

@@ -15,6 +15,7 @@ from typing import (
 )
 
 import sqlalchemy
+from packaging import version
 from sqlalchemy import MetaData, Table, UniqueConstraint, exists, select
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.schema import CreateIndex, CreateTable, DropTable
@@ -153,7 +154,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
             if os.environ.get("DEBUG_SHOW_SQL_QUERIES"):
                 import sys
 
-                db.set_trace_callback(sys.stderr.write)
+                db.set_trace_callback(lambda stmt: print(stmt, file=sys.stderr))
 
             load_usearch_extension(db)
 
@@ -345,45 +346,36 @@ class SQLiteIDGenerator(AbstractDBIDGenerator):
     def get_next_ids(self, uri: str, count: int) -> range:
         """Returns a range of IDs for the given URI."""
 
-        # NOTE: we can't use RETURNING clause here because it is only available
-        # in sqlalchemy v2, see
-        # https://github.com/sqlalchemy/sqlalchemy/issues/6195#issuecomment-1248700677
-        # After we upgrade to sqlalchemy v2, we can use the following code,
-        # leaving fallback to the current implementation for older versions of SQLite,
-        # which is still supported, for example, in Ubuntu 20.04 LTS (Focal Fossa),
-        # where SQLite version 3.31.1 is used.
-
-        # sqlite_version = version.parse(sqlite3.sqlite_version)
-        # if sqlite_version >= version.parse("3.35.0"):
-        #     # RETURNING is supported on SQLite 3.35.0 (2021-03-12) or newer
-        #     stmt = (
-        #         sqlite.insert(self._table)
-        #         .values(uri=uri, last_id=count)
-        #         .on_conflict_do_update(
-        #             index_elements=["uri"],
-        #             set_={"last_id": self._table.c.last_id + count},
-        #         )
-        #         .returning(self._table.c.last_id)
-        #     )
-        #     last_id = self._db.execute(stmt).fetchone()[0]
-        # else:
-        #     (fallback to the current implementation with a transaction)
-
-        # Transactions ensure no concurrency conflicts
-        with self._db.transaction() as conn:
-            # UPSERT syntax was added to SQLite with version 3.24.0 (2018-06-04).
-            stmt_ins = (
+        sqlite_version = version.parse(sqlite3.sqlite_version)
+        is_returning_supported = sqlite_version >= version.parse("3.35.0")
+        if is_returning_supported:
+            stmt = (
                 sqlite.insert(self._table)
                 .values(uri=uri, last_id=count)
                 .on_conflict_do_update(
                     index_elements=["uri"],
                     set_={"last_id": self._table.c.last_id + count},
                 )
+                .returning(self._table.c.last_id)
             )
-            self._db.execute(stmt_ins, conn=conn)
+            last_id = self._db.execute(stmt).fetchone()[0]
+        else:
+            # Older versions of SQLite are still the default under Ubuntu LTS,
+            # e.g. Ubuntu 20.04 LTS (Focal Fossa) uses 3.31.1
+            # Transactions ensure no concurrency conflicts
+            with self._db.transaction() as conn:
+                stmt_ins = (
+                    sqlite.insert(self._table)
+                    .values(uri=uri, last_id=count)
+                    .on_conflict_do_update(
+                        index_elements=["uri"],
+                        set_={"last_id": self._table.c.last_id + count},
+                    )
+                )
+                self._db.execute(stmt_ins, conn=conn)
 
-            stmt_sel = select(self._table.c.last_id).where(self._table.c.uri == uri)
-            last_id = self._db.execute(stmt_sel, conn=conn).fetchone()[0]
+                stmt_sel = select(self._table.c.last_id).where(self._table.c.uri == uri)
+                last_id = self._db.execute(stmt_sel, conn=conn).fetchone()[0]
 
         return range(last_id - count + 1, last_id + 1)
 

datachain/lib/arrow.py

@@ -1,8 +1,9 @@
 import re
 from collections.abc import Sequence
 from tempfile import NamedTemporaryFile
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
+import orjson
 import pyarrow as pa
 from pyarrow.dataset import CsvFileFormat, dataset
 from tqdm import tqdm
@@ -10,6 +11,7 @@ from tqdm import tqdm
 from datachain.lib.data_model import dict_to_data_model
 from datachain.lib.file import ArrowRow, File
 from datachain.lib.model_store import ModelStore
+from datachain.lib.signal_schema import SignalSchema
 from datachain.lib.udf import Generator
 
 if TYPE_CHECKING:
@@ -20,6 +22,9 @@ if TYPE_CHECKING:
     from datachain.lib.dc import DataChain
 
 
+DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY = b"DataChain SignalSchema"
+
+
 class ArrowGenerator(Generator):
     def __init__(
         self,
@@ -61,28 +66,35 @@ class ArrowGenerator(Generator):
                 path, filesystem=file.get_fs(), schema=self.input_schema, **self.kwargs
             )
         hf_schema = _get_hf_schema(ds.schema)
+        use_datachain_schema = (
+            bool(ds.schema.metadata)
+            and DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY in ds.schema.metadata
+        )
         index = 0
         with tqdm(desc="Parsed by pyarrow", unit=" rows") as pbar:
             for record_batch in ds.to_batches():
                 for record in record_batch.to_pylist():
-                    vals = list(record.values())
-                    if self.output_schema:
-                        fields = self.output_schema.model_fields
-                        vals_dict = {}
-                        for i, ((field, field_info), val) in enumerate(
-                            zip(fields.items(), vals)
-                        ):
-                            anno = field_info.annotation
-                            if hf_schema:
-                                from datachain.lib.hf import convert_feature
-
-                                feat = list(hf_schema[0].values())[i]
-                                vals_dict[field] = convert_feature(val, feat, anno)
-                            elif ModelStore.is_pydantic(anno):
-                                vals_dict[field] = anno(**val)  # type: ignore[misc]
-                            else:
-                                vals_dict[field] = val
-                        vals = [self.output_schema(**vals_dict)]
+                    if use_datachain_schema and self.output_schema:
+                        vals = [_nested_model_instantiate(record, self.output_schema)]
+                    else:
+                        vals = list(record.values())
+                        if self.output_schema:
+                            fields = self.output_schema.model_fields
+                            vals_dict = {}
+                            for i, ((field, field_info), val) in enumerate(
+                                zip(fields.items(), vals)
+                            ):
+                                anno = field_info.annotation
+                                if hf_schema:
+                                    from datachain.lib.hf import convert_feature
+
+                                    feat = list(hf_schema[0].values())[i]
+                                    vals_dict[field] = convert_feature(val, feat, anno)
+                                elif ModelStore.is_pydantic(anno):
+                                    vals_dict[field] = anno(**val)  # type: ignore[misc]
+                                else:
+                                    vals_dict[field] = val
+                            vals = [self.output_schema(**vals_dict)]
                     if self.source:
                         kwargs: dict = self.kwargs
                         # Can't serialize CsvFileFormat; may lose formatting options.
@@ -113,6 +125,9 @@ def schema_to_output(schema: pa.Schema, col_names: Optional[Sequence[str]] = Non
         )
     if not col_names:
         col_names = schema.names
+    signal_schema = _get_datachain_schema(schema)
+    if signal_schema:
+        return signal_schema.values
     columns = _convert_col_names(col_names)  # type: ignore[arg-type]
     hf_schema = _get_hf_schema(schema)
     if hf_schema:
@@ -197,3 +212,33 @@ def _get_hf_schema(
         features = schema_from_arrow(schema)
         return features, get_output_schema(features)
     return None
+
+
+def _get_datachain_schema(schema: "pa.Schema") -> Optional[SignalSchema]:
+    """Return a restored SignalSchema from parquet metadata, if any is found."""
+    if schema.metadata and DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY in schema.metadata:
+        serialized_signal_schema = orjson.loads(
+            schema.metadata[DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY]
+        )
+        return SignalSchema.deserialize(serialized_signal_schema)
+    return None
+
+
+def _nested_model_instantiate(
+    column_values: dict[str, Any], model: type["BaseModel"], prefix: str = ""
+) -> "BaseModel":
+    """Instantiate the given model and all sub-models/fields based on the provided
+    column values."""
+    vals_dict = {}
+    for field, field_info in model.model_fields.items():
+        anno = field_info.annotation
+        cur_path = f"{prefix}.{field}" if prefix else field
+        if ModelStore.is_pydantic(anno):
+            vals_dict[field] = _nested_model_instantiate(
+                column_values,
+                anno,  # type: ignore[arg-type]
+                prefix=cur_path,
+            )
+        elif cur_path in column_values:
+            vals_dict[field] = column_values[cur_path]
+    return model(**vals_dict)

datachain/lib/dc.py

@@ -16,6 +16,7 @@ from typing import (
     overload,
 )
 
+import orjson
 import pandas as pd
 import sqlalchemy
 from pydantic import BaseModel
@@ -58,7 +59,7 @@ from datachain.query.dataset import (
 from datachain.query.schema import DEFAULT_DELIMITER, Column, DatasetRow
 from datachain.sql.functions import path as pathfunc
 from datachain.telemetry import telemetry
-from datachain.utils import inside_notebook
+from datachain.utils import batched_it, inside_notebook
 
 if TYPE_CHECKING:
     from typing_extensions import Concatenate, ParamSpec, Self
@@ -71,6 +72,10 @@ C = Column
 
 _T = TypeVar("_T")
 D = TypeVar("D", bound="DataChain")
+UDFObjT = TypeVar("UDFObjT", bound=UDFBase)
+
+
+DEFAULT_PARQUET_CHUNK_SIZE = 100_000
 
 
 def resolve_columns(
@@ -819,7 +824,7 @@ class DataChain:
 
     def gen(
         self,
-        func: Optional[Callable] = None,
+        func: Optional[Union[Callable, Generator]] = None,
         params: Union[None, str, Sequence[str]] = None,
         output: OutputType = None,
         **signal_map,
@@ -931,12 +936,12 @@ class DataChain:
 
     def _udf_to_obj(
         self,
-        target_class: type[UDFBase],
-        func: Optional[Callable],
+        target_class: type[UDFObjT],
+        func: Optional[Union[Callable, UDFObjT]],
         params: Union[None, str, Sequence[str]],
         output: OutputType,
         signal_map,
-    ) -> UDFBase:
+    ) -> UDFObjT:
         is_generator = target_class.is_output_batched
         name = self.name or ""
 
@@ -1103,6 +1108,29 @@ class DataChain:
                 rows = (row_factory(db_signals, r) for r in rows)
             yield from rows
 
+    def to_columnar_data_with_names(
+        self, chunk_size: int = DEFAULT_PARQUET_CHUNK_SIZE
+    ) -> tuple[list[str], Iterator[list[list[Any]]]]:
+        """Returns column names and the results as an iterator that provides chunks,
+        with each chunk containing a list of columns, where each column contains a
+        list of the row values for that column in that chunk. Useful for columnar data
+        formats, such as parquet or other OLAP databases.
+        """
+        headers, _ = self._effective_signals_schema.get_headers_with_length()
+        column_names = [".".join(filter(None, header)) for header in headers]
+
+        results_iter = self.collect_flatten()
+
+        def column_chunks() -> Iterator[list[list[Any]]]:
+            for chunk_iter in batched_it(results_iter, chunk_size):
+                columns: list[list[Any]] = [[] for _ in column_names]
+                for row in chunk_iter:
+                    for i, col in enumerate(columns):
+                        col.append(row[i])
+                yield columns
+
+        return column_names, column_chunks()
+
     @overload
     def results(self) -> list[tuple[Any, ...]]: ...
 
@@ -1808,21 +1836,96 @@ class DataChain:
         self,
         path: Union[str, os.PathLike[str], BinaryIO],
         partition_cols: Optional[Sequence[str]] = None,
+        chunk_size: int = DEFAULT_PARQUET_CHUNK_SIZE,
         **kwargs,
     ) -> None:
-        """Save chain to parquet file.
+        """Save chain to parquet file with SignalSchema metadata.
 
         Parameters:
             path : Path or a file-like binary object to save the file.
             partition_cols : Column names by which to partition the dataset.
+            chunk_size : The chunk size of results to read and convert to columnar
+                data, to avoid running out of memory.
         """
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+
+        from datachain.lib.arrow import DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY
+
         _partition_cols = list(partition_cols) if partition_cols else None
-        return self.to_pandas().to_parquet(
-            path,
-            partition_cols=_partition_cols,
-            **kwargs,
+        signal_schema_metadata = orjson.dumps(
+            self._effective_signals_schema.serialize()
         )
 
+        column_names, column_chunks = self.to_columnar_data_with_names(chunk_size)
+
+        parquet_schema = None
+        parquet_writer = None
+        first_chunk = True
+
+        for chunk in column_chunks:
+            # pyarrow infers the best parquet schema from the python types of
+            # the input data.
+            table = pa.Table.from_pydict(
+                dict(zip(column_names, chunk)),
+                schema=parquet_schema,
+            )
+
+            # Preserve any existing metadata, and add the DataChain SignalSchema.
+            existing_metadata = table.schema.metadata or {}
+            merged_metadata = {
+                **existing_metadata,
+                DATACHAIN_SIGNAL_SCHEMA_PARQUET_KEY: signal_schema_metadata,
+            }
+            table = table.replace_schema_metadata(merged_metadata)
+            parquet_schema = table.schema
+
+            if _partition_cols:
+                # Write to a partitioned parquet dataset.
+                pq.write_to_dataset(
+                    table,
+                    root_path=path,
+                    partition_cols=_partition_cols,
+                    **kwargs,
+                )
+            else:
+                if first_chunk:
+                    # Write to a single parquet file.
+                    parquet_writer = pq.ParquetWriter(path, parquet_schema, **kwargs)
+                    first_chunk = False
+
+                assert parquet_writer
+                parquet_writer.write_table(table)
+
+        if parquet_writer:
+            parquet_writer.close()
+
+    def to_csv(
+        self,
+        path: Union[str, os.PathLike[str]],
+        delimiter: str = ",",
+        **kwargs,
+    ) -> None:
+        """Save chain to a csv (comma-separated values) file.
+
+        Parameters:
+            path : Path to save the file.
+            delimiter : Delimiter to use for the resulting file.
+        """
+        import csv
+
+        headers, _ = self._effective_signals_schema.get_headers_with_length()
+        column_names = [".".join(filter(None, header)) for header in headers]
+
+        results_iter = self.collect_flatten()
+
+        with open(path, "w", newline="") as f:
+            writer = csv.writer(f, delimiter=delimiter, **kwargs)
+            writer.writerow(column_names)
+
+            for row in results_iter:
+                writer.writerow(row)
+
     @classmethod
     def from_records(
         cls,

datachain/lib/udf.py

@@ -1,31 +1,33 @@
 import sys
 import traceback
-from typing import TYPE_CHECKING, Callable, Optional
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from pydantic import BaseModel
 
 from datachain.dataset import RowDict
 from datachain.lib.convert.flatten import flatten
-from datachain.lib.convert.unflatten import unflatten_to_json
 from datachain.lib.file import File
-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 UDFInputBatch
-from datachain.query.schema import ColumnParameter
-from datachain.query.udf import UDFBase as _UDFBase
-from datachain.query.udf import UDFProperties
+from datachain.query.batch import (
+    Batch,
+    BatchingStrategy,
+    NoBatching,
+    Partition,
+    RowsOutputBatch,
+    UDFInputBatch,
+)
+from datachain.query.schema import ColumnParameter, UDFParameter
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Iterator, Sequence
-
     from typing_extensions import Self
 
     from datachain.catalog import Catalog
+    from datachain.lib.udf_signature import UdfSignature
     from datachain.query.batch import RowsOutput, UDFInput
-    from datachain.query.udf import UDFResult
 
 
 class UdfError(DataChainParamsError):
@@ -33,14 +35,47 @@ class UdfError(DataChainParamsError):
         super().__init__(f"UDF error: {msg}")
 
 
-class UDFAdapter(_UDFBase):
+ColumnType = Any
+
+# Specification for the output of a UDF
+UDFOutputSpec = Mapping[str, ColumnType]
+
+# Result type when calling the UDF wrapper around the actual
+# Python function / class implementing it.
+UDFResult = dict[str, Any]
+
+
+@dataclass
+class UDFProperties:
+    """Container for basic UDF properties."""
+
+    params: list[UDFParameter]
+    output: UDFOutputSpec
+    batch: int = 1
+
+    def get_batching(self, use_partitioning: bool = False) -> BatchingStrategy:
+        if use_partitioning:
+            return Partition()
+        if self.batch == 1:
+            return NoBatching()
+        if self.batch > 1:
+            return Batch(self.batch)
+        raise ValueError(f"invalid batch size {self.batch}")
+
+    def signal_names(self) -> Iterable[str]:
+        return self.output.keys()
+
+
+class UDFAdapter:
     def __init__(
         self,
         inner: "UDFBase",
         properties: UDFProperties,
     ):
         self.inner = inner
-        super().__init__(properties)
+        self.properties = properties
+        self.signal_names = properties.signal_names()
+        self.output = properties.output
 
     def run(
         self,
@@ -51,20 +86,23 @@ class UDFAdapter(_UDFBase):
         cache: bool,
         download_cb: Callback = DEFAULT_CALLBACK,
         processed_cb: Callback = DEFAULT_CALLBACK,
-    ) -> "Iterator[Iterable[UDFResult]]":
-        self.inner._catalog = catalog
+    ) -> Iterator[Iterable[UDFResult]]:
+        self.inner.catalog = catalog
         if hasattr(self.inner, "setup") and callable(self.inner.setup):
             self.inner.setup()
 
-        yield from super().run(
-            udf_fields,
-            udf_inputs,
-            catalog,
-            is_generator,
-            cache,
-            download_cb,
-            processed_cb,
-        )
+        for batch in udf_inputs:
+            if isinstance(batch, RowsOutputBatch):
+                n_rows = len(batch.rows)
+                inputs: UDFInput = UDFInputBatch(
+                    [RowDict(zip(udf_fields, row)) for row in batch.rows]
+                )
+            else:
+                n_rows = 1
+                inputs = RowDict(zip(udf_fields, batch))
+            output = self.run_once(catalog, inputs, is_generator, cache, cb=download_cb)
+            processed_cb.relative_update(n_rows)
+            yield output
 
         if hasattr(self.inner, "teardown") and callable(self.inner.teardown):
             self.inner.teardown()
@@ -76,23 +114,46 @@ class UDFAdapter(_UDFBase):
         is_generator: bool = False,
         cache: bool = False,
         cb: Callback = DEFAULT_CALLBACK,
-    ) -> "Iterable[UDFResult]":
+    ) -> Iterable[UDFResult]:
         if isinstance(arg, UDFInputBatch):
             udf_inputs = [
                 self.bind_parameters(catalog, row, cache=cache, cb=cb)
                 for row in arg.rows
             ]
-            udf_outputs = self.inner(udf_inputs, cache=cache, download_cb=cb)
+            udf_outputs = self.inner.run_once(udf_inputs, cache=cache, download_cb=cb)
             return self._process_results(arg.rows, udf_outputs, is_generator)
         if isinstance(arg, RowDict):
             udf_inputs = self.bind_parameters(catalog, arg, cache=cache, cb=cb)
-            udf_outputs = self.inner(*udf_inputs, cache=cache, download_cb=cb)
+            udf_outputs = self.inner.run_once(udf_inputs, cache=cache, download_cb=cb)
             if not is_generator:
                 # udf_outputs is generator already if is_generator=True
                 udf_outputs = [udf_outputs]
             return self._process_results([arg], udf_outputs, is_generator)
         raise ValueError(f"Unexpected UDF argument: {arg}")
 
+    def bind_parameters(self, catalog: "Catalog", row: "RowDict", **kwargs) -> list:
+        return [p.get_value(catalog, row, **kwargs) for p in self.properties.params]
+
+    def _process_results(
+        self,
+        rows: Sequence["RowDict"],
+        results: Sequence[Sequence[Any]],
+        is_generator=False,
+    ) -> Iterable[UDFResult]:
+        """Create a list of dictionaries representing UDF results."""
+
+        # outputting rows
+        if is_generator:
+            # each row in results is a tuple of column values
+            return (dict(zip(self.signal_names, row)) for row in results)
+
+        # outputting signals
+        row_ids = [row["sys__id"] for row in rows]
+        return [
+            {"sys__id": row_id} | dict(zip(self.signal_names, signals))
+            for row_id, signals in zip(row_ids, results)
+        ]
+
 
 class UDFBase(AbstractUDF):
     """Base class for stateful user-defined functions.
@@ -146,14 +207,14 @@ class UDFBase(AbstractUDF):
     is_output_batched = False
     is_input_grouped = False
     params_spec: Optional[list[str]]
+    catalog: "Optional[Catalog]"
 
     def __init__(self):
         self.params = None
         self.output = None
         self.params_spec = None
         self.output_spec = None
-        self._contains_stream = None
-        self._catalog = None
+        self.catalog = None
         self._func = None
 
     def process(self, *args, **kwargs):
@@ -174,9 +235,9 @@ class UDFBase(AbstractUDF):
 
     def _init(
         self,
-        sign: UdfSignature,
+        sign: "UdfSignature",
         params: SignalSchema,
-        func: Callable,
+        func: Optional[Callable],
     ):
         self.params = params
         self.output = sign.output_schema
@@ -190,13 +251,13 @@ class UDFBase(AbstractUDF):
     @classmethod
     def _create(
         cls,
-        sign: UdfSignature,
+        sign: "UdfSignature",
         params: SignalSchema,
     ) -> "Self":
         if isinstance(sign.func, AbstractUDF):
             if not isinstance(sign.func, cls):  # type: ignore[unreachable]
                 raise UdfError(
-                    f"cannot create UDF: provided UDF '{sign.func.__name__}'"
+                    f"cannot create UDF: provided UDF '{type(sign.func).__name__}'"
                     f" must be a child of target class '{cls.__name__}'",
                 )
             result = sign.func
@@ -212,13 +273,6 @@ class UDFBase(AbstractUDF):
     def name(self):
         return self.__class__.__name__
 
-    def set_catalog(self, catalog):
-        self._catalog = catalog.copy(db=False)
-
-    @property
-    def catalog(self):
-        return self._catalog
-
     def to_udf_wrapper(self, batch: int = 1) -> UDFAdapter:
         assert self.params_spec is not None
         properties = UDFProperties(
@@ -229,11 +283,9 @@ class UDFBase(AbstractUDF):
     def validate_results(self, results, *args, **kwargs):
         return results
 
-    def __call__(self, *rows, cache, download_cb):
-        if self.is_input_grouped:
-            objs = self._parse_grouped_rows(rows[0], cache, download_cb)
-        elif self.is_input_batched:
-            objs = zip(*self._parse_rows(rows[0], cache, download_cb))
+    def run_once(self, rows, cache, download_cb):
+        if self.is_input_batched:
+            objs = zip(*self._parse_rows(rows, cache, download_cb))
         else:
             objs = self._parse_rows([rows], cache, download_cb)[0]
 
@@ -259,8 +311,8 @@ class UDFBase(AbstractUDF):
         ):
             res = list(res)
             assert len(res) == len(
-                rows[0]
-            ), f"{self.name} returns {len(res)} rows while len(rows[0]) expected"
+                rows
+            ), f"{self.name} returns {len(res)} rows while {len(rows)} expected"
 
         return res
 
@@ -283,41 +335,11 @@ class UDFBase(AbstractUDF):
             for obj in obj_row:
                 if isinstance(obj, File):
                     obj._set_stream(
-                        self._catalog, caching_enabled=cache, download_cb=download_cb
+                        self.catalog, caching_enabled=cache, download_cb=download_cb
                     )
             objs.append(obj_row)
         return objs
 
-    def _parse_grouped_rows(self, group, cache, download_cb):
-        spec_map = {}
-        output_map = {}
-        for name, (anno, subtree) in self.params.tree.items():
-            if ModelStore.is_pydantic(anno):
-                length = sum(1 for _ in self.params._get_flat_tree(subtree, [], 0))
-            else:
-                length = 1
-            spec_map[name] = anno, length
-            output_map[name] = []
-
-        for flat_obj in group:
-            position = 0
-            for signal, (cls, length) in spec_map.items():
-                slice = flat_obj[position : position + length]
-                position += length
-
-                if ModelStore.is_pydantic(cls):
-                    obj = cls(**unflatten_to_json(cls, slice))
-                else:
-                    obj = slice[0]
-
-                if isinstance(obj, File):
-                    obj._set_stream(
-                        self._catalog, caching_enabled=cache, download_cb=download_cb
-                    )
-                output_map[signal].append(obj)
-
-        return list(output_map.values())
-
     def process_safe(self, obj_rows):
         try:
             result_objs = self.process(*obj_rows)

datachain/lib/udf_signature.py

@@ -1,10 +1,11 @@
 import inspect
 from collections.abc import Generator, Iterator, Sequence
 from dataclasses import dataclass
-from typing import Callable, Optional, Union, get_args, get_origin
+from typing import Callable, Union, get_args, get_origin
 
 from datachain.lib.data_model import DataType, DataTypeNames, is_chain_type
 from datachain.lib.signal_schema import SignalSchema
+from datachain.lib.udf import UDFBase
 from datachain.lib.utils import AbstractUDF, DataChainParamsError
 
 
@@ -16,7 +17,7 @@ class UdfSignatureError(DataChainParamsError):
 
 @dataclass
 class UdfSignature:
-    func: Callable
+    func: Union[Callable, UDFBase]
     params: Sequence[str]
     output_schema: SignalSchema
 
@@ -27,7 +28,7 @@ class UdfSignature:
         cls,
         chain: str,
         signal_map: dict[str, Callable],
-        func: Optional[Callable] = None,
+        func: Union[None, UDFBase, Callable] = None,
         params: Union[None, str, Sequence[str]] = None,
         output: Union[None, DataType, Sequence[str], dict[str, DataType]] = None,
         is_generator: bool = True,
@@ -39,6 +40,7 @@ class UdfSignature:
                 f"multiple signals '{keys}' are not supported in processors."
                 " Chain multiple processors instead.",
             )
+        udf_func: Union[UDFBase, Callable]
         if len(signal_map) == 1:
             if func is not None:
                 raise UdfSignatureError(
@@ -53,7 +55,7 @@ class UdfSignature:
             udf_func = func
             signal_name = None
 
-        if not callable(udf_func):
+        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 = (
@@ -73,7 +75,7 @@ class UdfSignature:
             if not func_outs_sign:
                 raise UdfSignatureError(
                     chain,
-                    f"outputs are not defined in function '{udf_func.__name__}'"
+                    f"outputs are not defined in function '{udf_func}'"
                     " hints or 'output'",
                 )
 
@@ -154,7 +156,7 @@ class UdfSignature:
 
     @staticmethod
     def _func_signature(
-        chain: str, udf_func: Callable
+        chain: str, udf_func: Union[Callable, UDFBase]
     ) -> tuple[dict[str, type], Sequence[type], bool]:
         if isinstance(udf_func, AbstractUDF):
             func = udf_func.process  # type: ignore[unreachable]

datachain/query/dataset.py

@@ -42,6 +42,7 @@ from datachain.data_storage.schema import (
 )
 from datachain.dataset import DatasetStatus, RowDict
 from datachain.error import DatasetNotFoundError, QueryScriptCancelError
+from datachain.lib.udf import UDFAdapter
 from datachain.progress import CombinedDownloadCallback
 from datachain.sql.functions import rand
 from datachain.utils import (
@@ -53,7 +54,6 @@ from datachain.utils import (
 
 from .schema import C, UDFParamSpec, normalize_param
 from .session import Session
-from .udf import UDFBase
 
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ClauseElement
@@ -299,7 +299,7 @@ def adjust_outputs(
     return row
 
 
-def get_udf_col_types(warehouse: "AbstractWarehouse", udf: UDFBase) -> list[tuple]:
+def get_udf_col_types(warehouse: "AbstractWarehouse", udf: UDFAdapter) -> list[tuple]:
     """Optimization: Precompute UDF column types so these don't have to be computed
     in the convert_type function for each row in a loop."""
     dialect = warehouse.db.dialect
@@ -320,7 +320,7 @@ def process_udf_outputs(
     warehouse: "AbstractWarehouse",
     udf_table: "Table",
     udf_results: Iterator[Iterable["UDFResult"]],
-    udf: UDFBase,
+    udf: UDFAdapter,
     batch_size: int = INSERT_BATCH_SIZE,
     cb: Callback = DEFAULT_CALLBACK,
 ) -> None:
@@ -364,7 +364,7 @@ def get_generated_callback(is_generator: bool = False) -> Callback:
 
 @frozen
 class UDFStep(Step, ABC):
-    udf: UDFBase
+    udf: UDFAdapter
     catalog: "Catalog"
     partition_by: Optional[PartitionByType] = None
     parallel: Optional[int] = None
@@ -1465,7 +1465,7 @@ class DatasetQuery:
     @detach
     def add_signals(
         self,
-        udf: UDFBase,
+        udf: UDFAdapter,
         parallel: Optional[int] = None,
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,
@@ -1509,7 +1509,7 @@ class DatasetQuery:
     @detach
     def generate(
         self,
-        udf: UDFBase,
+        udf: UDFAdapter,
         parallel: Optional[int] = None,
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,

datachain/query/dispatch.py

@@ -13,6 +13,7 @@ from multiprocess import get_context
 
 from datachain.catalog import Catalog
 from datachain.catalog.loader import get_distributed_class
+from datachain.lib.udf import UDFAdapter, UDFResult
 from datachain.query.dataset import (
     get_download_callback,
     get_generated_callback,
@@ -27,7 +28,6 @@ from datachain.query.queue import (
     put_into_queue,
     unmarshal,
 )
-from datachain.query.udf import UDFBase, UDFResult
 from datachain.utils import batched_it
 
 DEFAULT_BATCH_SIZE = 10000
@@ -336,7 +336,7 @@ class ProcessedCallback(Callback):
 @attrs.define
 class UDFWorker:
     catalog: Catalog
-    udf: UDFBase
+    udf: UDFAdapter
     task_queue: "multiprocess.Queue"
     done_queue: "multiprocess.Queue"
     is_generator: bool

datachain/query/session.py

@@ -1,5 +1,8 @@
 import atexit
+import logging
+import os
 import re
+import sys
 from typing import TYPE_CHECKING, Optional
 from uuid import uuid4
 
@@ -9,6 +12,8 @@ from datachain.error import TableMissingError
 if TYPE_CHECKING:
     from datachain.catalog import Catalog
 
+logger = logging.getLogger("datachain")
+
 
 class Session:
     """
@@ -35,6 +40,7 @@ class Session:
 
     GLOBAL_SESSION_CTX: Optional["Session"] = None
     GLOBAL_SESSION: Optional["Session"] = None
+    ORIGINAL_EXCEPT_HOOK = None
 
     DATASET_PREFIX = "session_"
     GLOBAL_SESSION_NAME = "global"
@@ -58,6 +64,7 @@ class Session:
 
         session_uuid = uuid4().hex[: self.SESSION_UUID_LEN]
         self.name = f"{name}_{session_uuid}"
+        self.job_id = os.getenv("DATACHAIN_JOB_ID") or str(uuid4())
         self.is_new_catalog = not catalog
         self.catalog = catalog or get_catalog(
             client_config=client_config, in_memory=in_memory
@@ -67,6 +74,9 @@ class Session:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type:
+            self._cleanup_created_versions(self.name)
+
         self._cleanup_temp_datasets()
         if self.is_new_catalog:
             self.catalog.metastore.close_on_exit()
@@ -88,6 +98,21 @@ class Session:
         except TableMissingError:
             pass
 
+    def _cleanup_created_versions(self, job_id: str) -> None:
+        versions = self.catalog.metastore.get_job_dataset_versions(job_id)
+        if not versions:
+            return
+
+        datasets = {}
+        for dataset_name, version in versions:
+            if dataset_name not in datasets:
+                datasets[dataset_name] = self.catalog.get_dataset(dataset_name)
+            dataset = datasets[dataset_name]
+            logger.info(
+                "Removing dataset version %s@%s due to exception", dataset_name, version
+            )
+            self.catalog.remove_dataset_version(dataset, version)
+
     @classmethod
     def get(
         cls,
@@ -114,9 +139,23 @@ class Session:
                 in_memory=in_memory,
             )
             cls.GLOBAL_SESSION = cls.GLOBAL_SESSION_CTX.__enter__()
+
             atexit.register(cls._global_cleanup)
+            cls.ORIGINAL_EXCEPT_HOOK = sys.excepthook
+            sys.excepthook = cls.except_hook
+
         return cls.GLOBAL_SESSION
 
+    @staticmethod
+    def except_hook(exc_type, exc_value, exc_traceback):
+        Session._global_cleanup()
+        if Session.GLOBAL_SESSION_CTX is not None:
+            job_id = Session.GLOBAL_SESSION_CTX.job_id
+            Session.GLOBAL_SESSION_CTX._cleanup_created_versions(job_id)
+
+        if Session.ORIGINAL_EXCEPT_HOOK:
+            Session.ORIGINAL_EXCEPT_HOOK(exc_type, exc_value, exc_traceback)
+
     @classmethod
     def cleanup_for_tests(cls):
         if cls.GLOBAL_SESSION_CTX is not None:
@@ -125,6 +164,9 @@ class Session:
             cls.GLOBAL_SESSION_CTX = None
             atexit.unregister(cls._global_cleanup)
 
+        if cls.ORIGINAL_EXCEPT_HOOK:
+            sys.excepthook = cls.ORIGINAL_EXCEPT_HOOK
+
     @staticmethod
     def _global_cleanup():
         if Session.GLOBAL_SESSION_CTX is not None:

{datachain-0.5.0.dist-info → datachain-0.5.1.dist-info}/METADATA

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

{datachain-0.5.0.dist-info → datachain-0.5.1.dist-info}/RECORD

@@ -18,7 +18,7 @@ datachain/storage.py,sha256=RiSJLYdHUjnrEWkLBKPcETHpAxld_B2WxLg711t0aZI,3733
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
 datachain/utils.py,sha256=KeFSRHsiYthnTu4a6bH-rw04mX1m8krTX0f2NqfQGFI,12114
 datachain/catalog/__init__.py,sha256=g2iAAFx_gEIrqshXlhSEbrc8qDaEH11cjU40n3CHDz4,409
-datachain/catalog/catalog.py,sha256=FuKuIiCwPgN5Ea25hnFe_ZFZH9YEUZ2ma9k_Lczk-JU,63867
+datachain/catalog/catalog.py,sha256=BsMyk2RQibQYHgrmovFZeSEpPVMTwgb_7ntVYdc7t-E,64090
 datachain/catalog/datasource.py,sha256=D-VWIVDCM10A8sQavLhRXdYSCG7F4o4ifswEF80_NAQ,1412
 datachain/catalog/loader.py,sha256=-6VelNfXUdgUnwInVyA8g86Boxv2xqhTh9xNS-Zlwig,8242
 datachain/client/__init__.py,sha256=T4wiYL9KIM0ZZ_UqIyzV8_ufzYlewmizlV4iymHNluE,86
@@ -33,17 +33,17 @@ datachain/data_storage/__init__.py,sha256=cEOJpyu1JDZtfUupYucCDNFI6e5Wmp_Oyzq6rZ
 datachain/data_storage/db_engine.py,sha256=81Ol1of9TTTzD97ORajCnP366Xz2mEJt6C-kTUCaru4,3406
 datachain/data_storage/id_generator.py,sha256=lCEoU0BM37Ai2aRpSbwo5oQT0GqZnSpYwwvizathRMQ,4292
 datachain/data_storage/job.py,sha256=w-7spowjkOa1P5fUVtJou3OltT0L48P0RYWZ9rSJ9-s,383
-datachain/data_storage/metastore.py,sha256=NV4FJ_W16Q19Sx70i5Qtre-n4DC2kMD0qw0vBz3j7Ks,52228
+datachain/data_storage/metastore.py,sha256=HfCxk4lmDUg2Q4WsFNQGMWxllP0mToA00fxkFTwdNIE,52919
 datachain/data_storage/schema.py,sha256=AGbjyEir5UmRZXI3m0jChZogUh5wd8csj6-YlUWaAxQ,8383
 datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2krunWUum5o,927
-datachain/data_storage/sqlite.py,sha256=EBKJncuzcyQfcKFm2mUjvHjHRTODsteM-k_zndunBrw,28834
+datachain/data_storage/sqlite.py,sha256=fW08P7AbJ0cDbTbcTKuAGpvMXvBjg-QkGsKT_Dslyws,28383
 datachain/data_storage/warehouse.py,sha256=fXhVfao3NfWFGbbG5uJ-Ga4bX1FiKVfcbDyQgECYfk8,32122
 datachain/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datachain/lib/arrow.py,sha256=aUsoQmxDmuSnB8Ik9p57Y66gc_dgx6NBqkDDIfLsvno,7630
+datachain/lib/arrow.py,sha256=0R2CYsN82nNa5_03iS6jVix9EKeeqNZNAMgpSQP2hfo,9482
 datachain/lib/clip.py,sha256=lm5CzVi4Cj1jVLEKvERKArb-egb9j1Ls-fwTItT6vlI,6150
 datachain/lib/data_model.py,sha256=gHIjlow84GMRDa78yLL1Ud-N18or21fnTyPEwsatpXY,2045
 datachain/lib/dataset_info.py,sha256=srPPhI2UHf6hFPBecyFEVw2SS5aPisIIMsvGgKqi7ss,2366
-datachain/lib/dc.py,sha256=yTyHrKIswCzdlvl2n-wdEVZEEF5VQpkLJPzPfUL9CTU,72054
+datachain/lib/dc.py,sha256=HLOAkJEKFHJV_PqwSu0Pyl1m7JmUea8_wiMJFr14Nfk,75960
 datachain/lib/file.py,sha256=LjTW_-PDAnoUhvyB4bJ8Y8n__XGqrxvmd9mDOF0Gir8,14875
 datachain/lib/hf.py,sha256=cPnmLuprr0pYABH7KqA5FARQ1JGlywdDwD3yDzVAm4k,5920
 datachain/lib/image.py,sha256=AMXYwQsmarZjRbPCZY3M1jDsM2WAB_b3cTY4uOIuXNU,2675
@@ -56,8 +56,8 @@ datachain/lib/settings.py,sha256=39thOpYJw-zPirzeNO6pmRC2vPrQvt4eBsw1xLWDFsw,234
 datachain/lib/signal_schema.py,sha256=iqgubjCBRiUJB30miv05qFX4uU04dA_Pzi3DCUsHZGs,24177
 datachain/lib/tar.py,sha256=3WIzao6yD5fbLqXLTt9GhPGNonbFIs_fDRu-9vgLgsA,1038
 datachain/lib/text.py,sha256=UNHm8fhidk7wdrWqacEWaA6I9ykfYqarQ2URby7jc7M,1261
-datachain/lib/udf.py,sha256=nG7DDuPgZ5ZuijwvDoCq-OZMxlDM8vFNzyxMmik0Y1c,11716
-datachain/lib/udf_signature.py,sha256=gMStcEeYJka5M6cg50Z9orC6y6HzCAJ3MkFqqn1fjZg,7137
+datachain/lib/udf.py,sha256=oHhJWb0gVTxcybGzYDzAeN0Gb1IMhZBoGefncT88dIY,12339
+datachain/lib/udf_signature.py,sha256=GXw24A-Olna6DWCdgy2bC-gZh_gLGPQ-KvjuI6pUjC0,7281
 datachain/lib/utils.py,sha256=5-kJlAZE0D9nXXweAjo7-SP_AWGo28feaDByONYaooQ,463
 datachain/lib/vfile.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/lib/webdataset.py,sha256=o7SHk5HOUWsZ5Ln04xOM04eQqiBHiJNO7xLgyVBrwo8,6924
@@ -70,14 +70,13 @@ datachain/lib/convert/unflatten.py,sha256=Ogvh_5wg2f38_At_1lN0D_e2uZOOpYEvwvB2xd
 datachain/lib/convert/values_to_tuples.py,sha256=YOdbjzHq-uj6-cV2Qq43G72eN2avMNDGl4x5t6yQMl8,3931
 datachain/query/__init__.py,sha256=0NBOZVgIDpCcj1Ci883dQ9A0iiwe03xzmotkOCFbxYc,293
 datachain/query/batch.py,sha256=-vlpINJiertlnaoUVv1C95RatU0F6zuhpIYRufJRo1M,3660
-datachain/query/dataset.py,sha256=tLCTaj4K93BY93GgOPv9PknZByEF89zpHc7y9s8ZF_w,53610
-datachain/query/dispatch.py,sha256=CFAc09O6UllcyUSSEY1GUlEMPzeO8RYhXinNN4HBl9M,12405
+datachain/query/dataset.py,sha256=1c7y178ccFSeL_WIba0vT87Md_Oo4F8zaTVDjB9Bp3I,53641
+datachain/query/dispatch.py,sha256=JVcZ4REE_GOsqXbar_Cb_fk-pHgQoabQLzXwuu7IhOg,12409
 datachain/query/metrics.py,sha256=r5b0ygYhokbXp8Mg3kCH8iFSRw0jxzyeBe-C-J_bKFc,938
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
 datachain/query/queue.py,sha256=waqM_KzavU8C-G95-4211Nd4GXna_u2747Chgwtgz2w,3839
 datachain/query/schema.py,sha256=I8zLWJuWl5N332ni9mAzDYtcxMJupVPgWkSDe8spNEk,8019
-datachain/query/session.py,sha256=UPH5Z4fzCDsvj81ji0e8GA6Mgra3bOAEpVq4htqOtis,4317
-datachain/query/udf.py,sha256=HB2hbEuiGA4ch9P2mh9iLA5Jj9mRj-4JFy9VfjTLJ8U,3622
+datachain/query/session.py,sha256=kpFFJMfWBnxaMPojMGhJRbk-BOsSYI8Ckl6vvqnx7d0,5787
 datachain/remote/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/remote/studio.py,sha256=f5s6qSZ9uB4URGUoU_8_W1KZRRQQVSm6cgEBkBUEfuE,7226
 datachain/sql/__init__.py,sha256=A2djrbQwSMUZZEIKGnm-mnRA-NDSbiDJNpAmmwGNyIo,303
@@ -97,9 +96,9 @@ datachain/sql/sqlite/base.py,sha256=WLPHBhZbXbiqPoRV1VgDrXJqku4UuvJpBhYeQ0k5rI8,
 datachain/sql/sqlite/types.py,sha256=yzvp0sXSEoEYXs6zaYC_2YubarQoZH-MiUNXcpuEP4s,1573
 datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR0,469
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.5.0.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.5.0.dist-info/METADATA,sha256=tKSZNiHZY0WJ_w6irkpSF7qDfuOTfiYNEQ6St3eBs-M,17156
-datachain-0.5.0.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-datachain-0.5.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.5.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.5.0.dist-info/RECORD,,
+datachain-0.5.1.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.5.1.dist-info/METADATA,sha256=n8TFKjDmTzNBMaW5Oa6MUUUOAQbAjPzkAMaKCW3Y9NU,17156
+datachain-0.5.1.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+datachain-0.5.1.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.5.1.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.5.1.dist-info/RECORD,,

datachain/query/udf.py

@@ -1,126 +0,0 @@
-import typing
-from collections.abc import Iterable, Iterator, Sequence
-from dataclasses import dataclass
-from typing import (
-    TYPE_CHECKING,
-    Any,
-)
-
-from fsspec.callbacks import DEFAULT_CALLBACK, Callback
-
-from datachain.dataset import RowDict
-
-from .batch import (
-    Batch,
-    BatchingStrategy,
-    NoBatching,
-    Partition,
-    RowsOutputBatch,
-    UDFInputBatch,
-)
-from .schema import UDFParameter
-
-if TYPE_CHECKING:
-    from datachain.catalog import Catalog
-
-    from .batch import RowsOutput, UDFInput
-
-ColumnType = Any
-
-
-# Specification for the output of a UDF
-UDFOutputSpec = typing.Mapping[str, ColumnType]
-
-# Result type when calling the UDF wrapper around the actual
-# Python function / class implementing it.
-UDFResult = dict[str, Any]
-
-
-@dataclass
-class UDFProperties:
-    """Container for basic UDF properties."""
-
-    params: list[UDFParameter]
-    output: UDFOutputSpec
-    batch: int = 1
-
-    def get_batching(self, use_partitioning: bool = False) -> BatchingStrategy:
-        if use_partitioning:
-            return Partition()
-        if self.batch == 1:
-            return NoBatching()
-        if self.batch > 1:
-            return Batch(self.batch)
-        raise ValueError(f"invalid batch size {self.batch}")
-
-    def signal_names(self) -> Iterable[str]:
-        return self.output.keys()
-
-
-class UDFBase:
-    """A base class for implementing stateful UDFs."""
-
-    def __init__(
-        self,
-        properties: UDFProperties,
-    ):
-        self.properties = properties
-        self.signal_names = properties.signal_names()
-        self.output = properties.output
-
-    def run(
-        self,
-        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"]]:
-        for batch in udf_inputs:
-            if isinstance(batch, RowsOutputBatch):
-                n_rows = len(batch.rows)
-                inputs: UDFInput = UDFInputBatch(
-                    [RowDict(zip(udf_fields, row)) for row in batch.rows]
-                )
-            else:
-                n_rows = 1
-                inputs = RowDict(zip(udf_fields, batch))
-            output = self.run_once(catalog, inputs, is_generator, cache, cb=download_cb)
-            processed_cb.relative_update(n_rows)
-            yield output
-
-    def run_once(
-        self,
-        catalog: "Catalog",
-        arg: "UDFInput",
-        is_generator: bool = False,
-        cache: bool = False,
-        cb: Callback = DEFAULT_CALLBACK,
-    ) -> Iterable[UDFResult]:
-        raise NotImplementedError
-
-    def bind_parameters(self, catalog: "Catalog", row: "RowDict", **kwargs) -> list:
-        return [p.get_value(catalog, row, **kwargs) for p in self.properties.params]
-
-    def _process_results(
-        self,
-        rows: Sequence["RowDict"],
-        results: Sequence[Sequence[Any]],
-        is_generator=False,
-    ) -> Iterable[UDFResult]:
-        """Create a list of dictionaries representing UDF results."""
-
-        # outputting rows
-        if is_generator:
-            # each row in results is a tuple of column values
-            return (dict(zip(self.signal_names, row)) for row in results)
-
-        # outputting signals
-        row_ids = [row["sys__id"] for row in rows]
-        return [
-            {"sys__id": row_id} | dict(zip(self.signal_names, signals))
-            for row_id, signals in zip(row_ids, results)
-            if signals is not None  # skip rows with no output
-        ]