From 6521a71cd3981488d15547ebfbed0fbcec24245f Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sat, 24 Oct 2026 19:34:22 +0200 Subject: [PATCH] perf(reader): fold an i32 global dict column off its codes, not a per-row copy A dict chunk that spans several scan windows (a global dict column's codes chunk at 512Ki rows against 128Ki-row windows) was canonicalized once per covering flat: DictIntArray.materialize wrote a 4-byte value for every row, and each window then folded that copy. The shared flat now keeps an i32 dict as a dict over canonical codes and pool; each window slices only the codes, and DictIntArray.fold/materialize gather straight off the pool. For that to reach the fast path, the byte/short/int/long offset views report their window as their segment when the inner buffer covers it, and DictIntArray's gathers handle u8 codes as well as u16 (without u8, TPC-DS store_sales lost 4%). ParquetVsVortexReadBenchmark.vortexReadMultiColumn (PULocationID is a global dict): 162 -> 182 ops/s single fork on a quiet machine; 3 forks base vs change in one later, noisier session 133.8 -> 155.4 (+16%). RealFileScanBenchmark A/B on all seven real files: within noise (store_sales and appian-order re-checked at 3 forks, alternating). Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 1 + .../dfa1/vortex/reader/ScanIterator.java | 18 +++++++- .../vortex/reader/array/DictIntArray.java | 44 ++++++++++++++++--- .../vortex/reader/array/OffsetByteArray.java | 10 +++++ .../vortex/reader/array/OffsetIntArray.java | 9 ++++ .../vortex/reader/array/OffsetLongArray.java | 10 +++++ .../vortex/reader/array/OffsetShortArray.java | 10 +++++ .../vortex/reader/ScanIteratorSliceTest.java | 18 ++++++++ .../reader/array/DictRecordSmokeTest.java | 16 +++++++ .../vortex/reader/array/OffsetArrayTest.java | 31 +++++++++++++ 10 files changed, 161 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1208026dc..71e226c11 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/ScanIterator.java b/reader/src/main/java/io/github/dfa1/vortex/reader/ScanIterator.java index c8f2c3f39..34110300a 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/ScanIterator.java +++ b/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); diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/array/DictIntArray.java b/reader/src/main/java/io/github/dfa1/vortex/reader/array/DictIntArray.java index b13e9000b..70d747c50 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/array/DictIntArray.java +++ b/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++) { diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetByteArray.java b/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetByteArray.java index 8660e9cef..76b5cdeb6 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetByteArray.java +++ b/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 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 diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetIntArray.java b/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetIntArray.java index 3f19297e2..917cea985 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetIntArray.java +++ b/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 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 diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetLongArray.java b/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetLongArray.java index fda646c28..3f156d015 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetLongArray.java +++ b/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 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 diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetShortArray.java b/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetShortArray.java index 9cff6196e..6e5d93782 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/array/OffsetShortArray.java +++ b/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 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 diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/ScanIteratorSliceTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/ScanIteratorSliceTest.java index ed9dfb403..b52c8f660 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/ScanIteratorSliceTest.java +++ b/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 diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/array/DictRecordSmokeTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/array/DictRecordSmokeTest.java index 071380fbb..fe1c0cf37 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/array/DictRecordSmokeTest.java +++ b/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()) { diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/array/OffsetArrayTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/array/OffsetArrayTest.java index 71dab1e68..2bb9424c2 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/array/OffsetArrayTest.java +++ b/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()) {