review: ground-truth mixed-grid test + evict passed coarse flats (#221)

Addresses PR #223 review findings.

1. Ground-truth value test for the coarse-sub-chunk slice path. Neither writer
   can EMIT a disjoint per-column chunk grid — the Java VortexWriter's writeChunk
   requires all columns to agree on row count, and the JNI writer writes uniform
   row batches — so a genuine mixed grid ([3,3,2] vs [4,4]) is hand-assembled at
   the file-format level (reusing VortexReaderDecodeChunkTest's raw-file seam) and
   the full streaming scan is asserted value-for-value, including "a"'s middle
   chunk sliced at a nonzero offset and a nullable column crossing the MaskedArray
   slice. Rust-parity stays covered by RaincloudConformanceIntegrationTest.

2. Evict cached coarse flats once passed. Each coarse covering flat now decodes
   into its own confined arena; as the scan advances, a column whose covering flat
   changes from the previous decoded window releases the earlier flat's arena
   (the planner's per-column cursor is monotonic, so it is never revisited). The
   previous Chunk is already closed at that point, so no live slice references the
   freed arena. Bounds peak off-heap memory for wide disjoint-grid files instead
   of retaining every decoded chunk for the iterator's lifetime. The new streaming
   test iterates fully, proving later windows still decode correctly after earlier
   chunks are evicted.

3. Reword the columnZoneStats fallback javadoc: entries are now per scan window,
   with a coarse chunk's stats repeated (conservatively) across its sub-windows.

4. ScanIteratorChunkGridTest states column order via a LinkedHashMap helper rather
   than relying on Map.of iteration order.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: dfa1

reader/src/main/java/io/github/dfa1/vortex/reader/ScanIterator.java

@@ -87,8 +87,20 @@ public final class ScanIterator implements Iterator<Chunk>, AutoCloseable {
     private long rowsReturned;
     private Chunk openChunk;
     private boolean closed;
-    private Arena sharedArena;
-    private Map<Layout, Array> sharedFlats;
+    // Coarse chunks that span several split windows are decoded once and cached here (keyed by
+    // layout identity), each in its own confined arena so it can be released the moment the scan
+    // advances past its last covering window — see evictPassedFlats.
+    private Map<Layout, CachedFlat> sharedFlats;
+    // Per projected column, the covering flat of the most recently decoded window; drives eviction.
+    private Layout[] lastCoveringFlats;
+
+    /// A coarse covering flat decoded once and sliced across the windows it spans, together with
+    /// the confined [Arena] that owns its buffers (closed on eviction / iterator close).
+    ///
+    /// @param arena the confined arena owning this flat's decoded buffers
+    /// @param array the decoded flat, sliced per window
+    private record CachedFlat(Arena arena, Array array) {
+    }
 
     public ScanIterator(VortexHandle file, ScanOptions options) {
         this.file = file;
@@ -295,6 +307,8 @@ public Chunk next() {
         chunkIndex = peekedChunkIdx + 1;
         peekedChunkIdx = -1;
 
+        evictPassedFlats(spec);
+
         long remaining = options.limit() - rowsReturned;
         long chunkRows = Math.min(spec.rowCount(), remaining);
 
@@ -390,11 +404,14 @@ public long[] chunkRowCounts() {
     /// `null`, since the flat writer does not retain it). Either way a column that is absent
     /// or carries no stats yields [ArrayStats#empty()] per zone.
     ///
-    /// Zone granularity is the layout's, not the scan's. The fallback path is one entry per
-    /// chunk, positionally aligned with [#chunkRowCounts()]. The zone-map path is one entry per
-    /// zone of the stats table: this writer emits one zone per chunk (so the same alignment
-    /// holds), but a file from another writer may use a fixed zone length independent of chunk
-    /// boundaries, in which case the zone count need not match [#chunkRowCounts()].
+    /// Zone granularity is the layout's, not the scan's. The fallback path is one entry per scan
+    /// window, positionally aligned with [#chunkRowCounts()]. Each entry carries the stats of the
+    /// covering chunk of that window, so when columns chunk on different grids a coarse chunk's
+    /// stats repeat across every sub-window it spans (still a conservative min/max/null-count bound
+    /// for each sub-window). The zone-map path is one entry per zone of the stats table: this writer
+    /// emits one zone per chunk (so the same alignment holds), but a file from another writer may
+    /// use a fixed zone length independent of chunk boundaries, in which case the zone count need
+    /// not match [#chunkRowCounts()].
     ///
     /// This is the read-side surface for aggregate push-down (ADR 0013 §6): a reduction can
     /// fold whole zones from these rows and fall back to a streaming decode only for the
@@ -564,9 +581,10 @@ public void close() {
             openChunk.close();
         }
         openChunk = null;
-        if (sharedArena != null) {
-            sharedArena.close();
-            sharedArena = null;
+        if (sharedFlats != null) {
+            for (CachedFlat cached : sharedFlats.values()) {
+                cached.arena().close();
+            }
             sharedFlats = null;
         }
     }
@@ -608,6 +626,7 @@ private void initialize() {
 
         projectedNames = List.copyOf(columnDtypes.keySet());
         projectedDtypes = List.copyOf(columnDtypes.values());
+        lastCoveringFlats = new Layout[projectedNames.size()];
         chunks = buildChunks(columnFlats);
     }
 
@@ -682,8 +701,8 @@ private Array decodeOrSlice(int colIdx, Layout coveringFlat, long sliceOffset, l
             // no slice wrapper (aligned N-vs-N and per-chunk decode keep their zero-copy behavior).
             return decodeLayout(coveringFlat, dtype, arena);
         }
-        // Covering chunk spans several windows: decode it once into the shared arena and slice
-        // zero-copy per window, matching how Rust decodes a chunk once and slices each split range.
+        // Covering chunk spans several windows: decode it once (cached) and slice zero-copy per
+        // window, matching how Rust decodes a chunk once and slices each split range.
         Array full = sharedCoveringFlat(coveringFlat, dtype);
         return sliceArray(full, sliceOffset, rowCount, dtype);
     }
@@ -692,14 +711,37 @@ private static boolean isWholeFlat(Layout coveringFlat, long sliceOffset, long r
         return sliceOffset == 0 && rowCount == coveringFlat.rowCount();
     }
 
-    /// Decodes `flat` once into the iterator's shared arena and caches it by identity, so a coarse
-    /// chunk that covers several split windows is decoded a single time and then sliced per window.
+    /// Decodes `flat` once into its own confined arena and caches it by identity, so a coarse chunk
+    /// that covers several split windows is decoded a single time and then sliced per window. The
+    /// arena is released by [#evictPassedFlats] once the scan advances off the flat.
     private Array sharedCoveringFlat(Layout flat, DType dtype) {
-        if (sharedArena == null) {
-            sharedArena = Arena.ofConfined();
+        if (sharedFlats == null) {
             sharedFlats = new IdentityHashMap<>();
         }
-        return sharedFlats.computeIfAbsent(flat, f -> decodeLayout(f, dtype, sharedArena));
+        return sharedFlats.computeIfAbsent(flat, f -> {
+            Arena flatArena = Arena.ofConfined();
+            return new CachedFlat(flatArena, decodeLayout(f, dtype, flatArena));
+        }).array();
+    }
+
+    /// Releases every cached coarse flat the scan has moved past. As `spec` is decoded, a column
+    /// whose covering flat differs from the one it used in the previously decoded window can never
+    /// revisit that earlier flat — the planner's per-column cursor advances monotonically — so its
+    /// arena is closed and its cache entry dropped. Whole-window (uncached) flats are ignored. The
+    /// previously open [Chunk] is already closed here (enforced by [#next()]), so no live slice
+    /// still references the freed arena.
+    private void evictPassedFlats(ChunkSpec spec) {
+        Layout[] current = spec.columnLayouts();
+        for (int j = 0; j < current.length; j++) {
+            Layout previous = lastCoveringFlats[j];
+            if (previous != null && previous != current[j]) {
+                CachedFlat cached = sharedFlats == null ? null : sharedFlats.remove(previous);
+                if (cached != null) {
+                    cached.arena().close();
+                }
+            }
+            lastCoveringFlats[j] = current[j];
+        }
     }
 
     private static Array sliceArray(Array full, long offset, long length, DType dtype) {

reader/src/test/java/io/github/dfa1/vortex/reader/ScanIteratorChunkGridTest.java

@@ -26,7 +26,7 @@ class ScanIteratorChunkGridTest {
     @Test
     void alignedGridsDecodeEveryColumnAsAWholeChunk() {
         // Given two columns with the *same* grid [4, 4] over 8 rows — the aligned N-vs-N fast path.
-        var columnFlats = flats(Map.of(A, new long[]{4, 4}, B, new long[]{4, 4}));
+        var columnFlats = flats(A, new long[]{4, 4}, B, new long[]{4, 4});
 
         // When
         List<ChunkSpec> result = ScanIterator.buildChunks(columnFlats);
@@ -43,7 +43,7 @@ void alignedGridsDecodeEveryColumnAsAWholeChunk() {
     @Test
     void singleFlatColumnSharesTheChunkedGrid() {
         // Given one full-column flat [8] beside a chunked column [4, 4] — the 1-vs-N case.
-        var columnFlats = flats(Map.of(A, new long[]{8}, B, new long[]{4, 4}));
+        var columnFlats = flats(A, new long[]{8}, B, new long[]{4, 4});
 
         // When
         List<ChunkSpec> result = ScanIterator.buildChunks(columnFlats);
@@ -62,9 +62,9 @@ void nestedBoundariesSliceTheCoarseColumnAtTheFineGrid() {
         // Given the emotions-dataset-for-nlp shape scaled down: a coarse column [8, 8, 8, 4] (like
         // `label`'s [131072 ×3, 23593]) beside a fine column of 2-row chunks (like `text`'s 16384
         // grid, where 8 = 4 × 2 so every coarse boundary is also a fine boundary).
-        var columnFlats = flats(Map.of(
+        var columnFlats = flats(
                 A, new long[]{8, 8, 8, 4},
-                B, new long[]{2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2}));
+                B, new long[]{2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2});
 
         // When
         List<ChunkSpec> result = ScanIterator.buildChunks(columnFlats);
@@ -94,9 +94,9 @@ void disjointBoundariesSpliceAtTheMergedGrid() {
         // Given the uci-beijing-multi-site-air-quality shape scaled down: [3, 3, 2] vs [4, 4] over 8
         // rows. The boundaries do not nest (3 + 3 = 6 != 4, mirroring 40960 + 49152 != 131072), so
         // neither grid refines the other — the scan must emit windows at the union of boundaries.
-        var columnFlats = flats(Map.of(
+        var columnFlats = flats(
                 A, new long[]{3, 3, 2},
-                B, new long[]{4, 4}));
+                B, new long[]{4, 4});
 
         // When
         List<ChunkSpec> result = ScanIterator.buildChunks(columnFlats);
@@ -134,18 +134,24 @@ void emptyProjectionYieldsNoChunks() {
 
     // ── helpers ───────────────────────────────────────────────────────────────
 
-    private static Map<ColumnName, List<Layout>> flats(Map<ColumnName, long[]> grids) {
+    /// Builds an insertion-ordered `column -> flats` map for two columns. A [LinkedHashMap] states
+    /// the intended column order explicitly (the planner's output is order-independent and read
+    /// back by name here, but an unordered [Map#of] would hide that intent).
+    private static Map<ColumnName, List<Layout>> flats(ColumnName n1, long[] g1, ColumnName n2, long[] g2) {
         var out = new LinkedHashMap<ColumnName, List<Layout>>();
-        for (var entry : grids.entrySet()) {
-            var list = new ArrayList<Layout>();
-            for (long rows : entry.getValue()) {
-                list.add(new Layout(LayoutId.FLAT, rows, null, List.of(), List.of()));
-            }
-            out.put(entry.getKey(), list);
-        }
+        out.put(n1, toFlats(g1));
+        out.put(n2, toFlats(g2));
         return out;
     }
 
+    private static List<Layout> toFlats(long[] chunkRows) {
+        var list = new ArrayList<Layout>(chunkRows.length);
+        for (long rows : chunkRows) {
+            list.add(new Layout(LayoutId.FLAT, rows, null, List.of(), List.of()));
+        }
+        return list;
+    }
+
     private static List<Long> rowCounts(List<ChunkSpec> chunks) {
         var out = new ArrayList<Long>(chunks.size());
         for (ChunkSpec spec : chunks) {

reader/src/test/java/io/github/dfa1/vortex/reader/VortexReaderDecodeChunkTest.java

@@ -262,6 +262,54 @@ void decodeChunk_sharedSingleFlatColumn_isValueEquivalentAndSurvivesIteratorClos
         }
     }
 
+    // A genuinely DISJOINT per-column chunk grid (issue #221, "beijing" tier): column "a" chunks
+    // [3, 3, 2] and column "b" chunks [4, 4] over the same 8 rows. The boundaries do not nest
+    // ({0,3,6,8} vs {0,4,8}), so neither grid refines the other; the scan must split at the union
+    // {0,3,4,6,8}. "b" is nullable so the coarse-chunk slice also crosses the MaskedArray branch.
+    private static final long[] MIXED_A_ROWS = {3, 3, 2};
+    private static final long[][] MIXED_A = {{10, 11, 12}, {13, 14, 15}, {16, 17}};
+    private static final long[] MIXED_B_ROWS = {4, 4};
+    private static final long[][] MIXED_B = {{20, 21, 22, 23}, {24, 25, 26, 27}};
+    private static final boolean[][] MIXED_B_VALID = {{true, false, true, true}, {true, true, false, true}};
+    private static final List<Long> MIXED_A_EXPECTED = List.of(10L, 11L, 12L, 13L, 14L, 15L, 16L, 17L);
+    private static final List<Long> MIXED_B_EXPECTED =
+            java.util.Arrays.asList(20L, null, 22L, 23L, 24L, 25L, null, 27L);
+
+    @Test
+    void scan_disjointMixedGrid_streamsExactValuesAcrossMergedWindows(@TempDir Path tmp) throws Exception {
+        // This is the value-level ground truth for #221's coarse-sub-chunk slice path. Neither the
+        // Java VortexWriter (its writeChunk requires all columns to agree on row count) nor the JNI
+        // writer (it writes uniform row batches) can EMIT a disjoint per-column grid, so the file is
+        // hand-assembled here; Rust-parity of the fix is covered separately by the size-gated
+        // RaincloudConformanceIntegrationTest (uci-beijing-multi-site-air-quality). Streaming the
+        // whole scan exercises the merged-grid planner end to end AND, by fully iterating, drives
+        // the coarse-flat eviction path (earlier chunks are released while later windows must still
+        // decode correctly). The ground truth is the exact values written into each segment.
+
+        // Given a file whose two columns chunk on disjoint grids [3,3,2] vs [4,4]
+        Path file = writeMixedGridFile(tmp);
+        var windowRowCounts = new ArrayList<Long>();
+        var streamedA = new ArrayList<Long>();
+        var streamedB = new ArrayList<Long>();
+
+        // When streaming the whole scan
+        try (var reader = VortexReader.open(file, registry());
+                var iter = reader.scan(ScanOptions.all())) {
+            iter.forEachRemaining(chunk -> {
+                windowRowCounts.add(chunk.rowCount());
+                streamedA.addAll(values(chunk.column("a")));
+                streamedB.addAll(values(chunk.column("b")));
+            });
+        }
+
+        // Then the scan splits at the merged grid {0,3,4,6,8} -> windows of 3,1,2,2 rows, and every
+        // value (and null) survives the coarse-chunk sub-window slice — including "a"'s middle chunk
+        // sliced at a nonzero offset for window [4,6) and "b"'s first chunk sliced for [0,3)+[3,4).
+        assertThat(windowRowCounts).containsExactly(3L, 1L, 2L, 2L);
+        assertThat(streamedA).isEqualTo(MIXED_A_EXPECTED);
+        assertThat(streamedB).isEqualTo(MIXED_B_EXPECTED);
+    }
+
     // ── Helpers ────────────────────────────────────────────────────────────────
 
     private static ReadRegistry registry() {
@@ -302,6 +350,76 @@ private static Path writeMultiChunkFile(Path dir) throws Exception {
         return writeFile(dir, "multi.vtx", CHUNK_ROWS, COL_A, COL_B, COL_B_VALID);
     }
 
+    /// Assembles a struct file whose columns chunk on independent grids: "a" over [MIXED_A_ROWS]
+    /// (non-nullable) and "b" over [MIXED_B_ROWS] (nullable). Segment order: a-chunk0..a-chunkN,
+    /// then b-chunk0..b-chunkM. Layout: flat=0, chunked=1, struct=2; array specs: primitive=0,
+    /// bool=1.
+    private static Path writeMixedGridFile(Path dir) throws Exception {
+        List<byte[]> segments = new ArrayList<>();
+        for (long[] chunk : MIXED_A) {
+            segments.add(primitiveSegment(chunk, null));
+        }
+        int bBase = segments.size();
+        for (int k = 0; k < MIXED_B.length; k++) {
+            segments.add(primitiveSegment(MIXED_B[k], MIXED_B_VALID[k]));
+        }
+
+        long[] segOffsets = new long[segments.size()];
+        long[] segLengths = new long[segments.size()];
+        long off = 0;
+        for (int i = 0; i < segments.size(); i++) {
+            segOffsets[i] = off;
+            segLengths[i] = segments.get(i).length;
+            off += segments.get(i).length;
+        }
+
+        int[] aSeg = new int[MIXED_A_ROWS.length];
+        for (int k = 0; k < aSeg.length; k++) {
+            aSeg[k] = k;
+        }
+        int[] bSeg = new int[MIXED_B_ROWS.length];
+        for (int k = 0; k < bSeg.length; k++) {
+            bSeg[k] = bBase + k;
+        }
+
+        ByteBuffer footerBuf = MalformedFiles.buildFooter(
+                new String[]{"vortex.primitive", "vortex.bool"},
+                new String[]{"vortex.flat", "vortex.chunked", "vortex.struct"},
+                segOffsets, segLengths);
+        ByteBuffer dtypeBuf = buildStructDtype();
+        ByteBuffer layoutBuf = buildMixedLayout(aSeg, MIXED_A_ROWS, bSeg, MIXED_B_ROWS);
+
+        return writeVtxFile(dir, "mixed.vtx", segments, footerBuf, dtypeBuf, layoutBuf);
+    }
+
+    /// Layout for the disjoint-grid file: two chunked columns whose flats carry independent
+    /// row counts, so the reader's planner must split at the union of both grids.
+    private static ByteBuffer buildMixedLayout(int[] aSeg, long[] aRows, int[] bSeg, long[] bRows) {
+        var fbb = new FbsBuilder(1024);
+        long total = 0;
+        for (long r : aRows) {
+            total += r;
+        }
+        int chunkedA = buildChunkedColumn(fbb, aSeg, aRows, total);
+        int chunkedB = buildChunkedColumn(fbb, bSeg, bRows, total);
+        int structChildV = FbsLayout.createChildrenVector(fbb, new int[]{chunkedA, chunkedB});
+        int structOff = FbsLayout.createFbsLayout(fbb, 2, total, 0, structChildV, 0);
+        FbsLayout.finishFbsLayoutBuffer(fbb, structOff);
+        return MalformedFiles.slice(fbb);
+    }
+
+    /// Builds one `vortex.chunked` layout node whose children are one `vortex.flat` per chunk,
+    /// each flat carrying its own `rows[k]` row count and single segment index `segIdx[k]`.
+    private static int buildChunkedColumn(FbsBuilder fbb, int[] segIdx, long[] rows, long total) {
+        int[] flatOffs = new int[rows.length];
+        for (int k = 0; k < rows.length; k++) {
+            int segV = FbsLayout.createSegmentsVector(fbb, new long[]{segIdx[k]});
+            flatOffs[k] = FbsLayout.createFbsLayout(fbb, 0, rows[k], 0, 0, segV);
+        }
+        int childV = FbsLayout.createChildrenVector(fbb, flatOffs);
+        return FbsLayout.createFbsLayout(fbb, 1, total, 0, childV, 0);
+    }
+
     /// Streams the whole scan of the shared-column file, collecting the per-chunk column "a" and
     /// the per-chunk slice of the single-flat shared column "c" so a random-access decode can be
     /// compared element-by-element.