Route distinct_count_approx to APPROX_COUNT_DISTINCT

[vinaykpud]

Description

PPL's distinct_count_approx(x) is parsed as a SqlUserDefinedAggFunction named "APPROX_COUNT_DISTINCT", not Calcite's SqlStdOperatorTable.APPROX_COUNT_DISTINCT stdop. Substrait dispatch keys off operator identity, so the call never reached the approx_distinct binding and queries failed with:

UnsupportedOperationException: Unable to find binding for call APPROX_COUNT_DISTINCT($N)

This PR extends OpenSearchDistinctCountRule to rewrite the UDF marker to the stdop, alongside the existing COUNT(DISTINCT x) normalization. The rewritten aggregate is wrapped in a LogicalProject that casts back to the original nullable type — needed because Aggregate.typeMatchesInferred pins the new aggCall to the stdop's BIGINT NOT NULL while HepPlanner pins the replacement's row type to the original BIGINT (nullable).

Semantics

distinct_count_approx(field) returns an approximate distinct-value count using DataFusion's HyperLogLog++ sketch (native approx_distinct). Cardinality results may differ by ±1 at small group sizes; preferred over exact count(distinct x) for high-cardinality columns where the HLL sketch merge is dramatically cheaper than gather-then-dedupe.

Query shapes now supported

| stats distinct_count_approx(field)
| stats distinct_count_approx(field) as alias
| stats distinct_count_approx(field) by group_field
| stats distinct_count_approx(field) as alias by group_field

The same rewrite path also normalizes single-arg count(distinct field) (existing behaviour, unchanged).

Tests

• Unit: two new cases in AggregateRuleTests — one for the UDF → stdop rewrite with the cast-Project wrap, one for the already-stdop no-op path.
• Sandbox QA IT: new DistinctCountApproxIT over the shared calcs dataset covers the by-group, aliased, and global query shapes end-to-end against a parquet-backed index.

PPL Calcite tests in the SQL plugin unblocked

• CalcitePPLAggregationIT.testCountDistinctApprox
• CalcitePPLAggregationIT.testCountDistinctApproxWithAlias
• CalcitePPLAggregationPaginatingIT.testCountDistinctApprox
• CalcitePPLAggregationPaginatingIT.testCountDistinctApproxWithAlias

Check List

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

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

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 2ce8c21)

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 The matches() method returns a boolean but has no return statement. This will cause a compilation error. The method should return the result of the stream operation. public boolean matches(RelOptRuleCall ruleCall) { LogicalAggregate agg = ruleCall.rel(0); return agg.getAggCallList().stream().anyMatch(OpenSearchDistinctCountRule::needsRewriteToApprox); }

[github-actions]

PR Code Suggestions ✨

Latest suggestions up to 2ce8c21

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Add field count validation Add bounds checking before accessing newFields.get(i) to prevent potential IndexOutOfBoundsException. The code assumes origFields and newFields have the same size, but if the replacement aggregate has a different field count, the loop will fail when accessing newFields.get(i).getType(). sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchDistinctCountRule.java [74-96] private static RelNode projectToOriginalRowType(RelOptRuleCall ruleCall, LogicalAggregate original, LogicalAggregate replacement) { if (replacement.getRowType().equals(original.getRowType())) { return replacement; } RelBuilder relBuilder = ruleCall.builder(); relBuilder.push(replacement); RexBuilder rexBuilder = relBuilder.getRexBuilder(); List<RelDataTypeField> origFields = original.getRowType().getFieldList(); List<RelDataTypeField> newFields = replacement.getRowType().getFieldList(); + if (origFields.size() != newFields.size()) { + throw new IllegalStateException("Field count mismatch between original and replacement aggregates"); + } List<RexNode> projects = new ArrayList<>(origFields.size()); List<String> names = new ArrayList<>(origFields.size()); for (int i = 0; i < origFields.size(); i++) { RexNode ref = rexBuilder.makeInputRef(replacement, i); RelDataType targetType = origFields.get(i).getType(); if (!newFields.get(i).getType().equals(targetType)) { ref = rexBuilder.makeCast(targetType, ref); } projects.add(ref); names.add(origFields.get(i).getName()); } relBuilder.project(projects, names, /* forceProject */ true); return relBuilder.build(); } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential IndexOutOfBoundsException if origFields and newFields have different sizes. Adding validation improves robustness, though the scenario may be unlikely given the context where both aggregates should have matching field counts by design.	Medium

---

Previous suggestions

Suggestions up to commit e2657de

Category	Suggestion	Impact
Possible issue	Validate field list size match Add validation to ensure origFields and newFields have the same size before iterating. A size mismatch between original and replacement field lists could cause an IndexOutOfBoundsException when accessing newFields.get(i). sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchDistinctCountRule.java [74-96] private static RelNode projectToOriginalRowType(RelOptRuleCall ruleCall, LogicalAggregate original, LogicalAggregate replacement) { if (replacement.getRowType().equals(original.getRowType())) { return replacement; } RelBuilder relBuilder = ruleCall.builder(); relBuilder.push(replacement); RexBuilder rexBuilder = relBuilder.getRexBuilder(); List<RelDataTypeField> origFields = original.getRowType().getFieldList(); List<RelDataTypeField> newFields = replacement.getRowType().getFieldList(); + if (origFields.size() != newFields.size()) { + throw new IllegalStateException("Field count mismatch between original and replacement aggregates"); + } List<RexNode> projects = new ArrayList<>(origFields.size()); List<String> names = new ArrayList<>(origFields.size()); for (int i = 0; i < origFields.size(); i++) { RexNode ref = rexBuilder.makeInputRef(replacement, i); RelDataType targetType = origFields.get(i).getType(); if (!newFields.get(i).getType().equals(targetType)) { ref = rexBuilder.makeCast(targetType, ref); } projects.add(ref); names.add(origFields.get(i).getName()); } relBuilder.project(projects, names, /* forceProject */ true); return relBuilder.build(); } Suggestion importance[1-10]: 7 __ Why: The suggestion adds defensive validation to prevent IndexOutOfBoundsException when field counts mismatch between original and replacement aggregates. While this improves robustness, the scenario is unlikely given the controlled rewrite context where only aggregate call types change, not field counts.	Medium

Suggestions up to commit ede1e27

Category	Suggestion	Impact
Possible issue	Validate field count equality Add bounds checking before accessing newFields.get(i) to prevent IndexOutOfBoundsException when field counts differ. The method assumes origFields and newFields have the same size, but this isn't validated. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchDistinctCountRule.java [74-96] private static RelNode projectToOriginalRowType(RelOptRuleCall ruleCall, LogicalAggregate original, LogicalAggregate replacement) { if (replacement.getRowType().equals(original.getRowType())) { return replacement; } RelBuilder relBuilder = ruleCall.builder(); relBuilder.push(replacement); RexBuilder rexBuilder = relBuilder.getRexBuilder(); List<RelDataTypeField> origFields = original.getRowType().getFieldList(); List<RelDataTypeField> newFields = replacement.getRowType().getFieldList(); + if (origFields.size() != newFields.size()) { + throw new IllegalStateException("Field count mismatch: original=" + origFields.size() + ", replacement=" + newFields.size()); + } List<RexNode> projects = new ArrayList<>(origFields.size()); List<String> names = new ArrayList<>(origFields.size()); for (int i = 0; i < origFields.size(); i++) { RexNode ref = rexBuilder.makeInputRef(replacement, i); RelDataType targetType = origFields.get(i).getType(); if (!newFields.get(i).getType().equals(targetType)) { ref = rexBuilder.makeCast(targetType, ref); } projects.add(ref); names.add(origFields.get(i).getName()); } relBuilder.project(projects, names, /* forceProject */ true); return relBuilder.build(); } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential IndexOutOfBoundsException if origFields and newFields have different sizes. Adding validation improves robustness, though the scenario may be unlikely given the context where both aggregates should have matching field structures.	Medium

Suggestions up to commit 992afa1

Category	Suggestion	Impact
Possible issue	Validate field count equality Add bounds checking before accessing newFields.get(i) to prevent IndexOutOfBoundsException when field counts differ. The method assumes origFields and newFields have the same size, but this isn't validated, creating a potential crash if the aggregate transformation changes field count. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchDistinctCountRule.java [74-96] private static RelNode projectToOriginalRowType(RelOptRuleCall ruleCall, LogicalAggregate original, LogicalAggregate replacement) { if (replacement.getRowType().equals(original.getRowType())) { return replacement; } RelBuilder relBuilder = ruleCall.builder(); relBuilder.push(replacement); RexBuilder rexBuilder = relBuilder.getRexBuilder(); List<RelDataTypeField> origFields = original.getRowType().getFieldList(); List<RelDataTypeField> newFields = replacement.getRowType().getFieldList(); + if (origFields.size() != newFields.size()) { + throw new IllegalStateException("Field count mismatch: original=" + origFields.size() + ", replacement=" + newFields.size()); + } List<RexNode> projects = new ArrayList<>(origFields.size()); List<String> names = new ArrayList<>(origFields.size()); for (int i = 0; i < origFields.size(); i++) { RexNode ref = rexBuilder.makeInputRef(replacement, i); RelDataType targetType = origFields.get(i).getType(); if (!newFields.get(i).getType().equals(targetType)) { ref = rexBuilder.makeCast(targetType, ref); } projects.add(ref); names.add(origFields.get(i).getName()); } relBuilder.project(projects, names, /* forceProject */ true); return relBuilder.build(); } Suggestion importance[1-10]: 7 __ Why: The suggestion correctly identifies a potential IndexOutOfBoundsException when origFields and newFields have different sizes. However, given the context that this method is called only when row types differ (line 75 check), and the aggregate rewrite only changes type nullability (not field count), this scenario may be unlikely in practice. The validation adds defensive programming but addresses an edge case rather than a critical bug.	Medium

Suggestions up to commit 3aa4738

Category	Suggestion	Impact
Possible issue	Add field count validation Add bounds checking before accessing field lists by index. If origFields.size() differs from newFields.size(), the loop will cause an IndexOutOfBoundsException when accessing newFields.get(i). Validate that both field lists have the same size before iterating. sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchDistinctCountRule.java [74-96] private static RelNode projectToOriginalRowType(RelOptRuleCall ruleCall, LogicalAggregate original, LogicalAggregate replacement) { if (replacement.getRowType().equals(original.getRowType())) { return replacement; } RelBuilder relBuilder = ruleCall.builder(); relBuilder.push(replacement); RexBuilder rexBuilder = relBuilder.getRexBuilder(); List<RelDataTypeField> origFields = original.getRowType().getFieldList(); List<RelDataTypeField> newFields = replacement.getRowType().getFieldList(); + if (origFields.size() != newFields.size()) { + throw new IllegalStateException("Field count mismatch between original and replacement aggregates"); + } List<RexNode> projects = new ArrayList<>(origFields.size()); List<String> names = new ArrayList<>(origFields.size()); for (int i = 0; i < origFields.size(); i++) { RexNode ref = rexBuilder.makeInputRef(replacement, i); RelDataType targetType = origFields.get(i).getType(); if (!newFields.get(i).getType().equals(targetType)) { ref = rexBuilder.makeCast(targetType, ref); } projects.add(ref); names.add(origFields.get(i).getName()); } relBuilder.project(projects, names, /* forceProject */ true); return relBuilder.build(); } Suggestion importance[1-10]: 7 __ Why: Valid defensive programming suggestion that prevents potential IndexOutOfBoundsException when field counts mismatch. However, in the context of this rewrite rule where replacement is created from original with only aggregate calls modified, field counts should naturally match, making this a precautionary rather than critical fix.	Medium

[github-actions]

Persistent review updated to latest commit 992afa1

[github-actions]

✅ Gradle check result for 992afa1: SUCCESS

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.29%. Comparing base (2f2a25a) to head (2ce8c21).
⚠️ Report is 7 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff              @@
##               main   #22120      +/-   ##
============================================
- Coverage     73.44%   73.29%   -0.16%     
+ Complexity    75916    75813     -103     
============================================
  Files          6070     6070              
  Lines        344610   344610              
  Branches      49576    49576              
============================================
- Hits         253095   252565     -530     
- Misses        71373    71914     +541     
+ Partials      20142    20131      -11     

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

🚀 New features to boost your workflow:

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

[github-actions]

Persistent review updated to latest commit ede1e27

[github-actions]

✅ Gradle check result for ede1e27: SUCCESS

[github-actions]

Persistent review updated to latest commit e2657de

[github-actions]

✅ Gradle check result for e2657de: SUCCESS

[github-actions]

Persistent review updated to latest commit 2ce8c21

[github-actions]

❌ Gradle check result for 2ce8c21: FAILURE

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

[github-actions]

❌ Gradle check result for 2ce8c21: null

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

[github-actions]

✅ Gradle check result for 2ce8c21: SUCCESS