Skip to content

Commit 77233b5

Browse files
keanji-xkeanji-x
authored
[opt](Nereids) Replace Slot in Each Data Trait Separately (#36886)
To avoid replacing slots in each data trait repeatedly, we split the replace function into four functions and replaced them separately.
1 parent cc14cd5 commit 77233b5

File tree

10 files changed

+129
-54
lines changed

10 files changed

+129
-54
lines changed

‎fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DataTrait.java

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -321,11 +321,21 @@ public void pruneSlots(Set<Slot> outputSlots) {
321321
fdDgBuilder.removeNotContain(outputSlots);
322322
}
323323

324-
public void replace(Map<Slot, Slot> replaceMap) {
324+
public void replaceUniformBy(Map<Slot, Slot> replaceMap) {
325325
uniformSet.replace(replaceMap);
326+
}
327+
328+
public void replaceUniqueBy(Map<Slot, Slot> replaceMap) {
326329
uniqueSet.replace(replaceMap);
330+
}
331+
332+
public void replaceEqualSetBy(Map<Slot, Slot> replaceMap) {
327333
equalSetBuilder.replace(replaceMap);
328334
}
335+
336+
public void replaceFuncDepsBy(Map<Slot, Slot> replaceMap) {
337+
fdDgBuilder.replace(replaceMap);
338+
}
329339
}
330340

331341
static class NestedSet {

‎fe/fe-core/src/main/java/org/apache/doris/nereids/properties/FuncDepsDG.java

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,14 @@ public DGItem(DGItem dgItem) {
6868
this.parents = new HashSet<>(dgItem.parents);
6969
this.children = new HashSet<>(dgItem.children);
7070
}
71+
72+
public void replace(Map<Slot, Slot> replaceMap) {
73+
Set<Slot> newSlots = new HashSet<>();
74+
for (Slot slot : slots) {
75+
newSlots.add(replaceMap.getOrDefault(slot, slot));
76+
}
77+
this.slots = newSlots;
78+
}
7179
}
7280

7381
private final Map<Set<Slot>, Integer> itemMap; // Maps sets of slots to their indices in the dgItems list
@@ -211,6 +219,21 @@ public void addDeps(FuncDepsDG funcDepsDG) {
211219
}
212220
}
213221

222+
public void replace(Map<Slot, Slot> replaceSlotMap) {
223+
for (DGItem item : dgItems) {
224+
item.replace(replaceSlotMap);
225+
}
226+
Map<Set<Slot>, Integer> newItemMap = new HashMap<>();
227+
for (Entry<Set<Slot>, Integer> e : itemMap.entrySet()) {
228+
Set<Slot> key = new HashSet<>();
229+
for (Slot slot : e.getKey()) {
230+
key.add(replaceSlotMap.getOrDefault(slot, slot));
231+
}
232+
newItemMap.put(key, e.getValue());
233+
}
234+
this.itemMap = newItemMap;
235+
}
236+
214237
private DGItem getOrCreateNode(Set<Slot> slots) {
215238
if (!itemMap.containsKey(slots)) {
216239
itemMap.put(slots, dgItems.size());

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/AbstractPlan.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -203,7 +203,7 @@ public LogicalProperties computeLogicalProperties() {
203203
} else {
204204
Supplier<List<Slot>> outputSupplier = Suppliers.memoize(this::computeOutput);
205205
Supplier<DataTrait> fdSupplier = () -> this instanceof LogicalPlan
206-
? ((LogicalPlan) this).computeFuncDeps()
206+
? ((LogicalPlan) this).computeDataTrait()
207207
: DataTrait.EMPTY_TRAIT;
208208
return new LogicalProperties(outputSupplier, fdSupplier);
209209
}

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/BlockFuncDepsPropagation.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
*/
2929
public interface BlockFuncDepsPropagation extends LogicalPlan {
3030
@Override
31-
default DataTrait computeFuncDeps() {
31+
default DataTrait computeDataTrait() {
3232
return DataTrait.EMPTY_TRAIT;
3333
}
3434

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PropagateFuncDeps.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
*/
2929
public interface PropagateFuncDeps extends LogicalPlan {
3030
@Override
31-
default DataTrait computeFuncDeps() {
31+
default DataTrait computeDataTrait() {
3232
if (children().size() == 1) {
3333
// Note when changing function dependencies, we always clone it.
3434
// So it's safe to return a reference

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalExcept.java

Lines changed: 16 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -132,62 +132,42 @@ public ImmutableSet<FdItem> computeFdItems() {
132132
return builder.build();
133133
}
134134

135-
@Override
136-
public void computeUnique(Builder builder) {
137-
builder.addUniqueSlot(child(0).getLogicalProperties().getTrait());
138-
if (qualifier == Qualifier.DISTINCT) {
139-
builder.addUniqueSlot(ImmutableSet.copyOf(getOutput()));
140-
}
135+
Map<Slot, Slot> constructReplaceMapForChild(int index) {
141136
Map<Slot, Slot> replaceMap = new HashMap<>();
142137
List<Slot> output = getOutput();
143138
List<? extends Slot> originalOutputs = regularChildrenOutputs.isEmpty()
144-
? child(0).getOutput()
145-
: regularChildrenOutputs.get(0);
139+
? child(index).getOutput()
140+
: regularChildrenOutputs.get(index);
146141
for (int i = 0; i < output.size(); i++) {
147142
replaceMap.put(originalOutputs.get(i), output.get(i));
148143
}
149-
builder.replace(replaceMap);
144+
return replaceMap;
145+
}
146+
147+
@Override
148+
public void computeUnique(Builder builder) {
149+
builder.addUniqueSlot(child(0).getLogicalProperties().getTrait());
150+
if (qualifier == Qualifier.DISTINCT) {
151+
builder.addUniqueSlot(ImmutableSet.copyOf(getOutput()));
152+
}
153+
builder.replaceUniqueBy(constructReplaceMapForChild(0));
150154
}
151155

152156
@Override
153157
public void computeEqualSet(DataTrait.Builder builder) {
154158
builder.addEqualSet(child(0).getLogicalProperties().getTrait());
155-
Map<Slot, Slot> replaceMap = new HashMap<>();
156-
List<Slot> output = getOutput();
157-
List<? extends Slot> originalOutputs = regularChildrenOutputs.isEmpty()
158-
? child(0).getOutput()
159-
: regularChildrenOutputs.get(0);
160-
for (int i = 0; i < output.size(); i++) {
161-
replaceMap.put(originalOutputs.get(i), output.get(i));
162-
}
163-
builder.replace(replaceMap);
159+
builder.replaceEqualSetBy(constructReplaceMapForChild(0));
164160
}
165161

166162
@Override
167163
public void computeFd(DataTrait.Builder builder) {
168164
builder.addFuncDepsDG(child(0).getLogicalProperties().getTrait());
169-
Map<Slot, Slot> replaceMap = new HashMap<>();
170-
List<Slot> output = getOutput();
171-
List<? extends Slot> originalOutputs = regularChildrenOutputs.isEmpty()
172-
? child(0).getOutput()
173-
: regularChildrenOutputs.get(0);
174-
for (int i = 0; i < output.size(); i++) {
175-
replaceMap.put(originalOutputs.get(i), output.get(i));
176-
}
177-
builder.replace(replaceMap);
165+
builder.replaceFuncDepsBy(constructReplaceMapForChild(0));
178166
}
179167

180168
@Override
181169
public void computeUniform(Builder builder) {
182170
builder.addUniformSlot(child(0).getLogicalProperties().getTrait());
183-
Map<Slot, Slot> replaceMap = new HashMap<>();
184-
List<Slot> output = getOutput();
185-
List<? extends Slot> originalOutputs = regularChildrenOutputs.isEmpty()
186-
? child(0).getOutput()
187-
: regularChildrenOutputs.get(0);
188-
for (int i = 0; i < output.size(); i++) {
189-
replaceMap.put(originalOutputs.get(i), output.get(i));
190-
}
191-
builder.replace(replaceMap);
171+
builder.replaceUniformBy(constructReplaceMapForChild(0));
192172
}
193173
}

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalIntersect.java

Lines changed: 13 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
package org.apache.doris.nereids.trees.plans.logical;
1919

2020
import org.apache.doris.nereids.memo.GroupExpression;
21-
import org.apache.doris.nereids.properties.DataTrait;
2221
import org.apache.doris.nereids.properties.DataTrait.Builder;
2322
import org.apache.doris.nereids.properties.ExprFdItem;
2423
import org.apache.doris.nereids.properties.FdFactory;
@@ -109,22 +108,26 @@ public LogicalIntersect withNewOutputs(List<NamedExpression> newOutputs) {
109108
Optional.empty(), Optional.empty(), children);
110109
}
111110

112-
void replaceSlotInFuncDeps(DataTrait.Builder builder,
113-
List<Slot> originalOutputs, List<Slot> newOutputs) {
111+
Map<Slot, Slot> constructReplaceMap() {
114112
Map<Slot, Slot> replaceMap = new HashMap<>();
115-
for (int i = 0; i < newOutputs.size(); i++) {
116-
replaceMap.put(originalOutputs.get(i), newOutputs.get(i));
113+
for (int i = 0; i < children.size(); i++) {
114+
List<? extends Slot> originOutputs = this.regularChildrenOutputs.size() == children.size()
115+
? child(i).getOutput()
116+
: regularChildrenOutputs.get(i);
117+
for (int j = 0; j < originOutputs.size(); j++) {
118+
replaceMap.put(originOutputs.get(j), getOutput().get(j));
119+
}
117120
}
118-
builder.replace(replaceMap);
121+
return replaceMap;
119122
}
120123

121124
@Override
122125
public void computeUnique(Builder builder) {
123126
for (Plan child : children) {
124127
builder.addUniqueSlot(
125128
child.getLogicalProperties().getTrait());
126-
replaceSlotInFuncDeps(builder, child.getOutput(), getOutput());
127129
}
130+
builder.replaceUniqueBy(constructReplaceMap());
128131
if (qualifier == Qualifier.DISTINCT) {
129132
builder.addUniqueSlot(ImmutableSet.copyOf(getOutput()));
130133
}
@@ -135,26 +138,26 @@ public void computeUniform(Builder builder) {
135138
for (Plan child : children) {
136139
builder.addUniformSlot(
137140
child.getLogicalProperties().getTrait());
138-
replaceSlotInFuncDeps(builder, child.getOutput(), getOutput());
139141
}
142+
builder.replaceUniformBy(constructReplaceMap());
140143
}
141144

142145
@Override
143146
public void computeEqualSet(Builder builder) {
144147
for (Plan child : children) {
145148
builder.addEqualSet(
146149
child.getLogicalProperties().getTrait());
147-
replaceSlotInFuncDeps(builder, child.getOutput(), getOutput());
148150
}
151+
builder.replaceEqualSetBy(constructReplaceMap());
149152
}
150153

151154
@Override
152155
public void computeFd(Builder builder) {
153156
for (Plan child : children) {
154157
builder.addFuncDepsDG(
155158
child.getLogicalProperties().getTrait());
156-
replaceSlotInFuncDeps(builder, child.getOutput(), getOutput());
157159
}
160+
builder.replaceFuncDepsBy(constructReplaceMap());
158161
}
159162

160163
@Override

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalPlan.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ default LogicalPlan recomputeLogicalProperties() {
6262
* - BlockFDPropagation: clean the fd
6363
* - PropagateFD: propagate the fd
6464
*/
65-
default DataTrait computeFuncDeps() {
65+
default DataTrait computeDataTrait() {
6666
DataTrait.Builder fdBuilder = new DataTrait.Builder();
6767
computeUniform(fdBuilder);
6868
computeUnique(fdBuilder);

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalSubQueryAlias.java

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,7 @@ public void computeUnique(DataTrait.Builder builder) {
171171
for (int i = 0; i < outputs.size(); i++) {
172172
replaceMap.put(child(0).getOutput().get(i), outputs.get(i));
173173
}
174-
builder.replace(replaceMap);
174+
builder.replaceUniqueBy(replaceMap);
175175
}
176176

177177
@Override
@@ -182,7 +182,7 @@ public void computeUniform(DataTrait.Builder builder) {
182182
for (int i = 0; i < outputs.size(); i++) {
183183
replaceMap.put(child(0).getOutput().get(i), outputs.get(i));
184184
}
185-
builder.replace(replaceMap);
185+
builder.replaceUniformBy(replaceMap);
186186
}
187187

188188
@Override
@@ -199,12 +199,18 @@ public void computeEqualSet(DataTrait.Builder builder) {
199199
for (int i = 0; i < outputs.size(); i++) {
200200
replaceMap.put(child(0).getOutput().get(i), outputs.get(i));
201201
}
202-
builder.replace(replaceMap);
202+
builder.replaceEqualSetBy(replaceMap);
203203
}
204204

205205
@Override
206206
public void computeFd(DataTrait.Builder builder) {
207207
builder.addFuncDepsDG(child().getLogicalProperties().getTrait());
208+
Map<Slot, Slot> replaceMap = new HashMap<>();
209+
List<Slot> outputs = getOutput();
210+
for (int i = 0; i < outputs.size(); i++) {
211+
replaceMap.put(child(0).getOutput().get(i), outputs.get(i));
212+
}
213+
builder.replaceFuncDepsBy(replaceMap);
208214
}
209215

210216
public void setRelationId(RelationId relationId) {

‎fe/fe-core/src/test/java/org/apache/doris/nereids/properties/FuncDepsDGTest.java

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,9 @@
2525
import org.junit.jupiter.api.Assertions;
2626
import org.junit.jupiter.api.Test;
2727

28+
import java.util.HashMap;
29+
import java.util.Map;
30+
2831
class FuncDepsDGTest {
2932
@Test
3033
void testBasic() {
@@ -116,4 +119,54 @@ void testPruneTree() {
116119
FuncDeps res = dg.build().findValidFuncDeps(Sets.newHashSet(s1, s4, s3));
117120
Assertions.assertEquals(2, res.size());
118121
}
122+
123+
@Test
124+
void testReplaceTrans() {
125+
FuncDepsDG.Builder dg = new FuncDepsDG.Builder();
126+
Slot s1 = new SlotReference("s1", IntegerType.INSTANCE);
127+
Slot s2 = new SlotReference("s2", IntegerType.INSTANCE);
128+
Slot s3 = new SlotReference("s3", IntegerType.INSTANCE);
129+
Slot s5 = new SlotReference("s5", IntegerType.INSTANCE);
130+
dg.addDeps(Sets.newHashSet(s1), Sets.newHashSet(s2));
131+
dg.addDeps(Sets.newHashSet(s2), Sets.newHashSet(s3));
132+
Map<Slot, Slot> replaceMap = new HashMap<>();
133+
replaceMap.put(s1, s5);
134+
dg.replace(replaceMap);
135+
FuncDeps res = dg.build().findValidFuncDeps(Sets.newHashSet(s5, s3));
136+
System.out.println(res);
137+
Assertions.assertEquals(1, res.size());
138+
}
139+
140+
@Test
141+
void testReplaceCircle() {
142+
FuncDepsDG.Builder dg = new FuncDepsDG.Builder();
143+
Slot s1 = new SlotReference("s1", IntegerType.INSTANCE);
144+
Slot s2 = new SlotReference("s2", IntegerType.INSTANCE);
145+
Slot s5 = new SlotReference("s5", IntegerType.INSTANCE);
146+
dg.addDeps(Sets.newHashSet(s1), Sets.newHashSet(s2));
147+
dg.addDeps(Sets.newHashSet(s2), Sets.newHashSet(s1));
148+
Map<Slot, Slot> replaceMap = new HashMap<>();
149+
replaceMap.put(s1, s5);
150+
dg.replace(replaceMap);
151+
FuncDeps res = dg.build().findValidFuncDeps(Sets.newHashSet(s5, s2));
152+
Assertions.assertEquals(2, res.size());
153+
}
154+
155+
@Test
156+
void testReplaceTree() {
157+
FuncDepsDG.Builder dg = new FuncDepsDG.Builder();
158+
Slot s1 = new SlotReference("s1", IntegerType.INSTANCE);
159+
Slot s2 = new SlotReference("s2", IntegerType.INSTANCE);
160+
Slot s3 = new SlotReference("s3", IntegerType.INSTANCE);
161+
Slot s4 = new SlotReference("s4", IntegerType.INSTANCE);
162+
Slot s5 = new SlotReference("s5", IntegerType.INSTANCE);
163+
dg.addDeps(Sets.newHashSet(s1), Sets.newHashSet(s2));
164+
dg.addDeps(Sets.newHashSet(s2), Sets.newHashSet(s3));
165+
dg.addDeps(Sets.newHashSet(s2), Sets.newHashSet(s4));
166+
Map<Slot, Slot> replaceMap = new HashMap<>();
167+
replaceMap.put(s1, s5);
168+
dg.replace(replaceMap);
169+
FuncDeps res = dg.build().findValidFuncDeps(Sets.newHashSet(s5, s4, s3));
170+
Assertions.assertEquals(2, res.size());
171+
}
119172
}

0 commit comments

Comments
 (0)