Reduce peak memory in Delta Lake export flush path

[bdpedigo]

Summary

The Delta Lake export path (materializationengine/workflows/deltalake_export.py) accumulates a large in-memory buffer of Arrow RecordBatches and then transforms it wholesale through several representations (Arrow table → Polars → geometry-decoded Polars → per-spec derived Polars → Arrow again) before handing it to write_deltalake. Peak RSS is therefore some poorly-characterized multiple of flush_threshold_bytes (default 2 GiB), and large tables cannot be exported at all on current worker sizing.

This is not hypothetical — it OOMKilled the deltalake:write_deltalake_table worker three times in a row on a production export (details below). This issue tracks making the export path itself cheaper in memory. The helm-chart change that makes the existing knobs settable (CAVEconnectome/cave-helm-charts#4) is a mitigation, not a fix — see the cross-reference at the bottom.

Things to check

1. _flush_buffer multiplies the buffer several times over before anything is written

materializationengine/workflows/deltalake_export.py:1110-1165. This is where the whole problem lives. In order:

• materializationengine/workflows/deltalake_export.py:1120-1121 — pa.Table.from_batches(buffer) then buffer.clear(). The comment says "free buffer memory immediately", but from_batches is zero-copy: the resulting table still references the same underlying buffers, so clearing the Python list frees nothing. The comment is misleading and should go either way.
• materializationengine/workflows/deltalake_export.py:1129-1130 — pl.from_arrow(arrow_table) followed by del arrow_table. Whether the del actually frees anything depends on whether Polars copied. For column types where the conversion is zero-copy, the del is a no-op and both representations remain live. This needs measuring per column type rather than assuming either way; the relevant types here are int64 ids/root-ids, and large binary (WKB) geometry.
• materializationengine/workflows/deltalake_export.py:1087-1107 — _strip_arrow_extension_types rebuilds the table via set_column per extension column. Believed cheap (chunk .storage views are zero-copy) but worth confirming it doesn't force a copy.
• materializationengine/workflows/deltalake_export.py:1133-1134 — decode_geometry_columns (defined at materializationengine/workflows/deltalake_export.py:1018) adds decoded x/y/z Int32 columns on top of the existing frame. A synapse table typically has 3 geometry columns (pre/post/ctr positions), so this is 9 added columns and a significant addition, not a marginal one.
• materializationengine/workflows/deltalake_export.py:1136-1155 — the per-spec loop keeps df alive across all iterations because later specs need it, while add_morton_column (materializationengine/workflows/deltalake_export.py:965) and assign_partition (materializationengine/workflows/deltalake_export.py:835) each build additional frames on top of it. Note discover_default_output_specs (materializationengine/workflows/deltalake_export.py:254) emits one spec per indexed column, so on a synapse table this fan-out is several specs wide.
• materializationengine/workflows/deltalake_export.py:1162 — write_df.to_arrow() fully re-materializes the frame back into Arrow. There is already an inline # TODO check whether this to_arrow hurts memory usage on that exact line. Based on the incident this is very likely a real cost and plausibly the single largest contributor, though it has not been profiled.
• materializationengine/workflows/deltalake_export.py:1160-1165 — write_deltalake then does its own internal buffering and parquet encoding on top of everything above.

2. decode_geometry_columns has its own amplification (possibly the worst single spot)

Separate from the pipeline-stage issue, the geometry decode itself is expensive per flush:

• materializationengine/workflows/deltalake_export.py:1052-1054 — wkb_series.drop_nulls().to_list() materializes the entire WKB column as a Python list of bytes objects (list of pointers + per-object overhead), then shapely.from_wkb(...) builds an N-element object array of Shapely geometries, then shapely.get_coordinates(..., include_z=True) produces an (N, 3) float64 array. For ~10M rows × 3 geometry columns these intermediates are large in their own right — and note the comment at materializationengine/workflows/deltalake_export.py:1037-1038 explicitly claims to be avoiding materializing the column as Python objects, which the null-mask computation does avoid but the decode a few lines later does not.
• materializationengine/workflows/deltalake_export.py:1066-1074 — the has-nulls path is worse still: three full-length dtype=object numpy arrays plus .tolist() on each. add_morton_column has the same shape at materializationengine/workflows/deltalake_export.py:1011-1013.

Whichever direction the flush path goes, these want to be numpy/Arrow-native and chunked.

3. flush_threshold_bytes is not a memory bound and is hard to reason about

• Default 2 GiB (2 * 1024 * 1024 * 1024) in export_table_to_deltalake at materializationengine/workflows/deltalake_export.py:1173.
• Read from config as DELTALAKE_FLUSH_THRESHOLD_BYTES in write_deltalake_table at materializationengine/workflows/deltalake_export.py:2006-2008.
• Default also defined in materializationengine/config.py:67-69.

The threshold is measured on the input buffer (batch.nbytes, accumulated at materializationengine/workflows/deltalake_export.py:1236), while actual peak RSS is some multiple of it determined by the stages in (1) and (2). So it is a proxy knob, not a bound, and operators cannot derive a safe value from a memory limit. Ideally the fix in (5) makes the knob much less load-bearing; failing that, it would help to at least document the observed multiplier.

4. The view export path has the same shape and should be fixed together

• _flush_flat_buffer at materializationengine/workflows/deltalake_export.py:1281-1300 — identical structure to _flush_buffer (same from_batches / buffer.clear() / pl.from_arrow / del / decode_geometry_columns / df.to_arrow() sequence, minus the per-spec fan-out).
• _stream_view_to_flat_lake at materializationengine/workflows/deltalake_export.py:1303-1342 — same accumulate-then-flush loop.
• export_view_to_deltalake at materializationengine/workflows/deltalake_export.py:1506, with its own 2 GiB default at materializationengine/workflows/deltalake_export.py:1511.

Worth noting that pass 2 of the view path, _repartition_view_spec_from_lake at materializationengine/workflows/deltalake_export.py:1438 (loop at materializationengine/workflows/deltalake_export.py:1477-1493), already does the per-batch thing being proposed in (5): it iterates dataset.to_batches() and converts, morton-encodes, partitions, and writes one batch at a time. That is roughly the target shape for the table path too.

5. The post-export optimize step is a second, separate spike

optimize_deltalake at materializationengine/workflows/deltalake_export.py:1599, called per spec at the end of export_table_to_deltalake (materializationengine/workflows/deltalake_export.py:1261-1278) and in both phases of the view path. Z-ordering 639M rows is expensive and this is not addressed by fixing the flush path — a table can now finish streaming and then die in optimize.

DELTALAKE_OPTIMIZE_MAX_SPILL_SIZE_BYTES and DELTALAKE_OPTIMIZE_TARGET_SIZE_BYTES exist in materializationengine/config.py:76-88 as untested mitigations; there is an existing comment at materializationengine/config.py:81-83 noting that spilling to disk has not been tested on the mesh worker nodes and it is unclear how it behaves there out of the box. That needs an actual test at scale.

Proposed direction (to evaluate, not a settled design)

Rather than accumulating a large buffer and transforming it wholesale, hand write_deltalake a pyarrow.RecordBatchReader and do the geometry decode and partition assignment per batch inside that reader. Peak memory then becomes roughly one batch × number of output specs, instead of full buffer × number of pipeline stages.

The streaming source is already in place: stream_table_to_arrow at materializationengine/workflows/deltalake_export.py:883 already yields RecordBatches straight from cur.fetch_record_batch() (materializationengine/workflows/deltalake_export.py:929-931), and _repartition_view_spec_from_lake already demonstrates the per-batch transform shape.

Open questions that need answering before committing to this:

• Does write_deltalake with a RecordBatchReader still produce reasonably sized parquet files, or does per-batch writing create a small-files problem that pushes cost into optimize? The current large-flush design may exist partly to get decent file sizes; if so, the right answer may be "small transform batches, large write batches" rather than one batch end to end.
• Can the per-spec fan-out share a single pass instead of re-deriving frames per spec? Today each spec independently rebuilds morton/partition columns from the same df. With a reader-based design the source is single-pass by construction, so either the specs need to be written in one pass with a superset of derived columns, or the source needs to be re-read per spec (N streams from Postgres) — that tradeoff should be made deliberately.
• Whether pl.from_arrow / to_arrow should be in the path at all per batch, or whether the derived columns (x/y/z, morton, partition) can be computed with numpy + pa.RecordBatch directly and appended to the batch, skipping the Polars round trip.
• Batch size control: fetch_record_batch() batch sizes are whatever ADBC gives us; the design should not assume they are small.

Success criterion: synapses_v1dd (639M rows) exports to completion within a bounded, documented memory limit on the current node class, with peak RSS predictable from a config value.

Cross-reference

CAVEconnectome/cave-helm-charts#4 ("expose deltalake worker memory tuning options") makes DELTALAKE_FLUSH_THRESHOLD_BYTES, the two optimize knobs, and an optional deltalake worker memory limit settable from the chart — previously deployments were stuck on the built-in defaults, which is why the v1dd worker was running with a 2 GiB flush threshold and no memory limit.

To be clear about what that PR does and does not do: it makes the bound reachable (operators can lower the threshold) and failures diagnosable (a container memory limit means a clean container OOM rather than a node-level kill). It does not make the export cheaper — lowering the flush threshold just trades peak memory for more, smaller writes and more optimize work. This issue tracks the actual memory-efficiency work in the export path.