perf(calcite): fold GROUP BY on a dictionary string column per code

CHANGELOG.md

@@ -8,6 +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))
 - `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))
@@ -16,6 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Custom extension types: register an `ExtensionDecoder` on `ReadRegistry` and an `ExtensionEncoder` on `WriteRegistry`, with a `validate` hook on both sides, and `Chunk.as` decodes any registered type. ([#530](https://github.com/dfa1/vortex-java/issues/530))
 
 ### Changed
+- 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))
+- 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))
 - The Calcite adapter reads a table's schema and row count once instead of reopening the file for every planning call. ([#550](https://github.com/dfa1/vortex-java/pull/550))

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

@@ -5,6 +5,7 @@
 
 import com.google.common.collect.ImmutableList;
 import org.apache.calcite.interpreter.Bindables;
+import org.apache.calcite.plan.RelOptPlanner;
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptRuleCall;
 import org.apache.calcite.plan.RelOptRuleOperand;
@@ -29,12 +30,16 @@
 import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 
 /// Rewrites a whole-table `MIN`/`MAX`/`COUNT`/`SUM` aggregate over a [VortexTable] into a
 /// single-row [LogicalValues] computed from the footer zone-map statistics — answering the query
 /// without decoding a single data segment (ADR 0013 §6, ADR 0018 Phase 2).
 ///
+/// A `GROUP BY` on one dictionary-encoded string column is rewritten the same way, into one
+/// [LogicalValues] row per group, folded per dictionary code instead of from statistics.
+///
 /// Fires only when it can answer *every* aggregate from statistics: no `GROUP BY`, and each call
 /// is `COUNT(*)`, `COUNT(col)`, `MIN(col)`, `MAX(col)` (numeric or `VARCHAR`/`CHAR`), or `SUM(col)`
 /// over a numeric column, with no `DISTINCT` and no `FILTER (WHERE ...)` clause. `SUM` folds the per-zone `SUM` rows via [VortexTable#zoneSum(ColumnName)]; it
@@ -79,14 +84,28 @@ public final class VortexAggregatePushDownRule extends RelOptRule {
     public static final java.util.List<RelOptRule> RULES =
             java.util.List.of(WITH_PROJECT, NO_PROJECT, PROJECT_FILTER, FILTER_ONLY);
 
+    /// Plan-time folds per planning run. Calcite fires this rule several times per query (6 to 9
+    /// measured), once per matching operand and subset, and each firing would re-read the columns.
+    /// Keyed weakly by the planner, so a fold lives as long as its query's planning and the next
+    /// query computes its own.
+    private static final Map<RelOptPlanner, Map<List<Object>, Object>> FOLDS =
+            java.util.Collections.synchronizedMap(new java.util.WeakHashMap<>());
+
     private VortexAggregatePushDownRule(RelOptRuleOperand operand, String description) {
         super(operand, description);
     }
 
+    /// Returns the fold cached under `key` for `planner`'s run, computing it on the first call.
+    @SuppressWarnings("unchecked")
+    private static <T> T memo(RelOptPlanner planner, List<Object> key, java.util.function.Supplier<T> compute) {
+        Map<List<Object>, Object> folds = FOLDS.computeIfAbsent(planner, p -> new java.util.concurrent.ConcurrentHashMap<>());
+        return (T) folds.computeIfAbsent(key, k -> compute.get());
+    }
+
     @Override
     public void onMatch(RelOptRuleCall call) {
         Aggregate aggregate = call.rel(0);
-        if (aggregate.getGroupCount() != 0) {
+        if (aggregate.getGroupCount() > 1 || aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
             return;
         }
         // Statistics describe every row of a zone, so they cannot answer an aggregate over a subset
@@ -144,11 +163,17 @@ public void onMatch(RelOptRuleCall call) {
         List<RelDataType> outTypes = aggregate.getRowType().getFieldList().stream()
                 .map(f -> f.getType()).toList();
 
+        if (aggregate.getGroupCount() == 1) {
+            rewriteGrouped(call, aggregate, table, scanColumns, scanRowType, project, rexBuilder, outTypes,
+                    pushedFilter);
+            return;
+        }
+
         List<RexLiteral> row = new ArrayList<>();
         List<AggregateCall> calls = aggregate.getAggCallList();
         for (int i = 0; i < calls.size(); i++) {
             RexLiteral literal = evaluate(calls.get(i), outTypes.get(i), table, scanColumns, scanRowType,
-                    project, rexBuilder, pushedFilter);
+                    project, rexBuilder, pushedFilter, call.getPlanner());
             if (literal == null) {
                 return; // an aggregate we can't answer from stats — abandon the rewrite
             }
@@ -161,28 +186,113 @@ 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)]
+    /// 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
+    /// many groups).
+    private static void rewriteGrouped(RelOptRuleCall call, Aggregate aggregate, VortexTable table,
+                                       List<ColumnName> scanColumns, RelDataType scanRowType, Project project,
+                                       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())) {
+            return;
+        }
+        List<AggregateCall> calls = aggregate.getAggCallList();
+        List<ColumnName> callColumns = new ArrayList<>();
+        java.util.Set<ColumnName> aggColumns = new java.util.LinkedHashSet<>();
+        for (AggregateCall aggCall : calls) {
+            SqlKind kind = aggCall.getAggregation().getKind();
+            if (kind == SqlKind.COUNT && aggCall.getArgList().isEmpty()) {
+                callColumns.add(null);
+                continue;
+            }
+            boolean supported = kind == SqlKind.COUNT || kind == SqlKind.SUM || kind == SqlKind.MIN
+                    || kind == SqlKind.MAX;
+            ColumnName column = aggCall.getArgList().size() == 1
+                    ? resolveColumn(aggCall.getArgList().getFirst(), scanColumns, project) : null;
+            if (!supported || column == null) {
+                return;
+            }
+            callColumns.add(column);
+            aggColumns.add(column);
+        }
+        // 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 =
+                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);
+            if (key == null) {
+                return; // a key the column's charset cannot hold as a literal
+            }
+            List<RexLiteral> row = new ArrayList<>();
+            row.add(key);
+            for (int i = 0; i < calls.size(); i++) {
+                RexLiteral literal = groupedLiteral(calls.get(i), callColumns.get(i), entry.getValue(),
+                        outTypes.get(i + 1), scanRowType, rexBuilder);
+                if (literal == null) {
+                    return;
+                }
+                row.add(literal);
+            }
+            tuples.add(ImmutableList.copyOf(row));
+        }
+        call.transformTo(LogicalValues.create(aggregate.getCluster(), aggregate.getRowType(),
+                ImmutableList.copyOf(tuples)));
+    }
+
+    /// One aggregate call's literal for one group, or `null` when it cannot be built.
+    private static RexLiteral groupedLiteral(AggregateCall aggCall, ColumnName column, VortexTable.GroupFold fold,
+                                             RelDataType outType, RelDataType scanRowType, RexBuilder rexBuilder) {
+        if (column == null) {
+            return exact(rexBuilder, fold.rows(), outType);
+        }
+        VortexTable.FilteredFold columnFold = fold.columns().get(column);
+        long nulls = columnFold.nullCount();
+        return switch (aggCall.getAggregation().getKind()) {
+            case COUNT -> exact(rexBuilder, columnFold.rows() - nulls, outType);
+            case SUM -> sumLiteral(rexBuilder, outType, columnFold.sum(), nulls, columnFold.rows(),
+                    isNullable(scanRowType, column.value()));
+            case MIN, MAX -> {
+                Object value = aggCall.getAggregation().getKind() == SqlKind.MIN ? columnFold.min() : columnFold.max();
+                yield value == null ? rexBuilder.makeNullLiteral(outType) : minMaxLiteral(rexBuilder, value, outType);
+            }
+            default -> null;
+        };
+    }
+
     /// Evaluates one aggregate call from zone-map stats, returning a literal of `outType`, or
     /// `null` if it cannot be answered (so the caller abandons the rewrite). When `filter` is
     /// non-null the answer is folded from only the zones the filter fully selects, which exists
     /// only if the filter partitions the zones cleanly — otherwise the fold is empty and the call
     /// abandons.
     private static RexLiteral evaluate(AggregateCall agg, RelDataType outType, VortexTable table,
                                        List<ColumnName> scanColumns, RelDataType scanRowType,
-                                       Project project, RexBuilder rexBuilder, RowFilter filter) {
+                                       Project project, RexBuilder rexBuilder, RowFilter filter,
+                                       RelOptPlanner planner) {
         return switch (agg.getAggregation().getKind()) {
             case COUNT -> {
                 if (agg.getArgList().isEmpty()) {                  // COUNT(*)
                     if (filter == null) {
                         yield exact(rexBuilder, table.totalRows(), outType);
                     }
-                    Optional<VortexTable.FilteredFold> fold = table.filteredFold(filter, null);
+                    Optional<VortexTable.FilteredFold> fold = filteredFold(planner, table, filter, null);
                     yield fold.map(f -> exact(rexBuilder, f.rows(), outType)).orElse(null);
                 }
                 ColumnName col = resolveColumn(agg.getArgList().getFirst(), scanColumns, project);
                 if (col == null) {
                     yield null;
                 }
-                ColumnFold fold = columnFold(table, col, filter);
+                ColumnFold fold = columnFold(table, col, filter, planner);
                 if (fold == null) {
                     yield null;
                 }
@@ -203,7 +313,7 @@ private static RexLiteral evaluate(AggregateCall agg, RelDataType outType, Vorte
                 if (col == null) {
                     yield null;
                 }
-                ColumnFold fold = columnFold(table, col, filter);
+                ColumnFold fold = columnFold(table, col, filter, planner);
                 if (fold == null) {
                     yield null;
                 }
@@ -226,7 +336,7 @@ private static RexLiteral evaluate(AggregateCall agg, RelDataType outType, Vorte
                 if (col == null) {
                     yield null;
                 }
-                ColumnFold fold = columnFold(table, col, filter);
+                ColumnFold fold = columnFold(table, col, filter, planner);
                 if (fold == null) {
                     yield null;
                 }
@@ -244,18 +354,26 @@ yield sumLiteral(rexBuilder, outType, fold.sum(), fold.nullCount(), fold.rows(),
     private record ColumnFold(Number sum, Object min, Object max, Long nullCount, long rows) {
     }
 
-    private static ColumnFold columnFold(VortexTable table, ColumnName column, RowFilter filter) {
+    private static ColumnFold columnFold(VortexTable table, ColumnName column, RowFilter filter,
+                                         RelOptPlanner planner) {
         if (filter == null) {
             VortexTable.ColumnStats stats = table.statsAndRows(column);
             VortexTable.ZoneSum zoneSum = table.zoneSum(column);
             return new ColumnFold(zoneSum.sum(), stats.stats().min(), stats.stats().max(),
                     stats.stats().nullCount(), stats.totalRows());
         }
-        return table.filteredFold(filter, column)
+        return filteredFold(planner, table, filter, column)
                 .map(f -> new ColumnFold(f.sum(), f.min(), f.max(), f.nullCount(), f.rows()))
                 .orElse(null);
     }
 
+    /// [VortexTable#filteredFold(RowFilter, ColumnName)], computed once per planner run.
+    private static Optional<VortexTable.FilteredFold> filteredFold(RelOptPlanner planner, VortexTable table,
+                                                                    RowFilter filter, ColumnName column) {
+        List<Object> foldKey = List.of("filtered", table, filter, column == null ? List.of() : column);
+        return memo(planner, foldKey, () -> table.filteredFold(filter, column));
+    }
+
     /// Builds the literal for a `SUM`, resolving the all-null-vs-genuine-zero ambiguity from the
     /// null count. SQL `SUM` over zero non-null rows is `NULL`, but the writer records an all-null
     /// zone's sum as a sum-neutral `0` (validity placeholders are zero), so a `0` fold alone cannot