mihaibudiu commented on code in PR #4540:
URL: https://github.com/apache/calcite/pull/4540#discussion_r2364978363


##########
core/src/main/java/org/apache/calcite/rel/metadata/RelMdFunctionalDependency.java:
##########
@@ -64,212 +76,422 @@ protected RelMdFunctionalDependency() {}
     return BuiltInMetadata.FunctionalDependency.DEF;
   }
 
+  /**
+   * Determines if column is functionally dependent on key for a given rel 
node.
+   */
   public @Nullable Boolean determines(RelNode rel, RelMetadataQuery mq,
       int key, int column) {
-    return determinesImpl2(rel, mq, key, column);
+    return determinesSet(rel, mq, ImmutableBitSet.of(key), 
ImmutableBitSet.of(column));
   }
 
-  public @Nullable Boolean determines(SetOp rel, RelMetadataQuery mq,
-      int key, int column) {
-    return determinesImpl2(rel, mq, key, column);
+  /**
+   * Determines if a set of columns functionally determines another set of 
columns.
+   */
+  public Boolean determinesSet(RelNode rel, RelMetadataQuery mq,
+      ImmutableBitSet determinants, ImmutableBitSet dependents) {
+    FunctionalDependencySet fdSet = mq.getFunctionalDependencies(rel);
+    return fdSet.implies(determinants, dependents);
   }
 
-  public @Nullable Boolean determines(Join rel, RelMetadataQuery mq,
-      int key, int column) {
-    return determinesImpl2(rel, mq, key, column);
+  /**
+   * Returns the closure of a set of columns under all functional dependencies.
+   */
+  public ImmutableBitSet closure(RelNode rel, RelMetadataQuery mq, 
ImmutableBitSet attrs) {
+    FunctionalDependencySet fdSet = mq.getFunctionalDependencies(rel);
+    return fdSet.closure(attrs);
   }
 
-  public @Nullable Boolean determines(Correlate rel, RelMetadataQuery mq,
-      int key, int column) {
-    return determinesImpl2(rel, mq, key, column);
+  /**
+   * Returns candidate keys for the relation within the specified set of 
attributes.
+   */
+  public Set<ImmutableBitSet> candidateKeys(
+      RelNode rel, RelMetadataQuery mq, ImmutableBitSet attributes) {
+    FunctionalDependencySet fdSet = mq.getFunctionalDependencies(rel);
+    return fdSet.findCandidateKeys(attributes);
   }
 
-  public @Nullable Boolean determines(Aggregate rel, RelMetadataQuery mq,
-      int key, int column) {
-    return determinesImpl(rel, mq, key, column);
+  /**
+   * Main dispatch method for getFunctionalDependencies.
+   * Routes to appropriate handler based on RelNode type.
+   */
+  public FunctionalDependencySet getFunctionalDependencies(RelNode rel, 
RelMetadataQuery mq) {
+    if (rel instanceof TableScan) {
+      return getTableScanFD((TableScan) rel);
+    } else if (rel instanceof Project) {
+      return getProjectFD((Project) rel, mq);
+    } else if (rel instanceof Aggregate) {
+      return getAggregateFD((Aggregate) rel, mq);
+    } else if (rel instanceof Join) {
+      return getJoinFD((Join) rel, mq);
+    } else if (rel instanceof Calc) {
+      return getCalcFD((Calc) rel, mq);
+    } else if (rel instanceof SetOp) {
+      // TODO: Handle UNION, INTERSECT, EXCEPT functional dependencies
+      return new FunctionalDependencySet();
+    } else if (rel instanceof Correlate) {
+      // TODO: Handle CORRELATE functional dependencies
+      return new FunctionalDependencySet();
+    }
+    return getFD(rel.getInputs(), mq);
   }
 
-  public @Nullable Boolean determines(Calc rel, RelMetadataQuery mq,
-      int key, int column) {
-    return determinesImpl(rel, mq, key, column);
+  private static FunctionalDependencySet getFD(List<RelNode> inputs, 
RelMetadataQuery mq) {
+    FunctionalDependencySet result = new FunctionalDependencySet();
+    for (RelNode input : inputs) {
+      FunctionalDependencySet fdSet = mq.getFunctionalDependencies(input);
+      result = result.union(fdSet);
+    }
+    return result;
   }
 
-  public @Nullable Boolean determines(Project rel, RelMetadataQuery mq,
-      int key, int column) {
-    return determinesImpl(rel, mq, key, column);
-  }
+  private static FunctionalDependencySet getTableScanFD(TableScan rel) {
+    FunctionalDependencySet fdSet = new FunctionalDependencySet();
 
-  /**
-   * Checks if a column is functionally determined by a key column through 
expression analysis.
-   *
-   * @param rel The input relation
-   * @param mq Metadata query instance
-   * @param key Index of the determinant expression
-   * @param column Index of the dependent expression
-   * @return TRUE if column is determined by key,
-   *         FALSE if not determined,
-   *         NULL if undetermined
-   */
-  private static @Nullable Boolean determinesImpl(RelNode rel, 
RelMetadataQuery mq,
-      int key, int column) {
-    if (preCheck(rel, key, column)) {
-      return true;
+    RelOptTable table = rel.getTable();
+    List<ImmutableBitSet> keys = table.getKeys();
+    if (keys == null || keys.isEmpty()) {
+      return fdSet;
     }
 
-    ImmutableBitSet keyInputIndices = null;
-    ImmutableBitSet columnInputIndices = null;
-    if (rel instanceof Project || rel instanceof Calc) {
-      List<RexNode> exprs = null;
-      if (rel instanceof Project) {
-        Project project = (Project) rel;
-        exprs = project.getProjects();
-      } else {
-        Calc calc = (Calc) rel;
-        final RexProgram program = calc.getProgram();
-        exprs = program.expandList(program.getProjectList());
+    for (ImmutableBitSet key : keys) {
+      ImmutableBitSet allColumns = 
ImmutableBitSet.range(rel.getRowType().getFieldCount());
+      ImmutableBitSet dependents = allColumns.except(key);
+      if (!dependents.isEmpty()) {
+        fdSet.addFD(key, dependents);
       }
+    }
 
-      // TODO: Supports dependency analysis for all types of expressions
-      if (!(exprs.get(column) instanceof RexInputRef)) {
-        return false;
-      }
+    return fdSet;
+  }
 
-      RexNode keyExpr = exprs.get(key);
-      RexNode columnExpr = exprs.get(column);
+  private static FunctionalDependencySet getProjectFD(Project rel, 
RelMetadataQuery mq) {
+    return getProjectionFD(rel.getInput(), rel.getProjects(), mq);
+  }
 
-      // Identical expressions imply functional dependency
-      if (keyExpr.equals(columnExpr)) {
-        return true;
+  /**
+   * Common method to compute functional dependencies for projection 
operations.
+   * Used by both Project and Calc nodes.
+   *
+   * @param input the input relation
+   * @param projections the list of projection expressions
+   * @param mq the metadata query
+   * @return the functional dependency set for the projection
+   */
+  private static FunctionalDependencySet getProjectionFD(
+      RelNode input, List<RexNode> projections, RelMetadataQuery mq) {
+    FunctionalDependencySet inputFdSet = mq.getFunctionalDependencies(input);
+    FunctionalDependencySet projectionFdSet = new FunctionalDependencySet();
+    int fieldCount = projections.size();
+
+    // Create mapping from input column indices to project column indices
+    Mappings.TargetMapping inputToOutputMap =
+        RelOptUtil.permutation(projections, input.getRowType());
+
+    // Map input functional dependencies to project dependencies
+    mapInputFDs(inputFdSet, inputToOutputMap, projectionFdSet);
+
+    // For each pair of output columns, determine if one determines the other
+    for (int i = 0; i < fieldCount; i++) {
+      for (int j = i + 1; j < fieldCount; j++) {
+        RexNode expr1 = projections.get(i);
+        RexNode expr2 = projections.get(j);
+
+        // Handle identical expressions, they determine each other
+        if (expr1.equals(expr2) && RexUtil.isDeterministic(expr1)) {
+          projectionFdSet.addFD(i, j);
+          projectionFdSet.addFD(j, i);
+          continue;
+        }
+
+        // Handle literal constants, all columns determine literals
+        if (expr1 instanceof RexLiteral) {
+          projectionFdSet.addFD(j, i);
+        }
+        if (expr2 instanceof RexLiteral) {
+          projectionFdSet.addFD(i, j);
+        }
+
+        // For complex expressions, check if they have functional dependencies
+        if (!(expr1 instanceof RexLiteral) && !(expr2 instanceof RexLiteral)
+            && RexUtil.isDeterministic(expr1) && 
RexUtil.isDeterministic(expr2)) {
+          ImmutableBitSet inputs1 = RelOptUtil.InputFinder.bits(expr1);
+          ImmutableBitSet inputs2 = RelOptUtil.InputFinder.bits(expr2);
+
+          if (!inputs1.isEmpty() && !inputs2.isEmpty()) {
+            if (inputFdSet.implies(inputs1, inputs2)) {
+              projectionFdSet.addFD(i, j);
+            }
+            if (inputFdSet.implies(inputs2, inputs1)) {
+              projectionFdSet.addFD(j, i);
+            }
+          }
+        }
       }
+    }
 
-      keyInputIndices = extractDeterministicRefs(keyExpr);
-      columnInputIndices = extractDeterministicRefs(columnExpr);
-    } else if (rel instanceof Aggregate) {
-      Aggregate aggregate = (Aggregate) rel;
+    return projectionFdSet;
+  }
 
-      int groupByCnt = aggregate.getGroupCount();
-      if (key < groupByCnt && column >= groupByCnt) {
-        return false;
+  /**
+   * Maps input functional dependencies to output dependencies based on column 
mapping.
+   */
+  private static void mapInputFDs(FunctionalDependencySet inputFdSet,
+      Mappings.TargetMapping mapping, FunctionalDependencySet outputFdSet) {
+    for (FunctionalDependency inputFd : inputFdSet.getFDs()) {
+      ImmutableBitSet determinants = inputFd.getDeterminants();
+      ImmutableBitSet dependents = inputFd.getDependents();
+
+      // Skip this FD if any determinant column is unmappable
+      boolean allMappable =
+          determinants.stream().allMatch(col -> col >= 0
+              && col < mapping.getSourceCount()
+              && mapping.getTargetOpt(col) >= 0);
+      if (!allMappable) {
+        continue;
       }
 
-      keyInputIndices = extractDeterministicRefs(aggregate, key);
-      columnInputIndices = extractDeterministicRefs(aggregate, column);
-    } else {
-      throw new UnsupportedOperationException("Unsupported RelNode type: "
-          + rel.getClass().getSimpleName());
-    }
+      // Map all determinant columns
+      ImmutableBitSet mappedDeterminants = mapAllCols(determinants, mapping);
+      if (mappedDeterminants.isEmpty()) {
+        continue;
+      }
 
-    // Early return if invalid cases
-    if (keyInputIndices.isEmpty()
-        || columnInputIndices.isEmpty()) {
-      return false;
+      // Map only the dependent columns that can be mapped
+      ImmutableBitSet mappedDependents = mapAvailableCols(dependents, mapping);
+      if (!mappedDependents.isEmpty()) {
+        outputFdSet.addFD(mappedDeterminants, mappedDependents);
+      }
     }
+  }
 
-    // Currently only supports multiple (keyInputIndices) to one 
(columnInputIndices)
-    // dependency detection
-    for (Integer keyRef : keyInputIndices) {
-      if (Boolean.FALSE.equals(
-          mq.determines(rel.getInput(0), keyRef,
-          columnInputIndices.nextSetBit(0)))) {
-        return false;
+  /**
+   * Maps all columns in the set. Returns empty set if any column cannot be 
mapped.
+   */
+  private static ImmutableBitSet mapAllCols(
+      ImmutableBitSet columns, Mappings.TargetMapping mapping) {
+    ImmutableBitSet.Builder builder = ImmutableBitSet.builder();
+    for (int col : columns) {
+      if (col < 0 || col >= mapping.getSourceCount()) {
+        return ImmutableBitSet.of();
+      }
+      int mappedCol = mapping.getTargetOpt(col);
+      if (mappedCol >= 0) {
+        builder.set(mappedCol);
+      } else {
+        return ImmutableBitSet.of();
       }
     }
-
-    return true;
+    return builder.build();
   }
 
   /**
-   * determinesImpl2is similar to determinesImpl, but it doesn't need to 
handle the
-   * mapping between output and input columns.
+   * Maps only the columns that can be mapped, ignoring unmappable ones.
    */
-  private static @Nullable Boolean determinesImpl2(RelNode rel, 
RelMetadataQuery mq,
-      int key, int column) {
-    if (preCheck(rel, key, column)) {
-      return true;
+  private static ImmutableBitSet mapAvailableCols(
+      ImmutableBitSet columns, Mappings.TargetMapping mapping) {
+    ImmutableBitSet.Builder builder = ImmutableBitSet.builder();
+    for (int col : columns) {
+      if (col < 0 || col >= mapping.getSourceCount()) {
+        continue;
+      }
+      int mappedCol = mapping.getTargetOpt(col);
+      if (mappedCol >= 0) {
+        builder.set(mappedCol);
+      }
+    }
+    return builder.build();
+  }
+
+  private static FunctionalDependencySet getAggregateFD(Aggregate rel, 
RelMetadataQuery mq) {
+    FunctionalDependencySet fdSet = new FunctionalDependencySet();
+
+    FunctionalDependencySet inputFdSet = 
mq.getFunctionalDependencies(rel.getInput());
+
+    // Group set columns in the output
+    ImmutableBitSet groupSet = rel.getGroupSet();
+
+    // 1. Preserve input FDs that only involve group columns
+    for (FunctionalDependency inputFd : inputFdSet.getFDs()) {
+      ImmutableBitSet determinants = inputFd.getDeterminants();
+      ImmutableBitSet dependents = inputFd.getDependents();
+
+      // Only preserve if both determinants and dependents are within group 
columns
+      if (groupSet.contains(determinants) && groupSet.contains(dependents)) {
+        fdSet.addFD(determinants, dependents);
+      }
     }
 
-    if (rel instanceof TableScan) {
-      TableScan tableScan = (TableScan) rel;
-      RelOptTable table = tableScan.getTable();
-      List<ImmutableBitSet> keys = table.getKeys();
-      return keys != null
-          && keys.size() == 1
-          && keys.get(0).equals(ImmutableBitSet.of(column));
-    } else if (rel instanceof Join) {
-      Join join = (Join) rel;
-      // TODO Considering column mapping based on equality conditions in join
-      int leftFieldCnt = join.getLeft().getRowType().getFieldCount();
-      if (key < leftFieldCnt && column < leftFieldCnt) {
-        return mq.determines(join.getLeft(), key, column);
-      } else if (key >= leftFieldCnt && column >= leftFieldCnt) {
-        return mq.determines(join.getRight(), key - leftFieldCnt, column - 
leftFieldCnt);
+    // 2. Group keys determine all aggregate columns
+    if (!groupSet.isEmpty() && !rel.getAggCallList().isEmpty()) {
+      for (int i = rel.getGroupCount(); i < rel.getRowType().getFieldCount(); 
i++) {
+        fdSet.addFD(groupSet, ImmutableBitSet.of(i));
       }
-      return false;
-    } else if (rel instanceof Correlate) {
-      // TODO Support Correlate.
-      return false;
-    } else if (rel instanceof SetOp) {
-      // TODO Support SetOp
-      return false;
     }
 
-    return mq.determines(rel.getInput(0), key, column);
+    return fdSet;
   }
 
-  private static Boolean preCheck(RelNode rel, int key, int column) {
-    verifyIndex(rel, key, column);
+  private static FunctionalDependencySet getJoinFD(Join rel, RelMetadataQuery 
mq) {
+    FunctionalDependencySet leftFdSet = 
mq.getFunctionalDependencies(rel.getLeft());
+    FunctionalDependencySet rightFdSet = 
mq.getFunctionalDependencies(rel.getRight());
+
+    int leftFieldCount = rel.getLeft().getRowType().getFieldCount();
+    JoinRelType joinType = rel.getJoinType();
+
+    switch (joinType) {
+    case INNER:
+      // Inner join: preserve all FDs and derive cross-table dependencies
+      FunctionalDependencySet innerJoinFdSet
+          = leftFdSet.union(shiftFdSet(rightFdSet, leftFieldCount));
+      deriveTransitiveFDs(rel, innerJoinFdSet, leftFieldCount);
+      return innerJoinFdSet;
+    case LEFT:
+      // Left join: preserve left FDs, right FDs may be invalidated by NULLs
+      FunctionalDependencySet leftJoinFdSet = new 
FunctionalDependencySet(leftFdSet.getFDs());
+      deriveTransitiveFDs(rel, leftJoinFdSet, leftFieldCount);
+      return leftJoinFdSet;
+    case RIGHT:
+      // Right join: preserve right FDs, left FDs may be invalidated by NULLs
+      FunctionalDependencySet shiftedRightFdSet = shiftFdSet(rightFdSet, 
leftFieldCount);

Review Comment:
   sure, better safe than wrong



-- 
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]

Reply via email to