datachain 0.28.2__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")
 
@@ -2276,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,

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-0.28.2.dist-info → datachain-0.29.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datachain
-Version: 0.28.2
+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.2.dist-info → datachain-0.29.0.dist-info}/RECORD

@@ -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
@@ -86,7 +86,7 @@ datachain/lib/namespaces.py,sha256=it52UbbwB8dzhesO2pMs_nThXiPQ1Ph9sD9I3GQkg5s,2
 datachain/lib/projects.py,sha256=8lN0qV8czX1LGtWURCUvRlSJk-RpO9w9Rra_pOZus6g,2595
 datachain/lib/pytorch.py,sha256=S-st2SAczYut13KMf6eSqP_OQ8otWI5TRmzhK5fN3k0,7828
 datachain/lib/settings.py,sha256=n0YYhCVdgCdMkCSLY7kscJF9mUhlQ0a4ENWBsJFynkw,3809
-datachain/lib/signal_schema.py,sha256=JMsL8c4iCRH9PoRumvjimsOLQQslTjm_aDR2jh1zT2Q,38558
+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=IB1IKF5KyA-NiyfhVzmBPpF_aITPS3zSlrt24f_Ofjo,17956
@@ -103,8 +103,8 @@ 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=T5-b2LLCF0zYhXQjOgtzzr6cm5NfrKVGxcJTWn7tfNU,94164
+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
@@ -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.2.dist-info/licenses/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.28.2.dist-info/METADATA,sha256=dYo2qW8RMNNCyy6KOXztfXOIldyS4_mADxeAlCI9cKw,13766
-datachain-0.28.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-datachain-0.28.2.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.28.2.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.28.2.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,,