feat: support union type for basic types

python/cocoindex/convert.py

@@ -157,6 +157,9 @@ def decode_vector(value: Any) -> Any | None:
 
         return decode_vector
 
+    if src_type_kind == "Union":
+        return lambda value: value[1]
+
     return lambda value: value
 
 

python/cocoindex/tests/test_convert.py

@@ -477,6 +477,24 @@ def test_field_position_cases(
     assert decoder(engine_val) == PythonOrder(**expected_dict)
 
 
+def test_roundtrip_union_simple() -> None:
+    t = int | str | float
+    value = 10.4
+    validate_full_roundtrip(value, t)
+
+
+def test_roundtrip_union_with_active_uuid() -> None:
+    t = str | uuid.UUID | int
+    value = uuid.uuid4().bytes
+    validate_full_roundtrip(value, t)
+
+
+def test_roundtrip_union_with_inactive_uuid() -> None:
+    t = str | uuid.UUID | int
+    value = "5a9f8f6a-318f-4f1f-929d-566d7444a62d" # it's a string
+    validate_full_roundtrip(value, t)
+
+
 def test_roundtrip_ltable() -> None:
     t = list[Order]
     value = [Order("O1", "item1", 10.0), Order("O2", "item2", 20.0)]

python/cocoindex/typing.py

@@ -155,6 +155,7 @@ class AnalyzedTypeInfo:
 
     attrs: dict[str, Any] | None
     nullable: bool = False
+    union_variant_types: list[type] | None = None  # For Union
 
 
 def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
@@ -175,18 +176,6 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         if base_type is Annotated:
             annotations = t.__metadata__
             t = t.__origin__
-        elif base_type is types.UnionType:
-            possible_types = typing.get_args(t)
-            non_none_types = [
-                arg for arg in possible_types if arg not in (None, types.NoneType)
-            ]
-            if len(non_none_types) != 1:
-                raise ValueError(
-                    f"Expect exactly one non-None choice for Union type, but got {len(non_none_types)}: {t}"
-                )
-            t = non_none_types[0]
-            if len(possible_types) > 1:
-                nullable = True
         else:
             break
 
@@ -205,6 +194,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
 
     struct_type: type | None = None
     elem_type: ElementType | None = None
+    union_variant_types: typing.List[ElementType] = None
     key_type: type | None = None
     np_number_type: type | None = None
     if _is_struct_type(t):
@@ -254,6 +244,22 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         args = typing.get_args(t)
         elem_type = (args[0], args[1])
         kind = "KTable"
+    elif base_type is types.UnionType:
+        possible_types = typing.get_args(t)
+        non_none_types = [arg for arg in possible_types if arg not in (None, types.NoneType)]
+
+        if len(non_none_types) == 0:
+            return analyze_type_info(None)
+
+        nullable = len(non_none_types) < len(possible_types)
+
+        if len(non_none_types) == 1:
+            result = analyze_type_info(non_none_types[0])
+            result.nullable = nullable
+            return result
+
+        kind = 'Union'
+        union_variant_types = non_none_types
     elif kind is None:
         dtype_info = DtypeRegistry.get_by_dtype(t)
         if dtype_info is not None:
@@ -286,6 +292,7 @@ def analyze_type_info(t: Any) -> AnalyzedTypeInfo:
         kind=kind,
         vector_info=vector_info,
         elem_type=elem_type,
+        union_variant_types=union_variant_types,
         key_type=key_type,
         struct_type=struct_type,
         np_number_type=np_number_type,
@@ -345,6 +352,13 @@ def _encode_type(type_info: AnalyzedTypeInfo) -> dict[str, Any]:
         encoded_type["element_type"] = _encode_type(elem_type_info)
         encoded_type["dimension"] = type_info.vector_info.dim
 
+    elif type_info.kind == 'Union':
+        if type_info.union_variant_types is None:
+            raise ValueError("Union type must have a variant type list")
+        encoded_type['types'] = [
+            _encode_type(analyze_type_info(typ)) for typ in type_info.union_variant_types
+        ]
+
     elif type_info.kind in TABLE_TYPES:
         if type_info.elem_type is None:
             raise ValueError(f"{type_info.kind} type must have an element type")

src/base/json_schema.rs

@@ -2,7 +2,7 @@ use crate::prelude::*;
 
 use crate::utils::immutable::RefList;
 use schemars::schema::{
-    ArrayValidation, InstanceType, ObjectValidation, Schema, SchemaObject, SingleOrVec,
+    ArrayValidation, InstanceType, ObjectValidation, Schema, SchemaObject, SingleOrVec, SubschemaValidation,
 };
 use std::fmt::Write;
 
@@ -176,6 +176,16 @@ impl JsonSchemaBuilder {
                     ..Default::default()
                 }));
             }
+            schema::BasicValueType::Union(s) => {
+                schema.subschemas = Some(Box::new(SubschemaValidation {
+                    one_of: Some(
+                        s.types.iter()
+                            .map(|t| Schema::Object(self.for_basic_value_type(t, field_path)))
+                            .collect()
+                    ),
+                    ..Default::default()
+                }));
+            }
         }
         schema
     }

src/base/schema.rs

@@ -9,6 +9,11 @@ pub struct VectorTypeSchema {
     pub dimension: Option<usize>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub struct UnionTypeSchema {
+    pub types: Vec<BasicValueType>,
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
 #[serde(tag = "kind")]
 pub enum BasicValueType {
@@ -56,6 +61,9 @@ pub enum BasicValueType {
 
     /// A vector of values (usually numbers, for embeddings).
     Vector(VectorTypeSchema),
+
+    /// A union
+    Union(UnionTypeSchema),
 }
 
 impl std::fmt::Display for BasicValueType {
@@ -82,6 +90,17 @@ impl std::fmt::Display for BasicValueType {
                 }
                 write!(f, "]")
             }
+            BasicValueType::Union(s) => {
+                write!(f, "Union[")?;
+                for (i, typ) in s.types.iter().enumerate() {
+                    if i > 0 {
+                        // Add type delimiter
+                        write!(f, " | ")?;
+                    }
+                    write!(f, "{}", typ)?;
+                }
+                write!(f, "]")
+            }
         }
     }
 }

src/base/value.rs

             } 
 impl Serialize for TypedValue<'_> { 
     fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> { 
         match (self.t, self.v) { 
             (_, Value::Null) => serializer.serialize_none(), 
             (ValueType::Basic(_), v) => v.serialize(serializer), 
             (ValueType::Struct(s), Value::Struct(field_values)) => TypedFieldsValue { 
                 schema: &s.fields, 
                 values_iter: field_values.fields.iter(), 
             } 
@@ -379,6 +379,10 @@ pub enum BasicValue {
     TimeDelta(chrono::Duration),
     Json(Arc<serde_json::Value>),
     Vector(Arc<[BasicValue]>),
+    UnionVariant {
+        tag_id: usize,
+        value: Box<BasicValue>,
+    },
 }
 
 impl From<Bytes> for BasicValue {
@@ -496,7 +500,8 @@ impl BasicValue {
             | BasicValue::OffsetDateTime(_)
             | BasicValue::TimeDelta(_)
             | BasicValue::Json(_)
-            | BasicValue::Vector(_) => api_bail!("invalid key value type"),
+            | BasicValue::Vector(_)
+            | BasicValue::UnionVariant { .. } => api_bail!("invalid key value type"),
         };
         Ok(result)
     }
@@ -517,7 +522,8 @@ impl BasicValue {
             | BasicValue::OffsetDateTime(_)
             | BasicValue::TimeDelta(_)
             | BasicValue::Json(_)
-            | BasicValue::Vector(_) => api_bail!("invalid key value type"),
+            | BasicValue::Vector(_)
+            | BasicValue::UnionVariant { .. } => api_bail!("invalid key value type"),
         };
         Ok(result)
     }
@@ -539,6 +545,7 @@ impl BasicValue {
             BasicValue::TimeDelta(_) => "timedelta",
             BasicValue::Json(_) => "json",
             BasicValue::Vector(_) => "vector",
+            BasicValue::UnionVariant { .. } => "union",
         }
     }
 }
@@ -890,6 +897,12 @@ impl serde::Serialize for BasicValue {
             BasicValue::TimeDelta(v) => serializer.serialize_str(&v.to_string()),
             BasicValue::Json(v) => v.serialize(serializer),
             BasicValue::Vector(v) => v.serialize(serializer),
+            BasicValue::UnionVariant { tag_id, value } => {
+                let mut s = serializer.serialize_tuple(2)?;
+                s.serialize_element(tag_id)?;
+                s.serialize_element(value)?;
+                s.end()
+            }
         }
     }
 }
@@ -954,6 +967,33 @@ impl BasicValue {
                     .collect::<Result<Vec<_>>>()?;
                 BasicValue::Vector(Arc::from(vec))
             }
+            (v, BasicValueType::Union(typ)) => {
+                let obj: Vec<serde_json::Value> = serde_json::from_value(v) 
+                    .map_err(|_| anyhow::anyhow!("Invalid JSON value for union, expect array"))?;
+
+                if obj.len() != 2 {
+                    anyhow::bail!("Invalid union tuple: expect 2 values, received {}", obj.len());
+                }
+
+                let mut obj_iter = obj.into_iter();
+
+                // Take first element
+                let tag_id = obj_iter
+                    .next()
+                    .and_then(|value| value.as_u64().map(|num_u64| num_u64 as usize))
+                    .unwrap();
+
+                // Take second element
+                let value = obj_iter.next().unwrap();
+
+                let cur_type = typ.types.get(tag_id)
+                    .ok_or_else(|| anyhow::anyhow!("No type in `tag_id` \"{tag_id}\" found"))?;
+
+                BasicValue::UnionVariant {
+                    tag_id,
+                    value: Box::new(BasicValue::from_json(value, cur_type)?),
+                }
+            }
             (v, t) => {
                 anyhow::bail!("Value and type not matched.\nTarget type {t:?}\nJSON value: {v}\n")
             }

src/ops/storages/kuzu.rs

@@ -123,6 +123,14 @@ fn basic_type_to_kuzu(basic_type: &BasicValueType) -> Result<String> {
             t.dimension
                 .map_or_else(|| "".to_string(), |d| d.to_string())
         ),
+        BasicValueType::Union(t) => format!(
+            "UNION({})",
+            t.types.iter()
+                .enumerate()
+                .map(|(i,typ)| Ok(format!("val{} {}", i, basic_type_to_kuzu(typ)?)))
+                .collect::<Result<Vec<_>>>()?
+                .join(", "),
+        ),
         t @ (BasicValueType::Time | BasicValueType::Json) => {
             api_bail!("{t} is not supported in Kuzu")
         }
@@ -379,6 +387,11 @@ fn append_basic_value(cypher: &mut CypherBuilder, basic_value: &BasicValue) -> R
             }
             write!(cypher.query_mut(), "]")?;
         }
+        BasicValue::UnionVariant { tag_id, value } => {
+            write!(cypher.query_mut(), "union_value(val{}:=", tag_id)?;
+            append_basic_value(cypher, value)?;
+            write!(cypher.query_mut(), ")")?;
+        }
         v @ (BasicValue::Time(_) | BasicValue::Json(_)) => {
             bail!("value types are not supported in Kuzu: {}", v.kind());
         }

src/ops/storages/neo4j.rs

@@ -226,6 +226,15 @@ fn basic_value_to_bolt(value: &BasicValue, schema: &BasicValueType) -> Result<Bo
             _ => anyhow::bail!("Non-vector type got vector value: {}", schema),
         },
         BasicValue::Json(v) => json_value_to_bolt_value(v)?,
+        BasicValue::UnionVariant { tag_id, value } => match schema {
+            BasicValueType::Union(s) => {
+                let typ = s.types.get(*tag_id)
+                    .ok_or_else(|| anyhow::anyhow!("Invalid `tag_id`: {}", tag_id))?;
+
+                basic_value_to_bolt(value, typ)?
+            }
+            _ => anyhow::bail!("Non-union type got union value: {}", schema),
+        },
     };
     Ok(bolt_value)
 }

src/ops/storages/postgres.rs

@@ -154,6 +154,12 @@ fn bind_value_field<'arg>(
                     builder.push_bind(sqlx::types::Json(v));
                 }
             },
+            BasicValue::UnionVariant { .. } => {
+                builder.push_bind(sqlx::types::Json(TypedValue {
+                    t: &field_schema.value_type.typ,
+                    v: value,
+                }));
+            }
         },
         Value::Null => {
             builder.push("NULL");
@@ -383,6 +389,7 @@ fn to_column_type_sql(column_type: &ValueType) -> String {
                     "jsonb".into()
                 }
             }
+            BasicValueType::Union(_) => "jsonb".into(),
         },
         _ => "jsonb".into(),
     }

src/py/convert.rs

@@ -1,4 +1,5 @@
 use bytes::Bytes;
+use pyo3::exceptions::PyTypeError;
 use numpy::{PyArray1, PyArrayDyn, PyArrayMethods};
 use pyo3::IntoPyObjectExt;
 use pyo3::types::PyAny;
@@ -76,6 +77,9 @@ fn basic_value_to_py_object<'py>(
         value::BasicValue::TimeDelta(v) => v.into_bound_py_any(py)?,
         value::BasicValue::Json(v) => pythonize(py, v).into_py_result()?,
         value::BasicValue::Vector(v) => handle_vector_to_py(py, v)?,
+        value::BasicValue::UnionVariant { tag_id, value } => {
+            (*tag_id, basic_value_to_py_object(py, &value)?).into_bound_py_any(py)?
+        }
     };
     Ok(result)
 }
@@ -162,6 +166,24 @@ fn basic_value_from_py_object<'py>(
                 ))
             }
         }
+        schema::BasicValueType::Union(s) => {
+            let mut valid_value = None;
+
+            // Try parsing the value
+            for (i, typ) in s.types.iter().enumerate() {
+                if let Ok(value) = basic_value_from_py_object(typ, v) {
+                    valid_value = Some(value::BasicValue::UnionVariant {
+                        tag_id: i,
+                        value: Box::new(value),
+                    });
+                    break;
+                }
+            }
+
+            valid_value.ok_or_else(|| {
+                PyErr::new::<PyTypeError, _>(format!("invalid union value: {}, available types: {:?}", v, s.types))
+            })?
+        }
     };
     Ok(result)
 }