From 27f0c75a90d644e4bcb696598ebd6bb7e5ff4747 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sat, 24 Oct 2026 23:35:56 +0200 Subject: [PATCH] perf(reader): select a filtered chunk's rows through a Rust-style Mask (#602) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The filtered scan's selection is now a Mask, as Rust's vortex_mask::Mask: AllTrue, AllFalse, or Values holding the bitmap, whose sorted indices and runs of consecutive rows are built on first use and cached, so every column of the chunk reuses them (Rust's OnceLock fields). The copy picks as Rust's filter does at vortex-jni 0.89.0: runs in bulk when the mask is dense and its runs average at least 8 rows (MIN_SLICES_AVERAGE_RUN_LENGTH), a typed per-row gather otherwise. - RowSelection.toMask: one walk sets a bit per selected row; a bitmap selection with no residual test hands its words over with no walk. - Mask.limit truncates the bitmap, as Mask::limit does. - An all-true mask returns the column itself instead of copying it. - A filtered chunk past Integer.MAX_VALUE rows, or a fixed-size list whose element count would pass it, fails as VortexException; a fixed-size list of zero elements selects no elements. - Public surface: Mask, its three variants, fromBitmap, length, trueCount and limit; the index and run forms stay package-private. JavaVsJniFilterBenchmark.javaFilterClose, ohlc.vortex, -f 3, ops/s (the 0.001 label selects about 11% of rows on this file): selectivity int[] (main) Mask alloc/op int[] -> Mask 0.001 6729 ± 303 7130 ± 607 196 KB -> 81 KB 0.01 6806 ± 113 6843 ± 560 196 KB -> 81 KB 0.1 5326 ± 44 5860 ± 50 363 KB -> 118 KB 1.0 2458 ± 103 2347 ± 135 1247 KB -> 718 KB Long.compress/expand were considered for the bit walks and rejected: on this Apple M5 C2 inlines them as Java (no BEXT/BDEP without SVE2), while Long.bitCount and numberOfTrailingZeros, used here, are intrinsics. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01SwhK4W5cgj7y7s4ZPCif5x --- CHANGELOG.md | 2 +- .../io/github/dfa1/vortex/reader/Chunk.java | 9 +- .../dfa1/vortex/reader/ScanIterator.java | 31 +-- .../github/dfa1/vortex/reader/array/Mask.java | 229 ++++++++++++++++ .../vortex/reader/array/SelectedArrays.java | 217 ++++++++------- .../reader/compute/FusedFilterAggregate.java | 10 + .../vortex/reader/compute/RowSelection.java | 15 ++ .../dfa1/vortex/reader/array/MaskTest.java | 249 ++++++++++++++++++ .../reader/array/SelectedArraysTest.java | 149 ++++++++--- .../vortex/reader/compute/ComputeTest.java | 31 +++ 10 files changed, 789 insertions(+), 153 deletions(-) create mode 100644 reader/src/main/java/io/github/dfa1/vortex/reader/array/Mask.java create mode 100644 reader/src/test/java/io/github/dfa1/vortex/reader/array/MaskTest.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 1ccd04e8b..635e45d86 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Performance -- A filtered scan copies the matching rows into flat arrays, as Rust's `filter` does, instead of reading them through a per-row index lookup, and allocates up to 4.7× less. ([#602](https://github.com/dfa1/vortex-java/issues/602)) +- A filtered scan copies the matching rows into flat arrays, as Rust's `filter` does, instead of reading them through a per-row index lookup, and allocates up to 11× less. ([#602](https://github.com/dfa1/vortex-java/issues/602)) - 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)) diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/Chunk.java b/reader/src/main/java/io/github/dfa1/vortex/reader/Chunk.java index 088fe0997..4177ad2db 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/Chunk.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/Chunk.java @@ -4,6 +4,7 @@ import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.reader.array.Array; +import io.github.dfa1.vortex.reader.array.Mask; import io.github.dfa1.vortex.reader.array.SelectedArrays; import io.github.dfa1.vortex.core.model.ExtensionId; import io.github.dfa1.vortex.reader.extension.ExtensionDecoder; @@ -109,21 +110,21 @@ public Chunk slice(long offset, long length) { return new Chunk(length, names, positions, sliced, null, registry, _ -> { }); } - /// The rows `rows` of the columns `names`, in that order: the chunk a filtered scan returns. + /// The rows `rows` keeps, of the columns `names`: the chunk a filtered scan returns. /// It takes over this chunk's arena, so this chunk must not be used again; closing the result /// releases both. /// - /// @param rows the rows to keep, strictly ascending, each in `[0, rowCount())` + /// @param rows the rows to keep, a mask over this chunk's rows /// @param names the columns to keep, a subset of this chunk's /// @param positions each of `names` to its index /// @param onClose called when the result closes /// @return the selected rows - Chunk select(int[] rows, List names, Map positions, Consumer onClose) { + Chunk select(Mask rows, List names, Map positions, Consumer onClose) { IntFunction selected = i -> { Column column = column(position(names.get(i))); return new Column(SelectedArrays.select(column.array(), rows, arena), column.dtype()); }; - return new Chunk(rows.length, names, positions, selected, arena, registry, onClose); + return new Chunk(rows.trueCount(), names, positions, selected, arena, registry, onClose); } /// One decoded column: its zero-copy [Array] view and the [DType] it was declared with in the 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 8cbf4e00e..a3fc3a196 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 @@ -27,6 +27,7 @@ import io.github.dfa1.vortex.reader.array.LazyDecimalArray; import io.github.dfa1.vortex.reader.array.LazyDecimalBytePartsArray; import io.github.dfa1.vortex.reader.array.LongArray; +import io.github.dfa1.vortex.reader.array.Mask; import io.github.dfa1.vortex.reader.array.MaskedArray; import io.github.dfa1.vortex.reader.array.NullArray; import io.github.dfa1.vortex.reader.array.OffsetBoolArray; @@ -50,7 +51,6 @@ import java.lang.foreign.SegmentAllocator; import java.lang.foreign.ValueLayout; import java.util.ArrayList; -import java.util.Arrays; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -108,7 +108,7 @@ public final class ScanIterator implements Iterator, AutoCloseable { // A chunk hasNext() decoded to test the filter on, and its matching rows, until next() returns // them: a chunk with no match is skipped, as Rust's scan returns no empty batch. private Chunk pending; - private int[] pendingRows; + private Mask pendingRows; private List projectedNames; private List projectedDtypes; private Map projectedPositions; @@ -364,23 +364,18 @@ private boolean selectMatchingRows(ChunkSpec spec) { Chunk full = chunkFor(spec, spec.rowCount()); try { RowSelection selection = Compute.matcher(full, filter); - // the matching rows, grown as they come: memory follows the matches, not the chunk - int allowed = (int) Math.min(spec.rowCount(), limit - rowsReturned); - int[] rows = new int[Math.min(allowed, 1024)]; - int count = 0; - for (long row = selection.nextSelected(0); row >= 0 && count < allowed; - row = selection.nextSelected(row + 1)) { - if (count == rows.length) { - rows = Arrays.copyOf(rows, (int) Math.min(allowed, 2L * count)); - } - rows[count++] = (int) row; + if (spec.rowCount() > Integer.MAX_VALUE) { + throw new VortexException("scan: a filtered chunk of " + spec.rowCount() + + " rows exceeds the " + Integer.MAX_VALUE + " rows a row selection holds"); } - if (count == 0) { + int allowed = (int) Math.min(spec.rowCount(), limit - rowsReturned); + Mask mask = selection.toMask((int) spec.rowCount()).limit(allowed); + if (mask.trueCount() == 0) { full.close(); return false; } pending = full; - pendingRows = count == rows.length ? rows : Arrays.copyOf(rows, count); + pendingRows = mask; return true; } catch (RuntimeException e) { full.close(); @@ -426,11 +421,11 @@ public Chunk next() { Chunk chunk; if (pending != null) { chunk = pending; - int[] rows = pendingRows; + Mask rows = pendingRows; pending = null; pendingRows = null; chunk = selectedChunk(chunk, rows); - rowsReturned += rows.length; + rowsReturned += rows.trueCount(); } else { long chunkRows = Math.min(spec.rowCount(), limit - rowsReturned); chunk = chunkFor(spec, chunkRows); @@ -442,7 +437,7 @@ public Chunk next() { /// The `rows` of `full` over the output columns, or `full` itself when that is every row and /// every column. - private Chunk selectedChunk(Chunk full, int[] rows) { + private Chunk selectedChunk(Chunk full, Mask rows) { List names = outputNames; Map positions = outputPositions; if (singleColumnIsSyntheticWrapper) { @@ -450,7 +445,7 @@ private Chunk selectedChunk(Chunk full, int[] rows) { names = List.copyOf(full.columns().keySet()); positions = Chunk.positions(names); } - if (rows.length == full.rowCount() && names.size() == full.columns().size()) { + if (rows instanceof Mask.AllTrue && names.size() == full.columns().size()) { return full; } return full.select(rows, names, positions, this::onChunkClosed); diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/array/Mask.java b/reader/src/main/java/io/github/dfa1/vortex/reader/array/Mask.java new file mode 100644 index 000000000..5e0f46438 --- /dev/null +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/array/Mask.java @@ -0,0 +1,229 @@ +package io.github.dfa1.vortex.reader.array; + +import java.util.Arrays; + +/// The rows of an array a filter keeps, as Rust's `vortex_mask::Mask`: every row, no row, or some +/// rows held as a bitmap. The bitmap derives its row indices and its runs of consecutive rows on +/// first use and keeps them, so each column a filtered scan copies reuses them instead of walking +/// the bitmap again. +public sealed interface Mask permits Mask.AllTrue, Mask.AllFalse, Mask.Values { + + /// Returns the number of rows the mask covers. + /// + /// @return the row count, kept or not + int length(); + + /// Returns the number of rows the mask keeps. + /// + /// @return the kept row count + int trueCount(); + + /// The mask keeping row `i` when bit `i` of `words` is set, as Rust's `Mask::from_buffer`: a + /// bitmap keeping every row or none becomes [AllTrue] or [AllFalse]. + /// + /// @param words the bitmap, bit `i` for row `i`; it is taken over, not copied, and any bit past + /// `length` in its last word is cleared + /// @param length the row count; `words` holds `(length + 63) / 64` words + /// @return the mask + /// @throws IllegalArgumentException if `words` does not hold `length` bits + static Mask fromBitmap(long[] words, int length) { + if (length < 0 || words.length != (length + 63) >>> 6) { + throw new IllegalArgumentException(words.length + " words for " + length + " rows"); + } + // a stray bit past the last row would count as a kept row that does not exist + if ((length & 63) != 0) { + words[words.length - 1] &= (1L << length) - 1; + } + int trueCount = 0; + for (long word : words) { + trueCount += Long.bitCount(word); + } + if (trueCount == 0) { + return new AllFalse(length); + } + if (trueCount == length) { + return new AllTrue(length); + } + return new Values(words, length, trueCount); + } + + /// This mask keeping only its first `limit` kept rows, as Rust's `Mask::limit`. + /// + /// @param limit the most rows to keep + /// @return this mask when it keeps no more than `limit` rows, else the truncated mask + default Mask limit(int limit) { + if (trueCount() <= limit) { + return this; + } + if (limit <= 0) { + return new AllFalse(length()); + } + // past the checks above, the mask keeps more than `limit >= 1` rows: AllTrue or Values + long[] words = new long[(length() + 63) >>> 6]; + if (this instanceof Values v) { + int remaining = limit; + for (int w = 0; remaining > 0; w++) { + long word = v.words[w]; + int count = Long.bitCount(word); + if (count <= remaining) { + words[w] = word; + remaining -= count; + } else { + for (; remaining > 0; remaining--) { + long lowest = word & -word; + words[w] |= lowest; + word ^= lowest; + } + } + } + } else { + Arrays.fill(words, 0, limit >>> 6, -1L); + if ((limit & 63) != 0) { + words[limit >>> 6] = (1L << limit) - 1; + } + } + return new Values(words, length(), limit); + } + + /// Every row kept. + /// + /// @param length the row count + record AllTrue(int length) implements Mask { + @Override + public int trueCount() { + return length; + } + } + + /// No row kept. + /// + /// @param length the row count + record AllFalse(int length) implements Mask { + @Override + public int trueCount() { + return 0; + } + } + + /// Some rows kept, at least one and not all, as Rust's `MaskValues`. The index and run forms + /// are built from the bitmap on first use and published whole through a volatile field, so + /// threads sharing a mask may build one twice but never see a partial one. + final class Values implements Mask { + + private final long[] words; + private final int length; + private final int trueCount; + private volatile int[] indices; + private volatile int[] slices; + + private Values(long[] words, int length, int trueCount) { + this.words = words; + this.length = length; + this.trueCount = trueCount; + } + + @Override + public int length() { + return length; + } + + @Override + public int trueCount() { + return trueCount; + } + + /// The mask keeping the rows `indices`, as Rust's `Mask::from_indices`; the indices are + /// kept as the mask's index form. Package-private: the indices are trusted, not checked. + /// + /// @param length the row count + /// @param indices the kept rows, strictly ascending, each in `[0, length)`; taken over + /// @return the mask + static Mask fromIndices(int length, int[] indices) { + if (indices.length == 0) { + return new AllFalse(length); + } + if (indices.length == length) { + return new AllTrue(length); + } + long[] words = new long[(length + 63) >>> 6]; + for (int i : indices) { + words[i >>> 6] |= 1L << i; + } + Values values = new Values(words, length, indices.length); + values.indices = indices; + return values; + } + + /// The fraction of rows kept, in `(0, 1)`. + double density() { + return (double) trueCount / length; + } + + /// The kept rows, ascending, built from the bitmap a word at a time on first use. The array + /// is shared by every caller and must not be modified. + int[] indices() { + int[] result = indices; + if (result == null) { + result = new int[trueCount]; + int k = 0; + for (int w = 0; w < words.length; w++) { + long word = words[w]; + while (word != 0) { + result[k++] = (w << 6) + Long.numberOfTrailingZeros(word); + word &= word - 1; + } + } + indices = result; + } + return result; + } + + /// The runs of consecutive kept rows as `start, end` pairs, `end` exclusive, ascending and + /// not touching, built from the bitmap on first use. The array is shared by every caller and + /// must not be modified. + int[] slices() { + int[] result = slices; + if (result == null) { + int[] runs = new int[16]; + int count = 0; + int row = nextSet(0); + while (row >= 0) { + int end = nextClear(row); + if (count + 2 > runs.length) { + runs = Arrays.copyOf(runs, runs.length * 2); + } + runs[count++] = row; + runs[count++] = end; + row = end < length ? nextSet(end) : -1; + } + result = Arrays.copyOf(runs, count); + slices = result; + } + return result; + } + + private int nextSet(int from) { + int w = from >>> 6; + long word = words[w] & -1L << from; + while (word == 0) { + if (++w == words.length) { + return -1; + } + word = words[w]; + } + return (w << 6) + Long.numberOfTrailingZeros(word); + } + + private int nextClear(int from) { + int w = from >>> 6; + long word = ~words[w] & -1L << from; + while (word == 0) { + if (++w == words.length) { + return length; + } + word = ~words[w]; + } + return Math.min(length, (w << 6) + Long.numberOfTrailingZeros(word)); + } + } +} diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/array/SelectedArrays.java b/reader/src/main/java/io/github/dfa1/vortex/reader/array/SelectedArrays.java index 5aeb27493..9b8fdfc9c 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/array/SelectedArrays.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/array/SelectedArrays.java @@ -11,105 +11,123 @@ import java.util.ArrayList; import java.util.List; -/// Row selection: the given rows of another array, copied into flat buffers, as Rust's `filter` -/// does. Leaf arrays become materialized arrays, so the flat fast paths apply to a filtered column; -/// a dictionary copies only its codes; composite arrays select their children. This is how a -/// filtered scan returns only the rows that match, as Rust's scan does. -/// -/// The rows are a sorted `int[]`, one entry per selected row (roaring's array container): a copy -/// loop reads them straight, where a bitmap makes every walk pay for its empty words. +/// Row selection: the rows of another array a [Mask] keeps, copied into flat buffers, as Rust's +/// `filter` does. Leaf arrays become materialized arrays, so the flat fast paths apply to a filtered +/// column; a dictionary copies only its codes; composite arrays select their children. This is how +/// a filtered scan returns only the rows that match, as Rust's scan does. public final class SelectedArrays { + /// Rust's `MIN_SLICES_AVERAGE_RUN_LENGTH`: shorter runs copy slower than a per-row gather. + private static final int MIN_AVERAGE_RUN = 8; + private SelectedArrays() { } - /// The rows `rows[0] .. rows[rows.length - 1]` of `full`, in that order. + /// The rows of `full` that `mask` keeps, in order. /// /// @param full the array to select from - /// @param rows row indices into `full`, strictly ascending, each in `[0, full.length())` + /// @param mask the rows to keep; its length is `full`'s /// @param arena allocator for the copied rows - /// @return an array of `rows.length` rows with `full`'s dtype + /// @return `full` itself when `mask` keeps every row, else an array of `mask.trueCount()` rows + /// with `full`'s dtype /// @throws VortexException if `full`'s array type has no row selection - public static Array select(Array full, int[] rows, SegmentAllocator arena) { - return select(full, rows, 0, arena); + public static Array select(Array full, Mask mask, SegmentAllocator arena) { + return switch (mask) { + case Mask.AllTrue _ -> full; + case Mask.AllFalse _ -> Array.limited(full, 0); + case Mask.Values values -> select(full, values, 0, arena); + }; } - /// The rows `rows[k] + offset` of `full`: an offset slice selects from the array it slices, so a - /// dictionary under it stays one. - // ponytail: a flat buffer copies by a typed gather or by runs, by density, anything else by a - // per-row walk; Rust also has a SIMD compress (vortex-array filter/execute/buffer.rs), worth it - // if filtered scans profile hot here - private static Array select(Array full, int[] rows, long offset, SegmentAllocator arena) { - int n = rows.length; + /// The rows `r + offset` of `full`, for each row `r` that `mask` keeps: an offset slice selects + /// from the array it slices, so a dictionary under it stays one. + // ponytail: a flat buffer copies by runs or a per-row gather, anything else row by row through + // the mask's cached indices; Rust also has a SIMD compress (vortex-array filter/execute/buffer.rs), + // worth it if filtered scans profile hot here + private static Array select(Array full, Mask.Values mask, long offset, SegmentAllocator arena) { + int n = mask.trueCount(); return switch (full) { case MaskedArray m -> { BoolArray validity = null; if (m.validity() != null) { - validity = (BoolArray) select(m.validity(), rows, offset, arena); + validity = (BoolArray) select(m.validity(), mask, offset, arena); } - yield new MaskedArray(select(m.inner(), rows, offset, arena), validity); + yield new MaskedArray(select(m.inner(), mask, offset, arena), validity); } - case OffsetLongArray o -> select(o.inner(), rows, offset + o.offset(), arena); - case OffsetIntArray o -> select(o.inner(), rows, offset + o.offset(), arena); - case OffsetShortArray o -> select(o.inner(), rows, offset + o.offset(), arena); - case OffsetByteArray o -> select(o.inner(), rows, offset + o.offset(), arena); - case OffsetDoubleArray o -> select(o.inner(), rows, offset + o.offset(), arena); - case OffsetFloatArray o -> select(o.inner(), rows, offset + o.offset(), arena); - case OffsetBoolArray o -> select(o.inner(), rows, offset + o.offset(), arena); - case VarBinSlicedArray s -> select(s.inner(), rows, offset + s.offset(), arena); + case OffsetLongArray o -> select(o.inner(), mask, offset + o.offset(), arena); + case OffsetIntArray o -> select(o.inner(), mask, offset + o.offset(), arena); + case OffsetShortArray o -> select(o.inner(), mask, offset + o.offset(), arena); + case OffsetByteArray o -> select(o.inner(), mask, offset + o.offset(), arena); + case OffsetDoubleArray o -> select(o.inner(), mask, offset + o.offset(), arena); + case OffsetFloatArray o -> select(o.inner(), mask, offset + o.offset(), arena); + case OffsetBoolArray o -> select(o.inner(), mask, offset + o.offset(), arena); + case VarBinSlicedArray s -> select(s.inner(), mask, offset + s.offset(), arena); // A dictionary keeps its value pool and selects its codes, so the compute kernels' // dictionary lanes (code scans, per-entry string tests) still apply to the selection. case VarBinDictArray d -> new VarBinDictArray(d.dtype(), n, d.bytesSegment(), d.dictValOffsets(), - d.dictValOffPType(), copy(d.dictCodesSegs(), d.dictCodesPType().byteSize(), rows, offset, arena), + d.dictValOffPType(), copy(d.dictCodesSegs(), d.dictCodesPType().byteSize(), mask, offset, arena), d.dictCodesPType()); - case DictLongArray d -> DictLongArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, offset, arena)); - case DictIntArray d -> DictIntArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, offset, arena)); - case DictShortArray d -> DictShortArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, offset, arena)); - case DictByteArray d -> DictByteArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, offset, arena)); - case DictDoubleArray d -> DictDoubleArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, offset, arena)); - case DictFloatArray d -> DictFloatArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, offset, arena)); - case DictFloat16Array d -> DictFloat16Array.of(d.dtype(), n, d.values(), codes(d.codes(), rows, offset, arena)); - case LongArray a -> new MaterializedLongArray(a.dtype(), n, longs(a, rows, offset, arena)); - case IntArray a -> new MaterializedIntArray(a.dtype(), n, ints(a, rows, offset, arena)); - case ShortArray a -> new MaterializedShortArray(a.dtype(), n, shorts(a, rows, offset, arena)); - case ByteArray a -> new MaterializedByteArray(a.dtype(), n, bytes(a, rows, offset, arena)); - case DoubleArray a -> new MaterializedDoubleArray(a.dtype(), n, doubles(a, rows, offset, arena)); - case FloatArray a -> new MaterializedFloatArray(a.dtype(), n, floats(a, rows, offset, arena)); - case Float16Array a -> new MaterializedFloat16Array(a.dtype(), n, halves(a, rows, offset, arena)); - case BoolArray a -> new MaterializedBoolArray(a.dtype(), n, bools(a, rows, offset, arena)); - case VarBinArray a -> strings(a, rows, offset, arena); - case DecimalArray a -> decimals(a, rows, offset, arena); + case DictLongArray d -> DictLongArray.of(d.dtype(), n, d.values(), codes(d.codes(), mask, offset, arena)); + case DictIntArray d -> DictIntArray.of(d.dtype(), n, d.values(), codes(d.codes(), mask, offset, arena)); + case DictShortArray d -> DictShortArray.of(d.dtype(), n, d.values(), codes(d.codes(), mask, offset, arena)); + case DictByteArray d -> DictByteArray.of(d.dtype(), n, d.values(), codes(d.codes(), mask, offset, arena)); + case DictDoubleArray d -> DictDoubleArray.of(d.dtype(), n, d.values(), codes(d.codes(), mask, offset, arena)); + case DictFloatArray d -> DictFloatArray.of(d.dtype(), n, d.values(), codes(d.codes(), mask, offset, arena)); + case DictFloat16Array d -> DictFloat16Array.of(d.dtype(), n, d.values(), codes(d.codes(), mask, offset, arena)); + case LongArray a -> new MaterializedLongArray(a.dtype(), n, longs(a, mask, offset, arena)); + case IntArray a -> new MaterializedIntArray(a.dtype(), n, ints(a, mask, offset, arena)); + case ShortArray a -> new MaterializedShortArray(a.dtype(), n, shorts(a, mask, offset, arena)); + case ByteArray a -> new MaterializedByteArray(a.dtype(), n, bytes(a, mask, offset, arena)); + case DoubleArray a -> new MaterializedDoubleArray(a.dtype(), n, doubles(a, mask, offset, arena)); + case FloatArray a -> new MaterializedFloatArray(a.dtype(), n, floats(a, mask, offset, arena)); + case Float16Array a -> new MaterializedFloat16Array(a.dtype(), n, halves(a, mask, offset, arena)); + case BoolArray a -> new MaterializedBoolArray(a.dtype(), n, bools(a, mask.indices(), offset, arena)); + case VarBinArray a -> strings(a, mask.indices(), offset, arena); + case DecimalArray a -> decimals(a, mask.indices(), offset, arena); case NullArray a -> new NullArray(a.dtype(), n); case StructArray s -> { List fields = new ArrayList<>(s.fieldCount()); for (int i = 0; i < s.fieldCount(); i++) { - fields.add(select(s.field(i), rows, offset, arena)); + fields.add(select(s.field(i), mask, offset, arena)); } yield new StructArray((DType.Struct) s.dtype(), n, fields); } case FixedSizeListArray a -> { - // a row's elements sit at row * fixedSize, so select those element rows + // a row's elements sit at row * fixedSize, so keep those element rows int size = a.fixedSize(); - int[] elementRows = new int[Math.multiplyExact(n, size)]; + if (size == 0) { + // rows of no elements: the element mask would keep nothing, so there is no copy + yield new FixedSizeListArray((DType.FixedSizeList) a.dtype(), n, Array.limited(a.elements(), 0)); + } + long elementCount = (long) mask.length() * size; + if (elementCount > Integer.MAX_VALUE) { + throw new VortexException("scan: " + mask.length() + " rows of " + size + + " elements each exceed a chunk's row limit"); + } + int[] rows = mask.indices(); + int[] elementRows = new int[n * size]; for (int k = 0; k < n; k++) { - int start = Math.multiplyExact(rows[k], size); + int start = rows[k] * size; for (int j = 0; j < size; j++) { elementRows[k * size + j] = start + j; } } + // keeping at least one row and not all, and each row at least one element, the + // element mask keeps some elements and not all + Mask.Values elements = (Mask.Values) Mask.Values.fromIndices((int) elementCount, elementRows); yield new FixedSizeListArray((DType.FixedSizeList) a.dtype(), n, - select(a.elements(), elementRows, offset * size, arena)); + select(a.elements(), elements, offset * size, arena)); } - case ListArray a -> listView(a, rows, offset, arena); + case ListArray a -> listView(a, mask.indices(), offset, arena); case ListViewArray a -> new ListViewArray((DType.List) a.dtype(), n, a.elements(), - select(a.offsets(), rows, offset, arena), select(a.sizes(), rows, offset, arena)); - case MapArray a -> new MapArray((DType.Map) a.dtype(), n, select(a.entries(), rows, offset, arena)); + select(a.offsets(), mask, offset, arena), select(a.sizes(), mask, offset, arena)); + case MapArray a -> new MapArray((DType.Map) a.dtype(), n, select(a.entries(), mask, offset, arena)); case UnionArray a -> { List variants = new ArrayList<>(a.variantCount()); for (int i = 0; i < a.variantCount(); i++) { - variants.add(select(a.variant(i), rows, offset, arena)); + variants.add(select(a.variant(i), mask, offset, arena)); } - yield new UnionArray((DType.Union) a.dtype(), n, select(a.typeIds(), rows, offset, arena), variants); + yield new UnionArray((DType.Union) a.dtype(), n, select(a.typeIds(), mask, offset, arena), variants); } default -> throw new VortexException( "scan: cannot select rows of " + full.getClass().getSimpleName()); @@ -118,22 +136,23 @@ yield new FixedSizeListArray((DType.FixedSizeList) a.dtype(), n, /// The selected codes in a flat buffer of their own width, the shape the dictionary lanes scan /// fastest; a code array is one byte to eight per row, so the copy is small. - private static Array codes(Array codes, int[] rows, long offset, SegmentAllocator arena) { - int n = rows.length; + private static Array codes(Array codes, Mask.Values mask, long offset, SegmentAllocator arena) { + int n = mask.trueCount(); return switch (codes) { - case ByteArray c -> new MaterializedByteArray(c.dtype(), n, bytes(c, rows, offset, arena)); - case ShortArray c -> new MaterializedShortArray(c.dtype(), n, shorts(c, rows, offset, arena)); - case IntArray c -> new MaterializedIntArray(c.dtype(), n, ints(c, rows, offset, arena)); - case LongArray c -> new MaterializedLongArray(c.dtype(), n, longs(c, rows, offset, arena)); + case ByteArray c -> new MaterializedByteArray(c.dtype(), n, bytes(c, mask, offset, arena)); + case ShortArray c -> new MaterializedShortArray(c.dtype(), n, shorts(c, mask, offset, arena)); + case IntArray c -> new MaterializedIntArray(c.dtype(), n, ints(c, mask, offset, arena)); + case LongArray c -> new MaterializedLongArray(c.dtype(), n, longs(c, mask, offset, arena)); default -> throw new VortexException("dictionary codes of type " + codes.dtype() + " are not integers"); }; } - private static MemorySegment longs(LongArray a, int[] rows, long offset, SegmentAllocator arena) { + private static MemorySegment longs(LongArray a, Mask.Values mask, long offset, SegmentAllocator arena) { MemorySegment src = flat(a); if (src != null) { - return copy(src, 8, rows, offset, arena); + return copy(src, 8, mask, offset, arena); } + int[] rows = mask.indices(); MemorySegment out = arena.allocate(Math.max(8, rows.length * 8L), 8); for (int k = 0; k < rows.length; k++) { out.setAtIndex(VortexFormat.LE_LONG, k, a.getLong(rows[k] + offset)); @@ -141,11 +160,12 @@ private static MemorySegment longs(LongArray a, int[] rows, long offset, Segment return out; } - private static MemorySegment ints(IntArray a, int[] rows, long offset, SegmentAllocator arena) { + private static MemorySegment ints(IntArray a, Mask.Values mask, long offset, SegmentAllocator arena) { MemorySegment src = flat(a); if (src != null) { - return copy(src, 4, rows, offset, arena); + return copy(src, 4, mask, offset, arena); } + int[] rows = mask.indices(); MemorySegment out = arena.allocate(Math.max(4, rows.length * 4L), 4); for (int k = 0; k < rows.length; k++) { out.setAtIndex(VortexFormat.LE_INT, k, a.getInt(rows[k] + offset)); @@ -153,11 +173,12 @@ private static MemorySegment ints(IntArray a, int[] rows, long offset, SegmentAl return out; } - private static MemorySegment shorts(ShortArray a, int[] rows, long offset, SegmentAllocator arena) { + private static MemorySegment shorts(ShortArray a, Mask.Values mask, long offset, SegmentAllocator arena) { MemorySegment src = flat(a); if (src != null) { - return copy(src, 2, rows, offset, arena); + return copy(src, 2, mask, offset, arena); } + int[] rows = mask.indices(); MemorySegment out = arena.allocate(Math.max(2, rows.length * 2L), 2); for (int k = 0; k < rows.length; k++) { out.setAtIndex(VortexFormat.LE_SHORT, k, a.getShort(rows[k] + offset)); @@ -165,11 +186,12 @@ private static MemorySegment shorts(ShortArray a, int[] rows, long offset, Segme return out; } - private static MemorySegment bytes(ByteArray a, int[] rows, long offset, SegmentAllocator arena) { + private static MemorySegment bytes(ByteArray a, Mask.Values mask, long offset, SegmentAllocator arena) { MemorySegment src = flat(a); if (src != null) { - return copy(src, 1, rows, offset, arena); + return copy(src, 1, mask, offset, arena); } + int[] rows = mask.indices(); MemorySegment out = arena.allocate(Math.max(1, rows.length)); for (int k = 0; k < rows.length; k++) { out.set(ValueLayout.JAVA_BYTE, k, a.getByte(rows[k] + offset)); @@ -177,11 +199,12 @@ private static MemorySegment bytes(ByteArray a, int[] rows, long offset, Segment return out; } - private static MemorySegment doubles(DoubleArray a, int[] rows, long offset, SegmentAllocator arena) { + private static MemorySegment doubles(DoubleArray a, Mask.Values mask, long offset, SegmentAllocator arena) { MemorySegment src = flat(a); if (src != null) { - return copy(src, 8, rows, offset, arena); + return copy(src, 8, mask, offset, arena); } + int[] rows = mask.indices(); MemorySegment out = arena.allocate(Math.max(8, rows.length * 8L), 8); for (int k = 0; k < rows.length; k++) { out.setAtIndex(VortexFormat.LE_DOUBLE, k, a.getDouble(rows[k] + offset)); @@ -189,11 +212,12 @@ private static MemorySegment doubles(DoubleArray a, int[] rows, long offset, Seg return out; } - private static MemorySegment floats(FloatArray a, int[] rows, long offset, SegmentAllocator arena) { + private static MemorySegment floats(FloatArray a, Mask.Values mask, long offset, SegmentAllocator arena) { MemorySegment src = flat(a); if (src != null) { - return copy(src, 4, rows, offset, arena); + return copy(src, 4, mask, offset, arena); } + int[] rows = mask.indices(); MemorySegment out = arena.allocate(Math.max(4, rows.length * 4L), 4); for (int k = 0; k < rows.length; k++) { out.setAtIndex(VortexFormat.LE_FLOAT, k, a.getFloat(rows[k] + offset)); @@ -202,11 +226,12 @@ private static MemorySegment floats(FloatArray a, int[] rows, long offset, Segme } /// Re-encodes each selected half from its widened value, which is exact: every half is a float. - private static MemorySegment halves(Float16Array a, int[] rows, long offset, SegmentAllocator arena) { + private static MemorySegment halves(Float16Array a, Mask.Values mask, long offset, SegmentAllocator arena) { MemorySegment src = flat(a); if (src != null) { - return copy(src, 2, rows, offset, arena); + return copy(src, 2, mask, offset, arena); } + int[] rows = mask.indices(); MemorySegment out = arena.allocate(Math.max(2, rows.length * 2L), 2); for (int k = 0; k < rows.length; k++) { out.setAtIndex(VortexFormat.LE_SHORT, k, Float.floatToFloat16(a.getFloat(rows[k] + offset))); @@ -282,28 +307,28 @@ private static MemorySegment flat(Array a) { return null; } - /// The `width`-byte elements of `src` at the selected rows, in a new buffer, the way Rust's - /// filter picks by density: a sparse selection gathers each row with a typed load and store; a - /// dense one copies each run of consecutive rows in bulk, so it costs a few copies, not one a row. - private static MemorySegment copy(MemorySegment src, int width, int[] rows, long offset, SegmentAllocator arena) { - int n = rows.length; + /// The `width`-byte elements of `src` at the rows `mask` keeps, in a new buffer: a dense mask + /// whose runs average at least [#MIN_AVERAGE_RUN] rows copies each run in bulk; any other + /// gathers each row with a typed load and store. Rust's filter uses the same thresholds but only + /// on forms already cached; here the first column builds them, and they are cached on the mask + /// for the next. The result is the same rows either way. + private static MemorySegment copy(MemorySegment src, int width, Mask.Values mask, long offset, + SegmentAllocator arena) { + int n = mask.trueCount(); MemorySegment out = arena.allocate(Math.max(width, (long) n * width), Math.min(width, 16)); - if (n == 0) { - return out; - } - if (n * 2L <= (long) rows[n - 1] - rows[0] + 1) { - gather(src, width, rows, offset, out); - return out; - } - int k = 0; - while (k < n) { - int end = k + 1; - while (end < n && rows[end] == rows[end - 1] + 1) { - end++; + if (mask.density() > 0.5) { + int[] slices = mask.slices(); + if (n / (slices.length / 2) >= MIN_AVERAGE_RUN) { + long at = 0; + for (int s = 0; s < slices.length; s += 2) { + long bytes = (long) (slices[s + 1] - slices[s]) * width; + MemorySegment.copy(src, (slices[s] + offset) * width, out, at, bytes); + at += bytes; + } + return out; } - MemorySegment.copy(src, (rows[k] + offset) * width, out, (long) k * width, (long) (end - k) * width); - k = end; } + gather(src, width, mask.indices(), offset, out); return out; } diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/compute/FusedFilterAggregate.java b/reader/src/main/java/io/github/dfa1/vortex/reader/compute/FusedFilterAggregate.java index 915a4a476..ffdf6dcc6 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/compute/FusedFilterAggregate.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/compute/FusedFilterAggregate.java @@ -11,6 +11,7 @@ import io.github.dfa1.vortex.reader.array.BoolArray; import io.github.dfa1.vortex.reader.array.DoubleArray; import io.github.dfa1.vortex.reader.array.FloatArray; +import io.github.dfa1.vortex.reader.array.Mask; import io.github.dfa1.vortex.reader.array.VarBinDictArray; import io.github.dfa1.vortex.reader.array.VarBinOffsetArray; import io.github.dfa1.vortex.reader.array.VarBinSlicedArray; @@ -722,6 +723,15 @@ public boolean test(long row) { return (bits[(int) (row >>> 6)] & 1L << row) != 0 && (residual == null || residual.test(row)); } + /// The bitmap itself when it is the whole filter, with no walk. + @Override + public Mask toMask(int length) { + if (residual == null && length == rows) { + return Mask.fromBitmap(bits, length); + } + return RowSelection.super.toMask(length); + } + @Override public long nextSelected(long from) { long row = from; diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/compute/RowSelection.java b/reader/src/main/java/io/github/dfa1/vortex/reader/compute/RowSelection.java index bcd58f139..318f929db 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/compute/RowSelection.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/compute/RowSelection.java @@ -1,5 +1,7 @@ package io.github.dfa1.vortex.reader.compute; +import io.github.dfa1.vortex.reader.array.Mask; + import java.util.function.LongPredicate; /// The rows of one chunk a filter selects: a row test, plus a walk to the next selected row, which a @@ -12,4 +14,17 @@ public interface RowSelection extends LongPredicate { /// @param from the zero-based row to start at /// @return the selected row, or `-1` when no row from `from` on is selected long nextSelected(long from); + + /// Returns the selected rows among the first `length` as a [Mask], the form Rust's filter + /// evaluates to: one walk sets a bit per selected row. + /// + /// @param length the chunk's row count + /// @return the mask of the selected rows + default Mask toMask(int length) { + long[] words = new long[(length + 63) >>> 6]; + for (long row = nextSelected(0); row >= 0 && row < length; row = nextSelected(row + 1)) { + words[(int) (row >>> 6)] |= 1L << row; + } + return Mask.fromBitmap(words, length); + } } diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/array/MaskTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/array/MaskTest.java new file mode 100644 index 000000000..a08ebc887 --- /dev/null +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/array/MaskTest.java @@ -0,0 +1,249 @@ +package io.github.dfa1.vortex.reader.array; + +import org.junit.jupiter.api.Nested; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.stream.IntStream; +import java.util.stream.LongStream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/// A filtered scan copies each column through these forms, so a wrong index or run here returns +/// the wrong rows. Rows 0, 63, 64 and the last row sit on 64-bit word edges, where a shift or a +/// word index off by one shows. +class MaskTest { + + @Nested + class Construction { + + @Test + void aBitmapKeepingNoRow_isAllFalse() { + // Given + long[] words = new long[2]; + + // When + Mask result = Mask.fromBitmap(words, 100); + + // Then + assertThat(result).isEqualTo(new Mask.AllFalse(100)); + } + + @Test + void aBitmapKeepingEveryRow_isAllTrue() { + // Given: 70 rows, the last word partly filled + long[] words = {-1L, (1L << 6) - 1}; + + // When + Mask result = Mask.fromBitmap(words, 70); + + // Then + assertThat(result).isEqualTo(new Mask.AllTrue(70)); + } + + @Test + void bitsPastTheLastRow_areNotKeptRows() { + // Given: rows 0 and 1 of 3 kept, plus stray bits past row 2 that are no rows at all; + // counted, they would make the mask keep more rows than it has + long[] words = {0b11L | -1L << 3}; + + // When + Mask result = Mask.fromBitmap(words, 3); + + // Then + assertThat(result.trueCount()).isEqualTo(2); + assertThat(((Mask.Values) result).indices()).containsExactly(0, 1); + } + + @Test + void aBitmapWithTheWrongWordCount_isRejected() { + // Given: 65 rows need two words + long[] words = new long[1]; + + // When / Then + assertThatThrownBy(() -> Mask.fromBitmap(words, 65)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void indicesKeepingSomeRows_keepTheirOrderAndCount() { + // Given + int[] rows = {0, 63, 64, 129}; + + // When + Mask result = Mask.Values.fromIndices(130, rows); + + // Then + assertThat(result).isInstanceOf(Mask.Values.class); + Mask.Values values = (Mask.Values) result; + assertThat(values.trueCount()).isEqualTo(4); + assertThat(values.indices()).containsExactly(0, 63, 64, 129); + assertThat(values.density()).isEqualTo(4 / 130.0); + } + + @Test + void noIndices_isAllFalseAndEveryIndex_isAllTrue() { + // Given + int[] none = {}; + int[] all = {0, 1, 2}; + + // When + Mask noneResult = Mask.Values.fromIndices(3, none); + Mask allResult = Mask.Values.fromIndices(3, all); + + // Then + assertThat(noneResult).isEqualTo(new Mask.AllFalse(3)); + assertThat(allResult).isEqualTo(new Mask.AllTrue(3)); + } + } + + @Nested + class DerivedForms { + + @Test + void indices_walkTheBitmapAcrossWords() { + // Given: the same rows as a bitmap, so the indices come from the word walk + long[] words = new long[3]; + for (int row : new int[]{0, 63, 64, 129}) { + words[row >>> 6] |= 1L << row; + } + Mask.Values sut = (Mask.Values) Mask.fromBitmap(words, 130); + + // When + int[] result = sut.indices(); + + // Then + assertThat(result).containsExactly(0, 63, 64, 129); + assertThat(sut.indices()).isSameAs(result); + } + + @Test + void slices_coverRunsAcrossAWordEdgeAndEndingAtTheLastRow() { + // Given: rows 2, 60..69 (crossing the word edge at 64) and 125..129 (the last row) + int[] rows = IntStream.concat(IntStream.of(2), + IntStream.concat(IntStream.range(60, 70), IntStream.range(125, 130))).toArray(); + Mask.Values sut = (Mask.Values) Mask.Values.fromIndices(130, rows); + + // When + int[] result = sut.slices(); + + // Then + assertThat(result).containsExactly(2, 3, 60, 70, 125, 130); + assertThat(sut.slices()).isSameAs(result); + } + } + + @Nested + class Limit { + + @Test + void aLimitAtOrAboveTheKeptRows_keepsTheMask() { + // Given + Mask sut = Mask.Values.fromIndices(10, new int[]{1, 5}); + + // When + Mask result = sut.limit(2); + + // Then + assertThat(result).isSameAs(sut); + } + + @Test + void allTrue_keepsTheFirstRowsAcrossAWordEdge() { + // Given + Mask sut = new Mask.AllTrue(130); + + // When + Mask result = sut.limit(70); + + // Then + assertThat(result.trueCount()).isEqualTo(70); + assertThat(((Mask.Values) result).indices()).isEqualTo(IntStream.range(0, 70).toArray()); + } + + @Test + void values_keepTheFirstKeptRowsCuttingInsideAWord() { + // Given: the cut falls inside the second word, between rows 64 and 100 + Mask sut = Mask.Values.fromIndices(130, new int[]{3, 63, 64, 100, 129}); + + // When + Mask result = sut.limit(3); + + // Then + assertThat(result.length()).isEqualTo(130); + assertThat(((Mask.Values) result).indices()).containsExactly(3, 63, 64); + } + + @Test + void aZeroLimit_keepsNoRow() { + // Given + Mask sut = Mask.Values.fromIndices(10, new int[]{1, 5}); + + // When + Mask result = sut.limit(0); + + // Then + assertThat(result).isEqualTo(new Mask.AllFalse(10)); + } + } + + /// Random lengths around the word edges and random densities, so runs start and end on and + /// next to every word boundary; each form is checked against a plain list of the kept rows. + @ParameterizedTest(name = "seed {0}") + @MethodSource("seeds") + void everyForm_agreesWithTheKeptRows(long seed) { + // Given + Random random = new Random(seed); + int length = 1 + random.nextInt(300); + double density = random.nextDouble(); + long[] words = new long[(length + 63) >>> 6]; + List kept = new ArrayList<>(); + for (int row = 0; row < length; row++) { + if (random.nextDouble() < density) { + words[row >>> 6] |= 1L << row; + kept.add(row); + } + } + int limit = random.nextInt(length + 1); + + // When + Mask result = Mask.fromBitmap(words, length); + Mask limited = result.limit(limit); + + // Then + assertThat(result.trueCount()).isEqualTo(kept.size()); + assertThat(rowsOf(result)).isEqualTo(kept); + assertThat(rowsOf(limited)).isEqualTo(kept.subList(0, Math.min(limit, kept.size()))); + if (result instanceof Mask.Values values) { + List fromSlices = new ArrayList<>(); + int[] slices = values.slices(); + for (int s = 0; s < slices.length; s += 2) { + assertThat(slices[s]).as("empty run").isLessThan(slices[s + 1]); + if (s > 0) { + assertThat(slices[s]).as("touching runs").isGreaterThan(slices[s - 1]); + } + for (int row = slices[s]; row < slices[s + 1]; row++) { + fromSlices.add(row); + } + } + assertThat(fromSlices).isEqualTo(kept); + } + } + + static LongStream seeds() { + return LongStream.range(0, 30); + } + + private static List rowsOf(Mask mask) { + return switch (mask) { + case Mask.AllTrue(int length) -> IntStream.range(0, length).boxed().toList(); + case Mask.AllFalse _ -> List.of(); + case Mask.Values values -> IntStream.of(values.indices()).boxed().toList(); + }; + } +} diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/array/SelectedArraysTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/array/SelectedArraysTest.java index 032901c92..c671d29e5 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/array/SelectedArraysTest.java +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/array/SelectedArraysTest.java @@ -9,6 +9,7 @@ import java.lang.foreign.Arena; import java.lang.foreign.MemorySegment; +import java.lang.foreign.SegmentAllocator; import java.lang.foreign.ValueLayout; import java.math.BigDecimal; import java.math.BigInteger; @@ -33,16 +34,16 @@ class SelectedArraysTest { @Test void primitives_copyTheSelectedRowsIntoFlatBuffers() { // When - LongArray longs = (LongArray) SelectedArrays.select(TestArrays.longs(10, 11, 12, 13), ROWS, ARENA); - IntArray ints = (IntArray) SelectedArrays.select(TestArrays.ints(10, 11, 12, 13), ROWS, ARENA); - ShortArray shorts = (ShortArray) SelectedArrays.select( + LongArray longs = (LongArray) select(TestArrays.longs(10, 11, 12, 13), ROWS, ARENA); + IntArray ints = (IntArray) select(TestArrays.ints(10, 11, 12, 13), ROWS, ARENA); + ShortArray shorts = (ShortArray) select( TestArrays.shorts((short) 10, (short) 11, (short) 12, (short) 13), ROWS, ARENA); - ByteArray bytes = (ByteArray) SelectedArrays.select( + ByteArray bytes = (ByteArray) select( TestArrays.bytes((byte) 10, (byte) 11, (byte) 12, (byte) 13), ROWS, ARENA); - DoubleArray doubles = (DoubleArray) SelectedArrays.select(TestArrays.doubles(1.5, 2.5, 3.5, 4.5), ROWS, ARENA); - FloatArray floats = (FloatArray) SelectedArrays.select(TestArrays.floats(1.5f, 2.5f, 3.5f, 4.5f), ROWS, ARENA); - Float16Array halves = (Float16Array) SelectedArrays.select(TestArrays.float16(1.5f, 2.5f, 3.5f, 4.5f), ROWS, ARENA); - BoolArray bools = (BoolArray) SelectedArrays.select(TestArrays.bools(false, true, false, true), ROWS, ARENA); + DoubleArray doubles = (DoubleArray) select(TestArrays.doubles(1.5, 2.5, 3.5, 4.5), ROWS, ARENA); + FloatArray floats = (FloatArray) select(TestArrays.floats(1.5f, 2.5f, 3.5f, 4.5f), ROWS, ARENA); + Float16Array halves = (Float16Array) select(TestArrays.float16(1.5f, 2.5f, 3.5f, 4.5f), ROWS, ARENA); + BoolArray bools = (BoolArray) select(TestArrays.bools(false, true, false, true), ROWS, ARENA); // Then assertThat(longs.length()).isEqualTo(3); @@ -69,12 +70,12 @@ void sparseSelection_gathersEachRowAtEveryWidth() { int[] sparse = rows(0, 5, 9); // When - LongArray longs = (LongArray) SelectedArrays.select(TestArrays.longs(0, 1, 2, 3, 4, 5, 6, 7, 8, 9), sparse, ARENA); - IntArray ints = (IntArray) SelectedArrays.select(TestArrays.ints(0, 1, 2, 3, 4, 5, 6, 7, 8, 9), sparse, ARENA); - ShortArray shorts = (ShortArray) SelectedArrays.select(TestArrays.shorts( + LongArray longs = (LongArray) select(TestArrays.longs(0, 1, 2, 3, 4, 5, 6, 7, 8, 9), sparse, ARENA); + IntArray ints = (IntArray) select(TestArrays.ints(0, 1, 2, 3, 4, 5, 6, 7, 8, 9), sparse, ARENA); + ShortArray shorts = (ShortArray) select(TestArrays.shorts( (short) 0, (short) 1, (short) 2, (short) 3, (short) 4, (short) 5, (short) 6, (short) 7, (short) 8, (short) 9), sparse, ARENA); - ByteArray bytes = (ByteArray) SelectedArrays.select(TestArrays.bytes( + ByteArray bytes = (ByteArray) select(TestArrays.bytes( (byte) 0, (byte) 1, (byte) 2, (byte) 3, (byte) 4, (byte) 5, (byte) 6, (byte) 7, (byte) 8, (byte) 9), sparse, ARENA); @@ -96,7 +97,7 @@ void broadcastBuffer_copiesRowByRowNotByRuns() { Array constant = new MaterializedLongArray(DType.I64, 4, one); // When - LongArray result = (LongArray) SelectedArrays.select(constant, rows(1, 2, 3), ARENA); + LongArray result = (LongArray) select(constant, rows(1, 2, 3), ARENA); // Then assertThat(List.of(result.getLong(0), result.getLong(1), result.getLong(2))).isEqualTo(List.of(42L, 42L, 42L)); @@ -115,7 +116,7 @@ void bools_packEverySelectedBitAcrossBytes() { } // When - BoolArray result = (BoolArray) SelectedArrays.select(TestArrays.bools(values), selected, ARENA); + BoolArray result = (BoolArray) select(TestArrays.bools(values), selected, ARENA); // Then assertThat(result.length()).isEqualTo(19); @@ -128,7 +129,7 @@ void bools_packEverySelectedBitAcrossBytes() { @Test void materialize_writesTheSelectedRowsInOrder() { // Given - Array selected = SelectedArrays.select(TestArrays.longs(10, 11, 12, 13), ROWS, ARENA); + Array selected = select(TestArrays.longs(10, 11, 12, 13), ROWS, ARENA); // When MemorySegment result = selected.materialize(ARENA); @@ -140,7 +141,7 @@ void materialize_writesTheSelectedRowsInOrder() { @Test void float16Materialize_keepsTheHalfBits() { // Given - Float16Array selected = (Float16Array) SelectedArrays.select(TestArrays.float16(1.5f, 2.5f, 3.5f), rows(0, 2), ARENA); + Float16Array selected = (Float16Array) select(TestArrays.float16(1.5f, 2.5f, 3.5f), rows(0, 2), ARENA); // When MemorySegment result = selected.materialize(ARENA); @@ -156,7 +157,7 @@ void masked_selectsValuesAndValidityTogether() { Array masked = new MaskedArray(TestArrays.longs(10, 11, 12, 13), TestArrays.bools(true, false, true, true)); // When - MaskedArray result = (MaskedArray) SelectedArrays.select(masked, rows(1, 2), ARENA); + MaskedArray result = (MaskedArray) select(masked, rows(1, 2), ARENA); // Then assertThat(result.isValid(0)).isFalse(); @@ -167,7 +168,7 @@ void masked_selectsValuesAndValidityTogether() { @Test void strings_copyTheSelectedRows() { // When - VarBinArray result = (VarBinArray) SelectedArrays.select(strings("a", "bb", "ccc", "dddd"), ROWS, ARENA); + VarBinArray result = (VarBinArray) select(strings("a", "bb", "ccc", "dddd"), ROWS, ARENA); // Then assertThat(List.of(result.getString(0), result.getString(1), result.getString(2))) @@ -187,7 +188,7 @@ void stringDictionary_staysADictionaryOverTheSamePool() { Array dict = stringDict(new String[]{"x", "y"}, new byte[]{1, 0, 1, 0}); // When - Array result = SelectedArrays.select(dict, rows(0, 1), ARENA); + Array result = select(dict, rows(0, 1), ARENA); // Then assertThat(result).isInstanceOf(VarBinDictArray.class); @@ -204,7 +205,7 @@ void primitiveDictionary_staysADictionaryWithFlatCodes() { Array dict = DictLongArray.of(DType.I64, 4, pool, codes); // When - Array result = SelectedArrays.select(dict, rows(2, 3), ARENA); + Array result = select(dict, rows(2, 3), ARENA); // Then assertThat(result).isInstanceOf(DictLongArray.class); @@ -215,18 +216,78 @@ void primitiveDictionary_staysADictionaryWithFlatCodes() { @Test void offsetSlice_selectsFromTheArrayItSlicesSoADictionaryStaysOne() { - // Given: a slice starting at row 2 of a four-row dictionary + // Given: a slice starting at row 1 of a four-row dictionary; keeping slice rows 0 and 2 + // reads dictionary rows 1 and 3 only if the slice offset is applied Array dict = DictLongArray.of(DType.I64, 4, TestArrays.longs(100, 200), - codes((byte) 0, (byte) 0, (byte) 1, (byte) 0)); - Array slice = new OffsetLongArray(DType.I64, 2, (LongArray) dict, 2); + codes((byte) 0, (byte) 1, (byte) 0, (byte) 1)); + Array slice = new OffsetLongArray(DType.I64, 3, (LongArray) dict, 1); // When - Array result = SelectedArrays.select(slice, rows(0, 1), ARENA); + Array result = select(slice, rows(0, 2), ARENA); // Then assertThat(result).isInstanceOf(DictLongArray.class); LongArray selected = (LongArray) result; - assertThat(List.of(selected.getLong(0), selected.getLong(1))).isEqualTo(List.of(200L, 100L)); + assertThat(List.of(selected.getLong(0), selected.getLong(1))).isEqualTo(List.of(200L, 200L)); + } + + @Test + void allTrueMask_returnsTheArrayItself() { + // Given + Array full = TestArrays.longs(1, 2, 3); + + // When + Array result = SelectedArrays.select(full, new Mask.AllTrue(3), ARENA); + + // Then: no copy when every row is kept, as Rust's filter returns the array unchanged + assertThat(result).isSameAs(full); + } + + @Test + void denseSelectionWithLongRuns_copiesEachRunInBulk() { + // Given: 100 rows, all kept but rows 10 and 50: density 0.98, three runs averaging 32 + // rows, so the flat copy takes the run path; a run boundary off by one would show + long[] values = new long[100]; + int[] kept = new int[98]; + int k = 0; + for (int i = 0; i < values.length; i++) { + values[i] = i; + if (i != 10 && i != 50) { + kept[k++] = i; + } + } + + // When + LongArray result = (LongArray) select(TestArrays.longs(values), kept, ARENA); + + // Then + assertThat(result.length()).isEqualTo(98); + for (int j = 0; j < 98; j++) { + assertThat(result.getLong(j)).as("kept row %d", j).isEqualTo(kept[j]); + } + } + + @Test + void denseSelectionWithShortRuns_gathersEachRow() { + // Given: every row but each third kept: density 2/3 but runs of two, under the average + // run of eight a bulk copy needs, so the copy gathers row by row + long[] values = new long[30]; + int[] kept = new int[20]; + int k = 0; + for (int i = 0; i < values.length; i++) { + values[i] = 100 + i; + if (i % 3 != 0) { + kept[k++] = i; + } + } + + // When + LongArray result = (LongArray) select(TestArrays.longs(values), kept, ARENA); + + // Then + for (int j = 0; j < 20; j++) { + assertThat(result.getLong(j)).as("kept row %d", j).isEqualTo(100 + kept[j]); + } } @Test @@ -238,7 +299,7 @@ void offsetSliceOfAFlatArray_copiesFromItsShiftedRows() { Array masked = new MaskedArray(values, validity); // When - MaskedArray result = (MaskedArray) SelectedArrays.select(masked, rows(1, 2), ARENA); + MaskedArray result = (MaskedArray) select(masked, rows(1, 2), ARENA); // Then: slice rows 1 and 2 are rows 2 and 3 of the sliced arrays assertThat(((IntArray) result.inner()).getInt(1)).isEqualTo(13); @@ -253,7 +314,7 @@ void struct_selectsEveryField() { Array struct = new StructArray(dtype, 3, List.of(TestArrays.longs(1, 2, 3), strings("p", "q", "r"))); // When - StructArray result = (StructArray) SelectedArrays.select(struct, rows(0, 2), ARENA); + StructArray result = (StructArray) select(struct, rows(0, 2), ARENA); // Then assertThat(result.length()).isEqualTo(2); @@ -261,6 +322,21 @@ void struct_selectsEveryField() { assertThat(((VarBinArray) result.field(1)).getString(1)).isEqualTo("r"); } + @Test + void fixedSizeListOfNoElements_keepsTheSelectedRowCount() { + // Given: three rows of zero elements each, which Rust accepts; the element mask keeps + // nothing, so a copy that assumed some kept elements would fail + DType.FixedSizeList dtype = new DType.FixedSizeList(DType.I64, 0, false); + Array list = new FixedSizeListArray(dtype, 3, TestArrays.longs()); + + // When + FixedSizeListArray result = (FixedSizeListArray) select(list, rows(0, 2), ARENA); + + // Then + assertThat(result.length()).isEqualTo(2); + assertThat(result.elements().length()).isZero(); + } + @Test void fixedSizeList_selectsEachRowsElements() { // Given: three rows of two elements: [1,2] [3,4] [5,6] @@ -268,7 +344,7 @@ void fixedSizeList_selectsEachRowsElements() { Array list = new FixedSizeListArray(dtype, 3, TestArrays.longs(1, 2, 3, 4, 5, 6)); // When - FixedSizeListArray result = (FixedSizeListArray) SelectedArrays.select(list, rows(0, 2), ARENA); + FixedSizeListArray result = (FixedSizeListArray) select(list, rows(0, 2), ARENA); // Then LongArray elements = (LongArray) result.elements(); @@ -284,7 +360,7 @@ void list_becomesAListViewOverTheSameElements() { Array list = new ListArray(dtype, 3, TestArrays.longs(10, 20, 21, 30, 31, 32), TestArrays.ints(0, 1, 3, 6)); // When - ListViewArray result = (ListViewArray) SelectedArrays.select(list, rows(1, 2), ARENA); + ListViewArray result = (ListViewArray) select(list, rows(1, 2), ARENA); // Then: row 0 starts at element 1 and holds 2; row 1 starts at 3 and holds 3 assertThat(((LongArray) result.offsets()).getLong(0)).isEqualTo(1); @@ -305,7 +381,7 @@ void decimal_readsAndMaterializesTheSelectedRows() { Array decimals = new LazyDecimalArray(dtype, 3, buf, 4); // When - DecimalArray result = (DecimalArray) SelectedArrays.select(decimals, rows(0, 2), ARENA); + DecimalArray result = (DecimalArray) select(decimals, rows(0, 2), ARENA); // Then assertThat(result.getDecimal(0)).isEqualTo(new BigDecimal("1.00")); @@ -333,7 +409,7 @@ void wideDecimal_signExtendsEachValueToSixteenBytes() { Array decimals = new LazyDecimalArray(dtype, values.length, buf, 16); // When - DecimalArray result = (DecimalArray) SelectedArrays.select(decimals, rows(0, 2), ARENA); + DecimalArray result = (DecimalArray) select(decimals, rows(0, 2), ARENA); // Then assertThat(result.getDecimal(0)).isEqualTo(new BigDecimal(values[0])); @@ -343,7 +419,7 @@ void wideDecimal_signExtendsEachValueToSixteenBytes() { @Test void nullArray_keepsItsTypeWithTheSelectedLength() { // When - Array result = SelectedArrays.select(new NullArray(DType.I64.asNullable(), 10), ROWS, ARENA); + Array result = select(new NullArray(DType.I64.asNullable(), 10), ROWS, ARENA); // Then assertThat(result).isEqualTo(new NullArray(DType.I64.asNullable(), 3)); @@ -352,7 +428,7 @@ void nullArray_keepsItsTypeWithTheSelectedLength() { @Test void anEmptySelection_hasNoRows() { // When - Array result = SelectedArrays.select(TestArrays.longs(1, 2, 3), new int[0], ARENA); + Array result = select(TestArrays.longs(1, 2, 3), new int[0], ARENA); // Then assertThat(result.length()).isZero(); @@ -365,7 +441,7 @@ void anArrayWithoutRows_failsAsVortexException() { new MemorySegment[0], new Array[0]); // When / Then - assertThatThrownBy(() -> SelectedArrays.select(unknown, rows(0), ARENA)) + assertThatThrownBy(() -> select(unknown, rows(0), ARENA)) .isInstanceOf(VortexException.class) .hasMessageContaining("cannot select rows"); } @@ -374,6 +450,11 @@ private static int[] rows(int... rows) { return rows; } + /// Selects `rows` of `full` through the mask a filter would build from them. + private static Array select(Array full, int[] rows, SegmentAllocator arena) { + return SelectedArrays.select(full, Mask.Values.fromIndices((int) full.length(), rows), arena); + } + private static VarBinArray strings(String... values) { byte[] all = String.join("", values).getBytes(StandardCharsets.UTF_8); MemorySegment offsets = ARENA.allocate((values.length + 1) * 4L, 4); diff --git a/reader/src/test/java/io/github/dfa1/vortex/reader/compute/ComputeTest.java b/reader/src/test/java/io/github/dfa1/vortex/reader/compute/ComputeTest.java index 9297e5501..e8acf2fdb 100644 --- a/reader/src/test/java/io/github/dfa1/vortex/reader/compute/ComputeTest.java +++ b/reader/src/test/java/io/github/dfa1/vortex/reader/compute/ComputeTest.java @@ -8,8 +8,12 @@ import io.github.dfa1.vortex.reader.ScanFilter; import io.github.dfa1.vortex.reader.array.Array; import io.github.dfa1.vortex.reader.array.LazyDecimalBytePartsArray; +import io.github.dfa1.vortex.reader.array.LongArray; +import io.github.dfa1.vortex.reader.array.Mask; import io.github.dfa1.vortex.reader.array.MaskedArray; +import io.github.dfa1.vortex.reader.array.MaterializedLongArray; import io.github.dfa1.vortex.reader.array.NullArray; +import io.github.dfa1.vortex.reader.array.SelectedArrays; import io.github.dfa1.vortex.reader.array.VarBinArray; import io.github.dfa1.vortex.reader.array.VarBinDictArray; import io.github.dfa1.vortex.reader.array.VarBinSlicedArray; @@ -21,6 +25,8 @@ import org.junit.jupiter.params.provider.MethodSource; import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; import java.math.BigDecimal; import java.util.List; import java.util.Map; @@ -120,6 +126,31 @@ static Stream predicates() { new Predicate.Eq("AAPL"))), codes, List.of(1L, 3L, 4L)))); } + /// u8 and u16 codes hand their bitmap to the mask with no walk; i32 codes and the null + /// tests go through the default walk: both must keep the rows the row test selects. + @ParameterizedTest(name = "{0} over {1} codes") + @MethodSource("predicates") + void toMask_keepsTheRowsTheMatcherSelects(Predicate predicate, PType codesPType, List expected) { + // Given: the row numbers as a column, so the selected column reads back the kept rows + Chunk chunk = mock(Chunk.class); + given(chunk.column(SYM)).willReturn(nullableDict(codesPType)); + given(chunk.rowCount()).willReturn(6L); + RowSelection sut = Compute.matcher(chunk, new ScanFilter.Column(SYM, predicate)); + MemorySegment rowNumbers = ARENA.allocate(6 * 8L, 8); + for (int i = 0; i < 6; i++) { + rowNumbers.setAtIndex(ValueLayout.JAVA_LONG, i, i); + } + Array column = new MaterializedLongArray(DType.I64, 6, rowNumbers); + + // When + Mask result = sut.toMask(6); + + // Then + LongArray kept = (LongArray) SelectedArrays.select(column, result, ARENA); + assertThat(result.trueCount()).isEqualTo(expected.size()); + assertThat(LongStream.range(0, kept.length()).map(kept::getLong).boxed().toList()).isEqualTo(expected); + } + @ParameterizedTest(name = "{0} over {1} codes") @MethodSource("predicates") void selectsTheRowsMatchesSelects(Predicate predicate, PType codesPType, List expected) {