Fix VECTOR workaround: use custom metadata key

Author: bohutang

python/databend_udf/udf.py

@@ -49,7 +49,9 @@
 MAX_DECIMAL256_PRECISION = 76
 EXTENSION_KEY = b"Extension"
 ARROW_EXT_TYPE_VARIANT = b"Variant"
-ARROW_EXT_TYPE_VECTOR = b"Vector"
+# Use a custom key to avoid triggering Arrow Extension Type mechanism in Databend
+VECTOR_METADATA_KEY = b"x-databend-vector"
+
 
 TIMESTAMP_UINT = "us"
 _SCHEMA_METADATA_INPUT_COUNT_KEY = b"x-databend-udf-input-count"
@@ -1413,8 +1415,7 @@ def _type_str_to_arrow_field_inner(type_str: str) -> pa.Field:
             pa.list_(pa.field("item", pa.float32(), nullable=True)),
             nullable=False,
             metadata={
-                EXTENSION_KEY: ARROW_EXT_TYPE_VECTOR,
-                b"vector_size": str(dim).encode(),
+                VECTOR_METADATA_KEY: str(dim).encode(),
             },
         )
     else:
@@ -1440,8 +1441,8 @@ 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"))
+    if field.metadata and field.metadata.get(VECTOR_METADATA_KEY):
+        dim = int(field.metadata.get(VECTOR_METADATA_KEY))
         return f"VECTOR({dim})"
 
     if pa.types.is_boolean(t):

python/pyproject.toml

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

python/tests/test_vector_type.py

@@ -25,8 +25,7 @@ def test_vector_type_parsing():
     field = _type_str_to_arrow_field("VECTOR(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 field.metadata[b"x-databend-vector"] == b"1024"
     assert pa.types.is_float32(field.type.value_type)
     # Default is nullable
     assert field.nullable is True
@@ -43,8 +42,7 @@ def test_vector_type_formatting():
         pa.list_(pa.field("item", pa.float32(), nullable=True)),
         nullable=False,
         metadata={
-            b"Extension": b"Vector",
-            b"vector_size": b"1024",
+            b"x-databend-vector": b"1024",
         },
     )
     type_str = _field_type_to_string(field)
@@ -58,8 +56,7 @@ def test_vector_input_processing():
         pa.list_(pa.field("item", pa.float32(), nullable=True)),
         nullable=False,
         metadata={
-            b"Extension": b"Vector",
-            b"vector_size": b"3",
+            b"x-databend-vector": b"3",
         },
     )
     func = _input_process_func(field)
@@ -77,8 +74,7 @@ def test_vector_output_processing():
         pa.list_(pa.field("item", pa.float32(), nullable=True)),
         nullable=False,
         metadata={
-            b"Extension": b"Vector",
-            b"vector_size": b"3",
+            b"x-databend-vector": b"3",
         },
     )
     func = _output_process_func(field)