add shadow expression for ComparisonExpression

Author: shizy818

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/QueryPlanner.java

@@ -30,6 +30,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.analyzer.NodeRef;
 import org.apache.iotdb.db.queryengine.plan.relational.analyzer.RelationType;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.ir.GapFillStartAndEndTimeExtractVisitor;
+import org.apache.iotdb.db.queryengine.plan.relational.planner.ir.PredicateWithUncorrelatedScalarSubqueryReconstructor;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationNode;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.AggregationNode.Aggregation;
 import org.apache.iotdb.db.queryengine.plan.relational.planner.node.FilterNode;
@@ -763,9 +764,13 @@ private PlanBuilder filter(PlanBuilder subPlan, Expression predicate, Node node)
     }
 
     subPlan = subqueryPlanner.handleSubqueries(subPlan, predicate, analysis.getSubqueries(node));
-    return subPlan.withNewRoot(
-        new FilterNode(
-            queryIdAllocator.genPlanNodeId(), subPlan.getRoot(), subPlan.rewrite(predicate)));
+    PlanBuilder planBuilder =
+        subPlan.withNewRoot(
+            new FilterNode(
+                queryIdAllocator.genPlanNodeId(), subPlan.getRoot(), subPlan.rewrite(predicate)));
+    PredicateWithUncorrelatedScalarSubqueryReconstructor.getInstance()
+        .clearShadowExpression(predicate);
+    return planBuilder;
   }
 
   private PlanBuilder aggregate(PlanBuilder subPlan, QuerySpecification node) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/planner/ir/PredicateWithUncorrelatedScalarSubqueryReconstructor.java

@@ -34,6 +34,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.DoubleLiteral;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Expression;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.FunctionCall;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Identifier;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Literal;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LogicalExpression;
@@ -79,21 +80,19 @@ public void reconstructPredicateWithUncorrelatedScalarSubquery(
       ComparisonExpression comparisonExpression = (ComparisonExpression) expression;
       Expression left = comparisonExpression.getLeft();
       Expression right = comparisonExpression.getRight();
-      if (left instanceof Identifier && right instanceof SubqueryExpression) {
+      if ((left instanceof Identifier || left instanceof FunctionCall)
+          && right instanceof SubqueryExpression) {
         Optional<Literal> result =
             fetchUncorrelatedSubqueryResultForPredicate(
                 context, (SubqueryExpression) right, analysis.getWith());
         // If the subquery result is not present, we cannot reconstruct the predicate.
-        if (result.isPresent()) {
-          right = result.get();
-        }
-      } else if (right instanceof Identifier && left instanceof SubqueryExpression) {
+        result.ifPresent(comparisonExpression::setShadowRight);
+      } else if ((right instanceof Identifier || right instanceof FunctionCall)
+          && left instanceof SubqueryExpression) {
         Optional<Literal> result =
             fetchUncorrelatedSubqueryResultForPredicate(
                 context, (SubqueryExpression) left, analysis.getWith());
-        if (result.isPresent()) {
-          left = result.get();
-        }
+        result.ifPresent(comparisonExpression::setShadowLeft);
       }
       comparisonExpression.setLeft(left);
       comparisonExpression.setRight(right);
@@ -215,6 +214,23 @@ public Optional<Literal> fetchUncorrelatedSubqueryResultForPredicate(
     return Optional.empty();
   }
 
+  public void clearShadowExpression(Expression expression) {
+    if (expression instanceof LogicalExpression) {
+      LogicalExpression logicalExpression = (LogicalExpression) expression;
+      for (Expression term : logicalExpression.getTerms()) {
+        clearShadowExpression(term);
+      }
+    } else if (expression instanceof NotExpression) {
+      NotExpression notExpression = (NotExpression) expression;
+      clearShadowExpression(notExpression.getValue());
+    } else if (expression instanceof ComparisonExpression) {
+      ComparisonExpression comparisonExpression = (ComparisonExpression) expression;
+      comparisonExpression.clearShadow();
+      clearShadowExpression(comparisonExpression.getLeft());
+      clearShadowExpression(comparisonExpression.getRight());
+    }
+  }
+
   private static class PredicateWithUncorrelatedScalarSubqueryReconstructorHolder {
     private static PredicateWithUncorrelatedScalarSubqueryReconstructor INSTANCE =
         new PredicateWithUncorrelatedScalarSubqueryReconstructor();

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/sql/ast/ComparisonExpression.java

@@ -97,7 +97,9 @@ public Operator negate() {
 
   private final Operator operator;
   private Expression left;
+  private Expression shadowLeft;
   private Expression right;
+  private Expression shadowRight;
 
   public ComparisonExpression(Operator operator, Expression left, Expression right) {
     super(null);
@@ -127,11 +129,11 @@ public Operator getOperator() {
   }
 
   public Expression getLeft() {
-    return left;
+    return shadowLeft != null ? shadowLeft : left;
   }
 
   public Expression getRight() {
-    return right;
+    return shadowRight != null ? shadowRight : right;
   }
 
   public void setLeft(Expression left) {
@@ -152,6 +154,22 @@ public List<Node> getChildren() {
     return ImmutableList.of(left, right);
   }
 
+  // set by unfold of subquery
+  public void setShadowLeft(Expression shadowLeft) {
+    this.shadowLeft = shadowLeft;
+  }
+
+  // set by unfold of subquery
+  public void setShadowRight(Expression shadowRight) {
+    this.shadowRight = shadowRight;
+  }
+
+  // called after the stage is finished
+  public void clearShadow() {
+    this.shadowLeft = null;
+    this.shadowRight = null;
+  }
+
   @Override
   public boolean equals(Object o) {
     if (this == o) {

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/relational/planner/UncorrelatedSubqueryTest.java

@@ -500,4 +500,37 @@ public void testUncorrelatedNotExistsSubquery() {
                                         SINGLE,
                                         tableScan2)))))));
   }
+
+  @Test
+  public void testUncorrelatedHavingSubquery() {
+    String sql =
+        "SELECT min(time) as min FROM table1 group by s1 having min(time) > (select max(time) from table2)";
+    LogicalQueryPlan logicalQueryPlan = planTester.createPlan(sql);
+
+    PlanMatchPattern tableScan =
+        tableScan("testdb.table1", ImmutableList.of("time", "s1"), ImmutableSet.of("time", "s1"));
+    PlanMatchPattern agg =
+        aggregation(
+            singleGroupingSet("s1"),
+            ImmutableMap.of(
+                Optional.of("min"), aggregationFunction("min", ImmutableList.of("time"))),
+            ImmutableList.of(),
+            Optional.empty(),
+            SINGLE,
+            tableScan);
+
+    Expression filterPredicate =
+        new ComparisonExpression(GREATER_THAN, new SymbolReference("min"), new LongLiteral("1"));
+
+    // Verify full LogicalPlan
+    /*
+     *   └──OutputNode
+     *      └──FilterNode
+     *         ├──ProjectNode
+     *            └──Aggregation
+     *               └──TableScanNode
+     */
+
+    assertPlan(logicalQueryPlan, output(filter(filterPredicate, project(agg))));
+  }
 }