You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by da...@apache.org on 2026/07/27 20:45:07 UTC
(beam) 01/01: Decorrelate sub-queries using normalizeForVolcano pass
This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch backport-12c517dd
in repository https://gitbox.apache.org/repos/asf/beam.git
commit 4db4912bf72bd3ac4fcbeffba579efd197722b7e
Author: Danny McCormick <da...@google.com>
AuthorDate: Mon Jul 27 20:44:55 2026 +0000
Decorrelate sub-queries using normalizeForVolcano pass
---
.../extensions/sql/impl/CalciteQueryPlanner.java | 76 +++++++++++++++++++++-
1 file changed, 75 insertions(+), 1 deletion(-)
diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java
index aa6f4d12187..c45bfc5be3e 100644
--- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java
+++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/CalciteQueryPlanner.java
@@ -37,9 +37,13 @@ import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalciteSch
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.Contexts;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.ConventionTraitDef;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptCost;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptPlanner;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptPlanner.CannotPlanException;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptUtil;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelTraitDef;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelTraitSet;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.hep.HepPlanner;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.hep.HepProgramBuilder;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.prepare.CalciteCatalogReader;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelNode;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.RelRoot;
@@ -51,9 +55,12 @@ import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.metadata.Me
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.metadata.ReflectiveRelMetadataProvider;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.metadata.RelMetadataProvider;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.rules.CoreRules;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.type.RelDataType;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.type.RelDataTypeFactory;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexBuilder;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexDynamicParam;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexLiteral;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexNode;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexShuttle;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.schema.SchemaPlus;
@@ -64,6 +71,9 @@ import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlP
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParser;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParserImplFactory;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.util.SqlOperatorTables;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.validate.SqlConformance;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.validate.SqlConformanceEnum;
+import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql2rel.RelDecorrelator;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql2rel.SqlToRelConverter;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.tools.FrameworkConfig;
import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.tools.Frameworks;
@@ -203,7 +213,15 @@ public class CalciteQueryPlanner implements QueryPlanner {
relNode,
new ParameterBinder(root.rel.getCluster().getRexBuilder(), queryParameters));
}
- LOG.info("SQLPlan>\n{}", BeamSqlRelUtils.explainLazily(root.rel));
+ // Normalize correlated sub-queries into standard relational shapes (Join + Aggregate, etc.)
+ // BEFORE the Volcano program runs and BEFORE the metadata-provider swap below. Running this
+ // here keeps it off the cost-based metadata path (it uses stock Calcite metadata), so it
+ // cannot trigger the BeamCostModel / RelMdNodeStats recursion that the Volcano search guards
+ // against. The pre-pass is a no-op on trees that carry no referenced correlation variable
+ // (see normalizeForVolcano), which protects the UNNEST LogicalCorrelate(_, Uncollect) shape
+ // that BeamUnnestRule depends on.
+ relNode = normalizeForVolcano(relNode);
+ LOG.info("SQLPlan>\n{}", BeamSqlRelUtils.explainLazily(relNode));
RelTraitSet desiredTraits =
relNode
.getTraitSet()
@@ -246,6 +264,62 @@ public class CalciteQueryPlanner implements QueryPlanner {
return newRel.copy(newRel.getTraitSet(), newInputs);
}
+ /**
+ * Pre-Volcano normalization pass for correlated sub-queries.
+ *
+ * <p>The SqlToRel converter (driven by {@link Planner#rel}) already decorrelates most queries,
+ * but some correlated shapes (e.g. correlated EXISTS/IN inside a project or join condition)
+ * survive as a {@code RexSubQuery} or a residual {@code LogicalCorrelate}. The Beam Volcano
+ * ruleset has no converter rule for a general {@code LogicalCorrelate}, so such a residue would
+ * fail planning with a {@code CannotPlanException}. This pass rewrites those into standard
+ * relational nodes ({@code Join}, {@code Aggregate}, {@code Project}, {@code Filter}) that
+ * existing Beam converter rules already cover.
+ *
+ * <p>The pass:
+ *
+ * <ol>
+ * <li>runs a short-lived {@link HepPlanner} with the three {@code *_SUB_QUERY_TO_CORRELATE}
+ * rules to turn any un-expanded {@code RexSubQuery} into a {@code LogicalCorrelate}, then
+ * <li>runs {@link RelDecorrelator#decorrelateQuery(RelNode, RelBuilder)} to lower correlates
+ * into joins/aggregates, using the planner's configured {@link RelBuilder} so produced rels
+ * share the cluster's type factory.
+ * </ol>
+ *
+ * <p>It runs strictly before the Volcano {@code program.run(...)} and before the
+ * metadata-provider swap in {@link #convertToBeamRel(RelNode, QueryParameters)}, so it stays off
+ * the Beam cost path.
+ *
+ * <p>Pre-flight safety: the whole pass is gated on the tree actually <em>referencing</em> a
+ * correlation variable ({@link RelOptUtil#getVariablesUsed(RelNode)} non-empty). This makes it a
+ * strict no-op on trees without a referenced correlate. In particular the UNNEST shape {@code
+ * LogicalCorrelate(_, Uncollect)} <em>defines</em> a correlation id but does not
+ * <em>reference</em> one in its body, so {@code getVariablesUsed} is empty for it and the pass is
+ * skipped — leaving the correlate intact for {@code BeamUnnestRule}. (Note: {@code
+ * getVariablesSet} would be wrong here, as it is non-empty for UNNEST.)
+ */
+ private RelNode normalizeForVolcano(RelNode rel) {
+ // No-op unless the tree references a correlation variable. Protects the UNNEST
+ // LogicalCorrelate(_, Uncollect) shape, which defines but does not reference a correl var.
+ if (RelOptUtil.getVariablesUsed(rel).isEmpty()) {
+ return rel;
+ }
+
+ // (1) Convert any residual RexSubQuery (correlated EXISTS/IN in PROJECT/JOIN/FILTER that the
+ // SqlToRel converter left in place) into a LogicalCorrelate via a tiny HEP pass.
+ HepProgramBuilder hep =
+ new HepProgramBuilder()
+ .addRuleInstance(CoreRules.FILTER_SUB_QUERY_TO_CORRELATE)
+ .addRuleInstance(CoreRules.PROJECT_SUB_QUERY_TO_CORRELATE)
+ .addRuleInstance(CoreRules.JOIN_SUB_QUERY_TO_CORRELATE);
+ HepPlanner hepPlanner = new HepPlanner(hep.build());
+ hepPlanner.setRoot(rel);
+ RelNode noSubQuery = hepPlanner.findBestExp();
+
+ // (2) Decorrelate the LogicalCorrelate nodes into standard Join/Aggregate shapes. Uses the
+ // planner's RelBuilder so produced rels share the cluster's type factory + traits.
+ return RelDecorrelator.decorrelateQuery(noSubQuery, RelBuilder.create(config));
+ }
+
// It needs to be public so that the generated code in Calcite can access it.
public static class NonCumulativeCostImpl
implements MetadataHandler<BuiltInMetadata.NonCumulativeCost> {