Support VECTOR(N) type

Author: bohutang

python/databend_udf/udf.py

@@ -1179,7 +1179,7 @@ def _input_process_func(field: pa.Field) -> Callable:
     - Json=pa.large_binary(): bytes -> Any
     - Map=pa.map_(): list[tuple(k,v)] -> dict
     """
-    if pa.types.is_list(field.type):
+    if pa.types.is_list(field.type) or pa.types.is_fixed_size_list(field.type):
         func = _input_process_func(field.type.value_field)
         return (
             lambda array: [func(v) if v is not None else None for v in array]
@@ -1225,7 +1225,7 @@ def _output_process_func(field: pa.Field) -> Callable:
     - Json=pa.large_binary(): Any -> str
     - Map=pa.map_(): dict -> list[tuple(k,v)]
     """
-    if pa.types.is_list(field.type):
+    if pa.types.is_list(field.type) or pa.types.is_fixed_size_list(field.type):
         func = _output_process_func(field.type.value_field)
         return (
             lambda array: [func(v) if v is not None else None for v in array]
@@ -1402,6 +1402,10 @@ def _type_str_to_arrow_field_inner(type_str: str) -> pa.Field:
             type_str = type_str.strip()
             fields.append(_type_str_to_arrow_field_inner(type_str))
         return pa.field("", pa.struct(fields), False)
+    elif type_str.startswith("VECTOR"):
+        # VECTOR(1024)
+        dim = int(type_str[6:].strip("()").strip())
+        return pa.field("", pa.list_(pa.float32(), dim), False)
     else:
         raise ValueError(f"Unsupported type: {type_str}")
 
@@ -1460,6 +1464,8 @@ 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/tests/test_vector_type.py

@@ -0,0 +1,56 @@
+import pyarrow as pa
+import pytest
+from databend_udf.udf import (
+    _type_str_to_arrow_field,
+    _field_type_to_string,
+    _input_process_func,
+    _output_process_func,
+    _arrow_field_to_string,
+)
+
+def test_vector_sql_generation():
+    # Test nullable VECTOR (default)
+    field = _type_str_to_arrow_field("VECTOR(1024)")
+    sql_type = _arrow_field_to_string(field)
+    assert sql_type == "VECTOR(1024)"
+
+    # Test NOT NULL VECTOR
+    field_not_null = _type_str_to_arrow_field("VECTOR(1024) NOT NULL")
+    sql_type_not_null = _arrow_field_to_string(field_not_null)
+    assert sql_type_not_null == "VECTOR(1024) NOT NULL"
+
+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
+    assert pa.types.is_float32(field.type.value_type)
+    assert field.nullable is True
+
+def test_vector_type_formatting():
+    field = pa.field("", pa.list_(pa.float32(), 1024), nullable=True)
+    type_str = _field_type_to_string(field)
+    assert type_str == "VECTOR(1024)"
+
+def test_vector_input_processing():
+    field = pa.field("", pa.list_(pa.float32(), 3), nullable=True)
+    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
+
+def test_vector_output_processing():
+    field = pa.field("", pa.list_(pa.float32(), 3), nullable=True)
+    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