Layer 2: Storage

class SqlJournalStore(BaseJournalStore[SqlJournalWriter]):
    """SQL journal – an append-only heap per dataset in `JOURNAL_SCHEMA`.

    A flush renames the table to a timestamped segment and creates a fresh one
    in the same DDL transaction, streams the segment out as Arrow and drops it
    – cleanup is a catalog operation, never a ``DELETE``. Dialects:
    `SqliteJournalStore` / `PostgresJournalStore`, picked by `sql_journal`.
    """

    _writer_cls = SqlJournalWriter

    lock_timeout: str | None = None
    """Dialect bound on how long the rotation waits for in-flight writers."""

    def __init__(self, dataset: str, uri: str | None = None) -> None:
        super().__init__(dataset, uri)
        self.engine = self.make_engine()
        self.metadata = MetaData()
        self.table = journal_table(self.metadata, f"journal_{dataset}")
        self.metadata.create_all(self.engine, tables=[self.table], checkfirst=True)

    # -- dialect hooks

    def make_engine(self) -> Engine:
        # no idle connections: `get_journal` caches a store per dataset forever
        return create_engine(self.uri, hide_parameters=True, poolclass=NullPool)

    def connect(self) -> Any:
        """Open a connection for a writer's inserts."""
        raise NotImplementedError

    def acquire(self) -> Any:
        """Take a connection for a writer's inserts.

        A plain [`connect`][SqlJournalStore.connect]; `PostgresJournalStore`
        borrows from an ADBC pool instead.
        """
        return self.connect()

    def release(self, conn: Any) -> None:
        """Hand a writer's connection back by closing it.

        On postgres that checks it into the pool, rolled back, so the next
        writer never inherits an aborted transaction.
        """
        conn.close()

    def insert_batch(self, conn: Any, batch: pa.Table) -> None:
        """Append one packed batch to the live table."""
        raise NotImplementedError

    def read_segment(self, name: str) -> RecordBatches:
        """Stream a segment's rows."""
        raise NotImplementedError

    @contextmanager
    def flush_lock(self) -> Generator[bool, None, None]:
        """Hold this dataset's flush window; yields ``False`` if another flush has it.

        Without it a second flush would drain the first one's segment. The lock
        must release itself when its holder dies, or a crash strands a segment.
        """
        raise NotImplementedError
        yield True  # pragma: no cover - typing

    def _set_lock_timeout(self, conn: Any) -> None:
        if self.lock_timeout is not None:
            conn.exec_driver_sql(f"SET LOCAL lock_timeout = '{self.lock_timeout}'")

    # -- segments

    @property
    def _prefix(self) -> str:
        return f"{self.table.name}{SEGMENT_INFIX}"

    def _segment_name(self) -> str:
        """A fresh segment name – time-ordered, unique against a racing flush."""
        return f"{self._prefix}{utc_now().strftime('%Y%m%dT%H%M%S')}{uuid4().hex[:4]}"

    def _segments(self) -> list[str]:
        """Rotated segments, oldest first – orphans of a crashed flush included."""
        names = inspect(self.engine).get_table_names()
        return sorted(n for n in names if n.startswith(self._prefix))

    def _table_names(self) -> list[str]:
        """The live table plus every un-dropped segment."""
        return [self.table.name, *self._segments()]

    def _table(self, name: str) -> Table:
        return journal_table(MetaData(), name)

    def _rotate(self) -> None:
        """Rename the journal to a segment and recreate it in one DDL transaction.

        The rename's exclusive lock waits out in-flight inserts, so no row lands
        in the segment afterwards; blocked writers continue into the fresh table.
        `lock_timeout` bounds the wait – every new insert queues behind it – so a
        flush blocked by a long transaction fails, and the next flush retries.
        """
        name = self._segment_name()
        with self.engine.begin() as conn:
            self._set_lock_timeout(conn)
            conn.exec_driver_sql(f'ALTER TABLE "{self.table.name}" RENAME TO "{name}"')
            conn.execute(CreateTable(self.table))

    def _drop(self, name: str) -> None:
        with self.engine.begin() as conn:
            conn.exec_driver_sql(f'DROP TABLE IF EXISTS "{name}"')

    def _has_rows(self, name: str) -> bool:
        with self.engine.connect() as conn:
            res = conn.exec_driver_sql(f'SELECT 1 FROM "{name}" LIMIT 1')
            return res.first() is not None

    # -- flush

    def flush_batches(self) -> StatementTables:
        """Rotate the journal, then stream each segment as Arrow tables and drop it.

        Held under [`flush_lock`][SqlJournalStore.flush_lock] – a concurrent
        flush yields nothing. Segments left by a crashed flush are drained too,
        and a consumer that fails keeps the undrained rows for the next flush.
        Rows stream unordered, so a table spans shards and
        [`append`][ftm_lakehouse.storage.parquet.ParquetStore.append] writes one
        file per partition it touches.
        """
        with self.flush_lock() as acquired:
            if not acquired:
                log.warning(
                    "Another flush is draining this journal – skipping",
                    journal=self.table.name,
                )
                return
            if self._has_rows(self.table.name):
                self._rotate()
            for name in self._segments():
                yield from self._drain(name)

    def _drain(self, name: str) -> StatementTables:
        """Stream one segment in whole tables, then drop it.

        The consumer writes each table before asking for the next, so the drop
        only follows durable writes; a failure keeps the segment.
        """
        pending: list[pa.RecordBatch] = []
        rows = 0
        for chunk in self.read_segment(name):
            pending.append(chunk)
            rows += chunk.num_rows
            if rows >= settings.journal_drain_rows:
                yield pa.Table.from_batches(pending, schema=JOURNAL_SCHEMA)
                pending, rows = [], 0
        if pending:
            yield pa.Table.from_batches(pending, schema=JOURNAL_SCHEMA)
        # only after the reader is closed – its read transaction blocks the DROP
        self._drop(name)

    # -- reads

    def iterate_entity(self, entity_id: str) -> LakehouseStatements:
        """Iterate the live statements of one entity, across all segments.

        A full scan per call – the journal has no index (see
        [`journal_table`][ftm_lakehouse.model.statement.journal_table]).
        """
        with self.engine.connect() as conn:
            for name in self._table_names():
                table = self._table(name)
                q = (
                    select(table)
                    .where(table.c.entity_id == entity_id)
                    .where(table.c.deleted_at.is_(None))
                )
                for row in conn.execute(q):
                    yield _row_to_statement(row)

    def count(self) -> int:
        """Count rows for this dataset, across all segments."""
        total = 0
        with self.engine.connect() as conn:
            for name in self._table_names():
                res = conn.exec_driver_sql(f'SELECT count(*) FROM "{name}"').scalar()
                total += res or 0
        return total

    def clear(self) -> int:
        """Delete all rows for this dataset. Returns count of deleted rows."""
        count = self.count()
        with self.engine.begin() as conn:
            for name in self._segments():
                conn.exec_driver_sql(f'DROP TABLE IF EXISTS "{name}"')
            conn.execute(delete(self.table))
        return count

    def dispose(self) -> None:
        """Dispose the engine and close all pooled connections."""
        self.engine.dispose()

class ParquetStore:
    """The statement store of one dataset: one Delta table, partitioned by
    ``(shard, bucket, origin)``, at ``{uri}/statements`` – or at
    ``{LAKEHOUSE_STATEMENTS_URI}/{dataset}/statements`` when that is set. Locks
    and tags stay with the dataset at ``uri``."""

    def __init__(
        self,
        uri: Uri,
        dataset: str,
        shards: int | None = None,
    ) -> None:
        self.settings = Settings()
        root = uri
        if self.settings.statements_uri:
            root = join_uri(
                ensure_uri(self.settings.statements_uri),
                validate_dataset_name(dataset),
            )
        self.uri = join_uri(root, path.STATEMENTS)
        self.dataset = dataset
        self.shards = shards if shards is not None else DEFAULT_SHARDS
        self._store = get_store(uri)
        self._tags = TagStore(uri)
        self._snapshot_lock = RLock()
        self._snapshot: DeltaTable | None = None
        self._files: tuple[int, str, Partitions] | None = None
        self._lake = _LakeStore(
            uri=str(self.uri),
            dataset=self.dataset,
            partition_by=PARTITIONS,
            view_sqls={TABLE_RAW.name: raw_view_sql},
            duckdb_config=duckdb_config(),
            exists=lambda: self.exists,
        )
        self.log = get_logger(
            f"{self.dataset}.{self.__class__.__name__}",
            dataset=self.dataset,
            uri=mask_uri(self.uri),
        )

    @property
    def deltatable(self) -> DeltaTable:
        """A freshly loaded handle on the table, for callers outside the store."""
        return self._lake.deltatable

    @property
    def num_rows(self) -> int:
        """Physical row count from the snapshot's file statistics – no scan."""
        with self._snapshot_lock:
            snapshot = self._current_snapshot()
            return snapshot.count() if snapshot is not None else 0

    @property
    def version(self) -> int | None:
        """Current version of the main Delta table."""
        with self._snapshot_lock:
            snapshot = self._current_snapshot()
            return snapshot.version() if snapshot is not None else None

    @property
    def exists(self) -> bool:
        """Whether the Delta table exists."""
        with self._snapshot_lock:
            return self._current_snapshot() is not None

    def _current_snapshot(self) -> DeltaTable | None:
        """This process's Delta snapshot, advanced with ``update_incremental``.
        Callers hold `_snapshot_lock`.
        """
        if self._snapshot is not None:
            try:
                self._snapshot.update_incremental()
                return self._snapshot
            except DeltaError:
                self._snapshot = None
        try:
            self._snapshot = DeltaTable(
                str(self.uri), storage_options=storage_options()
            )
        except TableNotFoundError:
            return None
        return self._snapshot

    def _existing_snapshot(self) -> DeltaTable:
        """`_current_snapshot` of a table that must exist. Callers hold
        `_snapshot_lock`."""
        snapshot = self._current_snapshot()
        if snapshot is None:
            raise RuntimeError(f"Statement store vanished: `{self.uri}`")
        return snapshot

    def _partitions(self) -> tuple[str, Partitions]:
        """The table root and the snapshot's files per partition, cached per
        version; empty without a table."""
        with self._snapshot_lock:
            snapshot = self._current_snapshot()
            if snapshot is None:
                return "", {}
            version = snapshot.version()
            if self._files is None or self._files[0] != version:
                actions = pa.table(snapshot.get_add_actions(flatten=True))
                files: dict[Partition, Files] = {}
                for file, size, *partition in zip(
                    actions["path"].to_pylist(),
                    actions["size_bytes"].to_pylist(),
                    actions["partition.shard"].to_pylist(),
                    actions["partition.bucket"].to_pylist(),
                    actions["partition.origin"].to_pylist(),
                ):
                    key = cast(Partition, tuple(partition))
                    files.setdefault(key, []).append((unquote(file), size))
                partitions = {
                    key: (
                        sorted(fs),  # a merge's files sort in entity_id order
                        all(
                            posixpath.basename(f).startswith(MERGED_PREFIX)
                            for f, _ in fs
                        ),
                    )
                    for key, fs in sorted(files.items())
                }
                root = snapshot.table_uri.rstrip("/")
                self._files = (version, root, partitions)
            return self._files[1], self._files[2]

    def _pairs(self) -> tuple[str, Pairs]:
        """`_partitions`, grouped per ``(shard, bucket)`` pair."""
        root, partitions = self._partitions()
        pairs: Pairs = {}
        for partition, (files, clean) in partitions.items():
            pairs.setdefault(partition[:2], []).append((partition, files, clean))
        return root, pairs

    @property
    def source(self) -> SqlSource:
        return make_source(self.shards)

    def _statement_data(self, q: Query | None = None) -> Iterator[StatementDict]:
        """Statement dicts for ``q``, entity-contiguous – per ``(shard, bucket)``
        pair, or over all pairs at once when sorted or sliced."""
        sql = (q or Query()).compile(self.source)
        prune = self._prune_values(q, self.source)
        whole = q is not None and (q.sort is not None or q.slice is not None)
        sources = self._scoped_sources(prune, whole)
        yield from cast(Iterator[StatementDict], self._rows(sql, sources))

    def query(self, q: Query | None = None) -> StatementEntities:
        """Query entities from the store.

        Args:
            q: Filters, plus ordering / slicing.

        Yields:
            `StatementEntity` objects.
        """
        for data in aggregate_statement_payloads(self._statement_data(q), self.dataset):
            yield data.to_entity()

    def query_statements(self, q: Query | None = None) -> Statements:
        """Query statements from the store.

        Args:
            q: Filters, plus ordering / slicing.

        Yields:
            `LakehouseStatement` objects – with their ``fragment`` and ``role``,
            so one can be handed to
            [`delete_statement`][ftm_lakehouse.repository.EntityRepository.delete_statement].
        """
        for stmt_dict in self._statement_data(q):
            yield LakehouseStatement.from_dict(stmt_dict)

    def _lock(self, key: Uri) -> Lock:
        """A store lock file, waited for at most ``lock_max_retries`` times."""
        return Lock(self._store, key=key, max_retries=self.settings.lock_max_retries)

    def _fence_retry(self, attempt: Callable[[], T]) -> T:
        """Retry ``attempt`` until it stops raising – ``lock_max_retries``
        attempts, ``N²/2`` seconds in total, then the error propagates."""
        return cast(
            T,
            error_handler(
                max_retries=self.settings.lock_max_retries,
                backoff_factor=1,
                do_raise=True,
            )(attempt)(),
        )

    def merge_lock(self) -> Lock:
        """The ``.LOCK-MERGE`` – held by [`merge`][ParquetStore.merge], the export
        sweep and, with ``.LOCK``, the exclusive maintenance. Appends don't wait
        for it; it keeps a vacuum off files a running sweep reads."""
        return self._lock(path.LOCK_MERGE)

    def _await_unlocked(self) -> None:
        """Back off while ``.LOCK`` is held. An append that passed this may still be
        in flight when maintenance takes the lock: `delete_origin` retries its
        commit, `shard` runs with writers stopped."""

        def check() -> None:
            if self._store.exists(path.LOCK):
                raise RuntimeError(
                    f"Write fence busy: maintenance lock `{path.LOCK}` is "
                    "held. If a writer crashed, release the fence via "
                    "`ftm-lakehouse maintenance unlock`."
                )

        self._fence_retry(check)

    def _ensure_table(self) -> None:
        """Create the table as an empty commit, under ``.LOCK`` so two first
        imports cannot both commit version ``0``."""
        if self.exists:
            return
        with self._lock(path.LOCK):
            if self.exists:  # lost the create race - the table is there now
                return
            write_deltalake(
                str(self.uri),
                pa.Table.from_pylist([], schema=SHARDED_SCHEMA),
                partition_by=PARTITIONS,
                mode="overwrite",
                configuration=TABLE_CONFIGURATION,
                storage_options=storage_options(),
            )

    @contextmanager
    def _maintenance_fence(self) -> Iterator[None]:
        """``.LOCK`` plus ``.LOCK-MERGE``, for maintenance that rewrites or drops
        files in place."""
        with self._lock(path.LOCK), self.merge_lock():
            yield

    def unlock(self) -> bool:
        """Forcibly release ``.LOCK`` and ``.LOCK-MERGE`` after a crashed writer.

        Breaking a lock a live writer holds can corrupt its write – confirm no
        process is writing first.

        Returns:
            Whether a lock was released.
        """
        released = False
        for key in (path.LOCK, path.LOCK_MERGE):
            if self._store.exists(key):
                self._store.delete(key)
                released = True
        return released

    def evolve_schema(self) -> list[str]:
        """Add the `SHARDED_SCHEMA` columns the table was created without.

        A metadata-only commit; older files read the new columns as NULL.
        Additive only – a removed column needs a full rewrite.

        Returns:
            The columns added.
        """
        with self._snapshot_lock:
            snapshot = self._current_snapshot()
            if snapshot is None:
                return []
            known = {f.name for f in snapshot.schema().fields}
            missing = [
                f
                for f in Schema.from_arrow(SHARDED_SCHEMA).fields
                if f.name not in known
            ]
            if not missing:
                return []
            names = [f.name for f in missing]
            with self._maintenance_fence():
                snapshot.alter.add_columns(missing)
        self.log.info("Evolved parquet schema.", columns=names)
        return names

    def configure_table(self) -> dict[str, str]:
        """Apply `TABLE_CONFIGURATION`, then checkpoint and clean up the log under
        it.

        Returns:
            The properties that changed.
        """
        with self._snapshot_lock:
            snapshot = self._current_snapshot()
            if snapshot is None:
                return {}
            current = snapshot.metadata().configuration
            changed = {
                k: v for k, v in TABLE_CONFIGURATION.items() if current.get(k) != v
            }
            if not changed:
                return {}
            with self._maintenance_fence(), Took() as t:
                snapshot.alter.set_table_properties(changed)
                snapshot.update_incremental()
                snapshot.create_checkpoint()
                snapshot.cleanup_metadata()
        self.log.info("Configured delta table.", took=t.took, **changed)
        return changed

    def _with_shard(self, batch: pa.Table) -> pa.Table:
        """Derive ``shard`` from ``entity_id`` – the only place a row's partition is
        decided."""
        ids = pc.dictionary_encode(batch.column("entity_id").combine_chunks())
        shards = pa.array(
            [entity_shard(e, self.shards) for e in ids.dictionary.to_pylist()],
            pa.string(),
        )
        return batch.append_column(
            SHARDED_SCHEMA.field("shard"), pc.take(shards, ids.indices)
        ).select(SHARDED_SCHEMA.names)

    def append(self, batch: pa.Table) -> None:
        """Append `JOURNAL_SCHEMA` rows, deriving their ``shard``.

        One write per bucket, for its writer profile. Unsorted and lock-free –
        concurrent appends commit through Delta's optimistic concurrency; only
        ``.LOCK`` makes them wait. The touched partitions are dirty until the
        next [`merge`][ParquetStore.merge].
        """
        if len(batch) == 0:
            return

        batch = self._with_shard(batch)
        buckets = pc.unique(batch["bucket"]).to_pylist()
        shards = pc.unique(batch["shard"]).to_pylist()
        self.log.info(
            f"Flushing {len(batch)} statements to parquet ...",
            buckets=buckets,
            shards=shards,
        )
        self._ensure_table()
        for bucket in buckets:
            self._write(bucket, batch.filter(pc.equal(batch["bucket"], bucket)))

    def _write(
        self, bucket: str, table: pa.Table, target_file_size: int | None = None
    ) -> None:
        """Write a bucket from pyarrow Table to deltalake, exclusively via .LOCK"""
        # delta-rs concatenates the batches it is handed, and one entity's text
        # can pass the 2GiB 32-bit string offsets hold
        table = table.cast(
            pa.schema(
                f.with_type(pa.large_string()) if pa.types.is_string(f.type) else f
                for f in table.schema
            )
        )
        self._await_unlocked()
        with self._snapshot_lock:
            write_deltalake(
                self._existing_snapshot(),
                table,
                partition_by=PARTITIONS,
                mode="append",
                writer_properties=writer_for_bucket(bucket),
                target_file_size=target_file_size,
            )
        # stamped once committed: an entry-time stamp would predate an export
        # that started meanwhile and missed these rows, leaving it fresh
        self._tags.set(tag.STATEMENTS_UPDATED)

    def _by_partition(self, batch: pa.Table) -> Iterator[tuple[Partition, pa.Table]]:
        """`JOURNAL_SCHEMA` rows as one table per ``(shard, bucket, origin)``,
        ``shard`` derived."""
        batch = self._with_shard(batch)
        order = pc.sort_indices(batch, sort_keys=[(c, "ascending") for c in PARTITIONS])
        keys = pc.binary_join_element_wise(*(batch[c] for c in PARTITIONS), "/")
        runs = pc.value_counts(pc.take(keys, order))  # first seen first: sorted
        offset = 0
        for key, count in zip(
            runs.field("values").to_pylist(), runs.field("counts").to_pylist()
        ):
            shard, bucket, origin = key.split("/")
            # a copy per part – a slice would keep the whole batch alive
            yield (shard, bucket, origin), batch.take(order.slice(offset, count))
            offset += count

    def write_stream(
        self, tables: Iterable[pa.Table], budget: int | None = None
    ) -> int:
        """Append a stream of `JOURNAL_SCHEMA` tables as a ``part-*`` file per
        ``(shard, bucket, origin)``. This is not the journal drain but used for
        bulk imports via cli or scripts.

        A partition's rows are held until they fill a file (`MERGE_CHUNK_BYTES`)
        or the held rows exceed ``budget``; then the largest partitions are
        written, down to half the budget, in one commit per bucket of up to
        `WRITE_FLUSH_PARTITIONS` partitions.

        Args:
            tables: Stream of packed statement tables.
            budget: Arrow bytes to hold, default ``LAKEHOUSE_WRITE_BUFFER``.

        Returns:
            Number of rows written.
        """
        if budget is None:
            budget = parse_byte_size(self.settings.write_buffer)
        held: dict[Partition, list[pa.Table]] = {}
        sizes: dict[Partition, int] = {}
        rows = files = commits = 0

        def flush(keys: list[Partition]) -> None:
            nonlocal files, commits
            for bucket in sorted({key[1] for key in keys}):
                of_bucket = [key for key in keys if key[1] == bucket]
                for i in range(0, len(of_bucket), WRITE_FLUSH_PARTITIONS):
                    chunk = of_bucket[i : i + WRITE_FLUSH_PARTITIONS]
                    table = pa.concat_tables(
                        chain.from_iterable(held.pop(key) for key in chunk)
                    )
                    for key in chunk:
                        del sizes[key]
                    self.log.debug(
                        "Writing partitions ...",
                        bucket=bucket,
                        partitions=len(chunk),
                        rows=len(table),
                    )
                    self._write(bucket, table, target_file_size=1 << 62)  # a file each
                    files += len(chunk)
                    commits += 1

        self._ensure_table()
        for batch in tables:
            for key, part in self._by_partition(batch):
                held.setdefault(key, []).append(part)
                sizes[key] = sizes.get(key, 0) + part.nbytes
                rows += part.num_rows
                if sizes[key] >= MERGE_CHUNK_BYTES:
                    flush([key])
            if (left := sum(sizes.values())) > budget:
                largest: list[Partition] = []
                for key in sorted(sizes, key=sizes.__getitem__, reverse=True):
                    largest.append(key)
                    if (left := left - sizes[key]) <= budget // 2:
                        break
                flush(largest)
        flush(list(held))
        if commits:
            with self._snapshot_lock:
                self._existing_snapshot().create_checkpoint()
        self.log.info("Wrote statements.", rows=rows, files=files, commits=commits)
        return rows

    @property
    def needs_merge(self) -> bool:
        """Whether any partition holds a file [`merge`][ParquetStore.merge] did not
        write."""
        return not all(clean for _, clean in self._partitions()[1].values())

    def merge(self, force: bool = False) -> None:
        """Rewrite every dirty partition into canonical files, in ``entity_id``
        order – one per `MERGE_CHUNK_BYTES` of rows.

        Duplicates collapse, fragments supersede, ``first_seen`` folds and
        tombstones past ``LAKEHOUSE_GRACE_PERIOD_DAYS`` go with the rows they
        shadow ([`build_merge_sql`][ftm_lakehouse.logic.parquet.build_merge_sql]).
        Reads see the same rows before and after (``dataChange = false``).

        Partitions merge in ``LAKEHOUSE_WORKERS`` processes, committed in
        batches of `MERGE_COMMIT_BATCH`; a failed run still commits what
        finished, including the partitions still running when one failed.
        Holds only [`merge_lock`][ParquetStore.merge_lock], so appends
        keep flowing.

        Args:
            force: Rewrite every partition from scratch, clean ones included,
                re-sorting the last merge's files instead of trusting them. With
                a grace period of ``0`` it purges tombstones in otherwise idle
                partitions.
        """
        if not self.exists:
            return
        grace_cutoff = utc_now() - timedelta(days=self.settings.grace_period_days)
        workers = self.settings.workers
        merged = skipped = 0
        with self.merge_lock():
            root, partitions = self._partitions()
            tasks = [
                MergeTask(
                    partition,
                    files,
                    root,
                    grace_cutoff,
                    worker_duckdb_config(workers),
                    presorted=not force,
                )
                for partition, (files, clean) in partitions.items()
                if force or not clean
            ]
            skipped = len(partitions) - len(tasks)
            # advanced per partition, throughput in bytes read
            with (
                SyncProgressBar("Merging partitions", len(tasks)) as bar,
                process_map(workers, drain=True) as run,
            ):
                done: list[MergeResult] = []
                try:
                    for result in run(merge_partition, tasks):
                        bar.advance(size=sum(size for _, size in result.task.files))
                        shard, bucket, origin = result.task.partition
                        self.log.info(
                            f"Merged partition `{shard}/{bucket}/{origin}`.",
                            took=result.took,
                            shard=shard,
                            bucket=bucket,
                            origin=origin,
                            grace_period_days=self.settings.grace_period_days,
                        )
                        done.append(result)
                        if len(done) == MERGE_COMMIT_BATCH:
                            batch, done = done, []
                            self._commit_merged(batch)
                            merged += len(batch)
                finally:
                    # also on failure: commit what finished, so the next run resumes
                    if done:
                        self._commit_merged(done)
                        merged += len(done)
            if merged:
                with self._snapshot_lock, Took() as t:
                    self._existing_snapshot().create_checkpoint()
                self.log.info("Wrote checkpoint.", took=t.took)
        self.log.info(
            "Merge complete.",
            merged=merged,
            skipped=skipped,
            workers=workers,
            grace_period_days=self.settings.grace_period_days,
        )

    def _commit_merged(self, batch: list[MergeResult]) -> None:
        """Commit a batch of merged partitions as one ``dataChange = false``
        transaction."""
        now = int(utc_now().timestamp() * 1000)
        actions: list[AddAction | RemoveAction] = []
        for result in batch:
            task = result.task
            values: dict[str, str | None] = dict(zip(PARTITIONS, task.partition))
            for file, size, stats in result.files:
                actions.append(AddAction(file, size, values, now, False, stats))
            for file, size in task.files:
                actions.append(RemoveAction(file, False, now, size, values))
        with self._snapshot_lock, Took() as t:
            snapshot = self._existing_snapshot()
            snapshot.create_write_transaction(
                actions,
                mode="append",
                schema=snapshot.schema(),
                partition_by=PARTITIONS,
            )
            snapshot.update_incremental()
        self.log.info(
            "Committed merged partitions.", partitions=len(batch), took=t.took
        )

    def shard(self, shards: int) -> None:
        """Rewrite the store onto ``shards`` entity-hash shards.

        One streamed, atomic overwrite per ``(bucket, origin)`` group; every
        partition comes out dirty, so run ``optimize`` afterwards. Idempotent.
        Blocks parquet appends but not the journal – run it with writers stopped.

        Args:
            shards: Target shard count; ``<= 1`` means a single shard.
        """
        if self.exists:
            self._rewrite_shards(shards)
        self.shards = shards
        self.log.info("Re-shard complete.", shards=shards)

    def _rewrite_shards(self, shards: int) -> None:
        """Rewrite every ``(bucket, origin)`` group, under the maintenance fence,
        read in batches of `MERGE_BATCH_BYTES` (`_batch_rows`) – DuckDB's default
        of a million rows a batch overflows Arrow's 2GiB string buffer on
        document rows."""
        with self._maintenance_fence():
            root, partitions = self._partitions()
            groups: dict[tuple[str, str], list[tuple[Partition, Files]]] = {}
            for partition, (files, _) in partitions.items():
                _, bucket, origin = partition
                groups.setdefault((bucket, origin), []).append((partition, files))
            config = duckdb_config()
            for (bucket, origin), sources in groups.items():
                with (
                    Took() as t,
                    _connect(config) as con,
                    self._snapshot_lock,
                ):
                    # a batch's strings may pass 2GiB when `_batch_rows`
                    # underestimates
                    con.execute("SET arrow_large_buffer_size = true")
                    snapshot = self._existing_snapshot()
                    sqls = [
                        build_shard_sql(
                            *partition, shards, source=_relation(root, partition, files)
                        )
                        for partition, files in sources
                    ]
                    rows = _batch_rows(con, _relation(root, *sources[0]))
                    # lazily: executing the next query closes the current result
                    readers = (con.execute(sql).to_arrow_reader(rows) for sql in sqls)
                    first = next(readers)
                    batches = chain.from_iterable(chain([first], readers))
                    write_deltalake(
                        snapshot,
                        pa.RecordBatchReader.from_batches(first.schema, batches),
                        mode="overwrite",
                        partition_by=PARTITIONS,
                        predicate=(
                            f"bucket = '{bucket}' AND "
                            f"origin = '{validate_origin(origin)}'"
                        ),
                        writer_properties=writer_for_bucket(bucket),
                        target_file_size=shard_target_file_size(shards),
                    )
                    snapshot.update_incremental()
                self.log.info(
                    f"Re-sharded `{bucket}/{origin}`.",
                    took=t.took,
                    bucket=bucket,
                    origin=origin,
                    sources=len(sources),
                    shards=shards,
                )

    def delete_origin(self, origin: str) -> int:
        """Physically drop every row of one origin – whole partitions, at once,
        with no grace period.

        Args:
            origin: The origin to drop.

        Returns:
            Number of rows removed.

        Raises:
            ValueError: If ``origin`` is not a safe origin name.
            RuntimeError: When the write fence cannot be acquired.
        """
        origin = validate_origin(origin)
        if not self.exists:
            return 0
        with self._maintenance_fence(), Took() as t, self._snapshot_lock:
            # safe to interpolate: `validate_origin` rejects quotes
            metrics: dict[str, Any] = self._fence_retry(
                lambda: self._existing_snapshot().delete(f"origin = '{origin}'")
            )
            deleted = int(metrics.get("num_deleted_rows") or 0)
            if deleted:
                self._tags.set(tag.STATEMENTS_UPDATED)
            self.log.info(
                "Dropped origin.",
                took=t.took,
                origin=origin,
                deleted=deleted,
                **metrics,
            )
        return deleted

    def destroy(self) -> None:
        """Soft-delete the store: drop its ``_delta_log``. The parquet files stay
        on disk until a [`vacuum`][ParquetStore.vacuum] of the re-created table.

        Other processes keep the snapshot they hold – restart them.
        """
        if not self.exists:
            return
        with self._maintenance_fence(), Took() as t, self._snapshot_lock:
            table = get_store(self.uri)
            table._fs.rm(table._keys.to_fs_key("_delta_log"), recursive=True)
            self._snapshot = self._files = None
            self._tags.set(tag.STATEMENTS_UPDATED)
        self.log.warning("Destroyed statement store.", took=t.took)

    def vacuum(self, retention_hours: int = 0) -> None:
        """Delete the files the Delta log no longer references, and the staging
        tables a crashed merge left (`STAGING_DIR`).

        Args:
            retention_hours: Keep such files newer than this.
        """
        if not self.exists:
            return
        with self._maintenance_fence(), Took() as t, self._snapshot_lock:
            deleted = self._existing_snapshot().vacuum(
                retention_hours=retention_hours,
                dry_run=False,
                enforce_retention_duration=False,
                full=True,
            )
            # no merge stages under the fence
            table = get_store(self.uri)
            if table.exists(STAGING_DIR):
                table._fs.rm(table._keys.to_fs_key(STAGING_DIR), recursive=True)
        self.log.info("Vacuumed.", files=len(deleted), took=t.took)

    def sweep_sources(self) -> list[SweepSource]:
        """Every ``(shard, bucket)`` pair of one snapshot as a ``SweepSource``, the
        units an export sweeps (no entity spans two)."""
        root, pairs = self._pairs()
        return [
            SweepSource(
                key,
                [
                    SweepPartition(
                        relation=_relation(root, partition, files),
                        clean=clean,
                        presorted=clean,
                        size=sum(size for _, size in files),
                    )
                    for partition, files, clean in pairs[key]
                ],
            )
            for key in sorted(pairs)
        ]

    def deleted_candidates(self, since: datetime) -> Iterator[DeleteCandidate]:
        """Every entity with a tombstone at or after ``since`` – one raw scan over
        the whole store."""
        sql = deleted_candidates_select(since)
        for row in self._rows(sql, self._scoped_sources()):
            yield DeleteCandidate(
                id=row["entity_id"],
                deleted_at=row["deleted_at"],
                origins=frozenset(row["origins"] or ()),
                schemata=frozenset(row["schemata"] or ()),
                content_hash=bool(row["content_hash"]),
            )

    @staticmethod
    def _prune_values(q: Query | None, source: SqlSource) -> dict[str, set[str]]:
        """The partition values ``q`` confines its matches to, per prune column –
        only for a flat positive conjunction."""
        if q is None or not Sql(q, source)._is_flat_and:
            return {}
        return {
            column: set(values)
            for column, prune in source.prune.items()
            if (values := prune(q))
        }

    @staticmethod
    def _pruned_keys(
        pairs: Pairs, prune: dict[str, set[str]] | None
    ) -> list[tuple[str, str]]:
        """The keys of ``pairs`` that ``prune`` leaves in."""
        prune = prune or {}
        return [
            (s, b)
            for s, b in sorted(pairs)
            if s in prune.get("shard", {s}) and b in prune.get("bucket", {b})
        ]

    def _scoped_sources(
        self, prune: dict[str, set[str]] | None = None, whole: bool = False
    ) -> Iterator[tuple[str, bool]]:
        """One ``(relation, clean)`` per ``(shard, bucket)`` pair ``prune`` leaves in.

        A read pruned to shards puts all its pairs into one relation – an id
        names its shard but not its bucket – and so does a ``whole`` one, which
        sorts or slices across pairs.
        """
        root, pairs = self._pairs()
        keys = self._pruned_keys(pairs, prune)
        if whole or (prune and "shard" in prune):
            if keys:
                yield pair_source(root, [s for key in keys for s in pairs[key]])
            return
        for key in keys:
            yield pair_source(root, pairs[key])

    @contextmanager
    def cursor(
        self, raw: str = TABLE_RAW.name, live: str = TABLE.name
    ) -> Iterator[duckdb.DuckDBPyConnection]:
        """A cursor with the whole table registered as a read's views
        (`register_partition`): the live view is a plain scan when every
        partition is clean, the dedupe query otherwise.

        Args:
            raw: Name of the view over every physical row.
            live: Name of the view over the live rows.

        Yields:
            A DuckDB cursor of the store's connection.
        """
        with self._snapshot_lock:
            source = delta_scan_sql(self._existing_snapshot().table_uri)
        with self._cursor_over(source, not self.needs_merge, raw, live) as cur:
            yield cur

    @contextmanager
    def _cursor_over(
        self,
        source: str,
        clean: bool,
        raw: str = TABLE_RAW.name,
        live: str = TABLE.name,
    ) -> Iterator[duckdb.DuckDBPyConnection]:
        """A cursor of the shared connection, its views reading ``source``
        (`register_partition`)."""
        with self._lake.cursor() as cur:
            register_partition(cur, source, clean, raw, live)
            yield cur

    def _rows(
        self, sql: Select, sources: Iterable[tuple[str, bool]]
    ) -> Iterator[dict[str, Any]]:
        """``sql``'s rows over each ``(relation, clean)`` source in turn
        (`_scoped_sources`), in Arrow batches of `SWEEP_BATCH_SIZE`."""
        compiled = _compile(sql)
        for source, clean in sources:
            with self._cursor_over(source, clean) as cur:
                for batch in cur.execute(compiled).to_arrow_reader(SWEEP_BATCH_SIZE):
                    yield from batch.to_pylist()