Refactor aggregation function type checks to improve input validation

Author: CoollZzz

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/analyze/ExpressionTypeAnalyzer.java

@@ -361,50 +361,6 @@ public TSDataType visitFunctionExpression(
       }
 
       if (functionExpression.isBuiltInAggregationFunctionExpression()) {
-        // Additional type check for multi-input aggregation functions
-        String funcName = functionExpression.getFunctionName().toLowerCase();
-        if (funcName.equals(SqlConstant.CORR)
-            || funcName.equals(SqlConstant.COVAR_POP)
-            || funcName.equals(SqlConstant.COVAR_SAMP)
-            || funcName.equals(SqlConstant.REGR_SLOPE)
-            || funcName.equals(SqlConstant.REGR_INTERCEPT)) {
-          // Check both input parameters are numeric or timestamp
-          if (inputExpressions.size() >= 1) {
-            TSDataType firstInputType = expressionTypes.get(NodeRef.of(inputExpressions.get(0)));
-            if (firstInputType != null
-                && !firstInputType.isNumeric()
-                && firstInputType != TSDataType.TIMESTAMP) {
-              throw new SemanticException(
-                  String.format(
-                      "Aggregate functions [%s] only support numeric data types [INT32, INT64, FLOAT, DOUBLE, TIMESTAMP]",
-                      functionExpression.getFunctionName().toUpperCase()));
-            }
-          }
-          if (inputExpressions.size() >= 2) {
-            TSDataType secondInputType = expressionTypes.get(NodeRef.of(inputExpressions.get(1)));
-            if (secondInputType != null
-                && !secondInputType.isNumeric()
-                && secondInputType != TSDataType.TIMESTAMP) {
-              throw new SemanticException(
-                  String.format(
-                      "Aggregate functions [%s] only support numeric data types [INT32, INT64, FLOAT, DOUBLE, TIMESTAMP]",
-                      functionExpression.getFunctionName().toUpperCase()));
-            }
-          }
-        }
-        if (funcName.equals(SqlConstant.SKEWNESS) || funcName.equals(SqlConstant.KURTOSIS)) {
-          if (!inputExpressions.isEmpty()) {
-            TSDataType firstInputType = expressionTypes.get(NodeRef.of(inputExpressions.get(0)));
-            if (firstInputType != null
-                && !firstInputType.isNumeric()
-                && firstInputType != TSDataType.TIMESTAMP) {
-              throw new SemanticException(
-                  String.format(
-                      "Aggregate functions [%s] only support numeric data types [INT32, INT64, FLOAT, DOUBLE, TIMESTAMP]",
-                      functionExpression.getFunctionName().toUpperCase()));
-            }
-          }
-        }
 
         return setExpressionType(
             functionExpression,
@@ -587,15 +543,20 @@ private TSDataType getInputExpressionTypeForAggregation(
       case SqlConstant.VARIANCE:
       case SqlConstant.VAR_POP:
       case SqlConstant.VAR_SAMP:
+      case SqlConstant.SKEWNESS:
+      case SqlConstant.KURTOSIS:
+      case SqlConstant.MAX_BY:
+      case SqlConstant.MIN_BY:
+        return expressionTypes.get(NodeRef.of(inputExpressions.get(0)));
       case SqlConstant.CORR:
       case SqlConstant.COVAR_POP:
       case SqlConstant.COVAR_SAMP:
       case SqlConstant.REGR_SLOPE:
       case SqlConstant.REGR_INTERCEPT:
-      case SqlConstant.SKEWNESS:
-      case SqlConstant.KURTOSIS:
-      case SqlConstant.MAX_BY:
-      case SqlConstant.MIN_BY:
+        TypeInferenceUtils.verifyIsAggregationDataTypeMatchedForBothInputs(
+            aggregateFunctionName,
+            expressionTypes.get(NodeRef.of(inputExpressions.get(0))),
+            expressionTypes.get(NodeRef.of(inputExpressions.get(1))));
         return expressionTypes.get(NodeRef.of(inputExpressions.get(0)));
       default:
         throw new IllegalArgumentException(

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/utils/TypeInferenceUtils.java

@@ -200,23 +200,6 @@ private static void verifyIsAggregationDataTypeMatched(String aggrFuncName, TSDa
             "Aggregate functions [AVG, SUM, EXTREME, STDDEV, STDDEV_POP, STDDEV_SAMP, "
                 + "VARIANCE, VAR_POP, VAR_SAMP] only support "
                 + "numeric data types [INT32, INT64, FLOAT, DOUBLE]");
-      case SqlConstant.CORR:
-      case SqlConstant.COVAR_POP:
-      case SqlConstant.COVAR_SAMP:
-        if (dataType.isNumeric() || TSDataType.TIMESTAMP.equals(dataType)) {
-          return;
-        }
-        throw new SemanticException(
-            "Aggregate functions [CORR, COVAR_POP, COVAR_SAMP] only support "
-                + "numeric data types [INT32, INT64, FLOAT, DOUBLE, TIMESTAMP]");
-      case SqlConstant.REGR_SLOPE:
-      case SqlConstant.REGR_INTERCEPT:
-        if (dataType.isNumeric() || TSDataType.TIMESTAMP.equals(dataType)) {
-          return;
-        }
-        throw new SemanticException(
-            "Aggregate functions [REGR_SLOPE, REGR_INTERCEPT] only support "
-                + "numeric data types [INT32, INT64, FLOAT, DOUBLE, TIMESTAMP]");
       case SqlConstant.SKEWNESS:
       case SqlConstant.KURTOSIS:
         if (dataType.isNumeric() || TSDataType.TIMESTAMP.equals(dataType)) {
@@ -225,6 +208,11 @@ private static void verifyIsAggregationDataTypeMatched(String aggrFuncName, TSDa
         throw new SemanticException(
             "Aggregate functions [SKEWNESS, KURTOSIS] only support "
                 + "numeric data types [INT32, INT64, FLOAT, DOUBLE, TIMESTAMP]");
+      case SqlConstant.CORR:
+      case SqlConstant.COVAR_POP:
+      case SqlConstant.COVAR_SAMP:
+      case SqlConstant.REGR_SLOPE:
+      case SqlConstant.REGR_INTERCEPT:
       case SqlConstant.COUNT:
       case SqlConstant.COUNT_TIME:
       case SqlConstant.MIN_TIME:
@@ -249,6 +237,30 @@ private static void verifyIsAggregationDataTypeMatched(String aggrFuncName, TSDa
     }
   }
 
+  public static void verifyIsAggregationDataTypeMatchedForBothInputs(
+      String aggrFuncName, TSDataType firstDataType, TSDataType secondDataType) {
+    switch (aggrFuncName.toLowerCase()) {
+      case SqlConstant.CORR:
+      case SqlConstant.COVAR_POP:
+      case SqlConstant.COVAR_SAMP:
+      case SqlConstant.REGR_SLOPE:
+      case SqlConstant.REGR_INTERCEPT:
+        if ((firstDataType != null
+                && !firstDataType.isNumeric()
+                && !TSDataType.TIMESTAMP.equals(firstDataType))
+            || (secondDataType != null
+                && !secondDataType.isNumeric()
+                && !TSDataType.TIMESTAMP.equals(secondDataType))) {
+          throw new SemanticException(
+              "Aggregate functions [CORR, COVAR_POP, COVAR_SAMP, REGR_SLOPE, REGR_INTERCEPT] only support "
+                  + "numeric data types [INT32, INT64, FLOAT, DOUBLE, TIMESTAMP]");
+        }
+        return;
+      default:
+        break;
+    }
+  }
+
   /**
    * Bind Type for non-series input Expressions of AggregationFunction and check Semantic
    *