perf(calcite): fold GROUP BY on boolean and integer columns per value

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

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<RelDataType> 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<AggregateCall> 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<Object> foldKey = List.of("grouped", table, group, filter == null ? List.of() : filter, aggColumns);
-        Optional<Map<String, VortexTable.GroupFold>> folds =
+        Optional<Map<Object, VortexTable.GroupFold>> folds =
                 memo(call.getPlanner(), foldKey, () -> table.groupedFold(group, filter, aggColumns));
         if (folds.isEmpty()) {
             return;
         }
         RelDataType groupType = aggregate.getRowType().getFieldList().getFirst().getType();
         List<ImmutableList<RexLiteral>> tuples = new ArrayList<>();
-        for (Map.Entry<String, VortexTable.GroupFold> entry : folds.get().entrySet()) {
-            RexLiteral key = entry.getKey() == null ? rexBuilder.makeNullLiteral(groupType)
-                    : minMaxLiteral(rexBuilder, entry.getKey(), groupType);
+        for (Map.Entry<Object, VortexTable.GroupFold> 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
             }

calcite/src/main/java/io/github/dfa1/vortex/calcite/VortexTable.java

@@ -358,20 +358,21 @@ public Optional<FilteredFold> filteredFold(RowFilter filter, ColumnName aggColum
     public record GroupFold(long rows, Map<ColumnName, FilteredFold> 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<Map<String, GroupFold>> groupedFold(ColumnName group, RowFilter filter, Set<ColumnName> aggColumns) {
+    /// @return the fold per group value (`String`, `Boolean` or `Long`; `null` is the null group), or
+    ///         empty
+    public Optional<Map<Object, GroupFold>> groupedFold(ColumnName group, RowFilter filter, Set<ColumnName> aggColumns) {
         java.util.LinkedHashSet<ColumnName> columns = new java.util.LinkedHashSet<>();
         columns.add(group);
         if (filter != null) {
@@ -385,20 +386,20 @@ public Optional<Map<String, GroupFold>> groupedFold(ColumnName group, RowFilter
         // one pass per aggregate column (a null column counts rows only)
         List<ColumnName> passes = aggColumns.isEmpty() ? java.util.Collections.singletonList(null)
                 : List.copyOf(aggColumns);
-        Map<String, long[]> rows = new java.util.HashMap<>();
-        Map<String, Map<ColumnName, Fold>> byGroup = new java.util.HashMap<>();
+        Map<Object, long[]> rows = new java.util.HashMap<>();
+        Map<Object, Map<ColumnName, Fold>> 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<Map<String, FilteredAggregate>> part =
+                        Optional<Map<Object, FilteredAggregate>> part =
                                 Compute.groupedAggregate(chunk, group, filter, aggColumn);
                         if (part.isEmpty()) {
                             return Optional.empty();
                         }
-                        for (Map.Entry<String, FilteredAggregate> entry : part.get().entrySet()) {
+                        for (Map.Entry<Object, FilteredAggregate> 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<Map<String, GroupFold>> groupedFold(ColumnName group, RowFilter
         } catch (IOException e) {
             throw new UncheckedIOException("cannot fold groups of " + file, e);
         }
-        Map<String, GroupFold> result = new java.util.HashMap<>();
+        Map<Object, GroupFold> result = new java.util.HashMap<>();
         rows.forEach((key, count) -> {
             Map<ColumnName, FilteredFold> columnFolds = new java.util.HashMap<>();
             byGroup.getOrDefault(key, Map.of()).forEach((column, fold) -> columnFolds.put(column, fold.result()));

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",

writer/src/main/java/io/github/dfa1/vortex/writer/encode/LongIntMap.java → 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<Object, Integer>` 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;

writer/src/test/java/io/github/dfa1/vortex/writer/encode/LongIntMapTest.java → 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;

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`,

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<String> result = sqlRows(conn, sql);
+
+            // Then
+            assertThat(plan).containsIgnoringCase("Values").doesNotContain("TableScan");
+            assertThat(result).isEqualTo(sqlRows(conn, oracle)).hasSize(7);
+        }
+    }
+
     private static List<String> sqlRows(Connection conn, String sql) throws Exception {
         List<String> out = new java.util.ArrayList<>();
         try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery(sql)) {

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<Map<String, FilteredAggregate>> groupedAggregate(Chunk chunk, ColumnName groupColumn,
+    public static Optional<Map<Object, FilteredAggregate>> groupedAggregate(Chunk chunk, ColumnName groupColumn,
                                                                             RowFilter filter, ColumnName aggColumn) {
         Objects.requireNonNull(chunk, "chunk");
         Objects.requireNonNull(groupColumn, "groupColumn");