"""Parallel encode + shard-merge backend for :class:`~httk.store.backend.sql.bulk.BulkIngest`.
This module implements the ``workers > 1`` mode of
:meth:`~httk.store.backend.sql.store.SqlStore.bulk_ingest`. The serial ``workers = 1``
path in :mod:`httk.store.backend.sql.bulk` is untouched; :class:`~httk.store.backend.sql.bulk.BulkIngest`
delegates to the helpers here only when more than one worker is requested.
The design has three moving parts:
- **Workers** (``_worker_main``): forked processes that run the *pure* encoders
(``_encode_parent_row`` / ``_encode_child_rows`` from
:mod:`httk.store.backend.sql.store`) with a per-worker
:class:`~httk.store.store_common.SaveProjection`. Each worker owns a disjoint
sid block (``(worker_index + 1) << 26``) so its rows never collide with
another worker's before the merge. A worker deduplicates *content-addressed
records that carry no identity-excluded metadata* and *all by_value records*
within its own stream (bounding shard size); records that carry a metadata
plan are emitted per occurrence so the merge can verify every collision.
Workers never touch the database — they only write shard files.
- **Shards** (``_ParquetShardWriter``, ``_SqliteShardWriter``): per-worker,
per-table row files. DuckDB stores each flush as a pyarrow Parquet file
(``pyarrow`` imported lazily; its absence raises the documented ``parallel``
extra hint); SQLite stores one shard database per worker written with native
``executemany``. Shards live in a ``tempfile.TemporaryDirectory`` next to the
target database file when it is file-backed, else the tempfile default, and
are always removed.
- **Merge** (:func:`merge`): the main process, inside the ingest's spanning
transaction, loads every shard into the freshly created (index-less) record
tables under the workers' block sids, then collapses cross-worker duplicates
set-wise (content-id and by_value) in foreign-key dependency order. Because
referenced tables collapse before their referrers, two rows sharing a content
id then differ only in their identity-excluded (``IdentitySkip``) columns, so
each collision's metadata is verified with a single grouped scan per table
rather than by reconstructing every duplicate record (the dominant cost at
real-build scale); nested and ``descend`` conflicts surface at the target
table where the skip metadata lives. The merge then sweeps rows orphaned by a
collapsed duplicate's subtree and remaps the surviving block sids to a compact
``1..N`` range, rewriting every foreign-key column through the same map.
``workers > 1`` targets the offline *build* of a store: it requires a
physically empty target (no application table already holds rows). Incremental
appends into a populated store remain the serial path's domain, where the
per-record staging protocol and its metadata verification already live.
"""
import csv
import functools
import importlib
import math
import os
import pickle
import queue as queue_mod
import sqlite3
import tempfile
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any
import sqlalchemy
from httk.core.entry_ids import check_entry_id, check_immutable_id
from httk.core.storage import StorageProjectionCycleError, resolve_storage_record
from httk.store.backend.schema import TableSchema, resolve_schema
from httk.store.backend.sql.layout import actual_table_names
from httk.store.backend.sql.mapping import (
ALT_ID_COLUMN,
ALT_KIND_COLUMN,
CONTENT_ID_COLUMN,
DISPATCH_CONTENT_ID_COLUMN,
LOGICAL_ID_COLUMN,
ROLE_COLUMN,
SID_COLUMN,
STORE_TIMESTAMP_COLUMN,
backing_dispatch_column_name,
entry_dispatch_table_name,
)
from httk.store.backend.sql.store import (
SqlStore,
_encode_child_rows,
_encode_parent_row,
_encode_promoted_descendants,
_field_path,
)
from httk.store.store_common import (
EntryDispatchIntegrityError,
EntryIdConflictError,
EntryMetadataConflictError,
SaveProjection,
_metadata_plan,
)
if TYPE_CHECKING:
import httk.store.backend.sql.bulk
from httk.store.backend.sql.bulk import BulkIngest
__all__ = ["ParallelController", "merge"]
# Each worker allocates sids from a disjoint high block so rows never collide
# before the merge; the merge remaps every block sid to a compact 1..N value.
# The store's sid column is a 32-bit integer (SQLAlchemy ``Integer`` renders as
# DuckDB ``INTEGER``), so worker ``w`` bases its block at ``(w + 1) << 26``:
# that fits a signed 32-bit integer for up to ~30 workers and leaves 2**26
# (~67M) rows per worker per table before the next worker's block.
_SID_BLOCK_BITS = 26
_SID_BLOCK = 1 << _SID_BLOCK_BITS
_MAX_WORKERS = (1 << 31) // _SID_BLOCK - 1
# Per-worker task queue depth (bounds in-flight buffering). A module constant so
# tests can shrink it to force a saturated queue.
_QUEUE_MAXSIZE = 64
# Upper bound on how long the main process waits for a worker to make progress
# (report a result, or accept a stop sentinel) before declaring the pool stalled.
_WORKER_STALL_TIMEOUT = 300.0
_SQLITE_COPY_BATCH_SIZE = 1_000
def _worker_base(worker_index: int) -> int:
"""The first sid a worker may allocate (its block is ``[base, base + 2**26)``)."""
return (worker_index + 1) << _SID_BLOCK_BITS
def _references_reach(start: type, goal: type, seen: set[type] | None = None) -> bool:
"""Whether following reference edges from ``start`` reaches ``goal`` (a reference cycle test)."""
if start is goal:
return True
seen = set() if seen is None else seen
if start in seen:
return False
seen.add(start)
return any(_references_reach(referenced, goal, seen) for referenced in resolve_schema(start).referenced_classes())
@functools.cache
def unsupported_metadata_reason(record_type: type) -> str | None:
"""Why the parallel merge cannot verify ``record_type``'s identity-excluded metadata, or ``None``.
The set-wise merge verifies identity-excluded metadata with a grouped column
scan (see :meth:`_Merger._verify_collision_metadata`). That covers scalar
``IdentitySkip`` columns and skipped references to content-addressed or
by_value targets, and it delegates ``descend`` conflicts to the target
table's own collapse. Three shapes fall outside it and are rejected up front
(fail fast, naming ``workers=1``) rather than verified incorrectly:
- an identity-excluded **child sequence** (no parent column to group on);
- an identity-excluded **reference to a non-deduplicated** (``none``) record,
or a **descend into** one (the target is never collapsed, so its metadata
is never compared);
- a **self-referential** identity-excluded reference (the target table is the
one being collapsed, so its sids are not yet final when compared).
:param record_type: The record class to classify.
:return: A human-readable reason string, or ``None`` when the shape is supported.
"""
plan = _metadata_plan(record_type)
if plan is None:
return None
name = record_type.__name__
for spec in plan.skipped_nested:
if spec.role != "reference":
return f"{name}.{spec.field} is an identity-excluded child sequence"
if spec.target is None:
continue
if resolve_schema(spec.target).dedup not in ("content_id", "by_value"):
return f"{name}.{spec.field} is an identity-excluded reference to a non-deduplicated record"
if _references_reach(spec.target, record_type):
return f"{name}.{spec.field} is a self-referential identity-excluded reference"
for spec in plan.descend_specs:
if spec.target is not None and resolve_schema(spec.target).dedup == "none":
return f"{name}.{spec.field} descends into a non-deduplicated ('none') record's metadata"
return None
@functools.cache
def _plain_float_skip_fields(record_type: type) -> tuple[tuple[str, str], ...]:
"""The ``(field, column)`` pairs of ``record_type``'s plain-``float`` ``IdentitySkip`` fields.
Only the plain Python-``float`` codec is meant here — not the exact numeric
codecs (``fraction``, ``fracscalar``, tensor codecs) that also carry a float
column beside their exact text channel. A NaN in a plain-float column reads
back as ``NaN`` on DuckDB and as ``NULL`` on SQLite (which has no NaN), so it
cannot be told apart from a real ``None`` once stored; the worker therefore
flags NaN-bearing content ids while it still holds the source value, and the
merge treats a duplicated flagged content id as a conflict (serial's
``NaN != NaN``).
"""
plan = _metadata_plan(record_type)
if plan is None:
return ()
fields: list[tuple[str, str]] = []
for spec in plan.skipped_specs:
if spec.codec_name != "float":
continue
fields.extend((spec.field, column.name) for column in spec.columns if column.kind == "float")
return tuple(fields)
@dataclass(frozen=True)
class _WorkerConfig:
"""Immutable per-run settings handed to every worker (fork-inherited)."""
chunk_size: int
shard_dir: str
backend: str # "duckdb" or "sqlite"
track_sids: bool = True
store_timestamp: int | None = None
# Deferred Parquet builds persist auxiliary root/dispatch/diagnostic data
# alongside record rows. Parity merge deliberately keeps its established
# in-memory manifest protocol.
spill_deferred_auxiliary: bool = False
@dataclass
class _DispatchRow:
"""A buffered entry-dispatch row a worker produced (backing sid still a block sid)."""
dispatch_name: str
key: str
column: str
all_columns: tuple[str, ...]
ref_table: str
block_sid: int
family_name: str
@dataclass
class _WorkerManifest:
"""What a finished worker reports to the main process."""
worker_index: int
encoded_count: int
token_sid: dict[int, tuple[str, int]]
roots: list[tuple[str, int]]
late_role_roots: list[tuple[str, int]]
dispatch: list[_DispatchRow]
tables: list[str]
# DuckDB: table -> list of parquet file paths. SQLite: {"db": path}.
shards: dict[str, Any]
# (table, content_id, field) triples whose identity-excluded float held a NaN.
nan_content: list[tuple[str, str, str]] = field(default_factory=list)
# (table, stage_sid, column) cells whose float held a NaN (SQLite shards only;
# restored on PostgreSQL load where NaN survives, unlike SQLite's NULL).
nan_floats: list[tuple[str, int, str]] = field(default_factory=list)
# Fork-inherited handles the worker reads from module scope (never pickled).
_PARENT_STORE: SqlStore | None = None
_PARENT_CONFIG: _WorkerConfig | None = None
# --------------------------------------------------------------------- shard writers
def _pa_type(column: sqlalchemy.Column[Any], pa: Any) -> Any:
"""Map a record column's SQLAlchemy type to the pyarrow type of its shard column."""
type_ = column.type
if isinstance(type_, sqlalchemy.Boolean):
return pa.bool_()
if isinstance(type_, sqlalchemy.Integer):
return pa.int64()
if isinstance(type_, sqlalchemy.Float):
return pa.float64()
if isinstance(type_, sqlalchemy.LargeBinary):
return pa.binary()
# Text / String and everything else stringly-typed.
return pa.string()
class _ParquetShardWriter:
"""Write per-worker, per-table Parquet shards (the DuckDB backend hand-off)."""
def __init__(self, store: SqlStore, worker_index: int, shard_dir: str) -> None:
try:
self._pa = importlib.import_module("pyarrow")
self._pq = importlib.import_module("pyarrow.parquet")
except ImportError as error: # pragma: no cover - guarded before fork
raise ImportError(
"bulk_ingest(workers>1) on a DuckDB store needs pyarrow; "
"install the 'httk-store[parallel]' extra to use it"
) from error
self._store = store
self._worker_index = worker_index
self._dir = shard_dir
self._schemas: dict[str, Any] = {}
self._files: dict[str, list[str]] = {}
self._sequence = 0
def _schema_for(self, table_name: str) -> Any:
schema = self._schemas.get(table_name)
if schema is None:
table = self._store._table(table_name)
fields = [self._pa.field(column.name, _pa_type(column, self._pa)) for column in table.columns]
schema = self._pa.schema(fields)
self._schemas[table_name] = schema
return schema
def write(self, table_name: str, rows: list[dict[str, Any]]) -> None:
if not rows:
return
schema = self._schema_for(table_name)
columns = [field_.name for field_ in schema]
data = {name: self._pa.array([row.get(name) for row in rows], type=schema.field(name).type) for name in columns}
table = self._pa.table(data, schema=schema)
path = os.path.join(self._dir, f"w{self._worker_index}_{table_name}_{self._sequence}.parquet")
self._sequence += 1
self._pq.write_table(table, path)
self._files.setdefault(table_name, []).append(path)
def finalize(self) -> dict[str, Any]:
return dict(self._files)
def write_roots(self, rows: list[dict[str, Any]]) -> None:
"""Persist top-level roots instead of retaining them in a manifest."""
if not rows:
return
schema = self._pa.schema(
[
self._pa.field("token", self._pa.int64()),
self._pa.field("tbl", self._pa.string()),
self._pa.field("stage_sid", self._pa.int64()),
]
)
data = {field.name: self._pa.array([row[field.name] for row in rows], type=field.type) for field in schema}
path = os.path.join(self._dir, f"w{self._worker_index}_roots_{self._sequence}.parquet")
self._sequence += 1
self._pq.write_table(self._pa.table(data, schema=schema), path)
self._files.setdefault("_httk_roots", []).append(path)
def _write_auxiliary(self, name: str, schema: Any, rows: list[dict[str, Any]]) -> None:
"""Write one bounded auxiliary batch into the Parquet stage."""
if not rows:
return
data = {field.name: self._pa.array([row[field.name] for row in rows], type=field.type) for field in schema}
path = os.path.join(self._dir, f"w{self._worker_index}_{name}_{self._sequence}.parquet")
self._sequence += 1
self._pq.write_table(self._pa.table(data, schema=schema), path)
self._files.setdefault(name, []).append(path)
def write_dispatch(self, rows: list[dict[str, Any]]) -> None:
"""Persist deferred dispatch payloads rather than returning a manifest list."""
self._write_auxiliary(
"_httk_dispatch_payload",
self._pa.schema(
[
self._pa.field("dispatch_name", self._pa.string()),
self._pa.field("content_id", self._pa.string()),
self._pa.field("column", self._pa.string()),
self._pa.field("block_sid", self._pa.int64()),
]
),
rows,
)
def write_nan_content(self, rows: list[dict[str, Any]]) -> None:
"""Persist NaN conflict diagnostics rather than returning a manifest set."""
self._write_auxiliary(
"_httk_nan_content",
self._pa.schema(
[
self._pa.field("table_name", self._pa.string()),
self._pa.field("content_id", self._pa.string()),
self._pa.field("field_name", self._pa.string()),
]
),
rows,
)
class _SqliteShardWriter:
"""Write one native-SQLite shard database per worker (the SQLite backend hand-off)."""
def __init__(self, store: SqlStore, worker_index: int, shard_dir: str) -> None:
self._store = store
self._path = os.path.join(shard_dir, f"w{worker_index}.sqlite")
self._connection = sqlite3.connect(self._path)
self._connection.execute("PRAGMA journal_mode=OFF")
self._connection.execute("PRAGMA synchronous=OFF")
self._created: set[str] = set()
def _columns(self, table_name: str) -> list[str]:
return [column.name for column in self._store._table(table_name).columns]
def write(self, table_name: str, rows: list[dict[str, Any]]) -> None:
if not rows:
return
columns = self._columns(table_name)
if table_name not in self._created:
definitions = ", ".join(f'"{name}"' for name in columns)
self._connection.execute(f'CREATE TABLE "{table_name}" ({definitions})')
self._created.add(table_name)
placeholders = ", ".join("?" for _ in columns)
self._connection.executemany(
f'INSERT INTO "{table_name}" VALUES ({placeholders})',
[tuple(row.get(name) for name in columns) for row in rows],
)
def finalize(self) -> dict[str, Any]:
self._connection.commit()
self._connection.close()
return {"db": self._path, "tables": sorted(self._created)}
class _DuckdbStageWriter:
"""Stream a serial stage through quote-all CSV and DuckDB ``COPY``.
A CSV field is paired with an explicit boolean null bitmap. This avoids a
magic ``NULLSTR`` altogether: empty strings, arbitrary Unicode, and every
possible text value remain distinct from SQL ``NULL``. CSV files are
written as the encoder flushes; only native bulk ``COPY`` and a set-wise
typed projection run at stage finish.
"""
def __init__(self, store: SqlStore, worker_index: int, shard_dir: str) -> None:
self._store = store
self._path = os.path.join(shard_dir, f"w{worker_index}.duckdb")
self._dir = shard_dir
self._csv: dict[str, tuple[Any, Any, list[str]]] = {}
def write(self, table_name: str, rows: list[dict[str, Any]]) -> None:
if not rows:
return
writer, _file, columns = self._csv_for(table_name)
for row in rows:
encoded: list[object] = []
for name in columns:
value = row.get(name)
encoded.extend(("" if value is None else self._csv_value(value), value is None))
writer.writerow(encoded)
def finalize(self) -> dict[str, Any]:
for _writer, file, _columns in self._csv.values():
file.close()
engine = sqlalchemy.create_engine(f"duckdb:///{self._path}")
try:
with engine.begin() as connection:
for table_name, (_writer, _file, columns) in self._csv.items():
source = self._store._table(table_name)
raw = f"_httk_stage_raw_{table_name}"
raw_columns = ", ".join(
f'"v{index}" VARCHAR, "n{index}" BOOLEAN' for index, _name in enumerate(columns)
)
connection.execute(sqlalchemy.text(f'CREATE TABLE "{raw}" ({raw_columns})'))
path = os.path.join(self._dir, f"{table_name}.csv").replace("'", "''")
connection.execute(
sqlalchemy.text(
f"COPY \"{raw}\" FROM '{path}' (FORMAT CSV, HEADER TRUE, QUOTE '\"', ESCAPE '\"', "
f"FORCE_NOT_NULL ({', '.join(repr(f'v{index}') for index in range(len(columns)))}))"
)
)
stage = sqlalchemy.Table(
table_name,
sqlalchemy.MetaData(),
*(sqlalchemy.Column(column.name, column.type) for column in source.columns),
)
connection.execute(sqlalchemy.schema.CreateTable(stage))
select = ", ".join(
f'CASE WHEN "n{index}" THEN NULL ELSE CAST("v{index}" AS {column.type.compile(dialect=engine.dialect)}) END'
for index, column in enumerate(source.columns)
)
names = ", ".join(f'"{column.name}"' for column in source.columns)
connection.execute(
sqlalchemy.text(f'INSERT INTO "{table_name}" ({names}) SELECT {select} FROM "{raw}"')
)
connection.execute(sqlalchemy.text(f'DROP TABLE "{raw}"'))
finally:
engine.dispose()
return {"format": "duckdb", "db": self._path, "tables": sorted(self._csv)}
def _csv_for(self, table_name: str) -> tuple[Any, Any, list[str]]:
existing = self._csv.get(table_name)
if existing is not None:
return existing
path = os.path.join(self._dir, f"{table_name}.csv")
file = open(path, "w", newline="", encoding="utf-8") # noqa: SIM115 — handle is cached in self._csv and closed by the shard lifecycle, not per call
columns = [column.name for column in self._store._table(table_name).columns]
writer = csv.writer(file, quoting=csv.QUOTE_ALL, lineterminator="\n")
writer.writerow([item for index in range(len(columns)) for item in (f"v{index}", f"n{index}")])
built = (writer, file, columns)
self._csv[table_name] = built
return built
@staticmethod
def _csv_value(value: Any) -> str:
if isinstance(value, bytes):
return "".join(f"\\x{byte:02x}" for byte in value)
return str(value)
def _make_writer(store: SqlStore, worker_index: int, config: _WorkerConfig) -> Any:
if config.backend in {"duckdb", "parquet", "clickhousedb"}:
return _ParquetShardWriter(store, worker_index, config.shard_dir)
if config.backend == "duckdb-stage":
return _DuckdbStageWriter(store, worker_index, config.shard_dir)
return _SqliteShardWriter(store, worker_index, config.shard_dir)
# --------------------------------------------------------------------- worker encoder
class _WorkerEncoder:
"""Encode a worker's slice of the stream into shard rows with block sids.
A stripped connection-free counterpart of
:meth:`~httk.store.backend.sql.bulk.BulkIngest._encode_active`: no table DDL, sids from
the worker's own block, and no metadata verification (the merge verifies
every surviving collision). Content records that carry a metadata plan are
emitted per occurrence — never deduplicated in the worker — so the merge
sees, and can compare, all of them.
"""
def __init__(self, store: SqlStore, worker_index: int, config: _WorkerConfig) -> None:
self._store = store
self._config = config
self._chunk_size = config.chunk_size
self._writer = _make_writer(store, worker_index, config)
self._base = _worker_base(worker_index)
self._registered: set[type] = set()
self._next_sid: dict[str, int] = {}
self._rows: dict[str, list[dict[str, Any]]] = {}
self._content_index: dict[str, dict[str, int]] = {}
self._value_index: dict[str, dict[tuple[Any, ...], int]] = {}
self._token_sid: dict[int, tuple[str, int]] = {}
self._encoded_count = 0
self._roots: list[tuple[str, int]] = []
self._late_role_roots: list[tuple[str, int]] = []
self._root_rows: list[dict[str, Any]] = []
self._dispatch: list[_DispatchRow] = []
self._dispatch_payload: list[dict[str, Any]] = []
self._tables: set[str] = set()
self._nan_content: set[tuple[str, str, str]] = set()
self._nan_rows: list[dict[str, Any]] = []
# (table, stage_sid, column) cells that held a NaN float. SQLite shards
# store NaN as NULL; only PostgreSQL (which keeps NaN) restores them at
# load time so bulk matches a serial ``save()``. Recorded only for the
# SQLite shard backend (Parquet/DuckDB shards preserve NaN natively).
self._nan_floats: set[tuple[str, int, str]] = set()
self._since_flush = 0
@property
def _deduplicate_in_worker(self) -> bool:
"""Whether this encoder needs in-memory keys to preserve public provisional sids.
Finalization collapses content/by-value occurrences set-wise. With
``track_sids=False`` callers have opted out of stable provisional
identities, so retaining one client key per record defeats bounded mode.
"""
return self._config.track_sids
# -- encoding
def save(
self,
token: int,
obj: Any,
as_record: type | None,
promote: frozenset[type] = frozenset(),
) -> int:
record_type = resolve_storage_record(obj, as_record=as_record)
projection = SaveProjection(store_timestamp=self._config.store_timestamp)
occurrences: list[tuple[type, Any, int]] = []
seen: set[tuple[type, int, int]] = set()
sid = self._encode(record_type, obj, projection, "", promote, occurrences, seen)
root = (record_type, id(obj), sid)
if root not in seen:
occurrences.append((record_type, obj, sid))
for occurrence_type, source, occurrence_sid in occurrences:
self._promote_occurrence(token, occurrence_type, source, projection, occurrence_sid)
table_name = resolve_schema(record_type).table_name
if self._config.track_sids:
self._token_sid[token] = (table_name, sid)
self._since_flush += 1
if self._since_flush >= self._chunk_size:
self._flush()
self._encoded_count += 1
return sid
def _promote_occurrence(
self,
token: int,
record_type: type,
source: Any,
projection: SaveProjection,
sid: int,
) -> None:
table_name = resolve_schema(record_type).table_name
marked = self._promote_buffered_role(table_name, sid)
if self._config.track_sids and not self._config.spill_deferred_auxiliary:
self._roots.append((table_name, sid))
if not marked:
self._late_role_roots.append((table_name, sid))
if self._config.spill_deferred_auxiliary:
self._root_rows.append({"token": token, "tbl": table_name, "stage_sid": sid})
family = self._store._family_for_backing(record_type)
if (
(self._config.track_sids or self._config.spill_deferred_auxiliary)
and family is not None
and len(family.records) > 1
):
dispatch = _DispatchRow(
dispatch_name=entry_dispatch_table_name(family.name),
key=projection.content_id(record_type, source),
column=backing_dispatch_column_name(family.record_names[family.records.index(record_type)]),
all_columns=tuple(backing_dispatch_column_name(name) for name in family.record_names),
ref_table=table_name,
block_sid=sid,
family_name=family.name,
)
if self._config.spill_deferred_auxiliary:
self._dispatch_payload.append(
{
"dispatch_name": dispatch.dispatch_name,
"content_id": dispatch.key,
"column": dispatch.column,
"block_sid": dispatch.block_sid,
}
)
else:
self._dispatch.append(dispatch)
def _encode(
self,
record_type: type,
source: Any,
projection: SaveProjection,
path: str,
promote: frozenset[type],
occurrences: list[tuple[type, Any, int]],
seen: set[tuple[type, int, int]],
) -> int:
active_key = (record_type, id(source))
if active_key in projection.active:
raise StorageProjectionCycleError(path, record_type)
projection.active.add(active_key)
try:
sid = self._encode_active(record_type, source, projection, path, promote, occurrences, seen)
key = (record_type, id(source), sid)
if record_type in promote and key not in seen:
seen.add(key)
occurrences.append((record_type, source, sid))
return sid
finally:
projection.active.remove(active_key)
def _encode_active(
self,
record_type: type,
source: Any,
projection: SaveProjection,
path: str,
promote: frozenset[type],
occurrences: list[tuple[type, Any, int]],
seen: set[tuple[type, int, int]],
) -> int:
schema = resolve_schema(record_type)
self._register(record_type)
table_name = schema.table_name
self._next_sid.setdefault(table_name, self._base + 1)
projected = projection.projector(record_type, source)
validation_key = (record_type, id(source))
if type(source) is record_type and validation_key not in projection.validated:
validator = vars(record_type).get("__httk_validate__")
if validator is not None:
validator.__get__(None, record_type)(source)
projection.validated.add(validation_key)
def resolve_sid(referenced_type: type, value: Any, field_path: str) -> int:
return self._encode(referenced_type, value, projection, field_path, promote, occurrences, seen)
dedup_content = schema.dedup == "content_id" and _metadata_plan(record_type) is None
key: str | None = None
if schema.dedup == "content_id":
key = projection.content_id(record_type, source)
if dedup_content and self._deduplicate_in_worker:
existing = self._content_index.setdefault(table_name, {}).get(key)
if existing is not None:
_encode_promoted_descendants(
schema, source, projected, path, existing, promote, resolve_sid, references=True
)
return existing
values = _encode_parent_row(schema, source, projected, path, resolve_sid)
if record_type in self._store._entry_record_types:
entry_id = values.get("id")
immutable_id = values.get("immutable_id")
if entry_id is None and self._store._entry_ids is None:
raise ValueError(
f"{record_type.__name__} has no id and SqlStore(entry_ids=EntryIdScheme(...)) was not declared; "
"pass an explicit id or declare a scheme"
)
if entry_id is not None:
check_entry_id(entry_id)
if immutable_id is not None:
check_immutable_id(immutable_id)
if schema.dedup == "by_value":
value_tuple = tuple(sorted(values.items()))
if self._deduplicate_in_worker:
existing = self._value_index.setdefault(table_name, {}).get(value_tuple)
if existing is not None:
_encode_promoted_descendants(
schema, source, projected, path, existing, promote, resolve_sid, references=False
)
return existing
sid = self._next_sid[table_name]
if sid - self._base >= _SID_BLOCK:
raise RuntimeError(
f"bulk_ingest worker exceeded its {_SID_BLOCK} sid block for table {table_name!r}; "
"reduce the worker count or split the ingest"
)
self._next_sid[table_name] = sid + 1
# Bulk ingest is mains only: alt_id self-references (== sid), alt_kind
# stays NULL (omitted here, filled NULL by the column default).
row = {SID_COLUMN: sid, ROLE_COLUMN: 0, LOGICAL_ID_COLUMN: sid, ALT_ID_COLUMN: sid, **values}
if self._config.store_timestamp is not None:
row[STORE_TIMESTAMP_COLUMN] = self._config.store_timestamp
if key is not None:
row[CONTENT_ID_COLUMN] = key
for field_name, column_name in _plain_float_skip_fields(record_type):
candidate = values.get(column_name)
if isinstance(candidate, float) and math.isnan(candidate):
# Report every NaN field (not just the first): the merge picks
# the schema-order-first among them, deterministically.
if self._config.spill_deferred_auxiliary:
self._nan_rows.append({"table_name": table_name, "content_id": key, "field_name": field_name})
else:
self._nan_content.add((table_name, key, field_name))
if dedup_content and self._deduplicate_in_worker:
self._content_index[table_name][key] = sid
elif schema.dedup == "by_value" and self._deduplicate_in_worker:
self._value_index[table_name][tuple(sorted(values.items()))] = sid
self._buffer(table_name, row)
self._record_nan_floats(table_name, row)
for spec in schema.fields:
if spec.role != "child":
continue
assert spec.child is not None
child_rows = _encode_child_rows(
schema,
spec,
sid,
SqlStore._projected_value(record_type, source, projected, spec),
_field_path(path, spec.field),
resolve_sid,
)
for child_row in child_rows:
self._buffer(spec.child.table_name, child_row)
return sid
def _promote_buffered_role(self, table_name: str, sid: int) -> bool:
for row in self._rows.get(table_name, ()):
if row[SID_COLUMN] == sid:
row[ROLE_COLUMN] = 1
return True
return False
def _register(self, record_type: type) -> None:
if record_type in self._registered:
return
self._store._register_tables((record_type,))
self._registered.add(record_type)
def _buffer(self, table_name: str, row: dict[str, Any]) -> None:
self._rows.setdefault(table_name, []).append(row)
self._tables.add(table_name)
def _record_nan_floats(self, table_name: str, row: dict[str, Any]) -> None:
"""Note any NaN float cell so a PostgreSQL load can reinstate it.
A SQLite shard flattens NaN to NULL; recording ``(table, sid, column)``
lets the PostgreSQL stage/merge load restore ``NaN`` and stay identical
to a serial ``save()``. Only SQLite shards lose NaN -- both the SQLite
and PostgreSQL backends use them -- so this is keyed off the writer, not
the dialect name; the Parquet/DuckDB shard backend preserves NaN itself.
Scoped to sid-bearing parent rows, matching the ``_nan_content`` sidecar:
child element tables are keyed by ``(parent_sid, index)`` rather than a
sid, so a NaN inside a ``list[float]`` element is out of this path.
"""
if not isinstance(self._writer, _SqliteShardWriter) or SID_COLUMN not in row:
return
sid = row[SID_COLUMN]
for column_name, candidate in row.items():
if isinstance(candidate, float) and math.isnan(candidate):
self._nan_floats.add((table_name, sid, column_name))
def _flush(self) -> None:
for table_name, rows in self._rows.items():
if rows:
self._writer.write(table_name, rows)
rows.clear()
if self._root_rows:
self._writer.write_roots(self._root_rows)
self._root_rows.clear()
if self._config.spill_deferred_auxiliary:
self._writer.write_dispatch(self._dispatch_payload)
self._dispatch_payload.clear()
self._writer.write_nan_content(self._nan_rows)
self._nan_rows.clear()
self._since_flush = 0
def finish(self) -> _WorkerManifest:
self._flush()
return _WorkerManifest(
worker_index=-1, # filled by the caller
encoded_count=self._encoded_count,
token_sid=self._token_sid,
roots=self._roots,
late_role_roots=self._late_role_roots,
dispatch=self._dispatch,
tables=sorted(self._tables),
shards=self._writer.finalize(),
nan_content=sorted(self._nan_content),
nan_floats=sorted(self._nan_floats),
)
# --------------------------------------------------------------------- worker process
def _worker_main(worker_index: int, task_queue: Any, result_queue: Any) -> None:
"""Worker process entry point: encode tasks into shards, then report a manifest.
Tasks arrive as pickled ``(token, obj, as_record, promote)`` byte strings (the main
process pickles synchronously, so an unpicklable object fails the caller's
``save`` promptly instead of vanishing in a queue feeder thread). The worker
never touches the store's database. On completion (or failure) it flushes its
result onto ``result_queue`` and exits with :func:`os._exit` to skip
interpreter finalizers that might disturb the fork-inherited engine.
:param worker_index: The worker's index (its sid block and shard names).
:param task_queue: This worker's task queue of pickled tasks (``None`` stops).
:param result_queue: The queue the manifest or an exception is reported on.
:return: None.
"""
assert _PARENT_STORE is not None and _PARENT_CONFIG is not None
encoder = _WorkerEncoder(_PARENT_STORE, worker_index, _PARENT_CONFIG)
try:
while True:
item = task_queue.get()
if item is None:
break
token, obj, as_record, promote = pickle.loads(item)
encoder.save(token, obj, as_record, promote)
manifest = encoder.finish()
manifest.worker_index = worker_index
_report(result_queue, (worker_index, "ok", manifest))
except BaseException as error: # faithfully relayed to the caller
_report(result_queue, (worker_index, "error", _as_reportable(error)))
os._exit(0)
def _report(result_queue: Any, payload: Any) -> None:
result_queue.put(payload)
result_queue.close()
result_queue.join_thread()
def _as_reportable(error: BaseException) -> BaseException:
"""Return an exception that survives pickling back to the main process."""
try:
import pickle
pickle.loads(pickle.dumps(error))
except Exception: # fall back to a faithful-typed surrogate
return RuntimeError(f"{type(error).__name__}: {error}")
return error
# --------------------------------------------------------------------- pool controller
[docs]
class ParallelController:
"""Own the worker pool, task dispatch, and shard directory for one parallel ingest.
Each worker has its own task queue; ``dispatch`` routes token ``k`` to worker
``k % workers`` (deterministic round-robin), so the record order the caller
saves fully determines which worker encodes each record. A shared result
queue carries each worker's manifest (or exception) back.
"""
def __init__(
self,
store: SqlStore,
*,
workers: int,
chunk_size: int,
backend: str,
track_sids: bool = True,
store_timestamp: int | None = None,
spill_deferred_auxiliary: bool = False,
) -> None:
import multiprocessing
if workers > _MAX_WORKERS:
raise ValueError(f"bulk_ingest supports at most {_MAX_WORKERS} workers (sid-block limit)")
self._store = store
self._workers = workers
self._context = multiprocessing.get_context("fork")
self._temp = tempfile.TemporaryDirectory(prefix="httk_bulk_", dir=_shard_parent_dir(store))
self._config = _WorkerConfig(
chunk_size=chunk_size,
shard_dir=self._temp.name,
backend=backend,
track_sids=track_sids,
store_timestamp=store_timestamp,
spill_deferred_auxiliary=spill_deferred_auxiliary,
)
self._queues: list[Any] = [self._context.Queue(maxsize=_QUEUE_MAXSIZE) for _ in range(workers)]
self._result_queue: Any = self._context.Queue()
self._processes: list[Any] = []
# Every result consumed by health polling is cached here (not just errors),
# so a worker that reports and exits cleanly while a sibling's queue is full
# keeps its manifest and is not mistaken for a crash.
self._results_cache: dict[int, tuple[str, Any]] = {}
self._closed = False
[docs]
def start(self) -> None:
import warnings
global _PARENT_STORE, _PARENT_CONFIG
_PARENT_STORE = self._store
_PARENT_CONFIG = self._config
try:
for index in range(self._workers):
process = self._context.Process(
target=_worker_main,
args=(index, self._queues[index], self._result_queue),
daemon=True,
)
with warnings.catch_warnings():
# Forking is required: workers inherit the (unpicklable) store
# and never touch its database, so Python 3.12's multi-threaded
# fork() advisory does not apply here.
warnings.simplefilter("ignore", DeprecationWarning)
process.start()
self._processes.append(process)
finally:
# The children have forked; the parent no longer needs the globals.
_PARENT_STORE = None
_PARENT_CONFIG = None
[docs]
def dispatch(
self,
token: int,
obj: Any,
as_record: type | None,
promote: frozenset[type] = frozenset(),
) -> None:
"""Pickle the task synchronously and enqueue it on its worker (routed by token)."""
# Pickle here, in the caller's thread: an unpicklable object raises out of
# ``save`` promptly rather than being silently dropped by a queue feeder.
payload = pickle.dumps((token, obj, as_record, promote))
queue = self._queues[token % self._workers]
while True:
self._raise_if_worker_broken()
try:
queue.put(payload, timeout=0.5)
return
except queue_mod.Full:
continue
def _raise_if_worker_broken(self) -> None:
"""Raise if any worker reported an error or exited *without* reporting (a crash or kill).
Results are cached (both ``ok`` and ``error``), so a worker that reported
and then exited cleanly — e.g. it took its stop sentinel while a sibling's
queue was still full — is recognized as done, not misreported as a crash,
and its manifest survives for :meth:`finish`.
"""
self._drain_results(self._results_cache)
for status, payload in self._results_cache.values():
if status == "error":
raise _forward(payload)
for index, process in enumerate(self._processes):
if process.exitcode is not None and index not in self._results_cache:
raise RuntimeError(
"a bulk_ingest worker exited unexpectedly (crashed or was killed); the ingest is aborted"
)
[docs]
def finish(self) -> list[_WorkerManifest]:
"""Signal completion, collect every worker's manifest, and re-raise the first error.
A worker that exits without reporting (a crash or an external kill) is
detected by its exit code and aborts the ingest, so a lost task can never
reach the merge. Both the stop-sentinel sends and the result waits are
bounded and interleaved with health checks, so a worker that dies with a
full queue cannot deadlock the main process.
"""
import time
self._send_sentinels()
# Start from whatever health polling already consumed (sending the
# sentinels may have drained some workers' results into the cache).
results: dict[int, tuple[str, Any]] = dict(self._results_cache)
error: BaseException | None = None
last_progress = time.monotonic()
while len(results) < self._workers:
try:
worker_index, status, payload = self._result_queue.get(timeout=1.0)
results[worker_index] = (status, payload)
last_progress = time.monotonic()
except queue_mod.Empty:
if all(process.exitcode is not None for process in self._processes):
self._drain_results(results)
break
if time.monotonic() - last_progress > _WORKER_STALL_TIMEOUT:
error = RuntimeError("bulk_ingest workers stopped making progress; aborting")
break
if len(results) < self._workers and error is None:
error = RuntimeError("a bulk_ingest worker exited without reporting a result (crashed or was killed)")
for status, payload in results.values():
if status == "error" and error is None:
error = _forward(payload)
for process in self._processes:
process.join(timeout=30)
if error is not None:
raise error
return [payload for status, payload in results.values() if status == "ok"]
def _send_sentinels(self) -> None:
"""Put a stop sentinel on each worker queue, aborting if a worker died with a full queue."""
import time
pending = list(range(self._workers))
deadline = time.monotonic() + _WORKER_STALL_TIMEOUT
while pending:
still_pending: list[int] = []
for index in pending:
try:
self._queues[index].put(None, timeout=0.1)
except queue_mod.Full:
still_pending.append(index)
pending = still_pending
if not pending:
return
# A queue that will not accept the sentinel belongs to a worker that
# is no longer draining it — detect the crash/kill and abort.
self._raise_if_worker_broken()
if time.monotonic() > deadline:
raise RuntimeError("bulk_ingest could not signal completion to its workers; aborting")
def _drain_results(self, results: dict[int, tuple[str, Any]]) -> None:
"""Absorb any results still queued after every worker has exited (avoids a report/exit race)."""
try:
while True:
worker_index, status, payload = self._result_queue.get_nowait()
results[worker_index] = (status, payload)
except queue_mod.Empty:
return
[docs]
def close(self) -> None:
"""Terminate any live workers and remove the shard directory (idempotent)."""
if self._closed:
return
self._closed = True
for process in self._processes:
if process.is_alive():
process.terminate()
for process in self._processes:
process.join(timeout=10)
# Cancel each queue's feeder thread before closing: a queue left non-empty
# by an aborted ingest would otherwise block ``close`` on the feeder join.
for queue in (*self._queues, self._result_queue):
queue.cancel_join_thread()
queue.close()
self._temp.cleanup()
def _forward(payload: Any) -> BaseException:
return payload if isinstance(payload, BaseException) else RuntimeError(str(payload))
def _shard_parent_dir(store: SqlStore) -> str | None:
"""The directory shards are created in: next to a file-backed database, else the tempfile default."""
try:
database = store._database.engine.url.database
except Exception: # any odd URL falls back to the default temp root
return None
if not database or database == ":memory:":
return None
parent = os.path.dirname(os.path.abspath(database))
return parent if os.path.isdir(parent) else None
def _copy_sqlite_shard(
connection: sqlalchemy.Connection,
database: str,
tables: list[str] | tuple[str, ...],
destinations: dict[str, sqlalchemy.Table],
nan_cells: dict[str, dict[int, set[str]]] | None = None,
) -> None:
"""Stream one SQLite shard into tables on the destination connection.
``nan_cells`` (``{table: {sid: {column, ...}}}``) names cells that held a NaN
float the SQLite shard flattened to NULL. It is supplied only for a
PostgreSQL destination (which keeps NaN); each such cell is reinstated as
``NaN`` during the insert so bulk matches a serial ``save()`` and a NOT NULL
float column never sees the intermediate NULL.
"""
uri = Path(database).resolve().as_uri() + "?mode=ro"
source = sqlite3.connect(uri, uri=True)
try:
for table_name in tables:
destination = destinations[table_name]
columns = [column.name for column in destination.columns]
quoted = ", ".join('"' + name.replace('"', '""') + '"' for name in columns)
table = table_name.replace('"', '""')
table_nan = nan_cells.get(table_name) if nan_cells else None
cursor = source.execute(f'SELECT {quoted} FROM "{table}"')
while rows := cursor.fetchmany(_SQLITE_COPY_BATCH_SIZE):
records = [dict(zip(columns, row, strict=True)) for row in rows]
if table_nan:
for record in records:
restore = table_nan.get(record[SID_COLUMN])
if restore:
for column_name in restore:
record[column_name] = math.nan
connection.execute(sqlalchemy.insert(destination), records)
finally:
source.close()
def _nan_cells_from_manifests(manifests: list[_WorkerManifest]) -> dict[str, dict[int, set[str]]]:
"""Index every manifest's recorded NaN float cells as ``{table: {sid: {column}}}``."""
cells: dict[str, dict[int, set[str]]] = {}
for manifest in manifests:
for table_name, sid, column in manifest.nan_floats:
cells.setdefault(table_name, {}).setdefault(sid, set()).add(column)
return cells
# --------------------------------------------------------------------- merge (main process)
[docs]
def merge(ingest: "httk.store.backend.sql.bulk.BulkIngest", manifests: list[_WorkerManifest]) -> None:
"""Load every worker shard, collapse cross-worker duplicates, and compact the sids.
Runs in the main process inside the ingest's spanning transaction.
:param ingest: The owning bulk-ingest context (its connection and store).
:param manifests: One manifest per finished worker.
:return: None.
"""
_Merger(ingest, manifests).run()
def _rebuild_untracked_dispatch(
connection: sqlalchemy.Connection,
store: SqlStore,
active_tables: set[str],
) -> dict[str, int]:
"""Build multi-record dispatch tables set-wise from surviving top-level backings."""
counts: dict[str, int] = {}
for family in store.layout.families:
if len(family.records) < 2:
continue
dispatch_name = entry_dispatch_table_name(family.name)
if dispatch_name not in active_tables:
continue
dispatch = store._table(dispatch_name)
columns = [column.name for column in dispatch.columns]
candidates = []
for record_name, record_type in zip(family.record_names, family.records, strict=True):
backing = store._table(resolve_schema(record_type).table_name)
target_column = backing_dispatch_column_name(record_name)
candidates.append(
sqlalchemy.select(
*(
backing.c[CONTENT_ID_COLUMN].label(column.name)
if column.name == DISPATCH_CONTENT_ID_COLUMN
else (
backing.c[SID_COLUMN].label(column.name)
if column.name == target_column
else sqlalchemy.cast(sqlalchemy.null(), column.type).label(column.name)
)
for column in dispatch.columns
)
).where(backing.c[ROLE_COLUMN] == 1)
)
rows = sqlalchemy.union_all(*candidates).subquery()
conflict = connection.execute(
sqlalchemy.select(rows.c[DISPATCH_CONTENT_ID_COLUMN])
.group_by(rows.c[DISPATCH_CONTENT_ID_COLUMN])
.having(sqlalchemy.func.count() > 1)
.limit(1)
).first()
if conflict is not None:
raise EntryDispatchIntegrityError(
f"entry dispatch {family.name!r} maps content_id {conflict[0]!r} to a conflicting backing row"
)
count = int(connection.execute(sqlalchemy.select(sqlalchemy.func.count()).select_from(rows)).scalar_one())
if count:
connection.execute(
sqlalchemy.insert(dispatch).from_select(columns, sqlalchemy.select(*(rows.c[name] for name in columns)))
)
counts[dispatch_name] = count
return counts
class _Merger:
"""The set-wise shard merge for a parallel ingest (see :func:`merge`)."""
def __init__(self, ingest: "BulkIngest", manifests: list[_WorkerManifest]) -> None:
self._ingest = ingest
self._store = ingest._store
assert ingest._connection is not None
self._connection = ingest._connection
self._manifests = manifests
self._graph = ingest._logical_graph()
self._fk_columns = self._graph.sid_columns()
self._referrers = {name: list(self._graph.referrers(name)) for name in self._graph.tables}
# (table, block_sid) -> keep_sid after cross-worker collapse.
self._collapse: dict[tuple[str, int], int] = {}
# (table, sid) -> compact_sid after final renumbering.
self._compaction: dict[tuple[str, int], int] = {}
# table -> {content id -> set of fields} whose identity-excluded float held
# a NaN. A set (not last-manifest-wins) keeps attribution deterministic:
# the merge names the schema-order-first field among the reported set.
self._nan_content: dict[str, dict[str, set[str]]] = {}
for manifest in manifests:
for table_name, content_id, field_name in manifest.nan_content:
self._nan_content.setdefault(table_name, {}).setdefault(content_id, set()).add(field_name)
def run(self) -> None:
self._load_shards()
for name in self._graph.dependency_order(self._store._metadata.tables):
table = self._store._metadata.tables[name]
schema = self._ingest._parent_schema.get(table.name)
if schema is None:
continue
if schema.dedup == "content_id":
self._collapse_content(table, schema)
# A collapse in one by-value table can make a normalized key in a
# mutually-referential table equal only on the next pass. Iterate the
# full deterministic order to the graph-wide fixpoint before orphan
# sweeping (the former one-pass order under-collapsed A <-> B graphs).
by_value = [
(self._store._metadata.tables[name], self._ingest._parent_schema[name])
for name in self._graph.dependency_order(self._store._metadata.tables)
if name in self._ingest._parent_schema and self._ingest._parent_schema[name].dedup == "by_value"
]
while True:
changed = False
for table, schema in by_value:
changed = self._collapse_by_value(table, schema) or changed
if not changed:
break
self._promote_late_roles()
self._sweep_orphans()
self._verify_entry_id_conflicts()
# A declared family whose table was never written stays uncreated (deferred DDL only
# happens on write); the compaction pass must treat such a missing table as empty
# rather than querying it, matching the module's lazy-DDL read-path principle.
existing = actual_table_names(self._connection)
for name in self._graph.dependency_order(self._store._metadata.tables):
table = self._store._metadata.tables[name]
if name in existing and SID_COLUMN in table.c:
self._compact(table)
self._merge_dispatch()
self._populate_resolved_map()
def _verify_entry_id_conflicts(self) -> None:
"""Check surviving staged explicit ids across every backing of each family."""
existing_names = actual_table_names(self._connection)
for family in self._store.layout.families:
if family.definition_id is None:
continue
tables = [
self._store._table(resolve_schema(record).table_name)
for record in family.records
if resolve_schema(record).table_name in existing_names
]
if not tables:
continue
for field_name in ("id", "immutable_id"):
staged = sqlalchemy.union_all(
*(
sqlalchemy.select(table.c[field_name].label("value")).where(
table.c[SID_COLUMN] >= _SID_BLOCK, table.c[field_name].is_not(None)
)
for table in tables
)
).subquery()
duplicate = self._connection.execute(
sqlalchemy.select(staged.c.value)
.group_by(staged.c.value)
.having(sqlalchemy.func.count() > 1)
.limit(1)
).first()
if duplicate is not None:
raise EntryIdConflictError(family.name, str(duplicate[0]), None, None)
for staged_table in tables:
for owned_table in tables:
staged_alias = staged_table.alias("staged")
owned_alias = owned_table.alias("owned")
conflict = self._connection.execute(
sqlalchemy.select(staged_alias.c[field_name])
.join(owned_alias, staged_alias.c[field_name] == owned_alias.c[field_name])
.where(
staged_alias.c[SID_COLUMN] >= _SID_BLOCK,
staged_alias.c[field_name].is_not(None),
owned_alias.c[SID_COLUMN] < _SID_BLOCK,
)
.limit(1)
).first()
if conflict is not None:
raise EntryIdConflictError(owned_table.name, str(conflict[0]), None, None)
def _populate_resolved_map(self) -> None:
"""Map every sid ``save`` returned (a synthetic token) to its durable stored sid."""
ingest = self._ingest
for manifest in self._manifests:
for token, (table_name, block_sid) in manifest.token_sid.items():
ingest._resolved_map[(table_name, token)] = self._final_sid(table_name, block_sid)
# -- shard loading
def _load_shards(self) -> None:
backend = self._connection.dialect.name
if backend == "duckdb":
self._load_parquet_shards()
else:
self._load_sqlite_shards()
def _table_columns(self, table_name: str) -> list[str]:
return [column.name for column in self._store._table(table_name).columns]
def _load_parquet_shards(self) -> None:
files_by_table: dict[str, list[str]] = {}
for manifest in self._manifests:
for table_name, files in manifest.shards.items():
if table_name == "_httk_roots":
continue
files_by_table.setdefault(table_name, []).extend(files)
for table_name, files in files_by_table.items():
if not files:
continue
columns = ", ".join(f'"{name}"' for name in self._table_columns(table_name))
# Bind each shard path as a parameter — a path containing a quote must
# not break or inject into the SQL.
placeholders = ", ".join(f":f{index}" for index in range(len(files)))
statement = sqlalchemy.text(
f'INSERT INTO "{table_name}" ({columns}) SELECT {columns} FROM read_parquet([{placeholders}])'
).bindparams(**{f"f{index}": path for index, path in enumerate(files)})
self._connection.execute(statement)
def _load_sqlite_shards(self) -> None:
# Read each shard through its own connection: SQLite's compiled ATTACH
# ceiling is commonly ten and cannot always be raised at runtime. The
# destination inserts still belong to the merge's one spanning transaction.
nan_cells = (
_nan_cells_from_manifests(self._manifests) if self._connection.dialect.name == "postgresql" else None
)
for manifest in self._manifests:
database = manifest.shards.get("db")
tables = manifest.shards.get("tables", [])
if not database or not tables:
continue
destinations = {table_name: self._store._table(table_name) for table_name in tables}
_copy_sqlite_shard(self._connection, str(database), tables, destinations, nan_cells)
# -- cross-worker collapse
def _collapse_content(self, table: sqlalchemy.Table, schema: TableSchema) -> None:
keep = (
sqlalchemy.select(table.c[CONTENT_ID_COLUMN], sqlalchemy.func.min(table.c[SID_COLUMN]).label("keep"))
.group_by(table.c[CONTENT_ID_COLUMN])
.subquery()
)
statement = (
sqlalchemy.select(table.c[SID_COLUMN], keep.c.keep)
.join_from(table, keep, table.c[CONTENT_ID_COLUMN] == keep.c[CONTENT_ID_COLUMN])
.where(table.c[SID_COLUMN] != keep.c.keep)
)
pairs = [(int(row[0]), int(row[1])) for row in self._connection.execute(statement).all()]
if not pairs:
return
if self._ingest._verify_metadata and _metadata_plan(schema.cls) is not None:
self._verify_collision_metadata(table, schema)
self._apply_collapse(table, schema, pairs)
def _collapse_by_value(self, table: sqlalchemy.Table, schema: TableSchema) -> bool:
value_columns = [
column.name
for column in table.columns
if column.name
not in (SID_COLUMN, ROLE_COLUMN, STORE_TIMESTAMP_COLUMN, LOGICAL_ID_COLUMN, ALT_ID_COLUMN, ALT_KIND_COLUMN)
]
while True:
keep = (
sqlalchemy.select(
*(table.c[name] for name in value_columns),
sqlalchemy.func.min(table.c[SID_COLUMN]).label("keep"),
)
.group_by(*(table.c[name] for name in value_columns))
.subquery()
)
condition = sqlalchemy.and_(*(table.c[name].is_not_distinct_from(keep.c[name]) for name in value_columns))
statement = (
sqlalchemy.select(table.c[SID_COLUMN], keep.c.keep)
.join_from(table, keep, condition)
.where(table.c[SID_COLUMN] != keep.c.keep)
)
pairs = [(int(row[0]), int(row[1])) for row in self._connection.execute(statement).all()]
if not pairs:
return False
self._apply_collapse(table, schema, pairs)
# _apply_collapse rewrites references and deletes all pairs in one
# operation, so a second local pass is only needed for rows made
# equal by the remap.
return True
def _apply_collapse(self, table: sqlalchemy.Table, schema: TableSchema, pairs: list[tuple[int, int]]) -> None:
name = table.name
# Role is a monotone bookkeeping bit, not part of either identity. A
# canonical bulk row is main whenever any collapsed occurrence was a
# top-level record.
for old, keep in pairs:
old_role = self._connection.execute(
sqlalchemy.select(table.c[ROLE_COLUMN]).where(table.c[SID_COLUMN] == old)
).scalar_one()
if int(old_role) == 1:
self._connection.execute(
sqlalchemy.update(table).where(table.c[SID_COLUMN] == keep).values({ROLE_COLUMN: 1})
)
for old, keep in pairs:
self._collapse[(name, old)] = keep
child_links = {
(edge.target_table, edge.target_column)
for edge in self._graph.ownership()
if edge.source_table == name and edge.target_column is not None
}
map_table = self._make_map_table(pairs)
try:
for referrer_table, column in self._referrers.get(name, ()):
if (referrer_table, column) in child_links:
# A collapsed parent's own child rows are dropped, not repointed;
# the surviving parent already carries its own children.
self._delete_where_in_map(referrer_table, column, map_table)
else:
self._remap_column(referrer_table, column, map_table)
self._delete_where_in_map(name, SID_COLUMN, map_table)
finally:
self._drop_map_table(map_table)
def _verify_collision_metadata(self, table: sqlalchemy.Table, schema: TableSchema) -> None:
"""Set-wise verify that rows sharing a content id agree on their identity-excluded metadata.
Only the columns that actually carry identity-excluded metadata are
compared: the ``IdentitySkip`` scalar columns, and a skipped reference's
sid column when its target is content-addressed (equal content collapses
to one sid before this runs, so a differing sid is a differing skipped
reference). ``descend`` conflicts — a non-skipped reference whose target
carries the skip metadata — surface at that target table's own collapse,
where the metadata lives, because equal-content parents reference
equal-content (hence collapsed-together) targets. One grouped scan per
content table therefore replaces reconstructing each duplicate record,
the dominant cost at real-build scale.
:param table: The content-addressed record table being collapsed.
:param schema: The table's resolved schema (for the diagnostic record name).
:raises httk.store.store_common.EntryMetadataConflictError: If a content id occurs with differing metadata.
"""
compare_columns = self._metadata_compare_columns(schema)
if not compare_columns:
return
# NaN scan first: serial treats ``NaN != NaN`` as a conflict, but SQL
# equality groups a NaN with itself (DuckDB's total order) and SQLite
# stores no NaN at all, so the exact scan below cannot see it. Running it
# first also gives a NaN in an earlier schema field priority over a plain
# value difference in a later one, matching serial's field-order checks.
# A duplicated content id the workers flagged as NaN-bearing is a conflict;
# the field named is the schema-order-first of the reported set.
nan_by_content = self._nan_content.get(table.name)
if nan_by_content:
duplicated_nan = self._connection.execute(
sqlalchemy.select(table.c[CONTENT_ID_COLUMN])
.where(table.c[CONTENT_ID_COLUMN].in_(sorted(nan_by_content)))
.group_by(table.c[CONTENT_ID_COLUMN])
.having(sqlalchemy.func.count() > 1)
.limit(1)
).first()
if duplicated_nan is not None:
key = duplicated_nan[0]
reported = nan_by_content.get(key, set())
field_name = next(
(field for _column, field in compare_columns if field in reported),
schema.cls.__name__,
)
raise EntryMetadataConflictError(
f"metadata conflict for {schema.cls.__name__}.{field_name}: content id {key!r} occurs with "
"a NaN identity-excluded value that never equals itself"
)
# Exact-difference scan: SQL ``=`` matches serial's scalar equality for
# every finite value — ``-0.0 == 0.0`` and ``NULL``/``None`` both group as
# equal — so a content id whose group has more than one distinct tuple
# carries differing identity-excluded metadata.
column_names = [column for column, _field in compare_columns]
selected = [table.c[CONTENT_ID_COLUMN], *(table.c[name] for name in column_names)]
distinct_rows = sqlalchemy.select(*selected).distinct().subquery()
conflicting = self._connection.execute(
sqlalchemy.select(distinct_rows.c[CONTENT_ID_COLUMN])
.group_by(distinct_rows.c[CONTENT_ID_COLUMN])
.having(sqlalchemy.func.count() > 1)
.limit(1)
).first()
if conflicting is not None:
self._raise_metadata_conflict(table, schema, compare_columns, conflicting[0])
def _raise_metadata_conflict(
self, table: sqlalchemy.Table, schema: TableSchema, compare_columns: list[tuple[str, str]], key: str
) -> None:
differing_field = self._first_differing_field(table, compare_columns, key)
field_name = f"{schema.cls.__name__}.{differing_field}" if differing_field else schema.cls.__name__
raise EntryMetadataConflictError(
f"metadata conflict for {field_name}: content id {key!r} occurs with differing identity-excluded metadata"
)
@staticmethod
def _metadata_compare_columns(schema: TableSchema) -> list[tuple[str, str]]:
"""The ``(column, field)`` pairs whose within-group difference is an identity-excluded conflict."""
plan = _metadata_plan(schema.cls)
if plan is None:
return []
columns: list[tuple[str, str]] = []
for spec in plan.skipped_specs:
if spec.codec_name == "float":
# The plain-float codec stores an exact text companion for lossless
# reconstruction, but ``-0.0`` and ``0.0`` differ there while serial's
# ``_metadata_scalar_equal`` compares them with IEEE ``==``. Compare
# only the float column (whose SQL equality is IEEE) and drop the
# string companion; NaN is handled separately below. Exact numeric
# codecs (fraction, fracscalar, tensors) keep their exact text
# channel — dropping it would collapse them to a float approximation
# and silently accept e.g. Fraction(2**53) vs Fraction(2**53 + 1).
columns.extend((column.name, spec.field) for column in spec.columns if column.kind == "float")
else:
columns.extend((column.name, spec.field) for column in spec.columns)
for spec in plan.skipped_nested:
if (
spec.role == "reference"
and spec.target is not None
and resolve_schema(spec.target).dedup in ("content_id", "by_value")
):
columns.append((spec.columns[0].name, spec.field))
return columns
def _first_differing_field(
self, table: sqlalchemy.Table, compare_columns: list[tuple[str, str]], key: str
) -> str | None:
"""The schema field of the first compared column that differs within the conflicting group."""
for column, field_name in compare_columns:
distinct = self._connection.execute(
sqlalchemy.select(sqlalchemy.func.count(sqlalchemy.distinct(table.c[column]))).where(
table.c[CONTENT_ID_COLUMN] == key
)
).scalar_one()
if distinct is not None and int(distinct) > 1:
return field_name
return None
# -- orphan sweep
def _sweep_orphans(self) -> None:
"""Delete rows no longer reachable from a surviving top-level record.
A collapsed duplicate parent drops its subtree; descendants that other
surviving records also reach stay, but a duplicate's private ``dedup="none"``
(or otherwise non-deduplicated) descendants become unreachable and must go,
matching the per-record ``save()`` loop's result.
"""
reach = self._reachable_table()
if self._ingest._track_sids:
self._insert_reach(reach, self._survivor_seeds())
else:
self._insert_role_roots(reach)
self._close_reachability(reach)
self._delete_unreached(reach)
self._connection.execute(sqlalchemy.schema.DropTable(reach, if_exists=True))
def _survivor_seeds(self) -> list[tuple[str, int]]:
seeds: set[tuple[str, int]] = set()
for manifest in self._manifests:
for table_name, block_sid in manifest.roots:
seeds.add((table_name, self._collapse.get((table_name, block_sid), block_sid)))
return sorted(seeds)
def _promote_late_roles(self) -> None:
"""Mark roots whose worker row had already been flushed before promotion."""
by_table: dict[str, set[int]] = {}
for manifest in self._manifests:
for table_name, block_sid in manifest.late_role_roots:
by_table.setdefault(table_name, set()).add(self._collapse.get((table_name, block_sid), block_sid))
for table_name, sids in by_table.items():
table = self._store._table(table_name)
ordered = sorted(sids)
for start in range(0, len(ordered), _SQLITE_COPY_BATCH_SIZE):
batch = ordered[start : start + _SQLITE_COPY_BATCH_SIZE]
self._connection.execute(
sqlalchemy.update(table)
.where(table.c[SID_COLUMN].in_(batch), table.c[ROLE_COLUMN] == 0)
.values({ROLE_COLUMN: 1})
)
def _reachable_table(self) -> sqlalchemy.Table:
reach = sqlalchemy.Table(
"_httk_bulk_reach",
sqlalchemy.MetaData(),
sqlalchemy.Column("tbl", sqlalchemy.Text, nullable=False),
sqlalchemy.Column(SID_COLUMN, sqlalchemy.Integer, nullable=False),
)
self._connection.execute(sqlalchemy.schema.DropTable(reach, if_exists=True))
self._connection.execute(sqlalchemy.schema.CreateTable(reach))
self._ingest._staging_tables.add(reach.name)
return reach
def _insert_reach(self, reach: sqlalchemy.Table, seeds: list[tuple[str, int]]) -> None:
if not seeds:
return
self._connection.execute(
sqlalchemy.insert(reach), [{"tbl": table_name, SID_COLUMN: sid} for table_name, sid in seeds]
)
def _insert_role_roots(self, reach: sqlalchemy.Table) -> None:
"""Seed untracked reachability directly from top-level row roles."""
for table_name in self._ingest._parent_schema:
table = self._store._table(table_name)
source = sqlalchemy.select(sqlalchemy.literal(table_name).label("tbl"), table.c[SID_COLUMN]).where(
table.c[ROLE_COLUMN] == 1
)
self._connection.execute(sqlalchemy.insert(reach).from_select(["tbl", SID_COLUMN], source))
def _close_reachability(self, reach: sqlalchemy.Table) -> None:
store = self._store
# Forward edges: a reached row keeps every record it references, directly
# (reference columns) or through its child rows (child-element columns).
reference_edges: list[tuple[sqlalchemy.Table, str, str]] = []
child_edges: list[tuple[sqlalchemy.Table, str, str, str]] = []
ownership_columns = {
edge.target_table: edge.target_column for edge in self._graph.ownership() if edge.target_column
}
for edge in self._graph.edges:
if edge.kind == "reference" and edge.source_column is not None:
reference_edges.append((store._table(edge.source_table), edge.source_column, edge.target_table))
elif edge.kind == "child_element" and edge.source_column is not None:
parent_column = ownership_columns.get(edge.source_table)
if parent_column is not None:
child_edges.append(
(store._table(edge.source_table), parent_column, edge.source_column, edge.target_table)
)
while True:
before = self._connection.execute(
sqlalchemy.select(sqlalchemy.func.count()).select_from(reach)
).scalar_one()
for table, column, ref_table in reference_edges:
self._grow_reach(reach, table, table.name, column, ref_table)
for child, parent_column, element_column, ref_table in child_edges:
self._grow_reach_via_child(reach, child, parent_column, element_column, ref_table)
after = self._connection.execute(sqlalchemy.select(sqlalchemy.func.count()).select_from(reach)).scalar_one()
if after == before:
return
def _grow_reach(
self, reach: sqlalchemy.Table, table: sqlalchemy.Table, table_name: str, column: str, ref_table: str
) -> None:
already = sqlalchemy.select(reach.c[SID_COLUMN]).where(reach.c.tbl == ref_table)
source = (
sqlalchemy.select(sqlalchemy.literal(ref_table).label("tbl"), table.c[column].label(SID_COLUMN))
.join_from(
table, reach, sqlalchemy.and_(reach.c.tbl == table_name, reach.c[SID_COLUMN] == table.c[SID_COLUMN])
)
.where(table.c[column].is_not(None))
.where(table.c[column].not_in(already))
.distinct()
)
self._connection.execute(sqlalchemy.insert(reach).from_select(["tbl", SID_COLUMN], source))
def _grow_reach_via_child(
self,
reach: sqlalchemy.Table,
child: sqlalchemy.Table,
parent_column: str,
element_column: str,
ref_table: str,
) -> None:
parent_table = parent_column[: -(len(SID_COLUMN) + 1)]
already = sqlalchemy.select(reach.c[SID_COLUMN]).where(reach.c.tbl == ref_table)
source = (
sqlalchemy.select(sqlalchemy.literal(ref_table).label("tbl"), child.c[element_column].label(SID_COLUMN))
.join_from(
child,
reach,
sqlalchemy.and_(reach.c.tbl == parent_table, reach.c[SID_COLUMN] == child.c[parent_column]),
)
.where(child.c[element_column].is_not(None))
.where(child.c[element_column].not_in(already))
.distinct()
)
self._connection.execute(sqlalchemy.insert(reach).from_select(["tbl", SID_COLUMN], source))
def _delete_unreached(self, reach: sqlalchemy.Table) -> None:
store = self._store
ownership_by_parent: dict[str, list[tuple[str, str]]] = {}
for edge in self._graph.ownership():
if edge.target_column is not None:
ownership_by_parent.setdefault(edge.source_table, []).append((edge.target_table, edge.target_column))
for table_name in self._ingest._parent_schema:
table = store._table(table_name)
reached = sqlalchemy.select(reach.c[SID_COLUMN]).where(reach.c.tbl == table_name)
self._connection.execute(sqlalchemy.delete(table).where(table.c[SID_COLUMN].not_in(reached)))
for child_name, parent_column in ownership_by_parent.get(table_name, ()):
child = store._table(child_name)
surviving_parents = sqlalchemy.select(table.c[SID_COLUMN])
self._connection.execute(
sqlalchemy.delete(child).where(child.c[parent_column].not_in(surviving_parents))
)
# -- compaction
def _compact(self, table: sqlalchemy.Table) -> None:
name = table.name
start = self._ingest._initial_next_sid.get(name, 1)
row_number = sqlalchemy.func.row_number().over(order_by=table.c[SID_COLUMN])
statement = sqlalchemy.select(
table.c[SID_COLUMN].label("old"),
(row_number + (start - 1)).label("new"),
).where(table.c[SID_COLUMN] >= _SID_BLOCK)
pairs = [(int(row[0]), int(row[1])) for row in self._connection.execute(statement).all()]
surviving = self._connection.execute(sqlalchemy.select(sqlalchemy.func.count()).select_from(table)).scalar_one()
self._ingest._inserted_count[name] = int(surviving)
self._ingest._next_sid[name] = start + len(pairs)
if not pairs:
return
for old, new in pairs:
self._compaction[(name, old)] = new
map_table = self._make_map_table(pairs)
try:
for referrer_table, column in self._referrers.get(name, ()):
self._remap_column(referrer_table, column, map_table)
self._remap_column(name, SID_COLUMN, map_table)
# A bulk row's logical_id starts equal to its own block sid, so the
# same renumbering keeps the invariant logical_id == final sid. The
# map only contains block sids (>= _SID_BLOCK); pre-existing rows,
# whose logical_id may differ from their sid after replace(), carry
# ids below that floor and are never matched, hence never disturbed.
self._remap_column(name, LOGICAL_ID_COLUMN, map_table)
# alt_id likewise starts equal to the block sid for a bulk main, so
# the same remap preserves the invariant alt_id == final sid.
self._remap_column(name, ALT_ID_COLUMN, map_table)
self._mint_compacted_entry_ids(table)
finally:
self._drop_map_table(map_table)
def _mint_compacted_entry_ids(self, table: sqlalchemy.Table) -> None:
"""Fill ids after block-sid compaction, when final numbers are known."""
record_type = next(
(record for record in self._store._entry_record_types if resolve_schema(record).table_name == table.name),
None,
)
if record_type is None:
return
scheme = self._store._entry_ids
has_missing_id = self._connection.execute(
sqlalchemy.select(table.c[SID_COLUMN])
.where(
table.c[SID_COLUMN] >= self._ingest._initial_next_sid.get(table.name, 1),
table.c.id.is_(None),
)
.limit(1)
).first()
if scheme is None and has_missing_id is not None:
raise ValueError(
f"{record_type.__name__} has no id and SqlStore(entry_ids=EntryIdScheme(...)) was not declared; "
"pass an explicit id or declare a scheme"
)
entry_id: sqlalchemy.ColumnElement[Any]
if scheme is None:
entry_id = table.c.id
else:
base = scheme.base
if scheme.type_in_base:
base = f"{base}.{self._store._entry_record_types[record_type][0]}"
prefix = f"{base}-{self._ingest._id_series or scheme.series}-"
_entry_type, backing_count, backing_index = self._store._entry_record_types[record_type]
number = table.c[SID_COLUMN] * backing_count + backing_index
generated = sqlalchemy.literal(prefix).op("||")(sqlalchemy.cast(number, sqlalchemy.Text))
entry_id = sqlalchemy.case((table.c.id.is_(None), generated), else_=table.c.id)
self._connection.execute(
sqlalchemy.update(table)
.where(table.c[SID_COLUMN] >= self._ingest._initial_next_sid.get(table.name, 1))
.values(
{
"id": entry_id,
"immutable_id": sqlalchemy.case(
(table.c.immutable_id.is_(None), entry_id.op("||")(sqlalchemy.literal("~1"))),
else_=table.c.immutable_id,
),
}
)
)
# -- dispatch
def _merge_dispatch(self) -> None:
ingest = self._ingest
if not ingest._track_sids:
active = set(ingest._created_set) | set(ingest._preexisting)
ingest._inserted_count.update(_rebuild_untracked_dispatch(self._connection, self._store, active))
return
for manifest in self._manifests:
for row in manifest.dispatch:
final_sid = self._final_sid(row.ref_table, row.block_sid)
built: dict[str, Any] = {DISPATCH_CONTENT_ID_COLUMN: row.key}
for column in row.all_columns:
built[column] = None
built[row.column] = final_sid
bucket = ingest._dispatch_rows.setdefault(row.dispatch_name, {})
ingest._dispatch_family.setdefault(row.dispatch_name, self._family_named(row.family_name))
existing = bucket.get(row.key)
if existing is not None:
if existing != built:
raise EntryDispatchIntegrityError(
f"entry dispatch {row.family_name!r} maps content_id {row.key!r} "
f"to a conflicting backing row"
)
continue
bucket[row.key] = built
ingest._flush_dispatch()
def _family_named(self, family_name: str) -> Any:
for family in self._store.layout.families:
if family.name == family_name:
return family
raise KeyError(family_name) # pragma: no cover - families are declared up front
def _final_sid(self, table_name: str, block_sid: int) -> int:
keep = self._collapse.get((table_name, block_sid), block_sid)
return self._compaction.get((table_name, keep), keep)
# -- sid map helpers
def _make_map_table(self, pairs: list[tuple[int, int]]) -> sqlalchemy.Table:
map_table = sqlalchemy.Table(
"_httk_bulk_sidmap",
sqlalchemy.MetaData(),
sqlalchemy.Column("old", sqlalchemy.Integer, nullable=False),
sqlalchemy.Column("new", sqlalchemy.Integer, nullable=False),
)
self._connection.execute(sqlalchemy.schema.DropTable(map_table, if_exists=True))
self._connection.execute(sqlalchemy.schema.CreateTable(map_table))
self._ingest._staging_tables.add(map_table.name)
self._connection.execute(sqlalchemy.insert(map_table), [{"old": old, "new": new} for old, new in pairs])
return map_table
def _drop_map_table(self, map_table: sqlalchemy.Table) -> None:
self._connection.execute(sqlalchemy.schema.DropTable(map_table, if_exists=True))
def _remap_column(self, table_name: str, column: str, map_table: sqlalchemy.Table) -> None:
# A join-based ``UPDATE ... FROM`` (a single hash join), not a per-row
# correlated subquery: the latter is quadratic on DuckDB and dominates
# the whole merge on a real-scale build.
table = self._store._table(table_name)
self._connection.execute(
sqlalchemy.update(table).where(table.c[column] == map_table.c.old).values({column: map_table.c.new})
)
def _delete_where_in_map(self, table_name: str, column: str, map_table: sqlalchemy.Table) -> None:
table = self._store._table(table_name)
member = sqlalchemy.select(map_table.c.old)
self._connection.execute(sqlalchemy.delete(table).where(table.c[column].in_(member)))