diff --git a/CHANGELOG.md b/CHANGELOG.md index 71e226c11..1ccd04e8b 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 +- 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)) - 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 28082493f..088fe0997 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 @@ -113,12 +113,12 @@ public Chunk slice(long offset, long length) { /// 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, ascending, each in `[0, rowCount())` + /// @param rows the rows to keep, strictly ascending, each in `[0, rowCount())` /// @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(long[] rows, List names, Map positions, Consumer onClose) { + Chunk select(int[] 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()); 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 34110300a..8cbf4e00e 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 @@ -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 long[] pendingRows; + private int[] pendingRows; private List projectedNames; private List projectedDtypes; private Map projectedPositions; @@ -364,11 +364,16 @@ private boolean selectMatchingRows(ChunkSpec spec) { Chunk full = chunkFor(spec, spec.rowCount()); try { RowSelection selection = Compute.matcher(full, filter); - long[] rows = new long[(int) Math.min(spec.rowCount(), limit - rowsReturned)]; + // 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 < rows.length; + for (long row = selection.nextSelected(0); row >= 0 && count < allowed; row = selection.nextSelected(row + 1)) { - rows[count++] = row; + if (count == rows.length) { + rows = Arrays.copyOf(rows, (int) Math.min(allowed, 2L * count)); + } + rows[count++] = (int) row; } if (count == 0) { full.close(); @@ -421,7 +426,7 @@ public Chunk next() { Chunk chunk; if (pending != null) { chunk = pending; - long[] rows = pendingRows; + int[] rows = pendingRows; pending = null; pendingRows = null; chunk = selectedChunk(chunk, rows); @@ -437,7 +442,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, long[] rows) { + private Chunk selectedChunk(Chunk full, int[] rows) { List names = outputNames; Map positions = outputPositions; if (singleColumnIsSyntheticWrapper) { 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 c8e7e9361..5aeb27493 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 @@ -3,20 +3,21 @@ import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.core.io.VortexFormat; import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.core.model.PType; import java.lang.foreign.MemorySegment; import java.lang.foreign.SegmentAllocator; import java.lang.foreign.ValueLayout; -import java.math.BigDecimal; import java.util.ArrayList; -import java.util.Arrays; import java.util.List; -import java.util.function.IntConsumer; -import java.util.function.LongBinaryOperator; -/// Row selection: an array of some rows of another array, in the given order. Leaf arrays get an -/// index-mapped view, `get(i) = inner.get(rows[i])`, with no copy; composite arrays select their -/// children. This is how a filtered scan returns only the rows that match, as Rust's scan does. +/// 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. public final class SelectedArrays { private SelectedArrays() { @@ -25,424 +26,345 @@ private SelectedArrays() { /// The rows `rows[0] .. rows[rows.length - 1]` of `full`, in that order. /// /// @param full the array to select from - /// @param rows row indices into `full`, each in `[0, full.length())` - /// @param arena allocator for the per-row offsets a list selection rebuilds + /// @param rows row indices into `full`, strictly ascending, each in `[0, full.length())` + /// @param arena allocator for the copied rows /// @return an array of `rows.length` rows with `full`'s dtype /// @throws VortexException if `full`'s array type has no row selection - public static Array select(Array full, long[] rows, SegmentAllocator arena) { - long n = rows.length; + public static Array select(Array full, int[] rows, SegmentAllocator arena) { + return select(full, rows, 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; return switch (full) { case MaskedArray m -> { BoolArray validity = null; if (m.validity() != null) { - validity = (BoolArray) select(m.validity(), rows, arena); + validity = (BoolArray) select(m.validity(), rows, offset, arena); } - yield new MaskedArray(select(m.inner(), rows, arena), validity); + yield new MaskedArray(select(m.inner(), rows, offset, arena), validity); } - // A slice selects from the array it slices, so a dictionary under it stays one. - case OffsetLongArray o -> select(o.inner(), shifted(rows, o.offset()), arena); - case OffsetIntArray o -> select(o.inner(), shifted(rows, o.offset()), arena); - case OffsetShortArray o -> select(o.inner(), shifted(rows, o.offset()), arena); - case OffsetByteArray o -> select(o.inner(), shifted(rows, o.offset()), arena); - case OffsetDoubleArray o -> select(o.inner(), shifted(rows, o.offset()), arena); - case OffsetFloatArray o -> select(o.inner(), shifted(rows, o.offset()), arena); - case OffsetBoolArray o -> select(o.inner(), shifted(rows, o.offset()), arena); - case VarBinSlicedArray s -> select(s.inner(), shifted(rows, s.offset()), arena); + 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); // 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(), gather(d.dictCodesSegs(), d.dictCodesPType().byteSize(), rows, arena), + d.dictValOffPType(), copy(d.dictCodesSegs(), d.dictCodesPType().byteSize(), rows, offset, arena), d.dictCodesPType()); - case DictLongArray d -> DictLongArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, arena)); - case DictIntArray d -> DictIntArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, arena)); - case DictShortArray d -> DictShortArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, arena)); - case DictByteArray d -> DictByteArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, arena)); - case DictDoubleArray d -> DictDoubleArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, arena)); - case DictFloatArray d -> DictFloatArray.of(d.dtype(), n, d.values(), codes(d.codes(), rows, arena)); - case DictFloat16Array d -> DictFloat16Array.of(d.dtype(), n, d.values(), codes(d.codes(), rows, arena)); - case LongArray a -> new Longs(a, rows); - case IntArray a -> new Ints(a, rows); - case ShortArray a -> new Shorts(a, rows); - case ByteArray a -> new Bytes(a, rows); - case DoubleArray a -> new Doubles(a, rows); - case FloatArray a -> new Floats(a, rows); - case Float16Array a -> new Float16s(a, rows); - case BoolArray a -> new Bools(a, rows); - case VarBinArray a -> new VarBins(a, rows); - case DecimalArray a -> new Decimals(a, rows); + 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 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, arena)); + fields.add(select(s.field(i), rows, 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 int size = a.fixedSize(); - long[] elementRows = new long[Math.multiplyExact(rows.length, size)]; - for (int k = 0; k < rows.length; k++) { + int[] elementRows = new int[Math.multiplyExact(n, size)]; + for (int k = 0; k < n; k++) { + int start = Math.multiplyExact(rows[k], size); for (int j = 0; j < size; j++) { - elementRows[k * size + j] = rows[k] * size + j; + elementRows[k * size + j] = start + j; } } - yield new FixedSizeListArray((DType.FixedSizeList) a.dtype(), n, select(a.elements(), elementRows, arena)); + yield new FixedSizeListArray((DType.FixedSizeList) a.dtype(), n, + select(a.elements(), elementRows, offset * size, arena)); } - case ListArray a -> listView(a, rows, arena); + case ListArray a -> listView(a, rows, offset, arena); case ListViewArray a -> new ListViewArray((DType.List) a.dtype(), n, a.elements(), - select(a.offsets(), rows, arena), select(a.sizes(), rows, arena)); - case MapArray a -> new MapArray((DType.Map) a.dtype(), n, select(a.entries(), rows, arena)); + 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)); case UnionArray a -> { List variants = new ArrayList<>(a.variantCount()); for (int i = 0; i < a.variantCount(); i++) { - variants.add(select(a.variant(i), rows, arena)); + variants.add(select(a.variant(i), rows, offset, arena)); } - yield new UnionArray((DType.Union) a.dtype(), n, select(a.typeIds(), rows, arena), variants); + yield new UnionArray((DType.Union) a.dtype(), n, select(a.typeIds(), rows, offset, arena), variants); } default -> throw new VortexException( "scan: cannot select rows of " + full.getClass().getSimpleName()); }; } - private static long[] shifted(long[] rows, long offset) { - long[] out = new long[rows.length]; - for (int i = 0; i < rows.length; i++) { - out[i] = rows[i] + offset; - } - return out; - } - /// 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, long[] rows, SegmentAllocator arena) { - long n = rows.length; + private static Array codes(Array codes, int[] rows, long offset, SegmentAllocator arena) { + int n = rows.length; return switch (codes) { - case ByteArray c -> { - MemorySegment out = arena.allocate(Math.max(1, n)); - for (int i = 0; i < rows.length; i++) { - out.set(ValueLayout.JAVA_BYTE, i, c.getByte(rows[i])); - } - yield new MaterializedByteArray(c.dtype(), n, out); - } - case ShortArray c -> { - MemorySegment out = arena.allocate(Math.max(2, n * 2), 2); - for (int i = 0; i < rows.length; i++) { - out.setAtIndex(VortexFormat.LE_SHORT, i, c.getShort(rows[i])); - } - yield new MaterializedShortArray(c.dtype(), n, out); - } - case IntArray c -> { - MemorySegment out = arena.allocate(Math.max(4, n * 4), 4); - for (int i = 0; i < rows.length; i++) { - out.setAtIndex(VortexFormat.LE_INT, i, c.getInt(rows[i])); - } - yield new MaterializedIntArray(c.dtype(), n, out); - } - case LongArray c -> { - MemorySegment out = arena.allocate(Math.max(8, n * 8), 8); - for (int i = 0; i < rows.length; i++) { - out.setAtIndex(VortexFormat.LE_LONG, i, c.getLong(rows[i])); - } - yield new MaterializedLongArray(c.dtype(), n, out); - } + 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)); default -> throw new VortexException("dictionary codes of type " + codes.dtype() + " are not integers"); }; } - /// The `width`-byte elements of `src` at `rows`, in a new buffer. - private static MemorySegment gather(MemorySegment src, int width, long[] rows, SegmentAllocator arena) { - MemorySegment out = arena.allocate(Math.max(width, (long) rows.length * width), width); - for (int i = 0; i < rows.length; i++) { - MemorySegment.copy(src, rows[i] * width, out, (long) i * width, width); + private static MemorySegment longs(LongArray a, int[] rows, long offset, SegmentAllocator arena) { + MemorySegment src = flat(a); + if (src != null) { + return copy(src, 8, rows, offset, arena); } - return out; - } - - /// A list's rows address its elements through `offsets[row] .. offsets[row + 1]`, which a - /// selection breaks apart, so the selected rows become a list view over the same elements. - private static Array listView(ListArray list, long[] rows, SegmentAllocator arena) { - MemorySegment starts = arena.allocate(rows.length * 8L, 8); - MemorySegment sizes = arena.allocate(rows.length * 8L, 8); + MemorySegment out = arena.allocate(Math.max(8, rows.length * 8L), 8); for (int k = 0; k < rows.length; k++) { - long start = offsetAt(list.offsets(), rows[k]); - starts.setAtIndex(VortexFormat.LE_LONG, k, start); - sizes.setAtIndex(VortexFormat.LE_LONG, k, offsetAt(list.offsets(), rows[k] + 1) - start); + out.setAtIndex(VortexFormat.LE_LONG, k, a.getLong(rows[k] + offset)); } - return new ListViewArray((DType.List) list.dtype(), rows.length, list.elements(), - new MaterializedLongArray(DType.I64, rows.length, starts), - new MaterializedLongArray(DType.I64, rows.length, sizes)); - } - - private static long offsetAt(Array offsets, long i) { - boolean unsigned = offsets.dtype().isUnsigned(); - return switch (offsets) { - case LongArray a -> a.getLong(i); - case IntArray a when unsigned -> Integer.toUnsignedLong(a.getInt(i)); - case IntArray a -> a.getInt(i); - case ShortArray a when unsigned -> Short.toUnsignedLong(a.getShort(i)); - case ShortArray a -> a.getShort(i); - case ByteArray a when unsigned -> Byte.toUnsignedLong(a.getByte(i)); - case ByteArray a -> a.getByte(i); - default -> throw new VortexException("list offsets of type " + offsets.dtype() + " are not integers"); - }; + return out; } - private record Longs(LongArray inner, long[] rows) implements LongArray { - @Override - public long length() { - return rows.length; + private static MemorySegment ints(IntArray a, int[] rows, long offset, SegmentAllocator arena) { + MemorySegment src = flat(a); + if (src != null) { + return copy(src, 4, rows, offset, arena); } - - @Override - public DType dtype() { - return inner.dtype(); - } - - @Override - public long getLong(long i) { - return inner.getLong(rows[(int) i]); + 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)); } + return out; } - private record Ints(IntArray inner, long[] rows) implements IntArray { - @Override - public long length() { - return rows.length; + private static MemorySegment shorts(ShortArray a, int[] rows, long offset, SegmentAllocator arena) { + MemorySegment src = flat(a); + if (src != null) { + return copy(src, 2, rows, offset, arena); } - - @Override - public DType dtype() { - return inner.dtype(); - } - - @Override - public int getInt(long i) { - return inner.getInt(rows[(int) i]); + 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)); } + return out; } - private record Shorts(ShortArray inner, long[] rows) implements ShortArray { - @Override - public long length() { - return rows.length; + private static MemorySegment bytes(ByteArray a, int[] rows, long offset, SegmentAllocator arena) { + MemorySegment src = flat(a); + if (src != null) { + return copy(src, 1, rows, offset, arena); } - - @Override - public DType dtype() { - return inner.dtype(); - } - - @Override - public short getShort(long i) { - return inner.getShort(rows[(int) i]); - } - - @Override - public long fold(long identity, LongBinaryOperator op) { - long acc = identity; - for (long i = 0; i < rows.length; i++) { - acc = op.applyAsLong(acc, getShort(i)); - } - return acc; + 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)); } + return out; } - private record Bytes(ByteArray inner, long[] rows) implements ByteArray { - @Override - public long length() { - return rows.length; - } - - @Override - public DType dtype() { - return inner.dtype(); - } - - @Override - public byte getByte(long i) { - return inner.getByte(rows[(int) i]); + private static MemorySegment doubles(DoubleArray a, int[] rows, long offset, SegmentAllocator arena) { + MemorySegment src = flat(a); + if (src != null) { + return copy(src, 8, rows, offset, arena); } - - @Override - public long fold(long identity, LongBinaryOperator op) { - long acc = identity; - for (long i = 0; i < rows.length; i++) { - acc = op.applyAsLong(acc, getByte(i)); - } - return acc; + 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)); } + return out; } - private record Doubles(DoubleArray inner, long[] rows) implements DoubleArray { - @Override - public long length() { - return rows.length; - } - - @Override - public DType dtype() { - return inner.dtype(); + private static MemorySegment floats(FloatArray a, int[] rows, long offset, SegmentAllocator arena) { + MemorySegment src = flat(a); + if (src != null) { + return copy(src, 4, rows, offset, arena); } - - @Override - public double getDouble(long i) { - return inner.getDouble(rows[(int) i]); + 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)); } + return out; } - private record Floats(FloatArray inner, long[] rows) implements FloatArray { - @Override - public long length() { - return rows.length; - } - - @Override - public DType dtype() { - return inner.dtype(); + /// 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) { + MemorySegment src = flat(a); + if (src != null) { + return copy(src, 2, rows, offset, arena); } - - @Override - public float getFloat(long i) { - return inner.getFloat(rows[(int) i]); + 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))); } + return out; } - private record Float16s(Float16Array inner, long[] rows) implements Float16Array { - @Override - public long length() { - return rows.length; - } - - @Override - public DType dtype() { - return inner.dtype(); - } - - @Override - public float getFloat(long i) { - return inner.getFloat(rows[(int) i]); - } - - @Override - public Array limited(long n) { - return new Float16s(inner, Arrays.copyOf(rows, (int) n)); - } - - /// Re-encodes each selected half from its widened value, which is exact: every half is a - /// float. - @Override - public MemorySegment materialize(SegmentAllocator arena) { - MemorySegment out = arena.allocate(rows.length * 2L, 2); - for (int i = 0; i < rows.length; i++) { - out.setAtIndex(VortexFormat.LE_SHORT, i, Float.floatToFloat16(getFloat(i))); + private static MemorySegment bools(BoolArray a, int[] rows, long offset, SegmentAllocator arena) { + // a SegmentAllocator need not zero what it hands out, and only the set bits are written + MemorySegment out = arena.allocate(Math.max(1, (rows.length + 7L) >>> 3)).fill((byte) 0); + for (int k = 0; k < rows.length; k++) { + if (a.getBoolean(rows[k] + offset)) { + long at = (long) k >>> 3; + out.set(ValueLayout.JAVA_BYTE, at, (byte) (out.get(ValueLayout.JAVA_BYTE, at) | 1 << (k & 7))); } - return out; } + return out; } - private record Bools(BoolArray inner, long[] rows) implements BoolArray { - @Override - public long length() { - return rows.length; + /// The selected strings, bytes and I64 offsets in fresh buffers. + private static VarBinOffsetArray strings(VarBinArray a, int[] rows, long offset, SegmentAllocator arena) { + long totalBytes = 0; + for (int row : rows) { + totalBytes += a.getByteLength(row + offset); } - - @Override - public DType dtype() { - return inner.dtype(); - } - - @Override - public boolean getBoolean(long i) { - return inner.getBoolean(rows[(int) i]); + MemorySegment outBytes = arena.allocate(Math.max(1, totalBytes)); + MemorySegment outOffsets = arena.allocate((rows.length + 1L) * Long.BYTES, Long.BYTES); + outOffsets.setAtIndex(VortexFormat.LE_LONG, 0, 0L); + long at = 0; + for (int k = 0; k < rows.length; k++) { + byte[] b = a.getBytes(rows[k] + offset); + MemorySegment.copy(MemorySegment.ofArray(b), 0, outBytes, at, b.length); + at += b.length; + outOffsets.setAtIndex(VortexFormat.LE_LONG, k + 1L, at); } + return new VarBinOffsetArray(a.dtype(), rows.length, outBytes.asReadOnly(), outOffsets, PType.I64); } - private record VarBins(VarBinArray inner, long[] rows) implements VarBinArray { - @Override - public long length() { - return rows.length; - } - - @Override - public DType dtype() { - return inner.dtype(); - } - - @Override - public MemorySegment bytesSegment() { - return inner.bytesSegment(); - } - - @Override - public byte[] getBytes(long i) { - return inner.getBytes(rows[(int) i]); - } - - @Override - public String getString(long i) { - return inner.getString(rows[(int) i]); - } - - @Override - public int compareUnsigned(long i, byte[] other) { - return inner.compareUnsigned(rows[(int) i], other); - } - - @Override - public int getByteLength(long i) { - return inner.getByteLength(rows[(int) i]); - } - - @Override - public void forEachByteLength(IntConsumer c) { - for (int i = 0; i < rows.length; i++) { - c.accept(getByteLength(i)); + /// Writes each selected row's unscaled value as a little-endian two's-complement integer of the + /// dtype's width, read through the accessors: not every decimal array has a buffer to gather + /// from (`vortex.decimal_byte_parts` keeps its value in parts). + private static LazyDecimalArray decimals(DecimalArray a, int[] rows, long offset, SegmentAllocator arena) { + int width = ((DType.Decimal) a.dtype()).byteWidth(); + MemorySegment out = arena.allocate(Math.max(width, (long) rows.length * width), Math.min(width, 16)); + long at = 0; + for (int row : rows) { + long r = row + offset; + switch (width) { + case 1 -> out.set(ValueLayout.JAVA_BYTE, at, (byte) a.getUnscaledLong(r)); + case 2 -> out.set(VortexFormat.LE_SHORT, at, (short) a.getUnscaledLong(r)); + case 4 -> out.set(VortexFormat.LE_INT, at, (int) a.getUnscaledLong(r)); + case 8 -> out.set(VortexFormat.LE_LONG, at, a.getUnscaledLong(r)); + default -> { + // big-endian minimal bytes, reversed and sign-extended to the width + byte[] be = a.getDecimal(r).unscaledValue().toByteArray(); + byte sign = (byte) (be[0] < 0 ? -1 : 0); + for (int b = 0; b < width; b++) { + out.set(ValueLayout.JAVA_BYTE, at + b, b < be.length ? be[be.length - 1 - b] : sign); + } + } } + at += width; } + return new LazyDecimalArray(a.dtype(), rows.length, out, width); + } - @Override - public VarBinArray limited(long n) { - return new VarBins(inner, Arrays.copyOf(rows, (int) n)); + /// The flat buffer of a materialized primitive array, one element per row, or `null`: a bool + /// buffer packs eight rows a byte, and a broadcast constant holds one element for every row. + private static MemorySegment flat(Array a) { + if (a instanceof AbstractMaterializedArray m && !(a instanceof BoolArray) && m.elementCount == m.length) { + return m.buffer; } + return null; } - private record Decimals(DecimalArray inner, long[] rows) implements DecimalArray { - @Override - public long length() { - return rows.length; + /// 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; + MemorySegment out = arena.allocate(Math.max(width, (long) n * width), Math.min(width, 16)); + if (n == 0) { + return out; } - - @Override - public DType dtype() { - return inner.dtype(); + if (n * 2L <= (long) rows[n - 1] - rows[0] + 1) { + gather(src, width, rows, offset, out); + return out; } - - @Override - public BigDecimal getDecimal(long i) { - return inner.getDecimal(rows[(int) i]); + int k = 0; + while (k < n) { + int end = k + 1; + while (end < n && rows[end] == rows[end - 1] + 1) { + end++; + } + MemorySegment.copy(src, (rows[k] + offset) * width, out, (long) k * width, (long) (end - k) * width); + k = end; } + return out; + } - @Override - public long getUnscaledLong(long i) { - return inner.getUnscaledLong(rows[(int) i]); + private static void gather(MemorySegment src, int width, int[] rows, long offset, MemorySegment out) { + switch (width) { + case 1 -> { + for (int k = 0; k < rows.length; k++) { + out.set(ValueLayout.JAVA_BYTE, k, src.get(ValueLayout.JAVA_BYTE, rows[k] + offset)); + } + } + case 2 -> { + for (int k = 0; k < rows.length; k++) { + out.setAtIndex(VortexFormat.LE_SHORT, k, src.getAtIndex(VortexFormat.LE_SHORT, rows[k] + offset)); + } + } + case 4 -> { + for (int k = 0; k < rows.length; k++) { + out.setAtIndex(VortexFormat.LE_INT, k, src.getAtIndex(VortexFormat.LE_INT, rows[k] + offset)); + } + } + case 8 -> { + for (int k = 0; k < rows.length; k++) { + out.setAtIndex(VortexFormat.LE_LONG, k, src.getAtIndex(VortexFormat.LE_LONG, rows[k] + offset)); + } + } + default -> { + for (int k = 0; k < rows.length; k++) { + MemorySegment.copy(src, (rows[k] + offset) * width, out, (long) k * width, width); + } + } } + } - @Override - public Array limited(long n) { - return new Decimals(inner, Arrays.copyOf(rows, (int) n)); + /// A list's rows address its elements through `offsets[row] .. offsets[row + 1]`, which a + /// selection breaks apart, so the selected rows become a list view over the same elements. + private static Array listView(ListArray list, int[] rows, long offset, SegmentAllocator arena) { + int n = rows.length; + MemorySegment starts = arena.allocate(Math.max(8, n * 8L), 8); + MemorySegment sizes = arena.allocate(Math.max(8, n * 8L), 8); + for (int k = 0; k < n; k++) { + long row = rows[k] + offset; + long start = offsetAt(list.offsets(), row); + starts.setAtIndex(VortexFormat.LE_LONG, k, start); + sizes.setAtIndex(VortexFormat.LE_LONG, k, offsetAt(list.offsets(), row + 1) - start); } + return new ListViewArray((DType.List) list.dtype(), n, list.elements(), + new MaterializedLongArray(DType.I64, n, starts), + new MaterializedLongArray(DType.I64, n, sizes)); + } - /// Gathers each selected row's little-endian unscaled bytes from the inner buffer. - @Override - public MemorySegment materialize(SegmentAllocator arena) { - int width = ((DType.Decimal) inner.dtype()).byteWidth(); - MemorySegment src = inner.materialize(arena); - MemorySegment out = arena.allocate((long) rows.length * width, Math.min(width, 16)); - for (int i = 0; i < rows.length; i++) { - MemorySegment.copy(src, rows[i] * width, out, (long) i * width, width); - } - return out; - } + private static long offsetAt(Array offsets, long i) { + boolean unsigned = offsets.dtype().isUnsigned(); + return switch (offsets) { + case LongArray a -> a.getLong(i); + case IntArray a when unsigned -> Integer.toUnsignedLong(a.getInt(i)); + case IntArray a -> a.getInt(i); + case ShortArray a when unsigned -> Short.toUnsignedLong(a.getShort(i)); + case ShortArray a -> a.getShort(i); + case ByteArray a when unsigned -> Byte.toUnsignedLong(a.getByte(i)); + case ByteArray a -> a.getByte(i); + default -> throw new VortexException("list offsets of type " + offsets.dtype() + " are not integers"); + }; } } 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 2234cbb9f..032901c92 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 @@ -11,6 +11,7 @@ import java.lang.foreign.MemorySegment; import java.lang.foreign.ValueLayout; import java.math.BigDecimal; +import java.math.BigInteger; import java.nio.charset.StandardCharsets; import java.util.ArrayList; import java.util.List; @@ -19,16 +20,18 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; /// A filtered scan returns only the rows that match through these selections. Each family must -/// read back exactly the selected rows, in order; a dictionary must stay a dictionary, or the -/// compute kernels' dictionary lanes (and Calcite's grouped push-down) decline the selection. +/// read back exactly the selected rows, in order, from flat buffers (so the unfiltered scan's fast +/// paths apply); a dictionary must stay a dictionary, or the compute kernels' dictionary lanes (and +/// Calcite's grouped push-down) decline the selection. class SelectedArraysTest { private static final Arena ARENA = Arena.ofAuto(); - /// Out of order and repeating, so a view that ignored the index map or sorted it would show. - private static final long[] ROWS = {3, 0, 3, 1}; + /// Skips row 2 between selected rows, so a copy that ignored the bitmap, or stopped at the + /// first gap, would show. + private static final int[] ROWS = rows(0, 1, 3); @Test - void primitives_readTheSelectedRows() { + 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); @@ -42,21 +45,84 @@ void primitives_readTheSelectedRows() { BoolArray bools = (BoolArray) SelectedArrays.select(TestArrays.bools(false, true, false, true), ROWS, ARENA); // Then - assertThat(longs.length()).isEqualTo(4); - assertThat(List.of(longs.getLong(0), longs.getLong(1), longs.getLong(2), longs.getLong(3))) - .isEqualTo(List.of(13L, 10L, 13L, 11L)); - assertThat(List.of(ints.getInt(0), ints.getInt(1), ints.getInt(2), ints.getInt(3))) - .isEqualTo(List.of(13, 10, 13, 11)); - assertThat(List.of(shorts.getShort(0), shorts.getShort(3))).isEqualTo(List.of((short) 13, (short) 11)); - assertThat(shorts.fold(0, Long::sum)).isEqualTo(13 + 10 + 13 + 11); - assertThat(List.of(bytes.getByte(0), bytes.getByte(3))).isEqualTo(List.of((byte) 13, (byte) 11)); - assertThat(bytes.fold(0, Long::sum)).isEqualTo(13 + 10 + 13 + 11); - assertThat(List.of(doubles.getDouble(0), doubles.getDouble(1))).isEqualTo(List.of(4.5, 1.5)); - assertThat(List.of(floats.getFloat(0), floats.getFloat(1))).isEqualTo(List.of(4.5f, 1.5f)); - assertThat(List.of(halves.getFloat(0), halves.getFloat(1))).isEqualTo(List.of(4.5f, 1.5f)); - assertThat(List.of(bools.getBoolean(0), bools.getBoolean(1), bools.getBoolean(3))) - .isEqualTo(List.of(true, false, true)); + assertThat(longs.length()).isEqualTo(3); + assertThat(List.of(longs.getLong(0), longs.getLong(1), longs.getLong(2))).isEqualTo(List.of(10L, 11L, 13L)); + assertThat(List.of(ints.getInt(0), ints.getInt(1), ints.getInt(2))).isEqualTo(List.of(10, 11, 13)); + assertThat(List.of(shorts.getShort(0), shorts.getShort(2))).isEqualTo(List.of((short) 10, (short) 13)); + assertThat(shorts.fold(0, Long::sum)).isEqualTo(10 + 11 + 13); + assertThat(List.of(bytes.getByte(0), bytes.getByte(2))).isEqualTo(List.of((byte) 10, (byte) 13)); + assertThat(bytes.fold(0, Long::sum)).isEqualTo(10 + 11 + 13); + assertThat(List.of(doubles.getDouble(1), doubles.getDouble(2))).isEqualTo(List.of(2.5, 4.5)); + assertThat(List.of(floats.getFloat(1), floats.getFloat(2))).isEqualTo(List.of(2.5f, 4.5f)); + assertThat(List.of(halves.getFloat(1), halves.getFloat(2))).isEqualTo(List.of(2.5f, 4.5f)); + assertThat(List.of(bools.getBoolean(0), bools.getBoolean(1), bools.getBoolean(2))) + .isEqualTo(List.of(false, true, true)); assertThat(longs.dtype()).isEqualTo(DType.I64); + assertThat(List.of(longs, ints, shorts, bytes, doubles, floats, halves, bools)) + .allMatch(a -> a.segmentIfPresent().isPresent()); + } + + @Test + void sparseSelection_gathersEachRowAtEveryWidth() { + // Given: three of ten rows, density under one half, so the flat copy takes the typed + // per-row gather instead of the run copy (ROWS above takes the run copy) + 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( + (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( + (byte) 0, (byte) 1, (byte) 2, (byte) 3, (byte) 4, (byte) 5, (byte) 6, (byte) 7, (byte) 8, (byte) 9), + sparse, ARENA); + + // Then + assertThat(List.of(longs.getLong(0), longs.getLong(1), longs.getLong(2))).isEqualTo(List.of(0L, 5L, 9L)); + assertThat(List.of(ints.getInt(0), ints.getInt(1), ints.getInt(2))).isEqualTo(List.of(0, 5, 9)); + assertThat(List.of(shorts.getShort(0), shorts.getShort(1), shorts.getShort(2))) + .isEqualTo(List.of((short) 0, (short) 5, (short) 9)); + assertThat(List.of(bytes.getByte(0), bytes.getByte(1), bytes.getByte(2))) + .isEqualTo(List.of((byte) 0, (byte) 5, (byte) 9)); + } + + @Test + void broadcastBuffer_copiesRowByRowNotByRuns() { + // Given: a constant's one-element buffer standing for four rows; a run copy of rows 1..3 + // would read past the buffer's single element + MemorySegment one = ARENA.allocate(8, 8); + one.set(ValueLayout.JAVA_LONG, 0, 42L); + Array constant = new MaterializedLongArray(DType.I64, 4, one); + + // When + LongArray result = (LongArray) SelectedArrays.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)); + } + + @Test + void bools_packEverySelectedBitAcrossBytes() { + // Given: twenty rows, every third true, so the selection spans three output bytes + boolean[] values = new boolean[20]; + int[] selected = new int[19]; + for (int i = 0; i < values.length; i++) { + values[i] = i % 3 == 0; + if (i != 7) { + selected[i < 7 ? i : i - 1] = i; + } + } + + // When + BoolArray result = (BoolArray) SelectedArrays.select(TestArrays.bools(values), selected, ARENA); + + // Then + assertThat(result.length()).isEqualTo(19); + for (int k = 0; k < 19; k++) { + int row = k < 7 ? k : k + 1; + assertThat(result.getBoolean(k)).as("row %d", row).isEqualTo(row % 3 == 0); + } } @Test @@ -68,20 +134,20 @@ void materialize_writesTheSelectedRowsInOrder() { MemorySegment result = selected.materialize(ARENA); // Then - assertThat(result.toArray(ValueLayout.JAVA_LONG)).containsExactly(13, 10, 13, 11); + assertThat(result.asSlice(0, 24).toArray(ValueLayout.JAVA_LONG)).containsExactly(10, 11, 13); } @Test void float16Materialize_keepsTheHalfBits() { // Given - Float16Array selected = (Float16Array) SelectedArrays.select(TestArrays.float16(1.5f, 2.5f), new long[]{1, 0}, ARENA); + Float16Array selected = (Float16Array) SelectedArrays.select(TestArrays.float16(1.5f, 2.5f, 3.5f), rows(0, 2), ARENA); // When MemorySegment result = selected.materialize(ARENA); // Then - assertThat(Float.float16ToFloat(result.getAtIndex(ValueLayout.JAVA_SHORT_UNALIGNED, 0))).isEqualTo(2.5f); - assertThat(Float.float16ToFloat(result.getAtIndex(ValueLayout.JAVA_SHORT_UNALIGNED, 1))).isEqualTo(1.5f); + assertThat(Float.float16ToFloat(result.getAtIndex(ValueLayout.JAVA_SHORT_UNALIGNED, 0))).isEqualTo(1.5f); + assertThat(Float.float16ToFloat(result.getAtIndex(ValueLayout.JAVA_SHORT_UNALIGNED, 1))).isEqualTo(3.5f); } @Test @@ -90,7 +156,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, new long[]{1, 2}, ARENA); + MaskedArray result = (MaskedArray) SelectedArrays.select(masked, rows(1, 2), ARENA); // Then assertThat(result.isValid(0)).isFalse(); @@ -99,19 +165,19 @@ void masked_selectsValuesAndValidityTogether() { } @Test - void strings_readTheSelectedRows() { + void strings_copyTheSelectedRows() { // When VarBinArray result = (VarBinArray) SelectedArrays.select(strings("a", "bb", "ccc", "dddd"), ROWS, ARENA); // Then - assertThat(List.of(result.getString(0), result.getString(1), result.getString(2), result.getString(3))) - .isEqualTo(List.of("dddd", "a", "dddd", "bb")); - assertThat(result.getByteLength(0)).isEqualTo(4); - assertThat(result.getBytes(3)).isEqualTo("bb".getBytes(StandardCharsets.UTF_8)); - assertThat(result.compareUnsigned(1, "a".getBytes(StandardCharsets.UTF_8))).isZero(); + assertThat(List.of(result.getString(0), result.getString(1), result.getString(2))) + .isEqualTo(List.of("a", "bb", "dddd")); + assertThat(result.getByteLength(2)).isEqualTo(4); + assertThat(result.getBytes(1)).isEqualTo("bb".getBytes(StandardCharsets.UTF_8)); + assertThat(result.compareUnsigned(0, "a".getBytes(StandardCharsets.UTF_8))).isZero(); List lengths = new ArrayList<>(); result.forEachByteLength(lengths::add); - assertThat(lengths).isEqualTo(List.of(4, 1, 4, 2)); + assertThat(lengths).isEqualTo(List.of(1, 2, 4)); assertThat(result.limited(2).length()).isEqualTo(2); } @@ -121,7 +187,7 @@ void stringDictionary_staysADictionaryOverTheSamePool() { Array dict = stringDict(new String[]{"x", "y"}, new byte[]{1, 0, 1, 0}); // When - Array result = SelectedArrays.select(dict, new long[]{0, 1}, ARENA); + Array result = SelectedArrays.select(dict, rows(0, 1), ARENA); // Then assertThat(result).isInstanceOf(VarBinDictArray.class); @@ -138,13 +204,13 @@ void primitiveDictionary_staysADictionaryWithFlatCodes() { Array dict = DictLongArray.of(DType.I64, 4, pool, codes); // When - Array result = SelectedArrays.select(dict, new long[]{3, 2}, ARENA); + Array result = SelectedArrays.select(dict, rows(2, 3), ARENA); // Then assertThat(result).isInstanceOf(DictLongArray.class); DictLongArray selected = (DictLongArray) result; assertThat(selected.codes()).isInstanceOf(MaterializedByteArray.class); - assertThat(List.of(selected.getLong(0), selected.getLong(1))).isEqualTo(List.of(100L, 200L)); + assertThat(List.of(selected.getLong(0), selected.getLong(1))).isEqualTo(List.of(200L, 100L)); } @Test @@ -155,12 +221,29 @@ void offsetSlice_selectsFromTheArrayItSlicesSoADictionaryStaysOne() { Array slice = new OffsetLongArray(DType.I64, 2, (LongArray) dict, 2); // When - Array result = SelectedArrays.select(slice, new long[]{1, 0}, ARENA); + Array result = SelectedArrays.select(slice, rows(0, 1), ARENA); // Then assertThat(result).isInstanceOf(DictLongArray.class); LongArray selected = (LongArray) result; - assertThat(List.of(selected.getLong(0), selected.getLong(1))).isEqualTo(List.of(100L, 200L)); + assertThat(List.of(selected.getLong(0), selected.getLong(1))).isEqualTo(List.of(200L, 100L)); + } + + @Test + void offsetSliceOfAFlatArray_copiesFromItsShiftedRows() { + // Given: a slice starting at row 1, under a mask, so the slice offset has to reach the + // flat copy of both the values and the validity + Array values = new OffsetIntArray(DType.I32, 3, TestArrays.ints(10, 11, 12, 13), 1); + BoolArray validity = new OffsetBoolArray(DType.BOOL, 3, TestArrays.bools(true, true, false, true), 1); + Array masked = new MaskedArray(values, validity); + + // When + MaskedArray result = (MaskedArray) SelectedArrays.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); + assertThat(result.isValid(0)).isFalse(); + assertThat(result.isValid(1)).isTrue(); } @Test @@ -170,12 +253,12 @@ 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, new long[]{2, 0}, ARENA); + StructArray result = (StructArray) SelectedArrays.select(struct, rows(0, 2), ARENA); // Then assertThat(result.length()).isEqualTo(2); - assertThat(((LongArray) result.field(0)).getLong(0)).isEqualTo(3); - assertThat(((VarBinArray) result.field(1)).getString(1)).isEqualTo("p"); + assertThat(((LongArray) result.field(0)).getLong(1)).isEqualTo(3); + assertThat(((VarBinArray) result.field(1)).getString(1)).isEqualTo("r"); } @Test @@ -185,12 +268,13 @@ void fixedSizeList_selectsEachRowsElements() { Array list = new FixedSizeListArray(dtype, 3, TestArrays.longs(1, 2, 3, 4, 5, 6)); // When - FixedSizeListArray result = (FixedSizeListArray) SelectedArrays.select(list, new long[]{2, 0}, ARENA); + FixedSizeListArray result = (FixedSizeListArray) SelectedArrays.select(list, rows(0, 2), ARENA); // Then LongArray elements = (LongArray) result.elements(); + assertThat(elements.length()).isEqualTo(4); assertThat(List.of(elements.getLong(0), elements.getLong(1), elements.getLong(2), elements.getLong(3))) - .isEqualTo(List.of(5L, 6L, 1L, 2L)); + .isEqualTo(List.of(1L, 2L, 5L, 6L)); } @Test @@ -200,13 +284,13 @@ 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, new long[]{2, 1}, ARENA); + ListViewArray result = (ListViewArray) SelectedArrays.select(list, rows(1, 2), ARENA); - // Then: row 0 starts at element 3 and holds 3; row 1 starts at 1 and holds 2 - assertThat(((LongArray) result.offsets()).getLong(0)).isEqualTo(3); - assertThat(((LongArray) result.sizes()).getLong(0)).isEqualTo(3); - assertThat(((LongArray) result.offsets()).getLong(1)).isEqualTo(1); - assertThat(((LongArray) result.sizes()).getLong(1)).isEqualTo(2); + // 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); + assertThat(((LongArray) result.sizes()).getLong(0)).isEqualTo(2); + assertThat(((LongArray) result.offsets()).getLong(1)).isEqualTo(3); + assertThat(((LongArray) result.sizes()).getLong(1)).isEqualTo(3); assertThat(result.elements().length()).isEqualTo(6); } @@ -221,14 +305,39 @@ void decimal_readsAndMaterializesTheSelectedRows() { Array decimals = new LazyDecimalArray(dtype, 3, buf, 4); // When - DecimalArray result = (DecimalArray) SelectedArrays.select(decimals, new long[]{2, 0}, ARENA); + DecimalArray result = (DecimalArray) SelectedArrays.select(decimals, rows(0, 2), ARENA); // Then - assertThat(result.getDecimal(0)).isEqualTo(new BigDecimal("3.00")); - assertThat(result.getUnscaledLong(1)).isEqualTo(100); + assertThat(result.getDecimal(0)).isEqualTo(new BigDecimal("1.00")); + assertThat(result.getUnscaledLong(1)).isEqualTo(300); MemorySegment materialized = result.materialize(ARENA); - assertThat(materialized.getAtIndex(ValueLayout.JAVA_INT_UNALIGNED, 0)).isEqualTo(300); - assertThat(materialized.getAtIndex(ValueLayout.JAVA_INT_UNALIGNED, 1)).isEqualTo(100); + assertThat(materialized.getAtIndex(ValueLayout.JAVA_INT_UNALIGNED, 0)).isEqualTo(100); + assertThat(materialized.getAtIndex(ValueLayout.JAVA_INT_UNALIGNED, 1)).isEqualTo(300); + } + + @Test + void wideDecimal_signExtendsEachValueToSixteenBytes() { + // Given: a 16-byte decimal (precision 38) holding -5, 7 and -2^100: the copy writes the + // minimal big-endian bytes back as 16 little-endian ones, so a negative value whose + // minimal form is one byte only survives if the sign is extended + DType.Decimal dtype = new DType.Decimal((byte) 38, (byte) 0, false); + BigInteger[] values = {BigInteger.valueOf(-5), BigInteger.valueOf(7), BigInteger.TWO.pow(100).negate()}; + MemorySegment buf = ARENA.allocate(16L * values.length, 16); + for (int i = 0; i < values.length; i++) { + byte[] be = values[i].toByteArray(); + for (int b = 0; b < 16; b++) { + byte sign = (byte) (values[i].signum() < 0 ? -1 : 0); + buf.set(ValueLayout.JAVA_BYTE, i * 16L + b, b < be.length ? be[be.length - 1 - b] : sign); + } + } + Array decimals = new LazyDecimalArray(dtype, values.length, buf, 16); + + // When + DecimalArray result = (DecimalArray) SelectedArrays.select(decimals, rows(0, 2), ARENA); + + // Then + assertThat(result.getDecimal(0)).isEqualTo(new BigDecimal(values[0])); + assertThat(result.getDecimal(1)).isEqualTo(new BigDecimal(values[2])); } @Test @@ -237,13 +346,13 @@ void nullArray_keepsItsTypeWithTheSelectedLength() { Array result = SelectedArrays.select(new NullArray(DType.I64.asNullable(), 10), ROWS, ARENA); // Then - assertThat(result).isEqualTo(new NullArray(DType.I64.asNullable(), 4)); + assertThat(result).isEqualTo(new NullArray(DType.I64.asNullable(), 3)); } @Test void anEmptySelection_hasNoRows() { // When - Array result = SelectedArrays.select(TestArrays.longs(1, 2, 3), new long[0], ARENA); + Array result = SelectedArrays.select(TestArrays.longs(1, 2, 3), new int[0], ARENA); // Then assertThat(result.length()).isZero(); @@ -256,11 +365,15 @@ void anArrayWithoutRows_failsAsVortexException() { new MemorySegment[0], new Array[0]); // When / Then - assertThatThrownBy(() -> SelectedArrays.select(unknown, new long[]{0}, ARENA)) + assertThatThrownBy(() -> SelectedArrays.select(unknown, rows(0), ARENA)) .isInstanceOf(VortexException.class) .hasMessageContaining("cannot select rows"); } + private static int[] rows(int... rows) { + return rows; + } + private static VarBinArray strings(String... values) { byte[] all = String.join("", values).getBytes(StandardCharsets.UTF_8); MemorySegment offsets = ARENA.allocate((values.length + 1) * 4L, 4);