feat: ingest support for numeric typed ExpressionLambdaAggregatorFactory (#19508)

Author: clintropolis

docs/querying/aggregations.md

@@ -471,7 +471,11 @@ For these reasons, we have deprecated this aggregator and recommend using the Da
 
 ### Expression aggregator
 
-Aggregator applicable only at query time. Aggregates results using [Druid expressions](./math-expr.md) functions to facilitate building custom functions.
+Aggregates results using [Druid expressions](./math-expr.md) functions to facilitate building custom functions.
+
+The expression aggregator can be used at query time with any intermediate type. It can also be used at ingest time, but
+only when the type of `initialValue` is a primitive numeric type (`LONG` or `DOUBLE`) and matches the type of
+`initialCombineValue`. Other intermediate types, such as strings, arrays, and complex types, are query-time only.
 
 | Property | Description | Required |
 | --- | --- | --- |

embedded-tests/src/test/java/org/apache/druid/testing/embedded/compact/CompactionTaskTest.java

@@ -33,9 +33,12 @@
 import org.apache.druid.java.util.common.granularity.Granularity;
 import org.apache.druid.java.util.common.jackson.JacksonUtils;
 import org.apache.druid.query.Druids;
+import org.apache.druid.query.aggregation.CountAggregatorFactory;
+import org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory;
 import org.apache.druid.query.aggregation.datasketches.hll.HllSketchModule;
 import org.apache.druid.query.aggregation.datasketches.quantiles.DoublesSketchModule;
 import org.apache.druid.query.aggregation.datasketches.theta.SketchModule;
+import org.apache.druid.query.expression.TestExprMacroTable;
 import org.apache.druid.query.metadata.metadata.SegmentMetadataQuery;
 import org.apache.druid.segment.TestHelper;
 import org.apache.druid.testing.embedded.EmbeddedClusterApis;
@@ -55,6 +58,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
@@ -107,6 +111,65 @@ public class CompactionTaskTest extends CompactionTestBase
           "namespace", "continent", "country", "region", "city", "timestamp"
       );
 
+  /**
+   * Index task identical in shape to {@link MoreResources.Task#INDEX_TASK_WITH_AGGREGATORS} but with a pair of
+   * {@link ExpressionLambdaAggregatorFactory} metrics over the {@code added} long field. Used by
+   * {@link #testCompactionWithExpressionLambdaAggregator} to verify that an expression aggregator works correctly.
+   */
+  private static final Supplier<TaskBuilder.Index> INDEX_TASK_WITH_EXPR_AGG = () ->
+      TaskBuilder
+          .ofTypeIndex()
+          .jsonInputFormat()
+          .localInputSourceWithFiles(
+              Resources.DataFile.tinyWiki1Json(),
+              Resources.DataFile.tinyWiki2Json(),
+              Resources.DataFile.tinyWiki3Json()
+          )
+          .timestampColumn("timestamp")
+          .dimensions(
+              "page",
+              "language", "tags", "user", "unpatrolled", "newPage", "robot",
+              "anonymous", "namespace", "continent", "country", "region", "city"
+          )
+          .metricAggregates(
+              new CountAggregatorFactory("ingested_events"),
+              new ExpressionLambdaAggregatorFactory(
+                  "added_sum_expr",
+                  Set.of("added"),
+                  null,
+                  "0",
+                  null,
+                  null,
+                  false,
+                  false,
+                  "__acc + added",
+                  null,
+                  null,
+                  null,
+                  null,
+                  TestExprMacroTable.INSTANCE
+              ),
+              new ExpressionLambdaAggregatorFactory(
+                  "added_or_expr",
+                  Set.of("added"),
+                  null,
+                  "0",
+                  null,
+                  null,
+                  false,
+                  false,
+                  "bitwiseOr(\"__acc\", \"added\")",
+                  null,
+                  null,
+                  null,
+                  null,
+                  TestExprMacroTable.INSTANCE
+              )
+          )
+          .dynamicPartitionWithMaxRows(3)
+          .granularitySpec("DAY", "SECOND", true)
+          .appendToExisting(false);
+
   private String fullDatasourceName;
 
   @BeforeEach
@@ -259,6 +322,33 @@ public void testCompactionWithTimestampDimension() throws Exception
     loadDataAndCompact(INDEX_TASK_WITH_TIMESTAMP.get(), COMPACTION_TASK.get(), null);
   }
 
+  @Test
+  public void testCompactionWithExpressionLambdaAggregator() throws Exception
+  {
+    try (final Closeable ignored = unloader(fullDatasourceName)) {
+      runTask(INDEX_TASK_WITH_EXPR_AGG.get());
+      verifySegmentsCount(4);
+
+      // Snapshot metric values prior to compaction.
+      final String preCompact = cluster.runSql(
+          "SELECT SUM(added_sum_expr), SUM(added_or_expr) FROM %s",
+          fullDatasourceName
+      );
+
+      // Compact 4 segments -> 2; this performs cross-segment rollup which drives RowCombiningTimeAndDimsIterator
+      // into ExpressionLambdaAggregatorFactory.makeAggregateCombiner().
+      compactData(COMPACTION_TASK.get(), null, null);
+      verifySegmentsCount(2);
+
+      // Metric values must round-trip through compaction unchanged.
+      final String postCompact = cluster.runSql(
+          "SELECT SUM(added_sum_expr), SUM(added_or_expr) FROM %s",
+          fullDatasourceName
+      );
+      Assertions.assertEquals(preCompact, postCompact);
+    }
+  }
+
   private void loadDataAndCompact(
       TaskBuilder.Index indexTask,
       TaskBuilder.Compact compactionResource,

processing/src/main/java/org/apache/druid/query/aggregation/ExpressionLambdaAggregatorFactory.java

@@ -40,9 +40,11 @@
 import org.apache.druid.query.cache.CacheKeyBuilder;
 import org.apache.druid.segment.ColumnInspector;
 import org.apache.druid.segment.ColumnSelectorFactory;
+import org.apache.druid.segment.ColumnValueSelector;
 import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ColumnCapabilitiesImpl;
 import org.apache.druid.segment.column.ColumnType;
+import org.apache.druid.segment.column.ValueType;
 import org.apache.druid.segment.virtual.ExpressionPlan;
 import org.apache.druid.segment.virtual.ExpressionPlanner;
 import org.apache.druid.segment.virtual.ExpressionSelectors;
@@ -347,6 +349,86 @@ public Object combine(@Nullable Object lhs, @Nullable Object rhs)
     ).value();
   }
 
+  @Override
+  public AggregateCombiner makeAggregateCombiner()
+  {
+    final ColumnType intermediateType = getIntermediateType();
+    // The combiner delegates to combine(), which feeds inputs into combineExpression typed against initialCombineValue.
+    // If the fold-side intermediate type (what's stored in the segment column) differs from the combine-side type,
+    // the primitive selector would silently feed wrong-typed values into the expression. Fall through to UOE.
+    if (!intermediateType.equals(ExpressionType.toColumnType(initialCombineValue.get().type()))) {
+      return super.makeAggregateCombiner();
+    }
+    if (intermediateType.is(ValueType.LONG)) {
+      return new LongAggregateCombiner()
+      {
+        private long state;
+        private boolean isNull;
+
+        @Override
+        public void reset(ColumnValueSelector selector)
+        {
+          state = selector.getLong();
+          isNull = selector.isNull();
+        }
+
+        @Override
+        public void fold(ColumnValueSelector selector)
+        {
+          final Object combined = combine(isNull ? null : state, selector.getObject());
+          isNull = combined == null;
+          state = combined == null ? 0L : ((Number) combined).longValue();
+        }
+
+        @Override
+        public long getLong()
+        {
+          return state;
+        }
+
+        @Override
+        public boolean isNull()
+        {
+          return isNull;
+        }
+      };
+    } else if (intermediateType.is(ValueType.DOUBLE)) {
+      return new DoubleAggregateCombiner()
+      {
+        private double state;
+        private boolean isNull;
+
+        @Override
+        public void reset(ColumnValueSelector selector)
+        {
+          state = selector.getDouble();
+          isNull = selector.isNull();
+        }
+
+        @Override
+        public void fold(ColumnValueSelector selector)
+        {
+          final Object combined = combine(isNull ? null : state, selector.getObject());
+          isNull = combined == null;
+          state = combined == null ? 0.0 : ((Number) combined).doubleValue();
+        }
+
+        @Override
+        public double getDouble()
+        {
+          return state;
+        }
+
+        @Override
+        public boolean isNull()
+        {
+          return isNull;
+        }
+      };
+    }
+    return super.makeAggregateCombiner();
+  }
+
   @Override
   public Object deserialize(Object object)
   {