feat: make writer's global-dict retained-memory budget configurable

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- `WriteOptions.withGlobalDictMaxRetainedBytes(long)` makes the writer's global-dictionary retained-memory budget (default 256 MB) configurable. ([58ba2986](https://github.com/dfa1/vortex-java/commit/58ba2986))
+
 ### 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))

calcite/src/test/java/io/github/dfa1/vortex/calcite/AggregateSumNullTest.java

@@ -49,7 +49,7 @@ private SchemaPlus tableOf(long[] values, boolean[] valid) throws IOException {
                 List.of(ColumnName.of("v")), List.of(new DType.Primitive(PType.I64, true)), false);
         Path file = tmp.resolve("sum-nulls.vortex");
         // Large chunk so the whole column is one chunk; zone maps on so the SUM stat is emitted.
-        WriteOptions opts = new WriteOptions(1024, true, 0.90, 0, false, false);
+        WriteOptions opts = new WriteOptions(1024, true, 0.90, 0, false, false, 256L * 1024 * 1024);
         try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
              var writer = VortexWriter.create(ch, schema, opts)) {
             writer.writeChunk(Map.of(ColumnName.of("v"), new NullableData(values, valid)));

calcite/src/test/java/io/github/dfa1/vortex/calcite/AggregateWhereBoundaryTest.java

@@ -57,7 +57,7 @@ static void write() throws Exception {
                 .build();
         // enableZoneMaps=true emits the per-chunk min/max/sum/null-count the tier-1 fold reads and the
         // classify() step uses to find the boundary zones.
-        WriteOptions opts = new WriteOptions(CHUNK, true, 0.90, 0, true, false);
+        WriteOptions opts = new WriteOptions(CHUNK, true, 0.90, 0, true, false, 256L * 1024 * 1024);
         try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
              VortexWriter writer = VortexWriter.create(ch, schema, opts)) {
             for (int c = 0; c < CHUNKS; c++) {
@@ -456,7 +456,7 @@ private static Ground reduce(java.util.function.LongPredicate predicate) {
     private static void writeChunks(Path file, DType.Struct schema, Map<ColumnName, Object> chunk0,
                                     Map<ColumnName, Object> chunk1) throws Exception {
         // chunkSize large so each writeChunk is exactly one chunk (one zone).
-        WriteOptions opts = new WriteOptions(1024, true, 0.90, 0, false, false);
+        WriteOptions opts = new WriteOptions(1024, true, 0.90, 0, false, false, 256L * 1024 * 1024);
         try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
              VortexWriter writer = VortexWriter.create(ch, schema, opts)) {
             writer.writeChunk(chunk0);

calcite/src/test/java/io/github/dfa1/vortex/calcite/AggregateWhereCleanPartitionTest.java

@@ -56,7 +56,7 @@ static void write() throws Exception {
                 .field("val", DType.I64)
                 .build();
         // enableZoneMaps=true emits the per-chunk min/max/sum/null-count the fold reads.
-        WriteOptions opts = new WriteOptions(CHUNK, true, 0.90, 0, true, false);
+        WriteOptions opts = new WriteOptions(CHUNK, true, 0.90, 0, true, false, 256L * 1024 * 1024);
         try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
              VortexWriter writer = VortexWriter.create(ch, schema, opts)) {
             for (int c = 0; c < CHUNKS; c++) {
@@ -429,7 +429,7 @@ private static Path nullPartitionedFile(String name) throws Exception {
     private static void writeChunks(Path file, DType.Struct schema, Map<ColumnName, Object> chunk0,
                                     Map<ColumnName, Object> chunk1) throws Exception {
         // chunkSize large so each writeChunk is exactly one chunk (one zone).
-        WriteOptions opts = new WriteOptions(1024, true, 0.90, 0, false, false);
+        WriteOptions opts = new WriteOptions(1024, true, 0.90, 0, false, false, 256L * 1024 * 1024);
         try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
              VortexWriter writer = VortexWriter.create(ch, schema, opts)) {
             writer.writeChunk(chunk0);

calcite/src/test/java/io/github/dfa1/vortex/calcite/OhlcGenerator.java

@@ -20,7 +20,7 @@ private OhlcGenerator() {
     }
 
     static void write(Path file, int totalRows, int chunkSize) throws IOException {
-        WriteOptions opts = new WriteOptions(chunkSize, true, 0.90, 0, true, false);
+        WriteOptions opts = new WriteOptions(chunkSize, true, 0.90, 0, true, false, 256L * 1024 * 1024);
         try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
              var writer = VortexWriter.create(ch, OhlcData.SCHEMA, opts)) {
             for (OhlcData.Batch batch : OhlcData.generate(totalRows, chunkSize)) {

calcite/src/test/java/io/github/dfa1/vortex/calcite/UnsignedColumnTest.java

@@ -110,7 +110,7 @@ private Path write(String name, WriteOptions opts, int u32, long u64) throws Exc
     private static WriteOptions noZoneMaps() {
         // Same shape as the adapter coverage test's zone-maps-off options: the second flag disables
         // zone maps so no per-zone SUM exists and VortexAggregates falls back to scanSum.
-        return new WriteOptions(65_536, false, 0.90, 0, true, false);
+        return new WriteOptions(65_536, false, 0.90, 0, true, false, 256L * 1024 * 1024);
     }
 
     private static ReadRegistry registry() {

calcite/src/test/java/io/github/dfa1/vortex/calcite/VortexAdapterCoverageTest.java

@@ -278,7 +278,7 @@ void nonNumericColumn_throws() throws Exception {
         void noZoneMap_sumFallsBackToFullScan(@TempDir Path noStats) throws Exception {
             // Given — a file written with zone maps off, so no per-zone SUM exists to fold
             Path bare = noStats.resolve("nostats.vortex");
-            WriteOptions noZoneMaps = new WriteOptions(65_536, false, 0.90, 0, true, false);
+            WriteOptions noZoneMaps = new WriteOptions(65_536, false, 0.90, 0, true, false, 256L * 1024 * 1024);
             try (var ch = FileChannel.open(bare, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
                  var w = VortexWriter.create(ch, SCHEMA, noZoneMaps)) {
                 w.writeChunk(Map.ofEntries(

integration/src/test/java/io/github/dfa1/vortex/integration/JavaWritesRustReadsIntegrationTest.java

@@ -509,7 +509,7 @@ void javaWriter_jniReader_zoneMapped_multipleZones(@TempDir Path tmp) throws IOE
         // zone-map with one zone per chunk. The Rust reader must parse that layout and still
         // return every value (zones are a transparent pruning aux).
         Path file = tmp.resolve("java_zoned.vtx");
-        WriteOptions zoneMapped = new WriteOptions(4, true, 0.90, 0, true, false);
+        WriteOptions zoneMapped = new WriteOptions(4, true, 0.90, 0, true, false, 256L * 1024 * 1024);
         long[] ids = new long[20];
         double[] vals = new double[20];
         for (int i = 0; i < 20; i++) {

performance/src/main/java/io/github/dfa1/vortex/performance/CalciteBoundaryAggregateBenchmark.java

@@ -194,7 +194,7 @@ private static void writeFixture(Path file) throws IOException {
                 .field("val", DType.I64)
                 .build();
         // enableZoneMaps=true emits the per-chunk min/max/sum/null-count the interior-zone fold reads.
-        WriteOptions opts = new WriteOptions(CHUNK_SIZE, true, 0.90, 0, true, false);
+        WriteOptions opts = new WriteOptions(CHUNK_SIZE, true, 0.90, 0, true, false, 256L * 1024 * 1024);
         java.util.Random rng = new java.util.Random(SEED);
         try (FileChannel ch = FileChannel.open(file,
                 StandardOpenOption.CREATE, StandardOpenOption.WRITE, StandardOpenOption.TRUNCATE_EXISTING);

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

@@ -80,6 +80,12 @@
 ///     writer.writeChunk(Map.of(ColumnName.of("id"), idArray, ColumnName.of("value"), valueArray));
 /// }
 /// ```
+///
+/// With global dictionary encoding enabled (the default), candidate columns are buffered in the heap
+/// until `close()` so a shared dictionary can span every chunk. The aggregate memory this buffering may
+/// retain is bounded by [WriteOptions#globalDictMaxRetainedBytes()] (default 256 MB); once the budget is
+/// crossed the largest columns are demoted to per-chunk encoding, keeping writer memory bounded
+/// regardless of file size or column count.
 public final class VortexWriter implements Closeable {
 
     // Indices into layout_specs list in the FbsFooter
@@ -99,18 +105,6 @@ 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(),
@@ -141,9 +135,9 @@ public final class VortexWriter implements Closeable {
     // 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;
+    // Effective aggregate global-dict retention budget, configured via
+    // WriteOptions.globalDictMaxRetainedBytes(); see that field's javadoc for the rationale.
+    private final long dictRetainedBudget;
     private boolean firstChunkSeen = false;
 
     // Per-column zone-maps, populated by flushZoneMaps() in close() when enableZoneMaps is set.
@@ -173,6 +167,7 @@ private VortexWriter(
         this.channel = channel;
         this.schema = schema;
         this.options = options;
+        this.dictRetainedBudget = options.globalDictMaxRetainedBytes();
         this.encodings = encodings;
         this.defaultRegistry = buildRegistry(encodings);
         this.cascadeCodecs = buildCascadeCodecs(options);
@@ -182,12 +177,6 @@ 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();

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

@@ -13,19 +13,34 @@
 ///                                  file size by 10–15% on real-world datasets compared to ALP+bitpack alone.
 ///                                  Trade-off: Zstd decompression is ~6× slower than ALP decode;
 ///                                  prefer the default (`false`) for read-heavy workloads.
+/// @param globalDictMaxRetainedBytes aggregate byte budget for the raw data all global-dictionary
+///                                  candidate columns may buffer in the heap while waiting for
+///                                  `close()` (default 256 MB). A shared dictionary must see every
+///                                  chunk before it can be built, so a candidate column's raw arrays
+///                                  are held from its first chunk until the file is finished; when the
+///                                  running total across all such columns crosses this budget the
+///                                  largest-retained columns are demoted to per-chunk encoding until
+///                                  back under it, bounding writer memory regardless of file size or
+///                                  column count. Trade-off: demoted columns lose the shared-dictionary
+///                                  size benefit; raise the budget on memory-rich hosts to keep more
+///                                  wide, mis-detected columns dictionary-encoded.
 public record WriteOptions(
         int chunkSize,
         boolean enableZoneMaps,
         double compressionRatioThreshold,
         int allowedCascading,
         boolean globalDict,
-        boolean enableZstd
+        boolean enableZstd,
+        long globalDictMaxRetainedBytes
 ) {
+    /// Default aggregate retention budget (256 MB) for buffered global-dictionary candidate columns.
+    private static final long DEFAULT_GLOBAL_DICT_MAX_RETAINED_BYTES = 256L * 1024 * 1024;
+
     /// Default options: global dictionary encoding enabled, no cascading compression, Zstd disabled.
     ///
     /// @return default `WriteOptions`
     public static WriteOptions defaults() {
-        return new WriteOptions(65_536, true, 0.90, 0, true, false);
+        return new WriteOptions(65_536, true, 0.90, 0, true, false, DEFAULT_GLOBAL_DICT_MAX_RETAINED_BYTES);
     }
 
     /// Enable cascading compression with up to `depth` recursive levels.
@@ -34,23 +49,25 @@ public static WriteOptions defaults() {
     /// @param depth maximum cascade depth
     /// @return `WriteOptions` with cascading enabled at the given depth
     public static WriteOptions cascading(int depth) {
-        return new WriteOptions(65_536, true, 0.90, depth, true, false);
+        return new WriteOptions(65_536, true, 0.90, depth, true, false, DEFAULT_GLOBAL_DICT_MAX_RETAINED_BYTES);
     }
 
     /// Returns a copy of these options with zone-map statistics set to `enabled`.
     ///
     /// @param enabled `true` to write per-chunk min/max/sum statistics for zone-map pruning
     /// @return a new `WriteOptions` with the zone-map flag updated
     public WriteOptions withZoneMaps(boolean enabled) {
-        return new WriteOptions(chunkSize, enabled, compressionRatioThreshold, allowedCascading, globalDict, enableZstd);
+        return new WriteOptions(chunkSize, enabled, compressionRatioThreshold, allowedCascading, globalDict, enableZstd,
+                globalDictMaxRetainedBytes);
     }
 
     /// Returns a copy of these options with global dictionary encoding set to `enabled`.
     ///
     /// @param enabled `true` to enable global dictionary encoding across chunks
     /// @return a new `WriteOptions` with the global dict flag updated
     public WriteOptions withGlobalDict(boolean enabled) {
-        return new WriteOptions(chunkSize, enableZoneMaps, compressionRatioThreshold, allowedCascading, enabled, enableZstd);
+        return new WriteOptions(chunkSize, enableZoneMaps, compressionRatioThreshold, allowedCascading, enabled, enableZstd,
+                globalDictMaxRetainedBytes);
     }
 
     /// Returns a copy of these options with Zstandard compression set to `enabled`.
@@ -65,6 +82,21 @@ public WriteOptions withGlobalDict(boolean enabled) {
     /// @param enabled `true` to enable Zstd in the compression cascade
     /// @return a new `WriteOptions` with the Zstd flag updated
     public WriteOptions withZstd(boolean enabled) {
-        return new WriteOptions(chunkSize, enableZoneMaps, compressionRatioThreshold, allowedCascading, globalDict, enabled);
+        return new WriteOptions(chunkSize, enableZoneMaps, compressionRatioThreshold, allowedCascading, globalDict, enabled,
+                globalDictMaxRetainedBytes);
+    }
+
+    /// Returns a copy of these options with the global-dictionary retention budget set to `budgetBytes`.
+    ///
+    /// This is the aggregate byte budget across all global-dictionary candidate columns buffered in the
+    /// heap while the writer waits to build shared dictionaries at `close()`. Lower it to demote
+    /// mis-detected wide columns to per-chunk encoding sooner (bounding memory on huge files); raise it
+    /// on memory-rich hosts to keep more columns dictionary-encoded.
+    ///
+    /// @param budgetBytes aggregate retention budget in bytes for buffered global-dict candidate columns
+    /// @return a new `WriteOptions` with the global-dict retention budget updated
+    public WriteOptions withGlobalDictMaxRetainedBytes(long budgetBytes) {
+        return new WriteOptions(chunkSize, enableZoneMaps, compressionRatioThreshold, allowedCascading, globalDict,
+                enableZstd, budgetBytes);
     }
 }

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

@@ -224,7 +224,7 @@ void retainedBytesBudgetExceeded_utf8_demotesToPerChunkChunkedLayout(@TempDir Pa
         // 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
+        // budget. We lower the budget via WriteOptions 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.
@@ -234,11 +234,12 @@ void retainedBytesBudgetExceeded_utf8_demotesToPerChunkChunkedLayout(@TempDir Pa
         int chunkCount = 6;
         String[] expected = new String[rowsPerChunk * chunkCount];
 
+        // 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. Configured via the
+        // real public WriteOptions surface rather than a test-only seam.
+        WriteOptions opts = WriteOptions.cascading(3).withGlobalDictMaxRetainedBytes(120_000L);
         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);
+             var sut = VortexWriter.create(ch, SCHEMA, opts)) {
             // When
             for (int c = 0; c < chunkCount; c++) {
                 String[] data = new String[rowsPerChunk];

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

@@ -33,7 +33,7 @@ class NullCountPruningTest {
     // sizes and null patterns: 3 rows / 0 nulls, 2 rows / 1 null, 4 rows / all null.
     private Path write() throws IOException {
         Path file = tmp.resolve("nulls.vtx");
-        WriteOptions opts = new WriteOptions(1024, true, 0.90, 0, false, false);
+        WriteOptions opts = new WriteOptions(1024, true, 0.90, 0, false, false, 256L * 1024 * 1024);
         try (var ch = FileChannel.open(file, StandardOpenOption.CREATE, StandardOpenOption.WRITE);
              var sut = VortexWriter.create(ch, SCHEMA, opts)) {
             sut.writeChunk(Map.of(ColumnName.of("v"), new NullableData(