Skip to content

feat(analytics-engine): fix distributed COUNT(DISTINCT) / DC via HLL sketch merge + TopK optimization - #22013

Merged
mch2 merged 7 commits into
opensearch-project:mainfrom
sandeshkr419:dcr
Jun 6, 2026
Merged

feat(analytics-engine): fix distributed COUNT(DISTINCT) / DC via HLL sketch merge + TopK optimization#22013
mch2 merged 7 commits into
opensearch-project:mainfrom
sandeshkr419:dcr

Conversation

@sandeshkr419

Copy link
Copy Markdown
Member

Implements cross-shard dc() / distinct_count() using HyperLogLog sketch merge. Single-arg COUNT(DISTINCT x) is rewritten to APPROX_COUNT_DISTINCT at plan time, enabling per-shard partial HLL sketch computation with coordinator-side merge — eliminating the previous over-counting bug where per-shard distinct counts were summed.

For TopK queries (stats dc(x) by group | sort - dc(x) | head N), inserts a reduce_eval projection between the PARTIAL aggregate and the shard Sort to derive a sortable cardinality from opaque HLL state, enabling shard-side pruning that ships only top-K sketches over the wire.

Changes

Plan layer (Java)

  • OpenSearchDistinctCountRule: HEP rule rewrites single-arg COUNT(DISTINCT x) → APPROX_COUNT_DISTINCT(x) before aggregate marking
  • OpenSearchTopKRewriter: detects sort on engine-native-merge measures, inserts reduce_eval Project + strip Project around shard Sort
  • FragmentConversionDriver: layered substrait conversion for buried PARTIAL aggregates (scan → attachPartialAggOnTop → attachFragmentOnTop per operator above); emits SETUP_PARTIAL_AGGREGATE / SETUP_FINAL_AGGREGATE instructions via tree-walk
  • AggregateFunction SPI: isEngineNativeMerge(AggregateCall), reduceEvalName(), REDUCE_EVAL_OP — shared predicate and operator, no duplication across modules

Execution layer (Rust)

  • force_aggregate_mode: strips Final/Partial half of aggregates; remaps ProjectionExec Column references when child schema changes after strip
  • prepare_partial_plan / prepare_final_plan: activate Partial/Final execution mode on data-node/coordinator respectively
  • derive_schema_from_partial_plan: extracts Partial aggregate schema for StreamingTable registration; uses substrait Root.names for authoritative column naming (no string-pattern hacking)
  • reduce_eval UDF: evaluates HLL sketch state → UInt64 cardinality for TopK ranking

Test plan

  • ShardBucketOversamplingIT — all TopK tests including testDcByGroup_sortDesc_head10
  • TwoShardAggregationIT — all 36 reduce checks including distinct_count_by_cat
  • CoordinatorReduceIT — all 21 checks including testQ10ShapeAcrossShards, testGroupByCountMultiShard
  • AggregatePlanShapeTests — unit tests for rewrite + split shape
  • AggregateFunctionTests — SPI enum validation

Signed-off-by: Sandesh Kumar sandeshkr419@gmail.com

@sandeshkr419
sandeshkr419 requested a review from a team as a code owner June 5, 2026 13:23
@github-actions

github-actions Bot commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

PR Reviewer Guide 🔍

(Review updated until commit 64b8562)

Here are some key observations to aid the review process:

🧪 PR contains tests
🔒 No security concerns identified
✅ No TODO sections
🔀 No multiple PR themes
⚡ Recommended focus areas for review

Possible Issue

remap_column recursively remaps column names but returns the original expression on with_new_children failure via unwrap_or(fallback). If with_new_children fails after partial remapping of children, the returned expression is inconsistent—some children remapped, some not. This can cause schema mismatches when DataFusion validates column references. The failure scenario occurs when an expression type does not support child replacement (e.g., certain custom physical expressions).

fn remap_column(expr: Arc<dyn PhysicalExpr>, schema: &arrow::datatypes::SchemaRef) -> Arc<dyn PhysicalExpr> {
    if let Some(col) = expr.as_any().downcast_ref::<Column>() {
        return Arc::new(Column::new(schema.field(col.index()).name(), col.index()));
    }
    let children = expr.children();
    if children.is_empty() { return expr; }
    let new_children: Vec<_> = children.into_iter().map(|c| remap_column(c.clone(), schema)).collect();
    let fallback = expr.clone();
    expr.with_new_children(new_children).unwrap_or(fallback)
}
Possible Issue

derive_schema_from_partial_plan assumes declared_names.len() == partial_schema.fields().len() when both conditions are true, but does not verify this equality before zipping. If lengths differ (e.g., due to substrait plan inconsistencies or optimizer transformations), the zip silently truncates to the shorter length, producing a schema with fewer fields than expected. This causes downstream failures when the coordinator attempts to read all expected columns from the partial result.

if let Some(partial_schema) = crate::agg_mode::partial_aggregate_schema(&physical_plan) {
    let has_binary = partial_schema.fields().iter().any(|f| matches!(f.data_type(), arrow::datatypes::DataType::Binary));
    if has_binary && !declared_names.is_empty() && declared_names.len() == partial_schema.fields().len() {
        use arrow::datatypes::{Field, Schema};
Possible Issue

findBuriedPartialAggregate returns null when the fragment root is itself an OpenSearchAggregate, even if that aggregate is PARTIAL and engine-native-merge. The early-return if (fragment instanceof OpenSearchAggregate) return null; at line 602 prevents detection of top-level PARTIAL aggregates that should be handled by the buried-aggregate path. This causes the layered conversion logic (lines 428-448) to skip engine-native PARTIAL aggregates at the fragment root, falling back to the standard path that does not emit SETUP_PARTIAL_AGGREGATE instructions.

private static OpenSearchAggregate findBuriedPartialAggregate(RelNode fragment) {
    if (fragment instanceof OpenSearchAggregate) return null;
Index Out of Bounds

hasEngineNativeMergeMeasure computes sortIdx - groupCount to index into agg.getAggCallList() without verifying sortIdx >= groupCount. If the sort references a group-by column (sortIdx < groupCount), the subtraction yields a negative index, triggering an IndexOutOfBoundsException when accessing the aggregate call list.

private static boolean hasEngineNativeMergeMeasure(OpenSearchAggregate agg, RelCollation collation) {
    int groupCount = agg.getGroupSet().cardinality();
    for (RelFieldCollation fc : collation.getFieldCollations()) {
        int idx = fc.getFieldIndex();
        if (idx >= groupCount && AggregateFunction.isEngineNativeMerge(agg.getAggCallList().get(idx - groupCount))) {
            return true;
        }
    }
    return false;
}

@github-actions

github-actions Bot commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

PR Code Suggestions ✨

Latest suggestions up to 64b8562

Explore these optional code suggestions:

CategorySuggestion                                                                                                                                    Impact
Possible issue
Add bounds checking for aggregate list access

The expression sortIdx - groupCount could produce a negative index or exceed the
bounds of getAggCallList() if sortIdx is less than groupCount or if the aggregate
call list is shorter than expected. Add bounds checking before accessing the list.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [92-108]

 for (RelFieldCollation fc : sort.getCollation().getFieldCollations()) {
     int sortIdx = fc.getFieldIndex();
-    if (sortIdx >= groupCount && AggregateFunction.isEngineNativeMerge(partial.getAggCallList().get(sortIdx - groupCount))) {
+    int aggIdx = sortIdx - groupCount;
+    if (sortIdx >= groupCount 
+        && aggIdx >= 0 
+        && aggIdx < partial.getAggCallList().size()
+        && AggregateFunction.isEngineNativeMerge(partial.getAggCallList().get(aggIdx))) {
         ...
     }
 }
Suggestion importance[1-10]: 8

__

Why: Critical bounds checking issue. The code accesses partial.getAggCallList().get(sortIdx - groupCount) without verifying that sortIdx - groupCount is within valid bounds. If sortIdx references a group key field or the list is shorter than expected, this will throw IndexOutOfBoundsException. The suggestion correctly adds necessary validation.

Medium
Prevent infinite recursion in tree walk

The recursive tree walk has no depth limit or cycle detection. If the plan contains
a cycle or is extremely deep, this could cause a stack overflow. Consider adding a
visited set or depth limit to prevent infinite recursion.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java [256-266]

 private static boolean containsEngineNativeAggregate(RelNode root, AggregateMode mode) {
+    return containsEngineNativeAggregate(root, mode, new java.util.HashSet<>());
+}
+
+private static boolean containsEngineNativeAggregate(RelNode root, AggregateMode mode, java.util.Set<RelNode> visited) {
+    if (!visited.add(root)) return false;
     if (root instanceof OpenSearchAggregate agg
         && agg.getMode() == mode
         && agg.getAggCallList().stream().anyMatch(org.opensearch.analytics.spi.AggregateFunction::isEngineNativeMerge)) {
         return true;
     }
     for (RelNode child : root.getInputs()) {
-        if (containsEngineNativeAggregate(child, mode)) return true;
+        if (containsEngineNativeAggregate(child, mode, visited)) return true;
     }
     return false;
 }
Suggestion importance[1-10]: 7

__

Why: Valid concern about potential infinite recursion in tree traversal. While Calcite's RelNode graphs are typically DAGs without cycles, adding cycle detection is a defensive programming practice that prevents stack overflow in edge cases without significant performance cost.

Medium
Guard against empty inputs list

The condition checks node.getInputs().getFirst() without verifying that
node.getInputs() is non-empty. If node has no inputs, this will throw an exception.
Add a guard to ensure the inputs list is not empty before accessing the first
element.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java [551-555]

 if (node instanceof org.apache.calcite.rel.core.Project p
     && isPureReorderProject(p)
+    && !node.getInputs().isEmpty()
     && containsEngineNativeAggregate(node.getInputs().getFirst(), AggregateMode.FINAL)) {
     return innerBytes;
 }
Suggestion importance[1-10]: 3

__

Why: The suggestion adds a defensive check for empty inputs, but the context shows this code is within convertReduceNode which already validates single-input operators. The check node.getInputs().size() >= 2 earlier ensures non-empty inputs for this path, making the additional guard redundant but harmless.

Low

Previous suggestions

Suggestions up to commit 4956a4e
CategorySuggestion                                                                                                                                    Impact
Possible issue
Validate array index before access

The code accesses partial.getAggCallList().get(sortIdx - groupCount) without
verifying that sortIdx - groupCount is within bounds. If sortIdx is less than
groupCount or the index exceeds the list size, this will throw an
IndexOutOfBoundsException.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [92-94]

 for (RelFieldCollation fc : sort.getCollation().getFieldCollations()) {
     int sortIdx = fc.getFieldIndex();
-    if (sortIdx >= groupCount && AggregateFunction.isEngineNativeMerge(partial.getAggCallList().get(sortIdx - groupCount))) {
+    int aggIdx = sortIdx - groupCount;
+    if (sortIdx >= groupCount && aggIdx < partial.getAggCallList().size() 
+        && AggregateFunction.isEngineNativeMerge(partial.getAggCallList().get(aggIdx))) {
Suggestion importance[1-10]: 8

__

Why: Critical bounds check missing. The code assumes sortIdx - groupCount is valid without verifying it's within getAggCallList().size(). This could cause IndexOutOfBoundsException with malformed sort collations.

Medium
Prevent infinite loop in tree traversal

The while loop assumes walk will eventually equal partialAgg, but if the tree
structure is unexpected or partialAgg is not actually an ancestor, this creates an
infinite loop. Add a safety check to prevent unbounded iteration.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java [440-445]

 List<RelNode> aboveAgg = new ArrayList<>();
 RelNode walk = resolvedFragment;
-while (walk != partialAgg) {
+int maxDepth = 100;
+int depth = 0;
+while (walk != partialAgg && depth < maxDepth) {
     aboveAgg.add(walk);
+    if (walk.getInputs().isEmpty()) break;
     walk = walk.getInputs().getFirst();
+    depth++;
+}
+if (walk != partialAgg) {
+    throw new IllegalStateException("Could not find partialAgg in tree");
 }
Suggestion importance[1-10]: 7

__

Why: Valid concern about potential infinite loops if partialAgg is not found in the tree. Adding a depth limit and explicit error handling improves robustness against malformed plans.

Medium
Verify inputs exist before accessing

The condition checks if node is a Project but doesn't verify that node.getInputs()
is non-empty before calling getFirst(). If a Project node has no inputs, this will
throw an exception. Add a guard to ensure inputs exist.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java [551-555]

 if (node instanceof org.apache.calcite.rel.core.Project p
     && isPureReorderProject(p)
+    && !node.getInputs().isEmpty()
     && containsEngineNativeAggregate(node.getInputs().getFirst(), AggregateMode.FINAL)) {
     return innerBytes;
 }
Suggestion importance[1-10]: 3

__

Why: The suggestion correctly identifies a potential issue, but Project nodes in Calcite always have exactly one input by design. The check is defensive but unlikely to prevent real bugs in practice.

Low
General
Handle errors explicitly in remapping

Using unwrap_or with a cloned fallback silently ignores errors from
with_new_children. If the operation fails for a legitimate reason (e.g.,
incompatible child types), the error is masked. Consider propagating the error or
logging it.

sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [147-150]

 let new_children: Vec<_> = children.into_iter().map(|c| remap_column(c.clone(), schema)).collect();
-let fallback = expr.clone();
-expr.with_new_children(new_children).unwrap_or(fallback)
+match expr.with_new_children(new_children) {
+    Ok(new_expr) => new_expr,
+    Err(e) => {
+        log::warn!("Failed to remap column children: {}", e);
+        expr.clone()
+    }
+}
Suggestion importance[1-10]: 5

__

Why: The suggestion improves error visibility by logging failures instead of silently falling back. However, the current unwrap_or pattern is acceptable for non-critical remapping operations where fallback is reasonable.

Low
Suggestions up to commit f324e1e
CategorySuggestion                                                                                                                                    Impact
Possible issue
Add bounds check for list access

The code accesses partial.getAggCallList().get(sortIdx - groupCount) without
verifying that sortIdx - groupCount is within bounds. If sortIdx is less than
groupCount or the index exceeds the aggregate call list size, this will throw
IndexOutOfBoundsException. Add bounds checking before accessing the list.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [92-94]

 for (RelFieldCollation fc : sort.getCollation().getFieldCollations()) {
     int sortIdx = fc.getFieldIndex();
-    if (sortIdx >= groupCount && AggregateFunction.isEngineNativeMerge(partial.getAggCallList().get(sortIdx - groupCount))) {
+    int aggIdx = sortIdx - groupCount;
+    if (sortIdx >= groupCount 
+        && aggIdx >= 0 
+        && aggIdx < partial.getAggCallList().size()
+        && AggregateFunction.isEngineNativeMerge(partial.getAggCallList().get(aggIdx))) {
Suggestion importance[1-10]: 8

__

Why: The suggestion identifies a legitimate bounds-checking issue. The code accesses partial.getAggCallList().get(sortIdx - groupCount) without verifying the index is valid, which could cause IndexOutOfBoundsException if sortIdx - groupCount exceeds the list size.

Medium
Check inputs exist before accessing

The condition checks if node is a Project but then calls
containsEngineNativeAggregate on node.getInputs().getFirst() without verifying that
node has inputs. If node.getInputs() is empty, getFirst() will throw
NoSuchElementException. Add a guard to ensure inputs exist before accessing them.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java [551-555]

 if (node instanceof org.apache.calcite.rel.core.Project p
     && isPureReorderProject(p)
+    && !node.getInputs().isEmpty()
     && containsEngineNativeAggregate(node.getInputs().getFirst(), AggregateMode.FINAL)) {
     return innerBytes;
 }
Suggestion importance[1-10]: 7

__

Why: The suggestion correctly identifies a potential NoSuchElementException if node.getInputs() is empty. However, the context shows this code is within convertReduceNode which processes single-input operators, making an empty inputs list unlikely but still possible.

Medium
Validate node structure before processing

The loop iterates backwards through aboveAgg but doesn't verify that each element
has the expected single-input structure before calling stripSingleOperator. If any
element in the list has multiple inputs or is null, stripSingleOperator may fail.
Validate each node before processing to prevent runtime failures.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java [445-448]

 for (int i = aboveAgg.size() - 1; i >= 0; i--) {
-    current = convertor.attachFragmentOnTop(stripSingleOperator(aboveAgg.get(i)), current);
+    RelNode node = aboveAgg.get(i);
+    if (node == null || node.getInputs().isEmpty()) {
+        throw new IllegalStateException("Invalid node in aboveAgg chain at index " + i);
+    }
+    current = convertor.attachFragmentOnTop(stripSingleOperator(node), current);
 }
Suggestion importance[1-10]: 6

__

Why: The suggestion adds defensive validation for nodes in aboveAgg list. While the null check is reasonable, the empty inputs check may be overly strict since stripSingleOperator already handles single-input nodes and creates a placeholder child.

Low
Suggestions up to commit 6d3daf3
CategorySuggestion                                                                                                                                    Impact
Possible issue
Add bounds check for aggregate index

The index calculation sortIdx - groupCount is used to access
partial.getAggCallList() without bounds checking. If sortIdx is less than groupCount
or if the resulting index exceeds the aggregate call list size, this will throw an
IndexOutOfBoundsException. Add explicit bounds validation before the list access.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [92-95]

 List<RelFieldCollation> newCollations = new ArrayList<>();
 for (RelFieldCollation fc : sort.getCollation().getFieldCollations()) {
     int sortIdx = fc.getFieldIndex();
-    if (sortIdx >= groupCount && AggregateFunction.isEngineNativeMerge(partial.getAggCallList().get(sortIdx - groupCount))) {
+    int aggIdx = sortIdx - groupCount;
+    if (sortIdx >= groupCount && aggIdx >= 0 && aggIdx < partial.getAggCallList().size() 
+        && AggregateFunction.isEngineNativeMerge(partial.getAggCallList().get(aggIdx))) {
         ...
     }
 }
Suggestion importance[1-10]: 8

__

Why: Critical bounds checking issue. The code accesses partial.getAggCallList().get(sortIdx - groupCount) without verifying that sortIdx - groupCount is within valid bounds. This could cause IndexOutOfBoundsException if the sort index is invalid. The suggested bounds check prevents a potential runtime crash.

Medium
Prevent panic on out-of-bounds column index

The schema.field(col.index()) call will panic if col.index() is out of bounds for
the schema. This can occur if the column index from the old schema doesn't exist in
the new schema after stripping. Add bounds checking and return the original
expression or handle the error gracefully if the index is invalid.

sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [141-145]

 fn remap_column(expr: Arc<dyn PhysicalExpr>, schema: &arrow::datatypes::SchemaRef) -> Arc<dyn PhysicalExpr> {
     if let Some(col) = expr.as_any().downcast_ref::<Column>() {
-        return Arc::new(Column::new(schema.field(col.index()).name(), col.index()));
+        if col.index() < schema.fields().len() {
+            return Arc::new(Column::new(schema.field(col.index()).name(), col.index()));
+        }
+        return expr;
     }
     ...
 }
Suggestion importance[1-10]: 8

__

Why: Important safety check for Rust code. The schema.field(col.index()) call will panic if the index is out of bounds. After schema stripping operations, column indices may become invalid. The suggested bounds check prevents a panic and gracefully falls back to the original expression.

Medium
Prevent unbounded recursion in tree walk

The recursive tree walk in containsEngineNativeAggregate has no depth limit or cycle
detection. If the RelNode tree contains a cycle (malformed plan) or is extremely
deep, this could cause a stack overflow. Add a depth limit parameter or use an
iterative approach with a visited set to prevent unbounded recursion.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java [256-266]

 private static boolean containsEngineNativeAggregate(RelNode root, AggregateMode mode) {
+    return containsEngineNativeAggregateHelper(root, mode, new java.util.HashSet<>(), 0, 100);
+}
+
+private static boolean containsEngineNativeAggregateHelper(RelNode root, AggregateMode mode, java.util.Set<RelNode> visited, int depth, int maxDepth) {
+    if (depth > maxDepth || !visited.add(root)) return false;
     if (root instanceof OpenSearchAggregate agg
         && agg.getMode() == mode
         && agg.getAggCallList().stream().anyMatch(org.opensearch.analytics.spi.AggregateFunction::isEngineNativeMerge)) {
         return true;
     }
     for (RelNode child : root.getInputs()) {
-        if (containsEngineNativeAggregate(child, mode)) return true;
+        if (containsEngineNativeAggregateHelper(child, mode, visited, depth + 1, maxDepth)) return true;
     }
     return false;
 }
Suggestion importance[1-10]: 7

__

Why: Valid concern about potential stack overflow from unbounded recursion in containsEngineNativeAggregate. The RelNode tree structure should be acyclic by design, but adding depth protection is a reasonable defensive measure. The suggestion provides a working implementation with cycle detection and depth limiting.

Medium
Validate factory node creation success

The createPartialAggregateNode() and createFinalAggregateNode() calls are
conditional on containsEngineNativeAggregate returning true, but there's no
null-safety check on the factory methods. If the factory returns an empty Optional,
the instruction won't be added, which could silently skip critical aggregate setup.
Consider logging a warning or throwing an exception when the factory fails to create
the expected node for an engine-native aggregate.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java [246-251]

 if (leaf instanceof OpenSearchTableScan tableScan) {
     ...
     if (containsEngineNativeAggregate(resolvedFragment, AggregateMode.PARTIAL)) {
-        factory.createPartialAggregateNode().ifPresent(instructions::add);
+        factory.createPartialAggregateNode().ifPresentOrElse(
+            instructions::add,
+            () -> { throw new IllegalStateException("Factory failed to create PARTIAL aggregate node for engine-native merge"); }
+        );
     }
 } else if (leaf instanceof OpenSearchStageInputScan && containsEngineNativeAggregate(resolvedFragment, AggregateMode.FINAL)) {
-    factory.createFinalAggregateNode().ifPresent(instructions::add);
+    factory.createFinalAggregateNode().ifPresentOrElse(
+        instructions::add,
+        () -> { throw new IllegalStateException("Factory failed to create FINAL aggregate node for engine-native merge"); }
+    );
 }
Suggestion importance[1-10]: 5

__

Why: The suggestion correctly identifies that createPartialAggregateNode() and createFinalAggregateNode() return Optional values that could be empty. However, throwing an exception may be too strict if the factory legitimately returns empty for certain backends. A warning log would be more appropriate, making this a moderate improvement rather than a critical fix.

Low
Suggestions up to commit 4acb538
CategorySuggestion                                                                                                                                    Impact
General
Handle unrecognized aggregate function gracefully

The fromSqlAggFunction call can throw IllegalStateException if the aggregate
function is not recognized. This exception should be caught and handled
appropriately to prevent TopK rewrite failures when encountering unsupported
aggregate functions.

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTopKRewriter.java [94-96]

 if (sortIdx >= groupCount && AggregateFunction.isEngineNativeMerge(partial.getAggCallList().get(sortIdx - groupCount))) {
-    AggregateFunction aggFunc = AggregateFunction.fromSqlAggFunction(
-        partial.getAggCallList().get(sortIdx - groupCount).getAggregation());
+    AggregateFunction aggFunc;
+    try {
+        aggFunc = AggregateFunction.fromSqlAggFunction(
+            partial.getAggCallList().get(sortIdx - groupCount).getAggregation());
+    } catch (IllegalStateException e) {
+        continue;
+    }
Suggestion importance[1-10]: 5

__

Why: The suggestion correctly identifies that fromSqlAggFunction can throw IllegalStateException. However, silently continuing on error may hide legitimate issues. A better approach would be to log the error or fail fast if the function is truly unsupported.

Low
Ensure symmetric Final/FinalPartitioned matching

The logic treats FinalPartitioned as equivalent to Final only when the target is
Final, but doesn't handle the reverse case. If the target is FinalPartitioned and
the mode is Final, they should also match to ensure symmetric handling of both final
aggregate modes.

sandbox/plugins/analytics-backend-datafusion/rust/src/agg_mode.rs [67-69]

 let agg_is_target = *agg.mode() == target
-    || (target == AggregateMode::Final && *agg.mode() == AggregateMode::FinalPartitioned);
+    || (target == AggregateMode::Final && *agg.mode() == AggregateMode::FinalPartitioned)
+    || (target == AggregateMode::FinalPartitioned && *agg.mode() == AggregateMode::Final);
 if agg_is_target {
Suggestion importance[1-10]: 3

__

Why: The suggestion adds symmetric handling for Final/FinalPartitioned matching. While this could improve consistency, the PR's comment explicitly states that both modes are treated as the FINAL half when the target is Final, suggesting the asymmetry is intentional for the current use case.

Low

@github-actions

github-actions Bot commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

Persistent review updated to latest commit 6d3daf3

@github-actions

github-actions Bot commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

Persistent review updated to latest commit f324e1e

@github-actions

github-actions Bot commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

Persistent review updated to latest commit 4956a4e

sandeshkr419 and others added 7 commits June 5, 2026 22:19
…CT via HEP rule + strip aggregate state suffix in coord StreamingTable schema

Two complementary fixes for the analytics-engine cross-shard aggregate path.

(1) Plan-layer rewrite (Java, OpenSearchDistinctCountRule)

PPL `dc(x)` / `distinct_count(x)` parse to `COUNT(DISTINCT x)` at Calcite via
the SQL plugin's `AstExpressionBuilder.visitDistinctCountFunctionCall`. Without
intervention the call lands as `SqlStdOperatorTable.COUNT` with isDistinct=true,
the additive split rule decomposes it (PARTIAL count(distinct) + FINAL SUM-of-counts),
and cross-shard reduce over-counts any value present on more than one shard.

Replace 6d98's `AggregateFunction.resolveOperator` SPI hook (which polluted the
framework enum and rebuilt AggregateCalls with the original COUNT return type, risking
`Aggregate.typeMatchesInferred` mismatches against APPROX_COUNT_DISTINCT's
inferReturnType) with a dedicated HEP rule:

* New `OpenSearchDistinctCountRule` matches plain `LogicalAggregate` containing a
  single-arg `COUNT(DISTINCT x)` and rewrites it to `APPROX_COUNT_DISTINCT(x)`
  (isDistinct=false on the rewritten call). Built via the long-form `AggregateCall.create`
  with `(groupCount, input, type=null)` so Calcite re-infers the return type from
  `APPROX_COUNT_DISTINCT.inferReturnType(...)` — avoids the typeMatchesInferred mismatch
  the SPI rebuild would produce.
* Wired into the existing `PlannerImpl.decomposeAggregates` HEP phase alongside
  `OpenSearchAggregateReduceRule`, before `OpenSearchAggregateRule` marks the aggregate.
  Multi-arg `COUNT(DISTINCT a, b)` doesn't match — falls through to the residual
  `aggCall.isDistinct()` skip in `OpenSearchAggregateSplitRule`.

After this, the aggregate engages the existing
`AggregateFunction.APPROX_COUNT_DISTINCT(Type.APPROXIMATE, intermediateFields=[sketch:Binary, reducer==self])`
registration — capability resolution, structural split, `DistributedAggregateRewriter.overrideExchangeType`,
and the Substrait extension YAML `approx_distinct` alias all light up automatically.

* Drop `AggregateFunction.COUNT.resolveOperator` override + the default `resolveOperator`
  method on the enum.
* Drop `OpenSearchAggregateRule.resolveAggregateCall` + its unused `SqlAggFunction` import.
* Drop `Type.APPROXIMATE` from `OpenSearchAggregateSplitRule.shouldSkipPartialFinalSplit`
  (residual `aggCall.isDistinct()` skip stays for multi-arg fallback). APPROXIMATE now goes
  through the structural PARTIAL/FINAL split.

(2) Wire-layer schema bridge (Rust, derive_schema_from_partial_plan)

Commit `35ce14790c2` (opensearch-project#21690) folded the consumer's StreamingTable schema derivation into Rust
to eliminate per-cell coercion overhead; `coerceToDeclaredSchema` on the Java
`feedToSender` path was deleted in favour of running the producer's substrait through
DataFusion's substrait consumer + physical planner and using its output schema verbatim.

The DataFusion 53.x physical planner emits `AggregateExec(Mode::Partial)` columns with
state-suffixed names (`dc[hll_registers]`, `$f0[sum]`, `count(opt)[count]`), but the
FINAL substrait emitted by `attachPartialAggOnTop` declares the user-facing aliases
(`dc`, `$f0`, `count(opt)`). DataFusion's substrait consumer name-resolves the FINAL
Read against the registered StreamingTable and fails with
`Schema error: No field named dc. Valid fields are input-0.dc[hll_registers]`.

Add `strip_aggregate_state_suffix` after `coerce_inferred_schema` in
`derive_schema_from_partial_plan`. Splits each field name on the first `[` so the
StreamingTable's declared names match what the FINAL substrait expects. Wire data still
flows positionally via Arrow C Data — names don't affect runtime data movement;
`typesMatch` on the Java path is type-only by position.

This bridges the same physical-output ↔ declared-schema gap that pre-35ce14790c2 Java's
`coerceToDeclaredSchema` covered, but at the schema-declaration layer (per design §14.5.1
plan-authoring layer) rather than per-cell on every batch.

Tests:
* New `AggregatePlanShapeTests.testCountDistinct_1shard` and `testCountDistinct_2shard` pin
  the rewrite + structural split shape (1-shard SINGLE; 2-shard
  Aggregate(FINAL,APPROX_COUNT_DISTINCT) over Reducer over Aggregate(PARTIAL,APPROX_COUNT_DISTINCT)).
* `countDistinctCall` and `approxCountDistinctCall` helpers added to `BasePlannerRulesTests`
  for plan-shape construction.
* `testCountDistinctRewrittenToApproxCountDistinct` in `AggregateRuleTests` continues to
  validate the rewrite end-to-end through the planner — now via the HEP rule instead of
  the SPI hook.
* TwoShardAggregationIT 2-shard reduce checks: failures down from 17 → 5, all 5 remaining
  are HLL-specific Binary↔Int64 type mismatches at the FINAL substrait Read boundary
  (`Substrait error: Field 'dc' has a different type (Binary) than the corresponding
  field in the table schema (Int64)`). SUM/COUNT/AVG/MIN/MAX now all pass.
* Existing CoordinatorReduceIT `testDistinctCountAcrossShards` and
  `testDistinctCountCrossShardOverlap` / `testDistinctCountCrossShardOverlapKeyword`
  failures are HLL-specific and tracked separately — distributed sketch-merge requires
  activating the dormant SETUP_FINAL_AGGREGATE / prepareFinalPlan path.

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
…COUNT_DISTINCT

Coordinator merges per-shard HLL sketches via SETUP_FINAL_AGGREGATE +
force_aggregate_mode(Final), instead of gathering all rows then re-aggregating.
Three pieces:

(1) FragmentConversionDriver
  * Coord-side: emit SETUP_FINAL_AGGREGATE on engine-native-merge FINAL stages
    (containsEngineNativeFinalAggregate walks past Project/Sort wrappers).
  * Skip pure-reorder Projects above engine-native-merge FINAL when attaching
    fragments — DataFusion's substrait consumer can't bind their input field
    names against the FINAL aggregate's measure column. Java reduce sink
    consumes by name so the order shift is invisible.
  * Add isPureReorderProject + hasEngineNativeMergeFinalBelow helpers.

(2) Rust agg_mode
  * force_aggregate_mode(Final) now treats AggregateMode::FinalPartitioned as
    Final — the partitioned variant DataFusion picks for grouped aggregates
    consuming hash-repartitioned input. Without this, by-cat HLL stripped the
    only Final and shipped Binary state.
  * wrap_with_user_facing_names rebuilt structurally: walks to the topmost
    AggregateExec, derives expected names from each AggregateFunctionExpr's
    name() / state_fields() (single-state → final alias; multi-state → keep
    state names), and only wraps when the plan is AggregateExec(Partial).
    Replaces the prior name.contains('[') heuristic.

(3) Wire-in
  * Apply wrap_with_user_facing_names on every execute path that produces
    Partial-aggregate output: prepare_partial_plan, prepare_final_plan,
    LocalSession::execute_substrait, query_executor's two execute paths.
    No-op outside Partial mode.

Targeted IT pass: TwoShardAggregationIT all 36 reduce checks (incl.
distinct_count_by_cat), CoordinatorReduceIT 21/21 (incl. testQ10ShapeAcrossShards,
testGroupByCountMultiShard_*).

Signed-off-by: Sandesh Kumar <kusandes@amazon.com>
…UNT_DISTINCT

TopK shard-side oversampling previously failed for engine-native-merge
aggregates (APPROX_COUNT_DISTINCT/HLL) because partial state is Binary
sketch bytes that cannot be sorted directly by cardinality.

Insert reduce_eval("approx_distinct", sketch) Project between the PARTIAL
aggregate and the shard Sort to derive a sortable UInt64 cardinality from
opaque HLL state. A strip Project above Sort removes the extra column
before the wire ships [group, sketch:Binary] to the coord for final merge.

Three coordinated changes:

(1) OpenSearchTopKRewriter (Java)
  * Detect sort collation referencing engine-native-merge measures
  * Insert OpenSearchProject(reduce_eval) below Sort, OpenSearchProject(strip) above
  * Adjust collation to reference the new reduce_eval column index

(2) FragmentConversionDriver (Java)
  * Layered substrait conversion for buried PARTIAL aggregates: scan →
    attachPartialAggOnTop (INITIAL_TO_INTERMEDIATE) → attachFragmentOnTop
    per operator above, so derive_schema_from_partial_plan sees Binary type
  * containsEngineNativePartialAggregate tree-walk for SETUP_PARTIAL_AGGREGATE
  * strip() propagates stripped children through non-OpenSearch nodes
  * DAGBuilder.findFieldStorage walks past non-OpenSearch nodes

(3) agg_mode.rs (Rust)
  * force_aggregate_mode: when a ProjectionExec's child schema changes
    (names OR types), rebuild the projection with remapped Column references
    via remap_column_names — fixes the name/type mismatch between
    the reduce_eval Project and the state-suffixed Partial aggregate output

All ShardBucketOversamplingIT, TwoShardAggregationIT, CoordinatorReduceIT pass.

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
- Remove wrap_with_user_facing_names (wire is positional, names irrelevant)
- Replace strip_aggregate_state_suffix with substrait Root.names
- Deduplicate isEngineNativeMerge + REDUCE_EVAL_OP into AggregateFunction SPI
- Add reduceEvalName() — derive UDF name from enum, not hardcoded string
- Unify convert() paths (aggregate-at-top = degenerate buried case)
- Simplify partial_aggregate_schema to delegate to find_partial_input
- Remove DAGBuilder.findFieldStorage (direct cast, TopK only inserts OpenSearchRelNode)
- Remove Sort marker coupling (chain-length > 0 is sufficient)
- Trim bloated comments and remove debug loggers

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
- testMultiAgg_sortByDc_head10: dc + sum grouped, TopK sorted on dc
- dc_count_by_category: dc + count grouped without TopK (2-shard merge)

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
- testRewrite_dcByGroup_splitAndTopK: assert exact plan shape for
  dc + TopK (reduce_eval Project, strip Project, oversampled Sort)
- testRewrite_multiGroupByCount_splitAndTopK: assert PARTIAL/FINAL split
  fires for multi-group-by COUNT (was broken on main — stayed SINGLE)
- testMultiAgg_sortByDc_head10: IT for mixed dc + sum with TopK
- dc_count_by_category: golden-file IT for dc + count grouped (no TopK)
- testCountByGroup_having_sortDesc_head10: lenient assertion to tolerate
  TopK's approximate pruning while still catching double-counting bugs

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
@github-actions

github-actions Bot commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

Persistent review updated to latest commit 64b8562

@github-actions

github-actions Bot commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

❌ Gradle check result for 64b8562: 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?

@mch2
mch2 merged commit 490ba3d into opensearch-project:main Jun 6, 2026
15 of 20 checks passed
ahkcs added a commit to ahkcs/sql that referenced this pull request Jun 8, 2026
PPLBuiltinOperators.DISTINCT_COUNT_APPROX created its SqlAggFunction with the
runtime-resolution name "DISTINCT_COUNT_APPROX". The analytics-engine (DataFusion)
backend resolves aggregates by the Calcite/Substrait-standard name
APPROX_COUNT_DISTINCT, so distinct_count_approx() failed to bind on the analytics
route. Emit APPROX_COUNT_DISTINCT instead.

The Java field name stays DISTINCT_COUNT_APPROX (the PPL function name); only the
resolution string changes. The OpenSearch V3 path is unaffected — it overrides this
operator via the external HyperLogLog registration in OpenSearchExecutionEngine
(whose name is unchanged), so explain output and execution on that path are
identical (verified). The analytics-route binding is completed by
opensearch-project/OpenSearch#22013 (APPROX_COUNT_DISTINCT -> approx_distinct).

Per Sandesh Kumar.

Signed-off-by: Kai Huang <ahkcs@amazon.com>
ahkcs added a commit to ahkcs/sql that referenced this pull request Jun 8, 2026
distinct_count_approx() failed to bind on the analytics-engine (DataFusion) route
because the SqlAggFunction was named DISTINCT_COUNT_APPROX; the backend resolves
aggregates by the Calcite/Substrait-standard name APPROX_COUNT_DISTINCT. The Java
field name and PPL function name are unchanged. The OpenSearch V3 path is unaffected
(it overrides this via the external HLL registration). Analytics-route binding is
completed by opensearch-project/OpenSearch#22013. Per Sandesh Kumar.

Signed-off-by: Kai Huang <ahkcs@amazon.com>
ahkcs added a commit to opensearch-project/sql that referenced this pull request Jun 8, 2026
…stSimpleCount0 + APPROX_COUNT_DISTINCT name) (#5525)

* Use a parquet-backed index in CalcitePPLAggregationIT.testSimpleCount0

A bare auto-created index isn't composite/parquet-backed, so on the analytics-engine
route it doesn't route to the analytics engine. Switch to TEST_INDEX_BANK (loaded via
loadIndex, which injects parquet settings when the flag is set, 7 docs) so the test is
meaningful on both routes. Diagnosis by Sandesh Kumar.

Signed-off-by: Kai Huang <ahkcs@amazon.com>

* Emit APPROX_COUNT_DISTINCT as the distinct_count_approx runtime name

distinct_count_approx() failed to bind on the analytics-engine (DataFusion) route
because the SqlAggFunction was named DISTINCT_COUNT_APPROX; the backend resolves
aggregates by the Calcite/Substrait-standard name APPROX_COUNT_DISTINCT. The Java
field name and PPL function name are unchanged. The OpenSearch V3 path is unaffected
(it overrides this via the external HLL registration). Analytics-route binding is
completed by opensearch-project/OpenSearch#22013. Per Sandesh Kumar.

Signed-off-by: Kai Huang <ahkcs@amazon.com>

---------

Signed-off-by: Kai Huang <ahkcs@amazon.com>
mengweieric added a commit to mengweieric/OpenSearch that referenced this pull request Jun 10, 2026
…APPROX_COUNT_DISTINCT (opensearch-project#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/opensearch-project#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
mengweieric added a commit to mengweieric/OpenSearch that referenced this pull request Jun 11, 2026
…APPROX_COUNT_DISTINCT (opensearch-project#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/opensearch-project#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
mengweieric added a commit to mengweieric/OpenSearch that referenced this pull request Jun 12, 2026
…APPROX_COUNT_DISTINCT (opensearch-project#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/opensearch-project#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
mengweieric added a commit to mengweieric/OpenSearch that referenced this pull request Jun 12, 2026
…APPROX_COUNT_DISTINCT (opensearch-project#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/opensearch-project#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
mengweieric added a commit to mengweieric/OpenSearch that referenced this pull request Jun 16, 2026
…APPROX_COUNT_DISTINCT (opensearch-project#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/opensearch-project#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
mengweieric added a commit to mengweieric/OpenSearch that referenced this pull request Jun 16, 2026
…APPROX_COUNT_DISTINCT (opensearch-project#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/opensearch-project#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
mengweieric added a commit to mengweieric/OpenSearch that referenced this pull request Jun 16, 2026
…APPROX_COUNT_DISTINCT (opensearch-project#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/opensearch-project#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
sandeshkr419 pushed a commit that referenced this pull request Jun 18, 2026
…lytics-engine route + review fixes (#21975)

* [analytics-engine] Fix now()-family + enable now(fsp)/rand(seed)/sha2/span-month/ts-subtract/dc-approx on the analytics-engine route

Closes several PPL scalar/aggregate gaps on the analytics-engine (DataFusion)
route. Each was either an "Unable to convert call ..." Substrait failure or a
result-rendering gap; all are verified end-to-end via the CalcitePPL*IT remote
ITs on a parquet/composite (analytics-engine-routed) cluster.

now()-family (CalciteNowLikeFunctionIT 12/12):
- ArrowValues now converts Arrow Date32/Date64 -> LocalDate and
  Time{Sec,Milli,Micro,Nano} -> LocalTime so DATE/TIME results render as
  "uuuu-MM-dd" / "HH:mm:ss" via ExprValueUtils.fromObjectValue's existing
  temporal branch (mirrors how TIMESTAMP already arrives as LocalDateTime).
  Fixes current_date/curdate/current_time/curtime returning raw epoch-day /
  units-of-day integers.
- UTC_TIMESTAMP/UTC_DATE/UTC_TIME enum constants + adapter registrations route
  to the same DataFusion builtins as their CURRENT_* equivalents (cluster runs
  in UTC), reusing the existing now/currentDate/currentTime adapters.
- NowFspAdapter drops the optional fractional-seconds-precision arg so
  now(fsp)/current_timestamp(fsp)/sysdate(fsp) map to DataFusion's niladic
  now() instead of failing as "Unable to convert call now(i32)".

Other scalar/aggregate fixes:
- RandSeedAdapter drops the optional rand(seed) operand -> DataFusion random().
- Sha2FunctionAdapter raises a clear "Unsupported SHA2 algorithm [N]" for a
  concrete unsupported literal bit length, matching the SQL-plugin reference
  (CryptographicFunction) instead of surfacing a cryptic Substrait error.
- SpanAdapter supports variable-length month/quarter/year buckets.
- TimestampSubtractRewriter rewrites MINUS(timestamp, timestamp) to an
  epoch-second difference (to_unixtime), which is Substrait-convertible.
- PplAggregateCallRewriter + AggregateFunction map DISTINCT_COUNT_APPROX to
  APPROX_COUNT_DISTINCT (dc/distinct_count approx form).

Adds unit tests for the new adapters/rewriter and ArrowValues DATE/TIME
conversion; updates the SHA2 adapter test to assert the clear-error behavior.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Address review: RAND seed fail-clear, shared Substrait preprocess, adapter arity bounds, span hardening, SHA2 message

Review-driven hardening on top of the initial PR:

- RandSeedAdapter: stop silently dropping the rand(seed) operand (which turned a
  deterministic seeded call into a non-deterministic one). Niladic rand() still
  maps to DataFusion random(); seeded rand(seed) now fails with a clear
  unsupported-shape error until a seeded random is available on the backend.
- DataFusionFragmentConvertor: extract preprocessForSubstrait() and call it from
  both convertToSubstrait and convertStandalone so TimestampSubtractRewriter
  (and any future rewriter) runs on the wrapper/partial-aggregate path too, not
  only the top-level fragment path.
- NowFspAdapter / RandSeedAdapter: only normalize the valid 0-arg / 1-arg shapes;
  leave unexpected arities untouched instead of inventing a valid call.
- SpanAdapter: guard the month and second interval multiplications with
  Math.multiplyExact (clear "interval is too large" error on overflow), and use
  floored modulo/division for the month-index bucketing so pre-1970 (negative
  month index) timestamps snap down to the correct bucket start.
- Sha2FunctionAdapter: list the supported bit lengths (224, 256, 384, 512) in the
  unsupported-algorithm error.

Tests: add NowFsp unexpected-arity test, RAND seeded-reject test, span
month-bucket + overflow tests; update SHA2 test to assert the supported-values
message.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Fix grouped distinct_count_approx nullability + make TimestampSubtractRewriter identity-preserving

Two correctness fixes found while verifying the review round:

- PplAggregateCallRewriter (DISTINCT_COUNT_APPROX): pin the remapped aggregate's
  explicit return type to NOT NULL BIGINT. APPROX_COUNT_DISTINCT is a count, so its
  return-type inference is BIGINT NOT NULL; reusing the PPL call's nullable BIGINT
  left the declared and inferred types disagreeing and tripped Calcite's validity
  assertion ("aggCall type BIGINT vs inferred BIGINT NOT NULL") on the grouped
  two-phase path (e.g. `stats distinct_count_approx(x) by g`). This was a latent bug
  in the un-grouped-only test coverage; the grouped form is now verified green.

- TimestampSubtractRewriter: make it a true identity no-op for plans without a
  MINUS(timestamp, timestamp). Previously it walked every RelNode and called
  RelNode.accept(RexShuttle) unconditionally, which re-derives expression / aggCall
  types and can flip a cached nullable BIGINT to BIGINT NOT NULL — breaking unrelated
  shapes once the shared preprocessing pipeline runs it on the two-phase aggregate
  path. It now detects the target shape first (read-only visitors) and only applies
  the rewriting RexShuttle to nodes that actually contain it, returning the original
  object otherwise. Adds identity (assertSame) regression tests.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Drop TimestampSubtractRewriter — superseded by upstream MinusAdapter (#21978)

Upstream #21978 added MinusAdapter, the proper named ScalarFunctionAdapter for
timestamp/date subtraction (t1 - t2 -> from_unixtime(to_unixtime(t1) - to_unixtime(t2))).
It also coordinates with WidthBucketAdapter by deliberately leaving the
MINUS(MAX OVER(), MIN OVER()) binning shape untouched. Our pre-Substrait
TimestampSubtractRewriter rewrote every MINUS(timestamp, timestamp) including that
binning shape, so keeping both risked clobbering the binning path. Remove the
rewriter (and its test); MinusAdapter fully covers the case. The shared
preprocessForSubstrait helper is retained for the remaining rewriters.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Bind distinct_count_approx after SQL-plugin emits APPROX_COUNT_DISTINCT (#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Don't push filters below non-deterministic projects (fixes rand() predicates)

`eval r = rand() | where r > 0` failed with "Comparison performance-delegation requires
(RexInputRef, RexLiteral); got RAND()". Calcite's stock FILTER_PROJECT_TRANSPOSE only guards
against window functions (!containsOver()), so it pushed the filter below the rand() project,
inlining RAND() into the predicate — turning a delegatable ($ref > literal) comparison into
RAND() > literal on the scan, which Lucene performance-delegation cannot serialize (and which
must stay on the single in-memory engine, since each backend would draw different values).

Replace the stock rule with FILTER_PROJECT_TRANSPOSE_DETERMINISTIC: a FilterProjectTransposeRule
configured to refuse the transpose when the Project computes any non-deterministic expression
(RexUtil.isDeterministic). Keeping the Filter above the Project preserves the clean
($ref > literal) shape evaluated in memory.

Verified: CalcitePPLBuiltinFunctionIT.testRand passes on a force-routed analytics-engine cluster.
Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Restore distinct_count_approx binding on the analytics-engine route

Validation requested in PR review (sandeshkr419) found that sql#5525 alone does
NOT make `stats distinct_count_approx(x)` work on the analytics-engine route, so
the engine-side handling removed in the prior commit is still required. On a
fresh, force-routed cluster with sql#5525 deployed, CalcitePPLAggregationIT
testCountDistinctApprox + testCountDistinctApproxWithAlias fail in two stages:

  1. Planner stage: OpenSearchAggregateRule.resolveViableBackendsForCall reads
     aggCall.getAggregation().getName(), which is still "DISTINCT_COUNT_APPROX"
     (the DistinctCountApproxLogicalAggFunction marker). fromNameOrError then
     throws "No enum constant ...AggregateFunction.DISTINCT_COUNT_APPROX".
  2. Substrait stage: the PplAggregateCallRewriter case keyed only on
     "APPROX_COUNT_DISTINCT" never matched the marker, so the unbound op reached
     isthmus -> "Unable to find binding for call DISTINCT_COUNT_APPROX($1)".

sql#5525 only sets the substrait-emission name to APPROX_COUNT_DISTINCT, which is
a later stage than both of the above, so the marker's runtime name still reaches
them. Fixes:

  - Restore the DISTINCT_COUNT_APPROX -> APPROX_COUNT_DISTINCT alias in
    AggregateFunction.fromNameOrError (covers the planner path).
  - Match both "DISTINCT_COUNT_APPROX" and "APPROX_COUNT_DISTINCT" in the
    PplAggregateCallRewriter case (covers the substrait-binding path); keep the
    guard that skips the already-bound stock operator.
  - Add a fromNameOrError alias unit test.

Verified on a force-routed analytics-engine cluster (parquet-backed indices,
AE-routed): both tests pass; the surrounding 4-class IT sweep is unchanged at
6 pre-existing failures (date_format/strftime/percentile/nested-field), i.e. no
regressions.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Add qa-module ITs for distinct_count_approx + min/max(boolean)

Engine-side end-to-end regression coverage for the two analytics-engine binding
fixes in this PR, run against the production /_plugins/_ppl surface on the
qa-module's analytics-routed cluster (mirrors the SQL-repo CalcitePPLAggregationIT
cases that proved the fixes):

  - DistinctCountApproxAggregationIT: grouped / aliased / dc-shorthand / ungrouped
    distinct_count_approx (guards the planner enum alias + rewriter rebind).
  - MinMaxBooleanAggregationIT: ungrouped / grouped / mixed boolean+string
    min/max (guards the boolean min/max substrait overload).

Each IT creates its own parquet-backed (composite/parquet + lucene secondary)
index, bulk-ingests a small fixed dataset, and asserts oracle values. All 7
tests pass (4 + 3).

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Fix date_format %c month padding; correct stale RAND comments

Two small follow-ups surfaced while verifying the analytics-engine route:

  - date_format/strftime %c now emits the zero-padded month (Jan -> "01"),
    matching the PPL reference (SQL-plugin DateTimeFormatterUtil maps %c -> "MM"
    in date mode), which differs from stock MySQL's no-leading-zero %c. Time mode
    keeps the reference's single-"0" literal. Fixes the %c portion of
    CalciteDateTimeFunctionIT.testDateFormat and updates the Rust token-set unit
    test expectation accordingly. (The test's %U/%u/%V/%v week-number padding is a
    separate, unrelated date-format issue left for a follow-up.)

  - Correct two stale comments that described the old "drop the seed" RAND
    behavior; the adapter now rejects RAND(seed) with a clear error (addresses
    review feedback). No behavior change — comment/javadoc only.

Verified on a force-routed analytics-engine cluster: date_format(ts,'%c') returns
"01"; the now/sha2/dc/min-max ITs remain green (no regressions).

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Drop DISTINCT_COUNT_APPROX enum alias per review

Remove the DISTINCT_COUNT_APPROX -> APPROX_COUNT_DISTINCT alias in
AggregateFunction.fromNameOrError (and its unit test), and key the
PplAggregateCallRewriter case on "APPROX_COUNT_DISTINCT" only, per
@sandeshkr419's review. The SQL plugin (sql#5525) emits the operator named
APPROX_COUNT_DISTINCT, which the enum resolves directly, so the legacy-name
alias is no longer needed.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Address review: bulk NDJSON content-type, dc comments, error wording

Review follow-ups (sandeshkr419):

  - QA ITs: send _bulk as application/x-ndjson via the per-index endpoint with bare
    {"index": {}} action metadata, matching the existing CountFastPathIT pattern.
    setJsonEntity alone defaults to application/json, which the bulk endpoint rejects.
  - Clarify the distinct_count_approx marker naming across layers: current runtime name
    is APPROX_COUNT_DISTINCT (sql#5525), DISTINCT_COUNT_APPROX is the legacy/defensive
    spelling; OpenSearchDistinctCountRule is the primary planner rewrite and
    PplAggregateCallRewriter is a late Substrait-emission defensive fallback.
  - Document the phase-specific nullability handling: the early planner rewrite preserves
    the marker's nullable type (LogicalAggregate.copy re-validates row type); the late
    Substrait rewrite pins NOT NULL (Calcite validates against the stock op's inferred type).
  - SPAN overflow message now includes the operands.
  - Drop "yet" from the seeded-RAND error so it doesn't imply imminent support.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Empty commit to retrigger CI

No code change. The prior gradle-check failed in the Jenkins-trigger step
(jq parse error polling the Jenkins API, 10 retries, empty result) — a
transient CI-infra glitch unrelated to this PR's changes. Re-running checks.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Fix DatetimeCoverageIT %c assertions to match zero-padded month

The date_format %c code change (no-leading-zero -> zero-padded "MM") aligns the
analytics-engine Rust UDF with the legacy PPL contract: DateTimeFormatterUtil.DATE_HANDLERS
maps %c -> "MM" (date_format dispatches through DATE_HANDLERS), so January renders as "01".
The Rust code + unit test were updated for this, but the two DatetimeCoverageIT assertions
still expected the old unpadded "1" — the gradle-check / sandbox-check failures.

Update both assertions to "01", matching the canonical merged ITs (ppl/DateTimeFunctionIT
and sql/DateTimeFunctionIT both assert "Sat Jan 01 31st ..." for the same spec). Test-only.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Enable now()/rand(seed)/sha2/span-month on the analytics-engine route + review fixes

Enables these PPL functions on the force-routed analytics-engine (DataFusion)
path and addresses review feedback:

- now()-family: NowFspAdapter handles now([fsp]) / current_timestamp /
  localtimestamp, where fsp is MySQL fractional-seconds precision (0-6).
- rand([seed]): RandSeedAdapter.
- sha2(input, bitLen): Sha2FunctionAdapter -> encode(digest(input, 'shaN'), 'hex'),
  with a backend fail-clear guard for unsupported bit lengths (TODO to move
  validation to the frontend once the handoff carries function metadata).
- span() month bucketing: SpanAdapter.
- strftime: os_strftime.rs.
- date_format %c: zero-padded month (DatetimeCoverageIT updated to match the
  merged ppl/sql DateTimeFunctionIT expectations).
- PlannerImpl: don't push filters below non-deterministic projects; rand() in a
  pushed-down predicate would be re-evaluated and draw a fresh value. Documented
  as a semantic-correctness guard.
- DataFusionFragmentConvertor: shared Substrait preprocess entry point.
- New QA ITs: MinMaxBooleanAggregationIT.

Review fixes: removed the late defensive APPROX_COUNT_DISTINCT fallback in
PplAggregateCallRewriter (fail loud instead); restored the Substrait-layer
assessment TODO; clarified fsp in NowFspAdapter; tidied AggregateFunction.

distinct_count_approx on the analytics-engine route is handled upstream by #22120
and is not part of this PR.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Add 2-shard reduce coverage for boolean min/max

Addresses review feedback to cover the multi-shard case for boolean min/max.
min(flag)/max(flag) over the merge_coverage dataset's boolean field are added as
golden cases in the TwoShardReduceTestCase suite (#21951 convention), so they run
at 1 shard and 2 shards with a differential equality check plus a pinned golden,
exercising the two-phase per-shard-partial then coordinator-merge reduce path.
The merge_coverage flag field has 15 false and 15 true rows, so min=false and
max=true regardless of shard layout.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

---------

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
OVyshnevskyi pushed a commit to OVyshnevskyi/OpenSearch that referenced this pull request Jun 22, 2026
…lytics-engine route + review fixes (opensearch-project#21975)

* [analytics-engine] Fix now()-family + enable now(fsp)/rand(seed)/sha2/span-month/ts-subtract/dc-approx on the analytics-engine route

Closes several PPL scalar/aggregate gaps on the analytics-engine (DataFusion)
route. Each was either an "Unable to convert call ..." Substrait failure or a
result-rendering gap; all are verified end-to-end via the CalcitePPL*IT remote
ITs on a parquet/composite (analytics-engine-routed) cluster.

now()-family (CalciteNowLikeFunctionIT 12/12):
- ArrowValues now converts Arrow Date32/Date64 -> LocalDate and
  Time{Sec,Milli,Micro,Nano} -> LocalTime so DATE/TIME results render as
  "uuuu-MM-dd" / "HH:mm:ss" via ExprValueUtils.fromObjectValue's existing
  temporal branch (mirrors how TIMESTAMP already arrives as LocalDateTime).
  Fixes current_date/curdate/current_time/curtime returning raw epoch-day /
  units-of-day integers.
- UTC_TIMESTAMP/UTC_DATE/UTC_TIME enum constants + adapter registrations route
  to the same DataFusion builtins as their CURRENT_* equivalents (cluster runs
  in UTC), reusing the existing now/currentDate/currentTime adapters.
- NowFspAdapter drops the optional fractional-seconds-precision arg so
  now(fsp)/current_timestamp(fsp)/sysdate(fsp) map to DataFusion's niladic
  now() instead of failing as "Unable to convert call now(i32)".

Other scalar/aggregate fixes:
- RandSeedAdapter drops the optional rand(seed) operand -> DataFusion random().
- Sha2FunctionAdapter raises a clear "Unsupported SHA2 algorithm [N]" for a
  concrete unsupported literal bit length, matching the SQL-plugin reference
  (CryptographicFunction) instead of surfacing a cryptic Substrait error.
- SpanAdapter supports variable-length month/quarter/year buckets.
- TimestampSubtractRewriter rewrites MINUS(timestamp, timestamp) to an
  epoch-second difference (to_unixtime), which is Substrait-convertible.
- PplAggregateCallRewriter + AggregateFunction map DISTINCT_COUNT_APPROX to
  APPROX_COUNT_DISTINCT (dc/distinct_count approx form).

Adds unit tests for the new adapters/rewriter and ArrowValues DATE/TIME
conversion; updates the SHA2 adapter test to assert the clear-error behavior.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Address review: RAND seed fail-clear, shared Substrait preprocess, adapter arity bounds, span hardening, SHA2 message

Review-driven hardening on top of the initial PR:

- RandSeedAdapter: stop silently dropping the rand(seed) operand (which turned a
  deterministic seeded call into a non-deterministic one). Niladic rand() still
  maps to DataFusion random(); seeded rand(seed) now fails with a clear
  unsupported-shape error until a seeded random is available on the backend.
- DataFusionFragmentConvertor: extract preprocessForSubstrait() and call it from
  both convertToSubstrait and convertStandalone so TimestampSubtractRewriter
  (and any future rewriter) runs on the wrapper/partial-aggregate path too, not
  only the top-level fragment path.
- NowFspAdapter / RandSeedAdapter: only normalize the valid 0-arg / 1-arg shapes;
  leave unexpected arities untouched instead of inventing a valid call.
- SpanAdapter: guard the month and second interval multiplications with
  Math.multiplyExact (clear "interval is too large" error on overflow), and use
  floored modulo/division for the month-index bucketing so pre-1970 (negative
  month index) timestamps snap down to the correct bucket start.
- Sha2FunctionAdapter: list the supported bit lengths (224, 256, 384, 512) in the
  unsupported-algorithm error.

Tests: add NowFsp unexpected-arity test, RAND seeded-reject test, span
month-bucket + overflow tests; update SHA2 test to assert the supported-values
message.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Fix grouped distinct_count_approx nullability + make TimestampSubtractRewriter identity-preserving

Two correctness fixes found while verifying the review round:

- PplAggregateCallRewriter (DISTINCT_COUNT_APPROX): pin the remapped aggregate's
  explicit return type to NOT NULL BIGINT. APPROX_COUNT_DISTINCT is a count, so its
  return-type inference is BIGINT NOT NULL; reusing the PPL call's nullable BIGINT
  left the declared and inferred types disagreeing and tripped Calcite's validity
  assertion ("aggCall type BIGINT vs inferred BIGINT NOT NULL") on the grouped
  two-phase path (e.g. `stats distinct_count_approx(x) by g`). This was a latent bug
  in the un-grouped-only test coverage; the grouped form is now verified green.

- TimestampSubtractRewriter: make it a true identity no-op for plans without a
  MINUS(timestamp, timestamp). Previously it walked every RelNode and called
  RelNode.accept(RexShuttle) unconditionally, which re-derives expression / aggCall
  types and can flip a cached nullable BIGINT to BIGINT NOT NULL — breaking unrelated
  shapes once the shared preprocessing pipeline runs it on the two-phase aggregate
  path. It now detects the target shape first (read-only visitors) and only applies
  the rewriting RexShuttle to nodes that actually contain it, returning the original
  object otherwise. Adds identity (assertSame) regression tests.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Drop TimestampSubtractRewriter — superseded by upstream MinusAdapter (opensearch-project#21978)

Upstream opensearch-project#21978 added MinusAdapter, the proper named ScalarFunctionAdapter for
timestamp/date subtraction (t1 - t2 -> from_unixtime(to_unixtime(t1) - to_unixtime(t2))).
It also coordinates with WidthBucketAdapter by deliberately leaving the
MINUS(MAX OVER(), MIN OVER()) binning shape untouched. Our pre-Substrait
TimestampSubtractRewriter rewrote every MINUS(timestamp, timestamp) including that
binning shape, so keeping both risked clobbering the binning path. Remove the
rewriter (and its test); MinusAdapter fully covers the case. The shared
preprocessForSubstrait helper is retained for the remaining rewriters.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Bind distinct_count_approx after SQL-plugin emits APPROX_COUNT_DISTINCT (opensearch-project#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/opensearch-project#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Don't push filters below non-deterministic projects (fixes rand() predicates)

`eval r = rand() | where r > 0` failed with "Comparison performance-delegation requires
(RexInputRef, RexLiteral); got RAND()". Calcite's stock FILTER_PROJECT_TRANSPOSE only guards
against window functions (!containsOver()), so it pushed the filter below the rand() project,
inlining RAND() into the predicate — turning a delegatable ($ref > literal) comparison into
RAND() > literal on the scan, which Lucene performance-delegation cannot serialize (and which
must stay on the single in-memory engine, since each backend would draw different values).

Replace the stock rule with FILTER_PROJECT_TRANSPOSE_DETERMINISTIC: a FilterProjectTransposeRule
configured to refuse the transpose when the Project computes any non-deterministic expression
(RexUtil.isDeterministic). Keeping the Filter above the Project preserves the clean
($ref > literal) shape evaluated in memory.

Verified: CalcitePPLBuiltinFunctionIT.testRand passes on a force-routed analytics-engine cluster.
Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Restore distinct_count_approx binding on the analytics-engine route

Validation requested in PR review (sandeshkr419) found that sql#5525 alone does
NOT make `stats distinct_count_approx(x)` work on the analytics-engine route, so
the engine-side handling removed in the prior commit is still required. On a
fresh, force-routed cluster with sql#5525 deployed, CalcitePPLAggregationIT
testCountDistinctApprox + testCountDistinctApproxWithAlias fail in two stages:

  1. Planner stage: OpenSearchAggregateRule.resolveViableBackendsForCall reads
     aggCall.getAggregation().getName(), which is still "DISTINCT_COUNT_APPROX"
     (the DistinctCountApproxLogicalAggFunction marker). fromNameOrError then
     throws "No enum constant ...AggregateFunction.DISTINCT_COUNT_APPROX".
  2. Substrait stage: the PplAggregateCallRewriter case keyed only on
     "APPROX_COUNT_DISTINCT" never matched the marker, so the unbound op reached
     isthmus -> "Unable to find binding for call DISTINCT_COUNT_APPROX($1)".

sql#5525 only sets the substrait-emission name to APPROX_COUNT_DISTINCT, which is
a later stage than both of the above, so the marker's runtime name still reaches
them. Fixes:

  - Restore the DISTINCT_COUNT_APPROX -> APPROX_COUNT_DISTINCT alias in
    AggregateFunction.fromNameOrError (covers the planner path).
  - Match both "DISTINCT_COUNT_APPROX" and "APPROX_COUNT_DISTINCT" in the
    PplAggregateCallRewriter case (covers the substrait-binding path); keep the
    guard that skips the already-bound stock operator.
  - Add a fromNameOrError alias unit test.

Verified on a force-routed analytics-engine cluster (parquet-backed indices,
AE-routed): both tests pass; the surrounding 4-class IT sweep is unchanged at
6 pre-existing failures (date_format/strftime/percentile/nested-field), i.e. no
regressions.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Add qa-module ITs for distinct_count_approx + min/max(boolean)

Engine-side end-to-end regression coverage for the two analytics-engine binding
fixes in this PR, run against the production /_plugins/_ppl surface on the
qa-module's analytics-routed cluster (mirrors the SQL-repo CalcitePPLAggregationIT
cases that proved the fixes):

  - DistinctCountApproxAggregationIT: grouped / aliased / dc-shorthand / ungrouped
    distinct_count_approx (guards the planner enum alias + rewriter rebind).
  - MinMaxBooleanAggregationIT: ungrouped / grouped / mixed boolean+string
    min/max (guards the boolean min/max substrait overload).

Each IT creates its own parquet-backed (composite/parquet + lucene secondary)
index, bulk-ingests a small fixed dataset, and asserts oracle values. All 7
tests pass (4 + 3).

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Fix date_format %c month padding; correct stale RAND comments

Two small follow-ups surfaced while verifying the analytics-engine route:

  - date_format/strftime %c now emits the zero-padded month (Jan -> "01"),
    matching the PPL reference (SQL-plugin DateTimeFormatterUtil maps %c -> "MM"
    in date mode), which differs from stock MySQL's no-leading-zero %c. Time mode
    keeps the reference's single-"0" literal. Fixes the %c portion of
    CalciteDateTimeFunctionIT.testDateFormat and updates the Rust token-set unit
    test expectation accordingly. (The test's %U/%u/%V/%v week-number padding is a
    separate, unrelated date-format issue left for a follow-up.)

  - Correct two stale comments that described the old "drop the seed" RAND
    behavior; the adapter now rejects RAND(seed) with a clear error (addresses
    review feedback). No behavior change — comment/javadoc only.

Verified on a force-routed analytics-engine cluster: date_format(ts,'%c') returns
"01"; the now/sha2/dc/min-max ITs remain green (no regressions).

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Drop DISTINCT_COUNT_APPROX enum alias per review

Remove the DISTINCT_COUNT_APPROX -> APPROX_COUNT_DISTINCT alias in
AggregateFunction.fromNameOrError (and its unit test), and key the
PplAggregateCallRewriter case on "APPROX_COUNT_DISTINCT" only, per
@sandeshkr419's review. The SQL plugin (sql#5525) emits the operator named
APPROX_COUNT_DISTINCT, which the enum resolves directly, so the legacy-name
alias is no longer needed.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Address review: bulk NDJSON content-type, dc comments, error wording

Review follow-ups (sandeshkr419):

  - QA ITs: send _bulk as application/x-ndjson via the per-index endpoint with bare
    {"index": {}} action metadata, matching the existing CountFastPathIT pattern.
    setJsonEntity alone defaults to application/json, which the bulk endpoint rejects.
  - Clarify the distinct_count_approx marker naming across layers: current runtime name
    is APPROX_COUNT_DISTINCT (sql#5525), DISTINCT_COUNT_APPROX is the legacy/defensive
    spelling; OpenSearchDistinctCountRule is the primary planner rewrite and
    PplAggregateCallRewriter is a late Substrait-emission defensive fallback.
  - Document the phase-specific nullability handling: the early planner rewrite preserves
    the marker's nullable type (LogicalAggregate.copy re-validates row type); the late
    Substrait rewrite pins NOT NULL (Calcite validates against the stock op's inferred type).
  - SPAN overflow message now includes the operands.
  - Drop "yet" from the seeded-RAND error so it doesn't imply imminent support.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Empty commit to retrigger CI

No code change. The prior gradle-check failed in the Jenkins-trigger step
(jq parse error polling the Jenkins API, 10 retries, empty result) — a
transient CI-infra glitch unrelated to this PR's changes. Re-running checks.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Fix DatetimeCoverageIT %c assertions to match zero-padded month

The date_format %c code change (no-leading-zero -> zero-padded "MM") aligns the
analytics-engine Rust UDF with the legacy PPL contract: DateTimeFormatterUtil.DATE_HANDLERS
maps %c -> "MM" (date_format dispatches through DATE_HANDLERS), so January renders as "01".
The Rust code + unit test were updated for this, but the two DatetimeCoverageIT assertions
still expected the old unpadded "1" — the gradle-check / sandbox-check failures.

Update both assertions to "01", matching the canonical merged ITs (ppl/DateTimeFunctionIT
and sql/DateTimeFunctionIT both assert "Sat Jan 01 31st ..." for the same spec). Test-only.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Enable now()/rand(seed)/sha2/span-month on the analytics-engine route + review fixes

Enables these PPL functions on the force-routed analytics-engine (DataFusion)
path and addresses review feedback:

- now()-family: NowFspAdapter handles now([fsp]) / current_timestamp /
  localtimestamp, where fsp is MySQL fractional-seconds precision (0-6).
- rand([seed]): RandSeedAdapter.
- sha2(input, bitLen): Sha2FunctionAdapter -> encode(digest(input, 'shaN'), 'hex'),
  with a backend fail-clear guard for unsupported bit lengths (TODO to move
  validation to the frontend once the handoff carries function metadata).
- span() month bucketing: SpanAdapter.
- strftime: os_strftime.rs.
- date_format %c: zero-padded month (DatetimeCoverageIT updated to match the
  merged ppl/sql DateTimeFunctionIT expectations).
- PlannerImpl: don't push filters below non-deterministic projects; rand() in a
  pushed-down predicate would be re-evaluated and draw a fresh value. Documented
  as a semantic-correctness guard.
- DataFusionFragmentConvertor: shared Substrait preprocess entry point.
- New QA ITs: MinMaxBooleanAggregationIT.

Review fixes: removed the late defensive APPROX_COUNT_DISTINCT fallback in
PplAggregateCallRewriter (fail loud instead); restored the Substrait-layer
assessment TODO; clarified fsp in NowFspAdapter; tidied AggregateFunction.

distinct_count_approx on the analytics-engine route is handled upstream by opensearch-project#22120
and is not part of this PR.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Add 2-shard reduce coverage for boolean min/max

Addresses review feedback to cover the multi-shard case for boolean min/max.
min(flag)/max(flag) over the merge_coverage dataset's boolean field are added as
golden cases in the TwoShardReduceTestCase suite (opensearch-project#21951 convention), so they run
at 1 shard and 2 shards with a differential equality check plus a pinned golden,
exercising the two-phase per-shard-partial then coordinator-merge reduce path.
The merge_coverage flag field has 15 false and 15 true rows, so min=false and
max=true regardless of shard layout.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

---------

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
KhishorekumarBS pushed a commit to KhishorekumarBS/OpenSearch that referenced this pull request Jul 3, 2026
…sketch merge + TopK optimization (opensearch-project#22013)

* feat(analytics-engine): rewrite COUNT(DISTINCT) → APPROX_COUNT_DISTINCT via HEP rule + strip aggregate state suffix in coord StreamingTable schema

Two complementary fixes for the analytics-engine cross-shard aggregate path.

(1) Plan-layer rewrite (Java, OpenSearchDistinctCountRule)

PPL `dc(x)` / `distinct_count(x)` parse to `COUNT(DISTINCT x)` at Calcite via
the SQL plugin's `AstExpressionBuilder.visitDistinctCountFunctionCall`. Without
intervention the call lands as `SqlStdOperatorTable.COUNT` with isDistinct=true,
the additive split rule decomposes it (PARTIAL count(distinct) + FINAL SUM-of-counts),
and cross-shard reduce over-counts any value present on more than one shard.

Replace 6d98's `AggregateFunction.resolveOperator` SPI hook (which polluted the
framework enum and rebuilt AggregateCalls with the original COUNT return type, risking
`Aggregate.typeMatchesInferred` mismatches against APPROX_COUNT_DISTINCT's
inferReturnType) with a dedicated HEP rule:

* New `OpenSearchDistinctCountRule` matches plain `LogicalAggregate` containing a
  single-arg `COUNT(DISTINCT x)` and rewrites it to `APPROX_COUNT_DISTINCT(x)`
  (isDistinct=false on the rewritten call). Built via the long-form `AggregateCall.create`
  with `(groupCount, input, type=null)` so Calcite re-infers the return type from
  `APPROX_COUNT_DISTINCT.inferReturnType(...)` — avoids the typeMatchesInferred mismatch
  the SPI rebuild would produce.
* Wired into the existing `PlannerImpl.decomposeAggregates` HEP phase alongside
  `OpenSearchAggregateReduceRule`, before `OpenSearchAggregateRule` marks the aggregate.
  Multi-arg `COUNT(DISTINCT a, b)` doesn't match — falls through to the residual
  `aggCall.isDistinct()` skip in `OpenSearchAggregateSplitRule`.

After this, the aggregate engages the existing
`AggregateFunction.APPROX_COUNT_DISTINCT(Type.APPROXIMATE, intermediateFields=[sketch:Binary, reducer==self])`
registration — capability resolution, structural split, `DistributedAggregateRewriter.overrideExchangeType`,
and the Substrait extension YAML `approx_distinct` alias all light up automatically.

* Drop `AggregateFunction.COUNT.resolveOperator` override + the default `resolveOperator`
  method on the enum.
* Drop `OpenSearchAggregateRule.resolveAggregateCall` + its unused `SqlAggFunction` import.
* Drop `Type.APPROXIMATE` from `OpenSearchAggregateSplitRule.shouldSkipPartialFinalSplit`
  (residual `aggCall.isDistinct()` skip stays for multi-arg fallback). APPROXIMATE now goes
  through the structural PARTIAL/FINAL split.

(2) Wire-layer schema bridge (Rust, derive_schema_from_partial_plan)

Commit `35ce14790c2` (opensearch-project#21690) folded the consumer's StreamingTable schema derivation into Rust
to eliminate per-cell coercion overhead; `coerceToDeclaredSchema` on the Java
`feedToSender` path was deleted in favour of running the producer's substrait through
DataFusion's substrait consumer + physical planner and using its output schema verbatim.

The DataFusion 53.x physical planner emits `AggregateExec(Mode::Partial)` columns with
state-suffixed names (`dc[hll_registers]`, `$f0[sum]`, `count(opt)[count]`), but the
FINAL substrait emitted by `attachPartialAggOnTop` declares the user-facing aliases
(`dc`, `$f0`, `count(opt)`). DataFusion's substrait consumer name-resolves the FINAL
Read against the registered StreamingTable and fails with
`Schema error: No field named dc. Valid fields are input-0.dc[hll_registers]`.

Add `strip_aggregate_state_suffix` after `coerce_inferred_schema` in
`derive_schema_from_partial_plan`. Splits each field name on the first `[` so the
StreamingTable's declared names match what the FINAL substrait expects. Wire data still
flows positionally via Arrow C Data — names don't affect runtime data movement;
`typesMatch` on the Java path is type-only by position.

This bridges the same physical-output ↔ declared-schema gap that pre-35ce14790c2 Java's
`coerceToDeclaredSchema` covered, but at the schema-declaration layer (per design §14.5.1
plan-authoring layer) rather than per-cell on every batch.

Tests:
* New `AggregatePlanShapeTests.testCountDistinct_1shard` and `testCountDistinct_2shard` pin
  the rewrite + structural split shape (1-shard SINGLE; 2-shard
  Aggregate(FINAL,APPROX_COUNT_DISTINCT) over Reducer over Aggregate(PARTIAL,APPROX_COUNT_DISTINCT)).
* `countDistinctCall` and `approxCountDistinctCall` helpers added to `BasePlannerRulesTests`
  for plan-shape construction.
* `testCountDistinctRewrittenToApproxCountDistinct` in `AggregateRuleTests` continues to
  validate the rewrite end-to-end through the planner — now via the HEP rule instead of
  the SPI hook.
* TwoShardAggregationIT 2-shard reduce checks: failures down from 17 → 5, all 5 remaining
  are HLL-specific Binary↔Int64 type mismatches at the FINAL substrait Read boundary
  (`Substrait error: Field 'dc' has a different type (Binary) than the corresponding
  field in the table schema (Int64)`). SUM/COUNT/AVG/MIN/MAX now all pass.
* Existing CoordinatorReduceIT `testDistinctCountAcrossShards` and
  `testDistinctCountCrossShardOverlap` / `testDistinctCountCrossShardOverlapKeyword`
  failures are HLL-specific and tracked separately — distributed sketch-merge requires
  activating the dormant SETUP_FINAL_AGGREGATE / prepareFinalPlan path.

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>

* feat(analytics-engine): distributed HLL sketch-merge for dc()/APPROX_COUNT_DISTINCT

Coordinator merges per-shard HLL sketches via SETUP_FINAL_AGGREGATE +
force_aggregate_mode(Final), instead of gathering all rows then re-aggregating.
Three pieces:

(1) FragmentConversionDriver
  * Coord-side: emit SETUP_FINAL_AGGREGATE on engine-native-merge FINAL stages
    (containsEngineNativeFinalAggregate walks past Project/Sort wrappers).
  * Skip pure-reorder Projects above engine-native-merge FINAL when attaching
    fragments — DataFusion's substrait consumer can't bind their input field
    names against the FINAL aggregate's measure column. Java reduce sink
    consumes by name so the order shift is invisible.
  * Add isPureReorderProject + hasEngineNativeMergeFinalBelow helpers.

(2) Rust agg_mode
  * force_aggregate_mode(Final) now treats AggregateMode::FinalPartitioned as
    Final — the partitioned variant DataFusion picks for grouped aggregates
    consuming hash-repartitioned input. Without this, by-cat HLL stripped the
    only Final and shipped Binary state.
  * wrap_with_user_facing_names rebuilt structurally: walks to the topmost
    AggregateExec, derives expected names from each AggregateFunctionExpr's
    name() / state_fields() (single-state → final alias; multi-state → keep
    state names), and only wraps when the plan is AggregateExec(Partial).
    Replaces the prior name.contains('[') heuristic.

(3) Wire-in
  * Apply wrap_with_user_facing_names on every execute path that produces
    Partial-aggregate output: prepare_partial_plan, prepare_final_plan,
    LocalSession::execute_substrait, query_executor's two execute paths.
    No-op outside Partial mode.

Targeted IT pass: TwoShardAggregationIT all 36 reduce checks (incl.
distinct_count_by_cat), CoordinatorReduceIT 21/21 (incl. testQ10ShapeAcrossShards,
testGroupByCountMultiShard_*).

Signed-off-by: Sandesh Kumar <kusandes@amazon.com>

* feat(analytics-engine): wire reduce_eval into TopK for dc()/APPROX_COUNT_DISTINCT

TopK shard-side oversampling previously failed for engine-native-merge
aggregates (APPROX_COUNT_DISTINCT/HLL) because partial state is Binary
sketch bytes that cannot be sorted directly by cardinality.

Insert reduce_eval("approx_distinct", sketch) Project between the PARTIAL
aggregate and the shard Sort to derive a sortable UInt64 cardinality from
opaque HLL state. A strip Project above Sort removes the extra column
before the wire ships [group, sketch:Binary] to the coord for final merge.

Three coordinated changes:

(1) OpenSearchTopKRewriter (Java)
  * Detect sort collation referencing engine-native-merge measures
  * Insert OpenSearchProject(reduce_eval) below Sort, OpenSearchProject(strip) above
  * Adjust collation to reference the new reduce_eval column index

(2) FragmentConversionDriver (Java)
  * Layered substrait conversion for buried PARTIAL aggregates: scan →
    attachPartialAggOnTop (INITIAL_TO_INTERMEDIATE) → attachFragmentOnTop
    per operator above, so derive_schema_from_partial_plan sees Binary type
  * containsEngineNativePartialAggregate tree-walk for SETUP_PARTIAL_AGGREGATE
  * strip() propagates stripped children through non-OpenSearch nodes
  * DAGBuilder.findFieldStorage walks past non-OpenSearch nodes

(3) agg_mode.rs (Rust)
  * force_aggregate_mode: when a ProjectionExec's child schema changes
    (names OR types), rebuild the projection with remapped Column references
    via remap_column_names — fixes the name/type mismatch between
    the reduce_eval Project and the state-suffixed Partial aggregate output

All ShardBucketOversamplingIT, TwoShardAggregationIT, CoordinatorReduceIT pass.

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>

* refactor: harden and simplify dc/TopK implementation

- Remove wrap_with_user_facing_names (wire is positional, names irrelevant)
- Replace strip_aggregate_state_suffix with substrait Root.names
- Deduplicate isEngineNativeMerge + REDUCE_EVAL_OP into AggregateFunction SPI
- Add reduceEvalName() — derive UDF name from enum, not hardcoded string
- Unify convert() paths (aggregate-at-top = degenerate buried case)
- Simplify partial_aggregate_schema to delegate to find_partial_input
- Remove DAGBuilder.findFieldStorage (direct cast, TopK only inserts OpenSearchRelNode)
- Remove Sort marker coupling (chain-length > 0 is sufficient)
- Trim bloated comments and remove debug loggers

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>

* test: add dc combination coverage (mixed TopK + plain grouped merge)

- testMultiAgg_sortByDc_head10: dc + sum grouped, TopK sorted on dc
- dc_count_by_category: dc + count grouped without TopK (2-shard merge)

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>

* chore: apply spotless formatting

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>

* test: add dc/TopK plan shape + combination tests, fix HAVING flakiness

- testRewrite_dcByGroup_splitAndTopK: assert exact plan shape for
  dc + TopK (reduce_eval Project, strip Project, oversampled Sort)
- testRewrite_multiGroupByCount_splitAndTopK: assert PARTIAL/FINAL split
  fires for multi-group-by COUNT (was broken on main — stayed SINGLE)
- testMultiAgg_sortByDc_head10: IT for mixed dc + sum with TopK
- dc_count_by_category: golden-file IT for dc + count grouped (no TopK)
- testCountByGroup_having_sortDesc_head10: lenient assertion to tolerate
  TopK's approximate pruning while still catching double-counting bugs

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>

---------

Signed-off-by: Sandesh Kumar <sandeshkr419@gmail.com>
Signed-off-by: Sandesh Kumar <kusandes@amazon.com>
Co-authored-by: Sandesh Kumar <kusandes@amazon.com>
KhishorekumarBS pushed a commit to KhishorekumarBS/OpenSearch that referenced this pull request Jul 3, 2026
…lytics-engine route + review fixes (opensearch-project#21975)

* [analytics-engine] Fix now()-family + enable now(fsp)/rand(seed)/sha2/span-month/ts-subtract/dc-approx on the analytics-engine route

Closes several PPL scalar/aggregate gaps on the analytics-engine (DataFusion)
route. Each was either an "Unable to convert call ..." Substrait failure or a
result-rendering gap; all are verified end-to-end via the CalcitePPL*IT remote
ITs on a parquet/composite (analytics-engine-routed) cluster.

now()-family (CalciteNowLikeFunctionIT 12/12):
- ArrowValues now converts Arrow Date32/Date64 -> LocalDate and
  Time{Sec,Milli,Micro,Nano} -> LocalTime so DATE/TIME results render as
  "uuuu-MM-dd" / "HH:mm:ss" via ExprValueUtils.fromObjectValue's existing
  temporal branch (mirrors how TIMESTAMP already arrives as LocalDateTime).
  Fixes current_date/curdate/current_time/curtime returning raw epoch-day /
  units-of-day integers.
- UTC_TIMESTAMP/UTC_DATE/UTC_TIME enum constants + adapter registrations route
  to the same DataFusion builtins as their CURRENT_* equivalents (cluster runs
  in UTC), reusing the existing now/currentDate/currentTime adapters.
- NowFspAdapter drops the optional fractional-seconds-precision arg so
  now(fsp)/current_timestamp(fsp)/sysdate(fsp) map to DataFusion's niladic
  now() instead of failing as "Unable to convert call now(i32)".

Other scalar/aggregate fixes:
- RandSeedAdapter drops the optional rand(seed) operand -> DataFusion random().
- Sha2FunctionAdapter raises a clear "Unsupported SHA2 algorithm [N]" for a
  concrete unsupported literal bit length, matching the SQL-plugin reference
  (CryptographicFunction) instead of surfacing a cryptic Substrait error.
- SpanAdapter supports variable-length month/quarter/year buckets.
- TimestampSubtractRewriter rewrites MINUS(timestamp, timestamp) to an
  epoch-second difference (to_unixtime), which is Substrait-convertible.
- PplAggregateCallRewriter + AggregateFunction map DISTINCT_COUNT_APPROX to
  APPROX_COUNT_DISTINCT (dc/distinct_count approx form).

Adds unit tests for the new adapters/rewriter and ArrowValues DATE/TIME
conversion; updates the SHA2 adapter test to assert the clear-error behavior.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Address review: RAND seed fail-clear, shared Substrait preprocess, adapter arity bounds, span hardening, SHA2 message

Review-driven hardening on top of the initial PR:

- RandSeedAdapter: stop silently dropping the rand(seed) operand (which turned a
  deterministic seeded call into a non-deterministic one). Niladic rand() still
  maps to DataFusion random(); seeded rand(seed) now fails with a clear
  unsupported-shape error until a seeded random is available on the backend.
- DataFusionFragmentConvertor: extract preprocessForSubstrait() and call it from
  both convertToSubstrait and convertStandalone so TimestampSubtractRewriter
  (and any future rewriter) runs on the wrapper/partial-aggregate path too, not
  only the top-level fragment path.
- NowFspAdapter / RandSeedAdapter: only normalize the valid 0-arg / 1-arg shapes;
  leave unexpected arities untouched instead of inventing a valid call.
- SpanAdapter: guard the month and second interval multiplications with
  Math.multiplyExact (clear "interval is too large" error on overflow), and use
  floored modulo/division for the month-index bucketing so pre-1970 (negative
  month index) timestamps snap down to the correct bucket start.
- Sha2FunctionAdapter: list the supported bit lengths (224, 256, 384, 512) in the
  unsupported-algorithm error.

Tests: add NowFsp unexpected-arity test, RAND seeded-reject test, span
month-bucket + overflow tests; update SHA2 test to assert the supported-values
message.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Fix grouped distinct_count_approx nullability + make TimestampSubtractRewriter identity-preserving

Two correctness fixes found while verifying the review round:

- PplAggregateCallRewriter (DISTINCT_COUNT_APPROX): pin the remapped aggregate's
  explicit return type to NOT NULL BIGINT. APPROX_COUNT_DISTINCT is a count, so its
  return-type inference is BIGINT NOT NULL; reusing the PPL call's nullable BIGINT
  left the declared and inferred types disagreeing and tripped Calcite's validity
  assertion ("aggCall type BIGINT vs inferred BIGINT NOT NULL") on the grouped
  two-phase path (e.g. `stats distinct_count_approx(x) by g`). This was a latent bug
  in the un-grouped-only test coverage; the grouped form is now verified green.

- TimestampSubtractRewriter: make it a true identity no-op for plans without a
  MINUS(timestamp, timestamp). Previously it walked every RelNode and called
  RelNode.accept(RexShuttle) unconditionally, which re-derives expression / aggCall
  types and can flip a cached nullable BIGINT to BIGINT NOT NULL — breaking unrelated
  shapes once the shared preprocessing pipeline runs it on the two-phase aggregate
  path. It now detects the target shape first (read-only visitors) and only applies
  the rewriting RexShuttle to nodes that actually contain it, returning the original
  object otherwise. Adds identity (assertSame) regression tests.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Drop TimestampSubtractRewriter — superseded by upstream MinusAdapter (opensearch-project#21978)

Upstream opensearch-project#21978 added MinusAdapter, the proper named ScalarFunctionAdapter for
timestamp/date subtraction (t1 - t2 -> from_unixtime(to_unixtime(t1) - to_unixtime(t2))).
It also coordinates with WidthBucketAdapter by deliberately leaving the
MINUS(MAX OVER(), MIN OVER()) binning shape untouched. Our pre-Substrait
TimestampSubtractRewriter rewrote every MINUS(timestamp, timestamp) including that
binning shape, so keeping both risked clobbering the binning path. Remove the
rewriter (and its test); MinusAdapter fully covers the case. The shared
preprocessForSubstrait helper is retained for the remaining rewriters.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Bind distinct_count_approx after SQL-plugin emits APPROX_COUNT_DISTINCT (opensearch-project#5525)

sql#5525 changed PPLBuiltinOperators.DISTINCT_COUNT_APPROX to emit a user-defined
SqlAggFunction NAMED "APPROX_COUNT_DISTINCT" (DistinctCountApproxLogicalAggFunction).
isthmus binds aggregate sigs by operator identity, not name, so this custom op has no
substrait sig and grouped/ungrouped `stats distinct_count_approx(x)` fails with
"Unable to find binding for call APPROX_COUNT_DISTINCT". (The `dc`->COUNT(DISTINCT)
path works via OpenSearchDistinctCountRule/opensearch-project#22013; the named UDAF path does not.)

Re-key the PplAggregateCallRewriter case from the old name DISTINCT_COUNT_APPROX to
APPROX_COUNT_DISTINCT and remap the PPL marker to stock SqlStdOperatorTable.APPROX_COUNT_DISTINCT
(which ADDITIONAL_AGGREGATE_SIGS binds to DataFusion approx_distinct), guarding against
the stock operator itself. Remove the now-dead DISTINCT_COUNT_APPROX alias in
AggregateFunction.fromNameOrError — the enum constant is APPROX_COUNT_DISTINCT, which
valueOf resolves directly (addresses review comment).

Verified: CalcitePPLAggregationIT.testCountDistinctApprox + testCountDistinctApproxWithAlias
pass on a force-routed analytics-engine cluster.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Don't push filters below non-deterministic projects (fixes rand() predicates)

`eval r = rand() | where r > 0` failed with "Comparison performance-delegation requires
(RexInputRef, RexLiteral); got RAND()". Calcite's stock FILTER_PROJECT_TRANSPOSE only guards
against window functions (!containsOver()), so it pushed the filter below the rand() project,
inlining RAND() into the predicate — turning a delegatable ($ref > literal) comparison into
RAND() > literal on the scan, which Lucene performance-delegation cannot serialize (and which
must stay on the single in-memory engine, since each backend would draw different values).

Replace the stock rule with FILTER_PROJECT_TRANSPOSE_DETERMINISTIC: a FilterProjectTransposeRule
configured to refuse the transpose when the Project computes any non-deterministic expression
(RexUtil.isDeterministic). Keeping the Filter above the Project preserves the clean
($ref > literal) shape evaluated in memory.

Verified: CalcitePPLBuiltinFunctionIT.testRand passes on a force-routed analytics-engine cluster.
Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Restore distinct_count_approx binding on the analytics-engine route

Validation requested in PR review (sandeshkr419) found that sql#5525 alone does
NOT make `stats distinct_count_approx(x)` work on the analytics-engine route, so
the engine-side handling removed in the prior commit is still required. On a
fresh, force-routed cluster with sql#5525 deployed, CalcitePPLAggregationIT
testCountDistinctApprox + testCountDistinctApproxWithAlias fail in two stages:

  1. Planner stage: OpenSearchAggregateRule.resolveViableBackendsForCall reads
     aggCall.getAggregation().getName(), which is still "DISTINCT_COUNT_APPROX"
     (the DistinctCountApproxLogicalAggFunction marker). fromNameOrError then
     throws "No enum constant ...AggregateFunction.DISTINCT_COUNT_APPROX".
  2. Substrait stage: the PplAggregateCallRewriter case keyed only on
     "APPROX_COUNT_DISTINCT" never matched the marker, so the unbound op reached
     isthmus -> "Unable to find binding for call DISTINCT_COUNT_APPROX($1)".

sql#5525 only sets the substrait-emission name to APPROX_COUNT_DISTINCT, which is
a later stage than both of the above, so the marker's runtime name still reaches
them. Fixes:

  - Restore the DISTINCT_COUNT_APPROX -> APPROX_COUNT_DISTINCT alias in
    AggregateFunction.fromNameOrError (covers the planner path).
  - Match both "DISTINCT_COUNT_APPROX" and "APPROX_COUNT_DISTINCT" in the
    PplAggregateCallRewriter case (covers the substrait-binding path); keep the
    guard that skips the already-bound stock operator.
  - Add a fromNameOrError alias unit test.

Verified on a force-routed analytics-engine cluster (parquet-backed indices,
AE-routed): both tests pass; the surrounding 4-class IT sweep is unchanged at
6 pre-existing failures (date_format/strftime/percentile/nested-field), i.e. no
regressions.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Add qa-module ITs for distinct_count_approx + min/max(boolean)

Engine-side end-to-end regression coverage for the two analytics-engine binding
fixes in this PR, run against the production /_plugins/_ppl surface on the
qa-module's analytics-routed cluster (mirrors the SQL-repo CalcitePPLAggregationIT
cases that proved the fixes):

  - DistinctCountApproxAggregationIT: grouped / aliased / dc-shorthand / ungrouped
    distinct_count_approx (guards the planner enum alias + rewriter rebind).
  - MinMaxBooleanAggregationIT: ungrouped / grouped / mixed boolean+string
    min/max (guards the boolean min/max substrait overload).

Each IT creates its own parquet-backed (composite/parquet + lucene secondary)
index, bulk-ingests a small fixed dataset, and asserts oracle values. All 7
tests pass (4 + 3).

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Fix date_format %c month padding; correct stale RAND comments

Two small follow-ups surfaced while verifying the analytics-engine route:

  - date_format/strftime %c now emits the zero-padded month (Jan -> "01"),
    matching the PPL reference (SQL-plugin DateTimeFormatterUtil maps %c -> "MM"
    in date mode), which differs from stock MySQL's no-leading-zero %c. Time mode
    keeps the reference's single-"0" literal. Fixes the %c portion of
    CalciteDateTimeFunctionIT.testDateFormat and updates the Rust token-set unit
    test expectation accordingly. (The test's %U/%u/%V/%v week-number padding is a
    separate, unrelated date-format issue left for a follow-up.)

  - Correct two stale comments that described the old "drop the seed" RAND
    behavior; the adapter now rejects RAND(seed) with a clear error (addresses
    review feedback). No behavior change — comment/javadoc only.

Verified on a force-routed analytics-engine cluster: date_format(ts,'%c') returns
"01"; the now/sha2/dc/min-max ITs remain green (no regressions).

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Drop DISTINCT_COUNT_APPROX enum alias per review

Remove the DISTINCT_COUNT_APPROX -> APPROX_COUNT_DISTINCT alias in
AggregateFunction.fromNameOrError (and its unit test), and key the
PplAggregateCallRewriter case on "APPROX_COUNT_DISTINCT" only, per
@sandeshkr419's review. The SQL plugin (sql#5525) emits the operator named
APPROX_COUNT_DISTINCT, which the enum resolves directly, so the legacy-name
alias is no longer needed.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Address review: bulk NDJSON content-type, dc comments, error wording

Review follow-ups (sandeshkr419):

  - QA ITs: send _bulk as application/x-ndjson via the per-index endpoint with bare
    {"index": {}} action metadata, matching the existing CountFastPathIT pattern.
    setJsonEntity alone defaults to application/json, which the bulk endpoint rejects.
  - Clarify the distinct_count_approx marker naming across layers: current runtime name
    is APPROX_COUNT_DISTINCT (sql#5525), DISTINCT_COUNT_APPROX is the legacy/defensive
    spelling; OpenSearchDistinctCountRule is the primary planner rewrite and
    PplAggregateCallRewriter is a late Substrait-emission defensive fallback.
  - Document the phase-specific nullability handling: the early planner rewrite preserves
    the marker's nullable type (LogicalAggregate.copy re-validates row type); the late
    Substrait rewrite pins NOT NULL (Calcite validates against the stock op's inferred type).
  - SPAN overflow message now includes the operands.
  - Drop "yet" from the seeded-RAND error so it doesn't imply imminent support.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Empty commit to retrigger CI

No code change. The prior gradle-check failed in the Jenkins-trigger step
(jq parse error polling the Jenkins API, 10 retries, empty result) — a
transient CI-infra glitch unrelated to this PR's changes. Re-running checks.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Fix DatetimeCoverageIT %c assertions to match zero-padded month

The date_format %c code change (no-leading-zero -> zero-padded "MM") aligns the
analytics-engine Rust UDF with the legacy PPL contract: DateTimeFormatterUtil.DATE_HANDLERS
maps %c -> "MM" (date_format dispatches through DATE_HANDLERS), so January renders as "01".
The Rust code + unit test were updated for this, but the two DatetimeCoverageIT assertions
still expected the old unpadded "1" — the gradle-check / sandbox-check failures.

Update both assertions to "01", matching the canonical merged ITs (ppl/DateTimeFunctionIT
and sql/DateTimeFunctionIT both assert "Sat Jan 01 31st ..." for the same spec). Test-only.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Enable now()/rand(seed)/sha2/span-month on the analytics-engine route + review fixes

Enables these PPL functions on the force-routed analytics-engine (DataFusion)
path and addresses review feedback:

- now()-family: NowFspAdapter handles now([fsp]) / current_timestamp /
  localtimestamp, where fsp is MySQL fractional-seconds precision (0-6).
- rand([seed]): RandSeedAdapter.
- sha2(input, bitLen): Sha2FunctionAdapter -> encode(digest(input, 'shaN'), 'hex'),
  with a backend fail-clear guard for unsupported bit lengths (TODO to move
  validation to the frontend once the handoff carries function metadata).
- span() month bucketing: SpanAdapter.
- strftime: os_strftime.rs.
- date_format %c: zero-padded month (DatetimeCoverageIT updated to match the
  merged ppl/sql DateTimeFunctionIT expectations).
- PlannerImpl: don't push filters below non-deterministic projects; rand() in a
  pushed-down predicate would be re-evaluated and draw a fresh value. Documented
  as a semantic-correctness guard.
- DataFusionFragmentConvertor: shared Substrait preprocess entry point.
- New QA ITs: MinMaxBooleanAggregationIT.

Review fixes: removed the late defensive APPROX_COUNT_DISTINCT fallback in
PplAggregateCallRewriter (fail loud instead); restored the Substrait-layer
assessment TODO; clarified fsp in NowFspAdapter; tidied AggregateFunction.

distinct_count_approx on the analytics-engine route is handled upstream by opensearch-project#22120
and is not part of this PR.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

* [analytics-engine] Add 2-shard reduce coverage for boolean min/max

Addresses review feedback to cover the multi-shard case for boolean min/max.
min(flag)/max(flag) over the merge_coverage dataset's boolean field are added as
golden cases in the TwoShardReduceTestCase suite (opensearch-project#21951 convention), so they run
at 1 shard and 2 shards with a differential equality check plus a pinned golden,
exercising the two-phase per-shard-partial then coordinator-merge reduce path.
The merge_coverage flag field has 15 false and 15 true rows, so min=false and
max=true regardless of shard layout.

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>

---------

Signed-off-by: Eric Wei <mengwei.eric@gmail.com>
asifabashar pushed a commit to asifabashar/sql that referenced this pull request Jul 21, 2026
…stSimpleCount0 + APPROX_COUNT_DISTINCT name) (opensearch-project#5525)

* Use a parquet-backed index in CalcitePPLAggregationIT.testSimpleCount0

A bare auto-created index isn't composite/parquet-backed, so on the analytics-engine
route it doesn't route to the analytics engine. Switch to TEST_INDEX_BANK (loaded via
loadIndex, which injects parquet settings when the flag is set, 7 docs) so the test is
meaningful on both routes. Diagnosis by Sandesh Kumar.

Signed-off-by: Kai Huang <ahkcs@amazon.com>

* Emit APPROX_COUNT_DISTINCT as the distinct_count_approx runtime name

distinct_count_approx() failed to bind on the analytics-engine (DataFusion) route
because the SqlAggFunction was named DISTINCT_COUNT_APPROX; the backend resolves
aggregates by the Calcite/Substrait-standard name APPROX_COUNT_DISTINCT. The Java
field name and PPL function name are unchanged. The OpenSearch V3 path is unaffected
(it overrides this via the external HLL registration). Analytics-route binding is
completed by opensearch-project/OpenSearch#22013. Per Sandesh Kumar.

Signed-off-by: Kai Huang <ahkcs@amazon.com>

---------

Signed-off-by: Kai Huang <ahkcs@amazon.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants