Route distinct_count_approx to APPROX_COUNT_DISTINCT

sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchDistinctCountRule.java

@@ -10,18 +10,25 @@
 
 import org.apache.calcite.plan.RelOptRule;
 import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.rel.logical.LogicalAggregate;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeField;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.tools.RelBuilder;
 
 import java.util.ArrayList;
 import java.util.List;
 
 /**
- * Rewrites single-arg {@code COUNT(DISTINCT x)} → {@code APPROX_COUNT_DISTINCT(x)} before
- * the aggregate is marked by {@link OpenSearchAggregateRule}. Multi-arg distinct falls through
- * to coordinator-gather in {@link OpenSearchAggregateSplitRule}.
+ * Rewrites single-arg {@code COUNT(DISTINCT x)} and PPL's {@code distinct_count_approx(x)} UDAF
+ * marker to {@link SqlStdOperatorTable#APPROX_COUNT_DISTINCT} before the aggregate is marked by
+ * {@link OpenSearchAggregateRule}, so substrait dispatch resolves by operator identity. Multi-arg
+ * distinct falls through to coordinator-gather in {@link OpenSearchAggregateSplitRule}.
  *
  * @opensearch.internal
  */
@@ -34,7 +41,7 @@ public OpenSearchDistinctCountRule() {
     @Override
     public boolean matches(RelOptRuleCall ruleCall) {
         LogicalAggregate agg = ruleCall.rel(0);
-        return agg.getAggCallList().stream().anyMatch(OpenSearchDistinctCountRule::isSingleArgCountDistinct);
+        return agg.getAggCallList().stream().anyMatch(OpenSearchDistinctCountRule::needsRewriteToApprox);
     }
 
     @Override
@@ -43,21 +50,67 @@ public void onMatch(RelOptRuleCall ruleCall) {
         List<AggregateCall> rewritten = new ArrayList<>(agg.getAggCallList().size());
         boolean changed = false;
         for (AggregateCall call : agg.getAggCallList()) {
-            if (isSingleArgCountDistinct(call)) {
+            if (needsRewriteToApprox(call)) {
                 rewritten.add(rewriteToApprox(call, agg));
                 changed = true;
             } else {
                 rewritten.add(call);
             }
         }
         if (!changed) return;
-        ruleCall.transformTo(agg.copy(agg.getTraitSet(), agg.getInput(), agg.getGroupSet(), agg.getGroupSets(), rewritten));
+        LogicalAggregate replacement = (LogicalAggregate) agg.copy(
+            agg.getTraitSet(),
+            agg.getInput(),
+            agg.getGroupSet(),
+            agg.getGroupSets(),
+            rewritten
+        );
+        // Aggregate.typeMatchesInferred forces the new aggCall to BIGINT NOT NULL while HepPlanner
+        // requires the replacement's row type to equal the original's; bridge with a casting Project.
+        RelNode rewrittenNode = projectToOriginalRowType(ruleCall, agg, replacement);
+        ruleCall.transformTo(rewrittenNode);
+    }
+
+    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();
+        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();
+    }
+
+    /** True when the call is a single-arg COUNT(DISTINCT) or PPL's distinct_count_approx UDAF. */
+    private static boolean needsRewriteToApprox(AggregateCall call) {
+        return isSingleArgCountDistinct(call) || isPplDistinctCountApproxUdf(call);
     }
 
     private static boolean isSingleArgCountDistinct(AggregateCall call) {
         return call.getAggregation().getKind() == SqlKind.COUNT && call.isDistinct() && call.getArgList().size() == 1;
     }
 
+    /** PPL's distinct_count_approx is a UDF named "APPROX_COUNT_DISTINCT" that is not the stdop. */
+    private static boolean isPplDistinctCountApproxUdf(AggregateCall call) {
+        return call.getAggregation() != SqlStdOperatorTable.APPROX_COUNT_DISTINCT
+            && "APPROX_COUNT_DISTINCT".equals(call.getAggregation().getName())
+            && call.getArgList().size() == 1;
+    }
+
     private static AggregateCall rewriteToApprox(AggregateCall call, LogicalAggregate agg) {
         return AggregateCall.create(
             SqlStdOperatorTable.APPROX_COUNT_DISTINCT,

sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/AggregateRuleTests.java

@@ -11,7 +11,15 @@
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.core.AggregateCall;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlAggFunction;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.fun.SqlBasicAggFunction;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.OperandTypes;
+import org.apache.calcite.sql.type.ReturnTypes;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.opensearch.analytics.planner.rel.AggregateCallAnnotation;
 import org.opensearch.analytics.planner.rel.AggregateMode;
@@ -332,6 +340,80 @@ public void testCountDistinctRewrittenToApproxCountDistinct() {
         assertFalse("isDistinct must be cleared on the rewritten APPROX_COUNT_DISTINCT call", rebuilt.isDistinct());
     }
 
+    /**
+     * PPL's {@code distinct_count_approx} UDAF marker — a {@code SqlAggFunction} named
+     * {@code "APPROX_COUNT_DISTINCT"} but not the Calcite stdop, returning a NULLABLE BIGINT —
+     * is rewritten to {@code SqlStdOperatorTable.APPROX_COUNT_DISTINCT} (which infers BIGINT
+     * NOT NULL) and wrapped in a casting {@link OpenSearchProject} that restores the original
+     * nullable row type. The Project bridge is required because {@code Aggregate.typeMatchesInferred}
+     * pins the new aggCall's type to its operator's inferred type while {@code HepPlanner} pins
+     * the replacement's row type to the original's.
+     */
+    public void testPplDistinctCountApproxUdfRewrittenWithCastProject() {
+        SqlAggFunction pplUdfMarker = SqlBasicAggFunction.create(
+            "APPROX_COUNT_DISTINCT",
+            SqlKind.OTHER_FUNCTION,
+            ReturnTypes.BIGINT_FORCE_NULLABLE,
+            OperandTypes.ANY
+        );
+        RelDataType nullableBigint = typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true);
+        RelNode scan = stubScan(mockTable("test_index", "status", "size"));
+        AggregateCall pplApprox = AggregateCall.create(pplUdfMarker, /* distinct */ false, List.of(1), -1, scan, nullableBigint, "dca");
+
+        RelNode result = runPlanner(makeAggregate(pplApprox), defaultContext(1));
+        logger.info("Plan:\n{}", RelOptUtil.toString(result));
+        RelNode unwrapped = unwrapRootReducer(result);
+        assertTrue(
+            "Expected OpenSearchProject(OpenSearchAggregate(...)) wrap, got " + unwrapped.getClass().getSimpleName(),
+            unwrapped instanceof OpenSearchProject
+        );
+        OpenSearchProject project = (OpenSearchProject) unwrapped;
+        // Project preserves the original nullable BIGINT for the aggregated column.
+        RelDataType dcaType = project.getRowType().getFieldList().get(1).getType();
+        assertEquals("Project must restore original nullable BIGINT", nullableBigint, dcaType);
+        // Field 1 must be a non-trivial expression (the cast); field 0 stays as a passthrough ref.
+        // OpenSearchProjectRule wraps scalar calls in AnnotatedProjectExpression, so we search
+        // recursively for a CAST node.
+        assertTrue("Project must contain a CAST to bridge nullability", containsCast(project.getProjects().get(1)));
+
+        RelNode innerAgg = RelNodeUtils.unwrapHep(project.getInput());
+        assertTrue(
+            "Inner node must be OpenSearchAggregate, got " + innerAgg.getClass().getSimpleName(),
+            innerAgg instanceof OpenSearchAggregate
+        );
+        OpenSearchAggregate agg = (OpenSearchAggregate) innerAgg;
+        AggregateCall rebuilt = agg.getAggCallList().get(0);
+        assertSame(
+            "Inner aggregate must use SqlStdOperatorTable.APPROX_COUNT_DISTINCT",
+            SqlStdOperatorTable.APPROX_COUNT_DISTINCT,
+            rebuilt.getAggregation()
+        );
+        assertFalse("isDistinct must be cleared on the rewritten call", rebuilt.isDistinct());
+    }
+
+    /**
+     * Stdop {@code APPROX_COUNT_DISTINCT} (already canonical) must not be rewritten — the rule's
+     * predicate excludes the stdop, so no Project wrap is added and the result is a plain
+     * {@link OpenSearchAggregate}.
+     */
+    public void testStdopApproxCountDistinctNotRewritten() {
+        RelNode scan = stubScan(mockTable("test_index", "status", "size"));
+        OpenSearchAggregate agg = runAggregate(1, approxCountDistinctCall(scan));
+        AggregateCall call = agg.getAggCallList().get(0);
+        assertSame(SqlStdOperatorTable.APPROX_COUNT_DISTINCT, call.getAggregation());
+    }
+
+    /** Recursive search for a CAST node — Project rule wraps scalar calls in AnnotatedProjectExpression. */
+    private static boolean containsCast(RexNode node) {
+        if (node.getKind() == SqlKind.CAST) return true;
+        if (node instanceof RexCall call) {
+            for (RexNode operand : call.getOperands()) {
+                if (containsCast(operand)) return true;
+            }
+        }
+        return false;
+    }
+
     private OpenSearchAggregate runAggregate(int shardCount, AggregateCall aggCall) {
         RelNode result = runPlanner(makeAggregate(aggCall), defaultContext(shardCount));
         logger.info("Plan:\n{}", RelOptUtil.toString(result));

sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/DistinctCountApproxIT.java

@@ -0,0 +1,95 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.analytics.qa;
+
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * End-to-end coverage for PPL `stats distinct_count_approx(field)` over the shared
+ * `calcs` dataset. The PPL parser registers `distinct_count_approx` as a
+ * `SqlUserDefinedAggFunction` named {@code "APPROX_COUNT_DISTINCT"}; substrait
+ * emission keys off operator identity, so without `OpenSearchDistinctCountRule`
+ * rewriting the UDF marker to Calcite's stdop the call falls through with
+ * "Unable to find binding for call APPROX_COUNT_DISTINCT($N)".
+ */
+public class DistinctCountApproxIT extends AnalyticsRestTestCase {
+
+    private static final Dataset DATASET = new Dataset("calcs", "calcs");
+
+    private static boolean dataProvisioned = false;
+
+    @Override
+    protected void onBeforeQuery() throws IOException {
+        if (dataProvisioned == false) {
+            DatasetProvisioner.provision(client(), DATASET);
+            dataProvisioned = true;
+        }
+    }
+
+    public void testDistinctCountApproxByGroup() throws IOException {
+        // calcs has 17 rows with str0 ∈ {FURNITURE: 2, OFFICE SUPPLIES: 6, TECHNOLOGY: 9}
+        // distinct str1 values per group.
+        assertRowsEqual(
+            "source=" + DATASET.indexName + " | stats distinct_count_approx(str1) by str0 | sort str0",
+            row(2L, "FURNITURE"),
+            row(6L, "OFFICE SUPPLIES"),
+            row(9L, "TECHNOLOGY")
+        );
+    }
+
+    public void testDistinctCountApproxByGroupWithAlias() throws IOException {
+        assertRowsEqual(
+            "source=" + DATASET.indexName + " | stats distinct_count_approx(str1) as dca by str0 | sort str0",
+            row(2L, "FURNITURE"),
+            row(6L, "OFFICE SUPPLIES"),
+            row(9L, "TECHNOLOGY")
+        );
+    }
+
+    public void testDistinctCountApproxGlobal() throws IOException {
+        // `key` is the 17 distinct row identifiers (key00..key16).
+        assertRowsEqual(
+            "source=" + DATASET.indexName + " | stats distinct_count_approx(key) as dca",
+            row(17L)
+        );
+    }
+
+    private static List<Object> row(Object... values) {
+        return Arrays.asList(values);
+    }
+
+    @SafeVarargs
+    @SuppressWarnings("varargs")
+    private final void assertRowsEqual(String ppl, List<Object>... expected) throws IOException {
+        Map<String, Object> response = executePpl(ppl);
+        @SuppressWarnings("unchecked")
+        List<List<Object>> actualRows = (List<List<Object>>) response.get("datarows");
+        assertNotNull("Response missing 'datarows' for query: " + ppl, actualRows);
+        assertEquals("Row count mismatch for query: " + ppl, expected.length, actualRows.size());
+        for (int i = 0; i < expected.length; i++) {
+            List<Object> want = expected[i];
+            List<Object> got = actualRows.get(i);
+            assertEquals("Column count mismatch at row " + i + " for query: " + ppl, want.size(), got.size());
+            for (int j = 0; j < want.size(); j++) {
+                Object wantCell = want.get(j);
+                Object gotCell = got.get(j);
+                String label = "Cell mismatch at row " + i + ", col " + j + " for query: " + ppl;
+                if (wantCell instanceof Number && gotCell instanceof Number) {
+                    // Jackson boxes counts as Integer when they fit in 32 bits; compare numerically.
+                    assertEquals(label, ((Number) wantCell).longValue(), ((Number) gotCell).longValue());
+                } else {
+                    assertEquals(label, wantCell, gotCell);
+                }
+            }
+        }
+    }
+}