[BugFix]Fix multisearch UDT type loss through UNION (#5145, #5146, #5147)

core/src/main/java/org/opensearch/sql/calcite/utils/OpenSearchTypeFactory.java

@@ -45,6 +45,7 @@
 import org.apache.calcite.sql.SqlCollation;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.sql.type.SqlTypeUtil;
+import org.checkerframework.checker.nullness.qual.Nullable;
 import org.opensearch.sql.calcite.type.AbstractExprRelDataType;
 import org.opensearch.sql.calcite.type.ExprBinaryType;
 import org.opensearch.sql.calcite.type.ExprDateType;
@@ -377,6 +378,44 @@ public static boolean isNumericType(RelDataType fieldType) {
     return false;
   }
 
+  /**
+   * Preserves OpenSearch UDT types through set operations (UNION, INTERSECT, EXCEPT). When all
+   * input types share the same {@link AbstractExprRelDataType} with the same {@link
+   * AbstractExprRelDataType#getUdt() UDT}, the result retains the UDT wrapper instead of being
+   * downgraded to the underlying SQL type (e.g., VARCHAR). This is critical for operations like
+   * multisearch that use UNION ALL, where downstream operators (bin, span) rely on the UDT type to
+   * determine how to process the field. When inputs include non-UDT types or different UDTs, this
+   * method falls back to {@link super#leastRestrictive}.
+   *
+   * @param types the list of input {@link RelDataType} instances to find the least restrictive
+   *     common type for
+   * @return the least restrictive {@link RelDataType} preserving the UDT wrapper when all inputs
+   *     share the same UDT, or {@code null} if no common type exists (as determined by {@link
+   *     super#leastRestrictive})
+   */
+  @Override
+  public @Nullable RelDataType leastRestrictive(List<RelDataType> types) {
+    if (types.size() > 1) {
+      RelDataType first = types.get(0);
+      if (first instanceof AbstractExprRelDataType<?> firstUdt) {
+        boolean allSameUdt =
+            types.stream()
+                .allMatch(
+                    t ->
+                        t instanceof AbstractExprRelDataType<?> udt
+                            && udt.getUdt() == firstUdt.getUdt());
+        if (allSameUdt) {
+          boolean anyNullable = types.stream().anyMatch(RelDataType::isNullable);
+          if (anyNullable && !first.isNullable()) {
+            return firstUdt.createWithNullability(this, true);
+          }
+          return first;
+        }
+      }
+    }
+    return super.leastRestrictive(types);
+  }
+
   /**
    * Checks if the RelDataType represents a time-based field (timestamp, date, or time). Supports
    * both standard SQL time types (including TIMESTAMP, TIMESTAMP_WITH_LOCAL_TIME_ZONE, DATE, TIME,

core/src/test/java/org/opensearch/sql/calcite/utils/OpenSearchTypeFactoryTest.java

@@ -0,0 +1,131 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.calcite.utils;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.TYPE_FACTORY;
+
+import java.util.List;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.junit.jupiter.api.Test;
+import org.opensearch.sql.calcite.type.AbstractExprRelDataType;
+import org.opensearch.sql.calcite.utils.OpenSearchTypeFactory.ExprUDT;
+
+public class OpenSearchTypeFactoryTest {
+
+  @Test
+  public void testLeastRestrictivePreservesUdtWhenAllInputsSameUdt() {
+    RelDataType ts1 = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP);
+    RelDataType ts2 = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP);
+
+    RelDataType result = TYPE_FACTORY.leastRestrictive(List.of(ts1, ts2));
+
+    assertNotNull(result);
+    assertInstanceOf(AbstractExprRelDataType.class, result);
+    assertEquals(ExprUDT.EXPR_TIMESTAMP, ((AbstractExprRelDataType<?>) result).getUdt());
+  }
+
+  @Test
+  public void testLeastRestrictivePreservesUdtForDateType() {
+    RelDataType d1 = TYPE_FACTORY.createUDT(ExprUDT.EXPR_DATE);
+    RelDataType d2 = TYPE_FACTORY.createUDT(ExprUDT.EXPR_DATE);
+
+    RelDataType result = TYPE_FACTORY.leastRestrictive(List.of(d1, d2));
+
+    assertNotNull(result);
+    assertInstanceOf(AbstractExprRelDataType.class, result);
+    assertEquals(ExprUDT.EXPR_DATE, ((AbstractExprRelDataType<?>) result).getUdt());
+  }
+
+  @Test
+  public void testLeastRestrictivePreservesUdtForThreeInputs() {
+    RelDataType ts1 = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP);
+    RelDataType ts2 = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP);
+    RelDataType ts3 = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP);
+
+    RelDataType result = TYPE_FACTORY.leastRestrictive(List.of(ts1, ts2, ts3));
+
+    assertNotNull(result);
+    assertInstanceOf(AbstractExprRelDataType.class, result);
+    assertEquals(ExprUDT.EXPR_TIMESTAMP, ((AbstractExprRelDataType<?>) result).getUdt());
+  }
+
+  @Test
+  public void testLeastRestrictiveReturnsNullableWhenAnyInputIsNullable() {
+    RelDataType nonNullable = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP, false);
+    RelDataType nullable = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP, true);
+
+    RelDataType result = TYPE_FACTORY.leastRestrictive(List.of(nonNullable, nullable));
+
+    assertNotNull(result);
+    assertInstanceOf(AbstractExprRelDataType.class, result);
+    assertEquals(ExprUDT.EXPR_TIMESTAMP, ((AbstractExprRelDataType<?>) result).getUdt());
+    assertTrue(result.isNullable());
+  }
+
+  @Test
+  public void testLeastRestrictiveReturnsNullableWhenFirstNullableSecondNot() {
+    RelDataType nullable = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP, true);
+    RelDataType nonNullable = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP, false);
+
+    RelDataType result = TYPE_FACTORY.leastRestrictive(List.of(nullable, nonNullable));
+
+    assertNotNull(result);
+    assertInstanceOf(AbstractExprRelDataType.class, result);
+    assertTrue(result.isNullable());
+  }
+
+  @Test
+  public void testLeastRestrictiveFallsBackForMixedUdtAndNonUdt() {
+    RelDataType udt = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP);
+    RelDataType plain = TYPE_FACTORY.createSqlType(SqlTypeName.VARCHAR);
+
+    RelDataType result = TYPE_FACTORY.leastRestrictive(List.of(udt, plain));
+
+    // Falls back to super.leastRestrictive which may return a plain type or null
+    if (result != null) {
+      assertEquals(SqlTypeName.VARCHAR, result.getSqlTypeName());
+    }
+  }
+
+  @Test
+  public void testLeastRestrictiveFallsBackForDifferentUdts() {
+    RelDataType timestamp = TYPE_FACTORY.createUDT(ExprUDT.EXPR_TIMESTAMP);
+    RelDataType date = TYPE_FACTORY.createUDT(ExprUDT.EXPR_DATE);
+
+    RelDataType result = TYPE_FACTORY.leastRestrictive(List.of(timestamp, date));
+
+    // Different UDTs — falls back to super.leastRestrictive
+    if (result != null) {
+      assertEquals(SqlTypeName.VARCHAR, result.getSqlTypeName());
+    }
+  }
+
+  @Test
+  public void testLeastRestrictiveDelegatesToSuperForSingleType() {
+    RelDataType single = TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER);
+
+    RelDataType result = TYPE_FACTORY.leastRestrictive(List.of(single));
+
+    assertNotNull(result);
+    assertEquals(SqlTypeName.INTEGER, result.getSqlTypeName());
+  }
+
+  @Test
+  public void testLeastRestrictiveDelegatesToSuperForPlainTypes() {
+    RelDataType int1 = TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER);
+    RelDataType int2 = TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER);
+
+    RelDataType result = TYPE_FACTORY.leastRestrictive(List.of(int1, int2));
+
+    assertNotNull(result);
+    assertEquals(SqlTypeName.INTEGER, result.getSqlTypeName());
+  }
+}