feat(tools): ensure-sort-keys — canonical DuckLake sort keys for the tenant column

tests/integration/test_sort_keys_integration.py

@@ -0,0 +1,73 @@
+"""ensure-sort-keys: idempotent DuckLake SET SORTED BY maintenance.
+
+The canonical sort keys (ducklake_maintenance.SORT_KEYS) make the extension
+sort new writes per file and honor the sort in compaction merges, so parquet
+row-group min/max stats become tight for the tenant column — ClickHouse's
+row-group pruning then skips most of a file for team_id-filtered queries.
+"""
+
+import ducklake_maintenance
+import pytest
+
+
+def _sort_exprs(conn, table_name):
+    return conn.execute(
+        """
+        SELECT e.sort_key_index, e.expression, e.sort_direction
+        FROM __ducklake_metadata_lake.ducklake_sort_expression e
+        JOIN __ducklake_metadata_lake.ducklake_sort_info i
+          ON i.sort_id = e.sort_id AND i.table_id = e.table_id AND i.end_snapshot IS NULL
+        JOIN __ducklake_metadata_lake.ducklake_table t ON t.table_id = e.table_id
+        WHERE t.table_name = ?
+        ORDER BY e.sort_key_index
+        """,
+        [table_name],
+    ).fetchall()
+
+
+def test_ensure_sort_keys_applies_and_is_idempotent(ducklake_conn):
+    conn = ducklake_conn
+    conn.execute("CREATE TABLE lake.main.t (team_id BIGINT, ts VARCHAR)")
+    conn.execute("INSERT INTO lake.main.t VALUES (1, 'a'), (2, 'b')")
+
+    sort_keys = {("main", "t"): ["team_id"]}
+    applied = ducklake_maintenance.ensure_sort_keys(conn, sort_keys=sort_keys, dry_run=False)
+    assert applied == [("main", "t")]
+    assert _sort_exprs(conn, "t") == [(0, "team_id", "ASC")]
+
+    # second run: nothing to change
+    applied = ducklake_maintenance.ensure_sort_keys(conn, sort_keys=sort_keys, dry_run=False)
+    assert applied == []
+    assert _sort_exprs(conn, "t") == [(0, "team_id", "ASC")]
+
+
+def test_ensure_sort_keys_updates_changed_key(ducklake_conn):
+    conn = ducklake_conn
+    conn.execute("CREATE TABLE lake.main.t (team_id BIGINT, ts VARCHAR)")
+    conn.execute(f"ALTER TABLE lake.main.t SET SORTED BY (ts)")
+
+    applied = ducklake_maintenance.ensure_sort_keys(conn, sort_keys={("main", "t"): ["team_id"]}, dry_run=False)
+    assert applied == [("main", "t")]
+    assert _sort_exprs(conn, "t") == [(0, "team_id", "ASC")]
+
+
+def test_ensure_sort_keys_dry_run_changes_nothing(ducklake_conn):
+    conn = ducklake_conn
+    conn.execute("CREATE TABLE lake.main.t (team_id BIGINT, ts VARCHAR)")
+
+    applied = ducklake_maintenance.ensure_sort_keys(conn, sort_keys={("main", "t"): ["team_id"]}, dry_run=True)
+    assert applied == [("main", "t")]
+    assert _sort_exprs(conn, "t") == []
+
+
+def test_sort_keys_cover_megaduck_tables():
+    """The canonical set must name (schema, table) pairs of the megaduck tenant tables."""
+    expected = {
+        ("main", "events"): ["team_id", "timestamp"],
+        ("main", "events_nrt"): ["team_id", "timestamp"],
+        ("main", "heatmap_events"): ["team_id", "timestamp"],
+        ("main", "person"): ["team_id"],
+        ("main", "person_distinct_id"): ["team_id"],
+        ("main", "groups"): ["team_id"],
+    }
+    assert ducklake_maintenance.SORT_KEYS == expected

tools/ducklake_maintenance.py

@@ -1969,6 +1969,88 @@ def _merge_adjacent_call(schema_name: str | None, table_name: str, args: list[st
     )
 
 
+# Canonical DuckLake sort keys: (schema, table) -> sort columns. Applied to
+# megaduck on 2026-08-19; the cron keeps them in place (and applies them to
+# any catalog the routine runs against). The DuckLake extension sorts new
+# writes per file and honors the sort in compaction merges, so parquet
+# row-group min/max stats become tight for the tenant column — the
+# ClickHouse reader's row-group pruning then skips most of a file for
+# team_id-filtered queries (measured: 4/4 row groups reduced to 1 on sorted
+# files; before this, every file's stats spanned the full team_id range and
+# a `team_id = X` query was a full table scan).
+SORT_KEYS: dict[tuple[str, str], list[str]] = {
+    ("main", "events"): ["team_id", "timestamp"],
+    ("main", "events_nrt"): ["team_id", "timestamp"],
+    ("main", "heatmap_events"): ["team_id", "timestamp"],
+    ("main", "person"): ["team_id"],
+    ("main", "person_distinct_id"): ["team_id"],
+    ("main", "groups"): ["team_id"],
+}
+
+
+def _current_sort_keys(conn: duckdb.DuckDBPyConnection) -> dict[tuple[str, str], list[str]]:
+    """(schema, table) -> ordered sort expressions, from the catalog's sort tables."""
+    try:
+        rows = conn.execute(
+            f"""
+            SELECT sch.schema_name, t.table_name, e.expression
+            FROM {METADATA_SCHEMA}.ducklake_sort_expression e
+            JOIN {METADATA_SCHEMA}.ducklake_sort_info i
+              ON i.sort_id = e.sort_id AND i.table_id = e.table_id AND i.end_snapshot IS NULL
+            JOIN {METADATA_SCHEMA}.ducklake_table t ON t.table_id = e.table_id
+            JOIN {METADATA_SCHEMA}.ducklake_schema sch ON sch.schema_id = t.schema_id
+            WHERE t.end_snapshot IS NULL
+            ORDER BY sch.schema_name, t.table_name, e.sort_key_index
+            """
+        ).fetchall()
+    except duckdb.CatalogException:
+        # Catalog predates sorted tables (no ducklake_sort_* tables yet).
+        return {}
+    current: dict[tuple[str, str], list[str]] = {}
+    for schema_name, table_name, expression in rows:
+        current.setdefault((schema_name, table_name), []).append(expression)
+    return current
+
+
+def ensure_sort_keys(
+    conn: duckdb.DuckDBPyConnection,
+    sort_keys: dict[tuple[str, str], list[str]] | None = None,
+    dry_run: bool = False,
+) -> list[tuple[str, str]]:
+    """Apply SORT_KEYS to the catalog, idempotently. Returns the tables (re)applied.
+
+    A table whose current sort key already matches is skipped; a table with a
+    different sort key is re-set (SET SORTED BY replaces the sort spec).
+    Dry-run reports what would change without touching the catalog.
+    """
+    wanted = SORT_KEYS if sort_keys is None else sort_keys
+    current = _current_sort_keys(conn)
+    applied: list[tuple[str, str]] = []
+    for (schema_name, table_name), columns in sorted(wanted.items()):
+        existing = current.get((schema_name, table_name))
+        if existing == columns:
+            log.info("ensure-sort-keys: %s.%s already sorted by (%s), skipping", schema_name, table_name, ", ".join(columns))
+            continue
+        key_sql = ", ".join(f'"{c}"' if not c.isidentifier() else c for c in columns)
+        log.info(
+            "ensure-sort-keys: %s %s.%s SET SORTED BY (%s)%s",
+            "would apply" if dry_run else "applying",
+            schema_name,
+            table_name,
+            ", ".join(columns),
+            " (dry run)" if dry_run else "",
+        )
+        if not dry_run:
+            if existing:
+                # SET SORTED BY appends a new sort spec; it does not end the old one.
+                conn.execute(f'ALTER TABLE {ATTACH_NAME}."{schema_name}"."{table_name}" RESET SORTED BY')
+            conn.execute(
+                f'ALTER TABLE {ATTACH_NAME}."{schema_name}"."{table_name}" SET SORTED BY ({key_sql})'
+            )
+        applied.append((schema_name, table_name))
+    return applied
+
+
 def compact(
     conn: duckdb.DuckDBPyConnection,
     tier: int,
@@ -2422,6 +2504,18 @@ def build_parser() -> argparse.ArgumentParser:
     p.add_argument("--days", type=int, default=7)
     p.add_argument("--dry-run", action="store_true")
 
+    # ensure-sort-keys
+    p = sub.add_parser(
+        "ensure-sort-keys",
+        help="Apply the canonical DuckLake sort keys (SORT_KEYS) idempotently",
+        description=(
+            "Apply the canonical sort keys so new writes and compaction outputs are "
+            "sorted by the tenant column (tight parquet row-group stats -> ClickHouse "
+            "row-group pruning on team_id filters). Idempotent; safe to run every pass."
+        ),
+    )
+    p.add_argument("--dry-run", action="store_true")
+
     # checkpoint
     sub.add_parser("checkpoint", help="CHECKPOINT (merge + expire + cleanup)")
 
@@ -2556,6 +2650,8 @@ def main(argv: list[str] | None = None) -> None:
                 fsck(conn, args.dry_run, args.max_iterations)
             case "orphans":
                 orphans(conn, args.dry_run)
+            case "ensure-sort-keys":
+                ensure_sort_keys(conn, dry_run=args.dry_run)
             case "maintain":
                 maintain(conn, args.days, args.dry_run)
             case "checkpoint":

tools/justfile

@@ -426,9 +426,34 @@ bootstrap-index-file-partition-value-cover: _confirm-target
     @echo ">>> ducklake_file_partition_value_cover_idx"
     @{{ _psql }} -c "CREATE INDEX CONCURRENTLY IF NOT EXISTS ducklake_file_partition_value_cover_idx ON public.ducklake_file_partition_value USING btree (data_file_id, partition_key_index, partition_value);"
 
+# Table-scoped per-file stats fetch: `WHERE table_id = ? AND data_file_id IN
+# (...)`, the shape ClickHouse's DuckLake reader issues when listing a table's
+# files for a query (the catalog-side read behind
+# StorageObjectStorageSource file iteration). Composite so the IN-list
+# prunes inside the index without a table-wide scan.
+[group('bootstrap')]
+bootstrap-index-file-column-stats-table-file: _confirm-target
+    @echo ">>> ducklake_file_column_stats_table_file_idx"
+    @{{ _psql }} -c "CREATE INDEX CONCURRENTLY IF NOT EXISTS ducklake_file_column_stats_table_file_idx ON public.ducklake_file_column_stats USING btree (table_id, data_file_id);"
+
+# Same shape for partition values (ClickHouse partition-pruning fetches).
+[group('bootstrap')]
+bootstrap-index-file-partition-value-table-file: _confirm-target
+    @echo ">>> ducklake_file_partition_value_table_file_idx"
+    @{{ _psql }} -c "CREATE INDEX CONCURRENTLY IF NOT EXISTS ducklake_file_partition_value_table_file_idx ON public.ducklake_file_partition_value USING btree (table_id, data_file_id);"
+
 # Create every DuckLake catalog btree, sequentially (same-relation builds serialize anyway)
 [group('bootstrap')]
-bootstrap-indexes: bootstrap-index-data-file-compaction bootstrap-index-data-file-compaction-order bootstrap-index-data-file-snapshot-read bootstrap-index-delete-file-snapshot-read bootstrap-index-delete-file-table bootstrap-index-delete-file-metrics bootstrap-index-file-column-stats bootstrap-index-file-partition-value-file bootstrap-index-file-partition-value-table bootstrap-index-file-partition-value-cover
+bootstrap-indexes: bootstrap-index-data-file-compaction bootstrap-index-data-file-compaction-order bootstrap-index-data-file-snapshot-read bootstrap-index-delete-file-snapshot-read bootstrap-index-delete-file-table bootstrap-index-delete-file-metrics bootstrap-index-file-column-stats bootstrap-index-file-column-stats-table-file bootstrap-index-file-partition-value-file bootstrap-index-file-partition-value-table bootstrap-index-file-partition-value-table-file bootstrap-index-file-partition-value-cover
+
+# Apply the canonical DuckLake sort keys (tenant column first on every
+# tenant table): new writes and compaction outputs come out sorted, giving
+# tight parquet row-group stats, which is what the ClickHouse reader's
+# row-group pruning needs for team_id-filtered queries. Idempotent.
+[group('bootstrap')]
+ensure-sort-keys: _confirm-target
+    @echo ">>> ensure-sort-keys"
+    @env {{ _target_env }} python {{ ducklake_maintenance }} ensure-sort-keys
 
 # ----------------------------------------------------------------------------
 # state-metrics daemon