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

datachain/data_storage/sqlite.py

@@ -1,25 +1,29 @@
 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 wraps
+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 MetaData, Table, UniqueConstraint, exists, select
+from sqlalchemy import (
+    Column,
+    Integer,
+    MetaData,
+    Table,
+    UniqueConstraint,
+    exists,
+    select,
+)
 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
@@ -28,14 +32,18 @@ import datachain.sql.sqlite
 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
+from datachain.error import DataChainError, OutdatedDatabaseSchemaError
+from datachain.namespace import Namespace
+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
@@ -59,6 +67,14 @@ datachain.sql.sqlite.setup()
 quote_schema = sqlite_dialect.identifier_preparer.quote_schema
 quote = sqlite_dialect.identifier_preparer.quote
 
+# NOTE! This should be manually increased when we change our DB schema in codebase
+SCHEMA_VERSION = 1
+
+OUTDATED_SCHEMA_ERROR_MESSAGE = (
+    "You have an old version of the database schema. Please refer to the documentation"
+    " for more information."
+)
+
 
 def _get_in_memory_uri():
     return "file::memory:?cache=shared"
@@ -85,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:":
@@ -99,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__(
@@ -107,22 +123,24 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         engine: "Engine",
         metadata: "MetaData",
         db: sqlite3.Connection,
-        db_file: Optional[str] = None,
+        db_file: str | None = None,
+        max_variable_number: int | None = 999,
     ):
         self.engine = engine
         self.metadata = metadata
         self.db = db
         self.db_file = db_file
         self.is_closed = False
+        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,
-    ) -> tuple["Engine", "MetaData", sqlite3.Connection, str]:
+        db_file: str | None = None,
+    ) -> tuple["Engine", "MetaData", sqlite3.Connection, str, int]:
         try:
             if db_file == ":memory:":
                 # Enable multithreaded usage of the same in-memory db
@@ -149,6 +167,13 @@ class SQLiteDatabaseEngine(DatabaseEngine):
             db.execute("PRAGMA journal_mode = WAL")
             db.execute("PRAGMA synchronous = NORMAL")
             db.execute("PRAGMA case_sensitive_like = ON")
+
+            max_variable_number = 999  # minimum in old SQLite versions
+            for row in db.execute("PRAGMA compile_options;").fetchall():
+                option = row[0]
+                if option.startswith("MAX_VARIABLE_NUMBER="):
+                    max_variable_number = int(option.split("=")[1])
+
             if os.environ.get("DEBUG_SHOW_SQL_QUERIES"):
                 import sys
 
@@ -156,7 +181,7 @@ class SQLiteDatabaseEngine(DatabaseEngine):
 
             load_usearch_extension(db)
 
-            return engine, MetaData(), db, db_file
+            return engine, MetaData(), db, db_file, max_variable_number
         except RuntimeError:
             raise DataChainError("Can't connect to SQLite DB") from None
 
@@ -172,18 +197,25 @@ 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!")
-        engine, metadata, db, db_file = self._connect(db_file=self.db_file)
+        engine, metadata, db, db_file, max_variable_number = self._connect(
+            db_file=self.db_file
+        )
         self.engine = engine
         self.metadata = metadata
         self.db = db
         self.db_file = db_file
+        self.max_variable_number = max_variable_number
         self.is_closed = False
 
     def get_table(self, name: str) -> Table:
@@ -196,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:
@@ -215,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)
@@ -230,13 +262,27 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         return self.db.execute(sql, parameters)
 
     def insert_dataframe(self, table_name: str, df) -> int:
+        # Dynamically calculates chunksize by dividing max variable limit in a
+        # single SQL insert with number of columns in dataframe.
+        # This way we avoid error: sqlite3.OperationalError: too many SQL variables,
+        num_columns = df.shape[1]
+        if num_columns == 0:
+            num_columns = 1
+
+        if self.max_variable_number < num_columns:
+            raise RuntimeError(
+                "Number of columns exceeds DB maximum variables when inserting data"
+            )
+
+        chunksize = self.max_variable_number // num_columns
+
         return df.to_sql(
             table_name,
             self.db,
             if_exists="append",
             index=False,
             method="multi",
-            chunksize=1000,
+            chunksize=chunksize,
         )
 
     def cursor(self, factory=None):
@@ -245,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
 
@@ -276,7 +324,18 @@ class SQLiteDatabaseEngine(DatabaseEngine):
         )
         return bool(next(self.execute(query))[0])
 
-    def create_table(self, table: "Table", if_not_exists: bool = True) -> None:
+    @property
+    def table_names(self) -> list[str]:
+        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,
+        *,
+        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:
@@ -294,13 +353,15 @@ class SQLiteMetastore(AbstractDBMetastore):
     This is currently used for the local cli.
     """
 
-    db: "SQLiteDatabaseEngine"
+    META_TABLE = "meta"
+
+    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("")
@@ -315,7 +376,12 @@ class SQLiteMetastore(AbstractDBMetastore):
 
         self.db = db or SQLiteDatabaseEngine.from_db_file(db_file)
 
-        self._init_tables()
+        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."""
@@ -323,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("")
@@ -346,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,
@@ -356,8 +426,44 @@ class SQLiteMetastore(AbstractDBMetastore):
         (db_class, db_args, db_kwargs) = db_clone_params
         return cls(uri=uri, db=db_class(*db_args, **db_kwargs))
 
+    @cached_property
+    def _meta(self) -> Table:
+        return Table(self.META_TABLE, self.db.metadata, *self._meta_columns())
+
+    def _meta_select(self, *columns) -> "Select":
+        if not columns:
+            return self._meta.select()
+        return select(*columns)
+
+    def _meta_insert(self) -> "Insert":
+        return sqlite.insert(self._meta)
+
+    def _init_meta_table(self) -> None:
+        """Initializes meta table"""
+        # NOTE! needs to be called before _init_tables()
+        table_names = self.db.table_names
+        if table_names and self.META_TABLE not in table_names:
+            # this will happen on first run
+            raise OutdatedDatabaseSchemaError(OUTDATED_SCHEMA_ERROR_MESSAGE)
+
+        self.db.create_table(self._meta, if_not_exists=True)
+        self.default_table_names.append(self._meta.name)
+
+    def _init_meta_schema_value(self) -> None:
+        """Inserts current schema version value if not present in meta table yet"""
+        stmt = (
+            self._meta_insert()
+            .values(id=1, schema_version=SCHEMA_VERSION)
+            .on_conflict_do_nothing(index_elements=["id"])
+        )
+        self.db.execute(stmt)
+
     def _init_tables(self) -> None:
         """Initialize tables."""
+        self.db.create_table(self._namespaces, if_not_exists=True)
+        self.default_table_names.append(self._namespaces.name)
+        self.db.create_table(self._projects, if_not_exists=True)
+        self.default_table_names.append(self._projects.name)
         self.db.create_table(self._datasets, if_not_exists=True)
         self.default_table_names.append(self._datasets.name)
         self.db.create_table(self._datasets_versions, if_not_exists=True)
@@ -366,11 +472,61 @@ 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:
+        """
+        Creates local namespace and local project connected to it.
+        In local environment user cannot explicitly create other namespaces and
+        projects and all datasets user creates will be stored in those.
+        When pulling dataset from Studio, then other namespaces and projects will
+        be created implicitly though, to keep the same fully qualified name with
+        Studio dataset.
+        """
+        system_namespace = self.create_namespace(
+            Namespace.system(), "System namespace", validate=False
+        )
+        self.create_project(
+            system_namespace.name, Project.listing(), "Listing project", validate=False
+        )
+
+    def _check_schema_version(self) -> None:
+        """
+        Checks if current DB schema is up to date with latest DB model and schema
+        version. If not, OutdatedDatabaseSchemaError is raised.
+        """
+        schema_version = next(self.db.execute(self._meta_select()))[1]
+        if schema_version < SCHEMA_VERSION:
+            raise OutdatedDatabaseSchemaError(OUTDATED_SCHEMA_ERROR_MESSAGE)
+
+    #
+    # Dataset dependencies
+    #
+    @classmethod
+    def _meta_columns(cls) -> list["SchemaItem"]:
+        return [
+            Column("id", Integer, primary_key=True),
+            Column("schema_version", Integer, default=SCHEMA_VERSION),
+        ]
 
     @classmethod
     def _datasets_columns(cls) -> list["SchemaItem"]:
         """Datasets table columns."""
-        return [*super()._datasets_columns(), UniqueConstraint("name")]
+        return [*super()._datasets_columns(), UniqueConstraint("project_id", "name")]
+
+    @classmethod
+    def _namespaces_columns(cls) -> list["SchemaItem"]:
+        """Datasets table columns."""
+        return [*super()._namespaces_columns(), UniqueConstraint("name")]
+
+    def _namespaces_insert(self) -> "Insert":
+        return sqlite.insert(self._namespaces)
+
+    def _projects_insert(self) -> "Insert":
+        return sqlite.insert(self._projects)
 
     def _datasets_insert(self) -> "Insert":
         return sqlite.insert(self._datasets)
@@ -387,6 +543,8 @@ class SQLiteMetastore(AbstractDBMetastore):
 
     def _dataset_dependencies_select_columns(self) -> list["SchemaItem"]:
         return [
+            self._namespaces.c.name,
+            self._projects.c.name,
             self._datasets_dependencies.c.id,
             self._datasets_dependencies.c.dataset_id,
             self._datasets_dependencies.c.dataset_version_id,
@@ -395,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
     #
@@ -402,6 +580,31 @@ 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
+    #
+
+    @property
+    def default_namespace_name(self):
+        return Namespace.default()
+
+    #
+    # Projects
+    #
+
+    @property
+    def default_project_name(self):
+        return Project.default()
+
 
 class SQLiteWarehouse(AbstractWarehouse):
     """
@@ -409,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()
@@ -445,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,
@@ -468,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(
@@ -486,10 +693,10 @@ class SQLiteWarehouse(AbstractWarehouse):
         return table
 
     def get_dataset_sources(
-        self, dataset: DatasetRecord, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> list[StorageURI]:
         dr = self.dataset_rows(dataset, version)
-        query = dr.select(dr.c("source", object_name="file")).distinct()
+        query = dr.select(dr.c("source", column="file")).distinct()
         cur = self.db.cursor()
         cur.row_factory = sqlite3.Row  # type: ignore[assignment]
 
@@ -498,79 +705,26 @@ class SQLiteWarehouse(AbstractWarehouse):
             for row in self.db.execute(query, cursor=cur)
         ]
 
-    def merge_dataset_rows(
-        self,
-        src: DatasetRecord,
-        dst: DatasetRecord,
-        src_version: int,
-        dst_version: int,
-    ) -> None:
-        dst_empty = False
-
-        if not self.db.has_table(self.dataset_table_name(src.name, 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.name, dst_version)):
-            # destination table doesn't exist, create it
-            self.create_dataset_rows_table(
-                self.dataset_table_name(dst.name, 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 v.version < dst_version
-            ]
-            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: int) -> int:
+    def insert_dataset_rows(self, df, dataset: DatasetRecord, version: str) -> int:
         dr = self.dataset_rows(dataset, version)
         return self.db.insert_dataframe(dr.table.name, df)
 
@@ -595,7 +749,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         return col_type.python_type
 
     def dataset_table_export_file_names(
-        self, dataset: DatasetRecord, version: int
+        self, dataset: DatasetRecord, version: str
     ) -> list[str]:
         raise NotImplementedError("Exporting dataset table not implemented for SQLite")
 
@@ -603,7 +757,7 @@ class SQLiteWarehouse(AbstractWarehouse):
         self,
         bucket_uri: str,
         dataset: DatasetRecord,
-        version: int,
+        version: str,
         client_config=None,
     ) -> list[str]:
         raise NotImplementedError("Exporting dataset table not implemented for SQLite")
@@ -612,9 +766,17 @@ class SQLiteWarehouse(AbstractWarehouse):
         self,
         table: Table,
         query: Select,
-        progress_cb: Optional[Callable[[int], None]] = None,
+        progress_cb: Callable[[int], None] | None = None,
     ) -> None:
-        if len(query._group_by_clause) > 0:
+        col_id = (
+            query.selected_columns.sys__id
+            if "sys__id" in query.selected_columns
+            else None
+        )
+
+        # If there is no sys__id column, we cannot copy the table in batches,
+        # and we need to copy all rows at once. Same if there is a group by clause.
+        if col_id is None or len(query._group_by_clause) > 0:
             select_q = query.with_only_columns(
                 *[c for c in query.selected_columns if c.name != "sys__id"]
             )
@@ -622,12 +784,7 @@ class SQLiteWarehouse(AbstractWarehouse):
             self.db.execute(q)
             return
 
-        if "sys__id" in query.selected_columns:
-            col_id = query.selected_columns.sys__id
-        else:
-            col_id = sqlalchemy.column("sys__id")
         select_ids = query.with_only_columns(col_id)
-
         ids = self.db.execute(select_ids).fetchall()
 
         select_q = (
@@ -638,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)
@@ -693,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: