fix: incorrect total row count when folding uncorrelated scalar subquery (#17014)

Author: shizy818

integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/subquery/SubqueryDataSetUtils.java

@@ -110,6 +110,10 @@ public class SubqueryDataSetUtils {
         "INSERT INTO table3(time,device_id,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10) values (2024-09-24T06:14:30.000+00:00,'d01',40,40,40.0,40.0,false,'shanghai_huangpu_red_A_d01_40','shanghai_huangpu_red_A_d01_40',X'cafebabe40',2024-09-24T06:14:00.000+00:00,'2024-09-24')",
         "INSERT INTO table3(time,device_id,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10) values (2024-09-24T06:13:30.000+00:00,'d_null',30,30,30.0,30.0,true,'shanghai_huangpu_red_A_d01_30','shanghai_huangpu_red_A_d01_30',X'cafebabe30',2024-09-24T06:13:00.000+00:00,'2024-09-23')",
         "INSERT INTO table3(time,device_id,s2,s3,s4,s5,s6,s7,s8,s9,s10) values (2024-09-24T06:14:30.000+00:00,'d_null',40,40.0,40.0,false,'shanghai_huangpu_red_A_d01_40','shanghai_huangpu_red_A_d01_40',X'cafebabe40',2024-09-24T06:14:00.000+00:00,'2024-09-24')",
+        // table4
+        "CREATE TABLE IF NOT EXISTS table4(region STRING TAG, remark STRING ATTRIBUTE, name TEXT FIELD)",
+        "INSERT INTO table4 values(1, 'R001', 'good', 'Mary')",
+        "INSERT INTO table4 values(2, 'R002', 'bad', 'Tom')",
         "FLUSH",
         "CLEAR ATTRIBUTE CACHE",
       };

integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/subquery/uncorrelated/IoTDBUncorrelatedScalarSubqueryIT.java

@@ -265,6 +265,11 @@ public void testScalarSubqueryAfterComparisonLegalityCheck() {
         "701: Scalar sub-query has returned multiple rows.",
         DATABASE_NAME);
 
+    tableAssertTestFail(
+        "select s1 from table1 where time = (select time from table4)",
+        "701: Scalar sub-query has returned multiple rows.",
+        DATABASE_NAME);
+
     // Legality check: subquery can not be parsed
     tableAssertTestFail(
         "select s1 from table1 where s1 = (select s1 from)", "mismatched input", DATABASE_NAME);

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

@@ -162,6 +162,9 @@ public Optional<Literal> fetchUncorrelatedSubqueryResultForPredicate(
         return Optional.empty();
       }
 
+      Column column = null;
+      TSDataType dataType = null;
+      int rowCount = 0;
       while (coordinator.getQueryExecution(queryId).hasNextResult()) {
         final Optional<TsBlock> tsBlock;
         try {
@@ -174,40 +177,46 @@ public Optional<Literal> fetchUncorrelatedSubqueryResultForPredicate(
           continue;
         }
         final Column[] columns = tsBlock.get().getValueColumns();
+        // check column count
         checkArgument(columns.length == 1, "Scalar Subquery result should only have one column.");
+        // check row count
+        rowCount += tsBlock.get().getPositionCount();
         checkArgument(
-            tsBlock.get().getPositionCount() == 1 && !tsBlock.get().getColumn(0).isNull(0),
+            rowCount == 1 && !columns[0].isNull(0),
             "Scalar Subquery result should only have one row.");
+        column = columns[0];
 
         // column type
         DatasetHeader datasetHeader = coordinator.getQueryExecution(queryId).getDatasetHeader();
         List<TSDataType> dataTypes = datasetHeader.getRespDataTypes();
         checkArgument(dataTypes.size() == 1, "Scalar Subquery result should only have one column.");
-
-        switch (dataTypes.get(0)) {
-          case INT32:
-          case DATE:
-            return Optional.of(new LongLiteral(Long.toString(columns[0].getInt(0))));
-          case INT64:
-          case TIMESTAMP:
-            return Optional.of(new LongLiteral(Long.toString(columns[0].getLong(0))));
-          case FLOAT:
-            return Optional.of(new DoubleLiteral(Double.toString(columns[0].getFloat(0))));
-          case DOUBLE:
-            return Optional.of(new DoubleLiteral(Double.toString(columns[0].getDouble(0))));
-          case BOOLEAN:
-            return Optional.of(new BooleanLiteral(Boolean.toString(columns[0].getBoolean(0))));
-          case BLOB:
-            return Optional.of(new BinaryLiteral(columns[0].getBinary(0).toString()));
-          case TEXT:
-          case STRING:
-            return Optional.of(new StringLiteral(columns[0].getBinary(0).toString()));
-          default:
-            throw new IllegalArgumentException(
-                String.format(
-                    "Unsupported data type for scalar subquery result: %s",
-                    columns[0].getDataType()));
-        }
+        dataType = dataTypes.get(0);
+      }
+      checkArgument(
+          dataType != null && column != null,
+          "Scalar Subquery result should not get null dataType or null column.");
+      switch (dataType) {
+        case INT32:
+        case DATE:
+          return Optional.of(new LongLiteral(Long.toString(column.getInt(0))));
+        case INT64:
+        case TIMESTAMP:
+          return Optional.of(new LongLiteral(Long.toString(column.getLong(0))));
+        case FLOAT:
+          return Optional.of(new DoubleLiteral(Double.toString(column.getFloat(0))));
+        case DOUBLE:
+          return Optional.of(new DoubleLiteral(Double.toString(column.getDouble(0))));
+        case BOOLEAN:
+          return Optional.of(new BooleanLiteral(Boolean.toString(column.getBoolean(0))));
+        case BLOB:
+          return Optional.of(new BinaryLiteral(column.getBinary(0).toString()));
+        case TEXT:
+        case STRING:
+          return Optional.of(new StringLiteral(column.getBinary(0).toString()));
+        default:
+          throw new IllegalArgumentException(
+              String.format(
+                  "Unsupported data type for scalar subquery result: %s", column.getDataType()));
       }
     } catch (final Throwable throwable) {
       t = throwable;