Preserve element type for mvcombine ARRAY_AGG on the DataFusion route

[noCharger]

Description

PPL mvcombine lowers to Calcite ARRAY_AGG. On the analytics-engine (DataFusion) route this currently fails with No enum constant org.opensearch.analytics.spi.AggregateFunction.ARRAY_AGG — the aggregate is not recognized. Mapping it onto the existing list()/values() array-agg operator is not correct either: that operator casts every element to VARCHAR by design (the PPL list()/values() multivalue-string contract), so mvcombine over a non-string field would return ARRAY<VARCHAR> instead of preserving the source type.

This adds a dedicated type-preserving path so the collected array keeps the source element type (e.g. ARRAY<INT>, not ARRAY<VARCHAR>):

• AggregateFunction.fromNameOrError — recognize ARRAY_AGG as the LIST (collect-into-array) family for backend viability and the STATE_EXPANDING decomposition. Element typing is handled separately in the rewriter, so this mapping only affects viability/decomposition shape.
• DataFusionFragmentConvertor — add LOCAL_ARRAY_AGG_TYPED_OP, a LocalAggOp whose return type is derived from the operand (not hardcoded VARCHAR) and which does not cast the element. It drops null elements (matching ARRAY_AGG ... FILTER (IS NOT NULL)) via filtersNullArgs, and maps to the same type-generic DataFusion array_agg UDAF as LOCAL_ARRAY_AGG_OP.
• PplAggregateCallRewriter — route ARRAY_AGG to the typed op for the PARTIAL/SINGLE stage, and to the existing (already type-generic) list_merge for the FINAL merge.

list()/values() are unchanged (still ARRAY<VARCHAR>). No Rust change — the array_agg/list_merge UDAFs already derive the element type from the input.

Verified on the DataFusion route (composite/parquet index, analytics.planner.prefer_metadata_driver=false): source=<idx> | fields category, age | mvcombine age over an integer field returns unquoted integers ([30, 30], [25, 35], [40]) with the query profile reporting chosen_backend=datafusion; list()/values() still return ARRAY<VARCHAR>.

Related Issues

Enables mvcombine on the analytics-engine (DataFusion) route. The mvcombine command lives in the SQL plugin (opensearch-project/sql#4766).

Check List

• Functionality includes testing.
  • AggregateFunctionTests#testArrayAggResolvesToList — ARRAY_AGG resolves to the LIST family.
  • DataFusionFragmentConvertorTests#testAttachFragmentOnTop_ArrayAggOverScalar_MeasureArgIsRawFieldNotVarcharCast — the substrait measure arg is the raw field, not a VARCHAR cast (the type-preservation guarantee).
• Commits are signed per the DCO using --signoff.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit f6cfd03)

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 Duplicate Function Signature Both LOCAL_ARRAY_AGG_OP and LOCAL_ARRAY_AGG_TYPED_OP are registered under the same DataFusion name array_agg in ADDITIONAL_AGGREGATE_SIGS. If the mapping is keyed by name or by the Calcite op, this duplicate registration may cause lookups to be ambiguous or overwrite the earlier entry, potentially breaking the untyped list()/values() path or the new typed path depending on iteration order. Confirm the mapping infrastructure supports multiple Calcite ops mapping to the same backend function name. FunctionMappings.s(LOCAL_ARRAY_AGG_OP, "array_agg"), FunctionMappings.s(LOCAL_ARRAY_AGG_TYPED_OP, "array_agg"), Possible Issue When the input to ARRAY_AGG is already an array type (getComponentType() != null), the rewriter switches to LOCAL_LIST_MERGE_OP with filterArg=-1. However, LIST_MERGE is a FINAL-side merge operator, and using it in place of a PARTIAL/SINGLE ARRAY_AGG for an already-array input may produce incorrect semantics (flattening/un-nesting instead of collecting arrays into an array of arrays). Verify this branch is intended and covered by tests, since mvcombine over an array-typed field is a plausible input. if (arrayAggArg0.getComponentType() != null) { arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP; arrayAggType = arrayAggArg0; } else { // nullable element: NOT NULL trips Calcite typeMatchesInferred arrayAggOp = DataFusionFragmentConvertor.LOCAL_ARRAY_AGG_TYPED_OP; RelDataType arrayAggElem = arrayAggTf.createTypeWithNullability(arrayAggArg0, true); arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggTf.createArrayType(arrayAggElem, -1), true); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to f6cfd03

Explore these optional code suggestions:

Category	Suggestion	Impact
General	Ensure aggregate result type is nullable When the input is already an array (component type non-null), the resulting arrayAggType is assigned directly from the input field type, but the aggregate result should always be nullable (empty group -> null). Wrap arrayAggArg0 with createTypeWithNullability(..., true) to avoid a potential NOT NULL array tripping Calcite's typeMatchesInferred check. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [193-201] if (arrayAggArg0.getComponentType() != null) { arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP; - arrayAggType = arrayAggArg0; + arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true); } else { // nullable element: NOT NULL trips Calcite typeMatchesInferred arrayAggOp = DataFusionFragmentConvertor.LOCAL_ARRAY_AGG_TYPED_OP; RelDataType arrayAggElem = arrayAggTf.createTypeWithNullability(arrayAggArg0, true); arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggTf.createArrayType(arrayAggElem, -1), true); } Suggestion importance[1-10]: 5 __ Why: Making the aggregate result nullable is a reasonable defensive change since aggregates over empty groups may return null, but it's not clear this is causing an actual bug in the current PR since the input type may already be nullable.	Low
	Avoid masking future enum additions The special-case is placed before the valueOf fallback, but if an ARRAY_AGG enum constant is ever added later this shortcut will silently mask it. Consider first attempting the enum lookup and only mapping to LIST when the name isn't a recognized enum value, or add a comment/assertion documenting the intentional override. sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java [142-146] -if (name.equalsIgnoreCase("ARRAY_AGG")) { - return LIST; +// mvcombine's ARRAY_AGG uses the LIST decomposition; element type is preserved in +// PplAggregateCallRewriter (LOCAL_ARRAY_AGG_TYPED_OP). Intentional alias to LIST. +try { + return valueOf(name.toUpperCase(java.util.Locale.ROOT)); +} catch (IllegalArgumentException ignored) { + if (name.equalsIgnoreCase("ARRAY_AGG")) { + return LIST; + } + throw new IllegalStateException("Unrecognized aggregate function [" + name + "]"); } Suggestion importance[1-10]: 3 __ Why: Minor stylistic concern about future-proofing; the current placement is intentional and the comment already documents the override. Low impact.	Low

---

Previous suggestions

Suggestions up to commit 7fbfc40

Category	Suggestion	Impact
Possible issue	Ensure array return type is nullable When the input is already an array (component type non-null), the rewrite reuses the input's array type verbatim as the return type. If the input array is declared NOT NULL, this will trip Calcite's typeMatchesInferred check on the LIST_MERGE op just like the scalar branch warns about. Wrap arrayAggArg0 with createTypeWithNullability(..., true) to force a nullable array return type. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [193-196] if (arrayAggArg0.getComponentType() != null) { arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP; - arrayAggType = arrayAggArg0; + arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true); } else { Suggestion importance[1-10]: 6 __ Why: Reasonable defensive suggestion: if the input array is NOT NULL, reusing it as the return type could trip Calcite's typeMatchesInferred check, similar to the scalar branch. However, it's unclear from the diff whether the input array can be NOT NULL in practice.	Low
General	Handle null name argument safely Guard against a null name to avoid a NullPointerException from equalsIgnoreCase, matching the previous behavior where the valueOf call would throw a more specific error. Consider validating input before the case-insensitive comparison. sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java [141-146] public static AggregateFunction fromNameOrError(String name) { + if (name == null) { + throw new IllegalStateException("Unrecognized aggregate function [null]"); + } // mvcombine's ARRAY_AGG uses the LIST decomposition; element type is preserved in // PplAggregateCallRewriter (LOCAL_ARRAY_AGG_TYPED_OP). if (name.equalsIgnoreCase("ARRAY_AGG")) { return LIST; } Suggestion importance[1-10]: 3 __ Why: Minor defensive null check; the previous code would have thrown NPE on name.toUpperCase() anyway, so this doesn't materially change robustness.	Low

Suggestions up to commit f6cfd03

Category	Suggestion	Impact
Possible issue	Ensure aggregate return type is nullable When the input is already an array (component type not null), reusing the input's type as the aggregate return type may produce a NOT NULL array that trips Calcite's typeMatchesInferred check, mirroring the concern noted for the scalar branch. Force the returned array type to be nullable to keep it consistent with the op's inferred type. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [193-196] if (arrayAggArg0.getComponentType() != null) { arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP; - arrayAggType = arrayAggArg0; + arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true); } else { Suggestion importance[1-10]: 5 __ Why: The suggestion raises a plausible concern about type nullability consistency with Calcite's typeMatchesInferred check, mirroring the pattern used in the scalar branch. However, without evidence that the input array type is actually NOT NULL in practice, the impact is speculative.	Low

Suggestions up to commit 5dc833e

Category	Suggestion	Impact
General	Force nullable array type for merge branch When the input is already an array (component type != null), the resulting arrayAggType should still be forced nullable to avoid Calcite's typeMatchesInferred mismatch, since the merge op's inferred return type is nullable. Wrap arrayAggArg0 with createTypeWithNullability(..., true) in that branch as well. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [193-201] if (arrayAggArg0.getComponentType() != null) { arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP; - arrayAggType = arrayAggArg0; + arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true); } else { // nullable element: NOT NULL trips Calcite typeMatchesInferred arrayAggOp = DataFusionFragmentConvertor.LOCAL_ARRAY_AGG_TYPED_OP; RelDataType arrayAggElem = arrayAggTf.createTypeWithNullability(arrayAggArg0, true); arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggTf.createArrayType(arrayAggElem, -1), true); } Suggestion importance[1-10]: 5 __ Why: Reasonable defensive change to ensure nullability alignment with the merge op's inferred type, which could prevent Calcite typeMatchesInferred errors, though the actual necessity depends on LOCAL_LIST_MERGE_OP's return type inference.	Low
	Avoid NPE on null name argument Guard against a null name before invoking equalsIgnoreCase, otherwise callers passing a null aggregate name will hit a NullPointerException here instead of the previous IllegalStateException from valueOf. Consider validating input first or reversing the comparison to "ARRAY_AGG".equalsIgnoreCase(name). sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java [141-146] public static AggregateFunction fromNameOrError(String name) { // mvcombine's ARRAY_AGG uses the LIST decomposition; element type is preserved in // PplAggregateCallRewriter (LOCAL_ARRAY_AGG_TYPED_OP). - if (name.equalsIgnoreCase("ARRAY_AGG")) { + if ("ARRAY_AGG".equalsIgnoreCase(name)) { return LIST; } Suggestion importance[1-10]: 3 __ Why: Minor defensive change; swapping to "ARRAY_AGG".equalsIgnoreCase(name) avoids NPE on null input, but it's unclear whether null names are a valid concern here.	Low

Suggestions up to commit 6fc048c

Category	Suggestion	Impact
General	Ensure merged array type is nullable When merging a per-shard array in the FINAL stage, the returned array type should be nullable to match the aggregate result contract (empty groups → null). Reusing arrayAggArg0 directly may propagate a NOT NULL array type that trips Calcite's typeMatchesInferred check downstream. Wrap with createTypeWithNullability(..., true). sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java [197-200] if (arrayAggArg0.getComponentType() != null) { - // PARTIAL->FINAL merge: arg0 is already a typed array; preserve it. + // PARTIAL->FINAL merge: arg0 is already a typed array; preserve it (nullable array). arrayAggOp = DataFusionFragmentConvertor.LOCAL_LIST_MERGE_OP; - arrayAggType = arrayAggArg0; + arrayAggType = arrayAggTf.createTypeWithNullability(arrayAggArg0, true); } else { Suggestion importance[1-10]: 5 __ Why: Making the merged array type explicitly nullable is a reasonable defensive change to align with Calcite's inferred type expectations, though the input array from a prior aggregate is typically already nullable.	Low
	Null-guard the name comparison Guard against a null name argument before calling equalsIgnoreCase, otherwise a null input will throw NPE here instead of surfacing the existing "Unrecognized aggregate function" error path. This preserves the original error semantics for null inputs. sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AggregateFunction.java [145-147] -if (name.equalsIgnoreCase("ARRAY_AGG")) { +if (name != null && name.equalsIgnoreCase("ARRAY_AGG")) { return LIST; } Suggestion importance[1-10]: 3 __ Why: Null-guarding preserves original NPE-vs-IllegalStateException semantics, but the original code also would have thrown NPE on name.toUpperCase(...), so this is a minor behavioral consistency improvement.	Low

[github-actions]

✅ Gradle check result for 6fc048c: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 71.40%. Comparing base (03a4de3) to head (f6cfd03).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22617      +/-   ##
============================================
- Coverage     71.42%   71.40%   -0.02%     
- Complexity    76786    76789       +3     
============================================
  Files          6148     6148              
  Lines        357980   357980              
  Branches      52177    52177              
============================================
- Hits         255689   255624      -65     
- Misses        81937    82033      +96     
+ Partials      20354    20323      -31     

☔ View full report in Codecov by Harness.
📢 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 5dc833e

[github-actions]

❌ Gradle check result for 5dc833e: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit f6cfd03

[github-actions]

Persistent review updated to latest commit 7fbfc40

[github-actions]

❌ Gradle check result for 7fbfc40: FAILURE

Please examine the workflow log, locate, and copy-paste the failure(s) below, then iterate to green. Is the failure a flaky test unrelated to your change?

[github-actions]

Persistent review updated to latest commit f6cfd03

[github-actions]

✅ Gradle check result for f6cfd03: SUCCESS