[analytics-engine] Wire FINAL aggregate filter drop and join-condition adapter dispatch (opensearch-project#21911)

songkant-aws · web-flow · commit 4f6bcf58d52f · 2026-06-02T01:21:47.000-07:00
* [analytics-engine] Wire FINAL aggregate filter drop and join-condition adapter dispatch Two small fixes that together unblock CalciteTransposeCommandIT on the analytics-engine route end-to-end. Both surface only when a single PPL command produces (a) a non-prefix groupSet, (b) a FILTER aggCall, and (c) a Join whose condition carries PPL UDFs — which is how PPL transpose lowers via RelBuilder.unpivot()/pivot(). Today transpose is the only PPL command that hits this combination, but the bugs are general. 1. DistributedAggregateRewriter.buildOne — drop filterArg on FINAL ----------------------------------------------------------------- FINAL's input is PARTIAL output, laid out as [group keys, agg states]. The boolean column referenced by an aggCall's FILTER predicate exists only on the ORIGINAL child input that PARTIAL consumed; PARTIAL has already applied the filter while accumulating state. The Aggregate constructor's `isPredicate(input, filterArg)` check then fires when it reads filterArg=N against PARTIAL output that has fewer than N+1 columns (or whose Nth column is non-boolean). Set filterArg = -1 (Calcite's "no FILTER" sentinel — there's no create() overload that omits it) on the FINAL call. Semantically correct: FILTER is a row-level gate consumed once during accumulation; merging states never re-applies it. This generalises to multi-stage chains (PARTIAL → PARTIAL2 → FINAL): only the first stage that consumes raw rows keeps filterArg. Without this fix, transpose IT fails with `IllegalArgumentException: filter must be BOOLEAN NOT NULL` from Aggregate.<init>:178. 2. BackendPlanAdapter — dispatch OpenSearchJoin for adapter rewrite ------------------------------------------------------------------ adaptNode() walks Filter / Project / Aggregate(FINAL) and runs each RexNode through the backend's ScalarFunctionAdapter chain (e.g. ToStringFunctionAdapter rewrites NUMBER_TO_STRING to a plain CAST that isthmus understands). Calcite's FILTER_INTO_JOIN rule inlines an outer Filter's predicate into an inner Join's condition, so any PPL UDF that lived in the Project below the Filter rides into the Join condition. With Join missing from adaptNode's dispatch list, that copy of the UDF reaches isthmus unrewritten and trips "Unable to convert call NUMBER_TO_STRING(fp64?)" in RexExpressionConverter. Add an OpenSearchJoin branch that runs the join's condition through adaptRex with concatenated left+right field storage — same convention OpenSearchJoin#getOutputFieldStorage() uses on the output side. Verified -------- CalciteTransposeCommandIT (with `-Dtests.analytics.parquet_indices=true`) on 21804-merged main: * Without these fixes: 0/5 pass (5/5 hit "filter must be BOOLEAN NOT NULL") * With #1 only: 4/5 pass (testTransposeWithValueFieldNameCollision hits "Unable to convert call NUMBER_TO_STRING(fp64?)") * With #1 + #2: 5/5 pass Signed-off-by: Songkan Tang <songkant@amazon.com> * Read join field storage from the join node, not re-derived from children Address review feedback (expani): use OpenSearchJoin#getOutputFieldStorage() directly in adaptJoin instead of re-assembling it by unwrapping the children — same result (the node derives it from left ++ right child storage, which traces back to the FieldStorageInfo marked on the leaf OpenSearchTableScan), and consistent with how adaptFilter/adaptProject read storage off their node. Also drop the fieldStorage.isEmpty() short-circuit: adaptRex never indexes the storage list (it only hands it to scalar adapters, which no-op when a ref has no storage), so an empty list flows through harmlessly and yields the same result as the explicit guard — matching adaptFilter, which has no such check. BackendPlanAdapterTests 8/8 pass (incl. both join-condition tests). Signed-off-by: Songkan Tang <songkant@amazon.com> --------- Signed-off-by: Songkan Tang <songkant@amazon.com>
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/BackendPlanAdapter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/BackendPlanAdapter.java
@@ -28,6 +28,7 @@
 import org.opensearch.analytics.planner.rel.AggregateMode;
 import org.opensearch.analytics.planner.rel.OpenSearchAggregate;
 import org.opensearch.analytics.planner.rel.OpenSearchFilter;
+import org.opensearch.analytics.planner.rel.OpenSearchJoin;
 import org.opensearch.analytics.planner.rel.OpenSearchProject;
 import org.opensearch.analytics.planner.rel.OpenSearchRelNode;
 import org.opensearch.analytics.planner.rel.OperatorAnnotation;
@@ -104,6 +105,9 @@ private static RelNode adaptNode(RelNode node, Adapters adapters) {
         if (node instanceof OpenSearchProject project) {
             return adaptProject(project, adapters, adaptedChildren, childrenChanged);
         }
+        if (node instanceof OpenSearchJoin join) {
+            return adaptJoin(join, adapters, adaptedChildren, childrenChanged);
+        }
         if (node instanceof OpenSearchAggregate agg && agg.getMode() == AggregateMode.FINAL) {
             OpenSearchAggregate withAdaptedChildren = childrenChanged
                 ? (OpenSearchAggregate) agg.copy(agg.getTraitSet(), adaptedChildren)
@@ -114,6 +118,30 @@ private static RelNode adaptNode(RelNode node, Adapters adapters) {
         return childrenChanged ? node.copy(node.getTraitSet(), adaptedChildren) : node;
     }
 
+    /**
+     * Adapts {@link OpenSearchJoin#getCondition()} so PPL UDFs inlined by
+     * Calcite's FILTER_INTO_JOIN reach the fragment converter in their adapted shape.
+     * Field storage is left ++ right output storage (Calcite join row-type ordering).
+     */
+    private static RelNode adaptJoin(OpenSearchJoin join, Adapters adapters, List<RelNode> adaptedChildren, boolean childrenChanged) {
+        RelNode left = childrenChanged ? adaptedChildren.get(0) : join.getLeft();
+        RelNode right = childrenChanged ? adaptedChildren.get(1) : join.getRight();
+        List<FieldStorageInfo> fieldStorage = join.getOutputFieldStorage();
+        RexNode adaptedCondition = adaptRex(join.getCondition(), adapters, fieldStorage, join.getCluster());
+        if (adaptedCondition != join.getCondition() || childrenChanged) {
+            return new OpenSearchJoin(
+                join.getCluster(),
+                join.getTraitSet(),
+                left,
+                right,
+                adaptedCondition,
+                join.getJoinType(),
+                join.getViableBackends()
+            );
+        }
+        return join;
+    }
+
     private static RelNode adaptFilter(OpenSearchFilter filter, Adapters adapters, List<RelNode> adaptedChildren, boolean childrenChanged) {
         List<FieldStorageInfo> fieldStorage = filter.getOutputFieldStorage();
         RexNode adaptedCondition = adaptRex(filter.getCondition(), adapters, fieldStorage, filter.getCluster());
diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java
@@ -304,7 +304,7 @@ private static AggregateCall buildOne(
                 call.ignoreNulls(),
                 call.rexList,
                 List.copyOf(argList),
-                call.filterArg,
+                -1,  // filterArg dropped: FILTER consumed by PARTIAL on raw rows; FINAL only merges states.
                 call.distinctKeys,
                 call.collation,
                 hasEmptyGroup,
diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/BackendPlanAdapterTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/BackendPlanAdapterTests.java
@@ -11,7 +11,9 @@
 import org.apache.calcite.plan.RelOptTable;
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.JoinRelType;
 import org.apache.calcite.rel.logical.LogicalFilter;
+import org.apache.calcite.rel.logical.LogicalJoin;
 import org.apache.calcite.rel.logical.LogicalProject;
 import org.apache.calcite.rex.RexCall;
 import org.apache.calcite.rex.RexInputRef;
@@ -30,6 +32,7 @@
 import org.opensearch.analytics.planner.PlannerContext;
 import org.opensearch.analytics.planner.rel.AnnotatedPredicate;
 import org.opensearch.analytics.planner.rel.OpenSearchFilter;
+import org.opensearch.analytics.planner.rel.OpenSearchJoin;
 import org.opensearch.analytics.planner.rel.OperatorAnnotation;
 import org.opensearch.analytics.spi.FieldType;
 import org.opensearch.analytics.spi.ProjectCapability;
@@ -311,6 +314,94 @@ protected Map<ScalarFunction, ScalarFunctionAdapter> scalarFunctionAdapters() {
         );
     }
 
+    /**
+     * Join condition carries SIN(integer_column) — a PPL UDF that the SIN adapter rewrites to
+     * SIN(CAST(... AS DOUBLE)). Calcite's FILTER_INTO_JOIN rule inlines outer-Filter predicates
+     * into inner-Join conditions, so any PPL UDF that lived above a Join can ride into the
+     * Join's condition. Without dispatching OpenSearchJoin in adaptNode, that inlined UDF
+     * reaches the substrait converter unrewritten and fails (e.g. PPL transpose's
+     * `Unable to convert call NUMBER_TO_STRING(fp64?)`).
+     *
+     * <p>Verifies the adapter chain runs on the join condition: SIN's INPUT_REF operand
+     * should be wrapped in a CAST after adaptation, exactly as it would on a Filter.
+     */
+    public void testJoinConditionAdapterInsertsCastForIntegerField() {
+        MockDataFusionBackend dfWithAdapter = new MockDataFusionBackend() {
+            @Override
+            protected Map<ScalarFunction, ScalarFunctionAdapter> scalarFunctionAdapters() {
+                return Map.of(ScalarFunction.SIN, sinCastAdapter);
+            }
+        };
+
+        PlannerContext context = buildContext("parquet", 1, intFields(), List.of(dfWithAdapter));
+
+        RelOptTable leftTable = mockTable("test_index", "status", "size");
+        RelOptTable rightTable = mockTable("test_index", "status", "size");
+        RelNode leftScan = stubScan(leftTable);
+        RelNode rightScan = stubScan(rightTable);
+
+        // Join condition: SIN(left.$0) > 0.5. SIN is a PPL UDF whose adapter inserts
+        // CAST(... AS DOUBLE) around any RexInputRef operand whose type is INTEGER/BIGINT.
+        RexNode sinCall = rexBuilder.makeCall(SIN_FUNCTION, rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 0));
+        RexNode condition = rexBuilder.makeCall(
+            SqlStdOperatorTable.GREATER_THAN,
+            sinCall,
+            rexBuilder.makeLiteral(0.5, typeFactory.createSqlType(SqlTypeName.DOUBLE), true)
+        );
+        RelNode join = LogicalJoin.create(leftScan, rightScan, List.of(), condition, Set.of(), JoinRelType.INNER);
+
+        RelNode marked = runPlanner(join, context);
+
+        QueryDAG dag = DAGBuilder.build(marked, context.getCapabilityRegistry(), mockClusterService(), TEST_RESOLVER);
+        PlanForker.forkAll(dag, context.getCapabilityRegistry());
+        BackendPlanAdapter.adaptAll(dag, context.getCapabilityRegistry());
+
+        StagePlan plan = findStagePlanByFragmentType(dag, OpenSearchJoin.class);
+        OpenSearchJoin adaptedJoin = (OpenSearchJoin) plan.resolvedFragment();
+        RexCall sinResult = findCallByName(adaptedJoin.getCondition(), "SIN");
+        assertNotNull("SIN call should exist in adapted join condition", sinResult);
+        assertEquals(
+            "SIN operand should be CAST after adaptation in join condition",
+            SqlKind.CAST,
+            sinResult.getOperands().getFirst().getKind()
+        );
+    }
+
+    /**
+     * Join with no adapted UDFs in its condition — adapter chain runs but rewrites nothing,
+     * the Join's condition passes through structurally unchanged. Guards against the
+     * dispatcher accidentally rebuilding (and possibly mis-configuring) Joins that don't
+     * need adaptation.
+     */
+    public void testJoinConditionNoOpWhenNoAdaptedFunctions() {
+        PlannerContext context = buildContext("parquet", 1, intFields());
+
+        RelOptTable table = mockTable("test_index", "status", "size");
+        RelNode leftScan = stubScan(table);
+        RelNode rightScan = stubScan(table);
+
+        // Equi-join on left.$0 == right.$0 — no PPL UDF, nothing for adapter to rewrite.
+        RexNode condition = rexBuilder.makeCall(
+            SqlStdOperatorTable.EQUALS,
+            rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 0),
+            rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.INTEGER), 2)
+        );
+        RelNode join = LogicalJoin.create(leftScan, rightScan, List.of(), condition, Set.of(), JoinRelType.INNER);
+
+        RelNode marked = runPlanner(join, context);
+        QueryDAG dag = DAGBuilder.build(marked, context.getCapabilityRegistry(), mockClusterService(), TEST_RESOLVER);
+        PlanForker.forkAll(dag, context.getCapabilityRegistry());
+        BackendPlanAdapter.adaptAll(dag, context.getCapabilityRegistry());
+
+        StagePlan plan = findStagePlanByFragmentType(dag, OpenSearchJoin.class);
+        OpenSearchJoin adaptedJoin = (OpenSearchJoin) plan.resolvedFragment();
+        // Condition is structurally identical: EQUALS($0, $2) — both operands still INPUT_REF.
+        RexCall eq = (RexCall) adaptedJoin.getCondition();
+        assertEquals(SqlKind.EQUALS, eq.getKind());
+        assertEquals(SqlKind.INPUT_REF, eq.getOperands().get(0).getKind());
+        assertEquals(SqlKind.INPUT_REF, eq.getOperands().get(1).getKind());
+    }
+
     private static RexCall findCallByName(RexNode node, String name) {
         if (node instanceof AnnotatedPredicate annotated) return findCallByName(annotated.getOriginal(), name);
         if (node instanceof RexCall call) {
