fix/collect_keys_merge - #1419
Conversation
SkylordA
added
bug
feature
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1419 +/- ##
==========================================
+ Coverage 74.10% 75.23% +1.13%
==========================================
Files 214 215 +1
Lines 14140 14279 +139
==========================================
+ Hits 10478 10743 +265
+ Misses 3662 3536 -126 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Why DuckDB was investigated and ultimately not usedBefore settling on Polars, DuckDB was benchmarked as an alternative for this same merge step — it looked like a clear win in isolated, standalone benchmarks. It was not used in the end: a live-platform investigation found a critical failure mode where DuckDB's own internal thread pool stops dispatching work to its worker threads under real memory pressure, turning a merge that should take seconds into one that takes minutes (and, at larger scale in local stress testing, ones that never completed at all). Polars doesn't exhibit this problem under identical conditions. Standalone benchmark (
|
| Approach | CSV disjoint | CSV duplicated | Parquet disjoint | Parquet duplicated |
|---|---|---|---|---|
| current (old) | 62.65s / 4,979 MB | 11.08s / 3,685 MB | 2.50s / 2,304 MB — wrong | 2.47s / 2,297 MB |
| Polars | 8.91s / 8,160 MB | 2.52s / 881 MB | 7.13s / 8,943 MB | 2.21s / 964 MB |
| DuckDB | 3.43s / 4,033 MB | 1.72s / 1,959 MB | 1.89s / 4,855 MB | 0.98s / 1,602 MB |
("wrong" = the pre-existing parquet index-dedup bug this PR also fixes; disjoint content is the realistic case per the actual issue report.)
In this isolated setting, DuckDB was the fastest option across the board, with comparable memory usage. This is what made it worth investigating further.
Live-platform test (investigate_973.py)
The standalone win did not survive contact with the real deployment. Same config throughout — 50 lookup chunks, 1,000,000 rows/chunk, CSV, disjoint output, worker pod freshly restarted immediately before each run so neither approach had a "warmer" deployment than the other:
| Approach | collect_keys time |
|---|---|
| Polars | 39.8s |
| DuckDB | 266.4s |
That's DuckDB at ~6.7x slower than Polars on the real deployment, at the exact same test size — despite being the fastest option standalone. Both sides of this comparison are well-corroborated, not one-off numbers: DuckDB stayed in a tight 244-266s band across five separate runs, and Polars' 39.8s matches two independent fresh-pod measurements taken half an hour apart, and was reconfirmed again during a full re-run of this PR's own benchmark table (see PR description) that reproduced every OOM boundary exactly. (An earlier draft of this comment used 67.9s for Polars, matching what the PR description said at the time — that number didn't hold up under a controlled retest and the PR description above has been corrected to 39.8s accordingly.) Digging into why DuckDB stalls turned into most of the investigation below.
Root cause
A deep dive (thread-state sampling, memory instrumentation, reading DuckDB's own source, and a live kernel-level capture) traced this to DuckDB's TaskSchedulerPool: under sustained memory pressure, its worker threads are created successfully but never get dispatched work — confirmed directly via /proc/<tid>/wchan, showing them parked in futex_wait_queue (the wait path behind its internal moodycamel::LightweightSemaphore) for the full length of the stall, with memory usage completely flat, one thread doing all the work alone. The same capture against Polars shows its idle threads use the identical kernel wait state — that alone isn't the problem — but Polars' active-thread set keeps rotating as real work is dispatched, while DuckDB's stay parked. This reproduces in plain Docker under a real Celery worker (prefork pool, max-tasks-per-child=1, matching production's own setup), with no Kubernetes or minikube involved, so it isn't specific to this platform. It correlates with real memory pressure (a clean local dose-response experiment, and directly measured on this platform: available memory bottoming out at 3.75GB on both the DuckDB and Polars runs, with only DuckDB showing the stall).
Two mitigations were tried and both failed: giving the platform more memory (15GB → 20GB → 24GB) didn't reduce the stall or total time, and capping DuckDB's own memory_limit left the stall unchanged while making total runtime ~2x worse (forced disk spilling under the same pressure). One prior public report of a similar-sounding symptom exists (duckdb/duckdb#8726, 2023), but it's a different bug: their repro used DuckDB's implicit default connection (created at import time, before Celery forks its workers, so the fork loses its thread pool), fixed by calling duckdb.connect() explicitly inside the task. Our code already did that from the start - our threads are freshly created post-fork, confirmed via thread-count and the wchan capture - they just stop getting dispatched work under memory pressure. No upstream fix or tracked issue exists for either.
Conclusion: not fixable from this codebase — the defect is inside DuckDB's native thread pool. Sticking with Polars, which doesn't share this problem under any of the conditions tested.
Co-authored-by: carlfischerjba <carl.fischer@jbarisk.com>
4 participants
fix/collect_keys_merge
Problem
collect_keysmerges all parallel chunks by loading each fully into memory, concatenating, then deduping - peak memory scales with total data across chunks, not unique data. Reported on a model whose lookup produces substantial extra output beyondkeys.csv(multiple csv/parquet files, potentially several GB each, multiplied across chunks), exhausting memory and crashing or stalling the worker. Parquet has a second, independent bug: it dedupes by pandas index rather than row content, but chunk files carry no meaningful index (each chunk's is reset locally before its lookup runs) - two chunks' rows collide positionally even when entirely disjoint, silently dropping nearly all cross-chunk data.Change
scan_csv→unique→sink_csv) and Parquet (scan_parquet→unique→sink_parquet), deduping by row content in both cases - fixes the parquet bug in the same pass. Polars'unique()resolves hash collisions by comparing actual values, not trusting the hash alone.POLARS_MAX_THREADS, override viaOASIS_POLARS_MAX_THREADS). Polars sizes its thread pool from host CPU count at import time, which oversubscribes under celery's prefork concurrency (also host-CPU-sized by default) - a defensive default consistent with how the rest of the worker manages parallelism, not itself the source of the memory fix.Dependencies
Added
polars(>=1.40,<2.0) to worker requirements viapip-compile.Scope
This bounds memory to unique rows rather than total rows, and fixes the parquet bug - a real improvement, and consistently faster. It does not give unbounded scalability: for genuinely disjoint content (distinct output per chunk - the realistic case actually reported here), correct dedup still needs every unique row held somewhere, so old and new code eventually hit the same memory wall; new code just reaches it later, and is faster/lighter throughout. Content that overlaps heavily across chunks is a different story - memory there tracks the much smaller unique set, and new code keeps working far past where old code fails. Both are benchmarked below; disjoint is the one that matches the actual report.
Testing
Standalone benchmark (synthetic chunk files) and live-platform test (real analysis through a deployed worker, custom lookup module, peak memory from the container's cgroup), both against disjoint (distinct output per chunk - realistic) and duplicated (byte-identical per chunk - an artificial stress test, not a confirmed real pattern) content.
Standalone, 50 chunks, current → new:
Old parquet is wrong at every disjoint size tested - its lower time/memory reflects doing far less real work, not efficiency. Duplicated content can't expose the index bug (every chunk collides with itself either way), so old parquet is correct there.
Live platform, 50 lookup chunks, old → new:
"Killed (OOM)" is a genuine kernel-level SIGKILL (
WorkerLostError), not a code exception. For disjoint content, old and new hit the same real memory wall past ~150M unique rows on this test node (inherent, not a defect) - new code is faster and lighter throughout, and pushes the wall out by half a tier (succeeds at 75M where old already fails). For duplicated content the fix goes further: old still fails from 150M onward, new succeeds cleanly to 300M, since its memory tracks unique rows rather than total volume.closes #973