perf(reader): fold an i32 global dict column off its codes, not a per-row copy

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Performance
+- Scanning an integer column with a dict layout folds its codes against the dictionary instead of expanding every row first, about 16% faster on a two-column taxi scan. ([#603](https://github.com/dfa1/vortex-java/pull/603))
 - Folding dict-encoded float columns and scan windows reads their buffers directly, about 3% faster on a two-column taxi scan. ([#600](https://github.com/dfa1/vortex-java/pull/600))
 
 ### Added

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

@@ -14,6 +14,7 @@
 import io.github.dfa1.vortex.reader.array.BoolArray;
 import io.github.dfa1.vortex.reader.array.ByteArray;
 import io.github.dfa1.vortex.reader.array.CanonicalArrays;
+import io.github.dfa1.vortex.reader.array.DictIntArray;
 import io.github.dfa1.vortex.reader.array.DoubleArray;
 import io.github.dfa1.vortex.reader.array.FixedSizeListArray;
 import io.github.dfa1.vortex.reader.array.Float16Array;
@@ -975,10 +976,22 @@ private Array sharedCoveringFlat(Layout flat, DType dtype) {
         // slice then reads a flat buffer instead of re-resolving lazy composite children per row.
         return sharedFlats.computeIfAbsent(flat, f -> {
             Arena flatArena = Arena.ofConfined();
-            return new CachedFlat(flatArena, CanonicalArrays.of(decodeLayout(f, dtype, flatArena), flatArena));
+            return new CachedFlat(flatArena, canonicalShared(decodeLayout(f, dtype, flatArena), flatArena));
         }).array();
     }
 
+    /// Canonicalizes a shared flat, except that an `i32` dict stays a dict over flat codes and a
+    /// flat pool: each window slices just its codes, and a fold gathers off the pool, where
+    /// expanding it first wrote and re-read a 4-byte value for every row of the flat.
+    private static Array canonicalShared(Array array, Arena arena) {
+        return switch (array) {
+            case MaskedArray m -> new MaskedArray(canonicalShared(m.inner(), arena), m.validity());
+            case DictIntArray d -> DictIntArray.of(d.dtype(), d.length(),
+                    (IntArray) CanonicalArrays.of(d.values(), arena), CanonicalArrays.of(d.codes(), arena));
+            default -> CanonicalArrays.of(array, arena);
+        };
+    }
+
     /// 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
@@ -1010,6 +1023,9 @@ static Array sliceArray(Array full, long offset, long length, DType dtype) {
                 yield new MaskedArray(innerSlice, validitySlice);
             }
             case LongArray a -> new OffsetLongArray(dtype, length, a, offset);
+            // A dict window shares the pool and slices only its codes.
+            case DictIntArray d -> DictIntArray.of(dtype, length, d.values(),
+                    sliceArray(d.codes(), offset, length, d.codes().dtype()));
             case IntArray a -> new OffsetIntArray(dtype, length, a, offset);
             case DoubleArray a -> new OffsetDoubleArray(dtype, length, a, offset);
             case FloatArray a -> new OffsetFloatArray(dtype, length, a, offset);

reader/src/main/java/io/github/dfa1/vortex/reader/array/DictIntArray.java

@@ -6,6 +6,7 @@
 
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SegmentAllocator;
+import java.lang.foreign.ValueLayout;
 import java.util.function.IntBinaryOperator;
 import java.util.function.IntConsumer;
 
@@ -61,13 +62,21 @@ public MemorySegment materialize(SegmentAllocator arena) {
         MemorySegment dst = arena.allocate(n * 4L, 4);
         IntArray vals = values;
         Array flatCodes = CanonicalArrays.of(codes, arena);
-        // A flat pool and flat u16 codes, as a global dict column has them: gather off the two
-        // segments, skipping a virtual getShort and getInt per row.
+        // A flat pool and flat u8/u16 codes, as a global dict column has them: gather off the two
+        // segments, skipping a virtual getShort/getByte and getInt per row.
         MemorySegment pool = DictArrays.flat(vals, vals.length(), 4);
-        MemorySegment c = flatCodes instanceof ShortArray ? DictArrays.flat(flatCodes, n, 2) : null;
-        if (pool != null && c != null) {
+        MemorySegment c16 = flatCodes instanceof ShortArray ? DictArrays.flat(flatCodes, n, 2) : null;
+        if (pool != null && c16 != null) {
             for (long i = 0; i < n; i++) {
-                long code = Short.toUnsignedLong(c.getAtIndex(VortexFormat.LE_SHORT, i));
+                long code = Short.toUnsignedLong(c16.getAtIndex(VortexFormat.LE_SHORT, i));
+                dst.setAtIndex(VortexFormat.LE_INT, i, pool.getAtIndex(VortexFormat.LE_INT, code));
+            }
+            return dst.asReadOnly();
+        }
+        MemorySegment c8 = flatCodes instanceof ByteArray ? DictArrays.flat(flatCodes, n, 1) : null;
+        if (pool != null && c8 != null) {
+            for (long i = 0; i < n; i++) {
+                long code = Byte.toUnsignedLong(c8.get(ValueLayout.JAVA_BYTE, i));
                 dst.setAtIndex(VortexFormat.LE_INT, i, pool.getAtIndex(VortexFormat.LE_INT, code));
             }
             return dst.asReadOnly();
@@ -129,11 +138,36 @@ public void forEachInt(IntConsumer cons) {
         }
     }
 
+    /// Folds the gathered values in row order. With a flat pool and flat `u16` codes, as a scan
+    /// window of a global dict column has them, it gathers straight off the two segments instead
+    /// of through a virtual `getShort` and `getInt` per row. Codes were range-checked against the
+    /// pool when decoded.
+    ///
+    /// @param identity initial accumulator value
+    /// @param op       binary operator applied to the accumulator and each gathered value
+    /// @return the final accumulated result
     @Override
     public int fold(int identity, IntBinaryOperator op) {
         long n = length;
         IntArray vals = values;
         int result = identity;
+        MemorySegment pool = DictArrays.flat(vals, vals.length(), 4);
+        MemorySegment c16 = codes instanceof ShortArray ? DictArrays.flat(codes, n, 2) : null;
+        if (pool != null && c16 != null) {
+            for (long i = 0; i < n; i++) {
+                long code = Short.toUnsignedLong(c16.getAtIndex(VortexFormat.LE_SHORT, i));
+                result = op.applyAsInt(result, pool.getAtIndex(VortexFormat.LE_INT, code));
+            }
+            return result;
+        }
+        MemorySegment c8 = codes instanceof ByteArray ? DictArrays.flat(codes, n, 1) : null;
+        if (pool != null && c8 != null) {
+            for (long i = 0; i < n; i++) {
+                long code = Byte.toUnsignedLong(c8.get(ValueLayout.JAVA_BYTE, i));
+                result = op.applyAsInt(result, pool.getAtIndex(VortexFormat.LE_INT, code));
+            }
+            return result;
+        }
         switch (codes) {
             case ByteArray ba -> {
                 for (long i = 0; i < n; i++) {

reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetByteArray.java

@@ -4,6 +4,7 @@
 
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SegmentAllocator;
+import java.util.Optional;
 
 import java.util.function.LongBinaryOperator;
 
@@ -35,6 +36,15 @@ public long fold(long identity, LongBinaryOperator op) {
         return acc;
     }
 
+    /// The window of a flat inner buffer (see `OffsetArrays#window`), empty otherwise: a scan
+    /// window over a shared dict's codes then reads as flat codes, which the gathers read directly.
+    ///
+    /// @return the window's segment, if the inner buffer covers it
+    @Override
+    public Optional<MemorySegment> segmentIfPresent() {
+        return OffsetArrays.window(inner, offset, length, 1);
+    }
+
     /// Zero-copy over a flat inner buffer (see `OffsetArrays#window`), else the per-row default.
     ///
     /// @param arena allocator for the per-row fallback

reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetIntArray.java

@@ -42,6 +42,15 @@ public int fold(int identity, IntBinaryOperator op) {
         return result;
     }
 
+    /// The window of a flat inner buffer (see `OffsetArrays#window`), empty otherwise: a scan
+    /// window over a shared dict's codes then reads as flat codes, which the gathers read directly.
+    ///
+    /// @return the window's segment, if the inner buffer covers it
+    @Override
+    public Optional<MemorySegment> segmentIfPresent() {
+        return OffsetArrays.window(inner, offset, length, 4);
+    }
+
     /// Zero-copy over a flat inner buffer (see `OffsetArrays#window`), else the per-row default.
     ///
     /// @param arena allocator for the per-row fallback

reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetLongArray.java

@@ -4,6 +4,7 @@
 
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SegmentAllocator;
+import java.util.Optional;
 
 /// Sliced view over a [LongArray]: `getLong(i) = inner.getLong(i + offset)`.
 ///
@@ -24,6 +25,15 @@ public long getLong(long i) {
         return inner.getLong(i + offset);
     }
 
+    /// The window of a flat inner buffer (see `OffsetArrays#window`), empty otherwise: a scan
+    /// window over a shared dict's codes then reads as flat codes, which the gathers read directly.
+    ///
+    /// @return the window's segment, if the inner buffer covers it
+    @Override
+    public Optional<MemorySegment> segmentIfPresent() {
+        return OffsetArrays.window(inner, offset, length, 8);
+    }
+
     /// Zero-copy over a flat inner buffer (see `OffsetArrays#window`), else the per-row default.
     ///
     /// @param arena allocator for the per-row fallback

reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetShortArray.java

@@ -4,6 +4,7 @@
 
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SegmentAllocator;
+import java.util.Optional;
 
 import java.util.function.LongBinaryOperator;
 
@@ -35,6 +36,15 @@ public long fold(long identity, LongBinaryOperator op) {
         return acc;
     }
 
+    /// The window of a flat inner buffer (see `OffsetArrays#window`), empty otherwise: a scan
+    /// window over a shared dict's codes then reads as flat codes, which the gathers read directly.
+    ///
+    /// @return the window's segment, if the inner buffer covers it
+    @Override
+    public Optional<MemorySegment> segmentIfPresent() {
+        return OffsetArrays.window(inner, offset, length, 2);
+    }
+
     /// Zero-copy over a flat inner buffer (see `OffsetArrays#window`), else the per-row default.
     ///
     /// @param arena allocator for the per-row fallback

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

@@ -4,6 +4,8 @@
 import io.github.dfa1.vortex.core.model.DType;
 import io.github.dfa1.vortex.core.model.PType;
 import io.github.dfa1.vortex.reader.array.DecimalArray;
+import io.github.dfa1.vortex.reader.array.DictIntArray;
+import io.github.dfa1.vortex.reader.array.TestArrays;
 import io.github.dfa1.vortex.reader.array.IntArray;
 import io.github.dfa1.vortex.reader.array.LazyConstantDecimalArray;
 import io.github.dfa1.vortex.reader.array.LazyDecimalArray;
@@ -49,6 +51,22 @@ void bufferBackedDecimal_slicesItsBufferToTheWindow() {
         assertThat(result.getDecimal(1)).isEqualTo(new BigDecimal("3.00"));
     }
 
+    @Test
+    void intDict_sharesItsPoolAndSlicesItsCodesToAFlatWindow() {
+        // Given — a global dict column cached whole: the window must stay a dict (its fold gathers
+        // off the pool) whose codes read as a flat segment (the gather's fast path)
+        IntArray pool = TestArrays.ints(10, 20, 30);
+        var full = DictIntArray.of(DType.I32, 5, pool, TestArrays.shorts((short) 2, (short) 0, (short) 1, (short) 1, (short) 0));
+
+        // When
+        var result = (DictIntArray) ScanIterator.sliceArray(full, 1, 3, DType.I32);
+
+        // Then
+        assertThat(result.values()).isSameAs(pool);
+        assertThat(result.codes().segmentIfPresent()).hasValueSatisfying(seg -> assertThat(seg.byteSize()).isEqualTo(6));
+        assertThat(result.fold(0, (acc, v) -> acc * 100 + v)).isEqualTo(102020);
+    }
+
     @Test
     void constantDecimal_keepsItsValueAtTheWindowLength() {
         // Given

reader/src/test/java/io/github/dfa1/vortex/reader/array/DictRecordSmokeTest.java

@@ -836,6 +836,22 @@ void doubleFold_poolWithoutSegment_takesThePerRowPath() {
             }
         }
 
+        @Test
+        void intFold_u16Codes_matchesPerRow() {
+            try (Arena arena = Arena.ofConfined()) {
+                // Given
+                IntArray values = intArray(arena, I32, 10, 20, 30);
+                ShortArray codes = shortArray(arena, U16, (short) 2, (short) 2, (short) 0, (short) 1);
+                DictIntArray sut = DictIntArray.of(I32, 4, values, codes);
+
+                // When
+                int result = sut.fold(0, (acc, v) -> acc * 100 + v);
+
+                // Then
+                assertThat(result).isEqualTo(30301020);
+            }
+        }
+
         @Test
         void intMaterialize_u16Codes_matchesPerRow() {
             try (Arena arena = Arena.ofConfined()) {

reader/src/test/java/io/github/dfa1/vortex/reader/array/OffsetArrayTest.java

@@ -82,6 +82,37 @@ void getShortShiftsByOffset() {
             }
         }
 
+        @Test
+        void segmentIfPresent_flatInner_isTheWindow() {
+            try (Arena arena = Arena.ofConfined()) {
+                // Given — a sliced dict's codes: gathers read this window as flat codes
+                ShortArray inner = shortArray(arena, (short) 100, (short) 200, (short) 300);
+                var sut = new OffsetShortArray(U16, 2, inner, 1L);
+
+                // When
+                var result = sut.segmentIfPresent();
+
+                // Then
+                assertThat(result).hasValueSatisfying(seg -> {
+                    assertThat(seg.byteSize()).isEqualTo(4);
+                    assertThat(seg.getAtIndex(ValueLayout.JAVA_SHORT_UNALIGNED, 0)).isEqualTo((short) 200);
+                });
+            }
+        }
+
+        @Test
+        void segmentIfPresent_broadcastInner_isEmpty() {
+            // Given a one-element buffer standing for 4 rows: no segment holds the window's rows
+            MemorySegment one = Arena.ofAuto().allocate(2, 2);
+            var sut = new OffsetShortArray(U16, 2, new MaterializedShortArray(U16, 4, one), 1L);
+
+            // When
+            var result = sut.segmentIfPresent();
+
+            // Then
+            assertThat(result).isEmpty();
+        }
+
         @Test
         void getIntWidensUnsignedThroughOffset() {
             try (Arena arena = Arena.ofConfined()) {