# Copyright 2023 RisingWave Labs # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json import logging import inspect from concurrent.futures import ThreadPoolExecutor from typing import Iterator, Callable, Optional, Union, List, Dict import pyarrow as pa from pyarrow.flight import FlightServerBase, FlightInfo # comes from Databend MAX_DECIMAL128_PRECISION = 38 MAX_DECIMAL256_PRECISION = 76 EXTENSION_KEY = b"Extension" ARROW_EXT_TYPE_VARIANT = b"Variant" TIMESTAMP_UINT = "us" logger = logging.getLogger(__name__) class UserDefinedFunction: """ Base interface for user-defined function. """ _name: str _input_schema: pa.Schema _result_schema: pa.Schema def eval_batch(self, batch: pa.RecordBatch) -> Iterator[pa.RecordBatch]: """ Apply the function on a batch of inputs. """ return iter([]) class ScalarFunction(UserDefinedFunction): """ Base interface for user-defined scalar function. A user-defined scalar functions maps zero, one, or multiple scalar values to a new scalar value. """ _func: Callable _executor: Optional[ThreadPoolExecutor] def __init__( self, func, input_types, result_type, name=None, io_threads=None, skip_null=None ): self._func = func self._input_schema = pa.schema( field.with_name(arg_name) for arg_name, field in zip( inspect.getfullargspec(func)[0], [_to_arrow_field(t) for t in _to_list(input_types)], ) ) self._result_schema = pa.schema( [_to_arrow_field(result_type).with_name("output")] ) self._name = name or ( func.__name__ if hasattr(func, "__name__") else func.__class__.__name__ ) self._io_threads = io_threads self._executor = ( ThreadPoolExecutor(max_workers=self._io_threads) if self._io_threads is not None else None ) if skip_null and not self._result_schema.field(0).nullable: raise ValueError( f"Return type of function {self._name} must be nullable when skip_null is True" ) super().__init__() def eval_batch(self, batch: pa.RecordBatch) -> Iterator[pa.RecordBatch]: # Convert the RecordBatch to a list of arrays with Python objects inputs = [[v.as_py() for v in array] for array in batch] # Apply the input processing function to each array inputs = [ _input_process_func(_list_field(field))(array) for array, field in zip(inputs, self._input_schema) ] # Evaluate the function for the entire batch args = inputs column = self._func(*args) # Apply the output processing function to the result column = _output_process_func(_list_field(self._result_schema.field(0)))(column) # Convert the result to a PyArrow array and yield a RecordBatch array = pa.array(column, type=self._result_schema.types[0]) output_batch = pa.RecordBatch.from_arrays([array], schema=self._result_schema) yield output_batch def __call__(self, *args): return self._func(*args) def udf( input_types: Union[List[Union[str, pa.DataType]], Union[str, pa.DataType]], result_type: Union[str, pa.DataType], name: Optional[str] = None, ) -> Callable: return lambda f: ScalarFunction(f, input_types, result_type, name) class UDFServer(FlightServerBase): """ A server that provides user-defined functions to clients. Example: ``` server = UdfServer(location="0.0.0.0:8815") server.add_function(my_udf) server.serve() ``` """ _location: str _functions: Dict[str, UserDefinedFunction] def __init__(self, location="0.0.0.0:8815", **kwargs): super(UDFServer, self).__init__("grpc://" + location, **kwargs) self._location = location self._functions = {} def get_flight_info(self, context, descriptor): """Return the result schema of a function.""" func_name = descriptor.path[0].decode("utf-8") if func_name not in self._functions: raise ValueError(f"Function {func_name} does not exists") udf = self._functions[func_name] # return the concatenation of input and output schema full_schema = pa.schema(list(udf._input_schema) + list(udf._result_schema)) return FlightInfo( schema=full_schema, descriptor=descriptor, endpoints=[], total_records=len(full_schema), total_bytes=0, ) def do_exchange(self, context, descriptor, reader, writer): """Call a function from the client.""" func_name = descriptor.path[0].decode("utf-8") if func_name not in self._functions: raise ValueError(f"Function {func_name} does not exists") udf = self._functions[func_name] writer.begin(udf._result_schema) try: for batch in reader: for output_batch in udf.eval_batch(batch.data): writer.write_batch(output_batch) except Exception as e: logger.exception(e) raise e def add_function(self, udf: UserDefinedFunction): """Add a function to the server.""" name = udf._name if name in self._functions: raise ValueError("Function already exists: " + name) self._functions[name] = udf input_types = ", ".join( _arrow_field_to_string(field) for field in udf._input_schema ) output_type = _arrow_field_to_string(udf._result_schema[0]) sql = ( f"CREATE FUNCTION {name} ({input_types}) " f"RETURNS {output_type} LANGUAGE python " f"HANDLER = '{name}' ADDRESS = 'http://{self._location}';" ) logger.info(f"added function: {name}, SQL:\n{sql}\n") def serve(self): """Start the server.""" logger.info(f"listening on {self._location}") super(UDFServer, self).serve() def _input_process_func(field: pa.Field) -> Callable: """ Return a function to process input value. - Tuple=pa.struct(): dict -> tuple - Json=pa.large_binary(): bytes -> Any - Map=pa.map_(): list[tuple(k,v)] -> dict """ 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] if array 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 map is not None else None ) if pa.types.is_map(field.type): funcs = [ _input_process_func(field.type.key_field), _input_process_func(field.type.item_field), ] # 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 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 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. """ if isinstance(t, str): return _type_str_to_arrow_field(t) else: return pa.field("", t, False) def _type_str_to_arrow_field(type_str: str) -> pa.Field: """ Convert a SQL data type to `pyarrow.Field`. """ type_str = type_str.strip().upper() nullable = True if type_str.endswith("NULL"): type_str = type_str[:-4].strip() if type_str.endswith("NOT"): type_str = type_str[:-3].strip() nullable = False return _type_str_to_arrow_field_inner(type_str).with_nullable(nullable) def _type_str_to_arrow_field_inner(type_str: str) -> pa.Field: type_str = type_str.strip().upper() if type_str in ("BOOLEAN", "BOOL"): return pa.field("", pa.bool_(), False) elif type_str in ("TINYINT", "INT8"): return pa.field("", pa.int8(), False) elif type_str in ("SMALLINT", "INT16"): return pa.field("", pa.int16(), False) elif type_str in ("INT", "INTEGER", "INT32"): return pa.field("", pa.int32(), False) elif type_str in ("BIGINT", "INT64"): return pa.field("", pa.int64(), False) elif type_str in ("TINYINT UNSIGNED", "UINT8"): return pa.field("", pa.uint8(), False) elif type_str in ("SMALLINT UNSIGNED", "UINT16"): return pa.field("", pa.uint16(), False) elif type_str in ("INT UNSIGNED", "INTEGER UNSIGNED", "UINT32"): return pa.field("", pa.uint32(), False) elif type_str in ("BIGINT UNSIGNED", "UINT64"): return pa.field("", pa.uint64(), False) elif type_str in ("FLOAT", "FLOAT32"): return pa.field("", pa.float32(), False) elif type_str in ("FLOAT64", "DOUBLE"): return pa.field("", pa.float64(), False) elif type_str == "DATE": return pa.field("", pa.date32(), False) 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.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( "", pa.large_binary(), nullable=False, metadata={EXTENSION_KEY: ARROW_EXT_TYPE_VARIANT}, ) elif type_str.startswith("NULLABLE"): type_str = type_str[8:].strip("()").strip() return _type_str_to_arrow_field_inner(type_str).with_nullable(True) elif type_str.endswith("NULL"): type_str = type_str[:-4].strip() return _type_str_to_arrow_field_inner(type_str).with_nullable(True) elif type_str.startswith("DECIMAL"): # DECIMAL(precision, scale) str_list = type_str[7:].strip("()").split(",") precision = int(str_list[0].strip()) scale = int(str_list[1].strip()) if precision < 1 or precision > MAX_DECIMAL256_PRECISION: raise ValueError( f"Decimal precision must be between 1 and {MAX_DECIMAL256_PRECISION}" ) elif scale > precision: raise ValueError( f"Decimal scale must be between 0 and precision {precision}" ) if precision < MAX_DECIMAL128_PRECISION: return pa.field("", pa.decimal128(precision, scale), False) else: return pa.field("", pa.decimal256(precision, scale), False) elif type_str.startswith("ARRAY"): # ARRAY(INT) type_str = type_str[5:].strip("()").strip() return pa.field("", pa.list_(_type_str_to_arrow_field_inner(type_str)), False) elif type_str.startswith("MAP"): # MAP(STRING, INT) str_list = type_str[3:].strip("()").split(",") key_field = _type_str_to_arrow_field_inner(str_list[0].strip()) val_field = _type_str_to_arrow_field_inner(str_list[1].strip()) return pa.field("", pa.map_(key_field, val_field), False) elif type_str.startswith("TUPLE"): # TUPLE(STRING, INT, INT) str_list = type_str[5:].strip("()").split(",") fields = [] for type_str in str_list: type_str = type_str.strip() fields.append(_type_str_to_arrow_field_inner(type_str)) return pa.field("", pa.struct(fields), False) else: raise ValueError(f"Unsupported type: {type_str}") def _arrow_field_to_string(field: pa.Field) -> str: """ Convert a `pyarrow.Field` to a SQL data type string. """ 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 = _field_type_to_string(field) return f"{type_str} NULL" if field.nullable else type_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): return "TINYINT" elif pa.types.is_int16(t): return "SMALLINT" elif pa.types.is_int32(t): return "INT" elif pa.types.is_int64(t): return "BIGINT" elif pa.types.is_uint8(t): return "TINYINT UNSIGNED" elif pa.types.is_uint16(t): return "SMALLINT UNSIGNED" elif pa.types.is_uint32(t): return "INT UNSIGNED" elif pa.types.is_uint64(t): return "BIGINT UNSIGNED" elif pa.types.is_float32(t): return "FLOAT" elif pa.types.is_float64(t): return "DOUBLE" elif pa.types.is_decimal(t): return f"DECIMAL({t.precision}, {t.scale})" elif pa.types.is_date32(t): return "DATE" elif pa.types.is_timestamp(t): return "TIMESTAMP" elif pa.types.is_large_unicode(t) or pa.types.is_unicode(t): return "VARCHAR" elif pa.types.is_large_binary(t) or pa.types.is_binary(t): 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): return f"MAP({_inner_field_to_string(t.key_field)}, {_inner_field_to_string(t.item_field)})" elif pa.types.is_struct(t): args_str = ", ".join(_inner_field_to_string(field) for field in t) return f"TUPLE({args_str})" else: raise ValueError(f"Unsupported type: {t}") 

import logging from databend_udf import udf logger = logging.getLogger(__name__) @udf( input_types=["STRING"], result_type="STRING", ) def ping(inputs: list[str]) -> list[str]: logger.info(f"ping function called with {len(inputs)} inputs") results = [] try: for input_str in inputs: logger.info(f"Processing input: {input_str}") result = "pong" logger.info(f"Returning result: {result}") results.append(result) except Exception as e: logger.error(f"Error in ping function: {e}") results = [ "error" for _ in inputs ] # Return 'error' for all inputs in case of error return results # Debug: Print confirmation of UDF definition logger.info("Defined UDF: ping") 

select ping(data) from texts;