PostHog / PostHog/millpond

Add sort-key support for Iceberg writes (better Parquet row-group pruning)

Open
#57 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
10
Forks
1
Avg merge
10m
Merged PRs (30d)
16

Description

Why

Today millpond/iceberg.py creates tables with UNSORTED_SORT_ORDER (the default — we pass no sort_order= to create_table) and calls Table.append() with whatever Arrow row order Kafka produced. Within each Parquet file, row-group min/max statistics end up loose because adjacent rows aren't correlated on any meaningful column, so reader-side predicate pushdown can't prune row groups.

Sorting each flushed batch by a configurable key (e.g. _inserted_at, or a high-cardinality dimension like team_id) before append produces tight per-row-group stats. For time-range queries — the dominant access pattern for ingest data — this typically translates to large reductions in row groups scanned without changing the data itself.

Proposed shape

A new env-var-driven sort spec:

Env var Default Meaning
ICEBERG_SORT_BY unset Comma-separated column names; first listed is the primary sort key. Example: team_id,_inserted_at

Wiring:

  1. Config.iceberg_sort_by: tuple[str, ...] parsed from the env var in config.load().
  2. iceberg.connect() / iceberg._ensure_table() reads the spec; if non-empty, builds a PyIceberg SortOrder and passes it to catalog.create_table(sort_order=...). On an existing table whose sort order differs, decide between (a) leave it alone (table sort order persists across writes; operator must explicitly evolve it) or (b) table.update_sort_order() to match. Defer (b) until needed.
  3. iceberg.write() calls batch.sort_by(...) on the Arrow table before _add_metadata_columns so the partition cols inherit the sorted order. Cost: roughly tens to hundreds of ms per ~100MB batch — negligible at our flush cadence.

Validation

  • Column existence: validate at startup if all keys reference real columns in the source schema, OR validate at first write when the schema is known. Schema evolution may add columns after startup, so first-write or per-batch validation is safer.
  • Reserved-column collision: _inserted_at is a fine sort key, but the four partition columns (year/month/day/hour) are degenerate (every row in a batch has the same value); reject them with a clear error rather than silently no-op'ing.
  • Null handling: PyArrow's sort_by accepts null_placement="at_end" — pick a default and document.

Out of scope for this issue

  • Compaction / rewriting existing unsorted files. When sort order is added to an existing table, new files are written sorted; old files stay as-is. Iceberg's rewrite_data_files action handles back-filling — that's a maintenance-job concern, separate ticket.
  • Multi-key with mixed direction (ASC/DESC per column). SortOrder supports it but env-var syntax for a future team_id:asc,_inserted_at:desc form is overkill until needed.

Acceptance criteria

  • Unit tests covering: sort-spec parsing from env var, SortOrder construction, batch.sort_by ordering result, behavior when sort col is missing from the batch
  • Integration test that writes a batch with shuffled values, scans the resulting Parquet file, and asserts row-group min/max are tight (or at least: that the values are monotonically non-decreasing within the file)
  • No-op behavior when ICEBERG_SORT_BY is unset (today's behavior preserved)
  • Docs note in AGENT.md covering the tradeoff (write-time sort cost vs read-time pruning win) and the back-fill caveat

Follow-up to #56 (the Iceberg switch).

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with Config.iceberg_sort_by in config.load(), then trace iceberg.connect(), iceberg._ensure_table(), and iceberg.write() to understand table creation and batch handling. Add tests for parsing, SortOrder construction, sorting, missing columns, and the unset no-op case, plus the integration check described in the issue. Update AGENT.md with the write-cost and back-fill caveats.

Written by the indexing model from the issue text.

Assessment

Tech stack
kafka, python
Domain
data-engineering, databases
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
52/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.