datachain 0.30.6__py3-none-any.whl → 0.30.7__py3-none-any.whl

datachain/data_storage/sqlite.py

@@ -37,6 +37,7 @@ from datachain import semver
 from datachain.data_storage import AbstractDBMetastore, AbstractWarehouse
 from datachain.data_storage.db_engine import DatabaseEngine
 from datachain.data_storage.schema import DefaultSchema
+from datachain.data_storage.warehouse import INSERT_BATCH_SIZE
 from datachain.dataset import DatasetRecord, StorageURI
 from datachain.error import DataChainError, OutdatedDatabaseSchemaError
 from datachain.namespace import Namespace
@@ -44,7 +45,7 @@ from datachain.project import Project
 from datachain.sql.sqlite import create_user_defined_sql_functions, sqlite_dialect
 from datachain.sql.sqlite.base import load_usearch_extension
 from datachain.sql.types import SQLType
-from datachain.utils import DataChainDir, batched_it
+from datachain.utils import DataChainDir, batched, batched_it
 
 if TYPE_CHECKING:
     from sqlalchemy.dialects.sqlite import Insert
@@ -712,19 +713,21 @@ class SQLiteWarehouse(AbstractWarehouse):
     def prepare_entries(self, entries: "Iterable[File]") -> Iterable[dict[str, Any]]:
         return (e.model_dump() for e in entries)
 
-    def insert_rows(self, table: Table, rows: Iterable[dict[str, Any]]) -> None:
-        rows = list(rows)
-        if not rows:
-            return
-
-        with self.db.transaction() as conn:
-            # transactions speeds up inserts significantly as there is no separate
-            # transaction created for each insert row
-            self.db.executemany(
-                table.insert().values({f: bindparam(f) for f in rows[0]}),
-                rows,
-                conn=conn,
-            )
+    def insert_rows(
+        self,
+        table: Table,
+        rows: Iterable[dict[str, Any]],
+        batch_size: int = INSERT_BATCH_SIZE,
+    ) -> None:
+        for row_chunk in batched(rows, batch_size):
+            with self.db.transaction() as conn:
+                # transactions speeds up inserts significantly as there is no separate
+                # transaction created for each insert row
+                self.db.executemany(
+                    table.insert().values({f: bindparam(f) for f in row_chunk[0]}),
+                    row_chunk,
+                    conn=conn,
+                )
 
     def insert_dataset_rows(self, df, dataset: DatasetRecord, version: str) -> int:
         dr = self.dataset_rows(dataset, version)
@@ -797,7 +800,7 @@ class SQLiteWarehouse(AbstractWarehouse):
             .limit(None)
         )
 
-        for batch in batched_it(ids, 10_000):
+        for batch in batched_it(ids, INSERT_BATCH_SIZE):
             batch_ids = [row[0] for row in batch]
             select_q._where_criteria = (col_id.in_(batch_ids),)
             q = table.insert().from_select(list(select_q.selected_columns), select_q)

datachain/data_storage/warehouse.py

@@ -43,6 +43,7 @@ if TYPE_CHECKING:
 logger = logging.getLogger("datachain")
 
 SELECT_BATCH_SIZE = 100_000  # number of rows to fetch at a time
+INSERT_BATCH_SIZE = 10_000  # number of rows to insert at a time
 
 
 class AbstractWarehouse(ABC, Serializable):
@@ -415,7 +416,12 @@ class AbstractWarehouse(ABC, Serializable):
         """Convert File entries so they can be passed on to `insert_rows()`"""
 
     @abstractmethod
-    def insert_rows(self, table: sa.Table, rows: Iterable[dict[str, Any]]) -> None:
+    def insert_rows(
+        self,
+        table: sa.Table,
+        rows: Iterable[dict[str, Any]],
+        batch_size: int = INSERT_BATCH_SIZE,
+    ) -> None:
         """Does batch inserts of any kind of rows into table"""
 
     def insert_rows_done(self, table: sa.Table) -> None:

datachain/lib/dc/database.py

@@ -73,7 +73,7 @@ def to_database(
     table_name: str,
     connection: "ConnectionType",
     *,
-    batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
+    batch_size: int = DEFAULT_DATABASE_BATCH_SIZE,
     on_conflict: Optional[str] = None,
     conflict_columns: Optional[list[str]] = None,
     column_mapping: Optional[dict[str, Optional[str]]] = None,
@@ -124,7 +124,7 @@ def to_database(
                 table.create(conn, checkfirst=True)
 
                 rows_iter = chain._leaf_values()
-                for batch in batched(rows_iter, batch_rows):
+                for batch in batched(rows_iter, batch_size):
                     rows_affected = _process_batch(
                         conn,
                         table,

datachain/lib/dc/datachain.py

@@ -342,15 +342,15 @@ class DataChain:
 
     def settings(
         self,
-        cache=None,
-        parallel=None,
-        workers=None,
-        min_task_size=None,
-        prefetch: Optional[int] = None,
-        sys: Optional[bool] = None,
+        cache: Optional[bool] = None,
+        prefetch: Optional[Union[bool, int]] = None,
+        parallel: Optional[Union[bool, int]] = None,
+        workers: Optional[int] = None,
         namespace: Optional[str] = None,
         project: Optional[str] = None,
-        batch_rows: Optional[int] = None,
+        min_task_size: Optional[int] = None,
+        batch_size: Optional[int] = None,
+        sys: Optional[bool] = None,
     ) -> "Self":
         """Change settings for chain.
 
@@ -359,23 +359,23 @@ class DataChain:
 
         Parameters:
             cache : data caching. (default=False)
+            prefetch : number of workers to use for downloading files in advance.
+                      This is enabled by default and uses 2 workers.
+                      To disable prefetching, set it to 0 or False.
             parallel : number of thread for processors. True is a special value to
                 enable all available CPUs. (default=1)
             workers : number of distributed workers. Only for Studio mode. (default=1)
-            min_task_size : minimum number of tasks. (default=1)
-            prefetch : number of workers to use for downloading files in advance.
-                      This is enabled by default and uses 2 workers.
-                      To disable prefetching, set it to 0.
             namespace : namespace name.
             project : project name.
-            batch_rows : row limit per insert to balance speed and memory usage.
+            min_task_size : minimum number of tasks. (default=1)
+            batch_size : row limit per insert to balance speed and memory usage.
                       (default=2000)
 
         Example:
             ```py
             chain = (
                 chain
-                .settings(cache=True, parallel=8, batch_rows=300)
+                .settings(cache=True, parallel=8, batch_size=300)
                 .map(laion=process_webdataset(spec=WDSLaion), params="file")
             )
             ```
@@ -385,14 +385,14 @@ class DataChain:
         settings = copy.copy(self._settings)
         settings.add(
             Settings(
-                cache,
-                parallel,
-                workers,
-                min_task_size,
-                prefetch,
-                namespace,
-                project,
-                batch_rows,
+                cache=cache,
+                prefetch=prefetch,
+                parallel=parallel,
+                workers=workers,
+                namespace=namespace,
+                project=project,
+                min_task_size=min_task_size,
+                batch_size=batch_size,
             )
         )
         return self._evolve(settings=settings, _sys=sys)
@@ -745,7 +745,7 @@ class DataChain:
 
         return self._evolve(
             query=self._query.add_signals(
-                udf_obj.to_udf_wrapper(self._settings.batch_rows),
+                udf_obj.to_udf_wrapper(self._settings.batch_size),
                 **self._settings.to_dict(),
             ),
             signal_schema=self.signals_schema | udf_obj.output,
@@ -783,7 +783,7 @@ class DataChain:
             udf_obj.prefetch = prefetch
         return self._evolve(
             query=self._query.generate(
-                udf_obj.to_udf_wrapper(self._settings.batch_rows),
+                udf_obj.to_udf_wrapper(self._settings.batch_size),
                 **self._settings.to_dict(),
             ),
             signal_schema=udf_obj.output,
@@ -919,7 +919,7 @@ class DataChain:
         udf_obj = self._udf_to_obj(Aggregator, func, params, output, signal_map)
         return self._evolve(
             query=self._query.generate(
-                udf_obj.to_udf_wrapper(self._settings.batch_rows),
+                udf_obj.to_udf_wrapper(self._settings.batch_size),
                 partition_by=processed_partition_by,
                 **self._settings.to_dict(),
             ),
@@ -968,7 +968,7 @@ class DataChain:
 
         return self._evolve(
             query=self._query.add_signals(
-                udf_obj.to_udf_wrapper(self._settings.batch_rows, batch=batch),
+                udf_obj.to_udf_wrapper(self._settings.batch_size, batch=batch),
                 **self._settings.to_dict(),
             ),
             signal_schema=self.signals_schema | udf_obj.output,
@@ -2314,7 +2314,7 @@ class DataChain:
         table_name: str,
         connection: "ConnectionType",
         *,
-        batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
+        batch_size: int = DEFAULT_DATABASE_BATCH_SIZE,
         on_conflict: Optional[str] = None,
         conflict_columns: Optional[list[str]] = None,
         column_mapping: Optional[dict[str, Optional[str]]] = None,
@@ -2336,7 +2336,7 @@ class DataChain:
                 library. If a DBAPI2 object, only sqlite3 is supported. The user is
                 responsible for engine disposal and connection closure for the
                 SQLAlchemy connectable; str connections are closed automatically.
-            batch_rows: Number of rows to insert per batch for optimal performance.
+            batch_size: Number of rows to insert per batch for optimal performance.
                 Larger batches are faster but use more memory. Default: 10,000.
             on_conflict: Strategy for handling duplicate rows (requires table
                 constraints):
@@ -2417,7 +2417,7 @@ class DataChain:
             self,
             table_name,
             connection,
-            batch_rows=batch_rows,
+            batch_size=batch_size,
             on_conflict=on_conflict,
             conflict_columns=conflict_columns,
             column_mapping=column_mapping,

datachain/lib/dc/records.py

@@ -31,7 +31,7 @@ def read_records(
 
     Parameters:
         to_insert : records (or a single record) to insert. Each record is
-                    a dictionary of signals and theirs values.
+                    a dictionary of signals and their values.
         schema : describes chain signals and their corresponding types
 
     Example:
@@ -45,7 +45,6 @@ def read_records(
     """
     from datachain.query.dataset import adjust_outputs, get_col_types
     from datachain.sql.types import SQLType
-    from datachain.utils import batched
 
     from .datasets import read_dataset
 
@@ -96,7 +95,6 @@ def read_records(
         {c.name: c.type for c in columns if isinstance(c.type, SQLType)},
     )
     records = (adjust_outputs(warehouse, record, col_types) for record in to_insert)
-    for chunk in batched(records, READ_RECORDS_BATCH_SIZE):
-        warehouse.insert_rows(table, chunk)
+    warehouse.insert_rows(table, records, batch_size=READ_RECORDS_BATCH_SIZE)
     warehouse.insert_rows_done(table)
     return read_dataset(name=dsr.full_name, session=session, settings=settings)

datachain/lib/settings.py

@@ -1,111 +1,214 @@
+from typing import Any, Optional, Union
+
 from datachain.lib.utils import DataChainParamsError
-from datachain.utils import DEFAULT_CHUNK_ROWS
+
+DEFAULT_CACHE = False
+DEFAULT_PREFETCH = 2
+DEFAULT_BATCH_SIZE = 2_000
 
 
 class SettingsError(DataChainParamsError):
-    def __init__(self, msg):
+    def __init__(self, msg: str) -> None:
         super().__init__(f"Dataset settings error: {msg}")
 
 
 class Settings:
-    def __init__(
+    """Settings for datachain."""
+
+    _cache: Optional[bool]
+    _prefetch: Optional[int]
+    _parallel: Optional[Union[bool, int]]
+    _workers: Optional[int]
+    _namespace: Optional[str]
+    _project: Optional[str]
+    _min_task_size: Optional[int]
+    _batch_size: Optional[int]
+
+    def __init__(  # noqa: C901, PLR0912
         self,
-        cache=None,
-        parallel=None,
-        workers=None,
-        min_task_size=None,
-        prefetch=None,
-        namespace=None,
-        project=None,
-        batch_rows=None,
-    ):
-        self._cache = cache
-        self.parallel = parallel
-        self._workers = workers
-        self.min_task_size = min_task_size
-        self.prefetch = prefetch
-        self.namespace = namespace
-        self.project = project
-        self._chunk_rows = batch_rows
-
-        if not isinstance(cache, bool) and cache is not None:
-            raise SettingsError(
-                "'cache' argument must be bool"
-                f" while {cache.__class__.__name__} was given"
-            )
-
-        if not isinstance(parallel, int) and parallel is not None:
-            raise SettingsError(
-                "'parallel' argument must be int or None"
-                f" while {parallel.__class__.__name__} was given"
-            )
-
-        if (
-            not isinstance(workers, bool)
-            and not isinstance(workers, int)
-            and workers is not None
-        ):
-            raise SettingsError(
-                "'workers' argument must be int or bool"
-                f" while {workers.__class__.__name__} was given"
-            )
-
-        if min_task_size is not None and not isinstance(min_task_size, int):
-            raise SettingsError(
-                "'min_task_size' argument must be int or None"
-                f", {min_task_size.__class__.__name__} was given"
-            )
-
-        if batch_rows is not None and not isinstance(batch_rows, int):
-            raise SettingsError(
-                "'batch_rows' argument must be int or None"
-                f", {batch_rows.__class__.__name__} was given"
-            )
-
-        if batch_rows is not None and batch_rows <= 0:
-            raise SettingsError(
-                "'batch_rows' argument must be positive integer"
-                f", {batch_rows} was given"
-            )
+        cache: Optional[bool] = None,
+        prefetch: Optional[Union[bool, int]] = None,
+        parallel: Optional[Union[bool, int]] = None,
+        workers: Optional[int] = None,
+        namespace: Optional[str] = None,
+        project: Optional[str] = None,
+        min_task_size: Optional[int] = None,
+        batch_size: Optional[int] = None,
+    ) -> None:
+        if cache is None:
+            self._cache = None
+        else:
+            if not isinstance(cache, bool):
+                raise SettingsError(
+                    "'cache' argument must be bool"
+                    f" while {cache.__class__.__name__} was given"
+                )
+            self._cache = cache
+
+        if prefetch is None or prefetch is True:
+            self._prefetch = None
+        elif prefetch is False:
+            self._prefetch = 0  # disable prefetch (False == 0)
+        else:
+            if not isinstance(prefetch, int):
+                raise SettingsError(
+                    "'prefetch' argument must be int or bool"
+                    f" while {prefetch.__class__.__name__} was given"
+                )
+            if prefetch < 0:
+                raise SettingsError(
+                    "'prefetch' argument must be non-negative integer"
+                    f", {prefetch} was given"
+                )
+            self._prefetch = prefetch
+
+        if parallel is None or parallel is False:
+            self._parallel = None
+        elif parallel is True:
+            self._parallel = True
+        else:
+            if not isinstance(parallel, int):
+                raise SettingsError(
+                    "'parallel' argument must be int or bool"
+                    f" while {parallel.__class__.__name__} was given"
+                )
+            if parallel <= 0:
+                raise SettingsError(
+                    "'parallel' argument must be positive integer"
+                    f", {parallel} was given"
+                )
+            self._parallel = parallel
+
+        if workers is None:
+            self._workers = None
+        else:
+            if not isinstance(workers, int) or isinstance(workers, bool):
+                raise SettingsError(
+                    "'workers' argument must be int"
+                    f" while {workers.__class__.__name__} was given"
+                )
+            if workers <= 0:
+                raise SettingsError(
+                    f"'workers' argument must be positive integer, {workers} was given"
+                )
+            self._workers = workers
+
+        if namespace is None:
+            self._namespace = None
+        else:
+            if not isinstance(namespace, str):
+                raise SettingsError(
+                    "'namespace' argument must be str"
+                    f", {namespace.__class__.__name__} was given"
+                )
+            self._namespace = namespace
+
+        if project is None:
+            self._project = None
+        else:
+            if not isinstance(project, str):
+                raise SettingsError(
+                    "'project' argument must be str"
+                    f", {project.__class__.__name__} was given"
+                )
+            self._project = project
+
+        if min_task_size is None:
+            self._min_task_size = None
+        else:
+            if not isinstance(min_task_size, int) or isinstance(min_task_size, bool):
+                raise SettingsError(
+                    "'min_task_size' argument must be int"
+                    f", {min_task_size.__class__.__name__} was given"
+                )
+            if min_task_size <= 0:
+                raise SettingsError(
+                    "'min_task_size' argument must be positive integer"
+                    f", {min_task_size} was given"
+                )
+            self._min_task_size = min_task_size
+
+        if batch_size is None:
+            self._batch_size = None
+        else:
+            if not isinstance(batch_size, int) or isinstance(batch_size, bool):
+                raise SettingsError(
+                    "'batch_size' argument must be int"
+                    f", {batch_size.__class__.__name__} was given"
+                )
+            if batch_size <= 0:
+                raise SettingsError(
+                    "'batch_size' argument must be positive integer"
+                    f", {batch_size} was given"
+                )
+            self._batch_size = batch_size
+
+    @property
+    def cache(self) -> bool:
+        return self._cache if self._cache is not None else DEFAULT_CACHE
+
+    @property
+    def prefetch(self) -> Optional[int]:
+        return self._prefetch if self._prefetch is not None else DEFAULT_PREFETCH
+
+    @property
+    def parallel(self) -> Optional[Union[bool, int]]:
+        return self._parallel if self._parallel is not None else None
+
+    @property
+    def workers(self) -> Optional[int]:
+        return self._workers if self._workers is not None else None
+
+    @property
+    def namespace(self) -> Optional[str]:
+        return self._namespace if self._namespace is not None else None
 
     @property
-    def cache(self):
-        return self._cache if self._cache is not None else False
+    def project(self) -> Optional[str]:
+        return self._project if self._project is not None else None
 
     @property
-    def workers(self):
-        return self._workers if self._workers is not None else False
+    def min_task_size(self) -> Optional[int]:
+        return self._min_task_size if self._min_task_size is not None else None
 
     @property
-    def batch_rows(self):
-        return self._chunk_rows if self._chunk_rows is not None else DEFAULT_CHUNK_ROWS
+    def batch_size(self) -> int:
+        return self._batch_size if self._batch_size is not None else DEFAULT_BATCH_SIZE
 
-    def to_dict(self):
-        res = {}
+    def to_dict(self) -> dict[str, Any]:
+        res: dict[str, Any] = {}
         if self._cache is not None:
             res["cache"] = self.cache
-        if self.parallel is not None:
+        if self._prefetch is not None:
+            res["prefetch"] = self.prefetch
+        if self._parallel is not None:
             res["parallel"] = self.parallel
         if self._workers is not None:
             res["workers"] = self.workers
-        if self.min_task_size is not None:
+        if self._min_task_size is not None:
             res["min_task_size"] = self.min_task_size
-        if self.namespace is not None:
+        if self._namespace is not None:
             res["namespace"] = self.namespace
-        if self.project is not None:
+        if self._project is not None:
             res["project"] = self.project
-        if self._chunk_rows is not None:
-            res["batch_rows"] = self._chunk_rows
+        if self._batch_size is not None:
+            res["batch_size"] = self.batch_size
         return res
 
-    def add(self, settings: "Settings"):
-        self._cache = settings._cache or self._cache
-        self.parallel = settings.parallel or self.parallel
-        self._workers = settings._workers or self._workers
-        self.min_task_size = settings.min_task_size or self.min_task_size
-        self.namespace = settings.namespace or self.namespace
-        self.project = settings.project or self.project
-        if settings.prefetch is not None:
-            self.prefetch = settings.prefetch
-        if settings._chunk_rows is not None:
-            self._chunk_rows = settings._chunk_rows
+    def add(self, settings: "Settings") -> None:
+        if settings._cache is not None:
+            self._cache = settings._cache
+        if settings._prefetch is not None:
+            self._prefetch = settings._prefetch
+        if settings._parallel is not None:
+            self._parallel = settings._parallel
+        if settings._workers is not None:
+            self._workers = settings._workers
+        if settings._namespace is not None:
+            self._namespace = settings._namespace
+        if settings._project is not None:
+            self._project = settings._project
+        if settings._min_task_size is not None:
+            self._min_task_size = settings._min_task_size
+        if settings._batch_size is not None:
+            self._batch_size = settings._batch_size

datachain/lib/udf.py

@@ -54,23 +54,11 @@ UDFOutputSpec = Mapping[str, ColumnType]
 UDFResult = dict[str, Any]
 
 
-@attrs.define
-class UDFProperties:
-    udf: "UDFAdapter"
-
-    def get_batching(self, use_partitioning: bool = False) -> BatchingStrategy:
-        return self.udf.get_batching(use_partitioning)
-
-    @property
-    def batch_rows(self):
-        return self.udf.batch_rows
-
-
 @attrs.define(slots=False)
 class UDFAdapter:
     inner: "UDFBase"
     output: UDFOutputSpec
-    batch_rows: Optional[int] = None
+    batch_size: Optional[int] = None
     batch: int = 1
 
     def get_batching(self, use_partitioning: bool = False) -> BatchingStrategy:
@@ -83,11 +71,6 @@ class UDFAdapter:
             return Batch(self.batch)
         raise ValueError(f"invalid batch size {self.batch}")
 
-    @property
-    def properties(self):
-        # For backwards compatibility.
-        return UDFProperties(self)
-
     def run(
         self,
         udf_fields: "Sequence[str]",
@@ -237,13 +220,13 @@ class UDFBase(AbstractUDF):
 
     def to_udf_wrapper(
         self,
-        batch_rows: Optional[int] = None,
+        batch_size: Optional[int] = None,
         batch: int = 1,
     ) -> UDFAdapter:
         return UDFAdapter(
             self,
             self.output.to_udf_spec(),
-            batch_rows,
+            batch_size,
             batch,
         )
 

datachain/query/batch.py

@@ -81,8 +81,8 @@ class Batch(BatchingStrategy):
         # select rows in batches
         results = []
 
-        with contextlib.closing(execute(query, page_size=page_size)) as batch_rows:
-            for row in batch_rows:
+        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 :]

datachain/query/dataset.py

@@ -55,7 +55,6 @@ from datachain.query.udf import UdfInfo
 from datachain.sql.functions.random import rand
 from datachain.sql.types import SQLType
 from datachain.utils import (
-    batched,
     determine_processes,
     determine_workers,
     filtered_cloudpickle_dumps,
@@ -334,10 +333,10 @@ def process_udf_outputs(
     udf_results: Iterator[Iterable["UDFResult"]],
     udf: "UDFAdapter",
     cb: Callback = DEFAULT_CALLBACK,
+    batch_size: int = INSERT_BATCH_SIZE,
 ) -> None:
     # Optimization: Compute row types once, rather than for every row.
     udf_col_types = get_col_types(warehouse, udf.output)
-    batch_rows = udf.batch_rows or INSERT_BATCH_SIZE
 
     def _insert_rows():
         for udf_output in udf_results:
@@ -349,9 +348,7 @@ def process_udf_outputs(
                     cb.relative_update()
                     yield adjust_outputs(warehouse, row, udf_col_types)
 
-    for row_chunk in batched(_insert_rows(), batch_rows):
-        warehouse.insert_rows(udf_table, row_chunk)
-
+    warehouse.insert_rows(udf_table, _insert_rows(), batch_size=batch_size)
     warehouse.insert_rows_done(udf_table)
 
 
@@ -388,12 +385,13 @@ class UDFStep(Step, ABC):
     udf: "UDFAdapter"
     catalog: "Catalog"
     partition_by: Optional[PartitionByType] = None
+    is_generator = False
+    # Parameters from Settings
+    cache: bool = False
     parallel: Optional[int] = None
     workers: Union[bool, int] = False
     min_task_size: Optional[int] = None
-    is_generator = False
-    cache: bool = False
-    batch_rows: Optional[int] = None
+    batch_size: Optional[int] = None
 
     @abstractmethod
     def create_udf_table(self, query: Select) -> "Table":
@@ -450,6 +448,7 @@ class UDFStep(Step, ABC):
                         use_cache=self.cache,
                         is_generator=self.is_generator,
                         min_task_size=self.min_task_size,
+                        batch_size=self.batch_size,
                     )
                     udf_distributor()
                     return
@@ -486,6 +485,7 @@ class UDFStep(Step, ABC):
                         is_generator=self.is_generator,
                         cache=self.cache,
                         rows_total=rows_total,
+                        batch_size=self.batch_size or INSERT_BATCH_SIZE,
                     )
 
                     # Run the UDFDispatcher in another process to avoid needing
@@ -534,6 +534,7 @@ class UDFStep(Step, ABC):
                                 udf_results,
                                 self.udf,
                                 cb=generated_cb,
+                                batch_size=self.batch_size or INSERT_BATCH_SIZE,
                             )
                     finally:
                         download_cb.close()
@@ -595,7 +596,7 @@ class UDFStep(Step, ABC):
                 parallel=self.parallel,
                 workers=self.workers,
                 min_task_size=self.min_task_size,
-                batch_rows=self.batch_rows,
+                batch_size=self.batch_size,
             )
         return self.__class__(self.udf, self.catalog)
 
@@ -641,7 +642,16 @@ class UDFStep(Step, ABC):
 
 @frozen
 class UDFSignal(UDFStep):
+    udf: "UDFAdapter"
+    catalog: "Catalog"
+    partition_by: Optional[PartitionByType] = None
     is_generator = False
+    # Parameters from Settings
+    cache: bool = False
+    parallel: Optional[int] = None
+    workers: Union[bool, int] = False
+    min_task_size: Optional[int] = None
+    batch_size: Optional[int] = None
 
     def create_udf_table(self, query: Select) -> "Table":
         udf_output_columns: list[sqlalchemy.Column[Any]] = [
@@ -711,7 +721,16 @@ class UDFSignal(UDFStep):
 class RowGenerator(UDFStep):
     """Extend dataset with new rows."""
 
+    udf: "UDFAdapter"
+    catalog: "Catalog"
+    partition_by: Optional[PartitionByType] = None
     is_generator = True
+    # Parameters from Settings
+    cache: bool = False
+    parallel: Optional[int] = None
+    workers: Union[bool, int] = False
+    min_task_size: Optional[int] = None
+    batch_size: Optional[int] = None
 
     def create_udf_table(self, query: Select) -> "Table":
         warehouse = self.catalog.warehouse
@@ -1626,12 +1645,17 @@ class DatasetQuery:
     def add_signals(
         self,
         udf: "UDFAdapter",
+        partition_by: Optional[PartitionByType] = None,
+        # Parameters from Settings
+        cache: bool = False,
         parallel: Optional[int] = None,
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,
-        partition_by: Optional[PartitionByType] = None,
-        cache: bool = False,
-        batch_rows: Optional[int] = None,
+        batch_size: Optional[int] = None,
+        # Parameters are unused, kept only to match the signature of Settings.to_dict
+        prefetch: Optional[int] = None,
+        namespace: Optional[str] = None,
+        project: Optional[str] = None,
     ) -> "Self":
         """
         Adds one or more signals based on the results from the provided UDF.
@@ -1657,7 +1681,7 @@ class DatasetQuery:
                 workers=workers,
                 min_task_size=min_task_size,
                 cache=cache,
-                batch_rows=batch_rows,
+                batch_size=batch_size,
             )
         )
         return query
@@ -1672,14 +1696,17 @@ class DatasetQuery:
     def generate(
         self,
         udf: "UDFAdapter",
+        partition_by: Optional[PartitionByType] = None,
+        # Parameters from Settings
+        cache: bool = False,
         parallel: Optional[int] = None,
         workers: Union[bool, int] = False,
         min_task_size: Optional[int] = None,
-        partition_by: Optional[PartitionByType] = None,
+        batch_size: Optional[int] = None,
+        # Parameters are unused, kept only to match the signature of Settings.to_dict:
+        prefetch: Optional[int] = None,
         namespace: Optional[str] = None,
         project: Optional[str] = None,
-        cache: bool = False,
-        batch_rows: Optional[int] = None,
     ) -> "Self":
         query = self.clone()
         steps = query.steps
@@ -1692,7 +1719,7 @@ class DatasetQuery:
                 workers=workers,
                 min_task_size=min_task_size,
                 cache=cache,
-                batch_rows=batch_rows,
+                batch_size=batch_size,
             )
         )
         return query

datachain/query/dispatch.py

@@ -114,6 +114,7 @@ class UDFDispatcher:
         self.is_batching = udf_info["batching"].is_batching
         self.processes = udf_info["processes"]
         self.rows_total = udf_info["rows_total"]
+        self.batch_size = udf_info["batch_size"]
         self.buffer_size = buffer_size
         self.task_queue = None
         self.done_queue = None
@@ -142,6 +143,7 @@ class UDFDispatcher:
             self.table,
             self.cache,
             self.is_batching,
+            self.batch_size,
             self.udf_fields,
         )
 
@@ -232,6 +234,7 @@ class UDFDispatcher:
                     udf_results,
                     udf,
                     cb=generated_cb,
+                    batch_size=self.batch_size,
                 )
 
     def input_batch_size(self, n_workers: int) -> int:
@@ -385,6 +388,7 @@ class UDFWorker:
         table: "Table",
         cache: bool,
         is_batching: bool,
+        batch_size: int,
         udf_fields: Sequence[str],
     ) -> None:
         self.catalog = catalog
@@ -395,6 +399,7 @@ class UDFWorker:
         self.table = table
         self.cache = cache
         self.is_batching = is_batching
+        self.batch_size = batch_size
         self.udf_fields = udf_fields
 
         self.download_cb = DownloadCallback(self.done_queue)
@@ -420,6 +425,7 @@ class UDFWorker:
                     self.notify_and_process(udf_results),
                     self.udf,
                     cb=self.generated_cb,
+                    batch_size=self.batch_size,
                 )
         put_into_queue(self.done_queue, {"status": FINISHED_STATUS})
 

datachain/query/udf.py

@@ -21,6 +21,7 @@ class UdfInfo(TypedDict):
     is_generator: bool
     cache: bool
     rows_total: int
+    batch_size: int
 
 
 class AbstractUDFDistributor(ABC):
@@ -39,6 +40,7 @@ class AbstractUDFDistributor(ABC):
         use_cache: bool,
         is_generator: bool = False,
         min_task_size: Optional[Union[str, int]] = None,
+        batch_size: Optional[int] = None,
     ) -> None: ...
 
     @abstractmethod

datachain/utils.py

@@ -25,7 +25,7 @@ if TYPE_CHECKING:
     from typing_extensions import Self
 
 
-DEFAULT_CHUNK_ROWS = 2000
+DEFAULT_BATCH_SIZE = 2000
 
 logger = logging.getLogger("datachain")
 
@@ -228,7 +228,7 @@ _T_co = TypeVar("_T_co", covariant=True)
 
 def _dynamic_batched_core(
     iterable: Iterable[_T_co],
-    batch_rows: int,
+    batch_size: int,
 ) -> Iterator[list[_T_co]]:
     """Core batching logic that yields lists."""
 
@@ -236,7 +236,7 @@ def _dynamic_batched_core(
 
     for item in iterable:
         # Check if adding this item would exceed limits
-        if len(batch) >= batch_rows and batch:  # Yield current batch if we have one
+        if len(batch) >= batch_size and batch:  # Yield current batch if we have one
             yield batch
             batch = []
 
@@ -247,23 +247,22 @@ def _dynamic_batched_core(
         yield batch
 
 
-def batched(iterable: Iterable[_T_co], batch_rows: int) -> Iterator[tuple[_T_co, ...]]:
+def batched(iterable: Iterable[_T_co], batch_size: int) -> Iterator[tuple[_T_co, ...]]:
     """
-    Batch data into tuples of length batch_rows .
+    Batch data into tuples of length batch_size.
     The last batch may be shorter.
     """
-    yield from (tuple(batch) for batch in _dynamic_batched_core(iterable, batch_rows))
+    yield from (tuple(batch) for batch in _dynamic_batched_core(iterable, batch_size))
 
 
 def batched_it(
     iterable: Iterable[_T_co],
-    batch_rows: int = DEFAULT_CHUNK_ROWS,
+    batch_size: int = DEFAULT_BATCH_SIZE,
 ) -> Iterator[Iterator[_T_co]]:
     """
-    Batch data into iterators with dynamic sizing
-    based on row count and memory usage.
+    Batch data into iterators with dynamic sizing based on row count and memory usage.
     """
-    yield from (iter(batch) for batch in _dynamic_batched_core(iterable, batch_rows))
+    yield from (iter(batch) for batch in _dynamic_batched_core(iterable, batch_size))
 
 
 def flatten(items):

{datachain-0.30.6.dist-info → datachain-0.30.7.dist-info}/METADATA

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

{datachain-0.30.6.dist-info → datachain-0.30.7.dist-info}/RECORD

@@ -19,7 +19,7 @@ datachain/script_meta.py,sha256=V-LaFOZG84pD0Zc0NvejYdzwDgzITv6yHvAHggDCnuY,4978
 datachain/semver.py,sha256=UB8GHPBtAP3UJGeiuJoInD7SK-DnB93_Xd1qy_CQ9cU,2074
 datachain/studio.py,sha256=27750qCSNxIChEzhV02damIFreLMfr7UdiWqMFyk8AA,15361
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
-datachain/utils.py,sha256=RKe1-VuC9juQSIbIpMnELJ7QrsKQggj8l7Q8_FiCZHE,15664
+datachain/utils.py,sha256=5ehFeqXau7MFmGUQRsjRyPfDMPoOF1ojpfVciYUo5fE,15659
 datachain/catalog/__init__.py,sha256=9NBaywvAOaXdkyqiHjbBEiXs7JImR1OJsY9r8D5Q16g,403
 datachain/catalog/catalog.py,sha256=a1AN6eDHWWzII1wi46T_1JvTsW1AeMudwR_6sVQ4f7I,67588
 datachain/catalog/datasource.py,sha256=IkGMh0Ttg6Q-9DWfU_H05WUnZepbGa28HYleECi6K7I,1353
@@ -52,8 +52,8 @@ datachain/data_storage/job.py,sha256=ZkeXCNUj_VCkoKYx29hqB4AcfVUielnRjY-GYUcUxt4
 datachain/data_storage/metastore.py,sha256=aSeTRh43hmrOhULi9YD2VlgCj8B4bjE3jqCOvnb_HQs,53851
 datachain/data_storage/schema.py,sha256=o3JbURKXRg3IJyIVA4QjHHkn6byRuz7avbydU2FlvNY,9897
 datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2krunWUum5o,927
-datachain/data_storage/sqlite.py,sha256=edcTegzEoAEdEp62Rg9oERvHWXDcpg8d4onrD-P2xKM,30159
-datachain/data_storage/warehouse.py,sha256=sEbNiWKdB7yuLt88FuIfRur7U7WiOZrcHWhnBS_eMAg,32642
+datachain/data_storage/sqlite.py,sha256=1fIeIhmB3O8oQVzP8dDKap0KUIgI0n2TdBQSyv0R8J4,30345
+datachain/data_storage/warehouse.py,sha256=7jc69CtWdfQlc_9WbJ5l6yQooarpLFBrDk4fY-svi_0,32783
 datachain/diff/__init__.py,sha256=-OFZzgOplqO84iWgGY7kfe60NXaWR9JRIh9T-uJboAM,9668
 datachain/fs/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/fs/reference.py,sha256=A8McpXF0CqbXPqanXuvpKu50YLB3a2ZXA3YAPxtBXSM,914
@@ -85,11 +85,11 @@ datachain/lib/model_store.py,sha256=A0pSVQ7uaZ9RvANapzirF8Cqq9N6ysosPpMSkzdRPkU,
 datachain/lib/namespaces.py,sha256=I6gLC4ZzgyatFtHL85MWR4ml7-yuQOzxHE7IQNbt_ac,2107
 datachain/lib/projects.py,sha256=VJgmzHzKjmNPZD1tm0a1RNHmUQwn6WLWCLpKyc4UrSk,2605
 datachain/lib/pytorch.py,sha256=S-st2SAczYut13KMf6eSqP_OQ8otWI5TRmzhK5fN3k0,7828
-datachain/lib/settings.py,sha256=n0YYhCVdgCdMkCSLY7kscJF9mUhlQ0a4ENWBsJFynkw,3809
+datachain/lib/settings.py,sha256=xBQEPZfgaYKhHIFLd0u5CBTYDcJS8ZHCm47x7GJErFU,7666
 datachain/lib/signal_schema.py,sha256=YMMcc9gHIzBz88zfsreGa1nOoO_56HBtZlT6jf3V1WE,39224
 datachain/lib/tar.py,sha256=MLcVjzIgBqRuJacCNpZ6kwSZNq1i2tLyROc8PVprHsA,999
 datachain/lib/text.py,sha256=UNHm8fhidk7wdrWqacEWaA6I9ykfYqarQ2URby7jc7M,1261
-datachain/lib/udf.py,sha256=IB1IKF5KyA-NiyfhVzmBPpF_aITPS3zSlrt24f_Ofjo,17956
+datachain/lib/udf.py,sha256=08ia5T3gClen5ZQfIgop-swNnys2G-RIZpszqDnbc0w,17570
 datachain/lib/udf_signature.py,sha256=Yz20iJ-WF1pijT3hvcDIKFzgWV9gFxZM73KZRx3NbPk,7560
 datachain/lib/utils.py,sha256=RLji1gHnfDXtJCnBo8BcNu1obndFpVsXJ_1Vb-FQ9Qo,4554
 datachain/lib/video.py,sha256=ddVstiMkfxyBPDsnjCKY0d_93bw-DcMqGqN60yzsZoo,6851
@@ -103,15 +103,15 @@ datachain/lib/convert/unflatten.py,sha256=ysMkstwJzPMWUlnxn-Z-tXJR3wmhjHeSN_P-sD
 datachain/lib/convert/values_to_tuples.py,sha256=j5yZMrVUH6W7b-7yUvdCTGI7JCUAYUOzHUGPoyZXAB0,4360
 datachain/lib/dc/__init__.py,sha256=UrUzmDH6YyVl8fxM5iXTSFtl5DZTUzEYm1MaazK4vdQ,900
 datachain/lib/dc/csv.py,sha256=q6a9BpapGwP6nwy6c5cklxQumep2fUp9l2LAjtTJr6s,4411
-datachain/lib/dc/database.py,sha256=F6EOjPKwSdp26kJsOKGq49D9OxqyKEalINHEwLQav2s,14716
-datachain/lib/dc/datachain.py,sha256=2UtDhtBzx5VejkDE0UTS3t1517jCGr7YEKvO5wqNU-Q,99709
+datachain/lib/dc/database.py,sha256=sTpos1rE4BS5BTzzixykhWIO2JxVYKH1GTRncdpu4dU,14716
+datachain/lib/dc/datachain.py,sha256=AtsvBndqMyKrfW4yH8V0Nf__hfR0LN-NpA2munzfiPM,99888
 datachain/lib/dc/datasets.py,sha256=-Bvyyu4XXDXLiWa-bOnsp0Q11RSYXRO0j5DaX8ShaFs,15355
 datachain/lib/dc/hf.py,sha256=AP_MUHg6HJWae10PN9hD_beQVjrl0cleZ6Cvhtl1yoI,2901
 datachain/lib/dc/json.py,sha256=dNijfJ-H92vU3soyR7X1IiDrWhm6yZIGG3bSnZkPdAE,2733
 datachain/lib/dc/listings.py,sha256=V379Cb-7ZyquM0w7sWArQZkzInZy4GB7QQ1ZfowKzQY,4544
 datachain/lib/dc/pandas.py,sha256=ObueUXDUFKJGu380GmazdG02ARpKAHPhSaymfmOH13E,1489
 datachain/lib/dc/parquet.py,sha256=zYcSgrWwyEDW9UxGUSVdIVsCu15IGEf0xL8KfWQqK94,1782
-datachain/lib/dc/records.py,sha256=4N1Fq-j5r4GK-PR5jIO-9B2u_zTNX9l-6SmcRhQDAsw,3136
+datachain/lib/dc/records.py,sha256=IKf5MArify-cI1P4NgbIvrAi0UQ5cvofTI3u6_zKBP8,3069
 datachain/lib/dc/storage.py,sha256=OMJE-9ob9Ku5le8W6O8J1W-XJ0pwHt2PsO-ZCcee1ZA,7950
 datachain/lib/dc/utils.py,sha256=9OMiFu2kXIbtMqzJTEr1qbCoCBGpOmTnkWImVgFTKgo,4112
 datachain/lib/dc/values.py,sha256=7l1n352xWrEdql2NhBcZ3hj8xyPglWiY4qHjFPjn6iw,1428
@@ -125,15 +125,15 @@ datachain/model/ultralytics/bbox.py,sha256=C-aDiBhVa_ML2oERWvksRkyMU1XuYSpb6eItH
 datachain/model/ultralytics/pose.py,sha256=pvoXrWWUSWT_UBaMwUb5MBHAY57Co2HFDPigFYNZWUA,3392
 datachain/model/ultralytics/segment.py,sha256=v9_xDxd5zw_I8rXsbl7yQXgEdTs2T38zyY_Y4XGN8ok,3194
 datachain/query/__init__.py,sha256=7DhEIjAA8uZJfejruAVMZVcGFmvUpffuZJwgRqNwe-c,263
-datachain/query/batch.py,sha256=-goxLpE0EUvaDHu66rstj53UnfHpYfBUGux8GSpJ93k,4306
-datachain/query/dataset.py,sha256=OaGRBNSWYNaRbYn6avij0fiFN5DT-nwdM-wJ4yTfaYs,63317
-datachain/query/dispatch.py,sha256=f8IIvuLBJaCEwSRv7bWPMy1uXyc28W0LGqrBffjYf98,15831
+datachain/query/batch.py,sha256=ocPeNgrJM6Y_6SYCx3O2cwlCFAhNMfoYgB99GP6A1Bg,4294
+datachain/query/dataset.py,sha256=1eg5EE4vKI7c_Ng04or6zzKmFcOoEubMCoOaYmYPavE,64499
+datachain/query/dispatch.py,sha256=pygp7xg3lUDKlYHhecKxW5fB3zOSX1fPJfZBU4dfijk,16067
 datachain/query/metrics.py,sha256=DOK5HdNVaRugYPjl8qnBONvTkwjMloLqAr7Mi3TjCO0,858
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
 datachain/query/queue.py,sha256=v0UeK4ilmdiRoJ5OdjB5qpnHTYDxRP4vhVp5Iw_toaI,3512
 datachain/query/schema.py,sha256=qLpEyvnzKlNCOrThQiTNpUKTUEsVIHT9trt-0UMt6ko,6704
 datachain/query/session.py,sha256=gKblltJAVQAVSTswAgWGDgGbpmFlFzFVkIQojDCjgXM,6809
-datachain/query/udf.py,sha256=e753bDJzTNjGFQn1WGTvOAWSwjDbrFI1-_DDWkWN2ls,1343
+datachain/query/udf.py,sha256=jqutTpvkT6eHl96ZEgYiiTMAhI7vmTQA6JH9y4WCibI,1405
 datachain/query/utils.py,sha256=a2PTBZ3qsG6XlUcp9XsoGiQfKkca4Q3m-VzFgiGQPAc,1230
 datachain/remote/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/remote/studio.py,sha256=pDThxvEEpIKVGfa9rmtz_zeqHwrgzh0Lv-Pd4wzDx5k,15448
@@ -160,9 +160,9 @@ datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=ktGWzY4kyzjWyR86dhvzw-Zhl0lVk_LOX3NciTac6qo,2914
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.30.6.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.30.6.dist-info/METADATA,sha256=ZyXo8wdTrN08k--Soy3UHpCu_Jni_6ocO3_PbjCswCE,13898
-datachain-0.30.6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.30.6.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.30.6.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.30.6.dist-info/RECORD,,
+datachain-0.30.7.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.30.7.dist-info/METADATA,sha256=d6ClkSVhY7AFkjh7jgUFEwHpTa7LhpJU75_M8ufegcI,13898
+datachain-0.30.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+datachain-0.30.7.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.30.7.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.30.7.dist-info/RECORD,,