Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,10 @@

package org.apache.doris.nereids.processor.post;

import org.apache.doris.nereids.trees.expressions.ArrayItemReference;
import org.apache.doris.nereids.trees.expressions.ArrayItemReference.ArrayItemSlot;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;

import java.util.HashMap;
Expand All @@ -28,29 +31,38 @@
/**
* collect common expr
*/
public class CommonSubExpressionCollector extends ExpressionVisitor<Integer, Void> {
public class CommonSubExpressionCollector extends ExpressionVisitor<Integer, Boolean> {
public final Map<Integer, Set<Expression>> commonExprByDepth = new HashMap<>();
private final Map<Integer, Set<Expression>> expressionsByDepth = new HashMap<>();

public int collect(Expression expr) {
return expr.accept(this, expr instanceof Lambda);
}

@Override
public Integer visit(Expression expr, Void context) {
public Integer visit(Expression expr, Boolean inLambda) {
if (expr.children().isEmpty()) {
return 0;
}
return collectCommonExpressionByDepth(
expr.children()
.stream()
.map(child -> child.accept(this, context))
.map(child -> child.accept(this, inLambda == null || inLambda || child instanceof Lambda))
.reduce(Math::max)
.map(m -> m + 1)
.orElse(1),
expr
expr,
inLambda == null || inLambda
);
}

private int collectCommonExpressionByDepth(int depth, Expression expr) {
private int collectCommonExpressionByDepth(int depth, Expression expr, boolean inLambda) {
Set<Expression> expressions = getExpressionsFromDepthMap(depth, expressionsByDepth);
if (expressions.contains(expr)) {
// ArrayItemSlot and ArrayItemReference could not be common expressions
// TODO: could not extract common expression when expression contains same lambda expression
// because ArrayItemSlot in Lambda are not same.
if (expressions.contains(expr)
&& !(inLambda && expr.containsType(ArrayItemSlot.class, ArrayItemReference.class))) {
Set<Expression> commonExpression = getExpressionsFromDepthMap(depth, commonExprByDepth);
commonExpression.add(expr);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ private List<List<NamedExpression>> computeMultiLayerProjections(
List<List<NamedExpression>> multiLayers = Lists.newArrayList();
CommonSubExpressionCollector collector = new CommonSubExpressionCollector();
for (Expression expr : projects) {
expr.accept(collector, null);
collector.collect(expr);
}
// use linkedHashMap to make projects order stable
Map<Expression, Alias> aliasMap = new LinkedHashMap<>();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -20,14 +20,24 @@
import org.apache.doris.nereids.processor.post.CommonSubExpressionCollector;
import org.apache.doris.nereids.processor.post.CommonSubExpressionOpt;
import org.apache.doris.nereids.rules.expression.ExpressionRewriteTestHelper;
import org.apache.doris.nereids.trees.expressions.Add;
import org.apache.doris.nereids.trees.expressions.Alias;
import org.apache.doris.nereids.trees.expressions.And;
import org.apache.doris.nereids.trees.expressions.ArrayItemReference;
import org.apache.doris.nereids.trees.expressions.ExprId;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.NamedExpression;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.functions.scalar.ArrayMap;
import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda;
import org.apache.doris.nereids.trees.expressions.literal.Literal;
import org.apache.doris.nereids.trees.expressions.visitor.DefaultExpressionRewriter;
import org.apache.doris.nereids.types.ArrayType;
import org.apache.doris.nereids.types.IntegerType;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;

Expand All @@ -37,27 +47,40 @@
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;

public class CommonSubExpressionTest extends ExpressionRewriteTestHelper {
@Test
public void testExtractCommonExpr() {
List<NamedExpression> exprs = parseProjections("a+b, a+b+1, abs(a+b+1), a");
CommonSubExpressionCollector collector =
new CommonSubExpressionCollector();
CommonSubExpressionCollector collector = new CommonSubExpressionCollector();
exprs.forEach(expr -> collector.visit(expr, null));
System.out.println(collector.commonExprByDepth);
Assertions.assertEquals(2, collector.commonExprByDepth.size());
List<Expression> l1 = collector.commonExprByDepth.get(Integer.valueOf(1))
.stream().collect(Collectors.toList());
List<Expression> l2 = collector.commonExprByDepth.get(Integer.valueOf(2))
.stream().collect(Collectors.toList());
List<Expression> l1 = new ArrayList<>(collector.commonExprByDepth.get(1));
List<Expression> l2 = new ArrayList<>(collector.commonExprByDepth.get(2));
Assertions.assertEquals(1, l1.size());
assertExpression(l1.get(0), "a+b");
Assertions.assertEquals(1, l2.size());
assertExpression(l2.get(0), "a+b+1");
}

@Test
void testLambdaExpression() {
ArrayItemReference ref = new ArrayItemReference("x", new SlotReference(new ExprId(1), "y",
ArrayType.of(IntegerType.INSTANCE), true, ImmutableList.of()));
Expression add = new Add(ref.toSlot(), Literal.of(1));
Expression and = new And(add, add);
ArrayMap arrayMap = new ArrayMap(new Lambda(ImmutableList.of("x"), and, ImmutableList.of(ref)));
List<NamedExpression> exprs = Lists.newArrayList(
new Alias(new ExprId(10000), arrayMap, "c1"),
new Alias(new ExprId(10001), arrayMap, "c2")
);
CommonSubExpressionCollector collector = new CommonSubExpressionCollector();
exprs.forEach(expr -> collector.visit(expr, false));
Assertions.assertEquals(1, collector.commonExprByDepth.size());
Assertions.assertEquals(1, collector.commonExprByDepth.get(4).size());
Assertions.assertEquals(arrayMap, collector.commonExprByDepth.get(4).iterator().next());
}

@Test
public void testMultiLayers() throws Exception {
List<NamedExpression> exprs = parseProjections("a, a+b, a+b+1, abs(a+b+1), a");
Expand All @@ -68,15 +91,14 @@ public void testMultiLayers() throws Exception {
computeMultLayerProjectionsMethod.setAccessible(true);
List<List<NamedExpression>> multiLayers = (List<List<NamedExpression>>) computeMultLayerProjectionsMethod
.invoke(opt, inputSlots, exprs);
System.out.println(multiLayers);
Assertions.assertEquals(3, multiLayers.size());
List<NamedExpression> l0 = multiLayers.get(0);
Assertions.assertEquals(2, l0.size());
Assertions.assertTrue(l0.contains(ExprParser.INSTANCE.parseExpression("a")));
Assertions.assertTrue(l0.get(1) instanceof Alias);
Assertions.assertInstanceOf(Alias.class, l0.get(1));
assertExpression(l0.get(1).child(0), "a+b");
Assertions.assertEquals(multiLayers.get(1).size(), 3);
Assertions.assertEquals(multiLayers.get(2).size(), 5);
Assertions.assertEquals(3, multiLayers.get(1).size());
Assertions.assertEquals(5, multiLayers.get(2).size());
List<NamedExpression> l2 = multiLayers.get(2);
for (int i = 0; i < 5; i++) {
Assertions.assertEquals(exprs.get(i).getExprId().asInt(), l2.get(i).getExprId().asInt());
Expand Down
20 changes: 20 additions & 0 deletions regression-test/suites/nereids_rules_p0/cse/cse.groovy
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,26 @@ suite("cse") {
"""
contains("l1([k1#0, d#1, i#2, (k1 >= i)#9, cast(d as TEXT)#10, unix_timestamp(cast(d as TEXT)#10, '%Y%m%d') AS `unix_timestamp(cast(d as TEXT), '%Y%m%d')`#11])")
}

// cse should not extract expression use for lambda, such as ArrayItemSlot and ArrayItemReference
sql """
drop table if exists array_cse;
"""
sql """
create table array_cse(c1 int, c2 array<varchar(255)>) PROPERTIES ("replication_allocation" = "tag.location.default: 1");
"""
sql """
insert into array_cse values(1, [1,2,3]);
"""
sql """
sync
"""
sql """
SELECT array_map(x-> if(left(x, 5) = '12345', x, left(x, 5)), c2) FROM array_cse;
"""
sql """
SELECT c0, c0 FROM (SELECT ARRAY_MAP(x-> if(left(x, 5), x, left(x, 5)), `c2`) as `c0` FROM array_cse) t
"""

}

Loading