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

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) {