Workaround VECTOR type: use List with metadata

Author: BohuTANG

python/databend_udf/udf.py

@@ -49,6 +49,7 @@
 MAX_DECIMAL256_PRECISION = 76
 EXTENSION_KEY = b"Extension"
 ARROW_EXT_TYPE_VARIANT = b"Variant"
+ARROW_EXT_TYPE_VECTOR = b"Vector"
 
 TIMESTAMP_UINT = "us"
 _SCHEMA_METADATA_INPUT_COUNT_KEY = b"x-databend-udf-input-count"
@@ -1405,8 +1406,16 @@ def _type_str_to_arrow_field_inner(type_str: str) -> pa.Field:
     elif type_str.startswith("VECTOR"):
         # VECTOR(1024)
         dim = int(type_str[6:].strip("()").strip())
+        # Use List(float) with metadata to represent VECTOR(N)
+        # This is a workaround because Databend UDF client might not support FixedSizeList yet.
         return pa.field(
-            "", pa.list_(pa.field("item", pa.float32(), nullable=False), dim), False
+            "",
+            pa.list_(pa.field("item", pa.float32(), nullable=False)),
+            nullable=False,
+            metadata={
+                EXTENSION_KEY: ARROW_EXT_TYPE_VECTOR,
+                b"vector_size": str(dim).encode(),
+            },
         )
     else:
         raise ValueError(f"Unsupported type: {type_str}")
@@ -1431,6 +1440,10 @@ def _field_type_to_string(field: pa.Field) -> str:
     Convert a `pyarrow.DataType` to a SQL data type string.
     """
     t = field.type
+    if field.metadata and field.metadata.get(EXTENSION_KEY) == ARROW_EXT_TYPE_VECTOR:
+        dim = int(field.metadata.get(b"vector_size", b"0"))
+        return f"VECTOR({dim})"
+
     if pa.types.is_boolean(t):
         return "BOOLEAN"
     elif pa.types.is_int8(t):
@@ -1466,8 +1479,6 @@ def _field_type_to_string(field: pa.Field) -> str:
             return "VARIANT"
         else:
             return "BINARY"
-    elif pa.types.is_fixed_size_list(t):
-        return f"VECTOR({t.list_size})"
     elif pa.types.is_list(t):
         return f"ARRAY({_inner_field_to_string(t.value_field)})"
     elif pa.types.is_map(t):

python/pyproject.toml

@@ -7,7 +7,7 @@ classifiers = [
 description = "Databend UDF Server"
 license = { text = "Apache-2.0" }
 name = "databend-udf"
-version = "0.2.13"
+version = "0.2.14"
 readme = "README.md"
 requires-python = ">=3.7"
 dependencies = [

python/tests/test_vector_type.py

@@ -23,47 +23,67 @@ def test_vector_sql_generation():
 
 def test_vector_type_parsing():
     field = _type_str_to_arrow_field("VECTOR(1024)")
-    assert pa.types.is_fixed_size_list(field.type)
-    assert field.type.list_size == 1024
+    # Should be List type with metadata, not FixedSizeList
+    assert pa.types.is_list(field.type)
+    assert field.metadata[b"Extension"] == b"Vector"
+    assert field.metadata[b"vector_size"] == b"1024"
     assert pa.types.is_float32(field.type.value_type)
+    # Default is nullable
     assert field.nullable is True
 
+    # Test NOT NULL
+    field_not_null = _type_str_to_arrow_field("VECTOR(1024) NOT NULL")
+    assert field_not_null.nullable is False
+
 
 def test_vector_type_formatting():
+    # Test that a List with VECTOR metadata is formatted as VECTOR(N)
     field = pa.field(
         "",
-        pa.list_(pa.field("item", pa.float32(), nullable=False), 1024),
-        nullable=True,
+        pa.list_(pa.field("item", pa.float32(), nullable=False)),
+        nullable=False,
+        metadata={
+            b"Extension": b"Vector",
+            b"vector_size": b"1024",
+        },
     )
     type_str = _field_type_to_string(field)
     assert type_str == "VECTOR(1024)"
 
 
 def test_vector_input_processing():
+    # Input processing should handle List (which is what VECTOR is physically)
     field = pa.field(
-        "", pa.list_(pa.field("item", pa.float32(), nullable=False), 3), nullable=True
+        "",
+        pa.list_(pa.field("item", pa.float32(), nullable=False)),
+        nullable=False,
+        metadata={
+            b"Extension": b"Vector",
+            b"vector_size": b"3",
+        },
     )
     func = _input_process_func(field)
 
     # Input is a list of floats
-    input_data = [1.0, 2.0, 3.0]
-    result = func(input_data)
-    assert result == [1.0, 2.0, 3.0]
-
-    # Input is None
-    assert func(None) is None
+    data = [[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]]
+    processed = func(data)
+    assert processed == data
 
 
 def test_vector_output_processing():
+    # Output processing should handle List
     field = pa.field(
-        "", pa.list_(pa.field("item", pa.float32(), nullable=False), 3), nullable=True
+        "",
+        pa.list_(pa.field("item", pa.float32(), nullable=False)),
+        nullable=False,
+        metadata={
+            b"Extension": b"Vector",
+            b"vector_size": b"3",
+        },
     )
     func = _output_process_func(field)
 
     # Output is a list of floats
-    output_data = [1.0, 2.0, 3.0]
-    result = func(output_data)
-    assert result == [1.0, 2.0, 3.0]
-
-    # Output is None
-    assert func(None) is None
+    data = [[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]]
+    processed = func(data)
+    assert processed == data