logic

def duckdb_config() -> dict[str, str]:
    """DuckDB config from the lakehouse settings – memory limit, spill and
    extension directories.

    Call once per DuckDB instance: each call gets its own spill subdirectory,
    as instances sharing one overwrite each other's spill files. Warns once
    when that directory is RAM-backed. No ``TimeZone`` – at connect time it
    installs ``icu`` (fails offline); sessions pin UTC after connecting.
    """
    settings = Settings()
    config: dict[str, str] = {"memory_limit": settings.duckdb_memory_limit}
    if settings.duckdb_temp_directory:
        # DuckDB creates the spill directory but not its parents
        os.makedirs(settings.duckdb_temp_directory, exist_ok=True)
        _warn_ram_spill(settings.duckdb_temp_directory)
        config["temp_directory"] = os.path.join(
            settings.duckdb_temp_directory, ensure_uuid()
        )
    if settings.duckdb_extension_directory:
        config["extension_directory"] = settings.duckdb_extension_directory
    return config

def raw_view_sql(dt: DeltaTable) -> str:
    """SELECT body for the ``statement_raw`` view: every physical row,
    tombstones and duplicates included."""
    return f"SELECT * FROM {delta_scan_sql(dt.table_uri)}"

def live_rows_sql(source: str) -> str:
    """Live rows of a **clean** ``source`` – a filtered scan. ``canonical_id`` is
    ``entity_id``, for ftmq's query layer."""
    return f"SELECT *, entity_id AS canonical_id FROM {source} WHERE deleted_at IS NULL"

def dedupe_rows_sql(source: str) -> str:
    """Live rows of a ``source`` that may hold un-merged files: `_dedupe_sql`,
    ``canonical_id`` as in [`live_rows_sql`][live_rows_sql]. Row for row what
    the same partition reads once merged."""
    return _dedupe_sql(source, select="*, entity_id AS canonical_id")

def delta_scan_sql(table_uri: str) -> str:
    """``delta_scan('<uri>')``, the uri quote-escaped (no prepared parameters)."""
    return f"delta_scan('{_string_literal(table_uri)}')"

def build_merge_sql(
    shard: str,
    bucket: str,
    origin: str,
    grace_cutoff: datetime,
    source: str = TABLE_RAW.name,
    select: str = "*",
) -> str:
    """`_dedupe_sql` collapsing one partition for a merge, ordered by the file
    sort key. Tombstones within grace survive to keep shadowing.

    Args:
        shard: The partition's shard.
        bucket: The partition's bucket.
        origin: The partition's origin – validated before interpolation.
        grace_cutoff: Tombstones with ``deleted_at <= grace_cutoff`` are dropped.
        source: Relation holding the rows – ``statement_raw``, or one
            partition's files ([`partition_source_sql`][partition_source_sql]).
        select: Projection of the output, e.g. without the partition columns.
    """
    origin = validate_origin(origin)
    return _dedupe_sql(
        source=source,
        where=f"WHERE shard = '{shard}' AND bucket = '{bucket}' AND origin = '{origin}'",
        tombstone=(
            "(deleted_at IS NULL OR deleted_at > "
            f"TIMESTAMPTZ '{grace_cutoff.isoformat()}')"
        ),
        order_by="ORDER BY entity_id, fragment, role, prop, id, last_seen DESC",
        select=select,
    )

def partition_source_sql(
    files: Iterable[str], shard: str, bucket: str, origin: str
) -> str:
    """One partition's files as a `SHARDED_SCHEMA` relation, read without
    ``delta_scan`` – [`relation_source_sql`][relation_source_sql] over the
    files, columns they predate reading as ``NULL``.
    """
    files_sql = (
        f"(SELECT * FROM {read_parquet_sql(files)} UNION ALL BY NAME {_FILE_COLUMNS})"
    )
    return relation_source_sql(files_sql, shard, bucket, origin)

def relation_source_sql(relation: str, shard: str, bucket: str, origin: str) -> str:
    """``relation`` – rows of a partition's files, the partition columns absent –
    as a `SHARDED_SCHEMA` relation: the partition columns as constants,
    ``fragment`` coalesced to ``''``. ``origin`` is validated before
    interpolation."""
    origin = validate_origin(origin)
    constants = {"shard": shard, "bucket": bucket, "origin": origin}
    projection = ", ".join(
        (
            f"'{_string_literal(constants[f.name])}' AS {f.name}"
            if f.name in constants
            else (
                "COALESCE(fragment, '') AS fragment" if f.name == "fragment" else f.name
            )
        )
        for f in SHARDED_SCHEMA
    )
    return f"(SELECT {projection} FROM {relation})"

def read_parquet_sql(files: Iterable[str]) -> str:
    """``read_parquet`` over ``files``, unioned by name. No hive partitioning:
    it would parse a shard like ``10`` as an integer."""
    paths = ", ".join(f"'{_string_literal(f)}'" for f in files)
    return f"read_parquet([{paths}], union_by_name = true, hive_partitioning = false)"

def worker_duckdb_config(workers: int) -> dict[str, str]:
    """[`duckdb_config`][duckdb_config] for one of ``workers`` processes, one call
    per task.

    Memory limit and threads are split between the workers.
    """
    config = duckdb_config()
    budget = parse_byte_size(config["memory_limit"])
    config["memory_limit"] = f"{budget // workers}B"
    config["threads"] = str(max((os.cpu_count() or 1) // workers, 1))
    return config

def build_shard_sql(
    shard: str, bucket: str, origin: str, shards: int, source: str = TABLE_RAW.name
) -> str:
    """One partition's raw rows with ``shard`` recomputed for ``shards`` shards.

    Unordered and un-deduped – the follow-up merge does both. ``origin`` is
    validated before interpolation.
    """
    origin = validate_origin(origin)
    return (
        f"SELECT * REPLACE ({shard_expr_sql(shards)} AS shard) "
        f"FROM {source} "
        f"WHERE shard = '{shard}' AND bucket = '{bucket}' AND origin = '{origin}'"
    )

def shard_expr_sql(shards: int, column: str = "entity_id") -> str:
    """DuckDB expression for ``column``'s shard key – the SQL twin of
    `entity_shard`, pinned to agree by ``test_shard_expr_sql_parity``."""
    if shards <= 1:
        return "'0'"
    width = shard_hex_width(shards)
    return (
        f"printf('%0{width}x', "
        f"(('0x' || substr(sha1({column}), 1, 8))::BIGINT) % {int(shards)})"
    )

def entity_chunks(
    streams: Iterable[Iterator[pa.RecordBatch]], target: int
) -> Iterator[pa.Table]:
    """Cut ``entity_id``-sorted streams into entity-closed chunks of about
    ``target`` Arrow bytes, in ascending ``entity_id`` ranges.

    A chunk holds every row of each of its entities – across all streams – and
    no two chunks share an entity, so a query keyed on ``entity_id`` gives the
    same rows per chunk as over the whole. Rows within a chunk are not sorted.
    Memory is the chunk plus about one batch per stream.

    Each step advances the stream whose newest batch ends lowest; every row
    below the lowest such end is complete, as no stream has more of it to give.

    Args:
        streams: Record batches sharing one schema, each stream ascending by
            ``entity_id``.
        target: Arrow bytes to read between two chunks.

    Yields:
        ``pa.Table`` per chunk.
    """
    sources = [iter(stream) for stream in streams]
    ends: list[tuple[str, int]] = []  # (last entity_id read, stream) heap
    held: list[pa.RecordBatch] = []
    since = 0

    def pull(i: int) -> None:
        nonlocal since
        for batch in sources[i]:
            if batch.num_rows:
                held.append(batch)
                since += batch.nbytes
                heapq.heappush(ends, (batch.column("entity_id")[-1].as_py(), i))
                return

    for i in range(len(sources)):
        pull(i)
    while ends:
        _, i = heapq.heappop(ends)
        pull(i)
        if (ends and since < target) or not held:
            continue
        table = pa.Table.from_batches(held)
        since = 0
        if not ends:
            held = []
            yield table
            continue
        below = pc.less(table.column("entity_id"), ends[0][0])
        held = table.filter(pc.invert(below)).to_batches()
        chunk = table.filter(below)
        if chunk.num_rows:
            yield chunk