fix: string and binary type

.github/workflows/python.yml

@@ -34,6 +34,7 @@ jobs:
           python -m pip install build
           python -m build
       - uses: pypa/gh-action-pypi-publish@release/v1
+        if: github.event_name == 'push'
         with:
           packages-dir: python/dist/
           skip-existing: true

python/.gitignore

@@ -2,3 +2,4 @@ venv/
 *.egg-info/
 dist/
 __pycache__/
+Pipfile.lock

python/Pipfile

@@ -0,0 +1,14 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+databend-udf = {file = "."}
+
+[dev-packages]
+flake8 = "*"
+black = "*"
+
+[requires]
+python_version = "3.12"

python/databend_udf/udf.py

@@ -13,8 +13,8 @@
 # limitations under the License.
 
 import json
+import logging
 import inspect
-import traceback
 from concurrent.futures import ThreadPoolExecutor
 from typing import Iterator, Callable, Optional, Union, List, Dict
 
@@ -24,11 +24,13 @@
 # comes from Databend
 MAX_DECIMAL128_PRECISION = 38
 MAX_DECIMAL256_PRECISION = 76
-EXTENSION_KEY = "Extension"
-ARROW_EXT_TYPE_VARIANT = "Variant"
+EXTENSION_KEY = b"Extension"
+ARROW_EXT_TYPE_VARIANT = b"Variant"
 
 TIMESTAMP_UINT = "us"
 
+logger = logging.getLogger(__name__)
+
 
 class UserDefinedFunction:
     """
@@ -92,8 +94,8 @@ def __init__(
     def eval_batch(self, batch: pa.RecordBatch) -> Iterator[pa.RecordBatch]:
         inputs = [[v.as_py() for v in array] for array in batch]
         inputs = [
-            _process_func(pa.list_(type), False)(array)
-            for array, type in zip(inputs, self._input_schema.types)
+            _input_process_func(_list_field(field))(array)
+            for array, field in zip(inputs, self._input_schema)
         ]
         if self._executor is not None:
             # concurrently evaluate the function for each row
@@ -122,7 +124,7 @@ def eval_batch(self, batch: pa.RecordBatch) -> Iterator[pa.RecordBatch]:
                     for row in range(batch.num_rows)
                 ]
 
-        column = _process_func(pa.list_(self._result_schema.types[0]), True)(column)
+        column = _output_process_func(_list_field(self._result_schema.field(0)))(column)
 
         array = pa.array(column, type=self._result_schema.types[0])
         yield pa.RecordBatch.from_arrays([array], schema=self._result_schema)
@@ -231,7 +233,7 @@ def do_exchange(self, context, descriptor, reader, writer):
                 for output_batch in udf.eval_batch(batch.data):
                     writer.write_batch(output_batch)
         except Exception as e:
-            print(traceback.print_exc())
+            logger.exception(e)
             raise e
 
     def add_function(self, udf: UserDefinedFunction):
@@ -249,104 +251,138 @@ def add_function(self, udf: UserDefinedFunction):
             f"RETURNS {output_type} LANGUAGE python "
             f"HANDLER = '{name}' ADDRESS = 'http://{self._location}';"
         )
-        print(f"added function: {name}, corresponding SQL:\n{sql}\n")
+        logger.info(f"added function: {name}, SQL:\n{sql}\n")
 
     def serve(self):
         """Start the server."""
-        print(f"listening on {self._location}")
+        logger.info(f"listening on {self._location}")
         super(UDFServer, self).serve()
 
 
-def _null_func(*args):
-    return None
-
-
-def _process_func(type: pa.DataType, output: bool) -> Callable:
+def _input_process_func(field: pa.Field) -> Callable:
     """
-    Return a function to process input or output value.
-
-    For input type:
-        - String=pa.string(): bytes -> str
-        - Tuple=pa.struct(): dict -> tuple
-        - Json=pa.large_binary(): bytes -> Any
-        - Map=pa.map_(): list[tuple(k,v)] -> dict
+    Return a function to process input value.
 
-    For output type:
-        - Json=pa.large_binary(): Any -> str
-        - Map=pa.map_(): dict -> list[tuple(k,v)]
+    - Tuple=pa.struct(): dict -> tuple
+    - Json=pa.large_binary(): bytes -> Any
+    - Map=pa.map_(): list[tuple(k,v)] -> dict
     """
-    if pa.types.is_list(type):
-        func = _process_func(type.value_type, output)
+    if pa.types.is_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]
+            lambda array: [func(v) if v is not None else None for v in array]
             if array is not None
             else None
         )
-    if pa.types.is_struct(type):
-        funcs = [_process_func(field.type, output) for field in type]
-        if output:
-            return (
-                lambda tup: tuple(
-                    (func(v) if v is not None else None) for v, func in zip(tup, funcs)
-                )
-                if tup is not None
-                else None
-            )
-        else:
-            # the input value of struct type is a dict
-            # we convert it into tuple here
-            return (
-                lambda map: tuple(
-                    (func(v) if v is not None else None)
-                    for v, func in zip(map.values(), funcs)
-                )
-                if map is not None
-                else None
+    if pa.types.is_struct(field.type):
+        funcs = [_input_process_func(f) for f in field.type]
+        # the input value of struct type is a dict
+        # we convert it into tuple here
+        return (
+            lambda map: tuple(
+                func(v) if v is not None else None
+                for v, func in zip(map.values(), funcs)
             )
-    if pa.types.is_map(type):
+            if map is not None
+            else None
+        )
+    if pa.types.is_map(field.type):
         funcs = [
-            _process_func(type.key_type, output),
-            _process_func(type.item_type, output),
+            _input_process_func(field.type.key_field),
+            _input_process_func(field.type.item_field),
         ]
-        if output:
-            # dict -> list[tuple[k,v]]
-            return (
-                lambda map: [
-                    tuple(func(v) for v, func in zip(item, funcs))
-                    for item in map.items()
-                ]
-                if map is not None
-                else None
+        # list[tuple[k,v]] -> dict
+        return (
+            lambda array: dict(
+                tuple(func(v) for v, func in zip(item, funcs)) for item in array
             )
-        else:
-            # list[tuple[k,v]] -> dict
-            return (
-                lambda array: dict(
-                    tuple(func(v) for v, func in zip(item, funcs)) for item in array
-                )
-                if array is not None
-                else None
+            if array is not None
+            else None
+        )
+    if pa.types.is_large_binary(field.type):
+        if _field_is_variant(field):
+            return lambda v: json.loads(v) if v is not None else None
+
+    return lambda v: v
+
+
+def _output_process_func(field: pa.Field) -> Callable:
+    """
+    Return a function to process output value.
+
+    - Json=pa.large_binary(): Any -> str
+    - Map=pa.map_(): dict -> list[tuple(k,v)]
+    """
+    if pa.types.is_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]
+            if array is not None
+            else None
+        )
+    if pa.types.is_struct(field.type):
+        funcs = [_output_process_func(f) for f in field.type]
+        return (
+            lambda tup: tuple(
+                func(v) if v is not None else None for v, func in zip(tup, funcs)
             )
+            if tup is not None
+            else None
+        )
+    if pa.types.is_map(field.type):
+        funcs = [
+            _output_process_func(field.type.key_field),
+            _output_process_func(field.type.item_field),
+        ]
+        # dict -> list[tuple[k,v]]
+        return (
+            lambda map: [
+                tuple(func(v) for v, func in zip(item, funcs)) for item in map.items()
+            ]
+            if map is not None
+            else None
+        )
+    if pa.types.is_large_binary(field.type):
+        if _field_is_variant(field):
+            return lambda v: json.dumps(_ensure_str(v)) if v is not None else None
 
-    if pa.types.is_string(type) and not output:
-        # string type is converted to LargeBinary in Databend,
-        # we cast it back to string here
-        return lambda v: v.decode("utf-8") if v is not None else None
-    if pa.types.is_large_binary(type):
-        if output:
-            return lambda v: json.dumps(v) if v is not None else None
-        else:
-            return lambda v: json.loads(v) if v is not None else None
     return lambda v: v
 
 
+def _null_func(*args):
+    return None
+
+
+def _list_field(field: pa.Field) -> pa.Field:
+    return pa.field("", pa.list_(field))
+
+
 def _to_list(x):
     if isinstance(x, list):
         return x
     else:
         return [x]
 
 
+def _ensure_str(x):
+    if isinstance(x, bytes):
+        return x.decode("utf-8")
+    elif isinstance(x, list):
+        return [_ensure_str(v) for v in x]
+    elif isinstance(x, dict):
+        return {_ensure_str(k): _ensure_str(v) for k, v in x.items()}
+    else:
+        return x
+
+
+def _field_is_variant(field: pa.Field) -> bool:
+    if field.metadata is None:
+        return False
+    if field.metadata.get(EXTENSION_KEY) == ARROW_EXT_TYPE_VARIANT:
+        return True
+    return False
+
+
 def _to_arrow_field(t: Union[str, pa.DataType]) -> pa.Field:
     """
     Convert a string or pyarrow.DataType to pyarrow.Field.
@@ -401,7 +437,9 @@ def _type_str_to_arrow_field_inner(type_str: str) -> pa.Field:
     elif type_str in ("DATETIME", "TIMESTAMP"):
         return pa.field("", pa.timestamp(TIMESTAMP_UINT), False)
     elif type_str in ("STRING", "VARCHAR", "CHAR", "CHARACTER", "TEXT"):
-        return pa.field("", pa.string(), False)
+        return pa.field("", pa.large_utf8(), False)
+    elif type_str in ("BINARY"):
+        return pa.field("", pa.large_binary(), False)
     elif type_str in ("VARIANT", "JSON"):
         # In Databend, JSON type is identified by the "EXTENSION" key in the metadata.
         return pa.field(
@@ -460,20 +498,21 @@ def _arrow_field_to_string(field: pa.Field) -> str:
     """
     Convert a `pyarrow.Field` to a SQL data type string.
     """
-    type_str = _data_type_to_string(field.type)
+    type_str = _field_type_to_string(field)
     return f"{type_str} NOT NULL" if not field.nullable else type_str
 
 
 def _inner_field_to_string(field: pa.Field) -> str:
     # inner field default is NOT NULL in databend
-    type_str = _data_type_to_string(field.type)
+    type_str = _field_type_to_string(field)
     return f"{type_str} NULL" if field.nullable else type_str
 
 
-def _data_type_to_string(t: pa.DataType) -> str:
+def _field_type_to_string(field: pa.Field) -> str:
     """
     Convert a `pyarrow.DataType` to a SQL data type string.
     """
+    t = field.type
     if pa.types.is_boolean(t):
         return "BOOLEAN"
     elif pa.types.is_int8(t):
@@ -502,10 +541,13 @@ def _data_type_to_string(t: pa.DataType) -> str:
         return "DATE"
     elif pa.types.is_timestamp(t):
         return "TIMESTAMP"
-    elif pa.types.is_string(t):
+    elif pa.types.is_large_unicode(t):
         return "VARCHAR"
     elif pa.types.is_large_binary(t):
-        return "VARIANT"
+        if _field_is_variant(field):
+            return "VARIANT"
+        else:
+            return "BINARY"
     elif pa.types.is_list(t):
         return f"ARRAY({_inner_field_to_string(t.value_field)})"
     elif pa.types.is_map(t):