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

datachain/data_storage/sqlite.py

@@ -1,18 +1,11 @@
 import logging
 import os
 import sqlite3
-from collections.abc import Iterable, Sequence
+from collections.abc import Callable, Iterable, Sequence
 from contextlib import contextmanager
 from functools import cached_property, wraps
 from time import sleep
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    ClassVar,
-    Optional,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, ClassVar, Union
 
 import sqlalchemy
 from sqlalchemy import (
@@ -27,16 +20,19 @@ from sqlalchemy import (
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.schema import CreateIndex, CreateTable, DropTable
 from sqlalchemy.sql import func
-from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+from sqlalchemy.sql.elements import (
+    BinaryExpression,
+    BooleanClauseList,
+)
 from sqlalchemy.sql.expression import bindparam, cast
 from sqlalchemy.sql.selectable import Select
 from tqdm.auto import tqdm
 
 import datachain.sql.sqlite
-from datachain import semver
 from datachain.data_storage import AbstractDBMetastore, AbstractWarehouse
 from datachain.data_storage.db_engine import DatabaseEngine
 from datachain.data_storage.schema import DefaultSchema
+from datachain.data_storage.warehouse import INSERT_BATCH_SIZE
 from datachain.dataset import DatasetRecord, StorageURI
 from datachain.error import DataChainError, OutdatedDatabaseSchemaError
 from datachain.namespace import Namespace
@@ -44,9 +40,10 @@ from datachain.project import Project
 from datachain.sql.sqlite import create_user_defined_sql_functions, sqlite_dialect
 from datachain.sql.sqlite.base import load_usearch_extension
 from datachain.sql.types import SQLType
-from datachain.utils import DataChainDir, batched_it
+from datachain.utils import DataChainDir, batched, batched_it
 
 if TYPE_CHECKING:
+    from sqlalchemy import CTE, Subquery
     from sqlalchemy.dialects.sqlite import Insert
     from sqlalchemy.engine.base import Engine
     from sqlalchemy.schema import SchemaItem
@@ -104,8 +101,8 @@ def retry_sqlite_locks(func):
 
 
 def get_db_file_in_memory(
-    db_file: Optional[str] = None, in_memory: bool = False
-) -> Optional[str]:
+    db_file: str | None = None, in_memory: bool = False
+) -> str | None:
     """Get in-memory db_file and check that conflicting arguments are not provided."""
     if in_memory:
         if db_file and db_file != ":memory:":
@@ -118,7 +115,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
     dialect = sqlite_dialect
 
     db: sqlite3.Connection
-    db_file: Optional[str]
+    db_file: str | None
     is_closed: bool
 
     def __init__(
@@ -126,8 +123,8 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         engine: "Engine",
         metadata: "MetaData",
         db: sqlite3.Connection,
-        db_file: Optional[str] = None,
-        max_variable_number: Optional[int] = 999,
+        db_file: str | None = None,
+        max_variable_number: int | None = 999,
     ):
         self.engine = engine
         self.metadata = metadata
@@ -137,12 +134,12 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         self.max_variable_number = max_variable_number
 
     @classmethod
-    def from_db_file(cls, db_file: Optional[str] = None) -> "SQLiteDatabaseEngine":
+    def from_db_file(cls, db_file: str | None = None) -> "SQLiteDatabaseEngine":
         return cls(*cls._connect(db_file=db_file))
 
     @staticmethod
     def _connect(
-        db_file: Optional[str] = None,
+        db_file: str | None = None,
     ) -> tuple["Engine", "MetaData", sqlite3.Connection, str, int]:
         try:
             if db_file == ":memory:":
@@ -200,10 +197,14 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         """
         return (
             SQLiteDatabaseEngine.from_db_file,
-            [self.db_file],
+            [str(self.db_file)],
             {},
         )
 
+    @classmethod
+    def serialize_callable_name(cls) -> str:
+        return "sqlite.from_db_file"
+
     def _reconnect(self) -> None:
         if not self.is_closed:
             raise RuntimeError("Cannot reconnect on still-open DB!")
@@ -227,7 +228,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
     def execute(
         self,
         query,
-        cursor: Optional[sqlite3.Cursor] = None,
+        cursor: sqlite3.Cursor | None = None,
         conn=None,
     ) -> sqlite3.Cursor:
         if self.is_closed:
@@ -246,7 +247,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
 
     @retry_sqlite_locks
     def executemany(
-        self, query, params, cursor: Optional[sqlite3.Cursor] = None, conn=None
+        self, query, params, cursor: sqlite3.Cursor | None = None, conn=None
     ) -> sqlite3.Cursor:
         if cursor:
             return cursor.executemany(self.compile(query).string, params)
@@ -290,6 +291,8 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         return self.db.cursor(factory)
 
     def close(self) -> None:
+        if self.is_closed:
+            return
         self.db.close()
         self.is_closed = True
 
@@ -326,7 +329,13 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         query = "SELECT name FROM sqlite_master WHERE type='table';"
         return [r[0] for r in self.execute_str(query).fetchall()]
 
-    def create_table(self, table: "Table", if_not_exists: bool = True) -> None:
+    def create_table(
+        self,
+        table: "Table",
+        if_not_exists: bool = True,
+        *,
+        kind: str | None = None,
+    ) -> None:
         self.execute(CreateTable(table, if_not_exists=if_not_exists))
 
     def drop_table(self, table: "Table", if_exists: bool = False) -> None:
@@ -346,13 +355,13 @@ class SQLiteMetastore(AbstractDBMetastore):
 
     META_TABLE = "meta"
 
-    db: "SQLiteDatabaseEngine"
+    db: SQLiteDatabaseEngine
 
     def __init__(
         self,
-        uri: Optional[StorageURI] = None,
-        db: Optional["SQLiteDatabaseEngine"] = None,
-        db_file: Optional[str] = None,
+        uri: StorageURI | None = None,
+        db: SQLiteDatabaseEngine | None = None,
+        db_file: str | None = None,
         in_memory: bool = False,
     ):
         uri = uri or StorageURI("")
@@ -367,11 +376,12 @@ class SQLiteMetastore(AbstractDBMetastore):
 
         self.db = db or SQLiteDatabaseEngine.from_db_file(db_file)
 
-        self._init_meta_table()
-        self._init_meta_schema_value()
-        self._check_schema_version()
-        self._init_tables()
-        self._init_namespaces_projects()
+        with self._init_guard():
+            self._init_meta_table()
+            self._init_meta_schema_value()
+            self._check_schema_version()
+            self._init_tables()
+            self._init_namespaces_projects()
 
     def __exit__(self, exc_type, exc_value, traceback) -> None:
         """Close connection upon exit from context manager."""
@@ -379,7 +389,7 @@ class SQLiteMetastore(AbstractDBMetastore):
 
     def clone(
         self,
-        uri: Optional[StorageURI] = None,
+        uri: StorageURI | None = None,
         use_new_connection: bool = False,
     ) -> "SQLiteMetastore":
         uri = uri or StorageURI("")
@@ -402,6 +412,10 @@ class SQLiteMetastore(AbstractDBMetastore):
             },
         )
 
+    @classmethod
+    def serialize_callable_name(cls) -> str:
+        return "sqlite.metastore.init_after_clone"
+
     @classmethod
     def init_after_clone(
         cls,
@@ -458,6 +472,10 @@ class SQLiteMetastore(AbstractDBMetastore):
         self.default_table_names.append(self._datasets_dependencies.name)
         self.db.create_table(self._jobs, if_not_exists=True)
         self.default_table_names.append(self._jobs.name)
+        self.db.create_table(self._checkpoints, if_not_exists=True)
+        self.default_table_names.append(self._checkpoints.name)
+        self.db.create_table(self._dataset_version_jobs, if_not_exists=True)
+        self.default_table_names.append(self._dataset_version_jobs.name)
 
     def _init_namespaces_projects(self) -> None:
         """
@@ -535,6 +553,26 @@ class SQLiteMetastore(AbstractDBMetastore):
             self._datasets_versions.c.created_at,
         ]
 
+    def _dataset_dependency_nodes_select_columns(
+        self,
+        namespaces_subquery: "Subquery",
+        dependency_tree_cte: "CTE",
+        datasets_subquery: "Subquery",
+    ) -> list["ColumnElement"]:
+        return [
+            namespaces_subquery.c.name,
+            self._projects.c.name,
+            dependency_tree_cte.c.id,
+            dependency_tree_cte.c.dataset_id,
+            dependency_tree_cte.c.dataset_version_id,
+            datasets_subquery.c.name,
+            self._datasets_versions.c.version,
+            self._datasets_versions.c.created_at,
+            dependency_tree_cte.c.source_dataset_id,
+            dependency_tree_cte.c.source_dataset_version_id,
+            dependency_tree_cte.c.depth,
+        ]
+
     #
     # Jobs
     #
@@ -542,6 +580,15 @@ class SQLiteMetastore(AbstractDBMetastore):
     def _jobs_insert(self) -> "Insert":
         return sqlite.insert(self._jobs)
 
+    #
+    # Checkpoints
+    #
+    def _checkpoints_insert(self) -> "Insert":
+        return sqlite.insert(self._checkpoints)
+
+    def _dataset_version_jobs_insert(self) -> "Insert":
+        return sqlite.insert(self._dataset_version_jobs)
+
     #
     # Namespaces
     #
@@ -565,15 +612,15 @@ class SQLiteWarehouse(AbstractWarehouse):
     This is currently used for the local cli.
     """
 
-    db: "SQLiteDatabaseEngine"
+    db: SQLiteDatabaseEngine
 
     # Cache for our defined column types to dialect specific TypeEngine relations
     _col_python_type: ClassVar[dict[type, "TypeEngine"]] = {}
 
     def __init__(
         self,
-        db: Optional["SQLiteDatabaseEngine"] = None,
-        db_file: Optional[str] = None,
+        db: SQLiteDatabaseEngine | None = None,
+        db_file: str | None = None,
         in_memory: bool = False,
     ):
         self.schema: DefaultSchema = DefaultSchema()
@@ -601,6 +648,10 @@ class SQLiteWarehouse(AbstractWarehouse):
             {"db_clone_params": self.db.clone_params()},
         )
 
+    @classmethod
+    def serialize_callable_name(cls) -> str:
+        return "sqlite.warehouse.init_after_clone"
+
     @classmethod
     def init_after_clone(
         cls,
@@ -624,7 +675,7 @@ class SQLiteWarehouse(AbstractWarehouse):
             only=filter_tables,
         )
 
-    def is_ready(self, timeout: Optional[int] = None) -> bool:
+    def is_ready(self, timeout: int | None = None) -> bool:
         return True
 
     def create_dataset_rows_table(
@@ -654,77 +705,24 @@ class SQLiteWarehouse(AbstractWarehouse):
             for row in self.db.execute(query, cursor=cur)
         ]
 
-    def merge_dataset_rows(
-        self,
-        src: DatasetRecord,
-        dst: DatasetRecord,
-        src_version: str,
-        dst_version: str,
-    ) -> None:
-        dst_empty = False
-
-        if not self.db.has_table(self.dataset_table_name(src, src_version)):
-            # source table doesn't exist, nothing to do
-            return
-
-        src_dr = self.dataset_rows(src, src_version).table
-
-        if not self.db.has_table(self.dataset_table_name(dst, dst_version)):
-            # destination table doesn't exist, create it
-            self.create_dataset_rows_table(
-                self.dataset_table_name(dst, dst_version),
-                columns=src_dr.columns,
-            )
-            dst_empty = True
-
-        dst_dr = self.dataset_rows(dst, dst_version).table
-        merge_fields = [c.name for c in src_dr.columns if c.name != "sys__id"]
-        select_src = select(*(getattr(src_dr.columns, f) for f in merge_fields))
-
-        if dst_empty:
-            # we don't need union, but just select from source to destination
-            insert_query = sqlite.insert(dst_dr).from_select(merge_fields, select_src)
-        else:
-            dst_version_latest = None
-            # find the previous version of the destination dataset
-            dst_previous_versions = [
-                v.version
-                for v in dst.versions  # type: ignore [union-attr]
-                if semver.compare(v.version, dst_version) == -1
-            ]
-            if dst_previous_versions:
-                dst_version_latest = max(dst_previous_versions)
-
-            dst_dr_latest = self.dataset_rows(dst, dst_version_latest).table
-
-            select_dst_latest = select(
-                *(getattr(dst_dr_latest.c, f) for f in merge_fields)
-            )
-            union_query = sqlalchemy.union(select_src, select_dst_latest)
-            insert_query = (
-                sqlite.insert(dst_dr)
-                .from_select(merge_fields, union_query)
-                .prefix_with("OR IGNORE")
-            )
-
-        self.db.execute(insert_query)
-
     def prepare_entries(self, entries: "Iterable[File]") -> Iterable[dict[str, Any]]:
         return (e.model_dump() for e in entries)
 
-    def insert_rows(self, table: Table, rows: Iterable[dict[str, Any]]) -> None:
-        rows = list(rows)
-        if not rows:
-            return
-
-        with self.db.transaction() as conn:
-            # transactions speeds up inserts significantly as there is no separate
-            # transaction created for each insert row
-            self.db.executemany(
-                table.insert().values({f: bindparam(f) for f in rows[0]}),
-                rows,
-                conn=conn,
-            )
+    def insert_rows(
+        self,
+        table: Table,
+        rows: Iterable[dict[str, Any]],
+        batch_size: int = INSERT_BATCH_SIZE,
+    ) -> None:
+        for row_chunk in batched(rows, batch_size):
+            with self.db.transaction() as conn:
+                # transactions speeds up inserts significantly as there is no separate
+                # transaction created for each insert row
+                self.db.executemany(
+                    table.insert().values({f: bindparam(f) for f in row_chunk[0]}),
+                    row_chunk,
+                    conn=conn,
+                )
 
     def insert_dataset_rows(self, df, dataset: DatasetRecord, version: str) -> int:
         dr = self.dataset_rows(dataset, version)
@@ -768,7 +766,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         self,
         table: Table,
         query: Select,
-        progress_cb: Optional[Callable[[int], None]] = None,
+        progress_cb: Callable[[int], None] | None = None,
     ) -> None:
         col_id = (
             query.selected_columns.sys__id
@@ -797,7 +795,7 @@ class SQLiteWarehouse(AbstractWarehouse):
             .limit(None)
         )
 
-        for batch in batched_it(ids, 10_000):
+        for batch in batched_it(ids, INSERT_BATCH_SIZE):
             batch_ids = [row[0] for row in batch]
             select_q._where_criteria = (col_id.in_(batch_ids),)
             q = table.insert().from_select(list(select_q.selected_columns), select_q)
@@ -852,18 +850,20 @@ class SQLiteWarehouse(AbstractWarehouse):
                 if isinstance(c, BinaryExpression):
                     right_left_join = add_left_rows_filter(c)
 
-        union = sqlalchemy.union(left_right_join, right_left_join).subquery()
-        return sqlalchemy.select(*union.c).select_from(union)
+        union_cte = sqlalchemy.union(left_right_join, right_left_join).cte()
+        return sqlalchemy.select(*union_cte.c).select_from(union_cte)
+
+    def _system_row_number_expr(self):
+        return func.row_number().over()
+
+    def _system_random_expr(self):
+        return self._system_row_number_expr() * 1103515245 + 12345
 
     def create_pre_udf_table(self, query: "Select") -> "Table":
         """
         Create a temporary table from a query for use in a UDF.
         """
-        columns = [
-            sqlalchemy.Column(c.name, c.type)
-            for c in query.selected_columns
-            if c.name != "sys__id"
-        ]
+        columns = [sqlalchemy.Column(c.name, c.type) for c in query.selected_columns]
         table = self.create_udf_table(columns)
 
         with tqdm(desc="Preparing", unit=" rows", leave=False) as pbar: