Resolve type mismatch when WHEN result type differs from ELSE (INT32 vs INT64) (#17415)

Author: shuwenwei

integration-test/src/test/java/org/apache/iotdb/relational/it/db/it/IoTDBCaseWhenThenTableIT.java

@@ -88,6 +88,16 @@ public static void tearDown() throws Exception {
     EnvFactory.getEnv().cleanClusterEnvironment();
   }
 
+  @Test
+  public void testIfWithCastedDefaultType() {
+    String[] retArray = new String[] {"0,", "0,", "2,", "3,"};
+    tableResultSetEqualTest(
+        "select if(s2 > 1, s2, cast(0 as int64)) from table3 limit 4",
+        expectedHeader,
+        retArray,
+        DATABASE);
+  }
+
   @Test
   public void testKind1Basic() {
     String[] retArray = new String[] {"99,", "9999,", "9999,", "999,"};
@@ -161,7 +171,7 @@ public void testKind1OutputTypeRestrict() {
         DATABASE);
     tableAssertTestFail(
         "select case when s1<=0 then 0 when s1>1 then null end from table1",
-        "701: All result types must be the same:",
+        "701: All result types and default result type must be the same:",
         DATABASE);
 
     // TEXT and other types cannot exist at the same time

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

@@ -61,9 +61,12 @@
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SimpleCaseExpression;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.StringLiteral;
 import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.SymbolReference;
+import org.apache.iotdb.db.queryengine.plan.relational.sql.ast.WhenClause;
+import org.apache.iotdb.db.queryengine.plan.relational.type.TypeCoercionUtils;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.LinkedHashMultimap;
 import org.apache.tsfile.read.common.type.BlobType;
 import org.apache.tsfile.read.common.type.DateType;
 import org.apache.tsfile.read.common.type.RowType;
@@ -72,10 +75,14 @@
 import org.apache.tsfile.read.common.type.Type;
 
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
 import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import static com.google.common.base.Preconditions.checkArgument;
@@ -230,37 +237,31 @@ protected Type visitIfExpression(IfExpression node, Context context) {
 
     @Override
     protected Type visitSearchedCaseExpression(SearchedCaseExpression node, Context context) {
-      LinkedHashSet<Type> resultTypes =
-          node.getWhenClauses().stream()
-              .map(
-                  clause -> {
-                    Type operandType = process(clause.getOperand(), context);
-                    if (!operandType.equals(BOOLEAN)) {
-                      throw new SemanticException(
-                          String.format("When clause operand must be boolean: %s", operandType));
-                    }
-                    return setExpressionType(clause, process(clause.getResult(), context));
-                  })
-              .collect(Collectors.toCollection(LinkedHashSet::new));
+      for (WhenClause whenClause : node.getWhenClauses()) {
+        coerceType(
+            context,
+            whenClause.getOperand(),
+            BOOLEAN,
+            (actualType) -> String.format("When clause operand must be boolean: %s", actualType));
+      }
 
-      if (resultTypes.size() != 1) {
-        throw new SemanticException(
-            String.format("All result types must be the same: %s", resultTypes));
+      List<Expression> expressions = new ArrayList<>();
+      for (WhenClause whenClause : node.getWhenClauses()) {
+        expressions.add(whenClause.getResult());
       }
-      Type resultType = resultTypes.iterator().next();
-      node.getDefaultValue()
-          .ifPresent(
-              defaultValue -> {
-                Type defaultType = process(defaultValue, context);
-                if (!defaultType.equals(resultType)) {
-                  throw new SemanticException(
-                      String.format(
-                          "Default result type must be the same as WHEN result types: %s vs %s",
-                          defaultType, resultType));
-                }
-              });
+      node.getDefaultValue().ifPresent(expressions::add);
 
-      return setExpressionType(node, resultType);
+      Type type =
+          coerceToSingleType(
+              context, expressions, "All result types and default result type must be the same");
+      setExpressionType(node, type);
+
+      for (WhenClause whenClause : node.getWhenClauses()) {
+        Type whenClauseType = process(whenClause.getResult(), context);
+        setExpressionType(whenClause, whenClauseType);
+      }
+
+      return type;
     }
 
     @Override
@@ -485,6 +486,66 @@ protected Type visitNode(Node node, Context context) {
       throw new UnsupportedOperationException(
           "Not a valid IR expression: " + node.getClass().getName());
     }
+
+    // Only allow INT32 -> INT64 coercion to suppress some related bugs for now
+    private void coerceType(
+        Context context, Expression expression, Type expectedType, Function<Type, String> message) {
+      Type actualType = process(expression, context);
+      coerceType(expression, expectedType, actualType, message);
+    }
+
+    private Type coerceToSingleType(
+        Context context, List<Expression> expressions, String description) {
+      LinkedHashMultimap<Type, NodeRef<Expression>> typeExpressions = LinkedHashMultimap.create();
+
+      for (Expression expression : expressions) {
+        Type type = process(expression, context);
+        typeExpressions.put(type, NodeRef.of(expression));
+      }
+      Set<Type> types = typeExpressions.keySet();
+      Iterator<Type> iterator = types.iterator();
+      Type superType = iterator.next();
+      if (types.size() == 1) {
+        return superType;
+      }
+      while (iterator.hasNext()) {
+        Type current = iterator.next();
+        if (TypeCoercionUtils.canCoerceTo(current, superType)) {
+          continue;
+        }
+        if (TypeCoercionUtils.canCoerceTo(superType, current)) {
+          superType = current;
+        }
+        throw new SemanticException(String.format(description + ": %s vs %s", superType, current));
+      }
+      for (Type type : types) {
+        Set<NodeRef<Expression>> nodeRefs = typeExpressions.get(type);
+        if (type.equals(superType)) {
+          continue;
+        }
+        if (!TypeCoercionUtils.canCoerceTo(type, superType)) {
+          throw new SemanticException("Cannot coerce type " + type + " to " + superType);
+        }
+        addOrReplaceExpressionType(nodeRefs, superType);
+      }
+      return superType;
+    }
+
+    private void coerceType(
+        Expression expression,
+        Type actualType,
+        Type expectedType,
+        Function<Type, String> errorMsg) {
+      if (!TypeCoercionUtils.canCoerceTo(actualType, expectedType)) {
+        throw new SemanticException(errorMsg.apply(actualType));
+      }
+      setExpressionType(expression, actualType);
+    }
+
+    private void addOrReplaceExpressionType(
+        Collection<NodeRef<Expression>> expressions, Type superType) {
+      expressions.forEach(expression -> expressionTypes.put(expression, superType));
+    }
   }
 
   private static class Context {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/relational/type/TypeCoercionUtils.java

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.queryengine.plan.relational.type;
+
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import org.apache.tsfile.read.common.type.IntType;
+import org.apache.tsfile.read.common.type.LongType;
+import org.apache.tsfile.read.common.type.Type;
+
+import java.util.Collections;
+import java.util.Map;
+import java.util.Set;
+
+public class TypeCoercionUtils {
+
+  private static final Map<Type, Set<Type>> typeCoercionMap;
+
+  static {
+    typeCoercionMap = ImmutableMap.of(IntType.INT32, ImmutableSet.of(LongType.INT64));
+  }
+
+  public static boolean canCoerceTo(Type from, Type to) {
+    if (from.equals(to)) {
+      return true;
+    }
+    return typeCoercionMap.getOrDefault(from, Collections.emptySet()).contains(to);
+  }
+}