diff --git a/CHANGELOG.md b/CHANGELOG.md index fae7352f2..2271cda10 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] ### Added -- `Compute.groupedAggregate(Chunk, ColumnName, RowFilter, ColumnName)` folds an aggregate per value of a dictionary-encoded string column. ([#555](https://github.com/dfa1/vortex-java/pull/555)) +- `Compute.groupedAggregate(Chunk, ColumnName, RowFilter, ColumnName)` folds an aggregate per value of a dictionary-encoded string, boolean or integer column. ([#555](https://github.com/dfa1/vortex-java/pull/555)) - `Chunk.slice(offset, length)` gives a zero-copy view of a row range. ([#550](https://github.com/dfa1/vortex-java/pull/550)) - The Calcite adapter tells the planner each table's row count, so join order and costing use the real table size. ([#550](https://github.com/dfa1/vortex-java/pull/550)) - `Compute.matcher(Chunk, RowFilter)` lowers a filter once per chunk and returns a row test, for callers that test many rows. ([#552](https://github.com/dfa1/vortex-java/issues/552)) @@ -23,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - SQL `LIKE 'prefix%'` on a string column is about 5 times faster: it is pushed into the scan as a string range that prunes chunks, instead of being matched by Calcite row by row. ([#558](https://github.com/dfa1/vortex-java/pull/558)) - SQL `ORDER BY ... LIMIT n` is up to 9 times faster: the first rows are kept in a bounded heap during the scan instead of sorting the whole table. ([#557](https://github.com/dfa1/vortex-java/pull/557)) - SQL `GROUP BY` on a dictionary-encoded string column is 10 to 30 times faster: rows are counted and summed per dictionary code instead of being built one by one for Calcite. ([#555](https://github.com/dfa1/vortex-java/pull/555)) +- SQL `GROUP BY` on a boolean or integer column is 2 to 4 times faster: rows are folded per group value through a primitive map instead of being built one by one for Calcite. ([#561](https://github.com/dfa1/vortex-java/pull/561)) - A filtered SQL `COUNT`/`MIN`/`MAX` answered from zone maps decodes the zones its filter cuts once per query instead of up to six times. ([#555](https://github.com/dfa1/vortex-java/pull/555)) - The Calcite adapter pushes `IN` lists, comparisons with the literal on the left and `IS NULL` into the scan, so they prune chunks instead of being re-checked row by row. ([#550](https://github.com/dfa1/vortex-java/pull/550)) - The Calcite adapter answers a filtered `COUNT`, `MIN` or `MAX` from Rust's `vortex.zoned` zone maps, the default, decoding only the zones the filter cuts. ([#550](https://github.com/dfa1/vortex-java/pull/550)) diff --git a/calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexAggregatePushDownRule.java b/calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexAggregatePushDownRule.java index f2af628ac..da4d9ad27 100644 --- a/calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexAggregatePushDownRule.java +++ b/calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexAggregatePushDownRule.java @@ -186,8 +186,8 @@ public void onMatch(RelOptRuleCall call) { call.transformTo(values); } - /// Rewrites a `GROUP BY` on one string column into a [LogicalValues] with one row per group, - /// folded per dictionary code by [VortexTable#groupedFold(ColumnName, RowFilter, java.util.Set)] + /// Rewrites a `GROUP BY` on one string, integer or boolean column into a [LogicalValues] with one + /// row per group, folded by [VortexTable#groupedFold(ColumnName, RowFilter, java.util.Set)] /// instead of materializing every row for Calcite to hash. Leaves the plan alone when the group /// column is not a string, an aggregate is not `COUNT`/`SUM`/`MIN`/`MAX` over a plain column, or /// the fold declines (a chunk not dictionary-encoded, an aggregate column it cannot total, too @@ -197,7 +197,12 @@ private static void rewriteGrouped(RelOptRuleCall call, Aggregate aggregate, Vor RexBuilder rexBuilder, List outTypes, RowFilter filter) { ColumnName group = resolveColumn(aggregate.getGroupSet().nth(0), scanColumns, project); RelDataTypeField groupField = group == null ? null : scanRowType.getField(group.value(), true, false); - if (groupField == null || !SqlTypeName.CHAR_TYPES.contains(groupField.getType().getSqlTypeName())) { + if (groupField == null) { + return; + } + SqlTypeName groupSqlType = groupField.getType().getSqlTypeName(); + if (!SqlTypeName.CHAR_TYPES.contains(groupSqlType) && !SqlTypeName.INT_TYPES.contains(groupSqlType) + && groupSqlType != SqlTypeName.BOOLEAN) { return; } List calls = aggregate.getAggCallList(); @@ -221,16 +226,15 @@ private static void rewriteGrouped(RelOptRuleCall call, Aggregate aggregate, Vor } // ConcurrentHashMap keys cannot be null, so a missing filter is keyed as an empty list List foldKey = List.of("grouped", table, group, filter == null ? List.of() : filter, aggColumns); - Optional> folds = + Optional> folds = memo(call.getPlanner(), foldKey, () -> table.groupedFold(group, filter, aggColumns)); if (folds.isEmpty()) { return; } RelDataType groupType = aggregate.getRowType().getFieldList().getFirst().getType(); List> tuples = new ArrayList<>(); - for (Map.Entry entry : folds.get().entrySet()) { - RexLiteral key = entry.getKey() == null ? rexBuilder.makeNullLiteral(groupType) - : minMaxLiteral(rexBuilder, entry.getKey(), groupType); + for (Map.Entry entry : folds.get().entrySet()) { + RexLiteral key = VortexSortPushDownRule.literal(rexBuilder, entry.getKey(), groupType); if (key == null) { return; // a key the column's charset cannot hold as a literal } diff --git a/calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexTable.java b/calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexTable.java index aa4370afc..2556ba1f0 100644 --- a/calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexTable.java +++ b/calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexTable.java @@ -358,20 +358,21 @@ public Optional filteredFold(RowFilter filter, ColumnName aggColum public record GroupFold(long rows, Map columns) { } - /// Folds aggregates per value of a dictionary-encoded string column over the rows `filter` - /// selects, reading the group column's codes instead of building a row per record: each chunk + /// Folds aggregates per value of a dictionary-encoded string, boolean or integer column over the + /// rows `filter` selects, without building a row per record: each chunk /// goes through [Compute#groupedAggregate(Chunk, ColumnName, RowFilter, ColumnName)] and the /// chunks' results merge per group value. /// - /// The result is empty, leaving the query to the scan, when any chunk's group column is not - /// dictionary-encoded, an aggregate column has no long or double fold, or there are more than - /// [#MAX_GROUPS] groups. + /// The result is empty, leaving the query to the scan, when the kernel declines a chunk (a string + /// column not dictionary-encoded there, too many groups in it, an aggregate column with no long + /// or double fold) or there are more than [#MAX_GROUPS] groups in all. /// - /// @param group the string column to group by + /// @param group the column to group by /// @param filter the pushed predicate, or `null` for every row /// @param aggColumns the columns the aggregates reduce; empty when only rows are counted - /// @return the fold per group value (`null` is the null group), or empty - public Optional> groupedFold(ColumnName group, RowFilter filter, Set aggColumns) { + /// @return the fold per group value (`String`, `Boolean` or `Long`; `null` is the null group), or + /// empty + public Optional> groupedFold(ColumnName group, RowFilter filter, Set aggColumns) { java.util.LinkedHashSet columns = new java.util.LinkedHashSet<>(); columns.add(group); if (filter != null) { @@ -385,20 +386,20 @@ public Optional> groupedFold(ColumnName group, RowFilter // one pass per aggregate column (a null column counts rows only) List passes = aggColumns.isEmpty() ? java.util.Collections.singletonList(null) : List.copyOf(aggColumns); - Map rows = new java.util.HashMap<>(); - Map> byGroup = new java.util.HashMap<>(); + Map rows = new java.util.HashMap<>(); + Map> byGroup = new java.util.HashMap<>(); try (VortexReader reader = VortexReader.open(file); ScanIterator scan = reader.scan(options)) { while (scan.hasNext()) { try (Chunk chunk = scan.next()) { for (int p = 0; p < passes.size(); p++) { ColumnName aggColumn = passes.get(p); - Optional> part = + Optional> part = Compute.groupedAggregate(chunk, group, filter, aggColumn); if (part.isEmpty()) { return Optional.empty(); } - for (Map.Entry entry : part.get().entrySet()) { + for (Map.Entry entry : part.get().entrySet()) { FilteredAggregate aggregate = entry.getValue(); if (p == 0) { rows.computeIfAbsent(entry.getKey(), k -> new long[1])[0] += aggregate.selectedRows(); @@ -423,7 +424,7 @@ public Optional> groupedFold(ColumnName group, RowFilter } catch (IOException e) { throw new UncheckedIOException("cannot fold groups of " + file, e); } - Map result = new java.util.HashMap<>(); + Map result = new java.util.HashMap<>(); rows.forEach((key, count) -> { Map columnFolds = new java.util.HashMap<>(); byGroup.getOrDefault(key, Map.of()).forEach((column, fold) -> columnFolds.put(column, fold.result())); diff --git a/calcite/src/test/java/io/github/dfa1/vortex/calcite/AggregateRuleBranchTest.java b/calcite/src/test/java/io/github/dfa1/vortex/calcite/AggregateRuleBranchTest.java index 02f7f1161..0c50d9563 100644 --- a/calcite/src/test/java/io/github/dfa1/vortex/calcite/AggregateRuleBranchTest.java +++ b/calcite/src/test/java/io/github/dfa1/vortex/calcite/AggregateRuleBranchTest.java @@ -133,7 +133,7 @@ void groupedOnDictionaryString_rewritesToValues() { @ParameterizedTest @ValueSource(strings = { - // a numeric group key has no dictionary codes to fold by + // more distinct volumes than the 10,000 groups a rewrite emits as literal rows "select volume, count(*) from ohlc group by volume", // only a single group key is folded "select symbol, volume, count(*) from ohlc group by symbol, volume", diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/LongIntMap.java b/core/src/main/java/io/github/dfa1/vortex/core/compute/LongIntMap.java similarity index 98% rename from writer/src/main/java/io/github/dfa1/vortex/writer/encode/LongIntMap.java rename to core/src/main/java/io/github/dfa1/vortex/core/compute/LongIntMap.java index cebbe00b6..458eeae13 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/LongIntMap.java +++ b/core/src/main/java/io/github/dfa1/vortex/core/compute/LongIntMap.java @@ -1,11 +1,11 @@ -package io.github.dfa1.vortex.writer.encode; +package io.github.dfa1.vortex.core.compute; import java.util.Arrays; /// Open-addressing map from a 64-bit key to an `int` value, for the write path's per-row scans. /// /// The encode path kept growing private copies of this — a distinct-value counter in -/// [ArrayStats], a value-to-code index in `DictEncodingEncoder`, another in +/// `ArrayStats`, a value-to-code index in `DictEncodingEncoder`, another in /// `io.github.dfa1.vortex.writer.DictColumnState`, and a capped distinct set in the same class. /// They existed because probing a `HashMap` boxes a key per row, which profiled /// as the hottest frame of a write more than once. This is that structure, once. @@ -195,6 +195,8 @@ public Entry maxEntry() { public record Entry(long key, int value) { } + /// Returns how many distinct keys are stored. + /// /// @return the number of distinct keys stored public int size() { return size; diff --git a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/LongIntMapTest.java b/core/src/test/java/io/github/dfa1/vortex/core/compute/LongIntMapTest.java similarity index 99% rename from writer/src/test/java/io/github/dfa1/vortex/writer/encode/LongIntMapTest.java rename to core/src/test/java/io/github/dfa1/vortex/core/compute/LongIntMapTest.java index dbd140e11..755215f07 100644 --- a/writer/src/test/java/io/github/dfa1/vortex/writer/encode/LongIntMapTest.java +++ b/core/src/test/java/io/github/dfa1/vortex/core/compute/LongIntMapTest.java @@ -1,4 +1,4 @@ -package io.github.dfa1.vortex.writer.encode; +package io.github.dfa1.vortex.core.compute; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; diff --git a/docs/reference.md b/docs/reference.md index f74650526..0a029631d 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -637,11 +637,11 @@ column-vs-literal comparisons, `IN` lists, `IS [NOT] NULL` and prefix `LIKE` joi the rows of zones it cuts are decoded. `sum` needs a per-zone sum, which only the legacy `vortex.stats` zone map records (writes targeting an edition before `core2026.08.0`); with Rust's `vortex.zoned` — the default since #447, and every Rust-written file — `sum` is answered by a scan. -A `GROUP BY` on one string column, with `count`/`sum`/`min`/`max` of plain numeric columns and the -same kind of `WHERE`, is folded per dictionary code into one row per group instead of building a -row per record; it falls back to the scan when a chunk of the column is not dictionary-encoded, -there are more than 10,000 groups, or a group value is outside Latin-1 (Calcite's `VARCHAR` -charset). `DISTINCT` and `FILTER (WHERE ...)` aggregates are always computed by the scan. +A `GROUP BY` on one column (a dictionary-encoded string, a boolean, or an integer other than +`u64`), with `count`/`sum`/`min`/`max` of plain numeric columns and the same kind of `WHERE`, is +folded per group into one row per group instead of building a row per record; it falls back to the +scan when a chunk of a string column is not dictionary-encoded, there are more than 10,000 groups, +or a group value is outside Latin-1 (Calcite's `VARCHAR` charset). `DISTINCT` and `FILTER (WHERE ...)` aggregates are always computed by the scan. `ORDER BY ... LIMIT n` (any `OFFSET`, `n + offset` up to 1,000) over plain columns, with the same kind of `WHERE`, keeps the first rows in a bounded heap during the scan instead of sorting every row; it falls back to Calcite's sort when a returned value has no SQL literal of its type (`NaN`, diff --git a/integration/src/test/java/io/github/dfa1/vortex/integration/CalciteInteropIntegrationTest.java b/integration/src/test/java/io/github/dfa1/vortex/integration/CalciteInteropIntegrationTest.java index b22b1417d..2f3f6bdc4 100644 --- a/integration/src/test/java/io/github/dfa1/vortex/integration/CalciteInteropIntegrationTest.java +++ b/integration/src/test/java/io/github/dfa1/vortex/integration/CalciteInteropIntegrationTest.java @@ -399,6 +399,34 @@ void jniWriter_aggregateWhereOnSmallint_isPushedDownAndEqualsTheScan(String wher } } + @ParameterizedTest(name = "{0}") + @CsvSource(delimiter = '|', value = { + "select g, count(*), sum(id), min(n), max(n) from vtx.t group by g order by g " + + "| select g + 0, count(*), sum(id), min(n), max(n) from vtx.t group by g + 0 order by 1", + "select g, count(n) from vtx.t where id < 1500 group by g order by g " + + "| select g + 0, count(n) from vtx.t where id + 0 < 1500 group by g + 0 order by 1"}) + void jniWriter_groupByInteger_isPushedDownAndEqualsTheScan(String sql, String oracle, @TempDir Path tmp) + throws Exception { + // Given — a GROUP BY on an integer column folds per value through a primitive map instead of + // building a row per record; "g + 0" keeps the rule out, so Calcite groups the same rows + Path file = tmp.resolve("jni_group_int.vortex"); + writeIdsWithJni(file, BATCHES, BATCH_ROWS); + + try (Connection conn = VortexCalcite.connect("vtx", Map.of("t", file))) { + // When + String plan; + try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery("explain plan for " + sql)) { + rs.next(); + plan = rs.getString(1); + } + List result = sqlRows(conn, sql); + + // Then + assertThat(plan).containsIgnoringCase("Values").doesNotContain("TableScan"); + assertThat(result).isEqualTo(sqlRows(conn, oracle)).hasSize(7); + } + } + private static List sqlRows(Connection conn, String sql) throws Exception { List out = new java.util.ArrayList<>(); try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery(sql)) { diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/compute/Compute.java b/reader/src/main/java/io/github/dfa1/vortex/reader/compute/Compute.java index 6c82e4aeb..66116a92c 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/compute/Compute.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/compute/Compute.java @@ -116,24 +116,27 @@ public static LongPredicate matcher(Chunk chunk, RowFilter filter) { return FusedFilterAggregate.rowTest(chunk, filter); } - /// Folds `aggColumn` per value of a dictionary-encoded string `groupColumn` over the rows `filter` - /// selects: the per-group counterpart of [#filteredAggregate(Chunk, RowFilter, ColumnName)]. + /// Folds `aggColumn` per value of `groupColumn` over the rows `filter` selects: the per-group + /// counterpart of [#filteredAggregate(Chunk, RowFilter, ColumnName)]. /// - /// Rows accumulate per dictionary code, so a row costs a code read and a few array updates, and - /// each code becomes its string once per chunk. A null group value is its own group, under the - /// `null` key. Each group's [FilteredAggregate] follows the same null semantics as the ungrouped - /// fold. + /// Rows accumulate per group id in primitive arrays, so a row costs a key read and a few array + /// updates, and no key is boxed per row. The group column may be a dictionary-encoded string + /// column (keys are `String`s, each code turned into its string once per chunk), a boolean column + /// (`Boolean` keys) or an integer column other than `u64` (`Long` keys, unsigned values + /// zero-extended). A null group value is its own group, under the `null` key. Each group's + /// [FilteredAggregate] follows the same null semantics as the ungrouped fold. /// - /// The result is empty when the shapes do not qualify: `groupColumn` is not dictionary-encoded - /// in this chunk, or `aggColumn` is not an integer or floating column (`f16` and `u64` included, - /// which this fold does not total). The caller then groups the rows itself. + /// The result is empty when the shapes do not qualify: another group column type, a string + /// column not dictionary-encoded in this chunk, more than 65,536 groups in the chunk, or an + /// `aggColumn` that is not an integer or floating column (`f16` and `u64` included, which this + /// fold does not total). The caller then groups the rows itself. /// /// @param chunk the decoded chunk holding every referenced column - /// @param groupColumn the dictionary-encoded string column to group by + /// @param groupColumn the column to group by /// @param filter the rows to fold, or `null` for every row /// @param aggColumn the numeric column to reduce, or `null` to count rows only /// @return the fold per group value, or empty when the columns do not qualify - public static Optional> groupedAggregate(Chunk chunk, ColumnName groupColumn, + public static Optional> groupedAggregate(Chunk chunk, ColumnName groupColumn, RowFilter filter, ColumnName aggColumn) { Objects.requireNonNull(chunk, "chunk"); Objects.requireNonNull(groupColumn, "groupColumn"); diff --git a/reader/src/main/java/io/github/dfa1/vortex/reader/compute/GroupedAggregate.java b/reader/src/main/java/io/github/dfa1/vortex/reader/compute/GroupedAggregate.java index 15824b535..15b23bb36 100644 --- a/reader/src/main/java/io/github/dfa1/vortex/reader/compute/GroupedAggregate.java +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/compute/GroupedAggregate.java @@ -1,5 +1,6 @@ package io.github.dfa1.vortex.reader.compute; +import io.github.dfa1.vortex.core.compute.LongIntMap; import io.github.dfa1.vortex.core.error.VortexException; import io.github.dfa1.vortex.core.model.ColumnName; import io.github.dfa1.vortex.core.model.DType; @@ -14,22 +15,30 @@ import io.github.dfa1.vortex.reader.array.VarBinOffsetArray; import io.github.dfa1.vortex.reader.array.VarBinSlicedArray; +import java.util.Arrays; import java.util.HashMap; import java.util.Map; import java.util.Optional; +import java.util.function.LongFunction; import java.util.function.LongPredicate; +import java.util.function.LongUnaryOperator; -/// Grouped aggregation over a dictionary-encoded string column: rows are accumulated per dictionary -/// code in primitive arrays, and each code is turned into its string once, when the per-group -/// results are built. No row is materialized and no key is hashed per row. +/// Grouped aggregation that never builds a row or hashes a boxed key per row. Rows accumulate per +/// group id in primitive arrays, and each id becomes its key once, when the per-group results are +/// built. Three group columns qualify: +/// - a dictionary-encoded string column (or a slice of one): the group id is the dictionary code, +/// and each code is turned into its string once; +/// - a boolean column: two ids; +/// - an integer column (not `u64`, whose values past 2^63 have no signed-long key): ids assigned +/// through a primitive [LongIntMap], up to [#MAX_GROUPS] distinct values per chunk. /// -/// Only a [VarBinDictArray] group column (or a slice of one) qualifies, and only an aggregate column -/// the long or double lane can fold; anything else returns empty so the caller falls back to its -/// row path. +/// The aggregate column must be one the long or double lane can fold. Anything else returns empty +/// so the caller falls back to its row path. final class GroupedAggregate { - /// Most dictionary entries accumulated per chunk. The dictionary size comes from an untrusted - /// offsets buffer, so a larger one declines rather than size the per-code arrays from it. + /// Most groups accumulated per chunk. A dictionary's size comes from an untrusted offsets buffer, + /// and an integer column may hold a value per row, so more than this declines rather than size + /// the per-group arrays from either. private static final int MAX_GROUPS = 1 << 16; private GroupedAggregate() { @@ -38,49 +47,52 @@ private GroupedAggregate() { /// Folds `aggColumn` per value of `groupColumn` over the rows `filter` selects. /// /// @param chunk the decoded chunk holding every referenced column - /// @param groupColumn the dictionary-encoded string column to group by + /// @param groupColumn the column to group by /// @param filter the rows to fold, or `null` for every row /// @param aggColumn the numeric column to reduce, or `null` to count rows only - /// @return the fold per group value (the `null` key is the null group), or empty when the group - /// or aggregate column has a shape this kernel does not fold - static Optional> aggregate(Chunk chunk, ColumnName groupColumn, + /// @return the fold per group value (`String`, `Boolean` or `Long`; the `null` key is the null + /// group), or empty when the group or aggregate column has a shape this kernel does not fold + static Optional> aggregate(Chunk chunk, ColumnName groupColumn, RowFilter filter, ColumnName aggColumn) { + Lane lane = Lane.of(aggColumn == null ? null : chunk.column(aggColumn)); + if (lane == null) { + return Optional.empty(); + } Array group = chunk.column(groupColumn); - BoolArray groupValidity = NumericColumns.validityOf(group); - VarBinDictArray dict; - long offset; - switch (NumericColumns.unwrap(group)) { - case VarBinDictArray d -> { - dict = d; - offset = 0; - } - case VarBinSlicedArray(var _, var _, VarBinDictArray d, long sliceOffset) -> { - dict = d; - offset = sliceOffset; - } - default -> { - return Optional.empty(); + BoolArray validity = NumericColumns.validityOf(group); + LongPredicate selected = filter == null ? null : FusedFilterAggregate.rowTest(chunk, filter); + return switch (NumericColumns.unwrap(group)) { + case VarBinDictArray d -> byCode(d, 0, validity, selected, lane, chunk.rowCount()); + case VarBinSlicedArray(var _, var _, VarBinDictArray d, long offset) -> + byCode(d, offset, validity, selected, lane, chunk.rowCount()); + case BoolArray b -> byKey(i -> b.getBoolean(i) ? 1 : 0, validity, selected, lane, chunk.rowCount(), + key -> key != 0); + case Array data when NumericColumns.isLongDomain(data) + && data.dtype() instanceof DType.Primitive p && p.ptype() != PType.U64 -> { + boolean unsigned = p.ptype().isUnsigned(); + yield byKey(i -> NumericColumns.widenLong(data, unsigned, i), validity, selected, lane, + chunk.rowCount(), key -> key); } - } + default -> Optional.empty(); + }; + } + + /// Groups by dictionary code: one slot per entry, plus the null group. + private static Optional> byCode(VarBinDictArray dict, long offset, + BoolArray validity, LongPredicate selected, Lane lane, long n) { VarBinOffsetArray values = dict.values(); if (!(dict.dtype() instanceof DType.Utf8) || values.length() > MAX_GROUPS) { return Optional.empty(); } - Lane lane = Lane.of(aggColumn == null ? null : chunk.column(aggColumn)); - if (lane == null) { - return Optional.empty(); - } int groups = (int) values.length(); - LongPredicate selected = filter == null ? null : FusedFilterAggregate.rowTest(chunk, filter); // one slot per dictionary entry, plus the null group at index `groups` Accumulators acc = new Accumulators(groups + 1, lane); - long n = chunk.rowCount(); for (long i = 0; i < n; i++) { if (selected != null && !selected.test(i)) { continue; } int g; - if (groupValidity != null && !groupValidity.getBoolean(i)) { + if (validity != null && !validity.getBoolean(i)) { g = groups; } else { long code = dict.code(i + offset); @@ -92,7 +104,7 @@ case VarBinSlicedArray(var _, var _, VarBinDictArray d, long sliceOffset) -> { } acc.add(g, i); } - Map result = new HashMap<>(); + Map result = new HashMap<>(); for (int g = 0; g <= groups; g++) { if (acc.rows[g] > 0) { String key = g == groups ? null : values.getString(g); @@ -103,6 +115,58 @@ case VarBinSlicedArray(var _, var _, VarBinDictArray d, long sliceOffset) -> { return Optional.of(result); } + /// Groups by a 64-bit key read per row: ids are assigned in a first pass through a primitive + /// map, then the rows fold per id in a second, so no key is boxed per row. + private static Optional> byKey(LongUnaryOperator keyOf, BoolArray validity, + LongPredicate selected, Lane lane, long n, LongFunction boxKey) { + int[] ids = new int[Math.toIntExact(n)]; + LongIntMap idOf = new LongIntMap(64); + long[] keys = new long[16]; + int groups = 0; + boolean nulls = false; + for (int i = 0; i < n; i++) { + if (selected != null && !selected.test(i)) { + ids[i] = -1; + continue; + } + if (validity != null && !validity.getBoolean(i)) { + ids[i] = Integer.MAX_VALUE; // the null group, renumbered below + nulls = true; + continue; + } + long key = keyOf.applyAsLong(i); + int id = idOf.get(key); + if (id < 0) { + if (groups == MAX_GROUPS) { + return Optional.empty(); + } + id = groups++; + idOf.put(key, id); + if (id == keys.length) { + keys = Arrays.copyOf(keys, keys.length * 2); + } + keys[id] = key; + } + ids[i] = id; + } + // the null group takes the slot after the keyed groups + Accumulators acc = new Accumulators(groups + 1, lane); + for (int i = 0; i < n; i++) { + int id = ids[i]; + if (id >= 0) { + acc.add(id == Integer.MAX_VALUE ? groups : id, i); + } + } + Map result = new HashMap<>(); + for (int g = 0; g < groups; g++) { + result.put(boxKey.apply(keys[g]), acc.result(g)); + } + if (nulls) { + result.put(null, acc.result(groups)); + } + return Optional.of(result); + } + /// Folds two partial results of one group, as two dictionary entries with equal strings produce. private static FilteredAggregate combine(FilteredAggregate a, FilteredAggregate b) { Number sum; 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 ca57c8c89..e72b356ad 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 @@ -206,7 +206,7 @@ void countOnly_countsRowsPerGroupWithNullAsItsOwnGroup() { Chunk chunk = chunk(nullableGroups(), null); // When - Map result = Compute.groupedAggregate(chunk, G, null, null).orElseThrow(); + Map result = Compute.groupedAggregate(chunk, G, null, null).orElseThrow(); // Then assertThat(result).containsOnlyKeys("a", "b", null); @@ -222,7 +222,7 @@ void longColumn_foldsSumMinMaxSkippingNullValues() { new boolean[]{true, true, false, true, true, true}); // When - Map result = Compute.groupedAggregate(chunk(nullableGroups(), v), G, null, V) + Map result = Compute.groupedAggregate(chunk(nullableGroups(), v), G, null, V) .orElseThrow(); // Then — a: rows 0, 2, 5 with row 2 null; b: rows 1, 4; null group: row 3 @@ -253,7 +253,7 @@ void filter_foldsOnlyTheSelectedRows() { Chunk chunk = chunk(nullableGroups(), v); // When - Map result = Compute.groupedAggregate(chunk, G, RowFilter.gt(V, 4L), V) + Map result = Compute.groupedAggregate(chunk, G, RowFilter.gt(V, 4L), V) .orElseThrow(); // Then — groups with no selected row do not appear @@ -270,7 +270,7 @@ void slicedDictionary_readsCodesAtTheSliceOffset() { Chunk chunk = chunk(new VarBinSlicedArray(DType.UTF8, 3, dict, 2), null); // When - Map result = Compute.groupedAggregate(chunk, G, null, null).orElseThrow(); + Map result = Compute.groupedAggregate(chunk, G, null, null).orElseThrow(); // Then assertThat(result.get("b").selectedRows()).isEqualTo(2); @@ -285,7 +285,7 @@ void duplicateDictionaryValues_foldIntoOneGroup() { Array v = ComputeArrays.longArray(ARENA, 1, 10, 100, 1000); // When - Map result = Compute.groupedAggregate(chunk(group, v), G, null, V).orElseThrow(); + Map result = Compute.groupedAggregate(chunk(group, v), G, null, V).orElseThrow(); // Then assertThat(result.get("x")).isEqualTo(new FilteredAggregate(3, 3, 1011L, 1L, 1000L)); @@ -317,6 +317,65 @@ void codesBufferShorterThanRows_throwsVortexException() { .isInstanceOf(VortexException.class); } + @Test + void booleanGroupColumn_foldsFalseTrueAndNull() { + // Given — g = true, false, true, (null); v = 1, 2, 3, 4 + Array group = new MaskedArray(ComputeArrays.boolArray(ARENA, true, false, true, true), + ComputeArrays.boolArray(ARENA, true, true, true, false)); + Array v = ComputeArrays.longArray(ARENA, 1, 2, 3, 4); + + // When + Map result = Compute.groupedAggregate(chunk(group, v), G, null, V).orElseThrow(); + + // Then + assertThat(result).containsOnlyKeys(true, false, null); + assertThat(result.get(true)).isEqualTo(new FilteredAggregate(2, 2, 4L, 1L, 3L)); + assertThat(result.get(false)).isEqualTo(new FilteredAggregate(1, 1, 2L, 2L, 2L)); + assertThat(result.get(null)).isEqualTo(new FilteredAggregate(1, 1, 4L, 4L, 4L)); + } + + @Test + void integerGroupColumn_foldsEachDistinctValueWithAFilter() { + // Given — g = 7, -3, 7, 7, -3; keep v > 1 + Array group = ComputeArrays.longArray(ARENA, 7, -3, 7, 7, -3); + Array v = ComputeArrays.longArray(ARENA, 1, 2, 3, 4, 5); + + // When + Map result = Compute.groupedAggregate(chunk(group, v), G, RowFilter.gt(V, 1L), V) + .orElseThrow(); + + // Then + assertThat(result).containsOnlyKeys(7L, -3L); + assertThat(result.get(7L)).isEqualTo(new FilteredAggregate(2, 2, 7L, 3L, 4L)); + assertThat(result.get(-3L)).isEqualTo(new FilteredAggregate(2, 2, 7L, 2L, 5L)); + } + + /// A u64 key past 2^63 has no signed-long home, so the kernel declines rather than group it + /// under a negative key. + @Test + void u64GroupColumn_isEmpty() { + // Given + Chunk chunk = chunk(ComputeArrays.unsignedLongArray(ARENA, 1, -1), null); + + // When / Then + assertThat(Compute.groupedAggregate(chunk, G, null, null)).isEmpty(); + } + + /// An integer column may hold a value per row; past 65,536 groups in a chunk the per-group + /// arrays would grow with the rows, so the kernel declines. + @Test + void integerGroupColumnWithTooManyValues_isEmpty() { + // Given + long[] values = new long[70_000]; + for (int i = 0; i < values.length; i++) { + values[i] = i; + } + Chunk chunk = chunk(ComputeArrays.longArray(ARENA, values), null); + + // When / Then + assertThat(Compute.groupedAggregate(chunk, G, null, null)).isEmpty(); + } + /// The caller groups the rows itself when the kernel cannot: it must say so, not guess. @Test void groupColumnNotDictionaryEncoded_isEmpty() { diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java b/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java index d1a889e2d..b2a5d9f5d 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/DictColumnState.java @@ -4,7 +4,7 @@ import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.PType; import io.github.dfa1.vortex.core.proto.ProtoScalarValue; -import io.github.dfa1.vortex.writer.encode.LongIntMap; +import io.github.dfa1.vortex.core.compute.LongIntMap; import io.github.dfa1.vortex.writer.encode.MinMax; import io.github.dfa1.vortex.writer.encode.NullableData; import io.github.dfa1.vortex.writer.encode.PrimitiveEncodingEncoder; diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java index 69c5390d6..dcab9124c 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java @@ -1,5 +1,6 @@ package io.github.dfa1.vortex.writer.encode; +import io.github.dfa1.vortex.core.compute.LongIntMap; import io.github.dfa1.vortex.core.io.VortexFormat; import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.PType; diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ArrayStats.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ArrayStats.java index 193b58f9d..75c8b4ccc 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ArrayStats.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/ArrayStats.java @@ -1,5 +1,6 @@ package io.github.dfa1.vortex.writer.encode; +import io.github.dfa1.vortex.core.compute.LongIntMap; import io.github.dfa1.vortex.core.model.PType; import io.github.dfa1.vortex.core.simd.SimdOperationsSupport; diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java index 4c9c26df8..534f855d4 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/DictEncodingEncoder.java @@ -1,5 +1,6 @@ package io.github.dfa1.vortex.writer.encode; +import io.github.dfa1.vortex.core.compute.LongIntMap; import io.github.dfa1.vortex.core.compute.PrimitiveArrays; import io.github.dfa1.vortex.core.model.DType; import io.github.dfa1.vortex.core.model.PType; diff --git a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java index 1ff7ca22d..e9543add8 100644 --- a/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java +++ b/writer/src/main/java/io/github/dfa1/vortex/writer/encode/OnPairEncodingEncoder.java @@ -1,5 +1,6 @@ package io.github.dfa1.vortex.writer.encode; +import io.github.dfa1.vortex.core.compute.LongIntMap; import io.github.dfa1.vortex.core.compute.PrimitiveArrays; import io.github.dfa1.vortex.core.io.PTypeIO; import io.github.dfa1.vortex.core.model.DType;