[analytics-backend-datafusion] Implement partial/final aggregation mode for distributed execution

[sandeshkr419]

Introduces end-to-end distributed aggregate execution for the analytics plugin. Shards compute partial state, the coordinator merges to final results, and the whole path works across a multi-node multi-shard cluster. 39 of 43 ClickBench PPL queries now pass in that topology.

Architecture

The control split is layered to avoid per-function branches. Java owns the logical plan; Rust compiles each half (PARTIAL or FINAL) of whatever substrait it receives:

1. HEP marking: OpenSearchAggregateRule, OpenSearchProjectRule (with baseline scalar-op carve-out), and OpenSearchAggregateReduceRule (reduces AVG → SUM+COUNT+Project in a dedicated phase, before the Volcano split, using Calcite's testedmachinery).
2. Volcano CBO split: OpenSearchAggregateSplitRule turns SINGLE → FINAL(Exchange(PARTIAL(…))) structurally — no per-function logic.
3. Decomposition resolver (AggregateDecompositionResolver): single pass handling the three single-field shapes — pass-through (SUM/MIN/MAX), function-swap (COUNT → SUM at FINAL), engine-native merge (APPROX_COUNT_DISTINCT sketch). Rebuilds the exchange column row type with ArrowCalciteTypes (single-type bidirectional) and rewires parent projections.
4. DAG cut (DAGBuilder): structural only — cuts at OpenSearchExchangeReducer.
5. Backend conversion (DataFusionFragmentConvertor): mechanical Calcite → Substrait emission; APPROX_COUNT_DISTINCT routes to DataFusion's native HLL via the extension YAML + a subclassed AggregateFunctionConverter.
6. Instruction assembly (FragmentConversionDriver).
7. Execution handlers (PartialAggregateInstructionHandler, FinalAggregateInstructionHandler) → named FFI (preparePartialPlan, prepareFinalPlan, executeLocalPreparedPlan). No aggregate-mode enum on the wire.
8. Streaming (LocalStageScheduler): one-line arrowSchemaFromRowType(child.getFragment().getRowType()).
9. Physical (Rust agg_mode.rs): internal Mode {Default, Partial, Final}; force_aggregate_mode strips one half; CombinePartialFinalAggregate disabled on every SessionContext so Final(Partial(…)) survives to the strip pass.

Single source of truth for per-function decomposition

AggregateFunction enum (SPI) carries intermediateFields — the only authority for what the PARTIAL emits and how the FINAL reduces it. fromSqlAggFunction / toSqlAggFunction bridge between Calcite operators and enum constants. COUNT and APPROX_COUNT_DISTINCT declare their shapes there; SUM/MIN/MAX/AVG use null intermediates (pass-through post-reduction).

Multi-shard unblocks

• Empty-batch schema preservation: shards that match zero rows still emit one zero-row VSR carrying the declared schema, so the FINAL plan's first data frame carries types downstream even when WHERE filters out everything.
• Arrow IPC wire codec replaces row-oriented Object[] in FragmentExecutionResponse. Arrow's own serializer/deserializer round-trips every type natively (temporal, view, dictionary) without hand-rolled per-type dispatch. Removes the writeGenericValue(LocalDateTime) crash class that blocked all TIMESTAMP queries.
• Java↔Rust type-coercion bridge at DatafusionReduceSink.feedToSender reconciles DataFusion's physical emit types (e.g. Utf8View for string group keys) with Calcite's declared exchange schema (Utf8) before pushing into the FINAL plan's StreamingTable partition.

TODOs to revisit (tracked independently)

The goal is to unblock distributed query execution framework and avoid spilling more changes in this already huge PR.

1. Type coercion at the right layer: the Utf8View → Utf8 shim in coerceToDeclaredSchema is a runtime bridge; the principled fix is a Rust force_output_schema cast at the PARTIAL root, a Substrait view-type extension, or verified Arrow Java ↔Rust FFI symmetry for view types.
2. Aggregate decomposition catalog → per-backend SPI: AggregateFunction.intermediateFields hard-codes DataFusion's state shapes into the shared enum; the catalog should move to BackendCapabilityProvider.aggregateDecompositions() (engine keeps the structural resolver, backends declare their own metadata).
3. MultisearchCommandIT.testMultisearchCountEvalConditionalCount and AppendPipeCommandIT.testMultisearchCountEvalConditionalCount - commented out for now - need to fix and revert this test back.

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 Code Analyzer ❗

AI-powered 'Code-Diff-Analyzer' found issues on commit f87d308.

Path	Line	Severity	Description
sandbox/libs/dataformat-native/rust/Cargo.toml	31	high	New Rust dependency added: datafusion-physical-optimizer = "52.1.0". Per mandatory policy, all dependency additions must be flagged regardless of apparent legitimacy. Maintainers should verify this crate resolves to the expected artifact on crates.io and matches the datafusion ecosystem version already in use.
sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml	18	high	New Rust dependency added: datafusion-physical-optimizer = { workspace = true }. Per mandatory policy, all dependency additions must be flagged. Maintainers should confirm the workspace-level pin resolves to the correct datafusion-physical-optimizer crate and that no namespace hijacking or typosquatting is present.

The table above displays the top 10 most important findings.

Total: 2 | Critical: 0 | High: 2 | Medium: 0 | Low: 0

---

Pull Requests Author(s): Please update your Pull Request according to the report above.

Repository Maintainer(s): You can bypass diff analyzer by adding label skip-diff-analyzer after reviewing the changes carefully, then re-run failed actions. To re-enable the analyzer, remove the label, then re-run all actions.

---

⚠️ Note: The Code-Diff-Analyzer helps protect against potentially harmful code patterns. Please ensure you have thoroughly reviewed the changes beforehand.

Thanks.

[expani]

Thanks for getting this going. @sandeshkr419

Dropping some early comments on the approach.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 614a8c5)

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 Resource Leak When preparedState != null, the constructor catches exceptions during executeLocalPreparedPlan and closes senders allocated in the try block. However, if preparedState is non-null, the senders come from preparedState.senders() and are not allocated locally. The catch block at lines 125-138 closes senders only when preparedState == null, but if executeLocalPreparedPlan throws and preparedState != null, the streamPtr is closed but the state's senders are not. This could leak native resources if the state is not closed elsewhere. try { if (preparedState != null) { // Plan was already prepared by FinalAggregateInstructionHandler. The handler // registered senders in ctx.childInputs() iteration order; we re-index them // here by childStageId for lookup during feed(). int i = 0; for (Map.Entry<Integer, byte[]> child : childInputs.entrySet()) { senders.put(child.getKey(), preparedState.senders().get(i++)); } streamPtr = NativeBridge.executeLocalPreparedPlan(session.getPointer()); } else { // Legacy path (non-aggregate reduce): register partitions and execute the // fragment bytes directly. Used when no prior instruction prepared a plan. // // ctx.fragmentBytes() references each partition by its "input-<stageId>" name // (DataFusionFragmentConvertor names them this way during plan conversion). for (Map.Entry<Integer, byte[]> child : childInputs.entrySet()) { int childStageId = child.getKey(); byte[] schemaIpc = child.getValue(); long senderPtr = NativeBridge.registerPartitionStream(session.getPointer(), inputIdFor(childStageId), schemaIpc); senders.put(childStageId, new DatafusionPartitionSender(senderPtr)); } streamPtr = NativeBridge.executeLocalPlan(session.getPointer(), ctx.fragmentBytes()); } this.outStream = new StreamHandle(streamPtr, runtimeHandle); } catch (RuntimeException e) { if (streamPtr != 0) { NativeBridge.streamClose(streamPtr); } // Only close senders we allocated locally (legacy path). When preparedState // owns them, the state's close() will. if (preparedState == null) { for (DatafusionPartitionSender sender : senders.values()) { try { sender.close(); } catch (Throwable ignore) {} } session.close(); } throw e; } this.sendersByChildStageId = senders; Possible Double-Close At lines 96-105, if provider.createSink throws and backendContext is non-null, the code closes backendContext to avoid a leak. However, if backendContext is a DataFusionReduceState that was successfully prepared by the instruction handler, its session and senders are already registered and may be closed again by the state's own close method when the orchestrator cleans up. This could result in double-close on native handles if the state's close is idempotent but the native side is not. backendSink = provider.createSink(context, backendContext); } catch (Exception e) { // Sink creation failed — close backendContext to avoid resource leak. if (backendContext != null) { try { backendContext.close(); } catch (Exception closeFailure) { e.addSuppressed(closeFailure); } } Possible Resource Leak If prepareFinalPlan throws after some senders have been registered, the catch block at lines 56-63 closes all senders and the session. However, if a sender's close itself throws, the exception is silently ignored (line 59), and the session close at line 61 may not run if an earlier sender close threw an unchecked exception that bypassed the loop. This could leave the session open. for (DatafusionPartitionSender sender : senders) { try { sender.close(); } catch (Exception ignored) {} } session.close(); throw e; } Possible Issue At line 268, the resolver reads the canonical name from aggRowType.getFieldList().get(groupCount + i).getName() to align with DataFusion's $f<N> convention. However, if the aggregate call list is empty (no aggregates, only group keys), groupCount + i could be out of bounds. The loop condition i < agg.getAggCallList().size() prevents this, but if aggRowType has fewer fields than expected (e.g., due to a prior rewrite), the index could still be invalid. for (int i = 0; i < agg.getAggCallList().size(); i++) { AggregateCall call = agg.getAggCallList().get(i); String canonicalName = aggRowType.getFieldList().get(groupCount + i).getName(); CallRewrite rw = rewriteAggCall(call, finalColIdx, tf, canonicalName);

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 31e8fa2

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Validate column count before coercion Verify that batch.getFieldVectors().size() matches declaredSchema.getFields().size() before the loop to prevent IndexOutOfBoundsException when the batch has fewer columns than the declared schema. This mismatch can occur if the shard's partial aggregate emits a different column count than the coordinator expects. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [255-265] private static VectorSchemaRoot coerceToDeclaredSchema(VectorSchemaRoot batch, Schema declaredSchema, BufferAllocator alloc) { if (batch.getSchema().equals(declaredSchema)) { return batch; + } + if (batch.getFieldVectors().size() != declaredSchema.getFields().size()) { + batch.close(); + throw new IllegalStateException( + "Batch column count (" + batch.getFieldVectors().size() + ") does not match declared schema (" + declaredSchema.getFields().size() + ")" + ); } VectorSchemaRoot out = VectorSchemaRoot.create(declaredSchema, alloc); try { out.allocateNew(); int rows = batch.getRowCount(); for (int col = 0; col < declaredSchema.getFields().size(); col++) { FieldVector src = batch.getVector(col); FieldVector dst = out.getVector(col); ... } out.setRowCount(rows); } catch (RuntimeException e) { out.close(); throw e; } finally { batch.close(); } return out; } Suggestion importance[1-10]: 8 __ Why: The suggestion addresses a critical mismatch scenario where batch.getFieldVectors().size() differs from declaredSchema.getFields().size(), which would cause an IndexOutOfBoundsException in the loop. Validating upfront prevents silent failures and provides a clear diagnostic message.	Medium
	Validate prepared plan before execution The execute_prepared method panics if no plan was prepared, but this FFI function doesn't validate that condition before calling it. Add an explicit check and return an error code instead of allowing a panic across the FFI boundary. sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs [822-829] pub unsafe extern "C" fn df_execute_local_prepared_plan(session_ptr: i64) -> i64 { let session = &*(session_ptr as *const crate::local_executor::LocalSession); - ... + if session.prepared_plan.is_none() { + return ffm_error("No plan prepared on session".to_string()); + } + let mgr = get_rt_manager()?; + let _guard = mgr.io_runtime.enter(); let df_stream = session.execute_prepared().map_err(|e| e.to_string())?; Suggestion importance[1-10]: 8 __ Why: This is a valid safety concern for FFI boundaries. The execute_prepared method explicitly panics if no plan is prepared, and allowing a panic across FFI is dangerous. Adding validation prevents undefined behavior.	Medium
	Add bounds checking for array access Add bounds checking before accessing inputTypes.get(ref.getIndex()) in the RexShuttle to prevent IndexOutOfBoundsException when a RexInputRef references a column index beyond the replacement's row type field count. This can occur if the Project's expressions reference columns that were removed during aggregate rewriting. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/AggregateDecompositionResolver.java [210-214] private static RelNode replaceFirstWithRefRebinding(RelNode node, RelNode target, RelNode replacement) { if (node == target) return replacement; java.util.List<RelNode> newInputs = new java.util.ArrayList<>(); boolean changed = false; for (RelNode input : node.getInputs()) { RelNode newInput; if (input == target) { newInput = replacement; if (node instanceof Project proj) { RexBuilder rexBuilder = node.getCluster().getRexBuilder(); java.util.List<RelDataType> inputTypes = new java.util.ArrayList<>(); for (var f : replacement.getRowType().getFieldList()) { inputTypes.add(f.getType()); } RexShuttle rebind = new RexShuttle() { @Override public RexNode visitInputRef(RexInputRef ref) { + if (ref.getIndex() >= inputTypes.size()) { + throw new IllegalStateException( + "RexInputRef index " + ref.getIndex() + " out of bounds for replacement row type with " + inputTypes.size() + " fields" + ); + } RelDataType actual = inputTypes.get(ref.getIndex()); if (ref.getType().equals(actual)) return ref; return new RexInputRef(ref.getIndex(), actual); } }; java.util.List<RexNode> rebound = new java.util.ArrayList<>(proj.getProjects().size()); for (int i = 0; i < proj.getProjects().size(); i++) { RexNode expr = proj.getProjects().get(i).accept(rebind); RelDataType targetType = proj.getRowType().getFieldList().get(i).getType(); if (!expr.getType().equals(targetType)) { expr = rexBuilder.makeCast(targetType, expr); } rebound.add(expr); } return proj.copy(proj.getTraitSet(), replacement, rebound, proj.getRowType()); } } else { newInput = replaceFirstWithRefRebinding(input, target, replacement); } newInputs.add(newInput); if (newInput != input) changed = true; } return changed ? node.copy(node.getTraitSet(), newInputs) : node; } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential IndexOutOfBoundsException when ref.getIndex() exceeds inputTypes.size(). Adding bounds checking with a clear error message improves robustness and debuggability during aggregate rewriting edge cases.	Medium
	Add bounds check for array access Direct array indexing without bounds checking can panic if children() returns an empty slice. Verify that agg.children() is non-empty before accessing index 0, or use a safer method like first() with error handling. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [76-81] AggregateMode::Final => { // Current node is Partial; skip it, return its child // (the Final above will keep itself) - let child = agg.children()[0]; + let child = agg.children().first().ok_or_else(|| { + datafusion_common::DataFusionError::Internal( + "AggregateExec has no children".to_string() + ) + })?; force_aggregate_mode(Arc::clone(child), target) } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential panic from unchecked array indexing. However, in the context of AggregateExec, it's expected to have at least one child, making this a defensive programming improvement rather than a critical bug fix.	Medium
General	Ensure batch cleanup on concatenation failure Close all accumulated batches in the batches list if an exception occurs during the multi-batch concatenation phase (after the reader loop completes). Currently, only the reader-loop failures clean up batches; a failure in the concatenation logic below would leak the already-transferred batch roots. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/RowResponseCodec.java [82-110] List<VectorSchemaRoot> batches = new ArrayList<>(); Schema schema; try (ArrowStreamReader reader = new ArrowStreamReader(new ByteArrayInputStream(payload), allocator)) { VectorSchemaRoot readerRoot = reader.getVectorSchemaRoot(); schema = readerRoot.getSchema(); while (reader.loadNextBatch()) { VectorSchemaRoot batchRoot = VectorSchemaRoot.create(schema, allocator); int rowCount = readerRoot.getRowCount(); for (int i = 0; i < readerRoot.getFieldVectors().size(); i++) { readerRoot.getVector(i).makeTransferPair(batchRoot.getVector(i)).transfer(); } batchRoot.setRowCount(rowCount); batches.add(batchRoot); } } catch (IOException e) { for (VectorSchemaRoot b : batches) b.close(); throw new IllegalStateException("Failed to decode Arrow IPC payload from fragment response", e); } +if (batches.isEmpty()) { + return VectorSchemaRoot.create(schema, allocator); +} +if (batches.size() == 1) { + return batches.get(0); +} +VectorSchemaRoot combined = null; +try { + int totalRows = batches.stream().mapToInt(VectorSchemaRoot::getRowCount).sum(); + combined = VectorSchemaRoot.create(schema, allocator); + combined.allocateNew(); + ... + combined.setRowCount(totalRows); + return combined; +} catch (RuntimeException e) { + if (combined != null) combined.close(); + throw e; +} finally { + for (VectorSchemaRoot b : batches) + b.close(); +} + Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a resource leak when the concatenation phase (after the reader loop) throws an exception. The existing finally block in the concatenation section already closes batches, but wrapping the entire concatenation in a try-catch that also closes combined on failure is a defensive improvement.	Medium
	Prevent stack overflow from deep recursion The recursive call to force_aggregate_mode could lead to stack overflow for deeply nested plans. Consider adding a depth limit or converting to an iterative approach with an explicit stack for production robustness. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [59-64] -let new_children: Vec<Arc<dyn ExecutionPlan>> = agg - .children() - .into_iter() - .map(|c| force_aggregate_mode(Arc::clone(c), target)) - .collect::<Result<_>>()?; -return plan.with_new_children(new_children); +const MAX_RECURSION_DEPTH: usize = 1000; +fn force_aggregate_mode_impl( + plan: Arc<dyn ExecutionPlan>, + target: AggregateMode, + depth: usize, +) -> Result<Arc<dyn ExecutionPlan>> { + if depth > MAX_RECURSION_DEPTH { + return Err(datafusion_common::DataFusionError::Internal( + "Maximum recursion depth exceeded".to_string() + )); + } + // ... existing logic with depth + 1 passed to recursive calls +} Suggestion importance[1-10]: 5 __ Why: While recursion depth limits are good defensive programming, deeply nested execution plans are uncommon in practice. The suggestion addresses a theoretical concern rather than a likely real-world issue, and the proposed solution would require significant refactoring.	Low

---

Previous suggestions

Suggestions up to commit ee0a3d7

Category	Suggestion	Impact
Possible issue	Validate column count before coercion loop Verify that batch.getFieldVectors().size() matches declaredSchema.getFields().size() before the column loop. If the batch has fewer columns than the declared schema, accessing batch.getVector(col) will throw IndexOutOfBoundsException. Add an explicit check and throw a descriptive error if column counts mismatch. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [255-269] private static VectorSchemaRoot coerceToDeclaredSchema(VectorSchemaRoot batch, Schema declaredSchema, BufferAllocator alloc) { if (batch.getSchema().equals(declaredSchema)) { return batch; + } + if (batch.getFieldVectors().size() != declaredSchema.getFields().size()) { + batch.close(); + throw new IllegalStateException( + "Column count mismatch: batch has " + batch.getFieldVectors().size() + " columns, declared schema has " + declaredSchema.getFields().size() + ); } VectorSchemaRoot out = VectorSchemaRoot.create(declaredSchema, alloc); try { out.allocateNew(); int rows = batch.getRowCount(); for (int col = 0; col < declaredSchema.getFields().size(); col++) { FieldVector src = batch.getVector(col); FieldVector dst = out.getVector(col); ... } out.setRowCount(rows); } catch (RuntimeException e) { out.close(); throw e; } finally { batch.close(); } return out; } Suggestion importance[1-10]: 8 __ Why: This is a valid safety check that prevents IndexOutOfBoundsException when the batch and declared schema have mismatched column counts. The suggestion correctly identifies that the existing code assumes column count equality without verification. The improved code adds an explicit check with a clear error message, which is important for debugging schema mismatches at the Java→Rust boundary.	Medium
	Prevent panic across FFI boundary The execute_prepared method panics if no plan was prepared, but this panic will cross the FFI boundary unsafely. Add explicit validation and return an error code instead of allowing a panic to propagate through the C ABI. sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs [822-829] pub unsafe extern "C" fn df_execute_local_prepared_plan(session_ptr: i64) -> i64 { let session = &*(session_ptr as *const crate::local_executor::LocalSession); + if session.prepared_plan.is_none() { + return ffm_error("execute_local_prepared_plan called without a prepared plan"); + } ... let df_stream = session.execute_prepared().map_err(|e| e.to_string())?; Suggestion importance[1-10]: 8 __ Why: This is a critical safety issue. The execute_prepared method panics if no plan exists, and panics crossing FFI boundaries cause undefined behavior. The suggestion correctly proposes validating prepared_plan before calling execute_prepared.	Medium
	Add bounds check for input reference Add bounds checking before accessing inputTypes.get(ref.getIndex()) in the RexShuttle visitor. If the replacement's row type has fewer fields than the reference index, this will throw IndexOutOfBoundsException. Validate that ref.getIndex() < inputTypes.size() and handle the mismatch case explicitly. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/AggregateDecompositionResolver.java [210-214] -private static RelNode replaceFirstWithRefRebinding(RelNode node, RelNode target, RelNode replacement) { - if (node == target) return replacement; - java.util.List<RelNode> newInputs = new java.util.ArrayList<>(); - boolean changed = false; - for (RelNode input : node.getInputs()) { - RelNode newInput; - if (input == target) { - newInput = replacement; - if (node instanceof Project proj) { - RexBuilder rexBuilder = node.getCluster().getRexBuilder(); - java.util.List<RelDataType> inputTypes = new java.util.ArrayList<>(); - for (var f : replacement.getRowType().getFieldList()) { - inputTypes.add(f.getType()); - } - RexShuttle rebind = new RexShuttle() { - @Override - public RexNode visitInputRef(RexInputRef ref) { - RelDataType actual = inputTypes.get(ref.getIndex()); - if (ref.getType().equals(actual)) return ref; - return new RexInputRef(ref.getIndex(), actual); - } - }; - java.util.List<RexNode> rebound = new java.util.ArrayList<>(proj.getProjects().size()); - for (int i = 0; i < proj.getProjects().size(); i++) { - RexNode expr = proj.getProjects().get(i).accept(rebind); - RelDataType targetType = proj.getRowType().getFieldList().get(i).getType(); - if (!expr.getType().equals(targetType)) { - expr = rexBuilder.makeCast(targetType, expr); - } - rebound.add(expr); - } - return proj.copy(proj.getTraitSet(), replacement, rebound, proj.getRowType()); - } - } else { - newInput = replaceFirstWithRefRebinding(input, target, replacement); +RexShuttle rebind = new RexShuttle() { + @Override + public RexNode visitInputRef(RexInputRef ref) { + if (ref.getIndex() >= inputTypes.size()) { + throw new IllegalStateException( + "RexInputRef index " + ref.getIndex() + " out of bounds for replacement row type with " + inputTypes.size() + " fields" + ); } - newInputs.add(newInput); - if (newInput != input) changed = true; + RelDataType actual = inputTypes.get(ref.getIndex()); + if (ref.getType().equals(actual)) return ref; + return new RexInputRef(ref.getIndex(), actual); } - return changed ? node.copy(node.getTraitSet(), newInputs) : node; -} +}; Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential IndexOutOfBoundsException when accessing inputTypes.get(ref.getIndex()) without bounds checking. However, the scenario is somewhat unlikely in practice since the replacement's row type should match the Project's expectations if the decomposition logic is correct. The check adds defensive programming but may not catch real bugs if the root cause is incorrect decomposition logic upstream.	Medium
	Validate children before array access Direct array indexing without bounds checking can panic if children() returns an empty slice. Validate that children() has at least one element before accessing index 0 to prevent potential runtime panics. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [76-81] AggregateMode::Final => { // Current node is Partial; skip it, return its child // (the Final above will keep itself) - let child = agg.children()[0]; - force_aggregate_mode(Arc::clone(child), target) + let children = agg.children(); + if children.is_empty() { + return Err(datafusion_common::DataFusionError::Internal( + "AggregateExec has no children".to_string() + )); + } + force_aggregate_mode(Arc::clone(children[0]), target) } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential panic from unchecked array access at agg.children()[0]. Adding bounds checking would prevent runtime panics, though AggregateExec typically has at least one child in practice.	Medium
General	Close WriteChannel to prevent resource leak Close the WriteChannel explicitly before returning the response. Arrow's WriteChannel wraps a WritableByteChannel and may buffer data; not closing it risks incomplete writes or resource leaks. Add a channel.close() call in a finally block or use try-with-resources. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [246-247] FragmentExecutionResponse collectResponse(EngineResultStream stream, @Nullable AnalyticsShardTask task) { ByteArrayOutputStream baos = new ByteArrayOutputStream(); - WriteChannel channel = new WriteChannel(Channels.newChannel(baos)); Schema schema = null; int totalRows = 0; Iterator<EngineResultBatch> it = stream.iterator(); - try { + try (WriteChannel channel = new WriteChannel(Channels.newChannel(baos))) { while (it.hasNext()) { if (task != null && task.isCancelled()) { throw new TaskCancelledException("task cancelled: " + task.getReasonCancelled()); } EngineResultBatch batch = it.next(); VectorSchemaRoot root = batch.getArrowRoot(); try { if (schema == null) { schema = root.getSchema(); MessageSerializer.serialize(channel, schema); } try (ArrowRecordBatch recordBatch = new VectorUnloader(root).getRecordBatch()) { MessageSerializer.serialize(channel, recordBatch); } totalRows += root.getRowCount(); } finally { root.close(); } } if (schema != null) { ArrowStreamWriter.writeEndOfStream(channel, IpcOption.DEFAULT); } } catch (IOException e) { throw new IllegalStateException("Failed to serialize fragment output as Arrow IPC stream", e); } return new FragmentExecutionResponse(baos.toByteArray(), totalRows); } Suggestion importance[1-10]: 6 __ Why: The suggestion correctly identifies that WriteChannel should be closed to ensure buffered data is flushed. However, WriteChannel wraps a ByteArrayOutputStream which doesn't require explicit closing for correctness (it's an in-memory buffer). The try-with-resources pattern is still good practice for consistency and future-proofing, but the impact is lower than suggested since no actual resource leak occurs with ByteArrayOutputStream.	Low
	Add warning for unexpected plan structure The function silently returns None when encountering repartition/coalesce nodes with multiple children, which may hide unexpected plan structures. Consider logging a warning or returning an error to make debugging easier when the plan structure doesn't match expectations. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [113-123] fn find_partial_input(plan: Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> { ... if plan.as_any().downcast_ref::<RepartitionExec>().is_some() || plan .as_any() .downcast_ref::<CoalescePartitionsExec>() .is_some() { let children = plan.children(); if children.len() == 1 { return find_partial_input(Arc::clone(children[0])); } + // Unexpected: repartition/coalesce with != 1 child + log::warn!("find_partial_input: unexpected child count {} for {}", children.len(), plan.name()); } None } Suggestion importance[1-10]: 4 __ Why: While adding logging for unexpected plan structures could aid debugging, this is a minor enhancement. The function correctly returns None for unhandled cases, and the silent return is acceptable behavior for an optional search function.	Low

Suggestions up to commit 1e81404

Suggestions up to commit e648ee7

Category	Suggestion	Impact
Possible issue	Close channel to flush buffered writes Close the WriteChannel explicitly after serialization completes to ensure all buffered data is flushed and resources are released. Without closing the channel, buffered writes may not be flushed to the underlying ByteArrayOutputStream, potentially resulting in incomplete IPC payloads. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java [241-278] -FragmentExecutionResponse collectResponse(EngineResultStream stream, @Nullable AnalyticsShardTask task) { - ByteArrayOutputStream baos = new ByteArrayOutputStream(); - WriteChannel channel = new WriteChannel(Channels.newChannel(baos)); - Schema schema = null; - int totalRows = 0; - Iterator<EngineResultBatch> it = stream.iterator(); +try { + while (it.hasNext()) { + ... + } + if (schema != null) { + ArrowStreamWriter.writeEndOfStream(channel, IpcOption.DEFAULT); + } +} catch (IOException e) { + throw new IllegalStateException("Failed to serialize fragment output as Arrow IPC stream", e); +} finally { try { - while (it.hasNext()) { - if (task != null && task.isCancelled()) { - throw new TaskCancelledException("task cancelled: " + task.getReasonCancelled()); - } - EngineResultBatch batch = it.next(); - VectorSchemaRoot root = batch.getArrowRoot(); - try { - if (schema == null) { - schema = root.getSchema(); - MessageSerializer.serialize(channel, schema); - } - try (ArrowRecordBatch recordBatch = new VectorUnloader(root).getRecordBatch()) { - MessageSerializer.serialize(channel, recordBatch); - } - totalRows += root.getRowCount(); - } finally { - root.close(); - } - } - if (schema != null) { - ArrowStreamWriter.writeEndOfStream(channel, IpcOption.DEFAULT); - } + channel.close(); } catch (IOException e) { - throw new IllegalStateException("Failed to serialize fragment output as Arrow IPC stream", e); + // log or suppress } - return new FragmentExecutionResponse(baos.toByteArray(), totalRows); } Suggestion importance[1-10]: 9 __ Why: The suggestion identifies a resource leak where WriteChannel is not explicitly closed, potentially leaving buffered data unflushed. This could result in incomplete IPC payloads being sent, which is a critical correctness issue. Closing the channel in a finally block ensures proper resource cleanup.	High
	Validate column count and null vectors Verify that batch.getVector(col) does not return null before dereferencing. If the batch schema has fewer columns than declaredSchema, accessing batch.getVector(col) may return null or throw an exception. Add a null check and validate column count consistency. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [260-267] -private static VectorSchemaRoot coerceToDeclaredSchema(VectorSchemaRoot batch, Schema declaredSchema, BufferAllocator alloc) { - if (batch.getSchema().equals(declaredSchema)) { - return batch; +for (int col = 0; col < declaredSchema.getFields().size(); col++) { + if (col >= batch.getFieldVectors().size()) { + throw new IllegalStateException("Batch has fewer columns than declaredSchema: batch=" + batch.getFieldVectors().size() + ", declared=" + declaredSchema.getFields().size()); } - VectorSchemaRoot out = VectorSchemaRoot.create(declaredSchema, alloc); - try { - out.allocateNew(); - int rows = batch.getRowCount(); - for (int col = 0; col < declaredSchema.getFields().size(); col++) { - FieldVector src = batch.getVector(col); - FieldVector dst = out.getVector(col); - ... - } - out.setRowCount(rows); - } catch (RuntimeException e) { - out.close(); - throw e; - } finally { - batch.close(); + FieldVector src = batch.getVector(col); + if (src == null) { + throw new IllegalStateException("Batch vector at column " + col + " is null"); } - return out; + FieldVector dst = out.getVector(col); + ... } Suggestion importance[1-10]: 8 __ Why: The suggestion addresses a critical issue where accessing batch.getVector(col) without validating column count or null checks could cause NullPointerException or IndexOutOfBoundsException. This validation is essential for preventing runtime failures when batch and declared schemas mismatch.	Medium
	Prevent panic in FFI boundary The execute_prepared method panics if no plan was prepared, but this FFI function doesn't validate that condition before calling it. Add an explicit check to return a proper error instead of allowing a panic to cross the FFI boundary. sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs [822-829] pub unsafe extern "C" fn df_execute_local_prepared_plan(session_ptr: i64) -> i64 { let session = &*(session_ptr as *const crate::local_executor::LocalSession); - ... + if session.prepared_plan.is_none() { + return Err("No plan prepared. Call df_prepare_final_plan first.".to_string()); + } + let mgr = get_rt_manager()?; + let _guard = mgr.io_runtime.enter(); let df_stream = session.execute_prepared().map_err(|e| e.to_string())?; Suggestion importance[1-10]: 8 __ Why: Valid concern about panics crossing FFI boundaries. The execute_prepared method at line 190-195 in local_executor.rs explicitly panics if no plan is prepared, and this FFI function should validate that condition before calling it to return a proper error instead.	Medium
	Add bounds checking for index access Add bounds checking before accessing inputTypes by index in the RexShuttle visitor. If ref.getIndex() exceeds the size of inputTypes, the code will throw an IndexOutOfBoundsException at runtime. Validate that the index is within bounds and handle the error case explicitly. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/AggregateDecompositionResolver.java [208-214] -private static RelNode replaceFirstWithRefRebinding(RelNode node, RelNode target, RelNode replacement) { - if (node == target) return replacement; - java.util.List<RelNode> newInputs = new java.util.ArrayList<>(); - boolean changed = false; - for (RelNode input : node.getInputs()) { - RelNode newInput; - if (input == target) { - newInput = replacement; - if (node instanceof Project proj) { - RexBuilder rexBuilder = node.getCluster().getRexBuilder(); - java.util.List<RelDataType> inputTypes = new java.util.ArrayList<>(); - for (var f : replacement.getRowType().getFieldList()) { - inputTypes.add(f.getType()); - } - RexShuttle rebind = new RexShuttle() { - @Override - public RexNode visitInputRef(RexInputRef ref) { - RelDataType actual = inputTypes.get(ref.getIndex()); - if (ref.getType().equals(actual)) return ref; - return new RexInputRef(ref.getIndex(), actual); - } - }; - java.util.List<RexNode> rebound = new java.util.ArrayList<>(proj.getProjects().size()); - for (int i = 0; i < proj.getProjects().size(); i++) { - RexNode expr = proj.getProjects().get(i).accept(rebind); - RelDataType targetType = proj.getRowType().getFieldList().get(i).getType(); - if (!expr.getType().equals(targetType)) { - expr = rexBuilder.makeCast(targetType, expr); - } - rebound.add(expr); - } - return proj.copy(proj.getTraitSet(), replacement, rebound, proj.getRowType()); - } - } else { - newInput = replaceFirstWithRefRebinding(input, target, replacement); +RexShuttle rebind = new RexShuttle() { + @Override + public RexNode visitInputRef(RexInputRef ref) { + int idx = ref.getIndex(); + if (idx < 0 || idx >= inputTypes.size()) { + throw new IllegalStateException("RexInputRef index " + idx + " out of bounds for inputTypes size " + inputTypes.size()); } - newInputs.add(newInput); - if (newInput != input) changed = true; + RelDataType actual = inputTypes.get(idx); + if (ref.getType().equals(actual)) return ref; + return new RexInputRef(idx, actual); } - return changed ? node.copy(node.getTraitSet(), newInputs) : node; -} +}; Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential IndexOutOfBoundsException when accessing inputTypes.get(ref.getIndex()) without bounds checking. Adding validation prevents runtime crashes when the index is out of bounds, improving robustness.	Medium
	Validate array bounds before indexing Direct array indexing without bounds checking can panic if children() returns an empty slice. Validate that children() has at least one element before accessing index 0 to prevent potential runtime panics. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [76-81] AggregateMode::Final => { // Current node is Partial; skip it, return its child // (the Final above will keep itself) - let child = agg.children()[0]; - force_aggregate_mode(Arc::clone(child), target) + let children = agg.children(); + if children.is_empty() { + return Err(datafusion_common::DataFusionError::Internal( + "AggregateExec has no children".to_string() + )); + } + force_aggregate_mode(Arc::clone(children[0]), target) } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential panic from unchecked array indexing at line 79. Adding bounds validation would improve robustness, though AggregateExec typically has at least one child in practice.	Medium
General	Validate plan alternative count consistency Validate that parentStage.getPlanAlternatives() and childStage.getPlanAlternatives() have matching sizes before indexing into rewriteResults. The current code uses Math.min(i, rewriteResults.size() - 1) as a fallback, but if the parent has more alternatives than the child, this silently reuses the last child's rewrite result for multiple parent plans, which may not be semantically correct. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/AggregateDecompositionResolver.java [96-113] -private static void resolvePartialFinalPair(Stage parentStage, Stage childStage) { - List<StagePlan> resolvedChildPlans = new ArrayList<>(childStage.getPlanAlternatives().size()); - List<StagePlan> resolvedParentPlans = new ArrayList<>(parentStage.getPlanAlternatives().size()); - List<RewriteResult> rewriteResults = new ArrayList<>(); - - for (StagePlan childPlan : childStage.getPlanAlternatives()) { - OpenSearchAggregate partialAgg = findTopAggregate(childPlan.resolvedFragment(), AggregateMode.PARTIAL); - if (partialAgg == null) { - resolvedChildPlans.add(childPlan); - rewriteResults.add(null); - continue; - } - RewriteResult result = rewriteDecomposed(partialAgg); - rewriteResults.add(result); - RelNode newChildFragment = replaceFirst(childPlan.resolvedFragment(), partialAgg, result.newPartial(partialAgg)); - resolvedChildPlans.add(new StagePlan(newChildFragment, childPlan.backendId())); - } - - boolean anyChildRewritten = rewriteResults.stream().anyMatch(r -> r != null); - if (!anyChildRewritten) return; - - childStage.setPlanAlternatives(resolvedChildPlans); - - for (int i = 0; i < parentStage.getPlanAlternatives().size(); i++) { - StagePlan parentPlan = parentStage.getPlanAlternatives().get(i); - RewriteResult result = rewriteResults.get(Math.min(i, rewriteResults.size() - 1)); - ... - } - parentStage.setPlanAlternatives(resolvedParentPlans); +if (parentStage.getPlanAlternatives().size() != childStage.getPlanAlternatives().size()) { + throw new IllegalStateException("Parent and child stages must have the same number of plan alternatives for decomposition pairing"); +} +for (int i = 0; i < parentStage.getPlanAlternatives().size(); i++) { + StagePlan parentPlan = parentStage.getPlanAlternatives().get(i); + RewriteResult result = rewriteResults.get(i); + ... } Suggestion importance[1-10]: 6 __ Why: The suggestion points out that using Math.min(i, rewriteResults.size() - 1) silently reuses the last child's rewrite result when parent has more alternatives than child. While the current fallback prevents crashes, validating size equality upfront would make the assumption explicit and catch potential logic errors earlier.	Low
	Handle unexpected multi-child repartition nodes The function silently returns None when encountering repartition/coalesce nodes with multiple children, which may hide unexpected plan structures. Consider logging or returning an error for multi-child cases to aid debugging. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [113-123] fn find_partial_input(plan: Arc<dyn ExecutionPlan>) -> Option<Arc<dyn ExecutionPlan>> { ... if plan.as_any().downcast_ref::<RepartitionExec>().is_some() || plan .as_any() .downcast_ref::<CoalescePartitionsExec>() .is_some() { let children = plan.children(); if children.len() == 1 { return find_partial_input(Arc::clone(children[0])); } + // Unexpected: repartition/coalesce with != 1 child + return None; } None } Suggestion importance[1-10]: 3 __ Why: The suggestion points out that multi-child cases return None silently. However, the 'improved_code' is identical to the 'existing_code' (just adds a comment), so this offers minimal practical improvement. The current behavior is acceptable for the function's purpose.	Low

Suggestions up to commit e7cf5d1

Category	Suggestion	Impact
Possible issue	Prevent double-close on fast-path return The batch input is closed unconditionally in the finally block, even when the fast-path returns early (batch.getSchema().equals(declaredSchema)). This causes a double-close when the caller tries to use the returned batch. Move batch.close() inside the try block after the coercion loop completes, so the early-return path leaves batch open for the caller. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [255-297] private static VectorSchemaRoot coerceToDeclaredSchema(VectorSchemaRoot batch, Schema declaredSchema, BufferAllocator alloc) { if (batch.getSchema().equals(declaredSchema)) { return batch; } VectorSchemaRoot out = VectorSchemaRoot.create(declaredSchema, alloc); try { out.allocateNew(); int rows = batch.getRowCount(); for (int col = 0; col < declaredSchema.getFields().size(); col++) { FieldVector src = batch.getVector(col); FieldVector dst = out.getVector(col); if (src.getField().getType().equals(dst.getField().getType())) { src.makeTransferPair(dst).transfer(); continue; } ... } out.setRowCount(rows); + batch.close(); } catch (RuntimeException e) { out.close(); + batch.close(); throw e; - } finally { - batch.close(); } return out; } Suggestion importance[1-10]: 9 __ Why: The batch is closed unconditionally in the finally block even when the fast-path returns early at line 256. This causes a double-close when the caller tries to use the returned batch. Moving batch.close() inside the try block after the coercion loop completes prevents this critical resource management bug.	High
	Validate prepared plan before execution The function calls execute_prepared() which panics if no plan was prepared, but this panic will cross the FFI boundary unsafely. Add explicit validation that prepared_plan is Some before calling execute_prepared() and return a proper error code instead of allowing a panic. sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs [822-829] pub unsafe extern "C" fn df_execute_local_prepared_plan(session_ptr: i64) -> i64 { let session = &*(session_ptr as *const crate::local_executor::LocalSession); + if session.prepared_plan.is_none() { + return Err("No plan prepared. Call df_prepare_final_plan first.".to_string()); + } ... let df_stream = session.execute_prepared().map_err(|e| e.to_string())?; Suggestion importance[1-10]: 9 __ Why: The function calls execute_prepared() which panics if no plan was prepared. Since this crosses the FFI boundary, allowing a panic is unsafe and could crash the JVM. Adding validation prevents this critical issue.	High
	Add bounds check for children access The code directly accesses agg.children()[0] without bounds checking. If the aggregate node has no children, this will panic at runtime. Add a bounds check or use get(0) with proper error handling to prevent potential panics. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [76-81] AggregateMode::Final => { // Current node is Partial; skip it, return its child // (the Final above will keep itself) - let child = agg.children()[0]; - force_aggregate_mode(Arc::clone(child), target) + let children = agg.children(); + if children.is_empty() { + return Err(datafusion_common::DataFusionError::Internal( + "AggregateExec has no children".to_string() + )); + } + force_aggregate_mode(Arc::clone(children[0]), target) } Suggestion importance[1-10]: 8 __ Why: The code accesses agg.children()[0] without bounds checking, which could panic if the aggregate node has no children. This is a potential runtime crash that should be prevented with proper error handling.	Medium
	Return error instead of panicking Using expect() causes a panic when prepared_plan is None, which is problematic for FFI boundaries and error handling. Return a proper DataFusionError instead to allow callers to handle the error gracefully. sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs [190-196] pub fn execute_prepared(&self) -> Result<SendableRecordBatchStream, DataFusionError> { let plan = self .prepared_plan .as_ref() - .expect("execute_prepared called without a prepared plan"); + .ok_or_else(|| DataFusionError::Internal( + "execute_prepared called without a prepared plan".to_string() + ))?; datafusion::physical_plan::execute_stream(Arc::clone(plan), self.ctx.task_ctx()) } Suggestion importance[1-10]: 8 __ Why: Using expect() causes a panic when prepared_plan is None, which is problematic for FFI boundaries. Returning a proper DataFusionError allows callers to handle the error gracefully and prevents crashes.	Medium
General	Close session on constructor failure When preparedState is null, a new DatafusionLocalSession is created but never closed if an exception occurs during the constructor's loop (e.g., ArrowSchemaIpc.toBytes throws). Wrap the loop in a try-catch and close the session on failure to prevent native resource leaks. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/AbstractDatafusionReduceSink.java [103-124] protected AbstractDatafusionReduceSink( ExchangeSinkContext ctx, NativeRuntimeHandle runtimeHandle, DataFusionReduceState preparedState ) { this.ctx = ctx; this.runtimeHandle = runtimeHandle; this.preparedState = preparedState; - this.session = preparedState != null ? preparedState.session() : new DatafusionLocalSession(runtimeHandle.get()); + DatafusionLocalSession localSession = preparedState != null ? preparedState.session() : new DatafusionLocalSession(runtimeHandle.get()); Map<Integer, byte[]> inputs = new LinkedHashMap<>(ctx.childInputs().size()); Map<Integer, Schema> schemas = new LinkedHashMap<>(ctx.childInputs().size()); - for (ExchangeSinkContext.ChildInput child : ctx.childInputs()) { - inputs.put(child.childStageId(), ArrowSchemaIpc.toBytes(child.schema())); - schemas.put(child.childStageId(), child.schema()); + try { + for (ExchangeSinkContext.ChildInput child : ctx.childInputs()) { + inputs.put(child.childStageId(), ArrowSchemaIpc.toBytes(child.schema())); + schemas.put(child.childStageId(), child.schema()); + } + } catch (RuntimeException e) { + if (preparedState == null) { + localSession.close(); + } + throw e; } + this.session = localSession; this.childInputs = inputs; this.childSchemas = schemas; } Suggestion importance[1-10]: 7 __ Why: When preparedState is null, a new DatafusionLocalSession is created at line 115 but never closed if an exception occurs during the loop at lines 117-121. This causes a native resource leak. Wrapping the loop in try-catch and closing the session on failure is a valid resource management improvement.	Medium
	Re-evaluate group-set flag per call The hasEmptyGroup flag is computed once outside the stream but should be re-evaluated per call if the FINAL aggregate's group set changes during rewrite. If a rewrite inadvertently modifies the group set, the flag becomes stale and Calcite's type inference may produce incorrect nullability. Compute hasEmptyGroup inside the lambda using finalAgg.getGroupSet().isEmpty() to ensure correctness. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/AggregateDecompositionResolver.java [144-163] -private static RelNode rewriteParentFragment(RelNode fragment, RelDataType childRowType, int childStageId, RewriteResult result) { - ... - List<AggregateCall> rebuiltFinalCalls = result.newFinalCalls.stream() - .map( - c -> AggregateCall.create( - c.getAggregation(), - c.isDistinct(), - c.isApproximate(), - c.ignoreNulls(), - c.rexList, - c.getArgList(), - c.filterArg, - c.distinctKeys, - c.collation, - hasEmptyGroup, - newFinalInput, - null, - c.name - ) +List<AggregateCall> rebuiltFinalCalls = result.newFinalCalls.stream() + .map( + c -> AggregateCall.create( + c.getAggregation(), + c.isDistinct(), + c.isApproximate(), + c.ignoreNulls(), + c.rexList, + c.getArgList(), + c.filterArg, + c.distinctKeys, + c.collation, + finalAgg.getGroupSet().isEmpty(), + newFinalInput, + null, + c.name ) - .toList(); - ... -} + ) + .toList(); Suggestion importance[1-10]: 3 __ Why: The hasEmptyGroup flag is computed once at line 144 and reused for all calls. While the suggestion to re-evaluate per call is technically more defensive, the current code is correct because finalAgg.getGroupSet() is immutable during the stream operation. The suggestion adds redundant computation without fixing an actual bug.	Low

Suggestions up to commit 832e0b8

Category	Suggestion	Impact
Possible issue	Prevent resource leak on coercion failure The coerceToDeclaredSchema call can throw a RuntimeException after the early-return check but before the try-finally block that closes batch. If the coercion fails, the original batch reference is lost without being closed, causing a resource leak. Move the coercion inside the try-finally or ensure batch is closed on exception. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java [198-208] private void feedToSender(DatafusionPartitionSender sender, VectorSchemaRoot batch, Schema declaredSchema) { if (closed) { batch.close(); return; } BufferAllocator alloc = ctx.allocator(); - batch = coerceToDeclaredSchema(batch, declaredSchema, alloc); - ... + ArrowArray array = ArrowArray.allocateNew(alloc); + ArrowSchema arrowSchema = ArrowSchema.allocateNew(alloc); + try { + batch = coerceToDeclaredSchema(batch, declaredSchema, alloc); + Data.exportVectorSchemaRoot(alloc, batch, dictionaryProvider, array, arrowSchema); + ... + } finally { + batch.close(); + array.close(); + arrowSchema.close(); + } } Suggestion importance[1-10]: 8 __ Why: The coerceToDeclaredSchema call at line 208 can throw before the try-finally block that closes batch (lines 209-232). If coercion fails, the original batch is leaked. Moving the coercion inside the try or ensuring batch is closed on exception is critical for resource safety.	Medium
	Prevent panic across FFI boundary The function calls execute_prepared() which panics if no plan was prepared, but this panic will cross the FFI boundary unsafely. Add explicit validation to return an error code instead of allowing a panic to propagate through the C boundary. sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs [822-829] pub unsafe extern "C" fn df_execute_local_prepared_plan(session_ptr: i64) -> i64 { let session = &*(session_ptr as *const crate::local_executor::LocalSession); + if session.prepared_plan.is_none() { + return ffm_error("execute_local_prepared_plan called without a prepared plan"); + } let mgr = get_rt_manager()?; // DataFusion's execute_stream is sync, but kicks off RepartitionExec / stream // channels that require a Tokio reactor. Enter the IO runtime's context so those // operators can register with the reactor. let _guard = mgr.io_runtime.enter(); let df_stream = session.execute_prepared().map_err(|e| e.to_string())?; Suggestion importance[1-10]: 8 __ Why: This is a critical safety issue. Panics crossing FFI boundaries cause undefined behavior. The suggestion correctly identifies that execute_prepared() can panic and proposes checking prepared_plan before calling it, preventing unsafe behavior at the C boundary.	Medium
	Validate children before indexing Direct array indexing without bounds checking can panic if the aggregate has no children. Validate that children() is non-empty before accessing index 0 to prevent potential runtime panics. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [76-81] AggregateMode::Final => { // Current node is Partial; skip it, return its child // (the Final above will keep itself) - let child = agg.children()[0]; + let children = agg.children(); + if children.is_empty() { + return Err(datafusion_common::DataFusionError::Internal( + "AggregateExec has no children".to_string() + )); + } + let child = children[0]; force_aggregate_mode(Arc::clone(child), target) } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential panic from unchecked array indexing. However, in the context of AggregateExec, having no children would be a structural invariant violation, making this an edge case. The validation adds defensive programming but may not be critical in practice.	Medium
	Replace panic with error result Using expect() causes a panic when prepared_plan is None. Replace with proper error handling that returns a DataFusionError to maintain consistent error propagation and avoid unexpected panics. sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs [190-196] pub fn execute_prepared(&self) -> Result<SendableRecordBatchStream, DataFusionError> { let plan = self .prepared_plan .as_ref() - .expect("execute_prepared called without a prepared plan"); + .ok_or_else(|| DataFusionError::Internal( + "execute_prepared called without a prepared plan".to_string() + ))?; datafusion::physical_plan::execute_stream(Arc::clone(plan), self.ctx.task_ctx()) } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies that using expect() causes panics and proposes returning a proper DataFusionError instead. This improves error handling consistency, though the panic is documented and the caller should ensure a plan is prepared first. ...

[github-actions]

Persistent review updated to latest commit b696fc0

[sandeshkr419]

[outdated]

@expani Thanks for the detailed review. Here's what I changed structurally based on your feedback:

• AVG no longer has a custom decomposition. It's handled by Calcite's AggregateReduceFunctionsRule in the pre-marking phase so backends only ever see SUM and COUNT. Similarly, HLL no longer has a decomposition either; the intermediate Arrow type (Binary) is now declared directly on AggregateCapability.intermediateArrowType, which LocalStageScheduler reads, no more ad-hoc isApproximate() check in the scheduler.

• The COUNT(DISTINCT) → APPROX_COUNT_DISTINCT rewrite moved from OpenSearchAggregateRule into BackendPlanAdapter via a new aggregateCallAdapters() hook on BackendCapabilityProvider, keeping backend-specific logic out of the generic planner rule. stripAnnotations was also simplified to only strip annotation wrappers.

• On the Rust side, the StripFinalAggregateRule and StripPartialAggregateRule physical optimizer rules are replaced by FragmentConvertor-driven mode forcing. DataFusionFragmentConvertor now prepends a 1-byte mode prefix to plan bytes (attachPartialAggOnTop → partial, convertFinalAggFragment → final), and the Rust side calls force_aggregate_mode() directly after create_physical_plan, no plan inspection. The FragmentConvertor API is the authoritative signal for execution mode.

• One open item: DataFusion's Substrait consumer ignores aggregation_phase, which is why force_aggregate_mode exists. This could be a potential upstream contribution, if DataFusion adds aggregation_phase support, force_aggregate_mode can be removed entirely.

[github-actions]

Persistent review updated to latest commit fdcf13a

[github-actions]

✅ Gradle check result for fdcf13a: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.41%. Comparing base (36809cc) to head (ee0a3d7).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21457      +/-   ##
============================================
- Coverage     73.50%   73.41%   -0.09%     
+ Complexity    74644    74563      -81     
============================================
  Files          5980     5980              
  Lines        338777   338777              
  Branches      48848    48848              
============================================
- Hits         249011   248722     -289     
- Misses        69946    70252     +306     
+ Partials      19820    19803      -17     

☔ 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 f5d4771

[github-actions]

✅ Gradle check result for f5d4771: SUCCESS

[github-actions]

Persistent review updated to latest commit b89a7c8

[github-actions]

✅ Gradle check result for b89a7c8: SUCCESS

[github-actions]

Persistent review updated to latest commit c76d34d

[github-actions]

❌ Gradle check result for c76d34d: 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 ba512f3

[github-actions]

✅ Gradle check result for ba512f3: SUCCESS

[github-actions]

✅ Gradle check result for 614a8c5: SUCCESS

[github-actions]

✅ Gradle check result for e7cf5d1: SUCCESS

[github-actions]

✅ Gradle check result for e648ee7: SUCCESS

[github-actions]

✅ Gradle check result for ee0a3d7: SUCCESS