[zephyr] Switch internals to DataFusion

[loom-oa-dev]

TL;DR

PR #8651 replaces Polars inside Zephyr's shuffle with DataFusion, keeps Zephyr's external spill files on region-local GCS, and exposes DataFusion SQL for columnar map and reducer stages.

Against current main/Polars, the 13 SQL-migrated Datakit stages used 29.39% less CPU, 21.38% less summed stage time, and 49.93% less peak shard memory. Across all 16 shuffle-affected stages, including fuzzy verification's unchanged Python reducer, CPU fell 26.82%. Region-local output validators matched exactly.

Description

The SQL API mirrors the efficiency boundaries established by #8154 and #8256:

• Dataset.sql(...) runs batch-local map queries against an input view.
• Dataset.group_by(..., reducer=sql(...)) runs a query over the complete target shard after Zephyr's external merge.
• Scalar SQL functions provide expression-based routing keys without serializing DataFusion objects into a plan.

Datakit normalization, fuzzy graph construction, all connected-component iterations, and final fuzzy-attribute materialization use these paths. Fuzzy verification still uses the generic DataFusion shuffle with its existing Python reducer. Download, MinHash, memory-store loading, consolidation, tokenization, and cache-probe work do not execute the new SQL path.

DataFusion native disk spill stays disabled. Zephyr continues to write bounded sorted Parquet chunks and external merge runs to the stage filesystem, which is GCS on Iris. Eight target-shard ranges bound each DataFusion sort while preserving one selected file per mapper flush at reduce time.

Hypothesis or Goal

Remove Zephyr's Polars dependency and use DataFusion for shuffle sorting, reading, and columnar map/reduce computation without changing semantic output, moving external spill off GCS, or retaining the original DataFusion memory and CPU regressions.

Status

The implementation is complete on PR #8651, rebased on control commit dc584e765d with treatment head fbddc9de28.

The affected repository selection passed 1,529 tests with three skips and five expected failures. Ruff, Black, license checks, Pyrefly, documentation checks, and all remaining changed-file checks passed.

Current-main Datakit ferry

Workload: FineWeb-Edu sample/10BT at revision 87f0914, 14 input files, and the full experiments.ferries.datakit_ferry path used by the Polars columnar A/B. Both arms ran concurrently with identical non-preemptible resources in europe-west4. Inputs, outputs, status, caches, and Zephyr temporary files all used gs://marin-eu-west4; no cross-region GCS access was configured.

Control: main/Polars dc584e765d, /loom/zephyr-sql-ferry-20260828-control, 56m11.18s.

Treatment: DataFusion SQL fbddc9de28, /loom/zephyr-sql-ferry-20260828-treatment, 59m00.15s (+168.97s, +5.01%).

Operation / stage	CPU control	CPU treatment	CPU change	Elapsed control	Elapsed treatment	Elapsed change
SQL-migrated stages	17,530.04 s	12,378.79 s	-29.39%	903.15 s	710.07 s	-21.38%
All shuffle-affected stages	18,846.06 s	13,791.10 s	-26.82%	970.18 s	783.36 s	-19.26%
SQL · Normalize · stage0-Map → Scatter	2,035.54 s	990.56 s	-51.34%	161.77 s	80.67 s	-50.13%
SQL · Normalize · stage1-Reduce → Map	1,131.80 s	586.67 s	-48.16%	215.93 s	95.16 s	-55.93%
SQL · Fuzzy graph · stage0-Map → Scatter	5,352.49 s	2,201.18 s	-58.88%	147.19 s	104.56 s	-28.97%
SQL · Fuzzy graph · stage1-Reduce → Scatter	5,137.78 s	6,103.37 s	+18.79%	90.76 s	120.69 s	+32.98%
SQL · Fuzzy graph · stage2-Reduce → Write	1,129.47 s	1,062.20 s	-5.96%	25.58 s	39.37 s	+53.91%
SQL · CC iteration 1 · stage0-Map → Scatter	362.55 s	146.55 s	-59.58%	41.25 s	34.64 s	-16.02%
SQL · CC iteration 1 · stage1-Reduce → Map	420.52 s	237.70 s	-43.47%	24.57 s	28.18 s	+14.69%
SQL · CC iteration 2 · stage0-Map → Scatter	369.03 s	142.34 s	-61.43%	31.56 s	44.09 s	+39.71%
SQL · CC iteration 2 · stage1-Reduce → Map	424.50 s	242.73 s	-42.82%	16.77 s	32.74 s	+95.16%
SQL · CC iteration 3 · stage0-Map → Scatter	362.14 s	142.54 s	-60.64%	47.31 s	33.91 s	-28.32%
SQL · CC iteration 3 · stage1-Reduce → Map	412.52 s	233.07 s	-43.50%	27.89 s	35.69 s	+28.00%
SQL · Fuzzy attributes · stage0-Map → Scatter	252.46 s	110.66 s	-56.17%	58.96 s	35.86 s	-39.18%
SQL · Fuzzy attributes · stage1-Reduce	139.24 s	179.22 s	+28.71%	13.61 s	24.51 s	+80.13%
Shuffle only · Fuzzy verification · stage0-Map → Scatter	398.04 s	433.40 s	+8.88%	13.14 s	14.64 s	+11.35%
Shuffle only · Fuzzy verification · stage1-Reduce → Scatter	847.29 s	905.05 s	+6.82%	44.07 s	43.49 s	-1.30%
Shuffle only · Fuzzy verification · stage2-Reduce	70.69 s	73.86 s	+4.48%	9.82 s	15.16 s	+54.38%

Operation / stage	Peak memory control	Peak memory treatment	Memory change
SQL-migrated stages	7.94 GB	3.97 GB	-49.93%
All shuffle-affected stages	7.94 GB	3.97 GB	-49.93%
SQL · Normalize · stage0-Map → Scatter	7.94 GB	3.97 GB	-49.93%
SQL · Normalize · stage1-Reduce → Map	2.03 GB	1.74 GB	-14.57%
SQL · Fuzzy graph · stage0-Map → Scatter	2.43 GB	1.89 GB	-22.09%
SQL · Fuzzy graph · stage1-Reduce → Scatter	1.68 GB	2.15 GB	+27.61%
SQL · Fuzzy graph · stage2-Reduce → Write	2.21 GB	1.02 GB	-53.61%
SQL · CC iteration 1 · stage0-Map → Scatter	0.51 GB	0.55 GB	+8.16%
SQL · CC iteration 1 · stage1-Reduce → Map	0.52 GB	0.76 GB	+45.27%
SQL · CC iteration 2 · stage0-Map → Scatter	0.50 GB	0.55 GB	+8.42%
SQL · CC iteration 2 · stage1-Reduce → Map	0.52 GB	0.70 GB	+34.17%
SQL · CC iteration 3 · stage0-Map → Scatter	0.50 GB	0.55 GB	+8.12%
SQL · CC iteration 3 · stage1-Reduce → Map	0.52 GB	0.69 GB	+34.39%
SQL · Fuzzy attributes · stage0-Map → Scatter	0.57 GB	0.39 GB	-31.68%
SQL · Fuzzy attributes · stage1-Reduce	0.55 GB	0.66 GB	+18.96%
Shuffle only · Fuzzy verification · stage0-Map → Scatter	0.78 GB	0.72 GB	-7.59%
Shuffle only · Fuzzy verification · stage1-Reduce → Scatter	0.55 GB	0.76 GB	+38.47%
Shuffle only · Fuzzy verification · stage2-Reduce	0.53 GB	0.72 GB	+36.50%

The SQL-migrated stages processed the same 558,478,005 logical items; all shuffle-affected stages matched at 559,057,070. Byte telemetry is not comparable because the treatment counts Arrow buffer bytes while the control counts its row-oriented representation.

The control validator and treatment validator each found 9,672,101 downloaded rows, 9,268,156 normalized rows, 532,777 fuzzy cluster members with 246,814 representatives, 45,970 verified duplicates, and 9,222,186 consolidated and tokenized rows.

Measured CPU across the full ferry fell 9.04%, and summed stage time fell 4.09%. Root wall time rose 5.01% because unaffected runtime was noisier: tokenization elapsed rose from 333.92 to 474.38 seconds with CPU essentially unchanged at +0.69%, followed by a longer cache-probe scheduling delay. Unaffected-stage CPU was flat at -0.30%.

Earlier evidence

The range-only DataFusion A/B, before the SQL path, used 5.25% more CPU than Polars across the 16 shuffle-affected stages while lowering peak shard memory 43.28%. That isolated Python row materialization and reducer callbacks as the remaining actionable cost.

A control-only Polars target-range sweep showed that the memory mechanism is engine-independent:

Polars target ranges	Writer time	Change
1 (baseline)	2.600 s	-
2	2.649 s	+1.92%
4	2.770 s	+6.55%
8	3.319 s	+27.69%

Polars target ranges	Peak RSS	Change
1 (baseline)	1.434 GB	-
2	1.256 GB	-12.40%
4	1.080 GB	-24.70%
8	0.988 GB	-31.07%

An earlier pre-optimization Nemotron A/B raised affected-stage CPU 39.99% and peak memory from 7.88 to 13.42 GB. That run identified the original whole-buffer DataFusion SortExec problem and does not represent the current range-partitioned SQL treatment.

Links

• Implementation and review: PR #8651
• Per-stage A/B and historical detail: sticky PR comment
• Root-cause and benchmark investigation: Echo record

Decision Log

• DataFusion native disk spill stays disabled so its internal sorter remains in memory.
• Zephyr's explicit sorted chunks and external merge intermediates remain on the stage filesystem and therefore use region-local GCS on Iris.
• DataFusion uses eight target-shard ranges because two exhausted the sorter and four retained materially more memory.
• SQL is exposed as a serializable Zephyr plan object instead of exposing DataFusion runtime objects.
• Datakit uses SQL only where the transformation is columnar; fuzzy verification retains Python callbacks for full-text comparison and side effects.

Conclusion

The Python implementation removes the DataFusion memory and CPU regressions without modifying DataFusion. On the matched Datakit ferry, SQL-migrated DataFusion stages were faster than current-main Polars while using half the peak shard memory, and persisted outputs matched exactly.