fix failed test cases

apache · Aug 2, 2024 · 60d2df9 · 60d2df9
1 parent 96b7b26
commit 60d2df9
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 4 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/WindowExpression.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/WindowExpression.java
@@ -124,6 +124,12 @@ public WindowExpression withFunction(Expression function) {
  .orElseGet(() -> new WindowExpression(function, partitionKeys, orderKeys));
  }
 
+ public WindowExpression withFunctionPartitionKeysOrderKeys(Expression function,
+ List<Expression> partitionKeys, List<OrderExpression> orderKeys) {
+ return windowFrame.map(frame -> new WindowExpression(function, partitionKeys, orderKeys, frame))
+ .orElseGet(() -> new WindowExpression(function, partitionKeys, orderKeys));
+ }
+
  @Override
  public boolean nullable() {
  return function.nullable();

diff --git a/.../src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ExpressionVisitors.java b/.../src/main/java/org/apache/doris/nereids/trees/expressions/visitor/ExpressionVisitors.java
@@ -18,10 +18,13 @@
 package org.apache.doris.nereids.trees.expressions.visitor;
 
 import org.apache.doris.nereids.trees.expressions.Expression;
+import org.apache.doris.nereids.trees.expressions.OrderExpression;
 import org.apache.doris.nereids.trees.expressions.WindowExpression;
 import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
 import org.apache.doris.nereids.trees.expressions.functions.agg.MultiDistinction;
 
+import com.google.common.collect.Lists;
+
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -47,11 +50,27 @@ private static class ContainsAggregateChecker extends DefaultExpressionRewriter<
 
  @Override
  public Expression visitWindow(WindowExpression windowExpression, ContainsAggregateCheckerContext context) {
- List<Expression> children = new ArrayList<>();
- for (Expression child : windowExpression.getExpressionsInWindowSpec()) {
- children.add(child.accept(this, context));
+ Expression function = windowExpression.getFunction();
+ List<Expression> functionChildren = new ArrayList<>();
+ for (Expression child : function.children()) {
+ functionChildren.add(child.accept(this, context));
+ }
+ Expression newFunction = function.withChildren(functionChildren);
+
+ List<Expression> partitionKeys = windowExpression.getPartitionKeys();
+ List<Expression> newPartitionKeys = new ArrayList<>();
+ for (Expression key : partitionKeys) {
+ newPartitionKeys.add(key.accept(this, context));
  }
- return windowExpression.withChildren(children);
+
+ List<OrderExpression> orderKeys = windowExpression.getOrderKeys();
+ List<OrderExpression> newOrderKeys = new ArrayList<>();
+ for (OrderExpression orderKey : orderKeys) {
+ newOrderKeys.add((OrderExpression) orderKey
+ .withChildren(Lists.newArrayList(orderKey.child().accept(this, context))));
+ }
+
+ return windowExpression.withFunctionPartitionKeysOrderKeys(newFunction, newPartitionKeys, newOrderKeys);
  }
 
  @Override