datachain 0.28.1__tar.gz → 0.29.0__tar.gz

{datachain-0.28.1 → datachain-0.29.0}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.12.5'
+    rev: 'v0.12.7'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.28.1 → datachain-0.29.0}/PKG-INFO

@@ -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 → datachain-0.29.0}/src/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-0.29.0/src/datachain/lib/dc/database.py

@@ -0,0 +1,330 @@
+import contextlib
+import itertools
+import os
+import sqlite3
+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
+
+    import sqlalchemy.orm  # noqa: TC004
+
+    from datachain.lib.data_model import DataType
+    from datachain.query import Session
+
+    from .datachain import DataChain
+
+    ConnectionType = Union[
+        str,
+        sqlalchemy.engine.URL,
+        sqlalchemy.engine.interfaces.Connectable,
+        sqlalchemy.engine.Engine,
+        sqlalchemy.engine.Connection,
+        sqlalchemy.orm.Session,
+        sqlite3.Connection,
+    ]
+
+
+@contextlib.contextmanager
+def _connect(
+    connection: "ConnectionType",
+) -> "Iterator[sqlalchemy.engine.Connection]":
+    import sqlalchemy.orm
+
+    with contextlib.ExitStack() as stack:
+        engine_kwargs = {"echo": bool(os.environ.get("DEBUG_SHOW_SQL_QUERIES"))}
+        if isinstance(connection, (str, sqlalchemy.URL)):
+            engine = sqlalchemy.create_engine(connection, **engine_kwargs)
+            stack.callback(engine.dispose)
+            yield stack.enter_context(engine.connect())
+        elif isinstance(connection, sqlite3.Connection):
+            engine = sqlalchemy.create_engine(
+                "sqlite://", creator=lambda: connection, **engine_kwargs
+            )
+            # do not close the connection, as it is managed by the caller
+            yield engine.connect()
+        elif isinstance(connection, sqlalchemy.Engine):
+            yield stack.enter_context(connection.connect())
+        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 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.
+
+    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
+
+    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(
+    query: Union[str, "sqlalchemy.sql.expression.Executable"],
+    connection: "ConnectionType",
+    params: Union["Sequence[Mapping[str, Any]]", "Mapping[str, Any]", None] = None,
+    *,
+    output: Optional["dict[str, DataType]"] = None,
+    session: Optional["Session"] = None,
+    settings: Optional[dict] = None,
+    in_memory: bool = False,
+    infer_schema_length: Optional[int] = 100,
+) -> "DataChain":
+    """
+    Read the results of a SQL query into a DataChain, using a given database connection.
+
+    Args:
+        query:
+            The SQL query to execute. Can be a raw SQL string or a SQLAlchemy
+            `Executable` object.
+        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.
+        params: Parameters to pass to execute method.
+        output: A dictionary mapping column names to types, used to override the
+            schema inferred from the query results.
+        session: Session to use for the chain.
+        settings: Settings to use for the chain.
+        in_memory: If True, creates an in-memory session. Defaults to False.
+        infer_schema_length:
+            The maximum number of rows to scan for inferring schema.
+            If set to `None`, the full data may be scanned.
+            The rows used for schema inference are stored in memory,
+            so large values can lead to high memory usage.
+            Only applies if the `output` parameter is not set for the given column.
+
+    Examples:
+        Reading from a SQL query against a user-supplied connection:
+        ```python
+        query = "SELECT key, value FROM tbl"
+        chain = dc.read_database(query, connection, output={"value": float})
+        ```
+
+        Load data from a SQLAlchemy driver/engine:
+        ```python
+        from sqlalchemy import create_engine
+        engine = create_engine("postgresql+psycopg://myuser:mypassword@localhost:5432/mydb")
+        chain = dc.read_database("select * from tbl", engine)
+        ```
+
+        Load data from a parameterized SQLAlchemy query:
+        ```python
+        query = "SELECT key, value FROM tbl WHERE value > :value"
+        dc.read_database(query, engine, params={"value": 50})
+        ```
+
+    Notes:
+        - This function works with a variety of databases — including,
+        but not limited to, SQLite, DuckDB, PostgreSQL, and Snowflake,
+        provided the appropriate driver is installed.
+        - This call is blocking, and will execute the query and return once the
+          results are saved.
+    """
+    from datachain.lib.dc.records import read_records
+
+    output = output or {}
+    if isinstance(query, str):
+        query = sqlalchemy.text(query)
+    kw = {"execution_options": {"stream_results": True}}  # use server-side cursors
+    with _connect(connection) as conn, conn.execute(query, params, **kw) as result:
+        cols = result.keys()
+        to_infer = [k for k in cols if k not in output]  # preserve the order
+        rows, inferred_schema = _infer_schema(result, to_infer, infer_schema_length)
+        records = (row._asdict() for row in itertools.chain(rows, result))
+        return read_records(
+            records,
+            session=session,
+            settings=settings,
+            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-0.28.1 → datachain-0.29.0}/src/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-0.28.1 → datachain-0.29.0}/src/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)