polyzos commented on code in PR #29653:
URL: https://github.com/apache/flink/pull/29653#discussion_r4422968115
##########
flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/optimize/program/FlinkChangelogModeInferenceProgram.scala:
##########
@@ -1713,15 +1710,16 @@ class FlinkChangelogModeInferenceProgram extends
FlinkOptimizeProgram[StreamOpti
}
/**
- * Whether this calc's own output still has an upsert key after its
projection. A DELETE_BY_KEY
- * tombstone passed through this calc must still carry a key in its output,
otherwise nothing
- * downstream would know what to delete.
+ * Whether this calc's own output still has an upsert key after its
projection, including one
+ * recovered through an equi-join equivalence. A DELETE_BY_KEY tombstone
passed through this calc
+ * must still carry a key in its output, otherwise nothing downstream would
know what to delete.
*
* This method is an extra safety net, in case downstream consumers don't
require an upsert key.
*/
private def hasOutputUpsertKey(calc: StreamPhysicalCalcBase): Boolean = {
val fmq =
FlinkRelMetadataQuery.reuseOrCreate(calc.getCluster.getMetadataQuery)
- val upsertKeys = fmq.getUpsertKeys(calc)
+ val allColumns = ImmutableBitSet.range(calc.getRowType.getFieldCount)
+ val upsertKeys = UpsertKeyEquivalenceUtil.upsertKeysWithin(calc,
allColumns, fmq)
Review Comment:
A couple of sibling checks (join update-before, PTF update-before) still use
the raw upsert-key lookup instead of the new equivalence-aware one.
Intentional, or worth including too?
##########
flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/metadata/FlinkRelMdEquiJoinColumnPairs.java:
##########
@@ -0,0 +1,180 @@
+/*
+ * 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.flink.table.planner.plan.metadata;
+
+import
org.apache.flink.table.planner.plan.nodes.physical.stream.StreamPhysicalMultiJoin;
+
+import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.plan.hep.HepRelVertex;
+import org.apache.calcite.plan.volcano.RelSubset;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Calc;
+import org.apache.calcite.rel.core.Exchange;
+import org.apache.calcite.rel.core.Join;
+import org.apache.calcite.rel.core.JoinRelType;
+import org.apache.calcite.rel.core.Project;
+import org.apache.calcite.rel.metadata.MetadataDef;
+import org.apache.calcite.rel.metadata.MetadataHandler;
+import org.apache.calcite.rel.metadata.ReflectiveRelMetadataProvider;
+import org.apache.calcite.rel.metadata.RelMetadataProvider;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexInputRef;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.util.Util;
+import org.apache.calcite.util.mapping.IntPair;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/** A metadata handler for {@link FlinkMetadata.EquiJoinColumnPairs}. */
+public class FlinkRelMdEquiJoinColumnPairs
+ implements MetadataHandler<FlinkMetadata.EquiJoinColumnPairs> {
+ static final FlinkRelMdEquiJoinColumnPairs INSTANCE = new
FlinkRelMdEquiJoinColumnPairs();
+
+ public static final RelMetadataProvider SOURCE =
+ ReflectiveRelMetadataProvider.reflectiveSource(
+ FlinkMetadata.EquiJoinColumnPairs.METHOD, INSTANCE);
+
+ // ~ Constructors
-----------------------------------------------------------
+
+ private FlinkRelMdEquiJoinColumnPairs() {}
+
+ // ~ Methods
----------------------------------------------------------------
+
+ public MetadataDef<FlinkMetadata.EquiJoinColumnPairs> getDef() {
+ return FlinkMetadata.EquiJoinColumnPairs.DEF;
+ }
+
+ public List<EquiJoinColumnPair> getEquiJoinColumnPairs(Join join,
RelMetadataQuery mq) {
Review Comment:
Specialized join types (temporal, window, interval, delta) inherit this
handling untested - does the uniqueness assumption hold for them too?
##########
flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/utils/UpsertKeyEquivalenceUtil.scala:
##########
@@ -0,0 +1,381 @@
+/*
+ * 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.flink.table.planner.plan.utils
+
+import org.apache.flink.table.planner.plan.metadata.{EquiJoinColumnPair,
FlinkRelMetadataQuery}
+import
org.apache.flink.table.planner.plan.nodes.physical.stream.StreamPhysicalMultiJoin
+
+import org.apache.calcite.plan.hep.HepRelVertex
+import org.apache.calcite.plan.volcano.RelSubset
+import org.apache.calcite.rel.`type`.RelDataType
+import org.apache.calcite.rel.RelNode
+import org.apache.calcite.rel.core.{Calc, Exchange, Join, JoinRelType, Project}
+import org.apache.calcite.rex.{RexInputRef, RexNode}
+import org.apache.calcite.util.{ImmutableBitSet, Util}
+
+import java.util
+import java.util.{List => JList, Map => JMap, Set => JSet}
+
+import scala.collection.JavaConverters._
+
+/**
+ * Finds the upsert keys of a subtree that fall within a target column set,
treating the columns of
+ * [[FlinkRelMetadataQuery#getEquiJoinColumnPairs]] as interchangeable where
the substitution is
+ * null-safe (see `directionalSubstitutions`), also through chains of such
pairs. The target is
+ * fixed, so this stays bounded -- it never materializes equivalent key
variants the way eager key
+ * enrichment would. For a Calc that forwards deletes by key,
[[upsertKeyColumns]] gives the columns
+ * such keys can use, as one column set, so it does not enumerate keys either.
+ */
+object UpsertKeyEquivalenceUtil {
+
+ /**
+ * Upsert keys of `rel` (in its output columns) that are a subset of
`target`, both those reported
+ * by [[FlinkRelMetadataQuery#getUpsertKeys]] and those recovered by
null-safe equi-join
+ * substitution. Null if the upsert keys are undetermined.
+ */
+ def upsertKeysWithin(
+ rel: RelNode,
+ target: ImmutableBitSet,
+ mq: FlinkRelMetadataQuery): JSet[ImmutableBitSet] = {
+ val upsertKeys = mq.getUpsertKeys(rel)
+ if (upsertKeys == null) {
+ return null
+ }
+ val result = new util.LinkedHashSet[ImmutableBitSet]()
+ upsertKeys.asScala.foreach(key => if (target.contains(key))
result.add(key))
+ result.addAll(recoverWithin(rel, target, mq))
+ result
+ }
+
+ /**
+ * Every output column of `rel` that can be part of one of its upsert keys,
reported or recovered
+ * by equi-join substitution, for any target: a superset of the columns of
every key
+ * [[upsertKeysWithin]] can return. Recovered columns are pass-through join
columns, which every
+ * row a join emits carries: a join accepts deletes by key only from an
input whose unique key
+ * contains its join key, and a row an all-inner multi-join emits matched
every join condition.
+ */
+ def upsertKeyColumns(rel: RelNode, mq: FlinkRelMetadataQuery):
ImmutableBitSet = {
+ val upsertKeys = mq.getUpsertKeys(rel)
+ if (upsertKeys == null) {
+ return ImmutableBitSet.of()
+ }
+ keyColumns(upsertKeys).union(recoverableColumns(rel, mq))
+ }
+
+ /** The columns [[recoverWithin]] can use, over all targets. */
+ private def recoverableColumns(rel: RelNode, mq: FlinkRelMetadataQuery):
ImmutableBitSet =
+ rel match {
+ case subset: RelSubset =>
+ recoverableColumns(Util.first(subset.getBest, subset.getOriginal), mq)
+ case vertex: HepRelVertex =>
+ recoverableColumns(vertex.getCurrentRel, mq)
+ case calc: Calc =>
+ projectColumns(upsertKeyColumns(calc.getInput, mq), calcProjects(calc))
+ case project: Project =>
+ projectColumns(upsertKeyColumns(project.getInput, mq),
project.getProjects.asScala)
+ case exchange: Exchange =>
+ upsertKeyColumns(exchange.getInput, mq)
+ case join @ (_: Join | _: StreamPhysicalMultiJoin) =>
+ recoverableColumnsAtJoin(join, mq)
+ case _ =>
+ ImmutableBitSet.of()
+ }
+
+ /** [[recoverAtJoin]] over all targets. */
+ private def recoverableColumnsAtJoin(
+ join: RelNode,
+ mq: FlinkRelMetadataQuery): ImmutableBitSet = {
+ val columns = keyPreservingInputs(join,
mq).foldLeft(keyColumns(mq.getUpsertKeys(join))) {
+ case (acc, (input, offset)) => acc.union(upsertKeyColumns(input,
mq).shift(offset))
+ }
+ withSubstitutes(columns, substitutionsOf(join, mq))
+ }
+
+ /**
+ * Keys that fit `target` only after an equi-join substitution.
[[recoverableColumns]] walks the
+ * same nodes, so a node handled here must be handled there too.
+ */
+ private def recoverWithin(
+ rel: RelNode,
+ target: ImmutableBitSet,
+ mq: FlinkRelMetadataQuery): JSet[ImmutableBitSet] = rel match {
+ case subset: RelSubset =>
+ recoverWithin(Util.first(subset.getBest, subset.getOriginal), target, mq)
+ case vertex: HepRelVertex =>
+ recoverWithin(vertex.getCurrentRel, target, mq)
+ case calc: Calc =>
+ recoverThroughProjection(calcProjects(calc), calc.getInput, target, mq)
+ case project: Project =>
+ recoverThroughProjection(project.getProjects.asScala, project.getInput,
target, mq)
+ case exchange: Exchange =>
+ recoverWithin(exchange.getInput, target, mq)
+ case join @ (_: Join | _: StreamPhysicalMultiJoin) =>
+ recoverAtJoin(join, target, mq)
+ case _ =>
+ util.Collections.emptySet()
+ }
+
+ /** Maps the target down to the projection input, recovers keys there, and
maps them back up. */
+ private def recoverThroughProjection(
+ projects: Seq[RexNode],
+ input: RelNode,
+ target: ImmutableBitSet,
+ mq: FlinkRelMetadataQuery): JSet[ImmutableBitSet] = {
+ val inputToOutput = inputToOutputMap(projects)
+ val inputTarget = ImmutableBitSet.builder()
+ inputToOutput.asScala.foreach {
+ case (in, outs) => if (outs.asScala.exists(target.get(_)))
inputTarget.set(in)
+ }
+
+ val result = new util.LinkedHashSet[ImmutableBitSet]()
+ recoverWithin(input, inputTarget.build(), mq).asScala.foreach {
+ inputKey =>
+ val outputKey = ImmutableBitSet.builder()
+ var mappable = true
+ for (inIdx <- inputKey.toArray if mappable) {
+ inputToOutput.get(inIdx).asScala.find(target.get(_)) match {
+ case Some(out) => outputKey.set(out)
+ case None => mappable = false // not kept within the target ->
cannot be expressed
+ }
+ }
+ if (mappable) {
+ result.add(outputKey.build())
+ }
+ }
+ result
+ }
+
+ /**
+ * Repairs the join's own keys, then recurses into the inputs whose keys the
join preserves --
+ * which reaches a deeper join when an intermediate projection dropped the
equated column.
+ */
+ private def recoverAtJoin(
+ join: RelNode,
+ target: ImmutableBitSet,
+ mq: FlinkRelMetadataQuery): JSet[ImmutableBitSet] = {
+ val result = new util.LinkedHashSet[ImmutableBitSet]()
+ val substitutes = substitutionsOf(join, mq)
+ addRepaired(mq.getUpsertKeys(join), 0, target, substitutes, result)
+
+ val expanded = expandToward(target, substitutes)
+ keyPreservingInputs(join, mq).foreach {
+ case (input, offset) =>
+ val inputTarget = childTarget(expanded, offset,
input.getRowType.getFieldCount)
+ addRepaired(upsertKeysWithin(input, inputTarget, mq), offset, target,
substitutes, result)
+ }
+ result
+ }
+
+ /**
+ * The inputs of a join or multi-join whose upsert keys stay unique in its
output, each with the
+ * position of its first column there. Shared by [[recoverAtJoin]] and
+ * [[recoverableColumnsAtJoin]] so that [[upsertKeyColumns]] covers every
key [[upsertKeysWithin]]
+ * can return.
+ */
+ private def keyPreservingInputs(rel: RelNode, mq: FlinkRelMetadataQuery):
Seq[(RelNode, Int)] =
+ rel match {
+ case multiJoin: StreamPhysicalMultiJoin =>
+ // only all-inner (no padding); an input is preserved when every other
input is unique on
+ // its join key
+ if (multiJoin.getJoinTypes.asScala.forall(_ == JoinRelType.INNER)) {
+ val inputs = multiJoin.getInputs.asScala
+ val offsets = inputs.scanLeft(0)(_ + _.getRowType.getFieldCount)
+ val unique = inputs.indices.map {
+ j =>
+ isSideUnique(
+ mq.getUpsertKeys(inputs(j)),
+ ImmutableBitSet.of(multiJoin.getJoinKeyIndices(j): _*))
+ }
+ inputs.indices
+ .filter(i => inputs.indices.forall(j => j == i || unique(j)))
+ .map(i => (inputs(i), offsets(i)))
Review Comment:
Input 0 never has real join-key indices at runtime, so this check can
basically never mark any other input as key-preserving. Falls back safely to
upsertMaterialize, but it silently disables recovery through multi-join
children.
##########
flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/utils/UpsertKeyEquivalenceUtil.scala:
##########
@@ -0,0 +1,381 @@
+/*
+ * 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.flink.table.planner.plan.utils
+
+import org.apache.flink.table.planner.plan.metadata.{EquiJoinColumnPair,
FlinkRelMetadataQuery}
+import
org.apache.flink.table.planner.plan.nodes.physical.stream.StreamPhysicalMultiJoin
+
+import org.apache.calcite.plan.hep.HepRelVertex
+import org.apache.calcite.plan.volcano.RelSubset
+import org.apache.calcite.rel.`type`.RelDataType
+import org.apache.calcite.rel.RelNode
+import org.apache.calcite.rel.core.{Calc, Exchange, Join, JoinRelType, Project}
+import org.apache.calcite.rex.{RexInputRef, RexNode}
+import org.apache.calcite.util.{ImmutableBitSet, Util}
+
+import java.util
+import java.util.{List => JList, Map => JMap, Set => JSet}
+
+import scala.collection.JavaConverters._
+
+/**
+ * Finds the upsert keys of a subtree that fall within a target column set,
treating the columns of
+ * [[FlinkRelMetadataQuery#getEquiJoinColumnPairs]] as interchangeable where
the substitution is
+ * null-safe (see `directionalSubstitutions`), also through chains of such
pairs. The target is
+ * fixed, so this stays bounded -- it never materializes equivalent key
variants the way eager key
+ * enrichment would. For a Calc that forwards deletes by key,
[[upsertKeyColumns]] gives the columns
+ * such keys can use, as one column set, so it does not enumerate keys either.
+ */
+object UpsertKeyEquivalenceUtil {
+
+ /**
+ * Upsert keys of `rel` (in its output columns) that are a subset of
`target`, both those reported
+ * by [[FlinkRelMetadataQuery#getUpsertKeys]] and those recovered by
null-safe equi-join
+ * substitution. Null if the upsert keys are undetermined.
+ */
+ def upsertKeysWithin(
+ rel: RelNode,
+ target: ImmutableBitSet,
+ mq: FlinkRelMetadataQuery): JSet[ImmutableBitSet] = {
+ val upsertKeys = mq.getUpsertKeys(rel)
+ if (upsertKeys == null) {
+ return null
+ }
+ val result = new util.LinkedHashSet[ImmutableBitSet]()
+ upsertKeys.asScala.foreach(key => if (target.contains(key))
result.add(key))
+ result.addAll(recoverWithin(rel, target, mq))
+ result
+ }
+
+ /**
+ * Every output column of `rel` that can be part of one of its upsert keys,
reported or recovered
+ * by equi-join substitution, for any target: a superset of the columns of
every key
+ * [[upsertKeysWithin]] can return. Recovered columns are pass-through join
columns, which every
+ * row a join emits carries: a join accepts deletes by key only from an
input whose unique key
+ * contains its join key, and a row an all-inner multi-join emits matched
every join condition.
+ */
+ def upsertKeyColumns(rel: RelNode, mq: FlinkRelMetadataQuery):
ImmutableBitSet = {
+ val upsertKeys = mq.getUpsertKeys(rel)
+ if (upsertKeys == null) {
+ return ImmutableBitSet.of()
+ }
+ keyColumns(upsertKeys).union(recoverableColumns(rel, mq))
+ }
+
+ /** The columns [[recoverWithin]] can use, over all targets. */
+ private def recoverableColumns(rel: RelNode, mq: FlinkRelMetadataQuery):
ImmutableBitSet =
Review Comment:
This walks the same node types as its sibling recovery method and has to be
kept in sync by hand. Should this be merged into one traversal?
##########
flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/nodes/physical/stream/StreamPhysicalSink.scala:
##########
@@ -147,8 +147,9 @@ class StreamPhysicalSink(
val primaryKeys =
contextResolvedTable.getResolvedSchema.getPrimaryKeyIndexes
val pks = ImmutableBitSet.of(primaryKeys: _*)
val fmq = FlinkRelMetadataQuery.reuseOrCreate(getCluster.getMetadataQuery)
- val changeLogUpsertKeys = fmq.getUpsertKeys(getInput)
- changeLogUpsertKeys != null &&
changeLogUpsertKeys.asScala.exists(pks.contains)
+ // upsert keys within the pk, treating equi-join-equivalent columns as
interchangeable
+ val changeLogUpsertKeys =
UpsertKeyEquivalenceUtil.upsertKeysWithin(getInput, pks, fmq)
Review Comment:
The key actually passed to the sink's materializer is still computed with
the raw upsert-key lookup, unlike this check. Is this intentional?
##########
flink-table/flink-table-planner/src/main/scala/org/apache/flink/table/planner/plan/utils/UpsertKeyEquivalenceUtil.scala:
##########
@@ -0,0 +1,381 @@
+/*
+ * 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.flink.table.planner.plan.utils
+
+import org.apache.flink.table.planner.plan.metadata.{EquiJoinColumnPair,
FlinkRelMetadataQuery}
+import
org.apache.flink.table.planner.plan.nodes.physical.stream.StreamPhysicalMultiJoin
+
+import org.apache.calcite.plan.hep.HepRelVertex
+import org.apache.calcite.plan.volcano.RelSubset
+import org.apache.calcite.rel.`type`.RelDataType
+import org.apache.calcite.rel.RelNode
+import org.apache.calcite.rel.core.{Calc, Exchange, Join, JoinRelType, Project}
+import org.apache.calcite.rex.{RexInputRef, RexNode}
+import org.apache.calcite.util.{ImmutableBitSet, Util}
+
+import java.util
+import java.util.{List => JList, Map => JMap, Set => JSet}
+
+import scala.collection.JavaConverters._
+
+/**
+ * Finds the upsert keys of a subtree that fall within a target column set,
treating the columns of
+ * [[FlinkRelMetadataQuery#getEquiJoinColumnPairs]] as interchangeable where
the substitution is
+ * null-safe (see `directionalSubstitutions`), also through chains of such
pairs. The target is
+ * fixed, so this stays bounded -- it never materializes equivalent key
variants the way eager key
+ * enrichment would. For a Calc that forwards deletes by key,
[[upsertKeyColumns]] gives the columns
+ * such keys can use, as one column set, so it does not enumerate keys either.
+ */
+object UpsertKeyEquivalenceUtil {
+
+ /**
+ * Upsert keys of `rel` (in its output columns) that are a subset of
`target`, both those reported
+ * by [[FlinkRelMetadataQuery#getUpsertKeys]] and those recovered by
null-safe equi-join
+ * substitution. Null if the upsert keys are undetermined.
+ */
+ def upsertKeysWithin(
+ rel: RelNode,
+ target: ImmutableBitSet,
+ mq: FlinkRelMetadataQuery): JSet[ImmutableBitSet] = {
+ val upsertKeys = mq.getUpsertKeys(rel)
+ if (upsertKeys == null) {
+ return null
+ }
+ val result = new util.LinkedHashSet[ImmutableBitSet]()
+ upsertKeys.asScala.foreach(key => if (target.contains(key))
result.add(key))
+ result.addAll(recoverWithin(rel, target, mq))
+ result
+ }
+
+ /**
+ * Every output column of `rel` that can be part of one of its upsert keys,
reported or recovered
+ * by equi-join substitution, for any target: a superset of the columns of
every key
+ * [[upsertKeysWithin]] can return. Recovered columns are pass-through join
columns, which every
+ * row a join emits carries: a join accepts deletes by key only from an
input whose unique key
+ * contains its join key, and a row an all-inner multi-join emits matched
every join condition.
+ */
+ def upsertKeyColumns(rel: RelNode, mq: FlinkRelMetadataQuery):
ImmutableBitSet = {
+ val upsertKeys = mq.getUpsertKeys(rel)
+ if (upsertKeys == null) {
+ return ImmutableBitSet.of()
+ }
+ keyColumns(upsertKeys).union(recoverableColumns(rel, mq))
+ }
+
+ /** The columns [[recoverWithin]] can use, over all targets. */
+ private def recoverableColumns(rel: RelNode, mq: FlinkRelMetadataQuery):
ImmutableBitSet =
+ rel match {
+ case subset: RelSubset =>
+ recoverableColumns(Util.first(subset.getBest, subset.getOriginal), mq)
+ case vertex: HepRelVertex =>
+ recoverableColumns(vertex.getCurrentRel, mq)
+ case calc: Calc =>
+ projectColumns(upsertKeyColumns(calc.getInput, mq), calcProjects(calc))
+ case project: Project =>
+ projectColumns(upsertKeyColumns(project.getInput, mq),
project.getProjects.asScala)
+ case exchange: Exchange =>
+ upsertKeyColumns(exchange.getInput, mq)
+ case join @ (_: Join | _: StreamPhysicalMultiJoin) =>
+ recoverableColumnsAtJoin(join, mq)
+ case _ =>
+ ImmutableBitSet.of()
+ }
+
+ /** [[recoverAtJoin]] over all targets. */
+ private def recoverableColumnsAtJoin(
+ join: RelNode,
+ mq: FlinkRelMetadataQuery): ImmutableBitSet = {
+ val columns = keyPreservingInputs(join,
mq).foldLeft(keyColumns(mq.getUpsertKeys(join))) {
+ case (acc, (input, offset)) => acc.union(upsertKeyColumns(input,
mq).shift(offset))
+ }
+ withSubstitutes(columns, substitutionsOf(join, mq))
+ }
+
+ /**
+ * Keys that fit `target` only after an equi-join substitution.
[[recoverableColumns]] walks the
+ * same nodes, so a node handled here must be handled there too.
+ */
+ private def recoverWithin(
+ rel: RelNode,
+ target: ImmutableBitSet,
+ mq: FlinkRelMetadataQuery): JSet[ImmutableBitSet] = rel match {
+ case subset: RelSubset =>
+ recoverWithin(Util.first(subset.getBest, subset.getOriginal), target, mq)
+ case vertex: HepRelVertex =>
+ recoverWithin(vertex.getCurrentRel, target, mq)
+ case calc: Calc =>
+ recoverThroughProjection(calcProjects(calc), calc.getInput, target, mq)
+ case project: Project =>
+ recoverThroughProjection(project.getProjects.asScala, project.getInput,
target, mq)
+ case exchange: Exchange =>
+ recoverWithin(exchange.getInput, target, mq)
+ case join @ (_: Join | _: StreamPhysicalMultiJoin) =>
+ recoverAtJoin(join, target, mq)
+ case _ =>
+ util.Collections.emptySet()
+ }
+
+ /** Maps the target down to the projection input, recovers keys there, and
maps them back up. */
+ private def recoverThroughProjection(
+ projects: Seq[RexNode],
+ input: RelNode,
+ target: ImmutableBitSet,
+ mq: FlinkRelMetadataQuery): JSet[ImmutableBitSet] = {
+ val inputToOutput = inputToOutputMap(projects)
+ val inputTarget = ImmutableBitSet.builder()
+ inputToOutput.asScala.foreach {
+ case (in, outs) => if (outs.asScala.exists(target.get(_)))
inputTarget.set(in)
+ }
+
+ val result = new util.LinkedHashSet[ImmutableBitSet]()
+ recoverWithin(input, inputTarget.build(), mq).asScala.foreach {
+ inputKey =>
+ val outputKey = ImmutableBitSet.builder()
+ var mappable = true
+ for (inIdx <- inputKey.toArray if mappable) {
+ inputToOutput.get(inIdx).asScala.find(target.get(_)) match {
+ case Some(out) => outputKey.set(out)
+ case None => mappable = false // not kept within the target ->
cannot be expressed
+ }
+ }
+ if (mappable) {
+ result.add(outputKey.build())
+ }
+ }
+ result
+ }
+
+ /**
+ * Repairs the join's own keys, then recurses into the inputs whose keys the
join preserves --
+ * which reaches a deeper join when an intermediate projection dropped the
equated column.
+ */
+ private def recoverAtJoin(
+ join: RelNode,
+ target: ImmutableBitSet,
+ mq: FlinkRelMetadataQuery): JSet[ImmutableBitSet] = {
+ val result = new util.LinkedHashSet[ImmutableBitSet]()
+ val substitutes = substitutionsOf(join, mq)
+ addRepaired(mq.getUpsertKeys(join), 0, target, substitutes, result)
+
+ val expanded = expandToward(target, substitutes)
+ keyPreservingInputs(join, mq).foreach {
+ case (input, offset) =>
+ val inputTarget = childTarget(expanded, offset,
input.getRowType.getFieldCount)
+ addRepaired(upsertKeysWithin(input, inputTarget, mq), offset, target,
substitutes, result)
+ }
+ result
+ }
+
+ /**
+ * The inputs of a join or multi-join whose upsert keys stay unique in its
output, each with the
+ * position of its first column there. Shared by [[recoverAtJoin]] and
+ * [[recoverableColumnsAtJoin]] so that [[upsertKeyColumns]] covers every
key [[upsertKeysWithin]]
+ * can return.
+ */
+ private def keyPreservingInputs(rel: RelNode, mq: FlinkRelMetadataQuery):
Seq[(RelNode, Int)] =
+ rel match {
+ case multiJoin: StreamPhysicalMultiJoin =>
+ // only all-inner (no padding); an input is preserved when every other
input is unique on
+ // its join key
+ if (multiJoin.getJoinTypes.asScala.forall(_ == JoinRelType.INNER)) {
+ val inputs = multiJoin.getInputs.asScala
+ val offsets = inputs.scanLeft(0)(_ + _.getRowType.getFieldCount)
+ val unique = inputs.indices.map {
+ j =>
+ isSideUnique(
+ mq.getUpsertKeys(inputs(j)),
+ ImmutableBitSet.of(multiJoin.getJoinKeyIndices(j): _*))
+ }
+ inputs.indices
+ .filter(i => inputs.indices.forall(j => j == i || unique(j)))
+ .map(i => (inputs(i), offsets(i)))
+ } else {
+ Seq.empty
+ }
+ case join: Join if !join.getJoinType.projectsRight =>
+ // SEMI/ANTI output the left side only and never duplicate a left row,
so left keys survive
+ Seq((join.getLeft, 0))
+ case join: Join =>
+ // a side's keys survive iff the other side is unique on the join key
and this side isn't
+ // padded
+ val joinType = join.getJoinType
+ val joinInfo = join.analyzeCondition()
+ val left =
+ if (
+ !joinType.generatesNullsOnLeft &&
+ isSideUnique(mq.getUpsertKeys(join.getRight),
ImmutableBitSet.of(joinInfo.rightKeys))
+ ) Seq((join.getLeft, 0))
+ else Seq.empty
+ val right =
+ if (
+ !joinType.generatesNullsOnRight &&
+ isSideUnique(mq.getUpsertKeys(join.getLeft),
ImmutableBitSet.of(joinInfo.leftKeys))
+ ) Seq((join.getRight, join.getLeft.getRowType.getFieldCount))
+ else Seq.empty
+ left ++ right
+ case _ =>
+ Seq.empty
+ }
+
+ /**
+ * Repairs each key toward `target` (shifting up by `offset` first) and adds
the ones that fit.
+ */
+ private def addRepaired(
+ keys: JSet[ImmutableBitSet],
+ offset: Int,
+ target: ImmutableBitSet,
+ substitutes: JMap[Integer, JList[Integer]],
+ result: JSet[ImmutableBitSet]): Unit = {
+ if (keys == null) {
+ return
+ }
+ keys.asScala.foreach(
+ key => repairWithin(key.shift(offset), target,
substitutes).foreach(result.add))
+ }
+
+ /** The columns a child must carry (in its own space) to land within
`expanded`. */
+ private def childTarget(
+ expanded: ImmutableBitSet,
+ offset: Int,
+ fieldCount: Int): ImmutableBitSet =
+ expanded.intersect(ImmutableBitSet.range(offset, offset +
fieldCount)).shift(-offset)
+
+ private def isSideUnique(
+ sideKeys: JSet[ImmutableBitSet],
+ sideJoinKeys: ImmutableBitSet): Boolean =
+ sideKeys != null && sideKeys.asScala.exists(key =>
sideJoinKeys.contains(key))
+
+ private def substitutionsOf(
+ rel: RelNode,
+ mq: FlinkRelMetadataQuery): JMap[Integer, JList[Integer]] =
+ directionalSubstitutions(mq.getEquiJoinColumnPairs(rel), rel.getRowType)
+
+ private def calcProjects(calc: Calc): Seq[RexNode] =
+ calc.getProgram.getProjectList.asScala.map(calc.getProgram.expandLocalRef)
+
+ /** Each input column to the output positions where a projection passes it
through as-is. */
+ private def inputToOutputMap(projects: Seq[RexNode]): JMap[Integer,
JList[Integer]] = {
Review Comment:
This duplicates a projection-mapping helper that already exists elsewhere in
the codebase. Minor, not blocking.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
