Add schema-bound polymorphic aggregation across query engines

[xiangfu0]

PR flow

Derives aggregation argument and result types from input schema, binds immutable AggregateCallBinding to functions, and serializes it via Thrift for consistent polymorphic aggregation across query engines.

flowchart TD
  N0["Obtain table schema #40;F1#41;"]:::stAdded
  N1["Pass schema to expression override #40;F1#41;"]:::stAdded
  N2["Bind aggregation calls with schema #40;F1#41;"]:::stAdded
  N3["Resolve expression types using schema #40;F12#41;"]:::stAdded
  N4["Infer result type via SqlReturnTypeInference #40;F3#41;"]:::stAdded
  N5["Create immutable AggregateCallBinding #40;F3#41;"]:::stAdded
  N6["Attach binding to Function for serialization #40;F12#44; F6#41;"]:::stAdded
  N0 -->|"calls with schema"| N1
  N1 -->|"calls AggregationFunctionBinder#46;bind"| N2
  N2 -->|"calls bind with ExpressionTypeResolver"| N3
  N3 -->|"calls inferReturnType for aggregation"| N4
  N4 -->|"returns AggregateCallBinding"| N5
  N5 -->|"sets binding on Function and serializes"| N6
  classDef stAdded fill:#dafbe1,stroke:#1a7f37,color:#1f2328,stroke-width:2px
  classDef stModified fill:#fff8c5,stroke:#9a6700,color:#1f2328,stroke-width:2px
  classDef stRemoved fill:#ffebe9,stroke:#cf222e,color:#1f2328,stroke-width:2px
  classDef stUnchanged fill:#f6f8fa,stroke:#656d76,color:#1f2328,stroke-width:1px

Loading

AI-generated · Green: added · Yellow: modified · Red: removed · Gray: existing

Partial evidence: 52 file patches omitted; 0 truncated.

Diff evidence

• F1: pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java — before · after
• F3: pinot-common/src/main/java/org/apache/pinot/common/function/AggregationFunctionTypeResolver.java — after
• F6: pinot-common/src/main/java/org/apache/pinot/common/request/context/AggregateCallBinding.java — after
• F12: pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/AggregationFunctionBinder.java — after

• Regenerate PR flow

Polymorphic aggregates need a result schema before any rows arrive. Today, explicit SQL type arguments or data-block inference make that contract difficult to preserve through empty results, distributed merges, and logical BOOLEAN/TIMESTAMP values.

This PR derives logical argument and result types from the input schema, carries an immutable AggregateCallBinding through SSE requests and MSE plans, and constructs the aggregate with that binding. SQL names and operands remain unchanged.

SQL behavior

For an input schema with name STRING, enabled BOOLEAN, and eventTime TIMESTAMP:

Expression	Result type
MODE(name)	STRING
MODE(eventTime)	TIMESTAMP
FIRST_WITH_TIME(name, eventTime)	STRING
LAST_WITH_TIME(enabled, eventTime)	BOOLEAN
ANY_VALUE(enabled)	BOOLEAN
ARRAY_AGG(eventTime)	TIMESTAMP_ARRAY
ARRAY_AGG(name, true)	STRING_ARRAY, distinct values

Numeric MODE keeps its DOUBLE result and existing tie reducers; nonnumeric MODE supports MIN/MAX ties. Existing explicit-type FIRST/LAST and ARRAY_AGG overloads remain supported. The existing SSE EXPR_MIN/EXPR_MAX rewrite also preserves logical measuring/projection types, null projections, and all tied rows across serialized merges. Public EXPR_MIN/EXPR_MAX execution remains SSE-only.

Implementation and review order

The same binding must survive request serialization, stage splitting, runtime construction, and broker reduction: a FINAL-stage OBJECT accumulator is not the final SQL type. The cross-module changes also cover expression overrides, gapfill, post-aggregation, direct server SQL, pruned responses, and LIMIT 0.

Review area	Starting points
Logical type contract and inference	AggregateCallBinding, AggregationFunctionType
SSE schema binding and expression resolution	AggregationFunctionBinder, ExpressionTypeResolver
MSE stage and serialization contract	PinotRuleUtils, BoundAggregationFunction
Aggregate implementations and extension contract	Provider registry, extension guide

Providers are discovered once through ServiceLoader. Another aggregate still needs a declared AggregationFunctionType entry, its validation/type rule, and a provider with the required kernel/serialization support; it does not need new factory, planner-stage, or reducer switch branches. This does not introduce arbitrary SQL names or replace UDAF registration.

The diff includes 23 test/fixture files and two generated Thrift files (AggregationFunctionBinding.java and Function.java). Review their schema source in query.thrift; MSE metadata is declared in expressions.proto.

Compatibility and performance

• Binding metadata is optional in Thrift and protobuf. Existing explicit overloads and fixed numeric aggregate contracts retain their behavior.
• New inferred overloads and logical-type semantics require upgraded brokers and execution workers. Existing ANY_VALUE and SSE EXPR_MIN/EXPR_MAX calls can now expose logical BOOLEAN/TIMESTAMP result types. Do not rely on these semantics during a mixed-version rollout; optional metadata alone does not give older kernels these semantics.
• Existing aggregates with native transforms that lack schema-only metadata retain their legacy unbound path. New inferred FIRST/LAST and ARRAY_AGG overloads require resolvable metadata. Older physical segment types are read through conversion getters for the bound type.
• Inference is query setup/planning work. It adds expression traversal, metadata objects, and serialization; fixed aggregates can also incur setup traversal. No end-to-end latency or allocation benchmark has been run, so this PR makes no zero-overhead or latency-regression claim.

Validation

• 557 focused tests across 34 classes passed, including 13 runtime fixture queries in all three MSE execution modes (39 executions), both planners, raw/dictionary input, exact LONG/TIMESTAMP values, distributed merges, null/empty cases, and broker-only results. The final annotation-only edit also passed its class's 8 tests.
• Spotless, Checkstyle, and license checks passed for the affected modules. Local warning-enabled compilation used JDK 25 with an existing JetBrains annotation jar added only to the local compiler classpath for unchanged compression code; no project dependency changed.
• On code commit 7087c4f085, unit, integration, quickstart, Java 11 client, linter, and SSE/MSE query compatibility CI passed. Current head 51917aff8a is an empty CI-refresh commit with the identical source tree; the other checks are rerunning.
• Binary compatibility CI remains blocked by an upstream baseline mismatch. The current failed job built merge 6245210119 on master 810d7985b7 but compared it with the older f727b44f6e snapshot. It flags the TableConfig constructor changed by #16033. TableConfig's source blob is identical in the tested merge and its master parent, and this PR does not modify pinot-spi. A fresh CI event retained the same old baseline. This needs a CI baseline correction before merge; maintainer approval is also pending.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 80.97470% with 203 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.00%. Comparing base (7f79c65) to head (f6d20bc).

Files with missing lines	Patch %	Lines
...core/query/aggregation/ExpressionTypeResolver.java	69.59%	19 Missing and 26 partials ⚠️
...che/pinot/segment/spi/AggregationFunctionType.java	63.01%	5 Missing and 22 partials ⚠️
...mmon/function/AggregationFunctionTypeResolver.java	76.27%	8 Missing and 6 partials ⚠️
.../aggregation/GapfillAggregationFunctionBinder.java	87.34%	3 Missing and 7 partials ⚠️
...sthandler/BaseSingleStageBrokerRequestHandler.java	57.14%	7 Missing and 2 partials ⚠️
.../aggregation/function/ModeAggregationFunction.java	92.80%	2 Missing and 7 partials ⚠️
...ation/function/array/ArrayAggFunctionProvider.java	82.97%	2 Missing and 6 partials ⚠️
...apache/pinot/calcite/rel/rules/PinotRuleUtils.java	68.18%	4 Missing and 3 partials ⚠️
...inot/query/planner/logical/RexExpressionUtils.java	33.33%	6 Missing ⚠️
...t/common/request/context/AggregateCallBinding.java	72.22%	2 Missing and 3 partials ⚠️
... and 28 more

Additional details and impacted files

@@             Coverage Diff              @@
##             master   #19523      +/-   ##
============================================
+ Coverage     67.81%   68.00%   +0.19%     
  Complexity     1450     1450              
============================================
  Files          3506     3515       +9     
  Lines        227028   227861     +833     
  Branches      35881    36113     +232     
============================================
+ Hits         153969   154967     +998     
+ Misses        60903    60665     -238     
- Partials      12156    12229      +73     

Flag	Coverage Δ
integration	100.00% <ø> (ø)
integration1	100.00% <ø> (ø)
integration2	0.00% <ø> (ø)
java-25	68.00% <80.97%> (+0.19%)	⬆️
lane-a	100.00% <ø> (ø)
lane-b	0.00% <ø> (ø)
temurin	68.00% <80.97%> (+0.19%)	⬆️
unittests	68.00% <80.97%> (+0.19%)	⬆️
unittests1	58.33% <81.10%> (+0.36%)	⬆️
unittests2	39.66% <26.42%> (+0.09%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.