[fix][core] Union all cast.

nokiaMS · githubgxll · commit 2fddb55677c8 · 2025-08-27T17:11:10.000+08:00
diff --git a/core/src/main/java/org/apache/calcite/plan/RelOptNode.java b/core/src/main/java/org/apache/calcite/plan/RelOptNode.java
@@ -17,6 +17,7 @@
 package org.apache.calcite.plan;
 
 import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.sql.SqlOperator;
 
 import java.util.List;
 
@@ -65,6 +66,8 @@ public interface RelOptNode {
   // may not even be known at planning time.
   RelDataType getRowType();
 
+  RelDataType getRowTypeWithContext(SqlOperator.CallContext context);
+
   /**
    * Returns a string which describes the relational expression and, unlike
    * {@link #getDigest()}, also includes the identity. Typically returns
diff --git a/core/src/main/java/org/apache/calcite/rel/AbstractRelNode.java b/core/src/main/java/org/apache/calcite/rel/AbstractRelNode.java
@@ -28,12 +28,16 @@
 import org.apache.calcite.rel.core.CorrelationId;
 import org.apache.calcite.rel.hint.Hintable;
 import org.apache.calcite.rel.hint.RelHint;
+import org.apache.calcite.rel.logical.LogicalRepeatUnion;
+import org.apache.calcite.rel.logical.LogicalUnion;
 import org.apache.calcite.rel.metadata.Metadata;
 import org.apache.calcite.rel.metadata.MetadataFactory;
 import org.apache.calcite.rel.metadata.RelMetadataQuery;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rex.RexShuttle;
 import org.apache.calcite.sql.SqlExplainLevel;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.util.Litmus;
 import org.apache.calcite.util.Pair;
 import org.apache.calcite.util.Util;
@@ -168,14 +172,36 @@ protected static <T> T sole(List<T> collection) {
     return litmus.succeed();
   }
 
-  @Override public final RelDataType getRowType() {
+  @Override public RelDataType getRowType() {
     if (rowType == null) {
-      rowType = deriveRowType();
+      if(this instanceof LogicalUnion) {
+        if (((LogicalUnion)this).kind == SqlKind.UNION) {
+          rowType = deriveRowTypeWithContext(SqlOperator.CallContext.IN_UNION);
+        } else {
+          rowType = deriveRowType();
+        }
+      } else {
+        rowType = deriveRowType();
+      }
       assert rowType != null : this;
     }
     return rowType;
   }
 
+  @Override public final RelDataType getRowTypeWithContext(SqlOperator.CallContext context) {
+    if (rowType == null) {
+      rowType = deriveRowTypeWithContext(context);
+      assert rowType != null : this;
+    }
+    return rowType;
+  }
+
+  protected RelDataType deriveRowTypeWithContext(SqlOperator.CallContext context) {
+    // This method is only called if rowType is null, so you don't NEED to
+    // implement it if rowType is always set.
+    throw new UnsupportedOperationException();
+  }
+
   protected RelDataType deriveRowType() {
     // This method is only called if rowType is null, so you don't NEED to
     // implement it if rowType is always set.
diff --git a/core/src/main/java/org/apache/calcite/rel/RelNode.java b/core/src/main/java/org/apache/calcite/rel/RelNode.java
@@ -30,6 +30,7 @@
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexShuttle;
+import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.util.Litmus;
 
 import org.apiguardian.api.API;
@@ -115,6 +116,8 @@ public interface RelNode extends RelOptNode, Cloneable {
    */
   @Override RelDataType getRowType();
 
+  @Override RelDataType getRowTypeWithContext(SqlOperator.CallContext context);
+
   /**
    * Returns the type of the rows expected for an input. Defaults to
    * {@link #getRowType}.
diff --git a/core/src/main/java/org/apache/calcite/rel/core/SetOp.java b/core/src/main/java/org/apache/calcite/rel/core/SetOp.java
@@ -28,6 +28,7 @@
 import org.apache.calcite.rel.hint.RelHint;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.util.Util;
 
 import com.google.common.base.Preconditions;
@@ -113,10 +114,14 @@ public abstract SetOp copy(
   }
 
   @Override protected RelDataType deriveRowType() {
+    return deriveRowTypeWithContext(SqlOperator.CallContext.INVALID);
+  }
+
+  @Override protected RelDataType deriveRowTypeWithContext(SqlOperator.CallContext context) {
     final List<RelDataType> inputRowTypes =
         Util.transform(inputs, RelNode::getRowType);
     final RelDataType rowType =
-        getCluster().getTypeFactory().leastRestrictive(inputRowTypes);
+        getCluster().getTypeFactory().leastRestrictiveWithContext(inputRowTypes, context);
     if (rowType == null) {
       throw new IllegalArgumentException("Cannot compute compatible row type "
           + "for arguments to set op: "
diff --git a/core/src/main/java/org/apache/calcite/rel/type/RelDataTypeFactory.java b/core/src/main/java/org/apache/calcite/rel/type/RelDataTypeFactory.java
@@ -19,6 +19,7 @@
 import org.apache.calcite.avatica.util.TimeUnit;
 import org.apache.calcite.sql.SqlCollation;
 import org.apache.calcite.sql.SqlIntervalQualifier;
+import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.parser.SqlParserPos;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.sql.validate.SqlValidatorUtil;
@@ -208,6 +209,8 @@ RelDataType createTypeWithCharsetAndCollation(
    */
   @Nullable RelDataType leastRestrictive(List<RelDataType> types);
 
+  @Nullable RelDataType leastRestrictiveWithContext(List<RelDataType> types, SqlOperator.CallContext context);
+
   /**
    * Creates a SQL type with no precision or scale.
    *
diff --git a/core/src/main/java/org/apache/calcite/rel/type/RelDataTypeFactoryImpl.java b/core/src/main/java/org/apache/calcite/rel/type/RelDataTypeFactoryImpl.java
@@ -18,6 +18,7 @@
 
 import org.apache.calcite.linq4j.tree.Primitive;
 import org.apache.calcite.sql.SqlCollation;
+import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.type.ArraySqlType;
 import org.apache.calcite.sql.type.JavaToSqlTypeConversionRules;
 import org.apache.calcite.sql.type.MapSqlType;
@@ -212,18 +213,27 @@ private RelDataType createStructType(
         }, nullable);
   }
 
-  @Override public @Nullable RelDataType leastRestrictive(List<RelDataType> types) {
+  @Override public @Nullable RelDataType leastRestrictiveWithContext(List<RelDataType> types, SqlOperator.CallContext context) {
     assert types != null;
     assert types.size() >= 1;
     RelDataType type0 = types.get(0);
     if (type0.isStruct()) {
-      return leastRestrictiveStructuredType(types);
+      return leastRestrictiveStructuredTypeWithContext(types, context);
     }
     return null;
   }
 
+  @Override public @Nullable RelDataType leastRestrictive(List<RelDataType> types) {
+    return leastRestrictiveWithContext(types, SqlOperator.CallContext.INVALID);
+  }
+
   protected @Nullable RelDataType leastRestrictiveStructuredType(
-      final List<RelDataType> types) {
+          final List<RelDataType> types) {
+    return leastRestrictiveStructuredTypeWithContext(types, SqlOperator.CallContext.INVALID);
+  }
+
+  protected @Nullable RelDataType leastRestrictiveStructuredTypeWithContext(
+      final List<RelDataType> types, SqlOperator.CallContext context) {
     final RelDataType type0 = types.get(0);
     // precheck that fieldCount is present
     if (!type0.isStruct()) {
@@ -251,8 +261,8 @@ private RelDataType createStructType(
       // first type?
       final int k = j;
 
-      RelDataType type = leastRestrictive(
-          Util.transform(types, t -> t.getFieldList().get(k).getType())
+      RelDataType type = leastRestrictiveWithContext(
+          Util.transform(types, t -> t.getFieldList().get(k).getType()), context
       );
       if (type == null) {
         return null;
diff --git a/core/src/main/java/org/apache/calcite/sql/SqlOperator.java b/core/src/main/java/org/apache/calcite/sql/SqlOperator.java
@@ -85,7 +85,7 @@ public abstract class SqlOperator {
    */
   public static final int MDX_PRECEDENCE = 200;
 
-  public static enum CallContext{INVALID, IN_VALUES, IN_UPDATE_SOURCE};
+  public static enum CallContext{INVALID, IN_VALUES, IN_UPDATE_SOURCE, IN_UNION};
 
   //~ Instance fields --------------------------------------------------------
 
diff --git a/core/src/main/java/org/apache/calcite/sql/type/ReturnTypes.java b/core/src/main/java/org/apache/calcite/sql/type/ReturnTypes.java
@@ -28,6 +28,7 @@
 import org.apache.calcite.sql.SqlCollation;
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.SqlNodeList;
+import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.sql.SqlOperatorBinding;
 import org.apache.calcite.sql.SqlUtil;
 import org.apache.calcite.sql.validate.SqlValidatorNamespace;
@@ -440,8 +441,14 @@ public static SqlCall stripSeparator(SqlCall call) {
    * @see Glossary#SQL99 SQL:1999 Part 2 Section 9.3
    */
   public static final SqlReturnTypeInference LEAST_RESTRICTIVE =
-      opBinding -> opBinding.getTypeFactory().leastRestrictive(
-          opBinding.collectOperandTypes());
+      opBinding -> {
+    if(opBinding.getOperator().getKind() == SqlKind.UNION) {
+      return opBinding.getTypeFactory().leastRestrictiveWithContext(
+          opBinding.collectOperandTypes(), SqlOperator.CallContext.IN_UNION);
+    }
+        return opBinding.getTypeFactory().leastRestrictive(
+            opBinding.collectOperandTypes());
+      };
 
   /**
    * Returns the same type as the multiset carries. The multiset type returned
diff --git a/core/src/main/java/org/apache/calcite/sql/type/SetopOperandTypeChecker.java b/core/src/main/java/org/apache/calcite/sql/type/SetopOperandTypeChecker.java
@@ -19,6 +19,7 @@
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.sql.SqlCallBinding;
+import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.SqlNode;
 import org.apache.calcite.sql.SqlOperandCountRange;
 import org.apache.calcite.sql.SqlOperator;
@@ -108,8 +109,13 @@ public class SetopOperandTypeChecker implements SqlOperandTypeChecker {
         }
       };
 
+      SqlOperator.CallContext context = SqlOperator.CallContext.INVALID;
+      if(callBinding.getOperator().getKind() == SqlKind.UNION) {
+        context = SqlOperator.CallContext.IN_UNION;
+      }
+
       final RelDataType type =
-          callBinding.getTypeFactory().leastRestrictive(columnIthTypes);
+          callBinding.getTypeFactory().leastRestrictiveWithContext(columnIthTypes, context);
       if (type == null) {
         boolean coerced = false;
         if (callBinding.isTypeCoercionEnabled()) {
diff --git a/core/src/main/java/org/apache/calcite/sql/type/SqlTypeFactoryImpl.java b/core/src/main/java/org/apache/calcite/sql/type/SqlTypeFactoryImpl.java
@@ -23,6 +23,7 @@
 import org.apache.calcite.rel.type.RelDataTypeSystem;
 import org.apache.calcite.sql.SqlCollation;
 import org.apache.calcite.sql.SqlIntervalQualifier;
+import org.apache.calcite.sql.SqlOperator;
 import org.apache.calcite.util.Util;
 
 import org.checkerframework.checker.nullness.qual.Nullable;
@@ -238,6 +239,10 @@ private static void assertBasic(SqlTypeName typeName) {
   }
 
   private @Nullable RelDataType leastRestrictiveSqlType(List<RelDataType> types) {
+    return leastRestrictiveSqlTypeWithContext(types, SqlOperator.CallContext.INVALID);
+  }
+
+  private @Nullable RelDataType leastRestrictiveSqlTypeWithContext(List<RelDataType> types, SqlOperator.CallContext context) {
     RelDataType resultType = null;
     int nullCount = 0;
     int nullableCount = 0;
@@ -294,7 +299,7 @@ private static void assertBasic(SqlTypeName typeName) {
         resultType = type;
         SqlTypeName sqlTypeName = resultType.getSqlTypeName();
         if (sqlTypeName == SqlTypeName.ROW) {
-          return leastRestrictiveStructuredType(types);
+          return leastRestrictiveStructuredTypeWithContext(types, context);
         }
         if (sqlTypeName == SqlTypeName.ARRAY
             || sqlTypeName == SqlTypeName.MULTISET) {
