perf(reader): select a filtered chunk's rows through a Rust-style Mask (#602)

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))
 

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<ColumnName> names, Map<ColumnName, Integer> positions, Consumer<Chunk> onClose) {
+    Chunk select(Mask rows, List<ColumnName> names, Map<ColumnName, Integer> positions, Consumer<Chunk> onClose) {
         IntFunction<Column> 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

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<Chunk>, 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<ColumnName> projectedNames;
     private List<DType> projectedDtypes;
     private Map<ColumnName, Integer> 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,15 +437,15 @@ 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<ColumnName> names = outputNames;
         Map<ColumnName, Integer> positions = outputPositions;
         if (singleColumnIsSyntheticWrapper) {
             // the synthetic column was unpacked into the file's real columns
             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);

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));
+        }
+    }
+}