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 <[email protected]> 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> {
