feat!: Add multi-column support to AggFormula

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

@@ -93,13 +93,15 @@
 import io.deephaven.engine.table.impl.by.ssmcountdistinct.unique.ShortRollupUniqueOperator;
 import io.deephaven.engine.table.impl.by.ssmminmax.SsmChunkedMinMaxOperator;
 import io.deephaven.engine.table.impl.by.ssmpercentile.SsmChunkedPercentileOperator;
+import io.deephaven.engine.table.impl.select.FormulaColumn;
 import io.deephaven.engine.table.impl.sources.ReinterpretUtils;
 import io.deephaven.engine.table.impl.ssms.SegmentedSortedMultiSet;
 import io.deephaven.engine.table.impl.util.freezeby.FreezeByCountOperator;
 import io.deephaven.engine.table.impl.util.freezeby.FreezeByOperator;
 import io.deephaven.time.DateTimeUtils;
 import io.deephaven.util.annotations.FinalDefault;
 import io.deephaven.util.type.ArrayTypeUtils;
+import io.deephaven.vector.VectorFactory;
 import org.apache.commons.lang3.mutable.MutableBoolean;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -111,6 +113,7 @@
 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;
@@ -153,6 +156,13 @@ private enum Type {
     private final Collection<? extends Aggregation> aggregations;
     private final Type type;
 
+    /**
+     * For AggFormula we need a representation of the table definition with the column data types converted to
+     * {@link io.deephaven.vector.Vector vectors}. This can be computed once and re-used across all formula aggregations
+     * but does not need to be computed in their absence.
+     */
+    private Map<String, ColumnDefinition<?>> vectorColumnNameMap;
+
     /**
      * Convert a collection of {@link Aggregation aggregations} to an {@link AggregationContextFactory}.
      *
@@ -386,6 +396,10 @@ public final void visit(@NotNull final Aggregations aggregations) {
 
         @Override
         public final void visit(@NotNull final ColumnAggregation columnAgg) {
+            if (columnAgg.pair() == null) {
+                columnAgg.spec().walk(this);
+                return;
+            }
             resultPairs = List.of(columnAgg.pair());
             columnAgg.spec().walk(this);
             resultPairs = List.of();
@@ -707,6 +721,40 @@ public void visit(@NotNull final Partition partition) {
                     groupByColumnNames));
         }
 
+        @Override
+        public void visit(@NotNull final Formula formula) {
+            final String resultColumnName = formula.column().name();
+
+            final FormulaColumn formulaColumn =
+                    FormulaColumn.createFormulaColumn(resultColumnName, formula.formula());
+
+            // 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);
+                });
+            }
+
+            // Get the input column names from the formula and provide them to the groupBy operator
+            final String[] inputColumns =
+                    formulaColumn.initDef(vectorColumnNameMap, compilationProcessor).toArray(String[]::new);
+
+            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(
+                    groupByChunkedOperator,
+                    true,
+                    formulaColumn,
+                    resultColumnName,
+                    compilationProcessor);
+            addNoInputOperator(op);
+        }
+
         // -------------------------------------------------------------------------------------------------------------
         // AggSpec.Visitor
         // -------------------------------------------------------------------------------------------------------------
@@ -747,6 +795,7 @@ public void visit(@NotNull final AggSpecFormula formula) {
             unsupportedForBlinkTables("Formula");
             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));
@@ -860,6 +909,12 @@ default void visit(@NotNull final LastRowKey lastRowKey) {
             rollupUnsupported("LastRowKey");
         }
 
+        @Override
+        @FinalDefault
+        default void visit(@NotNull final Formula formula) {
+            rollupUnsupported("Formula");
+        }
+
         // -------------------------------------------------------------------------------------------------------------
         // AggSpec.Visitor for unsupported column aggregation specs
         // -------------------------------------------------------------------------------------------------------------

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

@@ -25,6 +25,7 @@
 import io.deephaven.engine.table.WritableColumnSource;
 import io.deephaven.util.SafeCloseable;
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import java.util.Arrays;
 import java.util.Collections;
@@ -74,7 +75,7 @@ class FormulaChunkedOperator implements IterativeChunkedAggregationOperator {
             @NotNull final GroupByChunkedOperator groupBy,
             final boolean delegateToBy,
             @NotNull final String formula,
-            @NotNull final String columnParamName,
+            @Nullable final String columnParamName,
             @NotNull final QueryCompilerRequestProcessor compilationProcessor,
             @NotNull final MatchPair... resultColumnPairs) {
         this.groupBy = groupBy;