fresh-borzoni commented on code in PR #29653:
URL: https://github.com/apache/flink/pull/29653#discussion_r4526613106
##########
flink-table/flink-table-planner/src/main/java/org/apache/flink/table/planner/plan/utils/UpsertKeyUtil.java:
##########
@@ -63,4 +78,97 @@ public static Optional<int[]> smallestKey(@Nullable
Set<ImmutableBitSet> upsertK
upsertKeys.stream().map(ImmutableBitSet::toArray).collect(Collectors.toList());
return Optional.of(UpsertKeyUtils.smallestKey(asArrays));
}
+
+ /**
+ * Returns the upsert keys of the given node, plus each key whose columns
outside {@code
+ * columns} can all be replaced by columns inside that are equal to them
on every row, e.g.
+ * through an inner equi-join condition. For example, for {@code t1 JOIN
t2 ON a1 = a2} and
+ * columns {a1, b2}, the upsert key {a2, b2} adds {a1, b2}.
+ *
+ * @return the upsert keys and the aligned ones, or null if the upsert
keys are unknown
+ */
+ public static @Nullable Set<ImmutableBitSet> getUpsertKeysAlignedTo(
+ RelNode rel, ImmutableBitSet columns, RelMetadataQuery mq) {
+ final FlinkRelMetadataQuery fmq =
FlinkRelMetadataQuery.reuseOrCreate(mq);
+ final Set<ImmutableBitSet> upsertKeys = fmq.getUpsertKeys(rel);
+ if (upsertKeys == null ||
upsertKeys.stream().allMatch(columns::contains)) {
+ return upsertKeys;
+ }
+
+ final Map<Integer, Integer> parents = getEqualColumnParents(rel, fmq);
+ final Map<Integer, Integer> replacements = new HashMap<>();
+ for (int column : columns) {
+ replacements.putIfAbsent(root(parents, column), column);
+ }
+ final Set<ImmutableBitSet> alignedKeys = new HashSet<>(upsertKeys);
+ for (ImmutableBitSet key : upsertKeys) {
+ final ImmutableBitSet.Builder alignedKey =
ImmutableBitSet.builder();
+ for (int column : key) {
+ alignedKey.set(
+ columns.get(column)
+ ? column
+ : replacements.getOrDefault(root(parents,
column), column));
+ }
+ final ImmutableBitSet aligned = alignedKey.build();
+ if (columns.contains(aligned)) {
+ alignedKeys.add(aligned);
+ }
+ }
+ return alignedKeys;
+ }
+
+ /**
+ * Returns {@code columns} and every column of the given node that is
equal to one of them on
+ * every row.
+ */
+ public static ImmutableBitSet getEqualColumns(
+ RelNode rel, ImmutableBitSet columns, RelMetadataQuery mq) {
+ final Map<Integer, Integer> parents = getEqualColumnParents(rel, mq);
+ final Set<Integer> roots =
+ columns.asList().stream()
+ .map(column -> root(parents, column))
+ .collect(Collectors.toSet());
+ final ImmutableBitSet.Builder equalColumns = columns.rebuild();
+ for (int column : parents.keySet()) {
+ final int root = root(parents, column);
+ if (roots.contains(root)) {
+ equalColumns.set(column).set(root);
+ }
+ }
+ return equalColumns.build();
+ }
+
+ /** Groups the columns that are equal on every row, as a union-find forest
of parents. */
+ private static Map<Integer, Integer> getEqualColumnParents(RelNode rel,
RelMetadataQuery mq) {
+ final Map<Integer, Integer> parents = new HashMap<>();
+ for (RexNode predicate :
mq.getPulledUpPredicates(rel).pulledUpPredicates) {
+ if (!predicate.isA(EQUALITY_KINDS)) {
+ continue;
+ }
+ final RexNode left = ((RexCall) predicate).getOperands().get(0);
+ final RexNode right = ((RexCall) predicate).getOperands().get(1);
+ // comparing different types may implicitly cast one side, which
needn't be injective,
+ // and FLOAT/DOUBLE compare -0.0 = 0.0 although they are different
keys
+ if (left instanceof RexInputRef
+ && right instanceof RexInputRef
+ && SqlTypeUtil.equalSansNullability(left.getType(),
right.getType())
+ && !SqlTypeUtil.isApproximateNumeric(left.getType())) {
Review Comment:
Good catch, applied. Covered by testJoinUpsertKeyEquivalenceMapNotSatisfied.
--
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]
