Special case casting for UNIONs

Author: evertlammerts

src/duckdb_py/native/python_conversion.cpp

@@ -13,6 +13,36 @@
 
 namespace duckdb {
 
+// Like DefaultCastAs, but handles UNION targets by finding the first compatible member. DefaultCastAs raises a
+// Conversion Error when multiple UNION members have the same type (e.g. UNION(u1 DOUBLE, u2 DOUBLE)), so for UNION
+// targets we resolve the member ourselves.
+static Value CastToTarget(Value val, const LogicalType &target_type) {
+	if (target_type.id() != LogicalTypeId::UNION) {
+		return val.DefaultCastAs(target_type);
+	}
+
+	auto member_count = UnionType::GetMemberCount(target_type);
+	auto &source_type = val.type();
+
+	// First pass: if there's an exact type match we use that
+	for (idx_t i = 0; i < member_count; i++) {
+		if (UnionType::GetMemberType(target_type, i) == source_type) {
+			return Value::UNION(UnionType::CopyMemberTypes(target_type), NumericCast<uint8_t>(i), std::move(val));
+		}
+	}
+
+	// Second pass: if there's a type we can implicitly cast to, we do that
+	for (idx_t i = 0; i < member_count; i++) {
+		auto member_type = UnionType::GetMemberType(target_type, i);
+		Value candidate = val;
+		if (candidate.DefaultTryCastAs(member_type)) {
+			return Value::UNION(UnionType::CopyMemberTypes(target_type), NumericCast<uint8_t>(i), std::move(candidate));
+		}
+	}
+	throw ConversionException("Could not convert value of type %s to %s", source_type.ToString(),
+	                          target_type.ToString());
+}
+
 static Value EmptyMapValue() {
 	auto map_type = LogicalType::MAP(LogicalType::SQLNULL, LogicalType::SQLNULL);
 	return Value::MAP(ListType::GetChildType(map_type), vector<Value>());
@@ -265,7 +295,7 @@ static Value TransformPythonLongToHugeInt(py::handle ele, const LogicalType &tar
 		if (target_type.id() == LogicalTypeId::UNKNOWN || target_type.id() == LogicalTypeId::HUGEINT) {
 			return val;
 		}
-		return val.DefaultCastAs(target_type);
+		return CastToTarget(std::move(val), target_type);
 	}
 	PyErr_Clear();
 
@@ -280,7 +310,7 @@ static Value TransformPythonLongToHugeInt(py::handle ele, const LogicalType &tar
 	if (target_type.id() == LogicalTypeId::UNKNOWN || target_type.id() == LogicalTypeId::UHUGEINT) {
 		return val;
 	}
-	return val.DefaultCastAs(target_type);
+	return CastToTarget(std::move(val), target_type);
 }
 
 void TransformPythonUnsigned(uint64_t value, Value &res) {
@@ -410,7 +440,7 @@ bool TryTransformPythonNumeric(Value &res, py::handle ele, const LogicalType &ta
 		if (!TrySniffPythonNumeric(res, value)) {
 			return false;
 		}
-		res = res.DefaultCastAs(target_type, true);
+		res = CastToTarget(std::move(res), target_type);
 		return true;
 	}
 	}
@@ -516,20 +546,20 @@ struct PythonValueConversion {
 			                          target_type.ToString());
 		}
 		default:
-			result = Value::DOUBLE(val).DefaultCastAs(target_type);
+			result = CastToTarget(Value::DOUBLE(val), target_type);
 			break;
 		}
 	}
 	static void HandleLongAsDouble(Value &result, const LogicalType &target_type, double val) {
 		auto cast_as = target_type.id() == LogicalTypeId::UNKNOWN ? LogicalType::DOUBLE : target_type;
-		result = Value::DOUBLE(val).DefaultCastAs(cast_as);
+		result = CastToTarget(Value::DOUBLE(val), cast_as);
 	}
 	static void HandleLongOverflow(Value &result, const LogicalType &target_type, py::handle ele) {
 		result = TransformPythonLongToHugeInt(ele, target_type);
 	}
 	static void HandleUnsignedBigint(Value &result, const LogicalType &target_type, uint64_t val) {
 		auto cast_as = target_type.id() == LogicalTypeId::UNKNOWN ? LogicalType::UBIGINT : target_type;
-		result = Value::UBIGINT(val).DefaultCastAs(cast_as);
+		result = CastToTarget(Value::UBIGINT(val), cast_as);
 	}
 	static void HandleBigint(Value &res, const LogicalType &target_type, int64_t value) {
 		switch (target_type.id()) {
@@ -545,7 +575,7 @@ struct PythonValueConversion {
 			break;
 		}
 		default:
-			res = Value::BIGINT(value).DefaultCastAs(target_type);
+			res = CastToTarget(Value::BIGINT(value), target_type);
 			break;
 		}
 	}
@@ -555,7 +585,7 @@ struct PythonValueConversion {
 		    (target_type.id() == LogicalTypeId::VARCHAR && !target_type.HasAlias())) {
 			result = Value(value);
 		} else {
-			result = Value(value).DefaultCastAs(target_type);
+			result = CastToTarget(Value(value), target_type);
 		}
 	}
 
@@ -692,7 +722,7 @@ struct PythonVectorConversion {
 			break;
 		}
 		default:
-			FallbackValueConversion(result, result_offset, Value::DOUBLE(val).DefaultCastAs(result.GetType()));
+			FallbackValueConversion(result, result_offset, CastToTarget(Value::DOUBLE(val), result.GetType()));
 			break;
 		}
 	}
@@ -716,7 +746,7 @@ struct PythonVectorConversion {
 			FlatVector::GetData<uint64_t>(result)[result_offset] = value;
 			break;
 		default:
-			FallbackValueConversion(result, result_offset, Value::UBIGINT(value));
+			FallbackValueConversion(result, result_offset, CastToTarget(Value::UBIGINT(value), result.GetType()));
 			break;
 		}
 	}
@@ -787,7 +817,7 @@ struct PythonVectorConversion {
 			break;
 		}
 		default:
-			FallbackValueConversion(result, result_offset, Value::BIGINT(value));
+			FallbackValueConversion(result, result_offset, CastToTarget(Value::BIGINT(value), result.GetType()));
 			break;
 		}
 	}

tests/fast/test_type_conversion.py

@@ -5,10 +5,11 @@
 Issue #330: Integers >64-bit lose precision via double conversion
 """
 
+import numpy as np
 import pytest
 
 import duckdb
-from duckdb.sqltypes import BIGINT, DOUBLE, HUGEINT, UHUGEINT, VARCHAR, DuckDBPyType
+from duckdb.sqltypes import BIGINT, DOUBLE, FLOAT, HUGEINT, UHUGEINT, VARCHAR, DuckDBPyType
 
 
 class TestIssue115FloatToUnion:
@@ -24,22 +25,55 @@ def test_udf_float_to_union_type(self):
         result = conn.sql("SELECT return_float()").fetchone()[0]
         assert result == 1.5
 
+    def test_udf_float_to_ambiguous_union_type(self):
+        """UNION with duplicate DOUBLE members (from np.float64 and float) must not raise ambiguity error."""
+        conn = duckdb.connect()
+        conn.create_function(
+            "return_float",
+            lambda: 1.5,
+            return_type=duckdb.union_type({"u1": VARCHAR, "u2": BIGINT, "u3": DOUBLE, "u4": FLOAT, "u5": DOUBLE}),
+        )
+        result = conn.sql("SELECT return_float()").fetchone()[0]
+        assert result == 1.5
+
     def test_udf_dict_with_float_in_union_struct(self):
-        """Original repro from issue #115."""
+        """Original repro from issue #115 with ambiguous UNION members."""
         conn = duckdb.connect()
 
         arr = [{"a": 1, "b": 1.2}, {"a": 3, "b": 2.4}]
 
         def test():
             return arr
 
-        return_type = DuckDBPyType(list[dict[str, int | float]])
+        return_type = DuckDBPyType(list[dict[str, str | int | np.float64 | np.float32 | float]])
         conn.create_function("test", test, return_type=return_type)
         result = conn.sql("SELECT test()").fetchone()[0]
         assert len(result) == 2
         assert result[0]["b"] == pytest.approx(1.2)
         assert result[1]["b"] == pytest.approx(2.4)
 
+    def test_udf_int_to_ambiguous_union_type(self):
+        """HandleBigint default branch: int into UNION with duplicate BIGINT members."""
+        conn = duckdb.connect()
+        conn.create_function(
+            "return_int",
+            lambda: 42,
+            return_type=duckdb.union_type({"u1": VARCHAR, "u2": BIGINT, "u3": BIGINT}),
+        )
+        result = conn.sql("SELECT return_int()").fetchone()[0]
+        assert result == 42
+
+    def test_udf_string_to_ambiguous_union_type(self):
+        """HandleString default branch: str into UNION with duplicate VARCHAR members."""
+        conn = duckdb.connect()
+        conn.create_function(
+            "return_str",
+            lambda: "hello",
+            return_type=duckdb.union_type({"u1": VARCHAR, "u2": BIGINT, "u3": VARCHAR}),
+        )
+        result = conn.sql("SELECT return_str()").fetchone()[0]
+        assert result == "hello"
+
 
 class TestIssue171DictKeyCaseSensitivity:
     """Dict keys differing only by case must preserve their individual values."""