perf(reader): gather dict and window folds straight off their segments

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Performance
+- 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
 - A CSV import schema can declare any integer type from `I8` to `U64`, range-checked per value, and `F32`; an unsupported type fails in `withSchema` instead of at the first chunk. ([#593](https://github.com/dfa1/vortex-java/pull/593))
 

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

@@ -2,6 +2,8 @@
 
 import io.github.dfa1.vortex.core.error.VortexException;
 
+import java.lang.foreign.MemorySegment;
+
 /// Package-private helpers shared by the `DictXxxArray` records.
 ///
 /// Centralizes the codes-type validation (so the four record `of`
@@ -26,6 +28,17 @@ static void validateCodes(Array codes, long expectedLength) {
         }
     }
 
+    /// The array's segment when it holds all `n` elements of `width` bytes, so a gather can read it
+    /// directly instead of through a virtual `getX` whose call site sees every pool type there is.
+    ///
+    /// @param array the pool or codes array
+    /// @param n     the elements it must hold
+    /// @param width the element width in bytes
+    /// @return the segment, or `null` for a lazy array or a broadcast buffer shorter than `n`
+    static MemorySegment flat(Array array, long n, int width) {
+        return array.segmentIfPresent().filter(s -> s.byteSize() >= n * width).orElse(null);
+    }
+
     /// Reads the code at row `i` from `codes`. One switch on the
     /// codes' Array type — JIT inline-caches on the stable concrete type.
     static long readCode(Array codes, long i) {

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

@@ -6,6 +6,7 @@
 
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SegmentAllocator;
+import java.lang.foreign.ValueLayout;
 import java.util.function.DoubleBinaryOperator;
 import java.util.function.DoubleConsumer;
 
@@ -117,11 +118,40 @@ public void forEachDouble(DoubleConsumer cons) {
         }
     }
 
+    /// Folds the gathered values in row order. With a flat pool and flat `u8`/`u16` codes, as a
+    /// decoded chunk has them, it gathers straight off the two segments: the per-row path paid a
+    /// virtual `getShort` and a virtual `getDouble`, each with its broadcast check, for every 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 double fold(double identity, DoubleBinaryOperator op) {
         long n = length;
         DoubleArray vals = values;
         double result = identity;
+        MemorySegment pool = DictArrays.flat(vals, vals.length(), 8);
+        if (pool != null && codes instanceof ShortArray) {
+            MemorySegment c = DictArrays.flat(codes, n, 2);
+            if (c != null) {
+                for (long i = 0; i < n; i++) {
+                    long code = Short.toUnsignedLong(c.getAtIndex(VortexFormat.LE_SHORT, i));
+                    result = op.applyAsDouble(result, pool.getAtIndex(VortexFormat.LE_DOUBLE, code));
+                }
+                return result;
+            }
+        }
+        if (pool != null && codes instanceof ByteArray) {
+            MemorySegment c = DictArrays.flat(codes, n, 1);
+            if (c != null) {
+                for (long i = 0; i < n; i++) {
+                    long code = Byte.toUnsignedLong(c.get(ValueLayout.JAVA_BYTE, i));
+                    result = op.applyAsDouble(result, pool.getAtIndex(VortexFormat.LE_DOUBLE, 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/DictIntArray.java

@@ -60,7 +60,19 @@ public MemorySegment materialize(SegmentAllocator arena) {
         long n = length;
         MemorySegment dst = arena.allocate(n * 4L, 4);
         IntArray vals = values;
-        switch (CanonicalArrays.of(codes, arena)) {
+        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.
+        MemorySegment pool = DictArrays.flat(vals, vals.length(), 4);
+        MemorySegment c = flatCodes instanceof ShortArray ? DictArrays.flat(flatCodes, n, 2) : null;
+        if (pool != null && c != null) {
+            for (long i = 0; i < n; i++) {
+                long code = Short.toUnsignedLong(c.getAtIndex(VortexFormat.LE_SHORT, i));
+                dst.setAtIndex(VortexFormat.LE_INT, i, pool.getAtIndex(VortexFormat.LE_INT, code));
+            }
+            return dst.asReadOnly();
+        }
+        switch (flatCodes) {
             case ByteArray ba -> {
                 for (long i = 0; i < n; i++) {
                     dst.setAtIndex(VortexFormat.LE_INT, i, vals.getInt(Byte.toUnsignedLong(ba.getByte(i))));

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

@@ -1,9 +1,12 @@
 package io.github.dfa1.vortex.reader.array;
 
+import io.github.dfa1.vortex.core.io.VortexFormat;
 import io.github.dfa1.vortex.core.model.DType;
 
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SegmentAllocator;
+import java.util.Optional;
+import java.util.function.DoubleBinaryOperator;
 
 /// Sliced view over a [DoubleArray]: `getDouble(i) = inner.getDouble(i + offset)`.
 ///
@@ -19,6 +22,26 @@ public double getDouble(long i) {
         return inner.getDouble(i + offset);
     }
 
+    /// Folds the window straight off a flat inner buffer (see `OffsetArrays#window`): a scan window
+    /// of a shared flat lands here, and the per-row default paid two virtual `getDouble` calls a row.
+    ///
+    /// @param identity initial accumulator value
+    /// @param op       binary operator applied to the accumulator and each double element
+    /// @return the final accumulated result
+    @Override
+    public double fold(double identity, DoubleBinaryOperator op) {
+        Optional<MemorySegment> window = OffsetArrays.window(inner, offset, length, 8);
+        if (window.isEmpty()) {
+            return DoubleArray.super.fold(identity, op);
+        }
+        MemorySegment seg = window.get();
+        double result = identity;
+        for (long i = 0; i < length; i++) {
+            result = op.applyAsDouble(result, seg.getAtIndex(VortexFormat.LE_DOUBLE, i));
+        }
+        return result;
+    }
+
     /// 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

@@ -1,9 +1,12 @@
 package io.github.dfa1.vortex.reader.array;
 
+import io.github.dfa1.vortex.core.io.VortexFormat;
 import io.github.dfa1.vortex.core.model.DType;
 
 import java.lang.foreign.MemorySegment;
 import java.lang.foreign.SegmentAllocator;
+import java.util.Optional;
+import java.util.function.IntBinaryOperator;
 
 /// Sliced view over an [IntArray]: `getInt(i) = inner.getInt(i + offset)`.
 ///
@@ -19,6 +22,26 @@ public int getInt(long i) {
         return inner.getInt(i + offset);
     }
 
+    /// Folds the window straight off a flat inner buffer (see `OffsetArrays#window`): a scan window
+    /// of a shared flat lands here, and the per-row default paid two virtual `getInt` calls a row.
+    ///
+    /// @param identity initial accumulator value
+    /// @param op       binary operator applied to the accumulator and each int element
+    /// @return the final accumulated result
+    @Override
+    public int fold(int identity, IntBinaryOperator op) {
+        Optional<MemorySegment> window = OffsetArrays.window(inner, offset, length, 4);
+        if (window.isEmpty()) {
+            return IntArray.super.fold(identity, op);
+        }
+        MemorySegment seg = window.get();
+        int result = identity;
+        for (long i = 0; i < length; i++) {
+            result = op.applyAsInt(result, seg.getAtIndex(VortexFormat.LE_INT, i));
+        }
+        return result;
+    }
+
     /// 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/decode/DictEncodingDecoder.java

@@ -10,6 +10,7 @@
 import io.github.dfa1.vortex.reader.array.Array;
 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.DictByteArray;
 import io.github.dfa1.vortex.reader.array.DictDoubleArray;
 import io.github.dfa1.vortex.reader.array.DictFloat16Array;
@@ -140,7 +141,14 @@ private static Array decodeRustProto(DecodeContext ctx, MemorySegment metaBuf) {
                 physicalBytes(rawValues), valPType.byteSize());
         validateCodesInRange(rawCodes, codePType, rawValues.length());
 
-        Array values = buildLazyDict(ctx.dtype(), valPType, rowCount, rawValues, rawCodes);
+        // A floating pool is decoded once, as Rust's dict canonicalizes its values before
+        // gathering: a lazy ALP pool re-ran its multiply on every row. An integer pool stays lazy:
+        // FoR costs an add a row, and decoding it shifted the gathers' type profile enough to cost
+        // the TPC-DS store_sales scan 3.5%.
+        Array pool = rawValues instanceof DoubleArray || rawValues instanceof FloatArray
+                ? CanonicalArrays.of(rawValues, ctx.arena())
+                : rawValues;
+        Array values = buildLazyDict(ctx.dtype(), valPType, rowCount, pool, rawCodes);
         BoolArray rowValidity = poolValidity == null
                 ? codesValidity
                 : rowValidity(ctx, ctx.materialize(rawCodes), codePType, codesValidity, poolValidity, rowCount);

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

@@ -783,6 +783,84 @@ void dictLongAcceptsNullableValuesViaWrapper() {
 
     // ── Helpers ───────────────────────────────────────────────────────────────
 
+    /// The flat-segment gathers must equal the per-row path they bypass: same values, same row
+    /// order. The fold operator is order-sensitive (`acc * 2 + v`) so a reordered gather fails.
+    @Nested
+    class FlatGather {
+
+        @Test
+        void doubleFold_u16Codes_matchesPerRow() {
+            try (Arena arena = Arena.ofConfined()) {
+                // Given
+                DoubleArray values = doubleArray(arena, 0.5, 1.5, 2.5);
+                ShortArray codes = shortArray(arena, U16, (short) 2, (short) 0, (short) 1, (short) 2);
+                DictDoubleArray sut = DictDoubleArray.of(F64, 4, values, codes);
+
+                // When
+                double result = sut.fold(1.0, (acc, v) -> acc * 2 + v);
+
+                // Then
+                assertThat(result).isEqualTo(perRowDoubleFold(sut));
+            }
+        }
+
+        @Test
+        void doubleFold_u8Codes_matchesPerRow() {
+            try (Arena arena = Arena.ofConfined()) {
+                // Given
+                DoubleArray values = doubleArray(arena, 0.5, 1.5, 2.5);
+                ByteArray codes = byteArray(arena, (byte) 1, (byte) 2, (byte) 0);
+                DictDoubleArray sut = DictDoubleArray.of(F64, 3, values, codes);
+
+                // When
+                double result = sut.fold(1.0, (acc, v) -> acc * 2 + v);
+
+                // Then
+                assertThat(result).isEqualTo(perRowDoubleFold(sut));
+            }
+        }
+
+        @Test
+        void doubleFold_poolWithoutSegment_takesThePerRowPath() {
+            try (Arena arena = Arena.ofConfined()) {
+                // Given — an offset view has no segment of its own, so the fast path must step aside
+                DoubleArray values = new OffsetDoubleArray(F64, 3, doubleArray(arena, 9.0, 0.5, 1.5, 2.5), 1);
+                ShortArray codes = shortArray(arena, U16, (short) 2, (short) 0, (short) 1);
+                DictDoubleArray sut = DictDoubleArray.of(F64, 3, values, codes);
+
+                // When
+                double result = sut.fold(1.0, (acc, v) -> acc * 2 + v);
+
+                // Then — ((1 * 2 + 2.5) * 2 + 0.5) * 2 + 1.5
+                assertThat(result).isEqualTo(20.5);
+            }
+        }
+
+        @Test
+        void intMaterialize_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
+                MemorySegment result = sut.materialize(arena);
+
+                // Then
+                assertThat(result.toArray(VortexFormat.LE_INT)).containsExactly(30, 30, 10, 20);
+            }
+        }
+
+        private static double perRowDoubleFold(DoubleArray array) {
+            double acc = 1.0;
+            for (long i = 0; i < array.length(); i++) {
+                acc = acc * 2 + array.getDouble(i);
+            }
+            return acc;
+        }
+    }
+
     private static LongArray longArray(Arena arena, long... values) {
         return longArray(arena, I64, values);
     }

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

@@ -244,6 +244,33 @@ void getIntShiftsByOffset() {
             assertThat(sut.getInt(0)).isEqualTo(30);
             assertThat(sut.getInt(1)).isEqualTo(40);
         }
+
+        @Test
+        void fold_flatInner_foldsTheWindowInOrder() {
+            // Given — an order-sensitive operator, so folding outside the window or out of order fails
+            IntArray inner = TestArrays.ints(1, 2, 3, 4, 5);
+            var sut = new OffsetIntArray(DType.I32, 3, inner, 1L);
+
+            // When
+            int result = sut.fold(0, (acc, v) -> acc * 10 + v);
+
+            // Then
+            assertThat(result).isEqualTo(234);
+        }
+
+        @Test
+        void fold_broadcastInner_fallsBackToPerRow() {
+            // Given a one-element buffer standing for 4 rows: a window read off it would overrun
+            MemorySegment one = Arena.ofAuto().allocate(4, 4);
+            one.setAtIndex(ValueLayout.JAVA_INT_UNALIGNED, 0, 7);
+            var sut = new OffsetIntArray(DType.I32, 3, new MaterializedIntArray(DType.I32, 4, one), 1L);
+
+            // When
+            int result = sut.fold(0, Integer::sum);
+
+            // Then
+            assertThat(result).isEqualTo(21);
+        }
     }
 
     @Nested
@@ -272,6 +299,33 @@ void getDoubleShiftsByOffset() {
             assertThat(sut.getDouble(0)).isEqualTo(2.5);
             assertThat(sut.getDouble(1)).isEqualTo(3.5);
         }
+
+        @Test
+        void fold_flatInner_foldsTheWindowInOrder() {
+            // Given — an order-sensitive operator, so folding outside the window or out of order fails
+            DoubleArray inner = TestArrays.doubles(1.0, 2.0, 3.0, 4.0);
+            var sut = new OffsetDoubleArray(DType.F64, 2, inner, 1L);
+
+            // When
+            double result = sut.fold(0.0, (acc, v) -> acc * 10 + v);
+
+            // Then
+            assertThat(result).isEqualTo(23.0);
+        }
+
+        @Test
+        void fold_broadcastInner_fallsBackToPerRow() {
+            // Given a one-element buffer standing for 4 rows: a window read off it would overrun
+            MemorySegment one = Arena.ofAuto().allocate(8, 8);
+            one.setAtIndex(ValueLayout.JAVA_DOUBLE_UNALIGNED, 0, 1.5);
+            var sut = new OffsetDoubleArray(DType.F64, 3, new MaterializedDoubleArray(DType.F64, 4, one), 1L);
+
+            // When
+            double result = sut.fold(0.0, java.lang.Double::sum);
+
+            // Then
+            assertThat(result).isEqualTo(4.5);
+        }
     }
 
     private static ByteArray byteArray(Arena arena, byte... vs) {