feat!: Add multi-column support to AggFormula

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

@@ -726,17 +726,14 @@ public void visit(@NotNull final Formula formula) {
                                 cd.getName(),
                                 VectorFactory.forElementType(cd.getDataType()).vectorType(),
                                 cd.getDataType())));
-                table.getColumnSourceMap().forEach((key, value) -> {
-                    final ColumnDefinition<?> columnDef = ColumnDefinition.fromGenericType(
-                            key, VectorFactory.forElementType(value.getType()).vectorType());
-                    vectorColumnDefinitions.put(key, columnDef);
-                });
             }
 
             // Get the input column names from the formula and provide them to the groupBy operator
             final String[] inputColumns =
                     selectColumn.initDef(vectorColumnDefinitions, compilationProcessor).toArray(String[]::new);
-
+            if (selectColumn.hasVirtualRowVariables()) {
+                throw new IllegalArgumentException("AggFormula does not support virtual row variables");
+            }
             final GroupByChunkedOperator groupByChunkedOperator = new GroupByChunkedOperator(table, false, null,
                     Arrays.stream(inputColumns).map(col -> MatchPair.of(Pair.parse(col)))
                             .toArray(MatchPair[]::new));

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

@@ -438,6 +438,7 @@ private void copyData(@NotNull final RowSequence rowSequence) {
         private void copyData(@NotNull final RowSequence rowSequence, @NotNull final boolean[] columnsMask) {
             try (final RowSequence.Iterator rowSequenceIterator = rowSequence.getRowSequenceIterator()) {
                 while (rowSequenceIterator.hasMore()) {
+                    sharedContext.reset();
                     final RowSequence rowSequenceSlice = rowSequenceIterator.getNextRowSequenceThrough(
                             calculateContainingBlockLastKey(rowSequenceIterator.peekNextKey()));
                     for (int ci = 0; ci < columnsToGetMask.length; ++ci) {

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

@@ -245,22 +245,13 @@ public UnaryOperator<ModifiedColumnSet> initializeRefreshing(@NotNull final Quer
         if (selectColumn.getColumns().isEmpty()) {
             return inputToResultModifiedColumnSetFactory = input -> ModifiedColumnSet.EMPTY;
         }
-
-        final ModifiedColumnSet updateMCS = new ModifiedColumnSet(resultTable.getModifiedColumnSetForUpdates());
         final ModifiedColumnSet resultMCS = resultTable.newModifiedColumnSet(selectColumn.getName());
         final String[] inputColumnNames = selectColumn.getColumns().toArray(String[]::new);
-        final ModifiedColumnSet[] modifiedColumnSets = selectColumn.getColumns().stream()
-                .map(ignored -> resultMCS).toArray(ModifiedColumnSet[]::new);
-        final ModifiedColumnSet.Transformer transformer = inputTable.newModifiedColumnSetTransformer(
-                inputColumnNames,
-                modifiedColumnSets);
-
+        final ModifiedColumnSet inputMCS = inputTable.newModifiedColumnSet(inputColumnNames);
         return inputToResultModifiedColumnSetFactory = input -> {
-            if (groupBy.getSomeKeyHasAddsOrRemoves()) {
+            if (groupBy.getSomeKeyHasAddsOrRemoves() ||
+                    (groupBy.getSomeKeyHasModifies() && input.containsAny(inputMCS))) {
                 return resultMCS;
-            } else if (groupBy.getSomeKeyHasModifies()) {
-                transformer.clearAndTransform(input, updateMCS);
-                return updateMCS;
             }
             return ModifiedColumnSet.EMPTY;
         };
@@ -383,6 +374,7 @@ private DataCopyContext() {
         private void copyData(@NotNull final RowSequence rowSequence) {
             try (final RowSequence.Iterator rowSequenceIterator = rowSequence.getRowSequenceIterator()) {
                 while (rowSequenceIterator.hasMore()) {
+                    sharedContext.reset();
                     final RowSequence rowSequenceSlice = rowSequenceIterator.getNextRowSequenceThrough(
                             calculateContainingBlockLastKey(rowSequenceIterator.peekNextKey()));
                     resultColumn.fillFromChunk(fillFromContext,

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

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

table-api/src/main/java/io/deephaven/api/agg/Aggregation.java

@@ -310,11 +310,11 @@ static Formula AggFormula(String formulaString) {
     /**
      * <p>
      * Create a {@link Formula formula} aggregation with the supplied {@code columnName} and {@code expression}. This
-     * variant requires the formula to provide the output column name and specific input column names in the following
+     * variant requires the formula to provide the output column name and the expression to evaluate in the following
      * format:
      * </p>
      * {@code
-     * AggFormula("output_col=(input_col1 + input_col2) * input_col3")
+     * AggFormula("output_col", "(input_col1 + input_col2) * input_col3")
      * }
      *
      * @param columnName The output column name