datachain 0.30.5__py3-none-any.whl → 0.39.0__py3-none-any.whl

datachain/data_storage/warehouse.py

@@ -1,30 +1,30 @@
 import glob
 import logging
 import posixpath
-import random
+import secrets
 import string
 from abc import ABC, abstractmethod
-from collections.abc import Generator, Iterable, Iterator, Sequence
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
+from typing import TYPE_CHECKING, Any, Union, cast
 from urllib.parse import urlparse
 
 import attrs
 import sqlalchemy as sa
-import ujson as json
 from sqlalchemy.sql.expression import true
 
+from datachain import json
 from datachain.client import Client
 from datachain.data_storage.schema import convert_rows_custom_column_types
 from datachain.data_storage.serializer import Serializable
 from datachain.dataset import DatasetRecord, StorageURI
 from datachain.lib.file import File
+from datachain.lib.model_store import ModelStore
 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
+from datachain.sql.types import SQLType
 from datachain.utils import sql_escape_like
 
 if TYPE_CHECKING:
@@ -33,6 +33,7 @@ if TYPE_CHECKING:
         _FromClauseArgument,
         _OnClauseArgument,
     )
+    from sqlalchemy.sql.selectable import FromClause
     from sqlalchemy.types import TypeEngine
 
     from datachain.data_storage import schema
@@ -43,6 +44,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):
@@ -69,12 +71,36 @@ class AbstractWarehouse(ABC, Serializable):
         return self
 
     def __exit__(self, exc_type, exc_value, traceback) -> None:
-        # Default behavior is to do nothing, as connections may be shared.
-        pass
+        """Default behavior is to do nothing, as connections may be shared."""
 
     def cleanup_for_tests(self):
         """Cleanup for tests."""
 
+    def _to_jsonable(self, obj: Any) -> Any:
+        """Recursively convert Python/Pydantic structures into JSON-serializable
+        objects.
+        """
+
+        if ModelStore.is_pydantic(type(obj)):
+            # Use Pydantic's JSON mode to ensure datetime and other non-JSON
+            # native types are serialized in a compatible way.
+            return obj.model_dump(mode="json")
+
+        if isinstance(obj, dict):
+            out: dict[str, Any] = {}
+            for k, v in obj.items():
+                if not isinstance(k, str):
+                    key_str = json.dumps(self._to_jsonable(k), ensure_ascii=False)
+                else:
+                    key_str = k
+                out[key_str] = self._to_jsonable(v)
+            return out
+
+        if isinstance(obj, (list, tuple, set)):
+            return [self._to_jsonable(i) for i in obj]
+
+        return obj
+
     def convert_type(  # noqa: PLR0911
         self,
         val: Any,
@@ -121,11 +147,13 @@ class AbstractWarehouse(ABC, Serializable):
             if col_python_type is dict or col_type_name == "JSON":
                 if value_type is str:
                     return val
-                if value_type in (dict, list):
-                    return json.dumps(val, ensure_ascii=False)
-                raise ValueError(
-                    f"Cannot convert value {val!r} with type {value_type} to JSON"
-                )
+                try:
+                    json_ready = self._to_jsonable(val)
+                    return json.dumps(json_ready, ensure_ascii=False)
+                except Exception as e:
+                    raise ValueError(
+                        f"Cannot convert value {val!r} with type {value_type} to JSON"
+                    ) from e
 
             if isinstance(val, col_python_type):
                 return val
@@ -173,12 +201,12 @@ class AbstractWarehouse(ABC, Serializable):
     #
 
     @abstractmethod
-    def is_ready(self, timeout: Optional[int] = None) -> bool: ...
+    def is_ready(self, timeout: int | None = None) -> bool: ...
 
     def dataset_rows(
         self,
         dataset: DatasetRecord,
-        version: Optional[str] = None,
+        version: str | None = None,
         column: str = "file",
     ):
         version = version or dataset.latest_version
@@ -227,7 +255,8 @@ class AbstractWarehouse(ABC, Serializable):
             while True:
                 if limit is not None:
                     limit -= num_yielded
-                    if limit == 0:
+                    num_yielded = 0
+                    if limit <= 0:
                         break
                     if limit < page_size:
                         paginated_query = paginated_query.limit(None).limit(limit)
@@ -245,6 +274,71 @@ class AbstractWarehouse(ABC, Serializable):
                     break  # no more results
                 offset += page_size
 
+    def _regenerate_system_columns(
+        self,
+        selectable: sa.Select,
+        keep_existing_columns: bool = False,
+        regenerate_columns: Iterable[str] | None = None,
+    ) -> sa.Select:
+        """
+        Return a SELECT that regenerates system columns deterministically.
+
+        If keep_existing_columns is True, existing system columns will be kept as-is
+        even when they are listed in ``regenerate_columns``.
+
+        Args:
+            selectable: Base SELECT
+            keep_existing_columns: When True, reuse existing system columns even if
+                they are part of the regeneration set.
+            regenerate_columns: Names of system columns to regenerate. Defaults to
+                {"sys__id", "sys__rand"}. Columns not listed are left untouched.
+        """
+        system_columns = {
+            sys_col.name: sys_col.type
+            for sys_col in self.schema.dataset_row_cls.sys_columns()
+        }
+        regenerate = set(regenerate_columns or system_columns)
+        generators = {
+            "sys__id": self._system_row_number_expr,
+            "sys__rand": self._system_random_expr,
+        }
+
+        base = cast("FromClause", selectable.subquery())
+
+        def build(name: str) -> sa.ColumnElement:
+            expr = generators[name]()
+            return sa.cast(expr, system_columns[name]).label(name)
+
+        columns: list[sa.ColumnElement] = []
+        present: set[str] = set()
+        changed = False
+
+        for col in base.c:
+            present.add(col.name)
+            regen = col.name in regenerate and not keep_existing_columns
+            columns.append(build(col.name) if regen else col)
+            changed |= regen
+
+        for name in regenerate - present:
+            columns.append(build(name))
+            changed = True
+
+        if not changed:
+            return selectable
+
+        inner = sa.select(*columns).select_from(base).subquery()
+        return sa.select(*inner.c).select_from(inner)
+
+    def _system_row_number_expr(self):
+        """Return an expression that produces deterministic row numbers."""
+
+        raise NotImplementedError
+
+    def _system_random_expr(self):
+        """Return an expression that produces deterministic random values."""
+
+        raise NotImplementedError
+
     #
     # Table Name Internal Functions
     #
@@ -303,21 +397,9 @@ class AbstractWarehouse(ABC, Serializable):
         table_name = self.dataset_table_name(dataset, version)
         table = sa.Table(table_name, self.db.metadata)
         self.db.drop_table(table, if_exists=if_exists)
-
-    @abstractmethod
-    def merge_dataset_rows(
-        self,
-        src: "DatasetRecord",
-        dst: "DatasetRecord",
-        src_version: str,
-        dst_version: str,
-    ) -> None:
-        """
-        Merges source dataset rows and current latest destination dataset rows
-        into a new rows table created for new destination dataset version.
-        Note that table for new destination version must be created upfront.
-        Merge results should not contain duplicates.
-        """
+        # Remove from metadata cache to allow recreation
+        if table_name in self.db.metadata.tables:
+            self.db.metadata.remove(self.db.metadata.tables[table_name])
 
     def dataset_rows_select(
         self,
@@ -341,7 +423,7 @@ class AbstractWarehouse(ABC, Serializable):
         """
         Fetch dataset rows from database using a list of IDs.
         """
-        if (id_col := get_query_id_column(query)) is None:
+        if (id_col := query.selected_columns.get("sys__id")) is None:
             raise RuntimeError("sys__id column not found in query")
 
         query = query._clone().offset(None).limit(None).order_by(None)
@@ -385,7 +467,7 @@ class AbstractWarehouse(ABC, Serializable):
 
     def dataset_stats(
         self, dataset: DatasetRecord, version: str
-    ) -> tuple[Optional[int], Optional[int]]:
+    ) -> tuple[int | None, int | None]:
         """
         Returns tuple with dataset stats: total number of rows and total dataset size.
         """
@@ -415,7 +497,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:
@@ -505,7 +592,7 @@ class AbstractWarehouse(ABC, Serializable):
         dr = dataset_rows
         columns = [c.name for c in query.selected_columns]
         for row in self.db.execute(query):
-            d = dict(zip(columns, row))
+            d = dict(zip(columns, row, strict=False))
             yield Node(**{dr.without_object(k): v for k, v in d.items()})
 
     def get_dirs_by_parent_path(
@@ -742,7 +829,7 @@ class AbstractWarehouse(ABC, Serializable):
     def size(
         self,
         dataset_rows: "DataTable",
-        node: Union[Node, dict[str, Any]],
+        node: Node | dict[str, Any],
         count_files: bool = False,
     ) -> tuple[int, int]:
         """
@@ -784,10 +871,10 @@ class AbstractWarehouse(ABC, Serializable):
         self,
         dataset_rows: "DataTable",
         parent_path: str,
-        fields: Optional[Sequence[str]] = None,
-        type: Optional[str] = None,
+        fields: Sequence[str] | None = None,
+        type: str | None = None,
         conds=None,
-        order_by: Optional[Union[str, list[str]]] = None,
+        order_by: str | list[str] | None = None,
         include_subobjects: bool = True,
     ) -> sa.Select:
         if not conds:
@@ -825,7 +912,7 @@ class AbstractWarehouse(ABC, Serializable):
         self,
         dataset_rows: "DataTable",
         node: Node,
-        sort: Union[list[str], str, None] = None,
+        sort: list[str] | str | None = None,
         include_subobjects: bool = True,
     ) -> Iterator[NodeWithPath]:
         """
@@ -883,20 +970,25 @@ class AbstractWarehouse(ABC, Serializable):
     def create_udf_table(
         self,
         columns: Sequence["sa.Column"] = (),
-        name: Optional[str] = None,
+        name: str | None = None,
     ) -> sa.Table:
         """
         Create a temporary table for storing custom signals generated by a UDF.
         SQLite TEMPORARY tables cannot be directly used as they are process-specific,
         and UDFs are run in other processes when run in parallel.
         """
+        columns = [
+            c
+            for c in columns
+            if c.name not in [col.name for col in self.dataset_row_cls.sys_columns()]
+        ]
         tbl = sa.Table(
             name or self.udf_table_name(),
             sa.MetaData(),
-            sa.Column("sys__id", Int, primary_key=True),
+            *self.dataset_row_cls.sys_columns(),
             *columns,
         )
-        self.db.create_table(tbl, if_not_exists=True)
+        self.db.create_table(tbl, if_not_exists=True, kind="udf")
         return tbl
 
     @abstractmethod
@@ -904,7 +996,7 @@ class AbstractWarehouse(ABC, Serializable):
         self,
         table: sa.Table,
         query: sa.Select,
-        progress_cb: Optional[Callable[[int], None]] = None,
+        progress_cb: Callable[[int], None] | None = None,
     ) -> None:
         """
         Copy the results of a query into a table.
@@ -917,6 +1009,8 @@ class AbstractWarehouse(ABC, Serializable):
         right: "_FromClauseArgument",
         onclause: "_OnClauseArgument",
         inner: bool = True,
+        full: bool = False,
+        columns=None,
     ) -> sa.Select:
         """
         Join two tables together.
@@ -953,7 +1047,5 @@ class AbstractWarehouse(ABC, Serializable):
 
 
 def _random_string(length: int) -> str:
-    return "".join(
-        random.choice(string.ascii_letters + string.digits)  # noqa: S311
-        for i in range(length)
-    )
+    alphabet = string.ascii_letters + string.digits
+    return "".join(secrets.choice(alphabet) for _ in range(length))