You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@calcite.apache.org by su...@apache.org on 2025/05/21 07:31:31 UTC
(calcite) branch main updated: [CALCITE-5638] Columns trimmer need to consider sub queries
This is an automated email from the ASF dual-hosted git repository.
suibianwanwan pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/calcite.git
The following commit(s) were added to refs/heads/main by this push:
new 00e09ac15e [CALCITE-5638] Columns trimmer need to consider sub queries
00e09ac15e is described below
commit 00e09ac15e09f2e1bd49a8b461818bb0ae47dfbc
Author: Evgeniy Stanilovsky <zs...@apache.com>
AuthorDate: Tue Apr 22 13:17:23 2025 +0300
[CALCITE-5638] Columns trimmer need to consider sub queries
---
.../java/org/apache/calcite/plan/RelOptUtil.java | 31 ++++++
.../calcite/rel/rules/SubQueryRemoveRule.java | 19 +++-
.../apache/calcite/sql2rel/RelFieldTrimmer.java | 122 ++++++++++++++++++++-
.../test/enumerable/EnumerableCorrelateTest.java | 18 +++
core/src/test/resources/sql/sub-query.iq | 28 ++++-
5 files changed, 210 insertions(+), 8 deletions(-)
diff --git a/core/src/main/java/org/apache/calcite/plan/RelOptUtil.java b/core/src/main/java/org/apache/calcite/plan/RelOptUtil.java
index 01e712c900..88f083ea2a 100644
--- a/core/src/main/java/org/apache/calcite/plan/RelOptUtil.java
+++ b/core/src/main/java/org/apache/calcite/plan/RelOptUtil.java
@@ -4552,6 +4552,37 @@ private void acceptFields(final List<RelDataTypeField> fields) {
}
}
+ /** Extension of {@link RelOptUtil.InputFinder} with optional subquery lookup. */
+ public static class SubQueryAwareInputFinder extends RelOptUtil.InputFinder {
+ boolean visitSubQuery;
+
+ public SubQueryAwareInputFinder(@Nullable Set<RelDataTypeField> extraFields,
+ boolean visitSubQuery) {
+ super(extraFields, ImmutableBitSet.builder());
+ this.visitSubQuery = visitSubQuery;
+ }
+
+ @Override public Void visitSubQuery(RexSubQuery subQuery) {
+ if (visitSubQuery && subQuery.getKind() == SqlKind.SCALAR_QUERY) {
+ subQuery.rel.accept(new RelHomogeneousShuttle() {
+ @Override public RelNode visit(LogicalProject project) {
+ project.getProjects().forEach(r -> r.accept(SubQueryAwareInputFinder.this));
+ return super.visit(project);
+ }
+
+ @Override public RelNode visit(LogicalFilter filter) {
+ filter.getCondition().accept(SubQueryAwareInputFinder.this);
+ return super.visit(filter);
+ }
+ });
+
+ return null;
+ } else {
+ return super.visitSubQuery(subQuery);
+ }
+ }
+ }
+
/**
* Visitor which builds a bitmap of the inputs used by an expression.
*/
diff --git a/core/src/main/java/org/apache/calcite/rel/rules/SubQueryRemoveRule.java b/core/src/main/java/org/apache/calcite/rel/rules/SubQueryRemoveRule.java
index 54a7543608..fd6e3d7d74 100644
--- a/core/src/main/java/org/apache/calcite/rel/rules/SubQueryRemoveRule.java
+++ b/core/src/main/java/org/apache/calcite/rel/rules/SubQueryRemoveRule.java
@@ -910,6 +910,24 @@ private static void matchJoin(SubQueryRemoveRule rule, RelOptRuleCall call) {
int nFieldsLeft = join.getLeft().getRowType().getFieldCount();
int nFieldsRight = join.getRight().getRowType().getFieldCount();
+ // Correlation columns should also be considered.
+ // For example:
+ // LogicalJoin
+ // left right
+ // | |
+ // LogicalProject.NONE.[0, 1] LogicalValues.NONE.[0]
+ // RecordType(INTEGER DEPTNO, CHAR(11) DNAME) RecordType(INTEGER DEPTNO)
+ //
+ // and subquery: $SCALAR_QUERY with correlate
+ // LogicalProject(DEPTNO=[$1])
+ // LogicalFilter(condition=[=(CAST($0):CHAR(11) NOT NULL, $cor0.DNAME)])
+ //
+ // In such a case $cor0.DNAME need to be accounted as input form left side.
+ final Set<CorrelationId> variablesSet = RelOptUtil.getVariablesUsed(e.rel);
+ for (CorrelationId id : variablesSet) {
+ ImmutableBitSet requiredColumns = RelOptUtil.correlationColumns(id, e.rel);
+ inputSet = ImmutableBitSet.union(ImmutableList.of(requiredColumns, inputSet));
+ }
boolean inputIntersectsLeftSide = inputSet.intersects(ImmutableBitSet.range(0, nFieldsLeft));
boolean inputIntersectsRightSide =
@@ -922,7 +940,6 @@ private static void matchJoin(SubQueryRemoveRule rule, RelOptRuleCall call) {
return;
}
- final Set<CorrelationId> variablesSet = RelOptUtil.getVariablesUsed(e.rel);
if (inputIntersectsLeftSide) {
builder.push(join.getLeft());
diff --git a/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java b/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
index effb5078f1..2dee1af505 100644
--- a/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
+++ b/core/src/main/java/org/apache/calcite/sql2rel/RelFieldTrimmer.java
@@ -23,6 +23,7 @@
import org.apache.calcite.rel.RelCollations;
import org.apache.calcite.rel.RelDistribution;
import org.apache.calcite.rel.RelFieldCollation;
+import org.apache.calcite.rel.RelHomogeneousShuttle;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.core.Aggregate;
import org.apache.calcite.rel.core.AggregateCall;
@@ -61,6 +62,7 @@
import org.apache.calcite.rex.RexSubQuery;
import org.apache.calcite.rex.RexUtil;
import org.apache.calcite.rex.RexVisitor;
+import org.apache.calcite.rex.RexVisitorImpl;
import org.apache.calcite.sql.SqlExplainFormat;
import org.apache.calcite.sql.SqlExplainLevel;
import org.apache.calcite.sql.SqlKind;
@@ -78,6 +80,8 @@
import org.apache.calcite.util.mapping.Mappings;
import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
import org.checkerframework.checker.nullness.qual.Nullable;
@@ -473,6 +477,77 @@ public TrimResult trimFields(
return result(newCalc, mapping, calc);
}
+ /**
+ * Shuttle that finds all {@link TableScan}s inside a given {@link RelNode}.
+ */
+ private static class TableScanCollector extends RelHomogeneousShuttle {
+ private ImmutableSet.Builder<List<String>> builder = ImmutableSet.builder();
+
+ /** Qualified names. */
+ Set<List<String>> tables() {
+ return builder.build();
+ }
+
+ @Override public RelNode visit(TableScan scan) {
+ builder.add(scan.getTable().getQualifiedName());
+ return super.visit(scan);
+ }
+ }
+
+ /**
+ * Shuttle that finds all {@link TableScan}`s inside a given {@link RexNode}.
+ */
+ private static class InputTablesVisitor extends RexVisitorImpl<Void> {
+ private ImmutableSet.Builder<List<String>> builder = ImmutableSet.builder();
+
+ protected InputTablesVisitor() {
+ super(false);
+ }
+
+ /** Qualified names. */
+ Set<List<String>> tables() {
+ return builder.build();
+ }
+
+ @Override public Void visitSubQuery(RexSubQuery subQuery) {
+ if (subQuery.getKind() == SqlKind.SCALAR_QUERY) {
+ subQuery.rel.accept(new RelHomogeneousShuttle() {
+ @Override public RelNode visit(TableScan scan) {
+ builder.add(scan.getTable().getQualifiedName());
+ return super.visit(scan);
+ }
+ });
+ }
+ return null;
+ }
+ }
+
+ private boolean inputContainsSubQueryTables(Project project, RelNode input) {
+ InputTablesVisitor inputSubQueryTablesCollector = new InputTablesVisitor();
+
+ RexUtil.apply(inputSubQueryTablesCollector, project.getProjects(), null);
+
+ Set<List<String>> subQueryTables = inputSubQueryTablesCollector.tables();
+
+ assert subQueryTables.isEmpty() || subQueryTables.size() == 1
+ : "unexpected different tables in subquery: " + subQueryTables;
+
+ TableScanCollector inputTablesCollector = new TableScanCollector();
+ input.accept(inputTablesCollector);
+
+ Set<List<String>> inputTables = inputTablesCollector.tables();
+ // Check for input and subquery tables intersection.
+ if (!subQueryTables.isEmpty()) {
+ for (List<String> t : inputTables) {
+ if (t.equals(Iterables.getOnlyElement(subQueryTables))) {
+ return true;
+ }
+ }
+ }
+
+ return false;
+ }
+
/**
* Variant of {@link #trimFields(RelNode, ImmutableBitSet, Set)} for
* {@link org.apache.calcite.rel.logical.LogicalProject}.
@@ -488,18 +563,24 @@ public TrimResult trimFields(
// Which fields are required from the input?
final Set<RelDataTypeField> inputExtraFields =
new LinkedHashSet<>(extraFields);
+
+ // Collect all the SubQueries in the projection list.
+ List<RexSubQuery> subQueries = RexUtil.SubQueryCollector.collect(project);
+ // Get all the correlationIds present in the SubQueries
+ Set<CorrelationId> correlationIds = RelOptUtil.getVariablesUsed(subQueries);
+ // Subquery lookup is required.
+ boolean subQueryLookUp =
+ !correlationIds.isEmpty() && inputContainsSubQueryTables(project, input);
+
RelOptUtil.InputFinder inputFinder =
- new RelOptUtil.InputFinder(inputExtraFields);
+ new RelOptUtil.SubQueryAwareInputFinder(inputExtraFields, subQueryLookUp);
+
for (Ord<RexNode> ord : Ord.zip(project.getProjects())) {
if (fieldsUsed.get(ord.i)) {
ord.e.accept(inputFinder);
}
}
- // Collect all the SubQueries in the projection list.
- List<RexSubQuery> subQueries = RexUtil.SubQueryCollector.collect(project);
- // Get all the correlationIds present in the SubQueries
- Set<CorrelationId> correlationIds = RelOptUtil.getVariablesUsed(subQueries);
ImmutableBitSet requiredColumns = ImmutableBitSet.of();
if (!correlationIds.isEmpty()) {
assert correlationIds.size() == 1;
@@ -547,6 +628,14 @@ public TrimResult trimFields(
if (fieldsUsed.get(ord.i)) {
mapping.set(ord.i, newProjects.size());
RexNode newProjectExpr = ord.e.accept(shuttle);
+ // Subquery need to be remapped
+ if (newProjectExpr instanceof RexSubQuery
+ && newProjectExpr.getKind() == SqlKind.SCALAR_QUERY
+ && !correlationIds.isEmpty()) {
+ newProjectExpr =
+ changeCorrelateReferences((RexSubQuery) newProjectExpr,
+ Iterables.getOnlyElement(correlationIds), newInput.getRowType(), inputMapping);
+ }
newProjects.add(newProjectExpr);
}
}
@@ -561,6 +650,24 @@ public TrimResult trimFields(
return result(newProject, mapping, project);
}
+ private RexNode changeCorrelateReferences(
+ RexSubQuery node,
+ CorrelationId corrId,
+ RelDataType rowType,
+ Mapping inputMapping) {
+ assert node.getKind() == SqlKind.SCALAR_QUERY : "Expected a SCALAR_QUERY, found "
+ + node.getKind();
+ RelNode subQuery = node.rel;
+ RexBuilder rexBuilder = relBuilder.getRexBuilder();
+
+ RexCorrelVariableMapShuttle rexVisitor =
+ new RexCorrelVariableMapShuttle(corrId, rowType, inputMapping, rexBuilder);
+ RelNode newSubQuery =
+ subQuery.accept(new RexRewritingRelShuttle(rexVisitor));
+
+ return RexSubQuery.scalar(newSubQuery);
+ }
+
/** Creates a project with a dummy column, to protect the parts of the system
* that cannot handle a relational expression with no columns.
*
@@ -1437,9 +1544,12 @@ static class RexCorrelVariableMapShuttle extends RexShuttle {
(RexCorrelVariable) fieldAccess.getReferenceExpr();
if (referenceExpr.id.equals(correlationId)) {
int oldIndex = fieldAccess.getField().getIndex();
+ int newIndex = mapping.getTarget(oldIndex);
+ if (newIndex == oldIndex) {
+ return super.visitFieldAccess(fieldAccess);
+ }
RexNode newCorrel =
rexBuilder.makeCorrel(newCorrelRowType, referenceExpr.id);
- int newIndex = mapping.getTarget(oldIndex);
return rexBuilder.makeFieldAccess(newCorrel, newIndex);
}
}
diff --git a/core/src/test/java/org/apache/calcite/test/enumerable/EnumerableCorrelateTest.java b/core/src/test/java/org/apache/calcite/test/enumerable/EnumerableCorrelateTest.java
index e1cd3dda83..19cbdb5b8f 100644
--- a/core/src/test/java/org/apache/calcite/test/enumerable/EnumerableCorrelateTest.java
+++ b/core/src/test/java/org/apache/calcite/test/enumerable/EnumerableCorrelateTest.java
@@ -168,6 +168,24 @@ class EnumerableCorrelateTest {
"empid=200");
}
+ /** Test case for
+ * <a href="https://issues.apache.org/jira/browse/CALCITE-5638">[CALCITE-5638]
+ * Columns trimmer need to consider sub queries</a>.
+ */
+ @Test void complexNestedCorrelatedSubquery() {
+ String sql = "SELECT empid, deptno, (SELECT count(*) FROM emps AS x "
+ + "WHERE x.salary>emps.salary and x.deptno<emps.deptno) FROM emps "
+ + "WHERE empid<salary ORDER BY 1,2,3";
+
+ tester(false, new HrSchema())
+ .query(sql)
+ .returnsOrdered(
+ "empid=100; deptno=10; EXPR$2=0",
+ "empid=110; deptno=10; EXPR$2=0",
+ "empid=150; deptno=10; EXPR$2=0",
+ "empid=200; deptno=20; EXPR$2=2");
+ }
+
/** Test case for
* <a href="https://issues.apache.org/jira/browse/CALCITE-2920">[CALCITE-2920]
* RelBuilder: new method to create an anti-join</a>. */
diff --git a/core/src/test/resources/sql/sub-query.iq b/core/src/test/resources/sql/sub-query.iq
index 59968cdec3..12f7e1e8cc 100644
--- a/core/src/test/resources/sql/sub-query.iq
+++ b/core/src/test/resources/sql/sub-query.iq
@@ -63,6 +63,32 @@ select * from dept where deptno not in (select deptno from emp);
(0 rows)
!ok
+
+# [CALCITE-5638] Assertion Failure during planning correlated query
+SELECT "hr"."emps"."empid", "hr"."emps"."deptno",
+ (SELECT count(*) FROM "hr"."emps" AS x WHERE x."salary">"hr"."emps"."salary" AND x."deptno"<"hr"."emps"."deptno")
+FROM "hr"."emps"
+ WHERE "hr"."emps"."empid"<"hr"."emps"."salary"
+ ORDER BY 1,2,3;
+ empid | deptno | EXPR$2
+-------+--------+--------
+ 100 | 10 | 0
+ 110 | 10 | 0
+ 150 | 10 | 0
+ 200 | 20 | 2
+(4 rows)
+
+!ok
+
+# [CALCITE-5638] Assertion Failure during planning correlated query
+SELECT t1.deptno FROM dept AS t0 JOIN emp AS t1 ON
+(t1.deptno = (SELECT inner_t1.deptno FROM emp AS inner_t1 WHERE inner_t1.ENAME = t0.DNAME));
+ DEPTNO
+--------
+(0 rows)
+
+!ok
+
select deptno, deptno in (select deptno from emp) from dept;
DEPTNO | EXPR$1
--------+--------
@@ -873,7 +899,7 @@ EnumerableCalc(expr#0..2=[{inputs}], proj#0..1=[{exprs}])
EnumerableCorrelate(correlation=[$cor0], joinType=[inner], requiredColumns=[{0}])
EnumerableValues(tuples=[[{ 1, 2 }]])
EnumerableAggregate(group=[{0}])
- EnumerableCalc(expr#0..7=[{inputs}], expr#8=[true], expr#9=[CAST($t7):INTEGER], expr#10=[$cor0], expr#11=[$t10.EXPR$0], expr#12=[=($t9, $t11)], i=[$t8], $condition=[$t12])
+ EnumerableCalc(expr#0..7=[{inputs}], expr#8=[true], expr#9=[CAST($t7):INTEGER], expr#10=[$cor0], expr#11=[$t10.A], expr#12=[=($t9, $t11)], i=[$t8], $condition=[$t12])
EnumerableTableScan(table=[[scott, EMP]])
!plan