fix: bound VortexWriter's global-dict retained memory to avoid OOM on huge files

VortexWriter buffers a global-dict candidate column's raw chunk data in
memory from its first chunk until close() -- a shared dictionary can only
be built once every chunk has been seen. On a huge, wide file (e.g. an
18.5M-row / 38-string-column real-world Parquet import) any column that
looks low-cardinality on an early chunk but never gets demoted keeps its
entire column pinned in the heap; with dozens of such columns the total
reaches several GB and the import throws OutOfMemoryError, independent of
how much heap is available (memory scales with file size x column count,
not with a bounded chunk size).

Add an aggregate retained-bytes budget (256 MB, GLOBAL_DICT_MAX_RETAINED_BYTES)
across all buffering dict-candidate columns. Each chunk appended to a
candidate updates a per-column and running-total estimate (estimateRetainedBytes);
crossing the budget demotes the largest-retained columns -- flushing their
already-buffered chunks as ordinary per-chunk segments and dropping them from
future global-dict candidacy -- until back under budget. This bounds writer
memory by the budget rather than by total file size, while still giving most
columns their shared dictionary in the common case.

Regression test (GlobalDictUtf8Test#retainedBytesBudgetExceeded_utf8_demotesToPerChunkChunkedLayout)
reproduces the bug shape at small scale via a test-only budget seam
(setDictRetainedBudgetForTest): a column with genuinely low, constant
cardinality (so the existing cardinality-ratio fallback never fires) whose
raw bytes alone cross a lowered budget, forcing mid-file demotion. Asserts
the demoted column lands as a plain per-chunk Chunked layout (not Dict) and
that every value still round-trips correctly across the demotion boundary.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: dfa1

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- `VortexWriter` no longer buffers a global-dictionary candidate column's raw data for the whole file, so importing a huge Parquet file (e.g. an 18.5M-row dataset) no longer exhausts the heap; a column whose retained bytes exceed a fixed budget is demoted to per-chunk encoding. ([a3b921b5](https://github.com/dfa1/vortex-java/commit/a3b921b5))
 - A chunked `List` column spanning several flat chunks now decodes into a single stitched array instead of throwing a raw `ClassCastException`; any other unhandled dtype now fails with `VortexException`. ([#268](https://github.com/dfa1/vortex-java/issues/268))
 - A chunked `Utf8`/`Binary` column with an entirely-null chunk (`NullArray`) no longer throws `chunk is not a VarBinArray`; the chunk materializes as an all-null run. ([#269](https://github.com/dfa1/vortex-java/issues/269))
 - A chunked `List` column with an entirely-null chunk (`NullArray`) no longer throws `chunk is not a ListArray`; the chunk's rows become zero-length lists with out-of-band nulls. ([#269](https://github.com/dfa1/vortex-java/issues/269))

writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java

@@ -99,6 +99,18 @@ public final class VortexWriter implements Closeable {
     // Kept low: global dict hurts high-cardinality F64 columns (ALP codes beat U16 dict codes).
     static final int GLOBAL_DICT_MAX_CARDINALITY = 2_048;
 
+    // Aggregate memory budget (bytes) for the raw data all columns may retain while buffering for a
+    // shared global dictionary. A global dict must see every chunk before it can be built, so a
+    // dict-candidate column's raw arrays are held from the first chunk until close(). On a huge,
+    // wide file (e.g. the 18.5M-row / 38-string-column NYC-311 Parquet import) any column
+    // mis-detected as low-cardinality on its first chunk — one whose distinct count grows only after
+    // millions of later rows — would otherwise pin its entire column in the heap; with dozens of such
+    // columns the total is several GB and the import OOMs. This budget bounds the SUM across all
+    // buffering columns: once the total is exceeded, the largest-retained columns are demoted (their
+    // buffered chunks flushed per-chunk, per-chunk encoding thereafter) until back under budget,
+    // keeping writer memory bounded by the budget rather than by total file size × column count.
+    static final long GLOBAL_DICT_MAX_RETAINED_BYTES = 256L * 1024 * 1024;
+
     private static final List<EncodingEncoder> DEFAULT_CODECS = List.of(
             new AlpEncodingEncoder(), new PrimitiveEncodingEncoder(), new BoolEncodingEncoder(),
             new DictEncodingEncoder(), new VarBinEncodingEncoder(), new ExtEncodingEncoder(),
@@ -124,6 +136,14 @@ public final class VortexWriter implements Closeable {
     private final Set<ColumnName> dictCandidates = new LinkedHashSet<>();
     private final Map<ColumnName, List<Object>> dictBuffers = new LinkedHashMap<>();
     private final Map<ColumnName, DictColRef> dictColRefs = new LinkedHashMap<>();
+    // Raw bytes retained per global-dict-candidate column, and their running sum; together they
+    // guard the aggregate memory budget (dictRetainedBudget) so no set of mis-detected columns can
+    // pin the heap. When the sum crosses the budget, the largest columns are demoted until under it.
+    private final Map<ColumnName, Long> dictRetainedBytes = new LinkedHashMap<>();
+    private long dictRetainedTotal = 0;
+    // Effective aggregate global-dict retention budget; the constant by default, lowered by tests
+    // to exercise the demotion path without allocating the full budget.
+    private long dictRetainedBudget = GLOBAL_DICT_MAX_RETAINED_BYTES;
     private boolean firstChunkSeen = false;
 
     // Per-column zone-maps, populated by flushZoneMaps() in close() when enableZoneMaps is set.
@@ -162,6 +182,12 @@ private VortexWriter(
         }
     }
 
+    // Test seam: lower the aggregate global-dict retention budget so the demotion path (see
+    // writeChunk) can be exercised without allocating GLOBAL_DICT_MAX_RETAINED_BYTES of column data.
+    void setDictRetainedBudgetForTest(long budgetBytes) {
+        this.dictRetainedBudget = budgetBytes;
+    }
+
     /// Builds a [WriteRegistry] from the given encoder list plus all built-in extension encoders.
     private static WriteRegistry buildRegistry(List<EncodingEncoder> encoders) {
         WriteRegistry.Builder b = WriteRegistry.builder();
@@ -496,7 +522,18 @@ public void writeChunk(Map<ColumnName, Object> columns) throws IOException {
             }
 
             if (dictCandidates.contains(colName)) {
-                dictBuffers.computeIfAbsent(colName, _ -> new ArrayList<>()).add(data);
+                List<Object> buffered = dictBuffers.computeIfAbsent(colName, _ -> new ArrayList<>());
+                buffered.add(data);
+                long delta = estimateRetainedBytes(data);
+                dictRetainedBytes.merge(colName, delta, Long::sum);
+                dictRetainedTotal += delta;
+                if (dictRetainedTotal > dictRetainedBudget) {
+                    // Aggregate budget exceeded — the buffered columns together no longer fit the
+                    // memory budget. Demote the largest-retained columns (flush their buffered chunks
+                    // per-chunk, encode per-chunk thereafter) until back under budget, so writer
+                    // memory stays bounded regardless of file size or column count.
+                    evictLargestDictColumnsUntilUnderBudget();
+                }
             } else {
                 long rowCount = arrayLength(data);
                 int segIdx = writeSegment(colDtype, data);
@@ -1164,6 +1201,82 @@ private static byte[] buildDictLayoutMetaBytes(PType codePType) {
 
     // ── Global dict helpers ───────────────────────────────────────────────────
 
+    /// Estimates the heap footprint of one chunk's worth of a column's raw data, used to bound the
+    /// per-column global-dict retention budget. Primitive arrays cost their element bytes; string
+    /// arrays cost each present string's UTF-16 char bytes plus a fixed per-element object-header
+    /// allowance (reference + `String`/`char[]` overhead), which dominates on wide, sparse columns.
+    ///
+    /// @param data the chunk data (primitive array, `String[]`, or a [NullableData] wrapper)
+    /// @return an approximate retained-byte count; never negative
+    private static long estimateRetainedBytes(Object data) {
+        Object values = data instanceof NullableData nd ? nd.values() : data;
+        long overhead = data instanceof NullableData nd ? (long) nd.validity().length : 0L;
+        return overhead + switch (values) {
+            case byte[] a -> (long) a.length;
+            case short[] a -> 2L * a.length;
+            case int[] a -> 4L * a.length;
+            case long[] a -> 8L * a.length;
+            case float[] a -> 4L * a.length;
+            case double[] a -> 8L * a.length;
+            case boolean[] a -> (long) a.length;
+            case String[] a -> {
+                long total = 0L;
+                for (String s : a) {
+                    // 48 bytes ~ String + char[] object headers plus the array reference slot.
+                    total += 48L + (s == null ? 0L : 2L * s.length());
+                }
+                yield total;
+            }
+            default -> 0L;
+        };
+    }
+
+    /// Demotes the largest-retained global-dict candidate columns to per-chunk encoding, one at a
+    /// time (largest first), until the aggregate retained bytes fall back under the budget. Demoting
+    /// the largest column frees the most memory per eviction, so the fewest columns lose their shared
+    /// dictionary. Called when a chunk pushes the running total over `dictRetainedBudget`.
+    ///
+    /// @throws IOException if writing a flushed segment fails
+    private void evictLargestDictColumnsUntilUnderBudget() throws IOException {
+        while (dictRetainedTotal > dictRetainedBudget && !dictRetainedBytes.isEmpty()) {
+            ColumnName largest = null;
+            long largestBytes = -1L;
+            for (Map.Entry<ColumnName, Long> e : dictRetainedBytes.entrySet()) {
+                if (e.getValue() > largestBytes) {
+                    largestBytes = e.getValue();
+                    largest = e.getKey();
+                }
+            }
+            demoteDictColumn(largest);
+        }
+    }
+
+    /// Abandons the shared global dictionary for one column whose retention pushed the aggregate over
+    /// the memory budget: flushes its already-buffered chunks as ordinary per-chunk segments (so no
+    /// data is lost) and removes it from the candidate set, so subsequent chunks encode per-chunk
+    /// too. The buffered chunks are released for GC and the running retained total is decremented.
+    ///
+    /// @param colName the column being demoted from global-dict to per-chunk encoding
+    /// @throws IOException if writing a flushed segment fails
+    private void demoteDictColumn(ColumnName colName) throws IOException {
+        List<Object> buffered = dictBuffers.remove(colName);
+        dictCandidates.remove(colName);
+        Long freed = dictRetainedBytes.remove(colName);
+        if (freed != null) {
+            dictRetainedTotal -= freed;
+        }
+        if (buffered == null) {
+            return;
+        }
+        DType colDtype = schema.fieldTypes().get(schema.fieldNames().indexOf(colName));
+        for (Object chunk : buffered) {
+            long rowCount = arrayLength(chunk);
+            int segIdx = writeSegment(colDtype, chunk);
+            colChunks.get(colName).add(
+                    new ChunkRef(segIdx, rowCount, lastStatsMin, lastStatsMax, lastStatsSum, lastNullCount));
+        }
+    }
+
     private void flushDictColumns() throws IOException {
         for (ColumnName colName : dictCandidates) {
             List<Object> chunks = dictBuffers.getOrDefault(colName, List.of());

writer/src/test/java/io/github/dfa1/vortex/writer/GlobalDictUtf8Test.java

@@ -216,6 +216,66 @@ void utf8_globalDict_disabled_byOptions(@TempDir Path tmp) throws IOException {
         }
     }
 
+    @Test
+    void retainedBytesBudgetExceeded_utf8_demotesToPerChunkChunkedLayout(@TempDir Path tmp) throws IOException {
+        // Given — a column that looks low-cardinality on its FIRST chunk (3 distinct values, well
+        // under the 50% ratio + 2048 cardinality gates) so it is admitted to the global dict and its
+        // raw data starts being buffered. This is exactly the nyc-311 OOM shape: a "category-ish"
+        // string column whose distinct set is small early but whose raw bytes accumulate across
+        // millions of rows because a global dict must hold every chunk until close(). Rather than
+        // pin the heap, the writer demotes the column once the aggregate retained bytes cross the
+        // budget. We lower the budget via the test seam so this triggers on a handful of small
+        // chunks instead of allocating the full budget. Cardinality never grows, so the OLD
+        // cardinality-fallback would never fire — only the new memory-budget demotion catches this,
+        // which is the bug under test.
+        Path file = tmp.resolve("retained_budget_utf8.vortex");
+        String[] dict = {"open", "closed", "delivered"};
+        int rowsPerChunk = 2_000;
+        int chunkCount = 6;
+        String[] expected = new String[rowsPerChunk * chunkCount];
+
+        try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
+             var sut = VortexWriter.create(ch, SCHEMA, WriteOptions.cascading(3))) {
+            // Budget of ~120 KB is crossed after ~2 chunks of 2000 rows (each row ~48 B object
+            // overhead + a short string), forcing demotion partway through the file.
+            sut.setDictRetainedBudgetForTest(120_000L);
+            // When
+            for (int c = 0; c < chunkCount; c++) {
+                String[] data = new String[rowsPerChunk];
+                for (int i = 0; i < rowsPerChunk; i++) {
+                    String value = dict[(c + i) % dict.length];
+                    data[i] = value;
+                    expected[c * rowsPerChunk + i] = value;
+                }
+                sut.writeChunk(Map.of(ColumnName.of("status"), data));
+            }
+        }
+
+        // Then — the demoted column is a plain Chunked-of-Flats layout with one Flat per chunk (the
+        // buffered-then-flushed chunks plus the per-chunk-encoded remainder), NOT a single Dict
+        // layout. A Dict here would mean the whole column was still buffered in memory — the OOM.
+        try (var vf = VortexReader.open(file, ReadRegistry.loadAll())) {
+            var columnLayout = unwrapZoned(vf.layout().children().getFirst());
+            assertThat(columnLayout.isDict()).as("demoted column must not be a global dict").isFalse();
+            assertThat(columnLayout.isChunked()).as("demoted column is a chunked layout").isTrue();
+            assertThat(columnLayout.children())
+                    .as("one Flat per written chunk after demotion")
+                    .hasSize(chunkCount)
+                    .allSatisfy(child -> assertThat(child.isFlat()).isTrue());
+
+            // And every value round-trips exactly across all chunks despite the mid-file demotion.
+            List<String> got = readAllStrings(vf, "status");
+            assertThat(got).containsExactly(expected);
+        }
+    }
+
+    /// Unwraps a column's [io.github.dfa1.vortex.reader.layout.Layout] Zoned/Stats wrapper (the
+    /// writer wraps every column in one for zone-map pruning) to reach the encoding layout beneath.
+    private static io.github.dfa1.vortex.reader.layout.Layout unwrapZoned(
+            io.github.dfa1.vortex.reader.layout.Layout layout) {
+        return layout.isZoned() ? layout.children().getFirst() : layout;
+    }
+
     /// Reads a nullable Utf8 column, mapping invalid rows to `null` so null positions are asserted
     /// alongside values. A nullable dict column decodes to a [MaskedArray] over the Utf8 payload.
     private static List<String> readNullableStrings(VortexReader vf, String col) {