[analytics-engine] Per-shard bucket oversampling for TopK aggregation queries

[sandeshkr419]

Description:

Adds a post-CBO rewriter that inserts a per-partition Sort+Limit between PARTIAL aggregate and ExchangeReducer, so each shard only ships top-N×factor groups instead of all groups.

• Index setting index.analytics.shard_bucket_oversampling_factor (default 0, opt-in)
• OpenSearchTopKRewriter finds bottommost Sort with collation above grouped FINAL→ER→PARTIAL and inserts per-shard Sort
• reduce_eval scalar UDF for evaluating opaque aggregate state (approx_distinct HLL) to a sortable scalar — generic pattern for future approximate aggregates
• OpenSearchSort.perPartition flag skips the singleton cost gate

Tests: TopKRewriterPlanShapeTests (2), ShardBucketOversamplingIT (7 on ClickBench multi-shard), reduce_eval, Rust unit test.

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 e770031)

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 Integer Overflow When coordLimit is close to Integer.MAX_VALUE and factor is greater than 1.0, the calculation (long) Math.ceil(coordLimit * factor) + coordLimit can exceed Integer.MAX_VALUE before the check on line 60. This causes shardSize to be cast to an int that wraps around to a negative value, bypassing the guard and producing an invalid negative limit literal. long coordLimit = (sort.fetch instanceof RexLiteral lit) ? RexLiteral.intValue(lit) : 10_000L; long shardSize = (long) Math.ceil(coordLimit * factor) + coordLimit; if (shardSize > Integer.MAX_VALUE) return Optional.empty(); Possible Issue replaceInTree performs a shallow object-identity check (root == oldNode) which may fail if the tree contains wrapped or copied nodes that are logically equivalent but not the same instance. If oldNode is never found, the function silently returns the original tree, and the rewrite does not occur. This can happen if the CBO or another pass creates a new node instance for finalAgg after findSortAboveFinal captures it. private static RelNode replaceInTree(RelNode root, RelNode oldNode, RelNode newNode) { if (root == oldNode) return newNode; List<RelNode> children = root.getInputs(); boolean changed = false; RelNode[] newChildren = new RelNode[children.size()]; for (int i = 0; i < children.size(); i++) { newChildren[i] = replaceInTree(children.get(i), oldNode, newNode); if (newChildren[i] != children.get(i)) changed = true; } if (!changed) return root; return root.copy(root.getTraitSet(), List.of(newChildren)); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to e770031

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Prevent arithmetic overflow in calculation The calculation Math.ceil(coordLimit * factor) + coordLimit can overflow when coordLimit is large. For example, if coordLimit is near Long.MAX_VALUE, multiplying by factor will overflow before the Integer.MAX_VALUE check. Add overflow protection before the multiplication. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [58-60] long coordLimit = (sort.fetch instanceof RexLiteral lit) ? RexLiteral.intValue(lit) : 10_000L; +if (coordLimit > Integer.MAX_VALUE / (factor + 1)) return Optional.empty(); long shardSize = (long) Math.ceil(coordLimit * factor) + coordLimit; if (shardSize > Integer.MAX_VALUE) return Optional.empty(); Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential overflow risk when coordLimit is large and multiplied by factor. However, the proposed check coordLimit > Integer.MAX_VALUE / (factor + 1) may not fully prevent overflow in all edge cases (e.g., when factor is very large). The suggestion is valid and improves safety, but the implementation could be more robust.	Medium
	Add null-safety for cluster state If context.getClusterState() or context.getClusterState().metadata() returns null, a NullPointerException will be thrown before the clusterSettings == null check. Add null-safety checks for the entire chain to prevent crashes. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [119-124] private static double resolveOversamplingFactor(PlannerContext context) { // TODO: Move to per-index setting once index-pattern/alias resolution is handled. + if (context.getClusterState() == null || context.getClusterState().metadata() == null) return 0.0; Settings clusterSettings = context.getClusterState().metadata().settings(); if (clusterSettings == null) return 0.0; return AnalyticsApproximationSettings.SHARD_BUCKET_OVERSAMPLING_FACTOR.get(clusterSettings); } Suggestion importance[1-10]: 6 __ Why: The suggestion correctly identifies a potential NullPointerException if getClusterState() or metadata() returns null. Adding null-safety checks improves robustness. However, this is a defensive programming improvement rather than a critical bug fix, as the likelihood depends on the framework's guarantees about these methods.	Low
General	Avoid creating accumulator per iteration Creating a new accumulator for each row is inefficient and may cause performance degradation for large batches. Consider reusing a single accumulator instance and resetting its state between iterations, or batch-process multiple states if the API supports it. sandbox/plugins/analytics-backend-datafusion/rust/src/udf/reduce_eval.rs [72-89] +let mut acc = approx_distinct_udaf().accumulator(AccumulatorArgs { + ... +})?; for i in 0..binary.len() { if binary.is_null(i) { results.push(0u64); continue; } - let mut acc = approx_distinct_udaf().accumulator(AccumulatorArgs { - ... - })?; let state_array: ArrayRef = Arc::new(BinaryArray::from(vec![binary.value(i)])); acc.merge_batch(&[state_array])?; + ... +} Suggestion importance[1-10]: 4 __ Why: The suggestion identifies a performance concern about creating a new accumulator per row. However, the improved code doesn't show how to reset the accumulator state between iterations, which is necessary for correctness. Without proper state reset, reusing the accumulator would produce incorrect results. The suggestion is conceptually valid but lacks implementation details.	Low

---

Previous suggestions

Suggestions up to commit ba0d773

Category	Suggestion	Impact
Possible issue	Prevent arithmetic overflow in limit calculation The calculation Math.ceil(coordLimit * factor) + coordLimit can overflow when coordLimit is large (e.g., close to Long.MAX_VALUE). This overflow would produce incorrect shardSize values before the Integer.MAX_VALUE check. Add overflow detection before the ceiling operation to prevent silent arithmetic errors. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [58-60] long coordLimit = (sort.fetch instanceof RexLiteral lit) ? RexLiteral.intValue(lit) : 10_000L; +if (coordLimit > (Long.MAX_VALUE / (factor + 1))) return Optional.empty(); long shardSize = (long) Math.ceil(coordLimit * factor) + coordLimit; if (shardSize > Integer.MAX_VALUE) return Optional.empty(); Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential overflow issue in the calculation Math.ceil(coordLimit * factor) + coordLimit. However, the overflow check coordLimit > (Long.MAX_VALUE / (factor + 1)) may not be entirely accurate for all edge cases, and the practical impact is limited since coordLimit is derived from RexLiteral.intValue() which returns an int, making overflow less likely in typical scenarios.	Medium
General	Handle null aggregate states correctly Returning 0u64 for null aggregate states may produce incorrect sort ordering in TopK queries, as null states should be treated distinctly from zero-count aggregates. Consider using a sentinel value or propagating nulls to maintain semantic correctness in the result set. sandbox/plugins/analytics-backend-datafusion/rust/src/udf/reduce_eval.rs [63-76] let binary = state_col.as_any().downcast_ref::<BinaryArray>() .ok_or_else(|| DataFusionError::Execution("reduce_eval(approx_distinct): expected Binary state".into()))?; ... +let mut null_mask = Vec::with_capacity(binary.len()); for i in 0..binary.len() { if binary.is_null(i) { results.push(0u64); + null_mask.push(false); continue; } + null_mask.push(true); Suggestion importance[1-10]: 6 __ Why: The suggestion raises a valid concern about returning 0u64 for null states potentially affecting sort ordering. However, the improved code only adds a null_mask vector without actually using it to propagate nulls in the result, making the suggested fix incomplete. The semantic correctness depends on whether 0u64 is an acceptable representation for null states in the TopK context.	Low
	Add logging for missing cluster settings When clusterSettings is null, returning 0.0 silently disables TopK optimization without logging. This makes debugging configuration issues difficult. Add logging to indicate when the feature is disabled due to missing settings, helping operators diagnose unexpected query behavior. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [119-124] private static double resolveOversamplingFactor(PlannerContext context) { // TODO: Move to per-index setting once index-pattern/alias resolution is handled. Settings clusterSettings = context.getClusterState().metadata().settings(); - if (clusterSettings == null) return 0.0; + if (clusterSettings == null) { + LOGGER.debug("Cluster settings unavailable, TopK oversampling disabled"); + return 0.0; + } return AnalyticsApproximationSettings.SHARD_BUCKET_OVERSAMPLING_FACTOR.get(clusterSettings); } Suggestion importance[1-10]: 4 __ Why: Adding debug logging when clusterSettings is null would improve observability and help with troubleshooting. However, this is a minor enhancement with limited impact, as the null check already handles the case gracefully by disabling the optimization, and debug-level logging may not always be enabled in production environments.	Low

Suggestions up to commit e27e229

Category	Suggestion	Impact
General	Optimize accumulator creation in loop Creating a new accumulator inside the loop for each row is inefficient and may cause performance issues with large datasets. Consider reusing a single accumulator instance or restructuring the logic to batch process states more efficiently. sandbox/plugins/analytics-backend-datafusion/rust/src/udf/reduce_eval.rs [72-87] +let field: Arc<Field> = Arc::new(Field::new("x", DataType::Int64, true)); +let schema = Arc::new(datafusion::arrow::datatypes::Schema::new(vec![field.as_ref().clone()])); +let expr: Arc<dyn datafusion::physical_plan::PhysicalExpr> = Arc::new(Column::new("x", 0)); +let ret_field: Arc<Field> = Arc::new(Field::new("r", DataType::UInt64, true)); + let mut results = Vec::with_capacity(binary.len()); for i in 0..binary.len() { if binary.is_null(i) { results.push(0u64); continue; } let mut acc = approx_distinct_udaf().accumulator(AccumulatorArgs { - ... + return_field: ret_field.clone(), + schema: &schema, + ignore_nulls: false, + order_bys: &[], + name: "x", + is_distinct: false, + exprs: &[expr.clone()], + expr_fields: &[field.clone()], + is_reversed: false, })?; ... } Suggestion importance[1-10]: 6 __ Why: Creating a new accumulator for each row is inefficient. However, the improved_code shown is identical to the existing_code and doesn't demonstrate the optimization. The suggestion correctly identifies a performance issue but fails to provide a meaningful solution, limiting its practical value.	Low
Possible issue	Handle potential NumberFormatException safely The RexLiteral.intValue() method can throw NumberFormatException if the literal value doesn't fit in an int. Add validation to ensure the literal is within valid bounds before extracting its value, or handle the potential exception. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [58] -long coordLimit = (sort.fetch instanceof RexLiteral lit) ? RexLiteral.intValue(lit) : 10_000L; +long coordLimit = 10_000L; +if (sort.fetch instanceof RexLiteral lit) { + try { + coordLimit = RexLiteral.intValue(lit); + } catch (NumberFormatException e) { + return Optional.empty(); + } +} Suggestion importance[1-10]: 4 __ Why: While RexLiteral.intValue() could theoretically throw NumberFormatException, this is unlikely in a well-formed query plan. The suggestion adds defensive error handling, but the improvement is marginal since the planner should ensure valid literals. The fallback to Optional.empty() is reasonable but not critical.	Low
	Prevent integer overflow in calculation The calculation can overflow when coordLimit is large. For example, if coordLimit is close to Integer.MAX_VALUE and factor is 2.0, the multiplication will overflow before the cast to long. Cast coordLimit to double before multiplication to prevent integer overflow. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [59-60] -long shardSize = (long) Math.ceil(coordLimit * factor) + coordLimit; +long shardSize = (long) Math.ceil((double) coordLimit * factor) + coordLimit; if (shardSize > Integer.MAX_VALUE) return Optional.empty(); Suggestion importance[1-10]: 3 __ Why: The concern about overflow is valid, but coordLimit is already a long (extracted from RexLiteral.intValue() which returns an int, then assigned to long). The multiplication coordLimit * factor is performed on a long and double, so no integer overflow occurs before the cast. The suggestion's cast to double is redundant and offers minimal improvement.	Low

Suggestions up to commit 1f61124

Category	Suggestion	Impact
Possible issue	Prevent potential arithmetic overflow The calculation may overflow when coordLimit is large. Consider checking for overflow before casting to int, or use Math.addExact() and Math.multiplyExact() to detect overflow conditions explicitly. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [59-60] -long shardSize = (long) Math.ceil(coordLimit * factor) + coordLimit; -if (shardSize > Integer.MAX_VALUE) return Optional.empty(); +try { + long shardSize = Math.addExact((long) Math.ceil(coordLimit * factor), coordLimit); + if (shardSize > Integer.MAX_VALUE) return Optional.empty(); +} catch (ArithmeticException e) { + return Optional.empty(); +} Suggestion importance[1-10]: 6 __ Why: The suggestion correctly identifies a potential overflow scenario when coordLimit is large. However, the risk is mitigated by the subsequent check shardSize > Integer.MAX_VALUE, which already handles the overflow case by returning empty. The suggested Math.addExact() approach adds explicit overflow detection but provides marginal improvement given the existing safeguard.	Low
	Handle empty qualified name list The method assumes a single table scan exists in the plan. For queries with multiple table scans or joins, this may return an incorrect or arbitrary index's setting. Consider validating that only one scan exists or handling multiple scans explicitly. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [115-123] private static double resolveOversamplingFactor(OpenSearchAggregate partial, PlannerContext context) { OpenSearchTableScan scan = findScan(partial); if (scan == null) return 0.0; - String indexName = scan.getTable().getQualifiedName().get(scan.getTable().getQualifiedName().size() - 1); + List<String> qualifiedName = scan.getTable().getQualifiedName(); + if (qualifiedName.isEmpty()) return 0.0; + String indexName = qualifiedName.get(qualifiedName.size() - 1); IndexMetadata meta = context.getClusterState().metadata().index(indexName); if (meta == null) return 0.0; if (meta.getSettings() == null) return 0.0; return AnalyticsApproximationSettings.INDEX_ANALYTICS_SHARD_BUCKET_OVERSAMPLING_FACTOR.get(meta.getSettings()); } Suggestion importance[1-10]: 5 __ Why: The suggestion adds a safety check for an empty qualifiedName list before accessing the last element. While this prevents a potential IndexOutOfBoundsException, the scenario is unlikely in practice as table scans typically have valid qualified names. The improvement is defensive but addresses a minor edge case.	Low

Suggestions up to commit 88fb3bb

Category	Suggestion	Impact
General	Optimize memory allocation for sparse data The capacity allocation Vec::with_capacity(arr.len()) may over-allocate when many rows are null. Consider tracking non-null count first or using a smaller initial capacity to reduce memory waste for sparse data. sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/state_shipping.rs [310-328] fn merge_state_binary(&mut self, binary_array: &ArrayRef) -> Result<()> { let arr = binary_array .as_any() .downcast_ref::<BinaryArray>() .ok_or_else(|| { DataFusionError::Execution( "StateShippingAccumulator: expected BinaryArray for state input".to_string(), ) })?; let num_state_cols = self.state_schema.fields().len(); + let non_null_count = arr.len() - arr.null_count(); let mut col_values: Vec<Vec<ScalarValue>> = - vec![Vec::with_capacity(arr.len()); num_state_cols]; + vec![Vec::with_capacity(non_null_count); num_state_cols]; for row_idx in 0..arr.len() { if arr.is_null(row_idx) { continue; } let bytes = arr.value(row_idx); let row_state = decode_state_from_ipc(bytes, &self.state_schema)?; Suggestion importance[1-10]: 6 __ Why: Valid optimization for sparse data scenarios. Using arr.null_count() to compute non_null_count reduces memory waste when many rows are null. The improvement is straightforward and beneficial for memory efficiency in sparse datasets.	Low
	Add safety limit to name generation The uniqueName method has an unbounded loop that could theoretically run indefinitely if the naming collision space is exhausted. Add a safety limit to prevent infinite loops in pathological cases where the namespace is saturated. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/OpenSearchSortExpressionRewriter.java [97-105] private static String uniqueName(String desired, Set<String> used) { if (used.add(desired)) return desired; int n = 1; - while (true) { + int maxAttempts = 10000; + while (n < maxAttempts) { String candidate = desired + "_" + n; if (used.add(candidate)) return candidate; n++; } + throw new IllegalStateException("Failed to generate unique name after " + maxAttempts + " attempts for: " + desired); } Suggestion importance[1-10]: 4 __ Why: The suggestion correctly identifies that the uniqueName method has an unbounded loop. While the pathological case (10,000+ naming collisions) is unlikely in practice, adding a safety limit is a defensive programming practice that prevents infinite loops in edge cases. The suggested limit of 10,000 is reasonable.	Low
	Cache plan traversal result The function traverses the entire plan tree for every call, which can be expensive for large plans. Consider caching the result or memoizing the computation if this function is called multiple times on the same plan during query execution. sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/state_shipping.rs [61-104] +// Consider adding a cached field to store the result after first computation, +// or pass the result through the execution context to avoid repeated traversals. pub fn substrait_has_partial_phase(plan: &substrait::proto::Plan) -> bool { use substrait::proto::plan_rel::RelType as PlanRelType; let phase_partial = substrait::proto::AggregationPhase::InitialToIntermediate as i32; fn rel_has_phase(rel: &substrait::proto::Rel, target: i32) -> bool { ... } for plan_rel in &plan.relations { let Some(rt) = plan_rel.rel_type.as_ref() else { continue }; let rel: &substrait::proto::Rel = match rt { PlanRelType::Rel(r) => r, PlanRelType::Root(root) => match root.input.as_ref() { Some(r) => r, None => continue, }, }; if rel_has_phase(rel, phase_partial) { return true; } } false } Suggestion importance[1-10]: 4 __ Why: The suggestion to cache the result is valid for performance optimization, but the impact is likely minimal since substrait_has_partial_phase is called once per query execution in the current codebase. The suggestion is correct but offers marginal improvement.	Low
	Add bounds check before indexing The code assumes new_children[0] exists without checking if new_children is non-empty. Although the outer condition checks new_children.is_empty(), add an explicit bounds check or assertion before indexing to prevent potential panics if the logic changes. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [110-130] } else if let Some(proj) = plan.as_any().downcast_ref::<datafusion::physical_plan::projection::ProjectionExec>() { + if new_children.is_empty() { + return exec_err!("ProjectionExec requires at least one child"); + } let exprs: Vec<(Arc<dyn datafusion::physical_plan::PhysicalExpr>, String)> = proj .expr() .iter() .map(|pe| (Arc::clone(&pe.expr), pe.alias.clone())) .collect(); Ok(Arc::new( datafusion::physical_plan::projection::ProjectionExec::try_new( exprs, Arc::clone(&new_children[0]), )?, )) Suggestion importance[1-10]: 3 __ Why: The suggestion adds defensive programming, but the outer condition at line 108 already checks new_children.is_empty() and returns early if true. The ProjectionExec branch is only reached when new_children is non-empty, making the additional check redundant. Minor defensive value only.	Low
Possible issue	Validate uniform merge strategy in aggregates The anyNativeMerge check determines whether to use SHARD_MERGE or FINAL mode for the shard-side aggregate. However, if the aggregate contains a mix of engine-native-merge and non-engine-native aggregates, the single shardMode applies to all calls uniformly. This may produce incorrect semantics for the non-native calls when SHARD_MERGE is chosen. Consider splitting the aggregate or validating that all calls share the same merge strategy. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java [110-117] ShardBucketHint hint = finalAgg.getShardBucketHint(); if (hint != null) { RelNode erChild = newFinalInput.getInputs().get(0); + boolean allNativeMerge = finalAgg.getAggCallList().stream().allMatch(c -> { + AggregateFunction fn = AggregateFunction.fromSqlAggFunction(c.getAggregation()); + return fn != null && fn.isEngineNativeMerge(); + }); boolean anyNativeMerge = finalAgg.getAggCallList().stream().anyMatch(c -> { AggregateFunction fn = AggregateFunction.fromSqlAggFunction(c.getAggregation()); return fn != null && fn.isEngineNativeMerge(); }); + if (anyNativeMerge && !allNativeMerge) { + throw new IllegalStateException("Mixed engine-native and non-native aggregates in shard-bucket hint path"); + } AggregateMode shardMode = anyNativeMerge ? AggregateMode.SHARD_MERGE : AggregateMode.FINAL; Suggestion importance[1-10]: 5 __ Why: The suggestion raises a valid concern about mixed engine-native and non-native aggregates in the same OpenSearchAggregate. The current code applies a single shardMode to all calls, which could produce incorrect semantics if the aggregate contains both types. However, the PR's design assumes that the shard-bucket rule only fires on aggregates where all calls share the same merge strategy (or the non-native calls are pass-through scalars). Adding an explicit validation would improve robustness.	Low
	Prevent integer overflow in shard size The overflow check shardSize > Integer.MAX_VALUE occurs after casting to long, but ShardBucketHint constructor expects an int. If shardSize exceeds Integer.MAX_VALUE, the subsequent cast to int in the hint constructor will silently truncate, producing incorrect shard-side limits. Move the overflow check before the hint is created and handle the edge case explicitly. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateShardBucketRule.java [101-103] long shardSize = (long) Math.ceil(Math.max(coordLimit, DEFAULT_COORD_LIMIT) * factor) + DEFAULT_COORD_LIMIT; -if (shardSize > Integer.MAX_VALUE) return; -if (shardSize <= coordLimit) return; +if (shardSize > Integer.MAX_VALUE || shardSize <= coordLimit) { + return; +} Suggestion importance[1-10]: 3 __ Why: The suggestion correctly identifies that the overflow check occurs after casting to long, but the current code already returns early if shardSize > Integer.MAX_VALUE (line 102), preventing the truncation. The suggested consolidation of the two checks is a minor readability improvement but doesn't fix a functional bug.	Low

Suggestions up to commit d462415

Category	Suggestion	Impact
General	Add bounds check before indexing child Verify that new_children has exactly one element before indexing new_children[0]. ProjectionExec should have exactly one child, but defensive bounds checking prevents potential panics if the plan structure is unexpected. sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [110-125] } else if let Some(proj) = plan.as_any().downcast_ref::<datafusion::physical_plan::projection::ProjectionExec>() { - // ProjectionExec::with_new_children reuses the old Projector (with cached - // output schema). When the child's schema changed (e.g. Float64 → Binary - // after stripping), the stale schema causes runtime Arrow type mismatches. - // Rebuild from expressions + new child to recompute the output schema. + if new_children.len() != 1 { + return exec_err!("ProjectionExec expected 1 child, got {}", new_children.len()); + } let exprs: Vec<(Arc<dyn datafusion::physical_plan::PhysicalExpr>, String)> = proj .expr() .iter() .map(|pe| (Arc::clone(&pe.expr), pe.alias.clone())) .collect(); Ok(Arc::new( datafusion::physical_plan::projection::ProjectionExec::try_new( exprs, Arc::clone(&new_children[0]), )?, )) Suggestion importance[1-10]: 7 __ Why: Good defensive programming suggestion. While ProjectionExec should always have exactly one child by design, adding an explicit bounds check prevents potential panics and provides clearer error messages if the plan structure is unexpected.	Medium
	Use nanoTime for elapsed time The polling loop uses System.currentTimeMillis() which can be affected by system clock adjustments. Use System.nanoTime() for elapsed-time measurements to avoid issues with clock skew or NTP corrections during the test run. sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ShardBucketOversamplingIT.java [312-317] -long deadline = System.currentTimeMillis() + 10_000L; -while (System.currentTimeMillis() < deadline) { +long deadlineNanos = System.nanoTime() + 10_000_000_000L; // 10 seconds in nanos +while (System.nanoTime() < deadlineNanos) { ... Thread.sleep(100); } Suggestion importance[1-10]: 5 __ Why: Using System.nanoTime() for elapsed-time measurements is a best practice to avoid clock-adjustment issues. The suggestion is correct and improves test reliability, though the impact is moderate since system clock adjustments during a 10-second test window are rare in practice.	Low
	Optimize memory allocation for non-null rows Pre-allocate col_values with the actual non-null count instead of arr.len() to avoid wasted capacity when many rows are null. Count non-nulls first, then allocate vectors with that exact capacity to reduce memory overhead. sandbox/plugins/analytics-backend-datafusion/rust/src/udaf/state_shipping.rs [310-322] fn merge_state_binary(&mut self, binary_array: &ArrayRef) -> Result<()> { let arr = binary_array .as_any() .downcast_ref::<BinaryArray>() .ok_or_else(|| { DataFusionError::Execution( "StateShippingAccumulator: expected BinaryArray for state input".to_string(), ) })?; let num_state_cols = self.state_schema.fields().len(); - // Row-major decode → column-major rebuild. Skip null rows. + let non_null_count = arr.len() - arr.null_count(); let mut col_values: Vec<Vec<ScalarValue>> = - vec![Vec::with_capacity(arr.len()); num_state_cols]; + vec![Vec::with_capacity(non_null_count); num_state_cols]; for row_idx in 0..arr.len() { if arr.is_null(row_idx) { continue; } let bytes = arr.value(row_idx); let row_state = decode_state_from_ipc(bytes, &self.state_schema)?; ... } Suggestion importance[1-10]: 5 __ Why: The suggestion correctly identifies a memory optimization opportunity by pre-allocating based on non_null_count. However, the impact is moderate since the vectors will grow to the correct size anyway, and the optimization only saves initial allocation overhead when many nulls are present.	Low
	Validate shard size is positive The shardFetch RexNode is created with makeExactLiteral using BigDecimal.valueOf(hint.shardSize()). If hint.shardSize() is negative or zero (due to a bug in the hint construction), the Sort will have an invalid fetch limit. Add a defensive check to ensure hint.shardSize() is positive before constructing the RexNode. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java [150-151] -OpenSearchSort shardSort = new OpenSearchSort( - shardLocalAgg.getCluster(), - shardLocalAgg.getTraitSet(), - shardLocalAgg, - hint.collation(), - null, - shardFetch, - shardLocalAgg.getViableBackends(), - true, - hint.sortExprs() -); +if (hint.shardSize() <= 0) { + throw new IllegalStateException("ShardBucketHint.shardSize must be positive, got: " + hint.shardSize()); +} +RexNode shardFetch = rexBuilder.makeExactLiteral(java.math.BigDecimal.valueOf(hint.shardSize()), intType); +OpenSearchSort shardSort = new OpenSearchSort(...); Suggestion importance[1-10]: 4 __ Why: Adding a defensive check for hint.shardSize() <= 0 is a reasonable safeguard against invalid hint construction. However, the ShardBucketHint is constructed in OpenSearchAggregateShardBucketRule where shardSize is already validated to be positive (via the calculation and the check shardSize <= coordLimit), so this is redundant defensive programming with low practical impact.	Low
	Cache null scalar to avoid repeated construction Each finalize_one call creates a fresh accumulator and decodes IPC state independently. For large arrays, this creates significant per-row overhead. Consider batching multiple rows through a single accumulator when possible to amortize initialization costs. sandbox/plugins/analytics-backend-datafusion/rust/src/udf/state_finalize.rs [181-196] fn finalize_array(&self, arr: &dyn Array) -> Result<ColumnarValue> { let bin = arr.as_any().downcast_ref::<BinaryArray>().ok_or_else(|| { DataFusionError::Execution(format!( "{}: expected BinaryArray, got {:?}", self.name, arr.data_type() )) })?; let mut scalars: Vec<ScalarValue> = Vec::with_capacity(bin.len()); + let null_scalar = null_scalar_for(&self.return_type)?; for i in 0..bin.len() { if bin.is_null(i) { - scalars.push(null_scalar_for(&self.return_type)?); + scalars.push(null_scalar.clone()); } else { scalars.push(self.finalize_one(bin.value(i))?); } } Suggestion importance[1-10]: 4 __ Why: Valid micro-optimization that avoids repeated null_scalar_for calls for null rows. The improvement is minor since null_scalar_for is likely inexpensive, but caching the null scalar is a reasonable optimization for arrays with many nulls.	Low
	Validate aggregate phase rewrite success The withAggregationPhaseRecursive method modifies the plan in-place by building a new list of roots. If the plan has no aggregates, changed remains false and the original plan is returned. However, if the plan structure is complex or contains nested aggregates, ensure that the visitor correctly handles all aggregate nodes. Consider adding validation to confirm that at least one aggregate was found and modified, or log a warning if no changes were made despite the method being called. sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java [461-462] Plan rephased = withAggregationPhaseRecursive(decoded, Expression.AggregationPhase.INITIAL_TO_INTERMEDIATE); +if (rephased == decoded) { + LOGGER.warn("withAggregationPhaseRecursive returned unchanged plan; no aggregates found or modified"); +} Suggestion importance[1-10]: 3 __ Why: Adding a warning when no aggregates are modified is a minor observability improvement. The suggestion is valid but has low impact since the method is designed to be a no-op when no aggregates are present, which is a legitimate case.	Low
Possible issue	Prevent silent arithmetic overflow The shardSize calculation can overflow when factor is very large or coordLimit is near Long.MAX_VALUE. Although the check shardSize > Integer.MAX_VALUE guards against casting to int, the multiplication and addition could silently overflow before the check. Use Math.multiplyExact and Math.addExact to detect overflow early and handle it explicitly. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateShardBucketRule.java [101-102] -long shardSize = (long) Math.ceil(Math.max(coordLimit, DEFAULT_COORD_LIMIT) * factor) + DEFAULT_COORD_LIMIT; -if (shardSize > Integer.MAX_VALUE) return; +try { + long maxLimit = Math.max(coordLimit, DEFAULT_COORD_LIMIT); + long scaled = Math.multiplyExact((long) Math.ceil(maxLimit * factor), 1L); + long shardSize = Math.addExact(scaled, DEFAULT_COORD_LIMIT); + if (shardSize > Integer.MAX_VALUE) return; + // ... rest of logic +} catch (ArithmeticException e) { + return; // overflow detected, skip hint +} Suggestion importance[1-10]: 6 __ Why: The suggestion correctly identifies a potential arithmetic overflow scenario when factor is very large. Using Math.multiplyExact and Math.addExact would catch overflow explicitly. However, the existing check shardSize > Integer.MAX_VALUE already guards against invalid values, and the practical risk is low given typical factor ranges (0.0 or ≥1.0).	Low

[github-actions]

❌ Gradle check result for 3d360bb: 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 b49b5a9

[github-actions]

✅ Gradle check result for b49b5a9: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.52%. Comparing base (ba395cb) to head (ba0d773).

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #21775      +/-   ##
============================================
+ Coverage     73.42%   73.52%   +0.10%     
- Complexity    75503    75564      +61     
============================================
  Files          6034     6034              
  Lines        342661   342661              
  Branches      49294    49294              
============================================
+ Hits         251584   251927     +343     
+ Misses        71063    70710     -353     
- Partials      20014    20024      +10     

☔ 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 2ab802d

[github-actions]

❌ Gradle check result for 2ab802d: 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 76d9b60

[github-actions]

Persistent review updated to latest commit ff75c9d

[github-actions]

✅ Gradle check result for ff75c9d: SUCCESS

[github-actions]

Persistent review updated to latest commit 4c9337d

[github-actions]

Persistent review updated to latest commit 475fc9f

[github-actions]

✅ Gradle check result for 475fc9f: SUCCESS

[github-actions]

Persistent review updated to latest commit 4fba1aa

[expani]

Thanks for the changes @sandeshkr419

Dropping some early comments, still going through the PR.

[github-actions]

PR Code Analyzer ❗

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

Path	Line	Severity	Description
sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/ShardBucketOversamplingIT.java	38	low	The ensureProvisioned() method sets a persistent cluster setting (analytics.shard_bucket_oversampling_factor=2.0) with no guaranteed teardown. If any test other than testFactorZero_queryWorks runs last, the setting persists in the cluster after the test suite exits. This is a test hygiene issue with minor operational impact but no evidence of malicious intent.

The table above displays the top 10 most important findings.

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

---

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.

[github-actions]

Persistent review updated to latest commit 9f92e07

[github-actions]

❌ Gradle check result for 9f92e07: 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 1f61124

[github-actions]

✅ Gradle check result for 1f61124: SUCCESS

[github-actions]

Persistent review updated to latest commit e27e229

[expani]

Thanks @sandeshkr419 it was fun iterating on this offline.

Really appreciate you sticking to the architecture and design intentions 🫡

Have a few comments around a few oddities that should be quick.

[github-actions]

❕ Gradle check result for ba0d773: UNSTABLE

Please review all flaky tests that succeeded after retry and create an issue if one does not already exist to track the flaky failure.

[github-actions]

Persistent review updated to latest commit e770031