From d31de14d1256dce95da3d0e590de930443b95a90 Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 18 Oct 2026 12:54:15 +0200 Subject: [PATCH 1/2] refactor: move LongIntMap to core.compute The reader needs the same primitive long-to-int map for grouped aggregation that the writer uses for dictionary encoding and stats, and the reader cannot depend on the writer. core.compute already holds the kernels both sides share. Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01WtuVGQ5h5Lk6Zn4JnkxF1K --- .../io/github/dfa1/vortex/core/compute}/LongIntMap.java | 6 ++++-- .../io/github/dfa1/vortex/core/compute}/LongIntMapTest.java | 2 +- .../java/io/github/dfa1/vortex/writer/DictColumnState.java | 2 +- .../dfa1/vortex/writer/encode/AlpRdEncodingEncoder.java | 1 + .../io/github/dfa1/vortex/writer/encode/ArrayStats.java | 1 + .../dfa1/vortex/writer/encode/DictEncodingEncoder.java | 1 + .../dfa1/vortex/writer/encode/OnPairEncodingEncoder.java | 1 + 7 files changed, 10 insertions(+), 4 deletions(-) rename {writer/src/main/java/io/github/dfa1/vortex/writer/encode => core/src/main/java/io/github/dfa1/vortex/core/compute}/LongIntMap.java (98%) rename {writer/src/test/java/io/github/dfa1/vortex/writer/encode => core/src/test/java/io/github/dfa1/vortex/core/compute}/LongIntMapTest.java (99%) 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/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; From ac5838501fa2a49dba1562738bb951a6a719532f Mon Sep 17 00:00:00 2001 From: Davide Angelocola Date: Sun, 18 Oct 2026 12:59:20 +0200 Subject: [PATCH 2/2] perf(calcite): fold GROUP BY on boolean and integer columns per value Only a dictionary-encoded string column grouped in the scan; a GROUP BY on a boolean or integer column built a row per record for Calcite to hash. Compute.groupedAggregate now also groups by a boolean column (two ids) and an integer column other than u64: a first pass assigns group ids through the primitive LongIntMap, a second folds the rows per id, so no key is boxed per row. More than 65,536 distinct values in a chunk declines, and so do more than 10,000 groups in all, as before. The result keys are Object (String, Boolean or Long); the method is unreleased, so the signature changes in place. Raincloud, best of 3, ms: appian-order group by order_isExpeditedShipped 116 -> 32 store_sales group by ss_store_sk, count, sum 154 -> 76 imdb title group by kind_id 93 -> 36 customer where segment = 'BUILDING' group by nation 30 -> 16 Co-Authored-By: Claude Sonnet 5 Claude-Session: https://claude.ai/code/session_01WtuVGQ5h5Lk6Zn4JnkxF1K --- CHANGELOG.md | 3 +- .../calcite/VortexAggregatePushDownRule.java | 18 ++- .../dfa1/vortex/calcite/VortexTable.java | 27 ++-- .../calcite/AggregateRuleBranchTest.java | 2 +- docs/reference.md | 10 +- .../CalciteInteropIntegrationTest.java | 28 ++++ .../dfa1/vortex/reader/compute/Compute.java | 25 ++-- .../reader/compute/GroupedAggregate.java | 134 +++++++++++++----- .../vortex/reader/compute/ComputeTest.java | 69 ++++++++- 9 files changed, 238 insertions(+), 78 deletions(-) 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/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() {