diff --git a/CHANGELOG.md b/CHANGELOG.md index 3d42e9f5..865ea977 100644 --- a/CHANGELOG.md +++ b/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)) 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 c0c80f6e..61f62af0 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 @@ -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 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, 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 memo(RelOptPlanner planner, List key, java.util.function.Supplier compute) { + Map, 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 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 row = new ArrayList<>(); List 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,6 +186,90 @@ 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 scanColumns, RelDataType scanRowType, Project project, + 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())) { + return; + } + List calls = aggregate.getAggCallList(); + List callColumns = new ArrayList<>(); + java.util.Set 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 foldKey = List.of("grouped", table, group, filter == null ? List.of() : filter, aggColumns); + 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); + if (key == null) { + return; // a key the column's charset cannot hold as a literal + } + List 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 @@ -168,21 +277,22 @@ public void onMatch(RelOptRuleCall call) { /// abandons. private static RexLiteral evaluate(AggregateCall agg, RelDataType outType, VortexTable table, List 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 fold = table.filteredFold(filter, null); + Optional 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 filteredFold(RelOptPlanner planner, VortexTable table, + RowFilter filter, ColumnName column) { + List 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 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 ad338536..a4cac876 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 @@ -39,7 +39,9 @@ import java.nio.file.Path; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Optional; +import java.util.Set; import java.util.concurrent.atomic.AtomicLong; /// A single Vortex file exposed to Calcite as a flat SQL table with column projection and @@ -345,6 +347,91 @@ public Optional filteredFold(RowFilter filter, ColumnName aggColum } } + /// Most groups [#groupedFold(ColumnName, RowFilter, Set)] answers. Each group becomes one literal + /// row of the plan, so a high-cardinality column is left to the scan. + private static final int MAX_GROUPS = 10_000; + + /// One group's fold, returned by [#groupedFold(ColumnName, RowFilter, Set)]. + /// + /// @param rows the rows of the group the filter selects + /// @param columns per aggregate column, its fold over those rows + 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 + /// 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. + /// + /// @param group the string 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) { + java.util.LinkedHashSet columns = new java.util.LinkedHashSet<>(); + columns.add(group); + if (filter != null) { + RexFilterTranslator.collectColumns(filter, columns); + } + columns.addAll(aggColumns); + ScanOptions options = ScanOptions.columns(columns.toArray(ColumnName[]::new)); + if (filter != null) { + options = options.withFilter(filter); + } + // 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<>(); + 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 = + Compute.groupedAggregate(chunk, group, filter, aggColumn); + if (part.isEmpty()) { + return Optional.empty(); + } + 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(); + } + if (aggColumn != null) { + Fold fold = byGroup.computeIfAbsent(entry.getKey(), k -> new java.util.HashMap<>()) + .computeIfAbsent(aggColumn, k -> new Fold()); + fold.keptRows += aggregate.selectedRows(); + // the kernel's min/max are null only when no selected row has a value + fold.addSum(aggregate.sum()); + fold.addMin(aggregate.min(), false); + fold.addMax(aggregate.max(), false); + fold.addNulls(aggregate.selectedRows() - aggregate.aggNonNullCount()); + } + } + } + if (rows.size() > MAX_GROUPS) { + return Optional.empty(); + } + } + } + } catch (IOException e) { + throw new UncheckedIOException("cannot fold groups of " + file, e); + } + 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())); + result.put(key, new GroupFold(count[0], Map.copyOf(columnFolds))); + }); + return Optional.of(java.util.Collections.unmodifiableMap(result)); + } + /// Decodes the chunks that hold boundary zones, each once, in row order, and folds a zone's rows /// through a [Chunk#slice(long, long)] view. Zones arrive in row order, so only the chunk under /// the current zone is kept open. 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 afdd1d75..02f7f116 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 @@ -124,10 +124,24 @@ void minOverComputedExpression_abandonsRewrite() { } @Test - void groupedAggregate_isLeftUntouched() { - // Given a GROUP BY — group count != 0, the rule returns immediately - // When / Then — Aggregate stays - assertThat(optimize("select symbol, max(high) from ohlc group by symbol")).contains("Aggregate"); + void groupedOnDictionaryString_rewritesToValues() { + // Given a GROUP BY on the dictionary-encoded symbol column + // When / Then — folded per dictionary code into one Values row per symbol + assertThat(optimize("select symbol, max(high), count(*) from ohlc group by symbol")) + .contains("LogicalValues").doesNotContain("Aggregate"); + } + + @ParameterizedTest + @ValueSource(strings = { + // a numeric group key has no dictionary codes to fold by + "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", + // AVG reaches the rule unreduced here; only COUNT/SUM/MIN/MAX fold + "select symbol, avg(high) from ohlc group by symbol"}) + void groupedAggregateTheFoldCannotAnswer_isLeftUntouched(String sql) { + // Given / When / Then — Aggregate stays + assertThat(optimize(sql)).contains("Aggregate"); } /// Zone statistics cover every row, so they cannot answer an aggregate over a subset of rows: diff --git a/docs/reference.md b/docs/reference.md index b1833afb..c3791c38 100644 --- a/docs/reference.md +++ b/docs/reference.md @@ -636,6 +636,11 @@ column-vs-literal comparisons, `IN` lists and `IS [NOT] NULL` joined by `AND`: z 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. `connect` is the supported entry point; `VortexSchema`/`VortexTable` (`io.github.dfa1.vortex.calcite`) are the underlying schema/table plumbing it assembles, exposed for callers building a Calcite schema tree 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 4d4c316c..6ea2bce6 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 @@ -264,6 +264,121 @@ void jniWriter_minMaxOfStringsOutsideLatin1_answer(@TempDir Path tmp) throws Exc assertThat(max).isEqualTo("\u03A9mega"); } + @ParameterizedTest(name = "{0}") + @org.junit.jupiter.params.provider.ValueSource(strings = {"", " where id < 1234", " where x > 0.5 and id >= 100"}) + void jniWriter_groupByDictionaryString_foldsEveryGroupLikeTheRows(String where, @TempDir Path tmp) + throws Exception { + // Given — Rust dict-encodes a low-cardinality string column with nulls. GROUP BY on it folds + // per dictionary code into a Values row per group; every group's aggregates must equal the + // ones computed from the rows written + assertGroupBy(where, new String[]{"AAPL", "MSFT", "NVDA", "caf\u00E9", "AMZN", "META", "TSLA", null}, + true, tmp); + } + + @Test + void jniWriter_groupByStringOutsideLatin1_answersThroughTheScan(@TempDir Path tmp) throws Exception { + // Given — a group key Latin-1 cannot hold has no VARCHAR literal, so the fold leaves the query + // to the scan; the answer must not change + assertGroupBy("", new String[]{"AAPL", "\u03A9mega", null}, false, tmp); + } + + private static void assertGroupBy(String where, String[] symbols, boolean folded, Path tmp) throws Exception { + Path file = tmp.resolve("jni_groups.vortex"); + List rows = writeGroupsWithJni(file, symbols); + java.util.function.Predicate selected = switch (where) { + case "" -> r -> true; + case " where id < 1234" -> r -> (long) r[1] < 1234; + default -> r -> (double) r[3] > 0.5 && (long) r[1] >= 100; + }; + Map> byGroup = new java.util.HashMap<>(); + rows.stream().filter(selected).forEach(r -> byGroup.computeIfAbsent((String) r[0], k -> new java.util.ArrayList<>()).add(r)); + Map expected = new java.util.HashMap<>(); + byGroup.forEach((group, groupRows) -> { + List ns = groupRows.stream().map(r -> (Long) r[2]).filter(java.util.Objects::nonNull).toList(); + expected.put(group, groupRows.size() + "|" + ns.size() + + "|" + groupRows.stream().mapToLong(r -> (long) r[1]).sum() + + "|" + groupRows.stream().mapToLong(r -> (long) r[1]).min().orElseThrow() + + "|" + (ns.isEmpty() ? null : java.util.Collections.max(ns)) + + "|" + groupRows.stream().mapToDouble(r -> (double) r[3]).max().orElseThrow()); + }); + String sql = "select s, count(*), count(n), sum(id), min(id), max(n), max(x) from vtx.t" + where + " group by s"; + + 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); + } + Map result = new java.util.HashMap<>(); + try (Statement st = conn.createStatement(); ResultSet rs = st.executeQuery(sql)) { + while (rs.next()) { + result.put(rs.getString(1), rs.getLong(2) + "|" + rs.getLong(3) + "|" + rs.getLong(4) + + "|" + rs.getLong(5) + "|" + rs.getObject(6) + "|" + rs.getDouble(7)); + } + } + + // Then — equal to the rows, and folded rather than scanned when the keys allow it + assertThat(result).isEqualTo(expected); + if (folded) { + assertThat(plan).containsIgnoringCase("Values").doesNotContain("TableScan"); + } + } + } + + /// Writes `s` drawn from `symbols`, `id`, `n` (null every 5th row, and always null for the first + /// symbol so its MAX(n) is SQL NULL) and `x`; returns the rows written. + private static List writeGroupsWithJni(Path file, String[] symbols) throws Exception { + Schema schema = new Schema(List.of( + Field.nullable("s", new ArrowType.Utf8()), + Field.notNullable("id", new ArrowType.Int(64, true)), + Field.nullable("n", new ArrowType.Int(64, true)), + Field.notNullable("x", new ArrowType.FloatingPoint(org.apache.arrow.vector.types.FloatingPointPrecision.DOUBLE)))); + List rows = new java.util.ArrayList<>(); + java.util.Random random = new java.util.Random(7); + try (VortexWriter writer = VortexWriter.builder(SESSION, file.toAbsolutePath().toUri().toString(), schema, + ALLOCATOR).build()) { + for (int batch = 0; batch < BATCHES; batch++) { + try (VectorSchemaRoot root = VectorSchemaRoot.create(schema, ALLOCATOR)) { + org.apache.arrow.vector.VarCharVector s = (org.apache.arrow.vector.VarCharVector) root.getVector("s"); + BigIntVector ids = (BigIntVector) root.getVector("id"); + BigIntVector n = (BigIntVector) root.getVector("n"); + org.apache.arrow.vector.Float8Vector x = (org.apache.arrow.vector.Float8Vector) root.getVector("x"); + s.allocateNew(BATCH_ROWS); + ids.allocateNew(BATCH_ROWS); + n.allocateNew(BATCH_ROWS); + x.allocateNew(BATCH_ROWS); + for (int i = 0; i < BATCH_ROWS; i++) { + long id = (long) batch * BATCH_ROWS + i; + String symbol = symbols[random.nextInt(symbols.length)]; + Long nValue = id % 5 == 0 || symbols[0].equals(symbol) ? null : id * 3; + double xValue = random.nextDouble(); + if (symbol == null) { + s.setNull(i); + } else { + s.setSafe(i, symbol.getBytes(java.nio.charset.StandardCharsets.UTF_8)); + } + ids.set(i, id); + if (nValue == null) { + n.setNull(i); + } else { + n.set(i, nValue); + } + x.set(i, xValue); + rows.add(new Object[]{symbol, id, nValue, xValue}); + } + root.setRowCount(BATCH_ROWS); + try (ArrowArray array = ArrowArray.allocateNew(ALLOCATOR); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(ALLOCATOR)) { + Data.exportVectorSchemaRoot(ALLOCATOR, root, null, array, arrowSchema); + writer.writeBatch(array.memoryAddress(), arrowSchema.memoryAddress()); + } + } + } + } + return rows; + } + private static void writeIdsWithJni(Path file, int batches, int batchRows) throws Exception { Schema schema = new Schema(List.of( Field.notNullable("id", new ArrowType.Int(64, true)), 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 ce9badda..6c82e4ae 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 @@ -5,7 +5,9 @@ import io.github.dfa1.vortex.reader.RowFilter; import io.github.dfa1.vortex.reader.array.Array; +import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.function.LongPredicate; /// The minimal, low-level entry point to the compute primitives of ADR 0013 — the fused, single-pass @@ -114,6 +116,30 @@ 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)]. + /// + /// 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. + /// + /// 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. + /// + /// @param chunk the decoded chunk holding every referenced column + /// @param groupColumn the dictionary-encoded string 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, + RowFilter filter, ColumnName aggColumn) { + Objects.requireNonNull(chunk, "chunk"); + Objects.requireNonNull(groupColumn, "groupColumn"); + return GroupedAggregate.aggregate(chunk, groupColumn, filter, aggColumn); + } + private static boolean evaluate(Chunk chunk, RowFilter filter, long row) { return switch (filter) { case RowFilter.And(var parts) -> { 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 new file mode 100644 index 00000000..15824b53 --- /dev/null +++ b/reader/src/main/java/io/github/dfa1/vortex/reader/compute/GroupedAggregate.java @@ -0,0 +1,225 @@ +package io.github.dfa1.vortex.reader.compute; + +import io.github.dfa1.vortex.core.error.VortexException; +import io.github.dfa1.vortex.core.model.ColumnName; +import io.github.dfa1.vortex.core.model.DType; +import io.github.dfa1.vortex.core.model.PType; +import io.github.dfa1.vortex.reader.Chunk; +import io.github.dfa1.vortex.reader.RowFilter; +import io.github.dfa1.vortex.reader.array.Array; +import io.github.dfa1.vortex.reader.array.BoolArray; +import io.github.dfa1.vortex.reader.array.DoubleArray; +import io.github.dfa1.vortex.reader.array.FloatArray; +import io.github.dfa1.vortex.reader.array.VarBinDictArray; +import io.github.dfa1.vortex.reader.array.VarBinOffsetArray; +import io.github.dfa1.vortex.reader.array.VarBinSlicedArray; + +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.function.LongPredicate; + +/// 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. +/// +/// 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. +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. + private static final int MAX_GROUPS = 1 << 16; + + 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 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, + RowFilter filter, ColumnName aggColumn) { + 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(); + } + } + 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)) { + g = groups; + } else { + long code = dict.code(i + offset); + if (code < 0 || code >= groups) { + throw new VortexException("dict code " + code + " at row " + (i + offset) + + " out of range for " + groups + " dictionary entries"); + } + g = (int) code; + } + acc.add(g, i); + } + Map result = new HashMap<>(); + for (int g = 0; g <= groups; g++) { + if (acc.rows[g] > 0) { + String key = g == groups ? null : values.getString(g); + // two entries may hold the same string; a group is a value, not a code + result.merge(key, acc.result(g), GroupedAggregate::combine); + } + } + 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; + if (a.sum() == null) { + sum = null; + } else if (a.sum() instanceof Double x) { + sum = x + b.sum().doubleValue(); + } else { + sum = a.sum().longValue() + b.sum().longValue(); + } + return new FilteredAggregate(a.selectedRows() + b.selectedRows(), + a.aggNonNullCount() + b.aggNonNullCount(), sum, + extreme(a.min(), b.min(), true), extreme(a.max(), b.max(), false)); + } + + @SuppressWarnings({"unchecked", "rawtypes"}) + private static Object extreme(Object a, Object b, boolean min) { + if (a == null) { + return b; + } + if (b == null) { + return a; + } + int cmp = ((Comparable) a).compareTo(b); + return (min ? cmp <= 0 : cmp >= 0) ? a : b; + } + + /// How the aggregate column's values are read: not at all (a row count), as longs, or as doubles. + private record Lane(Array data, BoolArray validity, boolean unsigned, boolean floating, boolean isFloat) { + + private static final Lane COUNT_ONLY = new Lane(null, null, false, false, false); + + /// Picks the lane for `column`, or `null` when no lane folds it: a non-numeric column, an + /// `f16`, or a `u64` whose values past 2^63 have no signed-long home. + static Lane of(Array column) { + if (column == null) { + return COUNT_ONLY; + } + Array data = NumericColumns.unwrap(column); + BoolArray validity = NumericColumns.validityOf(column); + if (data instanceof DoubleArray || data instanceof FloatArray) { + return new Lane(data, validity, false, true, data instanceof FloatArray); + } + if (NumericColumns.isLongDomain(data) + && data.dtype() instanceof DType.Primitive p && p.ptype() != PType.U64) { + return new Lane(data, validity, p.ptype().isUnsigned(), false, false); + } + return null; + } + } + + /// Per-group running totals, indexed by dictionary code. + private static final class Accumulators { + private final Lane lane; + private final long[] rows; + private final long[] nonNull; + private final long[] longSum; + private final long[] longMin; + private final long[] longMax; + private final double[] doubleSum; + private final double[] doubleMin; + private final double[] doubleMax; + + Accumulators(int size, Lane lane) { + this.lane = lane; + this.rows = new long[size]; + this.nonNull = new long[size]; + boolean longs = lane.data() != null && !lane.floating(); + boolean doubles = lane.data() != null && lane.floating(); + this.longSum = longs ? new long[size] : null; + this.longMin = longs ? new long[size] : null; + this.longMax = longs ? new long[size] : null; + this.doubleSum = doubles ? new double[size] : null; + this.doubleMin = doubles ? new double[size] : null; + this.doubleMax = doubles ? new double[size] : null; + } + + void add(int g, long i) { + rows[g]++; + Array data = lane.data(); + if (data == null || (lane.validity() != null && !lane.validity().getBoolean(i))) { + return; + } + boolean first = nonNull[g]++ == 0; + if (lane.floating()) { + double v = lane.isFloat() ? ((FloatArray) data).getFloat(i) : ((DoubleArray) data).getDouble(i); + doubleSum[g] += v; + // Double.compare orders like Rust (-0.0 < 0.0, NaN the maximum), as the fused kernel does + if (first || Double.compare(v, doubleMin[g]) < 0) { + doubleMin[g] = v; + } + if (first || Double.compare(v, doubleMax[g]) > 0) { + doubleMax[g] = v; + } + } else { + long v = NumericColumns.widenLong(data, lane.unsigned(), i); + longSum[g] += v; + if (first || v < longMin[g]) { + longMin[g] = v; + } + if (first || v > longMax[g]) { + longMax[g] = v; + } + } + } + + FilteredAggregate result(int g) { + if (lane.data() == null) { + return new FilteredAggregate(rows[g], 0L, null, null, null); + } + boolean any = nonNull[g] > 0; + if (lane.floating()) { + return new FilteredAggregate(rows[g], nonNull[g], doubleSum[g], + any ? doubleMin[g] : null, any ? doubleMax[g] : null); + } + return new FilteredAggregate(rows[g], nonNull[g], longSum[g], + any ? longMin[g] : null, any ? longMax[g] : null); + } + } +} 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 fcbbdf7d..ca57c8c8 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 @@ -8,6 +8,7 @@ import io.github.dfa1.vortex.reader.array.Array; import io.github.dfa1.vortex.reader.array.MaskedArray; import io.github.dfa1.vortex.reader.array.VarBinArray; +import io.github.dfa1.vortex.reader.array.VarBinDictArray; import io.github.dfa1.vortex.reader.array.VarBinSlicedArray; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; @@ -17,6 +18,7 @@ import java.lang.foreign.Arena; import java.util.List; +import java.util.Map; import java.util.function.LongPredicate; import java.util.stream.LongStream; import java.util.stream.Stream; @@ -175,4 +177,167 @@ void rejectsNullArguments() { assertThatNullPointerException().isThrownBy(() -> Compute.matcher(chunk, null)); } } + + /// [Compute#groupedAggregate(Chunk, ColumnName, RowFilter, ColumnName)] folds per dictionary + /// code; every expectation below is the plain per-row SQL answer, worked out by hand. + @Nested + class Grouped { + + private static final ColumnName G = ColumnName.of("g"); + private static final ColumnName V = ColumnName.of("v"); + + /// Groups a, b, a, (null), b, a: row 3's group is null although its code points at "a". + private static Array nullableGroups() { + Array dict = ComputeArrays.utf8DictArray(ARENA, new String[]{"a", "b"}, new int[]{0, 1, 0, 0, 1, 0}); + return new MaskedArray(dict, ComputeArrays.boolArray(ARENA, true, true, true, false, true, true)); + } + + private static Chunk chunk(Array group, Array agg) { + Chunk chunk = mock(Chunk.class); + given(chunk.rowCount()).willReturn(group.length()); + given(chunk.column(G)).willReturn(group); + given(chunk.column(V)).willReturn(agg); + return chunk; + } + + @Test + void countOnly_countsRowsPerGroupWithNullAsItsOwnGroup() { + // Given + Chunk chunk = chunk(nullableGroups(), null); + + // When + Map result = Compute.groupedAggregate(chunk, G, null, null).orElseThrow(); + + // Then + assertThat(result).containsOnlyKeys("a", "b", null); + assertThat(result.get("a").selectedRows()).isEqualTo(3); + assertThat(result.get("b").selectedRows()).isEqualTo(2); + assertThat(result.get(null).selectedRows()).isEqualTo(1); + } + + @Test + void longColumn_foldsSumMinMaxSkippingNullValues() { + // Given — v = 5, 7, (null), 2, -1, 9 + Array v = ComputeArrays.maskedLongArray(ARENA, new long[]{5, 7, 0, 2, -1, 9}, + new boolean[]{true, true, false, true, true, true}); + + // When + 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 + assertThat(result.get("a")).isEqualTo(new FilteredAggregate(3, 2, 14L, 5L, 9L)); + assertThat(result.get("b")).isEqualTo(new FilteredAggregate(2, 2, 6L, -1L, 7L)); + assertThat(result.get(null)).isEqualTo(new FilteredAggregate(1, 1, 2L, 2L, 2L)); + } + + @Test + void doubleColumn_ordersMinMaxLikeRust() { + // Given — group a holds -0.0, NaN and 0.0: Rust's order puts -0.0 first and NaN last + Array group = ComputeArrays.utf8DictArray(ARENA, new String[]{"a"}, new int[]{0, 0, 0}); + Array v = ComputeArrays.doubleArray(ARENA, -0.0, Double.NaN, 0.0); + + // When + FilteredAggregate result = Compute.groupedAggregate(chunk(group, v), G, null, V).orElseThrow().get("a"); + + // Then + assertThat(result.min()).isEqualTo(-0.0); + assertThat(result.max()).isEqualTo(Double.NaN); + assertThat(result.aggNonNullCount()).isEqualTo(3); + } + + @Test + void filter_foldsOnlyTheSelectedRows() { + // Given — keep v > 4: rows 0 (a, 5), 1 (b, 7), 5 (a, 9) + Array v = ComputeArrays.longArray(ARENA, 5, 7, 1, 2, -1, 9); + Chunk chunk = chunk(nullableGroups(), v); + + // When + Map result = Compute.groupedAggregate(chunk, G, RowFilter.gt(V, 4L), V) + .orElseThrow(); + + // Then — groups with no selected row do not appear + assertThat(result).containsOnlyKeys("a", "b"); + assertThat(result.get("a")).isEqualTo(new FilteredAggregate(2, 2, 14L, 5L, 9L)); + assertThat(result.get("b")).isEqualTo(new FilteredAggregate(1, 1, 7L, 7L, 7L)); + } + + /// A dictionary spanning many chunks reaches the kernel as a slice; codes are read at its offset. + @Test + void slicedDictionary_readsCodesAtTheSliceOffset() { + // Given — rows 2..4 of a, a, b, b, a + VarBinArray dict = ComputeArrays.utf8DictArray(ARENA, new String[]{"a", "b"}, new int[]{0, 0, 1, 1, 0}); + Chunk chunk = chunk(new VarBinSlicedArray(DType.UTF8, 3, dict, 2), null); + + // When + Map result = Compute.groupedAggregate(chunk, G, null, null).orElseThrow(); + + // Then + assertThat(result.get("b").selectedRows()).isEqualTo(2); + assertThat(result.get("a").selectedRows()).isEqualTo(1); + } + + /// A group is a value, not a code: two entries holding the same string fold into one group. + @Test + void duplicateDictionaryValues_foldIntoOneGroup() { + // Given — entries 0 and 2 are both "x" + Array group = ComputeArrays.utf8DictArray(ARENA, new String[]{"x", "y", "x"}, new int[]{0, 2, 1, 2}); + Array v = ComputeArrays.longArray(ARENA, 1, 10, 100, 1000); + + // When + Map result = Compute.groupedAggregate(chunk(group, v), G, null, V).orElseThrow(); + + // Then + assertThat(result.get("x")).isEqualTo(new FilteredAggregate(3, 3, 1011L, 1L, 1000L)); + } + + @Test + void codeOutsideDictionary_throwsVortexException() { + // Given — 2 entries, row 1 carries code 9 + Array group = ComputeArrays.utf8DictArray(ARENA, new String[]{"a", "b"}, new int[]{0, 9}); + Chunk chunk = chunk(group, null); + + // When / Then + assertThatThrownBy(() -> Compute.groupedAggregate(chunk, G, null, null)) + .isInstanceOf(VortexException.class); + } + + /// The codes buffer is untrusted: one shorter than the rows it claims must fail as malformed + /// input, not as an IndexOutOfBoundsException. + @Test + void codesBufferShorterThanRows_throwsVortexException() { + // Given — a dictionary array claiming 3 rows over a 1-code buffer + VarBinDictArray one = (VarBinDictArray) ComputeArrays.utf8DictArray(ARENA, new String[]{"a"}, new int[]{0}); + Array group = new VarBinDictArray(DType.UTF8, 3, one.bytesSegment(), one.dictValOffsets(), + one.dictValOffPType(), one.dictCodesSegs(), one.dictCodesPType()); + Chunk chunk = chunk(group, null); + + // When / Then + assertThatThrownBy(() -> Compute.groupedAggregate(chunk, G, null, null)) + .isInstanceOf(VortexException.class); + } + + /// The caller groups the rows itself when the kernel cannot: it must say so, not guess. + @Test + void groupColumnNotDictionaryEncoded_isEmpty() { + // Given + Chunk chunk = chunk(ComputeArrays.utf8Array(ARENA, "a", "b"), null); + + // When / Then + assertThat(Compute.groupedAggregate(chunk, G, null, null)).isEmpty(); + } + + /// u64 values past 2^63 have no signed-long home, and a string has no sum. + @Test + void aggregateColumnWithoutALane_isEmpty() { + // Given + Array group = ComputeArrays.utf8DictArray(ARENA, new String[]{"a"}, new int[]{0, 0}); + + // When / Then + assertThat(Compute.groupedAggregate(chunk(group, ComputeArrays.unsignedLongArray(ARENA, 1, 2)), G, null, V)) + .isEmpty(); + assertThat(Compute.groupedAggregate(chunk(group, ComputeArrays.utf8Array(ARENA, "p", "q")), G, null, V)) + .isEmpty(); + } + } }