finalize_meds: replace buffered repartition scatter with native sink_parquet(PartitionBy) - #206
Merged
Merged
Conversation
…parquet(PartitionBy) _repartition_pass previously scattered every input row into a per- (split, output_shard) buffer in Python, flushed via pq.ParquetWriter at REPARTITION_FLUSH_ROW_THRESHOLD-row intervals. Replaces it with a single lazy scan_parquet across all input shards feeding polars' native sink_parquet(pl.PartitionBy(key=["__split", "__out_shard"], ...)), which owns its own internal buffering. The three lazy lookups (subject_id remap, split, output shard) are all derived from one remapped_subject expression object in a single with_columns call, deliberately mirroring the prior eager implementation's Series-level dependency order. Based on a scratch benchmark (20M rows, 100 input files, 400 output destinations at real FINALIZE_SUBJECTS_PER_SHARD sizing): native sink 0.52-0.72s wall / 1.35-1.42GB peak RSS vs. buffered 4.62-4.77s wall / 1.86-1.88GB peak RSS, with an exhaustive (not sampled) row-set equality check across all 400 destinations confirming byte-identical output. Numbers and full benchmark description folded into the module docstring. _sort_and_finalize_shard now takes a list of paths per destination rather than one, since the sink is free to split a destination across more than one physical file (not observed at benchmark scale, but not assumed never to happen at real scale). Dropped the pyarrow-schema- enforced write in favor of explicit polars dtype casts in the lazy pipeline (MEDS_POLARS_DTYPES_FINAL replaces MEDS_ARROW_SCHEMA_FINAL) -- pyarrow is no longer imported by this module at all. Tests: updated the two existing _repartition_pass tests for the new signature/return shape, added a regression test for a real bug hit while building this (the error path was naming a null instead of the actual raw subject_id -- fixed by keeping the pre-remap value under its own column through the with_columns call), a focused test isolating the eager-vs-lazy replace_strict translation specifically, an exhaustive A/B test against a preserved reference copy of the old buffered algorithm, and a dedicated test pinning down the exact pl.PartitionBy surface this depends on (multi-column key, include_key=False, file_path_provider) -- pl.PartitionBy is marked unstable by polars itself, so a future polars bump that changes this API now fails the test suite, not a live GEMINI run. Branch-only: not merged to main or mirrored to gemini yet. finalize's current live run on the GEMINI node executes the server's already- checked-out copy and is unaffected by anything on this branch; merge and mirror wait for confirmation that run has completed. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
_repartition_pass's hand-rolled buffered per-(split, output_shard) scatter with a single lazyscan_parquetacross all input shards feeding polars' nativesink_parquet(pl.PartitionBy(...)). The subject_id remap + split + output-shard lookups are all derived from oneremapped_subjectexpression object in a singlewith_columnscall, deliberately mirroring the prior eager implementation's Series-level dependency order.FINALIZE_SUBJECTS_PER_SHARDsizing): native sink 0.52-0.72s wall / 1.35-1.42GB peak RSS vs. buffered 4.62-4.77s wall / 1.86-1.88GB peak RSS, with an exhaustive (not sampled) row-set equality check across all 400 destinations confirming byte-identical output. Numbers and full description folded into the module docstring._sort_and_finalize_shardnow takes a list of paths per destination (the sink can in principle split one destination across multiple physical files). Dropped the pyarrow-schema-enforced write in favor of explicit polars dtype casts (MEDS_POLARS_DTYPES_FINALreplacesMEDS_ARROW_SCHEMA_FINAL) -- pyarrow is no longer imported by this module at all.replace_stricttranslation test, an exhaustive A/B test against a preserved reference copy of the old buffered algorithm, a dedicated test pinning down the exactpl.PartitionBysurface used (pl.PartitionByis marked unstable by polars itself, so a future polars bump that changes this API now fails the test suite instead of a live GEMINI run), a multi-part-file concat test, and a regression test for a real bug caught while building this (the error path was naming a null instead of the actual raw subject_id).Not merging yet: gated on tonight's live
finalizerun on the GEMINI node completing -- this branch's push has no effect on that run (it executes the server's already-checked-out copy). Merge and mirror togemini/mainwait for Amrit's morning report.Test plan
🤖 Generated with Claude Code