This is an automated email from the ASF dual-hosted git repository.

rong pushed a commit to branch nested-operations
in repository https://gitbox.apache.org/repos/asf/iotdb.git

commit 420087949ea55f90d7eda0ac35bad9ad302992a7
Author: Steve Yurong Su <[email protected]>
AuthorDate: Fri Sep 17 12:00:11 2021 +0800

    fix tests
---
 .../apache/iotdb/db/qp/physical/crud/UDTFPlan.java | 44 ++----------------
 .../iotdb/db/query/expression/ResultColumn.java    | 11 +++++
 .../iotdb/db/query/udf/core/layer/DAGBuilder.java  |  6 +++
 .../org/apache/iotdb/db/service/TSServiceImpl.java |  2 +-
 .../iotdb/db/integration/IoTDBNestedQueryIT.java   | 52 ++++++++++++++++++++++
 .../integration/IoTDBUDTFAlignByTimeQueryIT.java   | 16 +++++++
 .../query/dataset/UDTFAlignByTimeDataSetTest.java  | 18 ++++----
 7 files changed, 98 insertions(+), 51 deletions(-)

diff --git 
a/server/src/main/java/org/apache/iotdb/db/qp/physical/crud/UDTFPlan.java 
b/server/src/main/java/org/apache/iotdb/db/qp/physical/crud/UDTFPlan.java
index 38c33cc..57aeff8 100644
--- a/server/src/main/java/org/apache/iotdb/db/qp/physical/crud/UDTFPlan.java
+++ b/server/src/main/java/org/apache/iotdb/db/qp/physical/crud/UDTFPlan.java
@@ -23,16 +23,12 @@ import 
org.apache.iotdb.db.exception.metadata.MetadataException;
 import org.apache.iotdb.db.metadata.PartialPath;
 import org.apache.iotdb.db.qp.logical.Operator;
 import org.apache.iotdb.db.qp.strategy.PhysicalGenerator;
-import org.apache.iotdb.db.query.expression.Expression;
 import org.apache.iotdb.db.query.expression.ResultColumn;
-import org.apache.iotdb.db.query.expression.binary.BinaryExpression;
 import org.apache.iotdb.db.query.expression.unary.FunctionExpression;
-import org.apache.iotdb.db.query.expression.unary.NegationExpression;
 import org.apache.iotdb.db.query.expression.unary.TimeSeriesOperand;
 import org.apache.iotdb.db.query.udf.core.executor.UDTFExecutor;
 import org.apache.iotdb.db.query.udf.service.UDFClassLoaderManager;
 import org.apache.iotdb.db.service.IoTDB;
-import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
 import org.apache.iotdb.tsfile.utils.Pair;
 
 import java.time.ZoneId;
@@ -49,11 +45,8 @@ public class UDTFPlan extends RawDataQueryPlan implements 
UDFPlan {
   protected final ZoneId zoneId;
 
   protected Map<String, UDTFExecutor> expressionName2Executor = new 
HashMap<>();
-  protected Map<Integer, UDTFExecutor> originalOutputColumnIndex2Executor = 
new HashMap<>();
-
   protected Map<Integer, Integer> datasetOutputIndexToResultColumnIndex = new 
HashMap<>();
-
-  protected Map<String, Integer> pathNameToReaderIndex;
+  protected Map<String, Integer> pathNameToReaderIndex = new HashMap<>();
 
   public UDTFPlan(ZoneId zoneId) {
     super();
@@ -72,7 +65,6 @@ public class UDTFPlan extends RawDataQueryPlan implements 
UDFPlan {
     }
     indexedPaths.sort(Comparator.comparing(pair -> pair.left));
 
-    pathNameToReaderIndex = new HashMap<>();
     Set<String> columnForReaderSet = new HashSet<>();
     Set<String> columnForDisplaySet = new HashSet<>();
 
@@ -109,13 +101,8 @@ public class UDTFPlan extends RawDataQueryPlan implements 
UDFPlan {
 
   @Override
   public void constructUdfExecutors(List<ResultColumn> resultColumns) {
-    for (int i = 0; i < resultColumns.size(); ++i) {
-      Expression expression = resultColumns.get(i).getExpression();
-      expression.constructUdfExecutors(expressionName2Executor, zoneId);
-      if (expression instanceof FunctionExpression) {
-        originalOutputColumnIndex2Executor.put(
-            i, expressionName2Executor.get(expression.toString()));
-      }
+    for (ResultColumn resultColumn : resultColumns) {
+      
resultColumn.getExpression().constructUdfExecutors(expressionName2Executor, 
zoneId);
     }
   }
 
@@ -130,31 +117,6 @@ public class UDTFPlan extends RawDataQueryPlan implements 
UDFPlan {
     }
   }
 
-  public TSDataType getOriginalOutputColumnDataType(int originalOutputColumn) {
-    Expression expression = 
resultColumns.get(originalOutputColumn).getExpression();
-    // UDF query
-    if (expression instanceof FunctionExpression) {
-      return getExecutorByOriginalOutputColumnIndex(originalOutputColumn)
-          .getConfigurations()
-          .getOutputDataType();
-    }
-    // arithmetic binary query
-    if (expression instanceof BinaryExpression) {
-      return TSDataType.DOUBLE;
-    }
-    // arithmetic negation query
-    if (expression instanceof NegationExpression) {
-      return getDeduplicatedDataTypes()
-          .get(getReaderIndex(((NegationExpression) 
expression).getExpression().toString()));
-    }
-    // raw query
-    return 
getDeduplicatedDataTypes().get(getReaderIndex(expression.toString()));
-  }
-
-  public UDTFExecutor getExecutorByOriginalOutputColumnIndex(int 
originalOutputColumn) {
-    return originalOutputColumnIndex2Executor.get(originalOutputColumn);
-  }
-
   public ResultColumn getResultColumnByDatasetOutputIndex(int 
datasetOutputIndex) {
     return 
resultColumns.get(datasetOutputIndexToResultColumnIndex.get(datasetOutputIndex));
   }
diff --git 
a/server/src/main/java/org/apache/iotdb/db/query/expression/ResultColumn.java 
b/server/src/main/java/org/apache/iotdb/db/query/expression/ResultColumn.java
index 9c50cc0..b36a56c 100644
--- 
a/server/src/main/java/org/apache/iotdb/db/query/expression/ResultColumn.java
+++ 
b/server/src/main/java/org/apache/iotdb/db/query/expression/ResultColumn.java
@@ -22,6 +22,7 @@ package org.apache.iotdb.db.query.expression;
 import org.apache.iotdb.db.exception.query.LogicalOptimizeException;
 import org.apache.iotdb.db.metadata.PartialPath;
 import org.apache.iotdb.db.qp.utils.WildcardsRemover;
+import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
 
 import java.util.ArrayList;
 import java.util.HashSet;
@@ -70,6 +71,8 @@ public class ResultColumn {
   private final Expression expression;
   private final String alias;
 
+  private TSDataType dataType;
+
   public ResultColumn(Expression expression, String alias) {
     this.expression = expression;
     this.alias = alias;
@@ -136,4 +139,12 @@ public class ResultColumn {
   public String getResultColumnName() {
     return alias != null ? alias : expression.getExpressionString();
   }
+
+  public void setDataType(TSDataType dataType) {
+    this.dataType = dataType;
+  }
+
+  public TSDataType getDataType() {
+    return dataType;
+  }
 }
diff --git 
a/server/src/main/java/org/apache/iotdb/db/query/udf/core/layer/DAGBuilder.java 
b/server/src/main/java/org/apache/iotdb/db/query/udf/core/layer/DAGBuilder.java
index a146826..ee637a0 100644
--- 
a/server/src/main/java/org/apache/iotdb/db/query/udf/core/layer/DAGBuilder.java
+++ 
b/server/src/main/java/org/apache/iotdb/db/query/udf/core/layer/DAGBuilder.java
@@ -22,6 +22,7 @@ package org.apache.iotdb.db.query.udf.core.layer;
 import org.apache.iotdb.db.exception.query.QueryProcessException;
 import org.apache.iotdb.db.qp.physical.crud.UDTFPlan;
 import org.apache.iotdb.db.query.expression.Expression;
+import org.apache.iotdb.db.query.expression.ResultColumn;
 import org.apache.iotdb.db.query.udf.core.reader.LayerPointReader;
 import org.apache.iotdb.tsfile.file.metadata.enums.TSDataType;
 
@@ -89,6 +90,11 @@ public class DAGBuilder {
                   memoryAssigner)
               .constructPointReader();
     }
+
+    for (ResultColumn resultColumn : udtfPlan.getResultColumns()) {
+      
resultColumn.setDataType(expressionDataTypeMap.get(resultColumn.getExpression()));
+    }
+
     return this;
   }
 
diff --git 
a/server/src/main/java/org/apache/iotdb/db/service/TSServiceImpl.java 
b/server/src/main/java/org/apache/iotdb/db/service/TSServiceImpl.java
index 6f6e25c..845ab56 100644
--- a/server/src/main/java/org/apache/iotdb/db/service/TSServiceImpl.java
+++ b/server/src/main/java/org/apache/iotdb/db/service/TSServiceImpl.java
@@ -965,7 +965,7 @@ public class TSServiceImpl implements TSIService.Iface {
         UDTFPlan udtfPlan = (UDTFPlan) plan;
         for (int i = 0; i < paths.size(); i++) {
           respColumns.add(resultColumns.get(i).getResultColumnName());
-          seriesTypes.add(udtfPlan.getOriginalOutputColumnDataType(i));
+          seriesTypes.add(resultColumns.get(i).getDataType());
         }
         break;
       default:
diff --git 
a/server/src/test/java/org/apache/iotdb/db/integration/IoTDBNestedQueryIT.java 
b/server/src/test/java/org/apache/iotdb/db/integration/IoTDBNestedQueryIT.java
index 3366db8..97fe79b 100644
--- 
a/server/src/test/java/org/apache/iotdb/db/integration/IoTDBNestedQueryIT.java
+++ 
b/server/src/test/java/org/apache/iotdb/db/integration/IoTDBNestedQueryIT.java
@@ -55,6 +55,7 @@ public class IoTDBNestedQueryIT {
     Class.forName(Config.JDBC_DRIVER_NAME);
     createTimeSeries();
     generateData();
+    registerUDF();
   }
 
   private static void createTimeSeries() throws MetadataException {
@@ -103,14 +104,37 @@ public class IoTDBNestedQueryIT {
     }
   }
 
+  private static void registerUDF() {
+    try (Connection connection =
+            DriverManager.getConnection(
+                Config.IOTDB_URL_PREFIX + "127.0.0.1:6667/", "root", "root");
+        Statement statement = connection.createStatement()) {
+      statement.execute("create function adder as 
\"org.apache.iotdb.db.query.udf.example.Adder\"");
+    } catch (SQLException throwable) {
+      fail(throwable.getMessage());
+    }
+  }
+
   @AfterClass
   public static void tearDown() throws Exception {
+    deregisterUDF();
     EnvironmentUtils.cleanEnv();
     
IoTDBDescriptor.getInstance().getConfig().setUdfCollectorMemoryBudgetInMB(100);
     
IoTDBDescriptor.getInstance().getConfig().setUdfTransformerMemoryBudgetInMB(100);
     
IoTDBDescriptor.getInstance().getConfig().setUdfReaderMemoryBudgetInMB(100);
   }
 
+  private static void deregisterUDF() {
+    try (Connection connection =
+            DriverManager.getConnection(
+                Config.IOTDB_URL_PREFIX + "127.0.0.1:6667/", "root", "root");
+        Statement statement = connection.createStatement()) {
+      statement.execute("drop function adder");
+    } catch (SQLException throwable) {
+      fail(throwable.getMessage());
+    }
+  }
+
   @Test
   public void testNestedArithmeticExpressions() {
     String sqlStr =
@@ -201,4 +225,32 @@ public class IoTDBNestedQueryIT {
       fail(throwable.getMessage());
     }
   }
+
+  @Test
+  public void testUDFWithMultiInputsInNestedExpressions() {
+    String sqlStr =
+        "select adder(d1.s1, d1.s2), -adder(d2.s1, d2.s2), adder(adder(d1.s1, 
d1.s2), -adder(d2.s1, d2.s2)), adder(adder(d1.s1, d1.s2), adder(d2.s1, d2.s2)) 
from root.vehicle";
+
+    try (Connection connection =
+            DriverManager.getConnection(
+                Config.IOTDB_URL_PREFIX + "127.0.0.1:6667/", "root", "root");
+        Statement statement = connection.createStatement()) {
+      ResultSet resultSet = statement.executeQuery(sqlStr);
+
+      assertEquals(1 + 4, resultSet.getMetaData().getColumnCount());
+
+      int count = 0;
+      while (resultSet.next()) {
+        ++count;
+        assertEquals(count, Double.parseDouble(resultSet.getString(1)), 1e-5);
+        assertEquals(2 * count, Double.parseDouble(resultSet.getString(2)), 
1e-5);
+        assertEquals(-2 * count, Double.parseDouble(resultSet.getString(3)), 
1e-5);
+        assertEquals(0, Double.parseDouble(resultSet.getString(4)), 1e-5);
+        assertEquals(4 * count, Double.parseDouble(resultSet.getString(5)), 
1e-5);
+      }
+      assertEquals(ITERATION_TIMES, count);
+    } catch (SQLException throwable) {
+      fail(throwable.getMessage());
+    }
+  }
 }
diff --git 
a/server/src/test/java/org/apache/iotdb/db/integration/IoTDBUDTFAlignByTimeQueryIT.java
 
b/server/src/test/java/org/apache/iotdb/db/integration/IoTDBUDTFAlignByTimeQueryIT.java
index 013e3fd..7fa0ce3 100644
--- 
a/server/src/test/java/org/apache/iotdb/db/integration/IoTDBUDTFAlignByTimeQueryIT.java
+++ 
b/server/src/test/java/org/apache/iotdb/db/integration/IoTDBUDTFAlignByTimeQueryIT.java
@@ -173,12 +173,28 @@ public class IoTDBUDTFAlignByTimeQueryIT {
 
   @AfterClass
   public static void tearDown() throws Exception {
+    deregisterUDF();
     EnvironmentUtils.cleanEnv();
     
IoTDBDescriptor.getInstance().getConfig().setUdfCollectorMemoryBudgetInMB(100);
     
IoTDBDescriptor.getInstance().getConfig().setUdfTransformerMemoryBudgetInMB(100);
     
IoTDBDescriptor.getInstance().getConfig().setUdfReaderMemoryBudgetInMB(100);
   }
 
+  private static void deregisterUDF() {
+    try (Connection connection =
+            DriverManager.getConnection(
+                Config.IOTDB_URL_PREFIX + "127.0.0.1:6667/", "root", "root");
+        Statement statement = connection.createStatement()) {
+      statement.execute("drop function udf");
+      statement.execute("drop function multiplier");
+      statement.execute("drop function max");
+      statement.execute("drop function terminate");
+      statement.execute("drop function validate");
+    } catch (SQLException throwable) {
+      fail(throwable.getMessage());
+    }
+  }
+
   @Test
   public void queryWithoutValueFilter1() {
     String sqlStr =
diff --git 
a/server/src/test/java/org/apache/iotdb/db/query/dataset/UDTFAlignByTimeDataSetTest.java
 
b/server/src/test/java/org/apache/iotdb/db/query/dataset/UDTFAlignByTimeDataSetTest.java
index 47e7ade..a318265 100644
--- 
a/server/src/test/java/org/apache/iotdb/db/query/dataset/UDTFAlignByTimeDataSetTest.java
+++ 
b/server/src/test/java/org/apache/iotdb/db/query/dataset/UDTFAlignByTimeDataSetTest.java
@@ -171,7 +171,7 @@ public class UDTFAlignByTimeDataSetTest {
         Path path = queryPlan.getPaths().get(i);
         String columnName =
             path == null
-                ? 
queryPlan.getExecutorByOriginalOutputColumnIndex(i).getExpression().toString()
+                ? 
queryPlan.getResultColumns().get(i).getExpression().toString()
                 : path.getFullPath();
         originalIndex2FieldIndex.add(path2Index.get(columnName));
       }
@@ -233,7 +233,7 @@ public class UDTFAlignByTimeDataSetTest {
         Path path = queryPlan.getPaths().get(i);
         String columnName =
             path == null
-                ? 
queryPlan.getExecutorByOriginalOutputColumnIndex(i).getExpression().toString()
+                ? 
queryPlan.getResultColumns().get(i).getExpression().toString()
                 : path.getFullPath();
         originalIndex2FieldIndex.add(path2Index.get(columnName));
       }
@@ -280,7 +280,7 @@ public class UDTFAlignByTimeDataSetTest {
         Path path = queryPlan.getPaths().get(i);
         String columnName =
             path == null
-                ? 
queryPlan.getExecutorByOriginalOutputColumnIndex(i).getExpression().toString()
+                ? 
queryPlan.getResultColumns().get(i).getExpression().toString()
                 : path.getFullPath();
         originalIndex2FieldIndex.add(path2Index.get(columnName));
       }
@@ -339,7 +339,7 @@ public class UDTFAlignByTimeDataSetTest {
         Path path = queryPlan.getPaths().get(i);
         String columnName =
             path == null
-                ? 
queryPlan.getExecutorByOriginalOutputColumnIndex(i).getExpression().toString()
+                ? 
queryPlan.getResultColumns().get(i).getExpression().toString()
                 : path.getFullPath();
         originalIndex2FieldIndex.add(path2Index.get(columnName));
       }
@@ -412,7 +412,7 @@ public class UDTFAlignByTimeDataSetTest {
         Path path = queryPlan.getPaths().get(i);
         String columnName =
             path == null
-                ? 
queryPlan.getExecutorByOriginalOutputColumnIndex(i).getExpression().toString()
+                ? 
queryPlan.getResultColumns().get(i).getExpression().toString()
                 : path.getFullPath();
         originalIndex2FieldIndex.add(path2Index.get(columnName));
       }
@@ -468,7 +468,7 @@ public class UDTFAlignByTimeDataSetTest {
         Path path = queryPlan.getPaths().get(i);
         String columnName =
             path == null
-                ? 
queryPlan.getExecutorByOriginalOutputColumnIndex(i).getExpression().toString()
+                ? 
queryPlan.getResultColumns().get(i).getExpression().toString()
                 : path.getFullPath();
         originalIndex2FieldIndex.add(path2Index.get(columnName));
       }
@@ -533,7 +533,7 @@ public class UDTFAlignByTimeDataSetTest {
         Path path = queryPlan.getPaths().get(i);
         String columnName =
             path == null
-                ? 
queryPlan.getExecutorByOriginalOutputColumnIndex(i).getExpression().toString()
+                ? 
queryPlan.getResultColumns().get(i).getExpression().toString()
                 : path.getFullPath();
         originalIndex2FieldIndex.add(path2Index.get(columnName));
       }
@@ -604,7 +604,7 @@ public class UDTFAlignByTimeDataSetTest {
         Path path = queryPlan.getPaths().get(i);
         String columnName =
             path == null
-                ? 
queryPlan.getExecutorByOriginalOutputColumnIndex(i).getExpression().toString()
+                ? 
queryPlan.getResultColumns().get(i).getExpression().toString()
                 : path.getFullPath();
         originalIndex2FieldIndex.add(path2Index.get(columnName));
       }
@@ -657,7 +657,7 @@ public class UDTFAlignByTimeDataSetTest {
         Path path = queryPlan.getPaths().get(i);
         String columnName =
             path == null
-                ? 
queryPlan.getExecutorByOriginalOutputColumnIndex(i).getExpression().toString()
+                ? 
queryPlan.getResultColumns().get(i).getExpression().toString()
                 : path.getFullPath();
         originalIndex2FieldIndex.add(path2Index.get(columnName));
       }

Reply via email to