diff --git a/CHANGELOG.md b/CHANGELOG.md index 35b2f411..8372fcc7 100644 --- a/CHANGELOG.md +++ b/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)) diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java b/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java index 289f926a..58d62fc3 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/VortexWriter.java +++ b/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 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 dictCandidates = new LinkedHashSet<>(); private final Map> dictBuffers = new LinkedHashMap<>(); private final Map 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 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 encoders) { WriteRegistry.Builder b = WriteRegistry.builder(); @@ -496,7 +522,18 @@ public void writeChunk(Map columns) throws IOException { } if (dictCandidates.contains(colName)) { - dictBuffers.computeIfAbsent(colName, _ -> new ArrayList<>()).add(data); + List 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 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 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 chunks = dictBuffers.getOrDefault(colName, List.of()); diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/GlobalDictUtf8Test.java b/writer/src/test/java/io/github/dfa1/vortex/writer/GlobalDictUtf8Test.java index 209bd325..0b7c9f9f 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/GlobalDictUtf8Test.java +++ b/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 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 readNullableStrings(VortexReader vf, String col) {