morrySnow commented on code in PR #67891:
URL: https://github.com/apache/doris/pull/67891#discussion_r4058920762


##########
fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/ForeignKeyContext.java:
##########
@@ -131,32 +133,113 @@ void putAllPrimaryKeys(TableIf table) {
         for (PrimaryKeyConstraint c : 
Env.getCurrentEnv().getConstraintManager()
                 .getPrimaryKeyConstraints(tableNameInfo)) {
             Set<QualifiedColumn> primaryKey = c.getPrimaryKeys(table).stream()
-                    .map(column -> new QualifiedColumn(table, 
column)).collect(Collectors.toSet());
-            primaryKeys.addAll(primaryKey);
+                    .map(column -> new QualifiedColumn(table, column))
+                    .collect(ImmutableSet.toImmutableSet());
+            declaredPrimaryKeys.add(primaryKey);
         }
     }
 
+    /** Return whether the slots form one declared foreign key from a single 
relation instance. */
     public boolean isForeignKey(Set<Slot> key) {
-        return foreignKeys.containsAll(
-                key.stream().map(s -> 
slotToColumn.get(s)).collect(Collectors.toSet()));
+        if (key.isEmpty()) {
+            return false;
+        }
+        RelationId relationId = slotToRelationId.get(key.iterator().next());
+        Set<QualifiedColumn> columns = key.stream()
+                .map(slotToColumn::get)
+                .collect(Collectors.toSet());
+        return relationId != null
+                && key.stream().allMatch(slot -> 
relationId.equals(slotToRelationId.get(slot)))
+                && key.size() == columns.size()
+                && !columns.contains(null)
+                && constraints.stream().anyMatch(constraint -> 
constraint.keySet().equals(columns));
     }
 
+    /** Return whether the slots form a complete primary key whose relation 
proof is still active. */
     public boolean isPrimaryKey(Set<Slot> key) {
-        return primaryKeys.containsAll(
-                key.stream().map(s -> 
slotToColumn.get(s)).collect(Collectors.toSet()));
+        if (key.isEmpty()) {
+            return false;
+        }
+        PrimaryKeyProof proof = 
slotToPrimaryKeyProof.get(key.iterator().next());
+        if (proof == null || key.stream().anyMatch(slot -> 
slotToPrimaryKeyProof.get(slot) != proof)) {
+            return false;
+        }
+        Set<QualifiedColumn> columns = key.stream()
+                .map(slotToColumn::get)
+                .collect(Collectors.toSet());
+        return key.size() == columns.size()
+                && !columns.contains(null)
+                && proof.columns.equals(columns);
     }
 
-    void putSlot(SlotReference slot, TableIf table) {
-        if (!slot.getOriginalColumn().isPresent()) {
-            return;
+    void putSlots(LogicalCatalogRelation relation, TableIf table) {
+        Map<QualifiedColumn, Slot> columnToSlot = new HashMap<>();
+        for (Slot slot : relation.getOutput()) {
+            if (!(slot instanceof SlotReference) || !((SlotReference) 
slot).getOriginalColumn().isPresent()) {
+                continue;
+            }
+            Column column = ((SlotReference) slot).getOriginalColumn().get();
+            QualifiedColumn qualifiedColumn = new QualifiedColumn(table, 
column);
+            slotToColumn.put(slot, qualifiedColumn);
+            slotToRelationId.put(slot, relation.getRelationId());
+            columnToSlot.put(qualifiedColumn, slot);
+        }
+
+        for (Set<QualifiedColumn> declaredPrimaryKey : declaredPrimaryKeys) {
+            if (!columnToSlot.keySet().containsAll(declaredPrimaryKey)) {
+                continue;
+            }
+            Set<Slot> primaryKey = declaredPrimaryKey.stream()
+                    .map(columnToSlot::get)

Review Comment:
   will update ConstraintManager to restrict create more than one PK on a table



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


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to