Skip to content

Commit 7f8a5eb

Browse files
songkant-awssongkant-aws
committed
Lower LITERAL_AGG and accept all join condition shapes in marking
Two complementary marking-phase fixes layered on top of PR opensearch-project#21795, both keeping Calcite as the source of truth: 1. PlannerImpl.removeSubQueries — after RelDecorrelator, run a Hep pass with ExtractLiteralAggRule (ported from Apache Impala) to lower LITERAL_AGG(literal) aggregate calls into Aggregate + Project. SubQueryRemoveRule.rewriteIn emits LITERAL_AGG for NOT IN's null-aware semantics; Calcite has no rule to remove it; each downstream engine handles it differently (Druid registers a native impl, Ignite wires an accumulator, Impala lowers to Aggregate + Project, Flink forks SubQueryRemoveRule). The Impala approach is a clean Calcite-rule-style rewrite — port it with attribution. 2. OpenSearchJoinRule.matches() — accept all join condition shapes, not just equi or mixed equi+non-equi. RelDecorrelator's output for correlated subqueries with non-equi correlation predicates (e.g. WHERE id > uid) is a pure non-equi LogicalJoin. DataFusion handles non-equi via NestedLoopJoin downstream, so there is no capability-level reason to reject these at the marking phase. End-to-end IT results on this branch (PR opensearch-project#21795 + these two fixes): ExistsSubqueryCommandIT : 11/11 InSubqueryCommandIT : 11/13 (1 row-order, 1 unrelated AggregateSplitRule bug) ScalarSubqueryCommandIT : 13/13 StreamstatsCommandIT : testStreamstatsReset passes Total : 36/38 = 95% The two remaining failures are unrelated to subquery decorrelation: * testTwoExpressionsInSubquery — result row ordering between hash join's bucketing and the IT's exact-row assertion. Test issue. * testTwoExpressionsNotInSubquery — OpenSearchAggregateSplitRule rejects nullable filter expressions ("filter must be BOOLEAN NOT NULL") on a multi-column NOT IN's COUNT FILTER. Independent AggregateSplitRule bug. Signed-off-by: Songkan Tang <songkant@amazon.com>
1 parent b8e2827 commit 7f8a5eb

3 files changed

Lines changed: 219 additions & 83 deletions

File tree

‎sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ public static RelNode runAllOptimizations(RelNode rawRelNode, PlannerContext con
8989

9090
RelNode modifiedRelNode = rawRelNode;
9191
modifiedRelNode = removeSubQueries(modifiedRelNode, listener);
92+
LOGGER.info("After removeSubQueries:\n{}", RelOptUtil.toString(modifiedRelNode));
9293
modifiedRelNode = reduceExpressions(modifiedRelNode, listener);
9394
modifiedRelNode = pushdownRules(modifiedRelNode, listener);
9495
modifiedRelNode = decomposeAggregates(modifiedRelNode, listener);
@@ -144,12 +145,37 @@ private static RelNode removeSubQueries(RelNode input, RuleProfilingListener lis
144145
RelNode withCorrelates = planner.findBestExp();
145146
// RexSubQuery removal introduces LogicalCorrelate; decorrelate back to a
146147
// straight join shape that the marking + capability rules already handle.
147-
return RelDecorrelator.decorrelateQuery(withCorrelates, RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null));
148+
RelNode decorrelated = RelDecorrelator.decorrelateQuery(
149+
withCorrelates,
150+
RelBuilder.proto(Contexts.empty()).create(input.getCluster(), null)
151+
);
152+
// SubQueryRemoveRule's NOT IN rewrite emits LITERAL_AGG, a Calcite-internal
153+
// aggregate that Calcite has no rule to remove. Lower it to an Aggregate +
154+
// Project before the marking phase observes the plan.
155+
return extractLiteralAgg(decorrelated);
148156
} finally {
149157
if (listener != null) listener.endPhase("subquery-remove");
150158
}
151159
}
152160

161+
/**
162+
* Hep pass that runs {@link org.opensearch.analytics.planner.rules.ExtractLiteralAggRule}
163+
* to lower {@code LITERAL_AGG(literal)} aggregate calls into an Aggregate + Project shape.
164+
*
165+
* <p>Calcite's {@code SubQueryRemoveRule.rewriteIn} emits {@code LITERAL_AGG} for
166+
* {@code NOT IN}'s null-aware semantics. Calcite itself has no rule to remove the function;
167+
* each consuming engine handles it differently (Druid registers a native impl, Ignite wires
168+
* an accumulator, Impala has {@code ExtractLiteralAgg} which we mirror). We pre-lower it
169+
* so DataFusion never sees the operator.
170+
*/
171+
private static RelNode extractLiteralAgg(RelNode input) {
172+
HepProgramBuilder builder = new HepProgramBuilder();
173+
builder.addRuleInstance(new org.opensearch.analytics.planner.rules.ExtractLiteralAggRule());
174+
HepPlanner planner = new HepPlanner(builder.build());
175+
planner.setRoot(input);
176+
return planner.findBestExp();
177+
}
178+
153179
/**
154180
* Phase 1a: constant-expression reduction on Filter and Project predicates. Kept in
155181
* its own phase so {@code ProjectReduceExpressionsRule} cannot use a downstream

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

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,173 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.analytics.planner.rules;
10+
11+
import com.google.common.base.Preconditions;
12+
import org.apache.calcite.plan.RelOptRule;
13+
import org.apache.calcite.plan.RelOptRuleCall;
14+
import org.apache.calcite.rel.RelNode;
15+
import org.apache.calcite.rel.core.AggregateCall;
16+
import org.apache.calcite.rel.logical.LogicalAggregate;
17+
import org.apache.calcite.rel.logical.LogicalProject;
18+
import org.apache.calcite.rex.RexBuilder;
19+
import org.apache.calcite.rex.RexLiteral;
20+
import org.apache.calcite.rex.RexNode;
21+
22+
import java.util.ArrayList;
23+
import java.util.HashSet;
24+
import java.util.List;
25+
26+
/**
27+
* Rewrites a {@link LogicalAggregate} that contains one or more {@code LITERAL_AGG}
28+
* calls into a {@code Project(literal)} layered on top of an aggregate without those
29+
* calls. Handles both the "single LITERAL_AGG" shape and the "mixed" shape where
30+
* {@code LITERAL_AGG} appears alongside regular aggregates.
31+
*
32+
* <p>{@code LITERAL_AGG} is a Calcite-internal aggregate that always returns the
33+
* same literal value supplied at plan time, regardless of the rows in each group.
34+
* It carries a stateless accumulator (state size 0; {@code send} is a no-op,
35+
* {@code end} returns the literal). Calcite's {@link
36+
* org.apache.calcite.rel.rules.SubQueryRemoveRule} emits it as an indicator column
37+
* for {@code NOT IN} / {@code SOME} / {@code ALL} subquery rewrites — the inner
38+
* aggregate's {@code i = LITERAL_AGG(true)} column ends up NULL when the inner
39+
* relation is empty (no group fired), letting the outer side distinguish "no
40+
* match" from "matched with a null".
41+
*
42+
* <p>Two shapes:
43+
* <ul>
44+
* <li>Single LITERAL_AGG (NOT IN's per-uid indicator):
45+
* <pre>
46+
* LogicalAggregate(group=[{0}], i=[LITERAL_AGG(true)])
47+
* </pre></li>
48+
* <li>Mixed (SOME / ALL / NOT EQUALS subqueries):
49+
* <pre>
50+
* LogicalAggregate(group=[{}],
51+
* c=[COUNT()], d=[COUNT(deptno)],
52+
* m=[MAX(deptno)], i=[LITERAL_AGG(true)])
53+
* </pre></li>
54+
* </ul>
55+
*
56+
* <p>Calcite has no rule to remove {@code LITERAL_AGG}; consuming engines are
57+
* expected to either implement the accumulator natively (Calcite's own Interpreter
58+
* and Enumerable codegen do; the SQL JDBC implementor inlines the literal directly
59+
* into the SQL output) or rewrite the call away. Other Calcite-using projects
60+
* (Druid, Ignite) wire it directly into their accumulator impls; Apache Impala's
61+
* {@code ExtractLiteralAgg} rule lowers it to {@code Aggregate + Project} but only
62+
* fires on the single-LITERAL_AGG shape. This rule extends the Impala approach to
63+
* the mixed shape by partitioning the agg call list and emitting the literal back
64+
* at the original output position.
65+
*
66+
* <p>Rewrite (mixed shape illustrated):
67+
* <pre>
68+
* LogicalAggregate(group=[{0}],
69+
* m=[MAX($2)], -- agg call 0, output position group_count+0=1
70+
* i=[LITERAL_AGG(true)], -- agg call 1, output position group_count+1=2
71+
* c=[COUNT()]) -- agg call 2, output position group_count+2=3
72+
*
73+
* becomes
74+
*
75+
* LogicalProject($0, $1, true_literal, $2)
76+
* LogicalAggregate(group=[{0}],
77+
* m=[MAX($2)], -- agg call 0, output position 1
78+
* c=[COUNT()]) -- agg call 1, output position 2
79+
* -- (renumbered down by 1 because
80+
* -- LITERAL_AGG was extracted)
81+
* </pre>
82+
*
83+
* <p>The Project re-emits the schema of the original Aggregate. Group keys pass
84+
* through. For each original agg call, the Project either references the rebuilt
85+
* Aggregate's column (regular agg) or substitutes the literal RexNode
86+
* (LITERAL_AGG). The output column order matches the original so any downstream
87+
* RexInputRef into the aggregate's row type stays valid.
88+
*
89+
* <p>This rule is adapted from Apache Impala's {@code ExtractLiteralAgg}
90+
* ({@code java/calcite-planner/src/main/java/org/apache/impala/calcite/rules/ExtractLiteralAgg.java})
91+
* — the operand pattern, the choice to emit a {@code Project} on top of a
92+
* stripped-down {@code Aggregate}, and the literal source ({@code AggregateCall.rexList.get(0)})
93+
* are the same. Generalised here to the mixed shape.
94+
*
95+
* @opensearch.internal
96+
*/
97+
public class ExtractLiteralAggRule extends RelOptRule {
98+
99+
public ExtractLiteralAggRule() {
100+
super(operand(LogicalAggregate.class, none()), "ExtractLiteralAggRule");
101+
}
102+
103+
@Override
104+
public void onMatch(RelOptRuleCall call) {
105+
final LogicalAggregate aggregate = call.rel(0);
106+
final List<AggregateCall> originalCalls = aggregate.getAggCallList();
107+
108+
// Partition: for each original agg call, either it stays (regular agg) or
109+
// it is extracted (LITERAL_AGG) and replaced by a literal in the Project.
110+
List<AggregateCall> keptCalls = new ArrayList<>(originalCalls.size());
111+
// Per-original-agg-call slot: null → regular agg (use rebuilt aggregate's column),
112+
// non-null → literal RexNode to substitute at this position.
113+
List<RexNode> literalAtSlot = new ArrayList<>(originalCalls.size());
114+
for (AggregateCall ac : originalCalls) {
115+
if (ac.getAggregation().getName().equals("LITERAL_AGG")) {
116+
Preconditions.checkState(
117+
ac.rexList.size() == 1 && ac.rexList.get(0) instanceof RexLiteral,
118+
"LITERAL_AGG must carry exactly one RexLiteral preOperand"
119+
);
120+
literalAtSlot.add(ac.rexList.get(0));
121+
} else {
122+
literalAtSlot.add(null);
123+
keptCalls.add(ac);
124+
}
125+
}
126+
// Nothing to do if no LITERAL_AGG present.
127+
if (literalAtSlot.stream().noneMatch(java.util.Objects::nonNull)) {
128+
return;
129+
}
130+
131+
// Build the rebuilt Aggregate without LITERAL_AGG calls. Same group keys.
132+
LogicalAggregate newAggregate = LogicalAggregate.create(
133+
aggregate.getInput(),
134+
aggregate.getHints(),
135+
aggregate.getGroupSet(),
136+
aggregate.getGroupSets(),
137+
keptCalls
138+
);
139+
140+
// Build the wrapping Project that reproduces the original Aggregate's row type.
141+
RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
142+
List<RexNode> projects = new ArrayList<>();
143+
144+
// Group keys pass through (positions 0 .. groupCount-1 are identical in
145+
// the rebuilt Aggregate's row type because group set was preserved).
146+
int groupCount = aggregate.getGroupSet().cardinality();
147+
for (int i = 0; i < groupCount; i++) {
148+
projects.add(rexBuilder.makeInputRef(newAggregate, i));
149+
}
150+
151+
// For each original agg call, either reference the rebuilt aggregate's
152+
// corresponding column (regular agg) or substitute the literal.
153+
// Walking original positions in order maintains the original schema.
154+
int rebuiltCallIdx = 0;
155+
for (RexNode literal : literalAtSlot) {
156+
if (literal != null) {
157+
projects.add(literal);
158+
} else {
159+
projects.add(rexBuilder.makeInputRef(newAggregate, groupCount + rebuiltCallIdx));
160+
rebuiltCallIdx++;
161+
}
162+
}
163+
164+
RelNode projectRelNode = LogicalProject.create(
165+
newAggregate,
166+
new ArrayList<>(),
167+
projects,
168+
aggregate.getRowType().getFieldNames(),
169+
new HashSet<>()
170+
);
171+
call.transformTo(projectRelNode);
172+
}
173+
}

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

Lines changed: 19 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -10,17 +10,11 @@
1010

1111
import org.apache.calcite.plan.RelOptRule;
1212
import org.apache.calcite.plan.RelOptRuleCall;
13-
import org.apache.calcite.plan.RelOptUtil;
1413
import org.apache.calcite.plan.RelTraitSet;
1514
import org.apache.calcite.plan.hep.HepRelVertex;
1615
import org.apache.calcite.rel.RelNode;
17-
import org.apache.calcite.rel.core.JoinInfo;
1816
import org.apache.calcite.rel.core.JoinRelType;
1917
import org.apache.calcite.rel.logical.LogicalJoin;
20-
import org.apache.calcite.rex.RexCall;
21-
import org.apache.calcite.rex.RexNode;
22-
import org.apache.calcite.sql.SqlKind;
23-
import org.apache.calcite.util.ImmutableBitSet;
2418
import org.opensearch.analytics.planner.PlannerContext;
2519
import org.opensearch.analytics.planner.RelNodeUtils;
2620
import org.opensearch.analytics.planner.rel.OpenSearchDistributionTraitDef;
@@ -39,15 +33,13 @@
3933
* gathered to the coordinator (enforced by the join's cost gate, which only accepts
4034
* SINGLETON inputs — Volcano inserts an {@link OpenSearchExchangeReducer} per side).
4135
*
42-
* <p>Accepts INNER / LEFT / RIGHT / FULL / SEMI / ANTI joins that have at least one
43-
* equi conjunct (so a hash join is viable; any remaining non-equi conjuncts ride
44-
* along as a residual filter). Cross joins (no condition at all) also match. Pure
45-
* non-equi joins (e.g. {@code t1.a < t2.b} as the only condition) are still
46-
* rejected — DataFusion would need a non-equi NestedLoopJoin path we don't enable
47-
* yet, but a mixed condition like
48-
* {@code AND(<=($8, $cor0.__stream_seq__), =($11, $cor0.__seg_id__), =($7, $cor0.DEPTNO))}
49-
* — produced by RelDecorrelator on streamstats reset's directly-built
50-
* LogicalCorrelate — survives via the equi conjuncts.
36+
* <p>Accepts INNER / LEFT / RIGHT / FULL / SEMI / ANTI joins of any condition
37+
* shape — pure equi (hash join), mixed equi + non-equi (hash on the equi keys
38+
* with the residual riding along as a filter), null-safe equi via
39+
* {@code IS_NOT_DISTINCT_FROM}, equi with one operand being a {@code RexCall}
40+
* expression, and pure non-equi (e.g. {@code t1.a < t2.b} alone) — DataFusion
41+
* picks the appropriate physical strategy (HashJoin / SortMergeJoin /
42+
* NestedLoopJoin) based on the condition shape downstream.
5143
*
5244
* @opensearch.internal
5345
*/
@@ -64,9 +56,9 @@ public OpenSearchJoinRule(PlannerContext context) {
6456
public boolean matches(RelOptRuleCall call) {
6557
LogicalJoin join = call.rel(0);
6658
JoinRelType joinType = join.getJoinType();
67-
// Accept INNER / LEFT / RIGHT / FULL / SEMI / ANTI equi-joins. FULL is needed
59+
// Accept INNER / LEFT / RIGHT / FULL / SEMI / ANTI joins. FULL is needed
6860
// by PPL's `appendcol` lowering (ROW_NUMBER pairing via a full outer join on the
69-
// row numbers). Pure non-equi joins are rejected below via JoinInfo.isEqui().
61+
// row numbers).
7062
if (joinType != JoinRelType.INNER
7163
&& joinType != JoinRelType.LEFT
7264
&& joinType != JoinRelType.RIGHT
@@ -75,71 +67,16 @@ public boolean matches(RelOptRuleCall call) {
7567
&& joinType != JoinRelType.ANTI) {
7668
return false;
7769
}
78-
// Accept equi-joins, cross joins, and mixed equi+non-equi conditions.
79-
// * Equi-only joins → JoinInfo.isEqui() true, leftKeys non-empty.
80-
// * Cross joins → condition is literal TRUE; JoinInfo.isEqui() true,
81-
// leftKeys empty.
82-
// * Mixed condition (e.g. AND(<=, =, =)) → isEqui() false because of the
83-
// non-equi conjunct, but leftKeys still carries the equi pairs and
84-
// nonEquiConditions carries the residual. A hash join on the equi keys
85-
// plus a residual filter is viable.
86-
// Pure non-equi joins (no equi conjunct, no leftKey) are still rejected —
87-
// the runtime hash-join path needs at least one equi key. Surfaces in
88-
// RelDecorrelator output for streamstats reset's directly-built
89-
// LogicalCorrelate, which always carries at least one equi correlation
90-
// alongside its range correlation.
91-
JoinInfo info = join.analyzeCondition();
92-
if (!info.leftKeys.isEmpty()) {
93-
// Equi or mixed equi+non-equi (hash key + residual filter both viable).
94-
return true;
95-
}
96-
if (info.isEqui()) {
97-
// Cross join (literal-true condition).
98-
return true;
99-
}
100-
// JoinInfo classifies IS_NOT_DISTINCT_FROM into nonEquiConditions in some
101-
// call paths even though splitJoinCondition handles it as a null-safe equi
102-
// when filterNulls is non-null. Walk the condition's conjuncts ourselves —
103-
// if any one is a binary EQUALS / IS_NOT_DISTINCT_FROM between RexInputRefs
104-
// straddling the inputs, the join is hashable on that pair (DataFusion
105-
// builds a hash join key, residual conditions ride along as a filter).
106-
if (hasInputRefEquiOrNullSafeConjunct(join)) {
107-
return true;
108-
}
109-
return false;
110-
}
111-
112-
/** Returns true if the join condition has at least one conjunct of the form
113-
* {@code RexInputRef = RexInputRef} or
114-
* {@code RexInputRef IS NOT DISTINCT FROM RexInputRef} where the two refs
115-
* straddle the left / right side. The dispatch in {@code Join#analyzeCondition()}
116-
* occasionally puts {@code IS_NOT_DISTINCT_FROM} into {@code nonEquiConditions}
117-
* even though {@code splitJoinCondition} logic recognises it as a null-safe
118-
* equi join key; this method is a robust fallback. */
119-
private static boolean hasInputRefEquiOrNullSafeConjunct(LogicalJoin join) {
120-
int leftFieldCount = join.getLeft().getRowType().getFieldCount();
121-
int rightFieldCount = join.getRight().getRowType().getFieldCount();
122-
ImmutableBitSet leftRange = ImmutableBitSet.range(0, leftFieldCount);
123-
ImmutableBitSet rightRange = ImmutableBitSet.range(leftFieldCount, leftFieldCount + rightFieldCount);
124-
for (RexNode conjunct : RelOptUtil.conjunctions(join.getCondition())) {
125-
if (!(conjunct instanceof RexCall call)) continue;
126-
SqlKind kind = call.getKind();
127-
if (kind != SqlKind.EQUALS && kind != SqlKind.IS_NOT_DISTINCT_FROM) continue;
128-
if (call.getOperands().size() != 2) continue;
129-
ImmutableBitSet op0Refs = RelOptUtil.InputFinder.bits(call.getOperands().get(0));
130-
ImmutableBitSet op1Refs = RelOptUtil.InputFinder.bits(call.getOperands().get(1));
131-
// Each operand must reference exactly one side; the two sides must be different.
132-
// DataFusion's hash-join evaluates each operand expression to derive the key —
133-
// RexCall arithmetic (e.g. {@code uid = salary + 1000}) is acceptable.
134-
boolean op0Left = leftRange.contains(op0Refs);
135-
boolean op0Right = rightRange.contains(op0Refs);
136-
boolean op1Left = leftRange.contains(op1Refs);
137-
boolean op1Right = rightRange.contains(op1Refs);
138-
if ((op0Left && op1Right) || (op0Right && op1Left)) {
139-
return true;
140-
}
141-
}
142-
return false;
70+
// All condition shapes — equi, mixed, null-safe, expression-key, pure
71+
// non-equi — are accepted. DataFusion's join planner picks the right
72+
// physical strategy:
73+
// * equi keys present → HashJoin
74+
// * IS_NOT_DISTINCT_FROM → HashJoin with null-safe key
75+
// * pure non-equi → NestedLoopJoin (RelDecorrelator output for
76+
// correlated subqueries with non-equi correlation predicates,
77+
// e.g. {@code outer.id > inner.uid}, lands here)
78+
// * literal-true (cross join) → CrossJoin
79+
return true;
14380
}
14481

14582
@Override

0 commit comments

Comments
 (0)