From f30927d31a524b3b95c61d47044f8ca82f4a5975 Mon Sep 17 00:00:00 2001 From: Kenneth Knowles Date: Thu, 27 Aug 2026 15:40:37 +0000 Subject: [PATCH] Beam SQL: name query composites after what was fused into them Beam SQL named each composite it expands a query into after the relational node class and a JVM-global counter, e.g. BeamCalcRel_57. That counter was global, so the same query produced different names on different runs, and the names said nothing about what the stage does. Name a stage after the operations fused into it instead, e.g. Filter;Project. The provenance label rides along as a RelHint so that it survives copy() and the conversion to physical rels, and is composed as a canonical union of the labels of the nodes a rule matched -- order independent, so transpose rules trading a pair back and forth cannot hand the planner an endless supply of rels it has not seen. Since Dataflow matches streaming pipelines for update by step name, the experiment legacy-sql-transform-names restores the old naming for pipelines that need to be updated across this change. Co-Authored-By: Claude Opus 4.8 --- CHANGES.md | 9 + .../sql/impl/CalciteQueryPlanner.java | 11 +- .../sdk/extensions/sql/impl/JdbcDriver.java | 3 +- .../sql/impl/rel/AbstractBeamCalcRel.java | 14 +- .../extensions/sql/impl/rel/BeamCalcRel.java | 19 +- .../sql/impl/rel/BeamSqlRelUtils.java | 47 ++- .../extensions/sql/impl/rel/StageName.java | 344 ++++++++++++++++++ .../sql/impl/rule/BeamCalcRule.java | 1 + .../sql/impl/rule/StageNameRule.java | 196 ++++++++++ .../sql/impl/rule/StageNameRuleCall.java | 212 +++++++++++ .../sql/impl/rel/StageNameTest.java | 248 +++++++++++++ .../sql/impl/rule/StageNameRuleTest.java | 123 +++++++ .../parquet/ParquetTableProviderTest.java | 29 +- 13 files changed, 1239 insertions(+), 17 deletions(-) create mode 100644 sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/StageName.java create mode 100644 sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRule.java create mode 100644 sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRuleCall.java create mode 100644 sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/StageNameTest.java create mode 100644 sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRuleTest.java diff --git a/CHANGES.md b/CHANGES.md index 73966a48313c..44e3aef29116 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -183,6 +183,15 @@ ## Breaking Changes +* (Java) Beam SQL now names the transforms a query expands into after the operations fused into them + (`Filter;Project`) instead of after the relational node class and a JVM-global counter + (`BeamCalcRel_57`). The old counter was JVM-global, so the old names differed between runs of the + same query. The new names do not, but two stages with the same provenance are still told apart by + an occurrence suffix (`Filter;Project #2`) numbered within the query, so reshaping a plan can + renumber the stages of that query. Since Dataflow matches streaming pipelines for update by step + name, a running streaming pipeline must either be drained or be started with + `--experiments=legacy-sql-transform-names` to keep the old names. That experiment exists only to + carry running pipelines over this change and is expected to be removed two releases from now. * (Python) Typehints of dataclass fields are honored during type inferences. To restore the behavior of fallback-to-any, use pipeline option `--exclude_infer_dataclass_field_type` ([#38797](https://github.com/apache/beam/issues/38797)). However fixing forward is recommended. 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 cedfc5dc6d1d..00164f14ed94 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 @@ -30,6 +30,8 @@ import org.apache.beam.sdk.extensions.sql.impl.rel.BeamLogicalConvention; import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode; import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils; +import org.apache.beam.sdk.extensions.sql.impl.rel.StageName; +import org.apache.beam.sdk.extensions.sql.impl.rule.StageNameRule; import org.apache.beam.sdk.extensions.sql.impl.udf.BeamBuiltinFunctionProvider; import org.apache.beam.vendor.calcite.v1_40_0.com.google.common.collect.Table; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.config.CalciteConnectionConfig; @@ -155,7 +157,9 @@ public FrameworkConfig defaultConfig(JdbcConnection connection, Collection\n{}", BeamSqlRelUtils.explainLazily(root.rel)); + StageName.register(relNode.getCluster()); + // Give every node a name to be composed from before any rule fuses it away. + relNode = StageName.backfill(relNode); + LOG.info("SQLPlan>\n{}", BeamSqlRelUtils.explainLazily(relNode)); RelTraitSet desiredTraits = relNode .getTraitSet() diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriver.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriver.java index ddc6c9e7500b..6f5bde501a83 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriver.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/JdbcDriver.java @@ -30,6 +30,7 @@ import java.util.function.Consumer; import org.apache.beam.sdk.extensions.sql.SqlTransform; import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRuleSets; +import org.apache.beam.sdk.extensions.sql.impl.rule.StageNameRule; import org.apache.beam.sdk.extensions.sql.meta.catalog.CatalogManager; import org.apache.beam.sdk.extensions.sql.meta.provider.TableProvider; import org.apache.beam.sdk.options.PipelineOptions; @@ -80,7 +81,7 @@ public class JdbcDriver extends Driver { planner -> { for (RuleSet ruleSet : BeamRuleSets.getRuleSets()) { for (RelOptRule rule : ruleSet) { - planner.addRule(rule); + StageNameRule.addTo(planner, rule); } } planner.removeRule(CoreRules.CALC_REMOVE); diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/AbstractBeamCalcRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/AbstractBeamCalcRel.java index cb2c9598f34f..081f3c15b9b7 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/AbstractBeamCalcRel.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/AbstractBeamCalcRel.java @@ -17,6 +17,7 @@ */ package org.apache.beam.sdk.extensions.sql.impl.rel; +import java.util.List; import org.apache.beam.sdk.annotations.Internal; import org.apache.beam.sdk.extensions.sql.impl.planner.BeamCostModel; import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery; @@ -26,10 +27,12 @@ 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.rel.RelNode; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.Calc; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.hint.RelHint; 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.rex.RexLocalRef; 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.RexProgram; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; /** BeamRelNode to replace {@code Project} and {@code Filter} node. */ @Internal @@ -40,7 +43,16 @@ public abstract class AbstractBeamCalcRel extends Calc implements BeamRelNode { public AbstractBeamCalcRel( RelOptCluster cluster, RelTraitSet traits, RelNode input, RexProgram program) { - super(cluster, traits, input, program); + this(cluster, traits, ImmutableList.of(), input, program); + } + + public AbstractBeamCalcRel( + RelOptCluster cluster, + RelTraitSet traits, + List hints, + RelNode input, + RexProgram program) { + super(cluster, traits, hints, input, program); } public boolean isInputSortRelAndLimitOnly() { diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java index b9525bd07dd0..29858ed99d9a 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamCalcRel.java @@ -90,6 +90,7 @@ 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.rel.RelNode; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.core.Calc; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.hint.RelHint; 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.rex.RexBuilder; import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rex.RexCall; @@ -151,12 +152,26 @@ public static long timestampToCalciteMillis(java.time.Instant instant) { } public BeamCalcRel(RelOptCluster cluster, RelTraitSet traits, RelNode input, RexProgram program) { - super(cluster, traits, input, program); + this(cluster, traits, ImmutableList.of(), input, program); + } + + public BeamCalcRel( + RelOptCluster cluster, + RelTraitSet traits, + List hints, + RelNode input, + RexProgram program) { + super(cluster, traits, hints, input, program); } @Override public Calc copy(RelTraitSet traitSet, RelNode input, RexProgram program) { - return new BeamCalcRel(getCluster(), traitSet, input, program); + return new BeamCalcRel(getCluster(), traitSet, hints, input, program); + } + + @Override + public RelNode withHints(List hintList) { + return new BeamCalcRel(getCluster(), traitSet, hintList, input, program); } @Override diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSqlRelUtils.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSqlRelUtils.java index 43e9b7ff333b..668a94846264 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSqlRelUtils.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/BeamSqlRelUtils.java @@ -24,6 +24,7 @@ import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.extensions.sql.impl.planner.BeamRelMetadataQuery; import org.apache.beam.sdk.extensions.sql.impl.planner.NodeStats; +import org.apache.beam.sdk.options.ExperimentalOptions; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.transforms.PTransform; import org.apache.beam.sdk.util.Preconditions; @@ -46,14 +47,14 @@ public class BeamSqlRelUtils { public static final String ERROR = "error"; public static PCollection toPCollection(Pipeline pipeline, BeamRelNode node) { - return toPCollection(pipeline, node, null, new HashMap()); + return toPCollection(pipeline, node, null, new HashMap(), new HashMap<>()); } public static PCollection toPCollection( Pipeline pipeline, BeamRelNode node, @Nullable PTransform, ? extends POutput> errorTransformer) { - return toPCollection(pipeline, node, errorTransformer, new HashMap()); + return toPCollection(pipeline, node, errorTransformer, new HashMap(), new HashMap<>()); } /** Transforms the inputs into a PInput. */ @@ -61,7 +62,8 @@ private static PCollectionList buildPCollectionList( List inputRels, Pipeline pipeline, @Nullable PTransform, ? extends POutput> errorTransformer, - Map> cache) { + Map> cache, + Map usedNames) { if (inputRels.isEmpty()) { return PCollectionList.empty(pipeline); } else { @@ -79,7 +81,7 @@ private static PCollectionList buildPCollectionList( beamRel = (BeamRelNode) input; } return BeamSqlRelUtils.toPCollection( - pipeline, beamRel, errorTransformer, cache); + pipeline, beamRel, errorTransformer, cache, usedNames); }) .collect(Collectors.toList())); } @@ -93,15 +95,17 @@ static PCollection toPCollection( Pipeline pipeline, BeamRelNode node, @Nullable PTransform, ? extends POutput> errorTransformer, - Map> cache) { + Map> cache, + Map usedNames) { PCollection output = cache.get(node.getId()); if (output != null) { return output; } - String name = node.getClass().getSimpleName() + "_" + node.getId(); + String name = uniqueName(usedNames, transformName(pipeline, node)); PCollectionList input = - buildPCollectionList(node.getPCollectionInputs(), pipeline, errorTransformer, cache); + buildPCollectionList( + node.getPCollectionInputs(), pipeline, errorTransformer, cache, usedNames); PTransform, PCollection> transform = node.buildPTransform(errorTransformer); output = Pipeline.applyTransform(name, input, transform); @@ -110,6 +114,35 @@ static PCollection toPCollection( return output; } + /** + * Names the composite that {@code node} expands into after the stage it was composed from, or + * after the node's own type when nothing composed it. + */ + private static String transformName(Pipeline pipeline, BeamRelNode node) { + if (ExperimentalOptions.hasExperiment(pipeline.getOptions(), StageName.LEGACY_EXPERIMENT)) { + return node.getClass().getSimpleName() + "_" + node.getId(); + } + String label = StageName.renderedName(node); + return label == null || label.isEmpty() ? node.getClass().getSimpleName() : label; + } + + /** + * Disambiguates repeated names in DFS order, within the plan being expanded. + * + *

Names no longer carry a rel id, so a plan can legitimately contain two stages with the same + * provenance. {@link Pipeline} would uniquify them itself, but then reports the pipeline as not + * having stable unique names, which is fatal under {@code --stableUniqueNames=ERROR}. + * + *

Counting per plan rather than per pipeline is what makes the numbering stable: every caller + * expands a plan either into a fresh pipeline or inside {@code SqlTransform}'s own composite, so + * names only have to be unique among the stages of one query. A counter shared across a pipeline + * would let an unrelated query added elsewhere renumber stages that did not themselves change. + */ + private static String uniqueName(Map usedNames, String name) { + int occurrence = usedNames.merge(name, 1, Integer::sum); + return occurrence == 1 ? name : name + " #" + occurrence; + } + public static BeamRelNode getBeamRelInput(RelNode input) { if (input instanceof RelSubset) { // go with known best input diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/StageName.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/StageName.java new file mode 100644 index 000000000000..0039ae7debe1 --- /dev/null +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rel/StageName.java @@ -0,0 +1,344 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sql.impl.rel; + +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptCluster; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptTable; +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.core.TableScan; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.hint.HintStrategyTable; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.hint.Hintable; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.hint.RelHint; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * Provenance label carried on a {@link RelNode} as a {@link RelHint}, used to name the composite + * {@link org.apache.beam.sdk.transforms.PTransform} the node expands into. + * + *

The label records which nodes were fused to produce this one rather than what the node + * computes. Planner rules collapse a whole chain of relational operations into a single {@code + * Calc}, so a label like {@code Filter;Project} tells the reader which of them ended up in the + * transform. + * + *

A {@link RelHint} is the carrier because hints are a field on every core Calcite rel type and + * survive both {@code copy()} and physical conversion. + * + *

Whether the planner can see a label

+ * + *

Only some rel types put their hints in the digest, and the two halves fail in opposite + * directions. + * + *

{@code Project}, {@code Filter} and {@code Join} override {@code deepEquals0}/{@code + * deepHashCode0} and compare the hints field, so for them two rels differing only in their label + * are two rels as far as the planner is concerned. A label must therefore be a canonical function + * of the set of nodes fused into it, never of the order they were matched in: {@code + * ProjectFilterTransposeRule} and {@code FilterProjectTransposeRule} swap a pair back and forth, + * and an order-sensitive label would hand the planner an endless supply of "new" rels to explore. + * That is what the position prefix is for -- see {@link #compose}. + * + *

{@code Calc}, {@code Aggregate}, {@code SetOp}, {@code Correlate} and {@code TableScan} do not + * override those methods, and {@code Calc.explainTerms} omits hints, so they fall back to a digest + * that ignores the label entirely. Relabelling one of those and calling {@code transformTo} finds a + * digest-equal rel already in the subset and discards the newly labelled object, leaving whichever + * label was registered first. Because {@link #compose} is a canonical union the two paths almost + * always agree, but they need not: a one-node rule such as {@code PROJECT_TO_CALC} labels its + * output {@code Project} where a two-node rule producing the same program would label it {@code + * Filter;Project}. This is the residual source of name variation, and it is why {@code + * StageNameTest} asserts that repeated planning of the same query produces identical names. + * + *

Labels over-claim

+ * + *

A label is a union that never splits. {@code StageNameRuleCall} stamps the composed label on + * every node a rule built, so {@code FilterProjectTransposeRule}, which rebuilds both halves of the + * pair it matched, leaves both of them labelled {@code Filter;Project}. If those halves end up in + * separate stages rather than being merged, both stages claim both operations and only the {@code + * #2} suffix tells them apart. That is the honest reason occurrences have to be numbered, and it is + * the price of the confluence that makes labelling terminate at all: a label that shrank when nodes + * separated would depend on the order rules fired, which is exactly what the previous paragraph + * rules out. + */ +@Internal +public class StageName { + + /** Hint name under which the label is stored. */ + private static final String HINT = "BEAM_STAGE"; + + /** + * Experiment restoring the pre-provenance {@code BeamCalcRel_57} style names. + * + *

Dataflow matches streaming pipelines for update by step name, so a pipeline already running + * cannot adopt the new names without being drained. + */ + public static final String LEGACY_EXPERIMENT = "legacy-sql-transform-names"; + + private static final String SEPARATOR = ";"; + private static final String POSITION_SEPARATOR = ":"; + private static final String ELLIPSIS = "..."; + + /** + * Longest rendered name. Beyond this the middle of the chain is elided. Names appear in runner + * UIs and in per-step metric keys, so they cannot grow with plan size. + */ + private static final int MAX_LENGTH = 100; + + private StageName() {} + + /** The only hint strategy table Beam installs; its identity marks a cluster as registered. */ + private static final HintStrategyTable STRATEGIES = + HintStrategyTable.builder().hintStrategy(HINT, (hint, rel) -> true).build(); + + /** + * Declares {@link #HINT} to {@code cluster}, which the planner requires before it will tolerate a + * node carrying it. + * + *

An unregistered hint trips an assertion the first time the planner tests a matched node for + * rule exclusion. The predicate admits every node type, since the label is about provenance + * rather than anything the node can act on. + * + *

{@link RelOptCluster#setHintStrategies} replaces the table rather than merging into it, and + * a cluster can reach here more than once. Registering again is a no-op, but finding strategies + * that Beam did not install is an error rather than something to silently discard: if SQL hint + * syntax is ever enabled, this has to become a merge, and it is better to find that out from a + * stack trace than from hints that quietly stop working. + */ + public static void register(RelOptCluster cluster) { + HintStrategyTable existing = cluster.getHintStrategies(); + if (existing == STRATEGIES) { + return; + } + checkState( + existing == HintStrategyTable.EMPTY, + "Cluster already has hint strategies configured; %s would discard them", + HINT); + cluster.setHintStrategies(STRATEGIES); + } + + /** Hints carrying {@code label}, for passing to a rel constructor or factory. */ + private static ImmutableList hintsFor(String label) { + return ImmutableList.of(RelHint.builder(HINT).hintOption(label).build()); + } + + /** The label on {@code node} in its stored form, or null if it carries none. */ + public static @Nullable String storedLabel(RelNode node) { + if (!(node instanceof Hintable)) { + return null; + } + for (RelHint hint : ((Hintable) node).getHints()) { + if (HINT.equals(hint.hintName) && !hint.listOptions.isEmpty()) { + return hint.listOptions.get(0); + } + } + return null; + } + + /** The name to show for {@code node}, or null if it carries no label. */ + public static @Nullable String renderedName(RelNode node) { + String label = storedLabel(node); + return label == null ? null : render(label); + } + + /** + * Returns {@code node} labelled as having come from a source operation called {@code name}, or + * {@code node} unchanged if it is already labelled or cannot be labelled. + * + *

Rels that do not override {@link Hintable#withHints} return themselves, so an un-plumbed rel + * type silently keeps the default naming instead of failing. + */ + private static RelNode stamp(RelNode node, String name) { + if (storedLabel(node) != null) { + return node; + } + // The node's own id orders it against the rest of the plan. Ids are handed out in construction + // order and backfill builds bottom-up, so ordering by id is data-flow order. Only the relative + // order matters, which is why an id being different on the next run is harmless. + // + // A node built later than its neighbours -- anything a pre-pass such as decorrelation left + // behind -- gets a larger id and so renders after nodes it actually feeds. The name is still + // complete; only its internal ordering is off. + return relabel(node, node.getId() + POSITION_SEPARATOR + sanitize(name)); + } + + /** + * Returns {@code node} carrying {@code label} in place of whatever label it had. + * + *

A rule that fuses nodes usually builds its result by copying one of them, which carries that + * one's label along. The result belongs to all of them, so the composed label has to win. + */ + public static RelNode relabel(RelNode node, String label) { + if (!(node instanceof Hintable)) { + return node; + } + // withHints() builds a whole new rel, and for a Calc that means re-validating its program, + // which renders the program to a string. Rules re-fire on nodes they have already labelled + // often enough -- transpose rules trade a pair back and forth, CalcMergeRule revisits merged + // inputs -- that skipping the no-op case is most of the labelling cost on a large plan. + if (label.equals(storedLabel(node))) { + return node; + } + List hints = new ArrayList<>(); + for (RelHint hint : ((Hintable) node).getHints()) { + if (!HINT.equals(hint.hintName)) { + hints.add(hint); + } + } + hints.addAll(hintsFor(label)); + return ((Hintable) node).withHints(hints); + } + + /** + * Labels every node in the tree rooted at {@code rel} that carries no label yet, after the kind + * of node it is. + * + *

A SQL query has no user-visible operator names for its stages to inherit, and neither does + * whatever the decorrelation pre-pass leaves behind, so a node is named for its own shape. + */ + public static RelNode backfill(RelNode rel) { + List inputs = new ArrayList<>(); + boolean rebuilt = false; + for (RelNode input : rel.getInputs()) { + RelNode labelled = backfill(input); + rebuilt |= labelled != input; + inputs.add(labelled); + } + RelNode result = rebuilt ? rel.copy(rel.getTraitSet(), inputs) : rel; + return stamp(result, structuralName(result)); + } + + private static String structuralName(RelNode rel) { + RelOptTable table = rel.getTable(); + if (rel instanceof TableScan && table != null) { + List qualifiedName = table.getQualifiedName(); + return "Scan(" + qualifiedName.get(qualifiedName.size() - 1) + ")"; + } + String name = rel.getClass().getSimpleName(); + return name.startsWith("Logical") ? name.substring("Logical".length()) : name; + } + + /** + * Merges {@code labels} into the label for a node fused from all of them. + * + *

The result depends only on the set of source operations, not on the order the rule + * happened to match them in: each is deduplicated by name and they are re-sorted by the position + * they held in the original plan. So the label of a fused node is stable no matter which sequence + * of rules assembled it, which is what keeps the planner from treating a re-derivation of the + * same node as a new one. + */ + public static String compose(List labels) { + Map positions = new LinkedHashMap<>(); + for (String label : labels) { + if (label == null) { + continue; + } + for (String part : label.split(SEPARATOR, -1)) { + if (!part.isEmpty()) { + positions.merge(nameOf(part), positionOf(part), Math::min); + } + } + } + return positions.entrySet().stream() + .sorted( + (a, b) -> { + int byPosition = Integer.compare(a.getValue(), b.getValue()); + return byPosition != 0 ? byPosition : a.getKey().compareTo(b.getKey()); + }) + .map(entry -> entry.getValue() + POSITION_SEPARATOR + entry.getKey()) + .reduce((a, b) -> a + SEPARATOR + b) + .orElse(""); + } + + /** The displayable name for a stored label: its source operations, in plan order. */ + static String render(String label) { + List names = new ArrayList<>(); + for (String part : label.split(SEPARATOR, -1)) { + if (!part.isEmpty()) { + names.add(nameOf(part)); + } + } + return truncate(names); + } + + private static String nameOf(String part) { + int separator = part.indexOf(POSITION_SEPARATOR); + return separator < 0 ? part : part.substring(separator + 1); + } + + /** + * The position prefix of a stored part. Unprefixed parts sort last; {@link #sanitize} makes them + * unreachable from {@link #stamp}, but {@link #compose} is public and its input is a string. + */ + private static int positionOf(String part) { + int separator = part.indexOf(POSITION_SEPARATOR); + if (separator < 0) { + return Integer.MAX_VALUE; + } + try { + return Integer.parseInt(part.substring(0, separator)); + } catch (NumberFormatException e) { + return Integer.MAX_VALUE; + } + } + + /** Keeps a name free of the characters that delimit the stored form. */ + static String sanitize(String name) { + return name.replace(SEPARATOR, "_").replace(POSITION_SEPARATOR, "_"); + } + + /** + * Joins {@code names}, eliding the middle if the result would exceed {@link #MAX_LENGTH}. The + * first and last are the recognizable ones, and the count of what was dropped signals how much + * fusion happened. + */ + private static String truncate(List names) { + String full = String.join(SEPARATOR, names); + if (full.length() <= MAX_LENGTH) { + return full; + } + if (names.size() < 3) { + return clip(full); + } + String elided = + names.get(0) + + SEPARATOR + + "+" + + (names.size() - 2) + + " more" + + SEPARATOR + + names.get(names.size() - 1); + return elided.length() <= MAX_LENGTH ? elided : clip(elided); + } + + /** + * Cuts {@code name} down to {@link #MAX_LENGTH}. These names reach runner UIs and per-step metric + * keys, so the marker is ASCII and the cut never lands inside a surrogate pair. + */ + private static String clip(String name) { + int end = MAX_LENGTH - ELLIPSIS.length(); + if (Character.isHighSurrogate(name.charAt(end - 1))) { + end--; + } + return name.substring(0, end) + ELLIPSIS; + } +} diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCalcRule.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCalcRule.java index 1319e20fdb4c..d6af320786cc 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCalcRule.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/BeamCalcRule.java @@ -69,6 +69,7 @@ public RelNode convert(RelNode rel) { return new BeamCalcRel( calc.getCluster(), calc.getTraitSet().replace(BeamLogicalConvention.INSTANCE), + calc.getHints(), RelOptRule.convert(input, input.getTraitSet().replace(BeamLogicalConvention.INSTANCE)), calc.getProgram()); } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRule.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRule.java new file mode 100644 index 000000000000..e1814626b385 --- /dev/null +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRule.java @@ -0,0 +1,196 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sql.impl.rule; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.function.Predicate; +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.extensions.sql.impl.rel.StageName; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.Convention; +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.RelOptRule; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRuleCall; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRuleOperand; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRuleOperandChildren; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelTrait; +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.convert.ConverterRule; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.rules.SubstitutionRule; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.rules.TransformationRule; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.tools.RuleSet; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.tools.RuleSets; + +/** + * Wraps a planner rule so that the nodes it produces inherit a {@link StageName} composed from the + * nodes it matched. + * + *

None of the rules that fuse relational nodes propagate hints, and they are Calcite's, not + * Beam's, to change. Beam does own the rule set, so it can hand each rule a {@link + * StageNameRuleCall} that labels whatever the rule produces. + * + *

The planner sees the wrapper: a rule looked up by identity in {@code + * RelOptPlanner#getRules()}, or excluded by instance, has to be looked up as the wrapper rather + * than as the original. Rule bodies are unaffected, because {@code + * StageNameRuleCall#getRule} hands back the delegate. + */ +@Internal +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class StageNameRule extends RelOptRule { + + private final RelOptRule delegate; + + /** The rule this one stands in for. */ + protected RelOptRule delegate() { + return delegate; + } + + /** + * Returns {@code rule} wrapped so its output is labelled, or {@code rule} itself where wrapping + * is not applicable. + * + *

The planner branches on {@link ConverterRule}, {@link SubstitutionRule} and {@link + * TransformationRule}, so the wrapper has to present the same markers. Converter rules cannot be + * wrapped at all -- the planner reads their in/out traits to register trait conversions -- so + * they carry hints through explicitly instead. + * + *

Wrapping happens here, at the boundary where rules are handed to a planner, rather than in + * {@code BeamRuleSets}. A wrapper is not an instance of the rule class it wraps, and callers + * select rules out of the published rule set by type. + */ + public static RelOptRule wrap(RelOptRule rule) { + if (rule instanceof ConverterRule) { + return rule; + } + if (rule instanceof SubstitutionRule) { + return new Substitution(rule); + } + if (rule instanceof TransformationRule) { + return new Transformation(rule); + } + return new StageNameRule(rule); + } + + /** {@code ruleSets} with every rule wrapped, for handing to a planner. */ + public static Collection wrapAll(Collection ruleSets) { + List wrapped = new ArrayList<>(); + for (RuleSet ruleSet : ruleSets) { + List rules = new ArrayList<>(); + for (RelOptRule rule : ruleSet) { + rules.add(wrap(rule)); + } + wrapped.add(RuleSets.ofList(rules)); + } + return wrapped; + } + + /** + * Adds {@code rule} to {@code planner} wrapped, dropping the unwrapped original. + * + *

A wrapper keeps its delegate's description and the planner requires descriptions to be + * unique. Beam's rule set overlaps Calcite's defaults, which are registered before Beam gets a + * chance, so a wrapper has to displace its original rather than sit alongside it -- otherwise + * both fire and the unlabelled result may win. + */ + public static void addTo(RelOptPlanner planner, RelOptRule rule) { + RelOptRule wrapped = wrap(rule); + if (wrapped instanceof StageNameRule) { + planner.removeRule(rule); + } + planner.addRule(wrapped); + } + + private StageNameRule(RelOptRule delegate) { + super( + CopiedOperand.copyOf(delegate.getOperand()), + delegate.relBuilderFactory, + delegate.toString()); + this.delegate = delegate; + } + + @Override + public boolean matches(RelOptRuleCall call) { + return delegate.matches(call); + } + + @Override + public void onMatch(RelOptRuleCall call) { + delegate.onMatch(new StageNameRuleCall(call, delegate)); + } + + @Override + public Convention getOutConvention() { + return delegate.getOutConvention(); + } + + @Override + public RelTrait getOutTrait() { + return delegate.getOutTrait(); + } + + private static class Transformation extends StageNameRule implements TransformationRule { + Transformation(RelOptRule delegate) { + super(delegate); + } + } + + private static class Substitution extends StageNameRule implements SubstitutionRule { + Substitution(RelOptRule delegate) { + super(delegate); + } + + @Override + public boolean autoPruneOld() { + return ((SubstitutionRule) delegate()).autoPruneOld(); + } + } + + /** + * A structural clone of another rule's operand tree. + * + *

Constructing a {@link RelOptRule} re-points its operands back at itself, so the wrapper must + * not share the delegate's operands: the delegate singletons are also registered directly with + * Calcite's own internal planners. + */ + private static class CopiedOperand extends RelOptRuleOperand { + private CopiedOperand( + Class clazz, + RelTrait trait, + Predicate predicate, + RelOptRuleOperandChildren children) { + super(clazz, trait, predicate, children); + } + + static RelOptRuleOperand copyOf(RelOptRuleOperand source) { + List children = new ArrayList<>(); + for (RelOptRuleOperand child : source.getChildOperands()) { + children.add(copyOf(child)); + } + return new CopiedOperand( + source.getMatchedClass(), + source.trait, + // Class and trait are re-checked by the copy's own matches() before this runs, so the + // delegate's predicate is only consulted for the part it alone knows. + source::matches, + new RelOptRuleOperandChildren(source.childPolicy, children)); + } + } +} diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRuleCall.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRuleCall.java new file mode 100644 index 000000000000..568110584a31 --- /dev/null +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRuleCall.java @@ -0,0 +1,212 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sql.impl.rule; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +import org.apache.beam.sdk.extensions.sql.impl.rel.StageName; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelHintsPropagator; +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.RelOptRule; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRuleCall; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRuleOperand; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.hep.HepRelVertex; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.volcano.RelSubset; +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.metadata.RelMetadataQuery; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.tools.RelBuilder; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * A {@link RelOptRuleCall} that labels the nodes a rule produces with the composed {@link + * StageName} of the nodes the rule matched, so that a fused node records what was fused into it. + * + *

Everything other than {@code transformTo} is delegated to the original call, except {@code + * getRule}, which hands back the rule being run rather than the {@link StageNameRule} wrapping it. + */ +@SuppressWarnings({ + "nullness" // TODO(https://github.com/apache/beam/issues/20497) +}) +public class StageNameRuleCall extends RelOptRuleCall { + private static final Logger LOG = LoggerFactory.getLogger(StageNameRuleCall.class); + + /** Keeps a rel type that cannot be labelled from filling the log, one line per JVM. */ + private static final AtomicBoolean LOGGED_FAILURE = new AtomicBoolean(); + + private final RelOptRuleCall originalCall; + private final RelOptRule delegate; + + /** + * Composed label of the matched nodes, computed on the first {@code transformTo}. + * + *

Volcano constructs one of these per {@code onMatch}, and most matches return without + * transforming anything. Composing eagerly would run a split, a map and a sort on every one of + * them. + */ + private @Nullable String label; + + StageNameRuleCall(RelOptRuleCall originalCall, RelOptRule delegate) { + super( + originalCall.getPlanner(), + originalCall.getOperand0(), + originalCall.rels, + ImmutableMap.of(), + null); + this.originalCall = originalCall; + this.delegate = delegate; + } + + private String label() { + if (label == null) { + List labels = new ArrayList<>(); + for (RelNode rel : originalCall.getRelList()) { + labels.add(StageName.storedLabel(rel)); + } + label = StageName.compose(labels); + } + return label; + } + + @Override + public void transformTo(RelNode rel, Map equiv) { + originalCall.transformTo(labelled(rel, equiv), equiv); + } + + @Override + public void transformTo( + RelNode rel, Map equiv, RelHintsPropagator hintsPropagator) { + originalCall.transformTo(labelled(rel, equiv), equiv, hintsPropagator); + } + + private RelNode labelled(RelNode rel, Map equiv) { + String composed = label(); + if (composed.isEmpty()) { + return rel; + } + try { + // An equivalence map keys off node instances inside `rel`. Rebuilding the subtree would + // invalidate those keys, so in that case only the root is labelled -- rewriting the root + // leaves its inputs untouched. + return equiv.isEmpty() ? labelSubtree(rel, composed) : StageName.relabel(rel, composed); + } catch (RuntimeException e) { + // Naming is cosmetic, so a rel type whose copy() or withHints() rejects this is not worth + // failing a query over. It is worth hearing about once: the fallback name is silent + // otherwise, and nobody reads debug in production. + if (LOGGED_FAILURE.compareAndSet(false, true)) { + LOG.warn( + "Could not label rule output with stage name '{}'; {} will fall back to default" + + " transform naming. Further occurrences are logged at debug.", + composed, + rel.getClass().getSimpleName(), + e); + } else { + LOG.debug("Could not label rule output with stage name '{}'", composed, e); + } + return rel; + } + } + + /** + * Labels {@code rel} and every descendant the rule built underneath it. + * + *

Several rules assemble a chain of nodes and hand only its root to {@code transformTo}, so + * labelling the root alone would leave the intermediates anonymous. {@link RelSubset} and {@link + * HepRelVertex} mark the boundary between what the rule built and the inputs it was given. + * + *

That boundary is not exact under Volcano, which hands a multi-operand rule concrete rels + * rather than subsets. A rule that splices a node it matched straight into its output has that + * node copied and relabelled here. The label it receives is the composed one, which already + * covers it, so the name stays correct -- but it is an extra rel in the memo. + */ + private RelNode labelSubtree(RelNode rel, String label) { + if (rel instanceof RelSubset || rel instanceof HepRelVertex) { + return rel; + } + List inputs = new ArrayList<>(); + boolean rebuilt = false; + for (RelNode input : rel.getInputs()) { + RelNode labelled = labelSubtree(input, label); + rebuilt |= labelled != input; + inputs.add(labelled); + } + RelNode result = rebuilt ? rel.copy(rel.getTraitSet(), inputs) : rel; + return StageName.relabel(result, label); + } + + // Methods that are delegated to originalCall. + + @Override + public RelOptRuleOperand getOperand0() { + return originalCall.getOperand0(); + } + + /** + * The rule whose {@code onMatch} is running, not the {@link StageNameRule} the planner dispatched + * through. Rule bodies cast this to their own type. + */ + @Override + public RelOptRule getRule() { + return delegate; + } + + @Override + public List getRelList() { + return originalCall.getRelList(); + } + + @Override + @SuppressWarnings("TypeParameterUnusedInFormals") + public T rel(int ordinal) { + return originalCall.rel(ordinal); + } + + @Override + public List getChildRels(RelNode rel) { + return originalCall.getChildRels(rel); + } + + @Override + public RelOptPlanner getPlanner() { + return originalCall.getPlanner(); + } + + @Override + public RelMetadataQuery getMetadataQuery() { + return originalCall.getMetadataQuery(); + } + + @Override + public List getParents() { + return originalCall.getParents(); + } + + @Override + public boolean isRuleExcluded() { + return originalCall.isRuleExcluded(); + } + + @Override + public RelBuilder builder() { + return originalCall.builder(); + } +} diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/StageNameTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/StageNameTest.java new file mode 100644 index 000000000000..59aa497c4d99 --- /dev/null +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rel/StageNameTest.java @@ -0,0 +1,248 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sql.impl.rel; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.empty; +import static org.hamcrest.Matchers.hasItem; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.not; + +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.beam.sdk.Pipeline; +import org.apache.beam.sdk.extensions.sql.meta.provider.test.TestBoundedTable; +import org.apache.beam.sdk.options.ExperimentalOptions; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.runners.TransformHierarchy; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.testing.CrashingRunner; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.values.PBegin; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.Row; +import org.junit.BeforeClass; +import org.junit.Test; + +/** Tests for provenance-based naming of the composites a Beam SQL plan expands into. */ +public class StageNameTest extends BaseRelTest { + + @BeforeClass + public static void prepare() { + registerTable( + "ORDERS", + TestBoundedTable.of( + Schema.FieldType.INT64, "order_id", + Schema.FieldType.INT32, "site_id", + Schema.FieldType.DECIMAL, "price") + .addRows(1L, 1, new BigDecimal(1.0), 2L, 2, new BigDecimal(2.0))); + } + + @Test + public void composeDeduplicatesAndRestoresPlanOrder() { + assertThat( + StageName.compose(Arrays.asList("7:Project", "3:Filter", "7:Project")), + is("3:Filter;7:Project")); + } + + @Test + public void composeSkipsMissingLabels() { + assertThat(StageName.compose(Arrays.asList(null, "", "3:Filter", null)), is("3:Filter")); + } + + /** + * A leaf is named after the table it scans, which could contain the characters the stored form + * uses to delimit parts and so make one label parse as several. + */ + @Test + public void separatorsInASourceNameAreNeutralized() { + assertThat(StageName.sanitize("a;b:c"), is("a_b_c")); + assertThat( + StageName.render(StageName.compose(Arrays.asList("3:" + StageName.sanitize("a;b:c")))), + is("a_b_c")); + } + + /** + * The property that keeps the planner terminating: transpose rules match the same pair of nodes + * in either order, and a label that changed with the order would be an endless supply of rels the + * planner has not seen before. + */ + @Test + public void composeIsIndependentOfMatchOrder() { + String forwards = StageName.compose(Arrays.asList("3:Filter", "7:Project")); + String backwards = StageName.compose(Arrays.asList("7:Project", "3:Filter")); + assertThat(backwards, is(forwards)); + assertThat(StageName.compose(Arrays.asList(forwards, backwards)), is(forwards)); + } + + @Test + public void aLongChainIsElidedInTheMiddle() { + List labels = new ArrayList<>(); + labels.add("0:first"); + for (int i = 0; i < 40; i++) { + labels.add((i + 1) + ":Project" + i); + } + labels.add("41:last"); + assertThat(StageName.render(StageName.compose(labels)), is("first;+40 more;last")); + } + + @Test + public void fusedCalcIsNamedAfterEverythingFusedIntoIt() { + assertThat( + topLevelNames("SELECT order_id + 1 FROM ORDERS WHERE site_id = 1"), + hasItem("Filter;Project")); + } + + @Test + public void repeatedNamesAreDisambiguatedInPlanOrder() { + List names = + topLevelNames( + "SELECT order_id FROM ORDERS WHERE site_id = 1 " + + "UNION ALL SELECT order_id FROM ORDERS WHERE site_id = 2"); + assertThat(names.stream().filter(n -> n.startsWith("Filter;Project")).count(), is(2L)); + assertThat(names, hasItem("Filter;Project #2")); + } + + @Test + public void legacyExperimentRestoresClassNameAndId() { + List names = + topLevelNames( + "SELECT order_id + 1 FROM ORDERS WHERE site_id = 1", StageName.LEGACY_EXPERIMENT); + assertThat(names, not(empty())); + assertThat(names.stream().allMatch(n -> n.matches("Beam\\w+Rel_\\d+")), is(true)); + } + + /** + * The one property the whole scheme is for. Labels are only invisible to the planner's digest for + * some rel types -- {@code Calc} among them -- so which of two equally valid labels survives can + * come down to the order rels were registered in. If that ever varies between runs, so do the + * names, and Dataflow streaming update breaks on a pipeline nobody edited. + */ + @Test + public void namesAreIdenticalAcrossRuns() { + String sql = + "SELECT order_id + 1, price * 2 FROM ORDERS WHERE site_id = 1 AND order_id > 0 " + + "UNION ALL SELECT order_id, price FROM ORDERS WHERE site_id = 2"; + List first = topLevelNames(sql); + for (int run = 0; run < 4; run++) { + assertThat(topLevelNames(sql), is(first)); + } + } + + /** + * {@code PROJECT_FILTER_TRANSPOSE} and {@code FILTER_PROJECT_TRANSPOSE} are both in Beam's rule + * set, so they swap a Filter/Project pair back and forth indefinitely. A label that depended on + * the order the pair was matched in would make each swap produce a rel the planner had not seen, + * and planning would never converge. Several stacked pairs give the planner room to do it. + */ + @Test(timeout = 120_000) + public void planningTerminatesWhenTransposeRulesCycle() { + List names = + topLevelNames( + "SELECT a + 1 AS a, b FROM (" + + " SELECT a, b FROM (" + + " SELECT order_id AS a, price AS b FROM ORDERS WHERE site_id = 1" + + " ) WHERE a > 0" + + ") WHERE b > 0"); + assertThat(names, not(empty())); + } + + /** + * Occurrence numbering restarts with each query. Every caller expands a plan either into a fresh + * pipeline or inside {@code SqlTransform}'s own composite, so names only have to be unique among + * the stages of one query -- and counting per query means a second query cannot renumber the + * stages of the first. + */ + @Test + public void numberingRestartsForEachQuery() { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(CrashingRunner.class); + Pipeline pipeline = Pipeline.create(options); + String sql = + "SELECT order_id FROM ORDERS WHERE site_id = 1 " + + "UNION ALL SELECT order_id FROM ORDERS WHERE site_id = 2"; + pipeline.apply("first", new Query(sql)); + pipeline.apply("second", new Query(sql)); + + List first = namesUnder(pipeline, "first"); + assertThat(first, hasItem("Filter;Project #2")); + assertThat(namesUnder(pipeline, "second"), is(first)); + } + + /** Applies a query the way {@code SqlTransform} does, inside a composite of its own. */ + private static class Query extends PTransform> { + private final String sql; + + Query(String sql) { + this.sql = sql; + } + + @Override + public PCollection expand(PBegin input) { + return BeamSqlRelUtils.toPCollection(input.getPipeline(), env.parseQuery(sql)); + } + } + + private static List namesUnder(Pipeline pipeline, String prefix) { + List names = new ArrayList<>(); + pipeline.traverseTopologically( + new Pipeline.PipelineVisitor.Defaults() { + @Override + public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) { + String full = node.getFullName(); + if (full.startsWith(prefix + "/")) { + String rest = full.substring(prefix.length() + 1); + if (!rest.contains("/")) { + names.add(rest); + } + } + return CompositeBehavior.ENTER_TRANSFORM; + } + }); + return names; + } + + private static List topLevelNames(String sql, String... experiments) { + PipelineOptions options = PipelineOptionsFactory.create(); + options.setRunner(CrashingRunner.class); + for (String experiment : experiments) { + ExperimentalOptions.addExperiment(options.as(ExperimentalOptions.class), experiment); + } + Pipeline pipeline = Pipeline.create(options); + compilePipeline(sql, pipeline); + return topLevelNames(pipeline); + } + + private static List topLevelNames(Pipeline pipeline) { + List names = new ArrayList<>(); + pipeline.traverseTopologically( + new Pipeline.PipelineVisitor.Defaults() { + @Override + public CompositeBehavior enterCompositeTransform(TransformHierarchy.Node node) { + if (!node.isRootNode() && !node.getFullName().contains("/")) { + names.add(node.getFullName()); + } + return CompositeBehavior.ENTER_TRANSFORM; + } + }); + return names; + } +} diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRuleTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRuleTest.java new file mode 100644 index 000000000000..7a85189d647e --- /dev/null +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/rule/StageNameRuleTest.java @@ -0,0 +1,123 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.beam.sdk.extensions.sql.impl.rule; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.hasItem; +import static org.hamcrest.Matchers.instanceOf; +import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.sameInstance; + +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRule; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.RelOptRuleOperand; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.plan.volcano.VolcanoPlanner; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.convert.ConverterRule; +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.rules.SubstitutionRule; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.rel.rules.TransformationRule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for the wrapper that gives a planner rule's output a provenance label. */ +@RunWith(JUnit4.class) +public class StageNameRuleTest { + + /** + * {@link RelOptRule}'s constructor walks its operand tree and re-points every operand at the rule + * being constructed. Handing it the delegate's own operands would therefore reassign them away + * from the delegate -- which is a singleton also registered with Calcite's internal planners. + */ + @Test + public void wrappingLeavesTheDelegateOperandsPointingAtTheDelegate() { + RelOptRule delegate = CoreRules.FILTER_INTO_JOIN; + RelOptRuleOperand delegateOperand = delegate.getOperand(); + + RelOptRule wrapper = StageNameRule.wrap(delegate); + + assertThat(delegateOperand.getRule(), is(sameInstance(delegate))); + assertThat(delegate.getOperand(), is(sameInstance(delegateOperand))); + assertThat(wrapper.getOperand(), is(not(sameInstance(delegateOperand)))); + assertThat(wrapper.getOperand().getRule(), is(sameInstance(wrapper))); + } + + @Test + public void copiedOperandTreeHasTheSameShape() { + RelOptRule wrapper = StageNameRule.wrap(CoreRules.FILTER_INTO_JOIN); + assertSameShape(CoreRules.FILTER_INTO_JOIN.getOperand(), wrapper.getOperand()); + } + + private static void assertSameShape(RelOptRuleOperand source, RelOptRuleOperand copy) { + assertThat(copy.getMatchedClass(), is(source.getMatchedClass())); + assertThat(copy.childPolicy, is(source.childPolicy)); + assertThat(copy.getChildOperands().size(), is(source.getChildOperands().size())); + for (int i = 0; i < source.getChildOperands().size(); i++) { + assertSameShape(source.getChildOperands().get(i), copy.getChildOperands().get(i)); + } + } + + /** + * The planner reads a converter rule's in and out traits to register trait conversions, which a + * wrapper cannot stand in for. Those rules propagate hints themselves instead. + */ + @Test + public void converterRulesAreLeftAlone() { + RelOptRule converter = BeamCalcRule.INSTANCE; + assertThat(converter, is(instanceOf(ConverterRule.class))); + assertThat(StageNameRule.wrap(converter), is(sameInstance(converter))); + } + + /** The planner branches on these markers, so a wrapper has to present the same ones. */ + @Test + public void ruleMarkersSurviveWrapping() { + assertThat(CoreRules.FILTER_INTO_JOIN, is(instanceOf(TransformationRule.class))); + assertThat( + StageNameRule.wrap(CoreRules.FILTER_INTO_JOIN), is(instanceOf(TransformationRule.class))); + + assertThat(CoreRules.PROJECT_REMOVE, is(instanceOf(SubstitutionRule.class))); + assertThat( + StageNameRule.wrap(CoreRules.PROJECT_REMOVE), is(instanceOf(SubstitutionRule.class))); + } + + /** + * A wrapper keeps its delegate's description, and the planner rejects two rules with the same + * description. Beam's rule set overlaps Calcite's defaults, which {@code JdbcDriver} finds + * already registered, so the wrapper has to displace the original rather than sit beside it. + */ + @Test + public void addToDisplacesAnAlreadyRegisteredRule() { + VolcanoPlanner planner = new VolcanoPlanner(); + RelOptRule rule = CoreRules.FILTER_INTO_JOIN; + planner.addRule(rule); + + StageNameRule.addTo(planner, rule); + + assertThat(planner.getRules(), not(hasItem(sameInstance(rule)))); + assertThat( + planner.getRules().stream() + .anyMatch(r -> r instanceof StageNameRule && r.toString().equals(rule.toString())), + is(true)); + } + + @Test + public void wrapperKeepsTheDelegateDescription() { + RelOptRule rule = CoreRules.FILTER_INTO_JOIN; + assertThat(StageNameRule.wrap(rule).toString(), is(rule.toString())); + } +} diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/parquet/ParquetTableProviderTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/parquet/ParquetTableProviderTest.java index 63197be0a45e..a8a5ee9096ee 100644 --- a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/parquet/ParquetTableProviderTest.java +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/meta/provider/parquet/ParquetTableProviderTest.java @@ -26,11 +26,14 @@ import org.apache.beam.sdk.PipelineResult; import org.apache.beam.sdk.PipelineResult.State; import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv; +import org.apache.beam.sdk.extensions.sql.impl.rel.BeamRelNode; import org.apache.beam.sdk.extensions.sql.impl.rel.BeamSqlRelUtils; import org.apache.beam.sdk.schemas.Schema; import org.apache.beam.sdk.testing.PAssert; import org.apache.beam.sdk.testing.TestPipeline; import org.apache.beam.sdk.transforms.Count; +import org.apache.beam.sdk.transforms.PTransform; +import org.apache.beam.sdk.values.PBegin; import org.apache.beam.sdk.values.PCollection; import org.apache.beam.sdk.values.Row; import org.junit.Before; @@ -114,7 +117,7 @@ public void testLocationPathConventions() { "CREATE EXTERNAL TABLE DirTable %s TYPE parquet LOCATION '%s'", FIELD_NAMES, writeLocation)); PCollection dirResult = - BeamSqlRelUtils.toPCollection(readPipeline, env.parseQuery("SELECT * FROM DirTable")); + readPipeline.apply("dir", new Query(env.parseQuery("SELECT * FROM DirTable"))); PAssert.that("Directory with '/' reads all files", dirResult).containsInAnyOrder(ALL_ROWS); String globPath = new File(destinationDir, "output-*").getAbsolutePath(); @@ -123,7 +126,7 @@ public void testLocationPathConventions() { "CREATE EXTERNAL TABLE GlobTable %s TYPE parquet LOCATION '%s'", FIELD_NAMES, globPath)); PCollection globResult = - BeamSqlRelUtils.toPCollection(readPipeline, env.parseQuery("SELECT * FROM GlobTable")); + readPipeline.apply("glob", new Query(env.parseQuery("SELECT * FROM GlobTable"))); PAssert.that("Glob 'output-*' reads all files", globResult).containsInAnyOrder(ALL_ROWS); File[] writtenFiles = destinationDir.listFiles((dir, name) -> name.startsWith("output-")); @@ -137,8 +140,8 @@ public void testLocationPathConventions() { "CREATE EXTERNAL TABLE SingleFileTable %s TYPE parquet LOCATION '%s'", FIELD_NAMES, singleFilePath)); PCollection singleFileResult = - BeamSqlRelUtils.toPCollection( - readPipeline, env.parseQuery("SELECT * FROM SingleFileTable")); + readPipeline.apply( + "singleFile", new Query(env.parseQuery("SELECT * FROM SingleFileTable"))); PCollection count = singleFileResult.apply(Count.globally()); PAssert.thatSingleton(count) @@ -151,4 +154,22 @@ public void testLocationPathConventions() { PipelineResult.State state = readPipeline.run().waitUntilFinish(); assertEquals(State.DONE, state); } + + /** + * Expands a plan in a composite of its own, the way {@code SqlTransform} does. Stage names are + * unique within the query that produced them, so several queries sharing a pipeline need a scope + * each. + */ + private static class Query extends PTransform> { + private final BeamRelNode plan; + + Query(BeamRelNode plan) { + this.plan = plan; + } + + @Override + public PCollection expand(PBegin input) { + return BeamSqlRelUtils.toPCollection(input.getPipeline(), plan); + } + } }