[analytics-engine] wire PPL TAKE / FIRST / LAST / LIST / VALUES end-to-end

[sandeshkr419]

Summary

Wires five PPL state-expanding aggregates through the analytics engine with correct multi-shard behavior:

• take(field, n) — bounded array of first N values
• first(field) / last(field) — first/last value per group
• list(field) / values(field) — array of values (values = distinct)

All work both single-shard and across shards.

Commits

1. SPI (analytics-framework): five new AggregateFunction enum entries plus IntermediateTypeResolver so the planner can derive the FINAL exchange-column type from the actual call's arg type rather than a constant.

2. Backend (analytics-backend-datafusion):

   • Rust UDAF for take(field, n) with bounded List<elem> state and coordinator-side merge in update_batch (DataFusion's substrait consumer ignores AggregationPhase, so FINAL is also a single-pass aggregate that needs to un-nest the per-shard state).
   • Substrait extensions (take, first_value, last_value, array_agg). take's 2-arg shape uses any1/any2 so isthmus's wildcard binding accepts mismatched arg types.
   • LOCAL_*_OP stubs + a pre-emit RelShuttle that rebinds PPL aggregations onto the stubs so isthmus resolves them through our extension catalog.
   • SubstraitPlanRewriter.visit(Aggregate) inlines literal Project columns as Substrait Expression.Literal on aggregate-call args (Calcite materializes constant aggregate-args as $f1 columns).

3. Planner (analytics-engine):

   • Move AggregateCallAnnotation out of AggregateCall.rexList onto an OpenSearchAggregate side-map. Calcite's AggCallBinding.preOperands would otherwise corrupt return-type inference for PPL's ARG0_ARRAY.
   • DistributedAggregateRewriter drives intermediate-field types through IntermediateTypeResolver and pins the FINAL aggCall return type to the StageInputScan column for state-expanding aggregates.
   • Cross-shard literal-arg forwarding for take's N: OpenSearchAggregateSplitRule captures literal aggregate-args from the original Project; the rewriter re-creates them as a constant Project below FINAL so the Rust accumulator preserves N.

4. Cross-shard LIST / VALUES (analytics-engine):

   • PPL STRING_ARRAY return type forces ARRAY<VARCHAR> regardless of input element. OpenSearchAggregateSplitRule repairs the PARTIAL aggCall return type to ARRAY<actual-arg0> (PARTIAL only — FINAL keeps the original to satisfy Volcano's parent row-type check).
   • DataFusion's substrait consumer ignores AggregationPhase, so a FINAL array_agg(state) re-wraps each shard's list rather than concatenating. Custom Rust UDAFs list_merge / list_merge_distinct un-nest the per-shard state. The convertor routes LIST/VALUES at FINAL (arg0 is already a list type) to the merge UDAFs; PARTIAL keeps array_agg.

Test plan

• Rust unit tests for take and list_merge accumulators (PARTIAL, FINAL
  un-nest, distinct dedup, null handling).
• CoordinatorReduceIT — single-shard + cross-shard for each of the five
  aggregates. 18/18 pass, 0 deferred.
• Full :sandbox:qa:analytics-engine-rest:integTest green.
• Existing planner unit tests updated for shape changes from the
  annotation side-map move.

Notes

The IT cluster runs with -da:org.apache.calcite... to silence Calcite's
Aggregate.typeMatchesInferred debug assertion. The wire schema is
derived in Rust from the producer's substrait plan (post #21690), not from
Calcite types, so this assertion is architecturally redundant for our setup
— enforced at runtime by Rust's derive_schema_from_partial_plan + Arrow's
ensure_field_compatibility. Production runs without -ea.

Check List

• Functionality includes testing.
• API changes companion pull request created, if applicable.
• Public documentation issue/PR created, if applicable.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit b674382)

Here are some key observations to aid the review process:

🧪 PR contains tests

🔒 No security concerns identified

✅ No TODO sections

🔀 No multiple PR themes

⚡ Recommended focus areas for review Possible Issue In the convert() override within the AggregateFunctionConverter anonymous class, the code attempts to inline literal Project columns into aggregate function arguments. However, the logic only checks if the input is a Project and if the argument is a FieldReference pointing to a literal. If the Project is not the immediate input (e.g., there's an intermediate node), or if the literal column is not at the expected offset due to other transformations, the inlining may fail silently or produce incorrect results. Additionally, the simpleStructOffset helper returns null for complex references, but the code does not validate that the offset is within bounds before accessing projects.get(offset), which could throw an IndexOutOfBoundsException if the Project's projection list has been modified elsewhere. public Optional<AggregateFunctionInvocation> convert( RelNode input, Type.Struct inputType, AggregateCall call, Function<RexNode, Expression> rexConverter ) { Optional<AggregateFunctionInvocation> bound = super.convert(input, inputType, call, rexConverter); if (bound.isEmpty() || !(input instanceof org.apache.calcite.rel.core.Project project)) { return bound; } AggregateFunctionInvocation fn = bound.get(); List<RexNode> projects = project.getProjects(); List<FunctionArg> args = fn.arguments(); List<FunctionArg> rewritten = null; for (int i = 0; i < args.size(); i++) { FunctionArg arg = args.get(i); if (!(arg instanceof io.substrait.expression.FieldReference fr)) continue; Integer offset = simpleStructOffset(fr); if (offset == null || offset < 0 || offset >= projects.size()) continue; if (!(projects.get(offset) instanceof RexLiteral rexLit)) continue; if (rewritten == null) rewritten = new ArrayList<>(args); rewritten.set(i, rexConverter.apply(rexLit)); } if (rewritten == null) return bound; return Optional.of(ImmutableAggregateFunctionInvocation.builder().from(fn).arguments(rewritten).build()); } Possible Issue The captureLiteralArgsForFinal method assumes that if an aggregate call has more than one argument and the function is STATE_EXPANDING, arguments 1+ are configuration literals. However, the code only checks if the Project column at that index is a RexLiteral. If the Project has been rewritten or if the argList indices no longer correspond to the expected Project columns (e.g., due to column reordering or intermediate transformations), the method may capture the wrong literals or miss them entirely. This could cause the FINAL aggregate to receive incorrect or missing configuration parameters, leading to runtime errors or incorrect results. private static Map<Integer, List<RexLiteral>> captureLiteralArgsForFinal(List<AggregateCall> aggCalls, RelNode child) { if (!(RelNodeUtils.unwrapHep(child) instanceof Project project)) { return Map.of(); } List<RexNode> projects = project.getProjects(); Map<Integer, List<RexLiteral>> captured = new LinkedHashMap<>(); for (int i = 0; i < aggCalls.size(); i++) { AggregateCall call = aggCalls.get(i); AggregateFunction fn = AggregateFunction.fromSqlAggFunction(call.getAggregation()); if (fn == null || fn.getType() != AggregateFunction.Type.STATE_EXPANDING) continue; List<Integer> args = call.getArgList(); if (args.size() < 2) continue; // arg 0 is the value/state column; args 1+ are the configuration literals. List<RexLiteral> literals = new ArrayList<>(args.size() - 1); boolean allLiteral = true; for (int a = 1; a < args.size(); a++) { int colIdx = args.get(a); if (colIdx < 0 || colIdx >= projects.size() || !(projects.get(colIdx) instanceof RexLiteral lit)) { allLiteral = false; break; } literals.add(lit); } if (allLiteral && !literals.isEmpty()) { captured.put(i, List.copyOf(literals)); } } return captured; } Possible Issue The rewrite method rebuilds the FINAL aggregate's input by injecting a Project node that materializes captured literal arguments as constant columns. The code constructs projectExprs and projectNames by first copying all existing columns from newFinalInput, then appending the literals. However, if newFinalInput's row type has already been modified (e.g., by a prior rewrite or if the StageInputScan's type override changed the column count), the loop that initializes projectExprs may produce a mismatch between the number of columns and the expected schema. Additionally, the code does not validate that the literal column indices in extraLiteralColIdxByCallIdx are within the bounds of the new Project's row type before using them in buildFinalCall, which could cause an IndexOutOfBoundsException when constructing the final argList. // Re-create captured literal aggregate-args (e.g. TAKE's N) as constant Project // columns. SubstraitPlanRewriter.visit(Aggregate) inlines them into the substrait. Map<Integer, List<RexLiteral>> extraLiterals = finalAgg.getFinalExtraLiteralArgs(); Map<Integer, List<Integer>> extraLiteralColIdxByCallIdx; if (extraLiterals.isEmpty()) { extraLiteralColIdxByCallIdx = Map.of(); } else { int origColCount = newFinalInput.getRowType().getFieldCount(); List<RexNode> projectExprs = new ArrayList<>(origColCount); List<String> projectNames = new ArrayList<>(origColCount); for (int idx = 0; idx < origColCount; idx++) { RelDataType fieldType = newFinalInput.getRowType().getFieldList().get(idx).getType(); projectExprs.add(new RexInputRef(idx, fieldType)); projectNames.add(newFinalInput.getRowType().getFieldList().get(idx).getName()); } java.util.LinkedHashMap<Integer, List<Integer>> idxMap = new java.util.LinkedHashMap<>(); for (Map.Entry<Integer, List<RexLiteral>> entry : extraLiterals.entrySet()) { List<Integer> colIdxs = new ArrayList<>(entry.getValue().size()); for (int litI = 0; litI < entry.getValue().size(); litI++) { RexLiteral lit = entry.getValue().get(litI); int colIdx = projectExprs.size(); projectExprs.add(lit); projectNames.add("$lit_call" + entry.getKey() + "_" + litI); colIdxs.add(colIdx); } idxMap.put(entry.getKey(), List.copyOf(colIdxs)); } RelDataTypeFactory.Builder rowTypeBuilder = tf.builder(); for (int idx = 0; idx < projectExprs.size(); idx++) { rowTypeBuilder.add(projectNames.get(idx), projectExprs.get(idx).getType()); } newFinalInput = new OpenSearchProject( newFinalInput.getCluster(), newFinalInput.getTraitSet(), newFinalInput, projectExprs, rowTypeBuilder.build(), finalAgg.getViableBackends() ); extraLiteralColIdxByCallIdx = idxMap; } List<AggregateCall> rebuiltCalls = new ArrayList<>(finalAgg.getAggCallList().size()); for (int i = 0; i < finalAgg.getAggCallList().size(); i++) { AggregateCall call = finalAgg.getAggCallList().get(i); IntermediateField field = perCallField.get(i); int stateColIdx = groupCount + i;

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to b674382

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Handle empty column case explicitly The condition n_col.len() == 0 allows proceeding with an unresolved limit when the column is empty. This could cause runtime errors in subsequent operations that expect a resolved limit. Return an error or set a default limit when the column is empty. sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/take.rs [238-246] fn resolve_limit_from(&mut self, n_col: &ArrayRef) -> Result<()> { - if self.limit.is_some() || n_col.len() == 0 { + if self.limit.is_some() { + return Ok(()); + } + if n_col.len() == 0 { + self.limit = Some(DEFAULT_LIMIT); return Ok(()); } let scalar = ScalarValue::try_from_array(n_col, 0)?; self.limit = Some(scalar_to_i64(&scalar)?); Ok(()) } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies that when n_col.len() == 0, the limit remains unresolved, which could cause issues in subsequent operations. Setting DEFAULT_LIMIT when the column is empty is a reasonable defensive measure that improves robustness.	Medium
General	Strengthen list type detection The check arg0Type.getComponentType() != null may incorrectly identify non-list types as lists if they have a component type. Use a more explicit type check against DataType list variants to ensure correct PARTIAL vs FINAL detection. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [112-119] -boolean arg0IsList = arg0Type.getComponentType() != null; +boolean arg0IsList = arg0Type instanceof org.apache.calcite.rel.type.ArraySqlType + || (arg0Type.getSqlTypeName() != null + && arg0Type.getSqlTypeName().getFamily() == org.apache.calcite.sql.type.SqlTypeFamily.ARRAY); if (arg0IsList) { targetOp = isValues ? DataFusionFragmentConvertor.LOCAL_LIST_MERGE_DISTINCT_OP : DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP; targetDistinct = false; explicitReturnType = arg0Type; } Suggestion importance[1-10]: 6 __ Why: The suggestion raises a valid concern about the robustness of the list type detection using getComponentType() != null. Using explicit checks against Calcite's type system (e.g., ArraySqlType or SqlTypeFamily.ARRAY) would be more precise and reduce the risk of false positives.	Low
	Avoid unnecessary clone for duplicates The value.clone() is unnecessary when the insert fails (duplicate detected). Move the clone inside the conditional to avoid allocating memory for duplicates that will be discarded. sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/list_merge.rs [136-143] fn push(&mut self, value: ScalarValue) { if let Some(seen) = self.seen.as_mut() { - if !seen.insert(value.clone()) { + if seen.contains(&value) { return; } + seen.insert(value.clone()); } self.buf.push(value); } Suggestion importance[1-10]: 5 __ Why: The suggestion correctly identifies that value.clone() is performed even when the value is a duplicate and will be discarded. Using contains() before insert() avoids the unnecessary clone, improving performance for duplicate-heavy workloads. However, this adds an extra hash lookup, so the net benefit depends on the duplicate rate.	Low
	Use immutable list construction Use List.of() for immutable list construction when the size is known and elements are available upfront. This avoids unnecessary ArrayList allocation and provides better performance for small, fixed-size lists. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java [195-197] -List<Integer> argList = new ArrayList<>(1 + extraColIdxs.size()); -argList.add(finalArgIdx); -argList.addAll(extraColIdxs); +List<Integer> argList = extraColIdxs.isEmpty() + ? List.of(finalArgIdx) + : Stream.concat(Stream.of(finalArgIdx), extraColIdxs.stream()).toList(); Suggestion importance[1-10]: 3 __ Why: The suggestion to use List.of() or Stream for immutable list construction is a minor optimization. However, the argList is passed to AggregateCall.create() which may expect a mutable list or perform internal operations. The current code is clear and safe, making this a low-impact style suggestion.	Low

---

Previous suggestions

Suggestions up to commit 0d26df5

Category	Suggestion	Impact
General	Initialize variable at declaration Initialize extraLiteralColIdxByCallIdx directly in the declaration to avoid potential uninitialized variable issues and improve code clarity. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java [99-103] Map<Integer, List<RexLiteral>> extraLiterals = finalAgg.getFinalExtraLiteralArgs(); -Map<Integer, List<Integer>> extraLiteralColIdxByCallIdx; -if (extraLiterals.isEmpty()) { - extraLiteralColIdxByCallIdx = Map.of(); -} else { +Map<Integer, List<Integer>> extraLiteralColIdxByCallIdx = Map.of(); +if (!extraLiterals.isEmpty()) { ... extraLiteralColIdxByCallIdx = idxMap; } Suggestion importance[1-10]: 4 __ Why: Reasonable code style improvement. Initializing extraLiteralColIdxByCallIdx to Map.of() at declaration and then reassigning in the non-empty case is clearer and avoids the uninitialized variable pattern, though the original code is also correct.	Low
	Use idiomatic empty check Check n_col.is_empty() instead of n_col.len() == 0 for better idiomatic Rust code and potential performance benefits with certain array implementations. sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/take.rs [239-246] fn resolve_limit_from(&mut self, n_col: &ArrayRef) -> Result<()> { - if self.limit.is_some() || n_col.len() == 0 { + if self.limit.is_some() || n_col.is_empty() { return Ok(()); } let scalar = ScalarValue::try_from_array(n_col, 0)?; self.limit = Some(scalar_to_i64(&scalar)?); Ok(()) } Suggestion importance[1-10]: 3 __ Why: Valid style improvement. Using is_empty() is more idiomatic in Rust than len() == 0, though the functional impact is minimal. The suggestion correctly identifies the location and provides a minor readability enhancement.	Low

Suggestions up to commit 2ddafc9

Category	Suggestion	Impact
General	Use idiomatic empty check Check if n_col.is_empty() is the idiomatic Rust way instead of n_col.len() == 0. This improves code readability and follows Rust conventions. sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/take.rs [240-246] fn resolve_limit_from(&mut self, n_col: &ArrayRef) -> Result<()> { - if self.limit.is_some() || n_col.len() == 0 { + if self.limit.is_some() || n_col.is_empty() { return Ok(()); } let scalar = ScalarValue::try_from_array(n_col, 0)?; self.limit = Some(scalar_to_i64(&scalar)?); Ok(()) } Suggestion importance[1-10]: 5 __ Why: Using is_empty() instead of len() == 0 is indeed more idiomatic in Rust and improves readability. However, this is a minor style improvement that doesn't affect functionality or performance significantly.	Low
	Use immutable list construction Use List.copyOf to create an immutable list directly instead of building a mutable ArrayList and then copying it. This is more efficient and prevents accidental mutation. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java [195-197] -List<Integer> argList = new ArrayList<>(1 + extraColIdxs.size()); -argList.add(finalArgIdx); -argList.addAll(extraColIdxs); +List<Integer> argList = extraColIdxs.isEmpty() + ? List.of(finalArgIdx) + : Stream.concat(Stream.of(finalArgIdx), extraColIdxs.stream()).toList(); Suggestion importance[1-10]: 4 __ Why: The suggestion to use immutable list construction is valid and could improve code clarity, but the proposed implementation using Stream.concat is more verbose and potentially less efficient than the current approach. The current code is clear and the final List.copyOf(argList) already creates an immutable copy, so the improvement is marginal.	Low
	Add defensive bounds checking Add a bounds check before accessing call.getArgList().get(0) to prevent potential IndexOutOfBoundsException if the list becomes empty between the check and access due to concurrent modification or logic errors. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java [125-126] -if (call.getArgList().isEmpty()) continue; -org.apache.calcite.rel.type.RelDataType arg0Type = input.getRowType().getFieldList().get(call.getArgList().get(0)).getType(); +List<Integer> argList = call.getArgList(); +if (argList.isEmpty()) continue; +org.apache.calcite.rel.type.RelDataType arg0Type = input.getRowType().getFieldList().get(argList.get(0)).getType(); Suggestion importance[1-10]: 3 __ Why: While defensive programming is generally good practice, the suggestion addresses a theoretical concern rather than a real issue. The isEmpty() check on the previous line already guards against empty lists, and concurrent modification is not a concern in this single-threaded Calcite rule context. The suggested change adds minimal value but does improve code clarity slightly.	Low

[github-actions]

✅ Gradle check result for 2ddafc9: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.45%. Comparing base (8f2d058) to head (2ddafc9).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21731      +/-   ##
============================================
- Coverage     73.46%   73.45%   -0.02%     
- Complexity    74825    74837      +12     
============================================
  Files          5997     6005       +8     
  Lines        339688   339767      +79     
  Branches      48961    48969       +8     
============================================
+ Hits         249558   249564       +6     
- Misses        70272    70321      +49     
- Partials      19858    19882      +24     

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Persistent review updated to latest commit 0d26df5

[sandeshkr419]

@expani Thanks for review, addressed your comments.

[github-actions]

Persistent review updated to latest commit b674382