datachain 0.28.1__py3-none-any.whl → 0.29.0__py3-none-any.whl

datachain/data_storage/warehouse.py

@@ -21,6 +21,7 @@ from datachain.lib.file import File
 from datachain.lib.signal_schema import SignalSchema
 from datachain.node import DirType, DirTypeGroup, Node, NodeWithPath, get_path
 from datachain.query.batch import RowsOutput
+from datachain.query.schema import ColumnMeta
 from datachain.query.utils import get_query_id_column
 from datachain.sql.functions import path as pathfunc
 from datachain.sql.types import Int, SQLType
@@ -400,7 +401,7 @@ class AbstractWarehouse(ABC, Serializable):
         expressions: tuple[_ColumnsClauseArgument[Any], ...] = (
             sa.func.count(table.c.sys__id),
         )
-        size_column_names = [s.replace(".", "__") + "__size" for s in file_signals]
+        size_column_names = [ColumnMeta.to_db_name(s) + "__size" for s in file_signals]
         size_columns = [c for c in table.columns if c.name in size_column_names]
 
         if size_columns:

datachain/lib/dc/database.py

@@ -6,6 +6,10 @@ from typing import TYPE_CHECKING, Any, Optional, Union
 
 import sqlalchemy
 
+from datachain.query.schema import ColumnMeta
+
+DEFAULT_DATABASE_BATCH_SIZE = 10_000
+
 if TYPE_CHECKING:
     from collections.abc import Iterator, Mapping, Sequence
 
@@ -30,7 +34,7 @@ if TYPE_CHECKING:
 @contextlib.contextmanager
 def _connect(
     connection: "ConnectionType",
-) -> "Iterator[Union[sqlalchemy.engine.Connection, sqlalchemy.orm.Session]]":
+) -> "Iterator[sqlalchemy.engine.Connection]":
     import sqlalchemy.orm
 
     with contextlib.ExitStack() as stack:
@@ -47,27 +51,184 @@ def _connect(
             yield engine.connect()
         elif isinstance(connection, sqlalchemy.Engine):
             yield stack.enter_context(connection.connect())
-        elif isinstance(connection, (sqlalchemy.Connection, sqlalchemy.orm.Session)):
+        elif isinstance(connection, sqlalchemy.Connection):
             # do not close the connection, as it is managed by the caller
             yield connection
+        elif isinstance(connection, sqlalchemy.orm.Session):
+            # For Session objects, get the underlying bind (Engine or Connection)
+            # Sessions don't support DDL operations directly
+            bind = connection.get_bind()
+            if isinstance(bind, sqlalchemy.Engine):
+                yield stack.enter_context(bind.connect())
+            else:
+                # bind is already a Connection
+                yield bind
         else:
             raise TypeError(f"Unsupported connection type: {type(connection).__name__}")
 
 
-def _infer_schema(
-    result: "sqlalchemy.engine.Result",
-    to_infer: list[str],
-    infer_schema_length: Optional[int] = 100,
-) -> tuple[list["sqlalchemy.Row"], dict[str, "DataType"]]:
-    from datachain.lib.convert.values_to_tuples import values_to_tuples
+def to_database(
+    chain: "DataChain",
+    table_name: str,
+    connection: "ConnectionType",
+    *,
+    batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
+    on_conflict: Optional[str] = None,
+    column_mapping: Optional[dict[str, Optional[str]]] = None,
+) -> None:
+    """
+    Implementation function for exporting DataChain to database tables.
 
-    if not to_infer:
-        return [], {}
+    This is the core implementation that handles the actual database operations.
+    For user-facing documentation, see DataChain.to_database() method.
+    """
+    from datachain.utils import batched
 
-    rows = list(itertools.islice(result, infer_schema_length))
-    values = {col: [row._mapping[col] for row in rows] for col in to_infer}
-    _, output_schema, _ = values_to_tuples("", **values)
-    return rows, output_schema
+    if on_conflict and on_conflict not in ("ignore", "update"):
+        raise ValueError(
+            f"on_conflict must be 'ignore' or 'update', got: {on_conflict}"
+        )
+
+    signals_schema = chain.signals_schema.clone_without_sys_signals()
+    all_columns = [
+        sqlalchemy.Column(c.name, c.type)  # type: ignore[union-attr]
+        for c in signals_schema.db_signals(as_columns=True)
+    ]
+
+    column_mapping = column_mapping or {}
+    normalized_column_mapping = _normalize_column_mapping(column_mapping)
+    column_indices_and_names, columns = _prepare_columns(
+        all_columns, normalized_column_mapping
+    )
+
+    with _connect(connection) as conn:
+        metadata = sqlalchemy.MetaData()
+        table = sqlalchemy.Table(table_name, metadata, *columns)
+
+        # Check if table already exists to determine if we should clean up on error.
+        inspector = sqlalchemy.inspect(conn)
+        assert inspector  # to satisfy mypy
+        table_existed_before = table_name in inspector.get_table_names()
+
+        try:
+            table.create(conn, checkfirst=True)
+            rows_iter = chain._leaf_values()
+            for batch in batched(rows_iter, batch_rows):
+                _process_batch(
+                    conn, table, batch, on_conflict, column_indices_and_names
+                )
+            conn.commit()
+        except Exception:
+            if not table_existed_before:
+                try:
+                    table.drop(conn, checkfirst=True)
+                    conn.commit()
+                except sqlalchemy.exc.SQLAlchemyError:
+                    pass
+            raise
+
+
+def _normalize_column_mapping(
+    column_mapping: dict[str, Optional[str]],
+) -> dict[str, Optional[str]]:
+    """
+    Convert column mapping keys from DataChain format (dots) to database format
+    (double underscores).
+
+    This allows users to specify column mappings using the intuitive DataChain
+    format like: {"nested_data.value": "data_value"} instead of
+    {"nested_data__value": "data_value"}
+    """
+    if not column_mapping:
+        return {}
+
+    normalized_mapping: dict[str, Optional[str]] = {}
+    original_keys: dict[str, str] = {}
+    for key, value in column_mapping.items():
+        db_key = ColumnMeta.to_db_name(key)
+        if db_key in normalized_mapping:
+            prev = original_keys[db_key]
+            raise ValueError(
+                "Column mapping collision: multiple keys map to the same "
+                f"database column name '{db_key}': '{prev}' and '{key}'. "
+            )
+        normalized_mapping[db_key] = value
+        original_keys[db_key] = key
+
+    # If it's a defaultdict, preserve the default factory
+    if hasattr(column_mapping, "default_factory"):
+        from collections import defaultdict
+
+        default_factory = column_mapping.default_factory
+        result: dict[str, Optional[str]] = defaultdict(default_factory)
+        result.update(normalized_mapping)
+        return result
+
+    return normalized_mapping
+
+
+def _prepare_columns(all_columns, column_mapping):
+    """Prepare column mapping and column definitions."""
+    column_indices_and_names = []  # List of (index, target_name) tuples
+    columns = []
+    for idx, col in enumerate(all_columns):
+        if col.name in column_mapping or hasattr(column_mapping, "default_factory"):
+            mapped_name = column_mapping[col.name]
+            if mapped_name:
+                columns.append(sqlalchemy.Column(mapped_name, col.type))
+                column_indices_and_names.append((idx, mapped_name))
+        else:
+            columns.append(col)
+            column_indices_and_names.append((idx, col.name))
+    return column_indices_and_names, columns
+
+
+def _process_batch(conn, table, batch, on_conflict, column_indices_and_names):
+    """Process a batch of rows with conflict resolution."""
+
+    def prepare_row(row_values):
+        """Convert a row tuple to a dictionary with proper DB column names."""
+        return {
+            target_name: row_values[idx]
+            for idx, target_name in column_indices_and_names
+        }
+
+    rows_to_insert = [prepare_row(row) for row in batch]
+
+    supports_conflict = on_conflict and conn.engine.name in ("postgresql", "sqlite")
+
+    if supports_conflict:
+        # Use dialect-specific insert for conflict resolution
+        if conn.engine.name == "postgresql":
+            from sqlalchemy.dialects.postgresql import insert as pg_insert
+
+            insert_stmt = pg_insert(table)
+        elif conn.engine.name == "sqlite":
+            from sqlalchemy.dialects.sqlite import insert as sqlite_insert
+
+            insert_stmt = sqlite_insert(table)
+    else:
+        insert_stmt = table.insert()
+
+    if supports_conflict:
+        if on_conflict == "ignore":
+            insert_stmt = insert_stmt.on_conflict_do_nothing()
+        elif on_conflict == "update":
+            update_values = {
+                col.name: insert_stmt.excluded[col.name] for col in table.columns
+            }
+            insert_stmt = insert_stmt.on_conflict_do_update(set_=update_values)
+    elif on_conflict:
+        import warnings
+
+        warnings.warn(
+            f"Database does not support conflict resolution. "
+            f"Ignoring on_conflict='{on_conflict}' parameter.",
+            UserWarning,
+            stacklevel=2,
+        )
+
+    conn.execute(insert_stmt, rows_to_insert)
 
 
 def read_database(
@@ -151,3 +312,19 @@ def read_database(
             in_memory=in_memory,
             schema=inferred_schema | output,
         )
+
+
+def _infer_schema(
+    result: "sqlalchemy.engine.Result",
+    to_infer: list[str],
+    infer_schema_length: Optional[int] = 100,
+) -> tuple[list["sqlalchemy.Row"], dict[str, "DataType"]]:
+    from datachain.lib.convert.values_to_tuples import values_to_tuples
+
+    if not to_infer:
+        return [], {}
+
+    rows = list(itertools.islice(result, infer_schema_length))
+    values = {col: [row._mapping[col] for row in rows] for col in to_infer}
+    _, output_schema, _ = values_to_tuples("", **values)
+    return rows, output_schema

datachain/lib/dc/datachain.py

@@ -58,6 +58,7 @@ from datachain.query.schema import DEFAULT_DELIMITER, Column
 from datachain.sql.functions import path as pathfunc
 from datachain.utils import batched_it, inside_notebook, row_to_nested_dict
 
+from .database import DEFAULT_DATABASE_BATCH_SIZE
 from .utils import (
     DatasetMergeError,
     DatasetPrepareError,
@@ -77,11 +78,23 @@ UDFObjT = TypeVar("UDFObjT", bound=UDFBase)
 DEFAULT_PARQUET_CHUNK_SIZE = 100_000
 
 if TYPE_CHECKING:
+    import sqlite3
+
     import pandas as pd
     from typing_extensions import ParamSpec, Self
 
     P = ParamSpec("P")
 
+    ConnectionType = Union[
+        str,
+        sqlalchemy.engine.URL,
+        sqlalchemy.engine.interfaces.Connectable,
+        sqlalchemy.engine.Engine,
+        sqlalchemy.engine.Connection,
+        "sqlalchemy.orm.Session",
+        sqlite3.Connection,
+    ]
+
 
 T = TypeVar("T", bound="DataChain")
 
@@ -324,6 +337,7 @@ class DataChain:
         sys: Optional[bool] = None,
         namespace: Optional[str] = None,
         project: Optional[str] = None,
+        batch_rows: Optional[int] = None,
     ) -> "Self":
         """Change settings for chain.
 
@@ -331,22 +345,24 @@ class DataChain:
         It returns chain, so, it can be chained later with next operation.
 
         Parameters:
-            cache : data caching (default=False)
+            cache : data caching. (default=False)
             parallel : number of thread for processors. True is a special value to
-                enable all available CPUs (default=1)
+                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.
+            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.
+            namespace : namespace name.
+            project : project name.
+            batch_rows : row limit per insert to balance speed and memory usage.
+                      (default=2000)
 
         Example:
             ```py
             chain = (
                 chain
-                .settings(cache=True, parallel=8)
+                .settings(cache=True, parallel=8, batch_rows=300)
                 .map(laion=process_webdataset(spec=WDSLaion), params="file")
             )
             ```
@@ -356,7 +372,14 @@ class DataChain:
         settings = copy.copy(self._settings)
         settings.add(
             Settings(
-                cache, parallel, workers, min_task_size, prefetch, namespace, project
+                cache,
+                parallel,
+                workers,
+                min_task_size,
+                prefetch,
+                namespace,
+                project,
+                batch_rows,
             )
         )
         return self._evolve(settings=settings, _sys=sys)
@@ -711,7 +734,7 @@ class DataChain:
 
         return self._evolve(
             query=self._query.add_signals(
-                udf_obj.to_udf_wrapper(),
+                udf_obj.to_udf_wrapper(self._settings.batch_rows),
                 **self._settings.to_dict(),
             ),
             signal_schema=self.signals_schema | udf_obj.output,
@@ -749,7 +772,7 @@ class DataChain:
             udf_obj.prefetch = prefetch
         return self._evolve(
             query=self._query.generate(
-                udf_obj.to_udf_wrapper(),
+                udf_obj.to_udf_wrapper(self._settings.batch_rows),
                 **self._settings.to_dict(),
             ),
             signal_schema=udf_obj.output,
@@ -885,7 +908,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(),
+                udf_obj.to_udf_wrapper(self._settings.batch_rows),
                 partition_by=processed_partition_by,
                 **self._settings.to_dict(),
             ),
@@ -917,11 +940,24 @@ class DataChain:
             )
             chain.save("new_dataset")
             ```
+
+        .. deprecated:: 0.29.0
+            This method is deprecated and will be removed in a future version.
+            Use `agg()` instead, which provides the similar functionality.
         """
+        import warnings
+
+        warnings.warn(
+            "batch_map() is deprecated and will be removed in a future version. "
+            "Use agg() instead, which provides the similar functionality.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         udf_obj = self._udf_to_obj(BatchMapper, func, params, output, signal_map)
+
         return self._evolve(
             query=self._query.add_signals(
-                udf_obj.to_udf_wrapper(batch),
+                udf_obj.to_udf_wrapper(self._settings.batch_rows, batch=batch),
                 **self._settings.to_dict(),
             ),
             signal_schema=self.signals_schema | udf_obj.output,
@@ -2253,6 +2289,97 @@ class DataChain:
         """
         self.to_json(path, fs_kwargs, include_outer_list=False)
 
+    def to_database(
+        self,
+        table_name: str,
+        connection: "ConnectionType",
+        *,
+        batch_rows: int = DEFAULT_DATABASE_BATCH_SIZE,
+        on_conflict: Optional[str] = None,
+        column_mapping: Optional[dict[str, Optional[str]]] = None,
+    ) -> None:
+        """Save chain to a database table using a given database connection.
+
+        This method exports all DataChain records to a database table, creating the
+        table if it doesn't exist and appending data if it does. The table schema
+        is automatically inferred from the DataChain's signal schema.
+
+        Parameters:
+            table_name: Name of the database table to create/write to.
+            connection: SQLAlchemy connectable, str, or a sqlite3 connection
+                Using SQLAlchemy makes it possible to use any DB supported by that
+                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.
+                Larger batches are faster but use more memory. Default: 10,000.
+            on_conflict: Strategy for handling duplicate rows (requires table
+                constraints):
+                - None: Raise error (`sqlalchemy.exc.IntegrityError`) on conflict
+                  (default)
+                - "ignore": Skip duplicate rows silently
+                - "update": Update existing rows with new values
+            column_mapping: Optional mapping to rename or skip columns:
+                - Dict mapping DataChain column names to database column names
+                - Set values to None to skip columns entirely, or use `defaultdict` to
+                  skip all columns except those specified.
+
+        Examples:
+            Basic usage with PostgreSQL:
+            ```py
+            import sqlalchemy as sa
+            import datachain as dc
+
+            chain = dc.read_storage("s3://my-bucket/")
+            engine = sa.create_engine("postgresql://user:pass@localhost/mydb")
+            chain.to_database("files_table", engine)
+            ```
+
+            Using SQLite with connection string:
+            ```py
+            chain.to_database("my_table", "sqlite:///data.db")
+            ```
+
+            Column mapping and renaming:
+            ```py
+            mapping = {
+                "user.id": "id",
+                "user.name": "name",
+                "user.password": None  # Skip this column
+            }
+            chain.to_database("users", engine, column_mapping=mapping)
+            ```
+
+            Handling conflicts (requires PRIMARY KEY or UNIQUE constraints):
+            ```py
+            # Skip duplicates
+            chain.to_database("my_table", engine, on_conflict="ignore")
+
+            # Update existing records
+            chain.to_database("my_table", engine, on_conflict="update")
+            ```
+
+            Working with different databases:
+            ```py
+            # MySQL
+            mysql_engine = sa.create_engine("mysql+pymysql://user:pass@host/db")
+            chain.to_database("mysql_table", mysql_engine)
+
+            # SQLite in-memory
+            chain.to_database("temp_table", "sqlite:///:memory:")
+            ```
+        """
+        from .database import to_database
+
+        to_database(
+            self,
+            table_name,
+            connection,
+            batch_rows=batch_rows,
+            on_conflict=on_conflict,
+            column_mapping=column_mapping,
+        )
+
     @classmethod
     def from_records(
         cls,
@@ -2340,7 +2467,7 @@ class DataChain:
     def setup(self, **kwargs) -> "Self":
         """Setup variables to pass to UDF functions.
 
-        Use before running map/gen/agg/batch_map to save an object and pass it as an
+        Use before running map/gen/agg to save an object and pass it as an
         argument to the UDF.
 
         The value must be a callable (a `lambda: <value>` syntax can be used to quickly

datachain/lib/dc/records.py

@@ -15,6 +15,8 @@ if TYPE_CHECKING:
 
     P = ParamSpec("P")
 
+READ_RECORDS_BATCH_SIZE = 10000
+
 
 def read_records(
     to_insert: Optional[Union[dict, Iterable[dict]]],
@@ -41,7 +43,7 @@ def read_records(
     Notes:
         This call blocks until all records are inserted.
     """
-    from datachain.query.dataset import INSERT_BATCH_SIZE, adjust_outputs, get_col_types
+    from datachain.query.dataset import adjust_outputs, get_col_types
     from datachain.sql.types import SQLType
     from datachain.utils import batched
 
@@ -94,7 +96,7 @@ 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, INSERT_BATCH_SIZE):
+    for chunk in batched(records, READ_RECORDS_BATCH_SIZE):
         warehouse.insert_rows(table, chunk)
     warehouse.insert_rows_done(table)
     return read_dataset(name=dsr.full_name, session=session, settings=settings)

datachain/lib/settings.py

@@ -1,4 +1,5 @@
 from datachain.lib.utils import DataChainParamsError
+from datachain.utils import DEFAULT_CHUNK_ROWS
 
 
 class SettingsError(DataChainParamsError):
@@ -16,6 +17,7 @@ class Settings:
         prefetch=None,
         namespace=None,
         project=None,
+        batch_rows=None,
     ):
         self._cache = cache
         self.parallel = parallel
@@ -24,6 +26,7 @@ class Settings:
         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(
@@ -53,6 +56,18 @@ class Settings:
                 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"
+            )
+
     @property
     def cache(self):
         return self._cache if self._cache is not None else False
@@ -61,6 +76,10 @@ class Settings:
     def workers(self):
         return self._workers if self._workers is not None else False
 
+    @property
+    def batch_rows(self):
+        return self._chunk_rows if self._chunk_rows is not None else DEFAULT_CHUNK_ROWS
+
     def to_dict(self):
         res = {}
         if self._cache is not None:
@@ -75,6 +94,8 @@ class Settings:
             res["namespace"] = self.namespace
         if self.project is not None:
             res["project"] = self.project
+        if self._chunk_rows is not None:
+            res["batch_rows"] = self._chunk_rows
         return res
 
     def add(self, settings: "Settings"):
@@ -86,3 +107,5 @@ class Settings:
         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

datachain/lib/signal_schema.py

@@ -34,7 +34,7 @@ from datachain.lib.data_model import DataModel, DataType, DataValue
 from datachain.lib.file import File
 from datachain.lib.model_store import ModelStore
 from datachain.lib.utils import DataChainParamsError
-from datachain.query.schema import DEFAULT_DELIMITER, Column
+from datachain.query.schema import DEFAULT_DELIMITER, Column, ColumnMeta
 from datachain.sql.types import SQLType
 
 if TYPE_CHECKING:
@@ -590,7 +590,7 @@ class SignalSchema:
 
         if name:
             if "." in name:
-                name = name.replace(".", "__")
+                name = ColumnMeta.to_db_name(name)
 
             signals = [
                 s

datachain/lib/udf.py

@@ -62,19 +62,21 @@ class UDFProperties:
         return self.udf.get_batching(use_partitioning)
 
     @property
-    def batch(self):
-        return self.udf.batch
+    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: 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:
@@ -233,10 +235,15 @@ class UDFBase(AbstractUDF):
     def signal_names(self) -> Iterable[str]:
         return self.output.to_udf_spec().keys()
 
-    def to_udf_wrapper(self, batch: int = 1) -> UDFAdapter:
+    def to_udf_wrapper(
+        self,
+        batch_rows: Optional[int] = None,
+        batch: int = 1,
+    ) -> UDFAdapter:
         return UDFAdapter(
             self,
             self.output.to_udf_spec(),
+            batch_rows,
             batch,
         )
 
@@ -418,11 +425,27 @@ class Mapper(UDFBase):
 
 
 class BatchMapper(UDFBase):
-    """Inherit from this class to pass to `DataChain.batch_map()`."""
+    """Inherit from this class to pass to `DataChain.batch_map()`.
+
+    .. deprecated:: 0.29.0
+        This class is deprecated and will be removed in a future version.
+        Use `Aggregator` instead, which provides the similar functionality.
+    """
 
     is_input_batched = True
     is_output_batched = True
 
+    def __init__(self):
+        import warnings
+
+        warnings.warn(
+            "BatchMapper is deprecated and will be removed in a future version. "
+            "Use Aggregator instead, which provides the similar functionality.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__()
+
     def run(
         self,
         udf_fields: Sequence[str],

datachain/query/dataset.py

@@ -333,32 +333,24 @@ def process_udf_outputs(
     udf_table: "Table",
     udf_results: Iterator[Iterable["UDFResult"]],
     udf: "UDFAdapter",
-    batch_size: int = INSERT_BATCH_SIZE,
     cb: Callback = DEFAULT_CALLBACK,
 ) -> None:
-    import psutil
-
-    rows: list[UDFResult] = []
     # 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
 
-    for udf_output in udf_results:
-        if not udf_output:
-            continue
-        with safe_closing(udf_output):
-            for row in udf_output:
-                cb.relative_update()
-                rows.append(adjust_outputs(warehouse, row, udf_col_types))
-                if len(rows) >= batch_size or (
-                    len(rows) % 10 == 0 and psutil.virtual_memory().percent > 80
-                ):
-                    for row_chunk in batched(rows, batch_size):
-                        warehouse.insert_rows(udf_table, row_chunk)
-                    rows.clear()
+    def _insert_rows():
+        for udf_output in udf_results:
+            if not udf_output:
+                continue
+
+            with safe_closing(udf_output):
+                for row in udf_output:
+                    cb.relative_update()
+                    yield adjust_outputs(warehouse, row, udf_col_types)
 
-    if rows:
-        for row_chunk in batched(rows, batch_size):
-            warehouse.insert_rows(udf_table, row_chunk)
+    for row_chunk in batched(_insert_rows(), batch_rows):
+        warehouse.insert_rows(udf_table, row_chunk)
 
     warehouse.insert_rows_done(udf_table)
 
@@ -401,6 +393,7 @@ class UDFStep(Step, ABC):
     min_task_size: Optional[int] = None
     is_generator = False
     cache: bool = False
+    batch_rows: Optional[int] = None
 
     @abstractmethod
     def create_udf_table(self, query: Select) -> "Table":
@@ -602,6 +595,7 @@ class UDFStep(Step, ABC):
                 parallel=self.parallel,
                 workers=self.workers,
                 min_task_size=self.min_task_size,
+                batch_rows=self.batch_rows,
             )
         return self.__class__(self.udf, self.catalog)
 
@@ -1633,6 +1627,7 @@ class DatasetQuery:
         min_task_size: Optional[int] = None,
         partition_by: Optional[PartitionByType] = None,
         cache: bool = False,
+        batch_rows: Optional[int] = None,
     ) -> "Self":
         """
         Adds one or more signals based on the results from the provided UDF.
@@ -1658,6 +1653,7 @@ class DatasetQuery:
                 workers=workers,
                 min_task_size=min_task_size,
                 cache=cache,
+                batch_rows=batch_rows,
             )
         )
         return query
@@ -1679,6 +1675,7 @@ class DatasetQuery:
         namespace: Optional[str] = None,
         project: Optional[str] = None,
         cache: bool = False,
+        batch_rows: Optional[int] = None,
     ) -> "Self":
         query = self.clone()
         steps = query.steps
@@ -1691,6 +1688,7 @@ class DatasetQuery:
                 workers=workers,
                 min_task_size=min_task_size,
                 cache=cache,
+                batch_rows=batch_rows,
             )
         )
         return query

datachain/utils.py

@@ -11,7 +11,6 @@ import time
 from collections.abc import Iterable, Iterator, Sequence
 from contextlib import contextmanager
 from datetime import date, datetime, timezone
-from itertools import chain, islice
 from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
 from uuid import UUID
 
@@ -26,6 +25,8 @@ if TYPE_CHECKING:
     from typing_extensions import Self
 
 
+DEFAULT_CHUNK_ROWS = 2000
+
 logger = logging.getLogger("datachain")
 
 NUL = b"\0"
@@ -225,30 +226,44 @@ def get_envs_by_prefix(prefix: str) -> dict[str, str]:
 _T_co = TypeVar("_T_co", covariant=True)
 
 
-def batched(iterable: Iterable[_T_co], n: int) -> Iterator[tuple[_T_co, ...]]:
-    """Batch data into tuples of length n. The last batch may be shorter."""
-    # Based on: https://docs.python.org/3/library/itertools.html#itertools-recipes
-    # batched('ABCDEFG', 3) --> ABC DEF G
-    if n < 1:
-        raise ValueError("Batch size must be at least one")
-    it = iter(iterable)
-    while batch := tuple(islice(it, n)):
+def _dynamic_batched_core(
+    iterable: Iterable[_T_co],
+    batch_rows: int,
+) -> Iterator[list[_T_co]]:
+    """Core batching logic that yields lists."""
+
+    batch: list[_T_co] = []
+
+    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
+            yield batch
+            batch = []
+
+        batch.append(item)
+
+    # Yield any remaining items
+    if batch:
         yield batch
 
 
-def batched_it(iterable: Iterable[_T_co], n: int) -> Iterator[Iterator[_T_co]]:
-    """Batch data into iterators of length n. The last batch may be shorter."""
-    # batched('ABCDEFG', 3) --> ABC DEF G
-    if n < 1:
-        raise ValueError("Batch size must be at least one")
-    it = iter(iterable)
-    while True:
-        chunk_it = islice(it, n)
-        try:
-            first_el = next(chunk_it)
-        except StopIteration:
-            return
-        yield chain((first_el,), chunk_it)
+def batched(iterable: Iterable[_T_co], batch_rows: int) -> Iterator[tuple[_T_co, ...]]:
+    """
+    Batch data into tuples of length batch_rows .
+    The last batch may be shorter.
+    """
+    yield from (tuple(batch) for batch in _dynamic_batched_core(iterable, batch_rows))
+
+
+def batched_it(
+    iterable: Iterable[_T_co],
+    batch_rows: int = DEFAULT_CHUNK_ROWS,
+) -> Iterator[Iterator[_T_co]]:
+    """
+    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))
 
 
 def flatten(items):

{datachain-0.28.1.dist-info → datachain-0.29.0.dist-info}/METADATA

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

{datachain-0.28.1.dist-info → datachain-0.29.0.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=-BmKLVNBLPFveUgVVE2So3aaiGndO2jK2qbHZ0zBDd8,15239
 datachain/telemetry.py,sha256=0A4IOPPp9VlP5pyW9eBfaTK3YhHGzHl7dQudQjUAx9A,994
-datachain/utils.py,sha256=DNqOi-Ydb7InyWvD9m7_yailxz6-YGpZzh00biQaHNo,15305
+datachain/utils.py,sha256=Gp5JVr_m7nVWQGDOjrGnZjRXF9-Ai-MBxiPJIcpPvWQ,15451
 datachain/catalog/__init__.py,sha256=cMZzSz3VoUi-6qXSVaHYN-agxQuAcz2XSqnEPZ55crE,353
 datachain/catalog/catalog.py,sha256=QTWCXy75iWo-0MCXyfV_WbsKeZ1fpLpvL8d60rxn1ws,65528
 datachain/catalog/datasource.py,sha256=IkGMh0Ttg6Q-9DWfU_H05WUnZepbGa28HYleECi6K7I,1353
@@ -53,7 +53,7 @@ datachain/data_storage/metastore.py,sha256=Qw332arvhgXB4UY0yX-Hu8Vgl3smU12l6bvxr
 datachain/data_storage/schema.py,sha256=o3JbURKXRg3IJyIVA4QjHHkn6byRuz7avbydU2FlvNY,9897
 datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2krunWUum5o,927
 datachain/data_storage/sqlite.py,sha256=TTQjdDXUaZSr3MEaxZjDhsVIkIJqxFNA-sD25TO3m_4,30228
-datachain/data_storage/warehouse.py,sha256=nhF8yfpdJpstpXnv_sj7WFzU97JkvSeqetqJQp33cyE,32563
+datachain/data_storage/warehouse.py,sha256=66PETLzfkgSmj-EF604m62xmFMQBXaRZSw8sdKGMam8,32613
 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=dkL2rcT5ag-kbgkhQPL_byEs-TCYr29qvdltroL5NxM,
 datachain/lib/namespaces.py,sha256=it52UbbwB8dzhesO2pMs_nThXiPQ1Ph9sD9I3GQkg5s,2099
 datachain/lib/projects.py,sha256=8lN0qV8czX1LGtWURCUvRlSJk-RpO9w9Rra_pOZus6g,2595
 datachain/lib/pytorch.py,sha256=S-st2SAczYut13KMf6eSqP_OQ8otWI5TRmzhK5fN3k0,7828
-datachain/lib/settings.py,sha256=9wi0FoHxRxNiyn99pR28IYsMkoo47jQxeXuObQr2Ar0,2929
-datachain/lib/signal_schema.py,sha256=JMsL8c4iCRH9PoRumvjimsOLQQslTjm_aDR2jh1zT2Q,38558
+datachain/lib/settings.py,sha256=n0YYhCVdgCdMkCSLY7kscJF9mUhlQ0a4ENWBsJFynkw,3809
+datachain/lib/signal_schema.py,sha256=FmsfEAdRDeAzv1ApQnRXzkkyNeY9fTaXpjMzSMhDh7M,38574
 datachain/lib/tar.py,sha256=MLcVjzIgBqRuJacCNpZ6kwSZNq1i2tLyROc8PVprHsA,999
 datachain/lib/text.py,sha256=UNHm8fhidk7wdrWqacEWaA6I9ykfYqarQ2URby7jc7M,1261
-datachain/lib/udf.py,sha256=SUnJWRDC3TlLhvpi8iqqJbeZGn5DChot7DyH-0Q-z20,17305
+datachain/lib/udf.py,sha256=IB1IKF5KyA-NiyfhVzmBPpF_aITPS3zSlrt24f_Ofjo,17956
 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=TFci5HTvYGjBesNUxDAnXaX36PnzPEUSn5a6JxB9o0U,872
 datachain/lib/dc/csv.py,sha256=q6a9BpapGwP6nwy6c5cklxQumep2fUp9l2LAjtTJr6s,4411
-datachain/lib/dc/database.py,sha256=g5M6NjYR1T0vKte-abV-3Ejnm-HqxTIMir5cRi_SziE,6051
-datachain/lib/dc/datachain.py,sha256=U2CV8-ewfu-sW1D2BysdqCtbnEA7uNL1ZhYLWPAFB1o,93298
+datachain/lib/dc/database.py,sha256=MPE-KzwcR2DhWLCEbl1gWFp63dLqjWuiJ1iEfC2BrJI,12443
+datachain/lib/dc/datachain.py,sha256=_C9PZjUHVewpdp94AR2GS3QEI96Svsyx52dLJVM4tm4,98143
 datachain/lib/dc/datasets.py,sha256=P6CIJizD2IYFwOQG5D3VbQRjDmUiRH0ysdtb551Xdm8,15098
 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=FpPbApWopUri1gIaSMsfXN4fevja4mjmfb6Q5eiaGxI,3116
+datachain/lib/dc/records.py,sha256=4N1Fq-j5r4GK-PR5jIO-9B2u_zTNX9l-6SmcRhQDAsw,3136
 datachain/lib/dc/storage.py,sha256=FXroEdxOZfbuEBIWfWTkbGwrI0D4_mrLZSRsIQm0WFE,7693
 datachain/lib/dc/utils.py,sha256=VawOAlJSvAtZbsMg33s5tJe21TRx1Km3QggI1nN6tnw,3984
 datachain/lib/dc/values.py,sha256=7l1n352xWrEdql2NhBcZ3hj8xyPglWiY4qHjFPjn6iw,1428
@@ -126,7 +126,7 @@ datachain/model/ultralytics/pose.py,sha256=pBlmt63Qe68FKmexHimUGlNbNOoOlMHXG4fzX
 datachain/model/ultralytics/segment.py,sha256=63bDCj43E6iZ0hFI5J6uQfksdCmjEp6sEm1XzVaE8pw,2986
 datachain/query/__init__.py,sha256=7DhEIjAA8uZJfejruAVMZVcGFmvUpffuZJwgRqNwe-c,263
 datachain/query/batch.py,sha256=-goxLpE0EUvaDHu66rstj53UnfHpYfBUGux8GSpJ93k,4306
-datachain/query/dataset.py,sha256=cYNrg1QyrZpO-oup3mqmSYHUvgEYBKe8RgkVbyQa6p0,62777
+datachain/query/dataset.py,sha256=OJZ_YwpS5i4B0wVmosMmMNW1qABr6zyOmqNHQdAWir4,62704
 datachain/query/dispatch.py,sha256=A0nPxn6mEN5d9dDo6S8m16Ji_9IvJLXrgF2kqXdi4fs,15546
 datachain/query/metrics.py,sha256=DOK5HdNVaRugYPjl8qnBONvTkwjMloLqAr7Mi3TjCO0,858
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
@@ -158,9 +158,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.28.1.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.28.1.dist-info/METADATA,sha256=9rZc1mFjNj6S3v6FjgrhM7bUdi6kO_5606CB7HQCfeo,13766
-datachain-0.28.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.28.1.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.28.1.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.28.1.dist-info/RECORD,,
+datachain-0.29.0.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.29.0.dist-info/METADATA,sha256=g5YmnSXxBvUz_ZO1ZoEPHkzRyQGW5ZbPc8a4ZRJqHXE,13766
+datachain-0.29.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+datachain-0.29.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.29.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.29.0.dist-info/RECORD,,