fix: with in scalar subquery

Author: shizy818

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/Analysis.java

@@ -69,6 +69,7 @@
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Table;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.TableFunctionInvocation;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.WindowFrame;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.With;
 import org.apache.iotdb.db.queryengine.plan.statement.component.FillPolicy;
 
 import com.google.common.collect.ArrayListMultimap;
@@ -123,6 +124,8 @@ public class Analysis implements IAnalysis {
 
   private final Map<NodeRef<Table>, Query> namedQueries = new LinkedHashMap<>();
 
+  private With with;
+
   // map expandable query to the node being the inner recursive reference
   private final Map<NodeRef<Query>, Node> expandableNamedQueries = new LinkedHashMap<>();
 
@@ -282,6 +285,14 @@ public Map<NodeRef<Table>, Query> getNamedQueries() {
     return namedQueries;
   }
 
+  public With getWith() {
+    return with;
+  }
+
+  public void setWith(With with) {
+    this.with = with;
+  }
+
   public boolean isAnalyzed(Expression expression) {
     return expression instanceof DataType || types.containsKey(NodeRef.of(expression));
   }

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/ExpressionAnalyzer.java

@@ -1717,12 +1717,12 @@ private Type analyzePredicateWithSubquery(
       return subqueryType;
     }
 
-    private Type analyzeSubquery(
-        SubqueryExpression node, StackableAstVisitorContext<Context> context) {
+    private Type analyzeSubquery(SubqueryExpression node, StackableAstVisitorContext<Context> ctx) {
       StatementAnalyzer analyzer =
-          statementAnalyzerFactory.apply(node, context.getContext().getCorrelationSupport());
-      Scope subqueryScope = Scope.builder().withParent(context.getContext().getScope()).build();
+          statementAnalyzerFactory.apply(node, ctx.getContext().getCorrelationSupport());
+      Scope subqueryScope = Scope.builder().withParent(ctx.getContext().getScope()).build();
       Scope queryScope = analyzer.analyze(node.getQuery(), subqueryScope);
+      context.getSubQueryTables().put(node.getQuery(), queryScope.getTables());
 
       ImmutableList.Builder<RowType.Field> fields = ImmutableList.builder();
       for (int i = 0; i < queryScope.getRelationType().getAllFieldCount(); i++) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/analyzer/StatementAnalyzer.java

@@ -920,6 +920,7 @@ private Scope analyzeWith(Query node, Optional<Scope> scope) {
 
       // analyze WITH clause
       With with = node.getWith().get();
+      analysis.setWith(with);
       Scope.Builder withScopeBuilder = scopeBuilder(scope);
 
       for (WithQuery withQuery : with.getQueries()) {
@@ -4455,11 +4456,10 @@ private Scope createAndAssignScope(Node node, Optional<Scope> parentScope, List<
 
     private Scope createAndAssignScope(
         Node node, Optional<Scope> parentScope, RelationType relationType) {
-      Scope scope =
-          scopeBuilder(parentScope)
-              .withRelationType(RelationId.of(node), relationType)
-              .withTables(parentScope.get().getTables())
-              .build();
+      Scope.Builder scopeBuilder =
+          scopeBuilder(parentScope).withRelationType(RelationId.of(node), relationType);
+      parentScope.ifPresent(scope -> scopeBuilder.withTables(scope.getTables()));
+      Scope scope = scopeBuilder.build();
 
       analysis.setScope(node, scope);
       return scope;

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

@@ -74,15 +74,15 @@ public class CteMaterializer {
 
   private static final Coordinator coordinator = Coordinator.getInstance();
 
-  public void materializeCTE(MPPQueryContext context, Analysis analysis, Query mainQuery) {
+  public void materializeCTE(Analysis analysis, MPPQueryContext context) {
     analysis
         .getNamedQueries()
         .forEach(
             (tableRef, query) -> {
               Table table = tableRef.getNode();
               if (query.isMaterialized() && !query.isDone()) {
                 CteDataStore dataStore =
-                    fetchCteQueryResult(context, table, query, mainQuery.getWith().orElse(null));
+                    fetchCteQueryResult(context, table, query, analysis.getWith());
                 if (dataStore == null) {
                   // CTE query execution failed. Use inline instead of materialization
                   // in the outer query

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

@@ -158,7 +158,7 @@ public PlanBuilder handleSubqueries(
   private void tryFoldUncorrelatedScalarSubqueryInPredicate(
       Expression expression, MPPQueryContext context) {
     PredicateWithUncorrelatedScalarSubqueryReconstructor.getInstance()
-        .reconstructPredicateWithUncorrelatedScalarSubquery(expression, context);
+        .reconstructPredicateWithUncorrelatedScalarSubquery(context, analysis, expression);
   }
 
   /**

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

@@ -374,7 +374,7 @@ private RelationPlan createRelationPlan(Analysis analysis, PipeEnriched pipeEnri
   private RelationPlan createRelationPlan(Analysis analysis, Query query) {
     // materialize cte if needed
     if (!queryContext.isSubquery()) {
-      CteMaterializer.getInstance().materializeCTE(queryContext, analysis, query);
+      CteMaterializer.getInstance().materializeCTE(analysis, queryContext);
     }
     return getRelationPlanner(analysis).process(query, null);
   }

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

@@ -28,6 +28,7 @@
 import org.apache.iotdb.db.queryengine.plan.Coordinator;
 import org.apache.iotdb.db.queryengine.plan.execution.ExecutionResult;
 import org.apache.iotdb.db.queryengine.plan.planner.LocalExecutionPlanner;
+import org.apache.iotdb.db.queryengine.plan.relational.analyzer.Analysis;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.BinaryLiteral;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.BooleanLiteral;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.ComparisonExpression;
@@ -38,17 +39,22 @@
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LogicalExpression;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.LongLiteral;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.NotExpression;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.Query;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.StringLiteral;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SubqueryExpression;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.With;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.WithQuery;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.parser.SqlParser;
 import org.apache.iotdb.rpc.TSStatusCode;
 
+import com.google.common.collect.ImmutableList;
 import org.apache.tsfile.block.column.Column;
 import org.apache.tsfile.enums.TSDataType;
 import org.apache.tsfile.read.common.block.TsBlock;
 
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkArgument;
 
@@ -59,29 +65,32 @@ public class PredicateWithUncorrelatedScalarSubqueryReconstructor {
   private static final Coordinator coordinator = Coordinator.getInstance();
 
   public void reconstructPredicateWithUncorrelatedScalarSubquery(
-      Expression expression, MPPQueryContext context) {
+      MPPQueryContext context, Analysis analysis, Expression expression) {
     if (expression instanceof LogicalExpression) {
       LogicalExpression logicalExpression = (LogicalExpression) expression;
       for (Expression term : logicalExpression.getTerms()) {
-        reconstructPredicateWithUncorrelatedScalarSubquery(term, context);
+        reconstructPredicateWithUncorrelatedScalarSubquery(context, analysis, term);
       }
     } else if (expression instanceof NotExpression) {
       NotExpression notExpression = (NotExpression) expression;
-      reconstructPredicateWithUncorrelatedScalarSubquery(notExpression.getValue(), context);
+      reconstructPredicateWithUncorrelatedScalarSubquery(
+          context, analysis, notExpression.getValue());
     } else if (expression instanceof ComparisonExpression) {
       ComparisonExpression comparisonExpression = (ComparisonExpression) expression;
       Expression left = comparisonExpression.getLeft();
       Expression right = comparisonExpression.getRight();
       if (left instanceof Identifier && right instanceof SubqueryExpression) {
         Optional<Literal> result =
-            fetchUncorrelatedSubqueryResultForPredicate((SubqueryExpression) right, context);
+            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) {
         Optional<Literal> result =
-            fetchUncorrelatedSubqueryResultForPredicate((SubqueryExpression) left, context);
+            fetchUncorrelatedSubqueryResultForPredicate(
+                context, (SubqueryExpression) left, analysis.getWith());
         if (result.isPresent()) {
           left = result.get();
         }
@@ -97,14 +106,40 @@ public void reconstructPredicateWithUncorrelatedScalarSubquery(
    *     valid result.
    */
   public Optional<Literal> fetchUncorrelatedSubqueryResultForPredicate(
-      SubqueryExpression subqueryExpression, MPPQueryContext context) {
+      MPPQueryContext context, SubqueryExpression subqueryExpression, With with) {
     final long queryId = SessionManager.getInstance().requestQueryId();
     Throwable t = null;
 
     try {
+      Query query = subqueryExpression.getQuery();
+      Query q = query;
+      if (with != null) {
+        List<Identifier> tables =
+            context.getSubQueryTables().getOrDefault(query, ImmutableList.of());
+        List<WithQuery> withQueries =
+            with.getQueries().stream()
+                .filter(
+                    x ->
+                        tables.contains(x.getName())
+                            && !x.getQuery().isMaterialized()
+                            && !x.getQuery().isDone())
+                .collect(Collectors.toList());
+
+        if (!withQueries.isEmpty()) {
+          With w = new With(with.getLocation().orElse(null), with.isRecursive(), withQueries);
+          q =
+              new Query(
+                  Optional.of(w),
+                  query.getQueryBody(),
+                  query.getFill(),
+                  query.getOrderBy(),
+                  query.getOffset(),
+                  query.getLimit());
+        }
+      }
       final ExecutionResult executionResult =
           coordinator.executeForTableModel(
-              subqueryExpression.getQuery(),
+              q,
               relationSqlParser,
               SessionManager.getInstance().getCurrSession(),
               queryId,

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

@@ -105,8 +105,8 @@ private void mockSubquery() {
 
     Mockito.doAnswer(
             (InvocationOnMock invocation) -> {
-              Table table = invocation.getArgument(0);
-              Query query = invocation.getArgument(1);
+              Table table = invocation.getArgument(1);
+              Query query = invocation.getArgument(2);
               List<SelectItem> selectItems =
                   ((QuerySpecification) query.getQueryBody()).getSelect().getSelectItems();
               List<ColumnSchema> columnsSchemas =

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

@@ -75,7 +75,7 @@ private void mockPredicateWithUncorrelatedScalarSubquery() {
         Mockito.spy(new PredicateWithUncorrelatedScalarSubqueryReconstructor());
     Mockito.when(
             predicateWithUncorrelatedScalarSubquery.fetchUncorrelatedSubqueryResultForPredicate(
-                Mockito.any(), Mockito.any()))
+                Mockito.any(), Mockito.any(), Mockito.any()))
         .thenReturn(Optional.of(new LongLiteral("1")));
     PredicateWithUncorrelatedScalarSubqueryReconstructor.setInstance(
         predicateWithUncorrelatedScalarSubquery);