feat!: Add multi-column support to AggFormula

engine/table/src/main/java/io/deephaven/engine/table/impl/by/AggregationProcessor.java

@@ -111,7 +111,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.function.BiFunction;
@@ -159,7 +158,7 @@ private enum Type {
      * types converted to {@link io.deephaven.vector.Vector vectors}. This can be computed once and re-used across all
      * formula aggregations.
      */
-    private Map<String, ColumnDefinition<?>> vectorColumnNameMap;
+    private Map<String, ColumnDefinition<?>> vectorColumnDefinitions;
 
     /**
      * Convert a collection of {@link Aggregation aggregations} to an {@link AggregationContextFactory}.
@@ -717,28 +716,34 @@ public void visit(@NotNull final Partition partition) {
 
         @Override
         public void visit(@NotNull final Formula formula) {
-            final SelectColumn selectColumn =
-                    SelectColumn.of(Selectable.of(formula.column(), RawString.of(formula.formula())));
+            final SelectColumn selectColumn = SelectColumn.of(formula.selectable());
 
             // Get or create a column definition map composed of vectors of the original column types.
-            if (vectorColumnNameMap == null) {
-                vectorColumnNameMap = new HashMap<>();
-                table.getColumnSourceMap().forEach((key, value) -> {
-                    final ColumnDefinition<?> columnDef = ColumnDefinition.fromGenericType(
-                            key, VectorFactory.forElementType(value.getType()).vectorType());
-                    vectorColumnNameMap.put(key, columnDef);
-                });
+            if (vectorColumnDefinitions == null) {
+                vectorColumnDefinitions = table.getDefinition().getColumnStream().collect(Collectors.toMap(
+                        ColumnDefinition::getName,
+                        (final ColumnDefinition<?> cd) -> ColumnDefinition.fromGenericType(
+                                cd.getName(),
+                                VectorFactory.forElementType(cd.getDataType()).vectorType(),
+                                cd.getDataType())));
             }
 
             // Get the input column names from the formula and provide them to the groupBy operator
             final String[] inputColumns =
-                    selectColumn.initDef(vectorColumnNameMap, compilationProcessor).toArray(String[]::new);
-
+                    selectColumn.initDef(vectorColumnDefinitions, compilationProcessor).toArray(String[]::new);
+            if (!selectColumn.getColumnArrays().isEmpty()) {
+                throw new IllegalArgumentException("AggFormula does not support column arrays ("
+                        + selectColumn.getColumnArrays() + ")");
+            }
+            if (selectColumn.hasVirtualRowVariables()) {
+                throw new IllegalArgumentException("AggFormula does not support virtual row variables");
+            }
+            // TODO: re-use shared groupBy operators (https://github.com/deephaven/deephaven-core/issues/6363)
             final GroupByChunkedOperator groupByChunkedOperator = new GroupByChunkedOperator(table, false, null,
                     Arrays.stream(inputColumns).map(col -> MatchPair.of(Pair.parse(col)))
                             .toArray(MatchPair[]::new));
 
-            final FormulaMultiColumnChunkedOperator op = new FormulaMultiColumnChunkedOperator(
+            final FormulaMultiColumnChunkedOperator op = new FormulaMultiColumnChunkedOperator(table,
                     groupByChunkedOperator, true, selectColumn);
             addNoInputOperator(op);
         }
@@ -781,9 +786,9 @@ public void visit(@NotNull final AggSpecFirst first) {
         @Override
         public void visit(@NotNull final AggSpecFormula formula) {
             unsupportedForBlinkTables("Formula");
+            // TODO: re-use shared groupBy operators (https://github.com/deephaven/deephaven-core/issues/6363)
             final GroupByChunkedOperator groupByChunkedOperator = new GroupByChunkedOperator(table, false, null,
                     resultPairs.stream().map(pair -> MatchPair.of((Pair) pair.input())).toArray(MatchPair[]::new));
-
             final FormulaChunkedOperator formulaChunkedOperator = new FormulaChunkedOperator(groupByChunkedOperator,
                     true, formula.formula(), formula.paramToken(), compilationProcessor,
                     MatchPair.fromPairs(resultPairs));

engine/table/src/main/java/io/deephaven/engine/table/impl/by/FormulaChunkedOperator.java

@@ -291,9 +291,8 @@ public UnaryOperator<ModifiedColumnSet> initializeRefreshing(@NotNull final Quer
             resultColumnModifiedColumnSets[ci] = resultTable.newModifiedColumnSet(resultColumnNames[ci]);
         }
         if (delegateToBy) {
-            // We cannot use the groupBy's result MCS factory, because the result column names are not guaranteed to be
-            // the
-            // same.
+            // We cannot use the groupBy's result MCS factory, because the result column names are not guaranteed
+            // to be the same.
             groupBy.initializeRefreshing(resultTable, aggregationUpdateListener);
         }
         // Note that we also use the factory in propagateUpdates to identify the set of modified columns to handle.
@@ -377,7 +376,7 @@ public SingletonContext makeSingletonContext(final int size) {
     private class DataFillerContext implements SafeCloseable {
 
         private final boolean[] columnsToFillMask;
-        protected final FillFromContext[] fillFromContexts;
+        final FillFromContext[] fillFromContexts;
 
         private DataFillerContext(final boolean @NotNull [] columnsToFillMask) {
             this.columnsToFillMask = columnsToFillMask;
@@ -448,6 +447,7 @@ private void copyData(@NotNull final RowSequence rowSequence, @NotNull final boo
                                     rowSequenceSlice);
                         }
                     }
+                    sharedContext.reset();
                 }
             }
         }

engine/table/src/main/java/io/deephaven/engine/table/impl/by/FormulaMultiColumnChunkedOperator.java

@@ -30,9 +30,9 @@
  */
 class FormulaMultiColumnChunkedOperator implements IterativeChunkedAggregationOperator {
 
+    private final QueryTable inputTable;
+
     private final GroupByChunkedOperator groupBy;
-    // TODO: `delegateToBy` is always true, for this class and FormulaChunkedOperator. Can we remove this and all
-    // the checks?
     private final boolean delegateToBy;
     private final SelectColumn selectColumn;
     private final WritableColumnSource<?> resultColumn;
@@ -50,18 +50,20 @@ class FormulaMultiColumnChunkedOperator implements IterativeChunkedAggregationOp
     private ModifiedColumnSet updateUpstreamModifiedColumnSet;
 
     /**
-     * Construct an operator for applying a formula to a grouped table.
+     * Construct an operator for applying a formula to slot-vectors over an aggregated table..
      *
      * @param groupBy The {@link GroupByChunkedOperator} to use for tracking indices
      * @param delegateToBy Whether this operator is responsible for passing methods through to {@code groupBy}. Should
-     *        be false if {@code groupBy} is updated by the helper (and {@code groupBy} must come before this operator
-     *        if so), or if this is not the first operator sharing {@code groupBy}.
+     *        be false if {@code groupBy} is updated by the helper, or if this is not the first operator sharing
+     *        {@code groupBy}.
      * @param selectColumn The formula column that will produce the results
      */
     FormulaMultiColumnChunkedOperator(
+            @NotNull final QueryTable inputTable,
             @NotNull final GroupByChunkedOperator groupBy,
             final boolean delegateToBy,
             @NotNull final SelectColumn selectColumn) {
+        this.inputTable = inputTable;
         this.groupBy = groupBy;
         this.delegateToBy = delegateToBy;
         this.selectColumn = selectColumn;
@@ -238,11 +240,21 @@ public UnaryOperator<ModifiedColumnSet> initializeRefreshing(@NotNull final Quer
             // guaranteed to be the same.
             groupBy.initializeRefreshing(resultTable, aggregationUpdateListener);
         }
-        // Note that we also use the factory in propagateUpdates to identify the set of modified columns to handle.
-        final ModifiedColumnSet resultModifiedColumnSet = resultTable
-                .newModifiedColumnSet(getResultColumns().keySet().toArray(String[]::new));
 
-        return inputToResultModifiedColumnSetFactory = upstreamModifiedColumnSet -> resultModifiedColumnSet;
+        // Note that we also use the factory in propagateUpdates to identify the set of modified columns to handle.
+        if (selectColumn.getColumns().isEmpty()) {
+            return inputToResultModifiedColumnSetFactory = input -> ModifiedColumnSet.EMPTY;
+        }
+        final ModifiedColumnSet resultMCS = resultTable.newModifiedColumnSet(selectColumn.getName());
+        final String[] inputColumnNames = selectColumn.getColumns().toArray(String[]::new);
+        final ModifiedColumnSet inputMCS = inputTable.newModifiedColumnSet(inputColumnNames);
+        return inputToResultModifiedColumnSetFactory = input -> {
+            if (groupBy.getSomeKeyHasAddsOrRemoves() ||
+                    (groupBy.getSomeKeyHasModifies() && input.containsAny(inputMCS))) {
+                return resultMCS;
+            }
+            return ModifiedColumnSet.EMPTY;
+        };
     }
 
     @Override
@@ -288,14 +300,10 @@ public void propagateUpdates(@NotNull final TableUpdate downstream,
             if (removesToProcess && !resultColumn.getType().isPrimitive()) {
                 dataCopyContext.clearObjectColumnData(downstream.removed());
             }
-            if (modifiesToProcess && addsToProcess) {
-                // Union the rowsets and do a single pass.
-                try (final RowSequence combinedRowSequence = downstream.modified().union(downstream.added())) {
-                    dataCopyContext.copyData(combinedRowSequence);
-                }
-            } else if (modifiesToProcess) {
+            if (modifiesToProcess) {
                 dataCopyContext.copyData(downstream.modified());
-            } else {
+            }
+            if (addsToProcess) {
                 dataCopyContext.copyData(downstream.added());
             }
         }
@@ -324,7 +332,7 @@ public SingletonContext makeSingletonContext(final int size) {
      */
     private class DataFillerContext implements SafeCloseable {
 
-        protected final FillFromContext fillFromContext;
+        final FillFromContext fillFromContext;
 
         private DataFillerContext() {
             fillFromContext = resultColumn.makeFillFromContext(BLOCK_SIZE);
@@ -370,6 +378,7 @@ private void copyData(@NotNull final RowSequence rowSequence) {
                             calculateContainingBlockLastKey(rowSequenceIterator.peekNextKey()));
                     resultColumn.fillFromChunk(fillFromContext,
                             formulaDataSource.getChunk(getContext, rowSequenceSlice), rowSequenceSlice);
+                    sharedContext.reset();
                 }
             }
         }

engine/table/src/main/java/io/deephaven/engine/table/impl/by/GroupByChunkedOperator.java

@@ -614,4 +614,12 @@ public boolean requiresRowKeys() {
     public boolean unchunkedRowSet() {
         return true;
     }
+
+    boolean getSomeKeyHasAddsOrRemoves() {
+        return someKeyHasAddsOrRemoves;
+    }
+
+    boolean getSomeKeyHasModifies() {
+        return someKeyHasModifies;
+    }
 }

engine/table/src/main/java/io/deephaven/engine/table/impl/rangejoin/SupportedRangeJoinAggregations.java

@@ -40,11 +40,11 @@ public static void validate(
         }
     }
 
-    private boolean lastSupported;
+    private boolean hasUnsupportedAggs;
 
     private boolean isSupported(@NotNull final Aggregation aggregation) {
         aggregation.walk(this);
-        return lastSupported;
+        return !hasUnsupportedAggs;
     }
 
     @Override
@@ -54,36 +54,36 @@ public void visit(@NotNull final Aggregations aggregations) {
 
     @Override
     public void visit(@NotNull final ColumnAggregation columnAgg) {
-        lastSupported = columnAgg.spec() instanceof AggSpecGroup;
+        hasUnsupportedAggs |= !(columnAgg.spec() instanceof AggSpecGroup);
     }
 
     @Override
     public void visit(@NotNull final ColumnAggregations columnAggs) {
-        lastSupported = columnAggs.spec() instanceof AggSpecGroup;
+        hasUnsupportedAggs |= !(columnAggs.spec() instanceof AggSpecGroup);
     }
 
     @Override
     public void visit(@NotNull final Count count) {
-        lastSupported = false;
+        hasUnsupportedAggs = true;
     }
 
     @Override
     public void visit(@NotNull final FirstRowKey firstRowKey) {
-        lastSupported = false;
+        hasUnsupportedAggs = true;
     }
 
     @Override
     public void visit(@NotNull final LastRowKey lastRowKey) {
-        lastSupported = false;
+        hasUnsupportedAggs = true;
     }
 
     @Override
     public void visit(@NotNull final Partition partition) {
-        lastSupported = false;
+        hasUnsupportedAggs = true;
     }
 
     @Override
     public void visit(@NotNull final Formula formula) {
-        lastSupported = false;
+        hasUnsupportedAggs = true;
     }
 }