[Enhancement] support more pushdown functions (#50559)

Signed-off-by: Murphy <[email protected]> (cherry picked from commit 66e314a)
StarRocks · Sep 2, 2024 · 04d2130 · 04d2130
1 parent f7f5af4
commit 04d2130
Show file tree

Hide file tree

Showing 6 changed files with 289 additions and 122 deletions.
diff --git a/.../src/main/java/com/starrocks/sql/optimizer/rule/tree/prunesubfield/PruneSubfieldRule.java b/.../src/main/java/com/starrocks/sql/optimizer/rule/tree/prunesubfield/PruneSubfieldRule.java
@@ -56,13 +56,21 @@ public class PruneSubfieldRule extends TransformationRule {
             .add(FunctionSet.JSON_LENGTH)
             .build();
 
-    public static final List<String> SUPPORT_FUNCTIONS = ImmutableList.<String>builder()
+    public static final List<String> PRUNE_FUNCTIONS = ImmutableList.<String>builder()
             .add(FunctionSet.MAP_KEYS, FunctionSet.MAP_SIZE)
             .add(FunctionSet.ARRAY_LENGTH)
             .add(FunctionSet.CARDINALITY)
             .addAll(SUPPORT_JSON_FUNCTIONS)
             .build();
 
+    public static final List<String> PUSHDOWN_FUNCTIONS = ImmutableList.<String>builder()
+            .addAll(PRUNE_FUNCTIONS)
+            .add(FunctionSet.ARRAY_CONTAINS, FunctionSet.ARRAY_CONTAINS_ALL)
+            .add(FunctionSet.ARRAY_MAX, FunctionSet.ARRAY_MIN, FunctionSet.ARRAY_SUM, FunctionSet.ARRAY_AVG)
+            .add(FunctionSet.ARRAY_POSITION)
+            .add(FunctionSet.ARRAY_JOIN)
+            .build();
+
     public PruneSubfieldRule() {
         super(RuleType.TF_PRUNE_SUBFIELD, Pattern.create(OperatorType.LOGICAL_PROJECT, OperatorType.PATTERN_SCAN));
     }

diff --git a/...c/main/java/com/starrocks/sql/optimizer/rule/tree/prunesubfield/PushDownSubfieldRule.java b/...c/main/java/com/starrocks/sql/optimizer/rule/tree/prunesubfield/PushDownSubfieldRule.java
@@ -222,7 +222,7 @@ public OptExpression visitLogicalProject(OptExpression optExpression, Context co
             ColumnRefSet allUsedColumns = new ColumnRefSet();
             context.pushDownExprUseColumns.values().forEach(allUsedColumns::union);
 
-            SubfieldExpressionCollector collector = new SubfieldExpressionCollector();
+            SubfieldExpressionCollector collector = SubfieldExpressionCollector.buildPushdownCollector();
             for (ScalarOperator value : lpo.getColumnRefMap().values()) {
                 // check repeat put complex column, like that
                 //      project( columnB: structA.b.c.d )

diff --git a/...ava/com/starrocks/sql/optimizer/rule/tree/prunesubfield/SubfieldAccessPathNormalizer.java b/...ava/com/starrocks/sql/optimizer/rule/tree/prunesubfield/SubfieldAccessPathNormalizer.java
@@ -192,7 +192,7 @@ public Optional<AccessPath> visitCollectionElement(CollectionElementOperator col
 
         @Override
         public Optional<AccessPath> visitCall(CallOperator call, List<Optional<AccessPath>> childrenAccessPaths) {
-            if (!PruneSubfieldRule.SUPPORT_FUNCTIONS.contains(call.getFnName())) {
+            if (!PruneSubfieldRule.PRUNE_FUNCTIONS.contains(call.getFnName())) {
                 return Optional.empty();
             }
 

diff --git a/...java/com/starrocks/sql/optimizer/rule/tree/prunesubfield/SubfieldExpressionCollector.java b/...java/com/starrocks/sql/optimizer/rule/tree/prunesubfield/SubfieldExpressionCollector.java
@@ -15,6 +15,7 @@
 package com.starrocks.sql.optimizer.rule.tree.prunesubfield;
 
 import com.google.common.collect.Lists;
+import com.google.common.collect.Sets;
 import com.starrocks.catalog.Type;
 import com.starrocks.sql.optimizer.operator.scalar.CallOperator;
 import com.starrocks.sql.optimizer.operator.scalar.CollectionElementOperator;
@@ -24,13 +25,14 @@
 import com.starrocks.sql.optimizer.operator.scalar.SubfieldOperator;
 
 import java.util.List;
+import java.util.Set;
 
 /*
  * collect all complex expressions, such as: MAP_KEYS, MAP_VALUES, map['key'], struct.a.b.c ...
  */
 public class SubfieldExpressionCollector extends ScalarOperatorVisitor<Void, Void> {
     private final List<ScalarOperator> complexExpressions = Lists.newArrayList();
-
+    private Set<String> checkFunctions;
     private final boolean enableJsonCollect;
 
     public List<ScalarOperator> getComplexExpressions() {
@@ -43,6 +45,19 @@ public SubfieldExpressionCollector() {
 
     public SubfieldExpressionCollector(boolean enableJsonCollect) {
         this.enableJsonCollect = enableJsonCollect;
+        this.checkFunctions = Sets.newHashSet(PruneSubfieldRule.PRUNE_FUNCTIONS);
+    }
+
+    public static SubfieldExpressionCollector buildPruneCollector() {
+        SubfieldExpressionCollector collector = new SubfieldExpressionCollector();
+        collector.checkFunctions = Sets.newHashSet(PruneSubfieldRule.PRUNE_FUNCTIONS);
+        return collector;
+    }
+
+    public static SubfieldExpressionCollector buildPushdownCollector() {
+        SubfieldExpressionCollector collector = new SubfieldExpressionCollector();
+        collector.checkFunctions = Sets.newHashSet(PruneSubfieldRule.PUSHDOWN_FUNCTIONS);
+        return collector;
     }
 
     @Override
@@ -85,7 +100,7 @@ public Void visitCall(CallOperator call, Void context) {
             return null;
         }
 
-        if (!PruneSubfieldRule.SUPPORT_FUNCTIONS.contains(call.getFnName())) {
+        if (!checkFunctions.contains(call.getFnName())) {
             return visit(call, context);
         }