fresh-borzoni commented on code in PR #29653: URL: https://github.com/apache/flink/pull/29653#discussion_r4423430988
########## 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: I don't think input 0 is empty: createJoinAttributeMap adds a {-1, -1} -> {0, col} entry for input 0, so getJoinKeyIndices(0) returns its join columns, and recovery through input 1 works. But the check is wrong in a different way: for input i > 0 it accepts input 0 if it's unique on all its join columns, even ones joined to an input not reached yet. E.g. t0 pk (k, y), t1 pk (k1, b1), t2 pk (k2, y2), t0.k = t1.k1, t0.k = t2.k2 AND t0.y = t2.y2: a t1 row matches every y, but its key passes. FlinkRelMdUpsertKeys has the same issue for multi-join and already reports {k1, b1} on master, so I'll open a separate ticket. -- 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]
