Big refactoring of the way parameters / return values work

Author: kesmit13

accel.c

@@ -36,6 +36,7 @@
 #define NUMPY_DATETIME 13
 #define NUMPY_OBJECT 14
 #define NUMPY_BYTES 15
+#define NUMPY_FIXED_STRING 16
 
 #define MYSQL_FLAG_NOT_NULL 1
 #define MYSQL_FLAG_PRI_KEY 2
@@ -2745,6 +2746,13 @@ static NumpyColType get_numpy_col_type(PyObject *py_array) {
             goto error;
         }
         break;
+    case 'U':
+        out.type = NUMPY_FIXED_STRING;
+        out.length = (Py_ssize_t)strtol(str + 2, NULL, 10);
+        if (out.length < 0) {
+            goto error;
+        }
+        break;
     default:
         goto error;
     }
@@ -3845,7 +3853,7 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
             case MYSQL_TYPE_MEDIUM_BLOB:
             case MYSQL_TYPE_LONG_BLOB:
             case MYSQL_TYPE_BLOB:
-                if  (col_types[i].type != NUMPY_OBJECT) {
+                if  (col_types[i].type != NUMPY_OBJECT && col_types[i].type != NUMPY_FIXED_STRING) {
                     PyErr_SetString(PyExc_ValueError, "unsupported numpy data type for character output types");
                     goto error;
                 }
@@ -3856,6 +3864,24 @@ static PyObject *dump_rowdat_1_numpy(PyObject *self, PyObject *args, PyObject *k
                     memcpy(out+out_idx, &i64, 8);
                     out_idx += 8;
 
+                } else if (col_types[i].type == NUMPY_FIXED_STRING) {
+                    void *bytes = (void*)(cols[i] + j * 8);
+
+                    if (bytes == NULL) {
+                        CHECKMEM(8);
+                        i64 = 0;
+                        memcpy(out+out_idx, &i64, 8);
+                        out_idx += 8;
+                    } else {
+                        Py_ssize_t str_l = strnlen(bytes, col_types[i].length);
+                        CHECKMEM(8+str_l);
+                        i64 = str_l;
+                        memcpy(out+out_idx, &i64, 8);
+                        out_idx += 8;
+                        memcpy(out+out_idx, bytes, str_l);
+                        out_idx += str_l;
+                    }
+
                 } else {
                     u64 = *(uint64_t*)(cols[i] + j * 8);
 

singlestoredb/functions/decorator.py

@@ -1,142 +1,56 @@
-import dataclasses
-import datetime
+from __future__ import annotations
+
 import functools
 import inspect
 from typing import Any
 from typing import Callable
-from typing import Dict
 from typing import List
 from typing import Optional
-from typing import Tuple
+from typing import Type
 from typing import Union
 
-from . import dtypes
-from .dtypes import DataType
-from .signature import simplify_dtype
-
-try:
-    import pydantic
-    has_pydantic = True
-except ImportError:
-    has_pydantic = False
-
-python_type_map: Dict[Any, Callable[..., str]] = {
-    str: dtypes.TEXT,
-    int: dtypes.BIGINT,
-    float: dtypes.DOUBLE,
-    bool: dtypes.BOOL,
-    bytes: dtypes.BINARY,
-    bytearray: dtypes.BINARY,
-    datetime.datetime: dtypes.DATETIME,
-    datetime.date: dtypes.DATE,
-    datetime.timedelta: dtypes.TIME,
-}
-
-
-def listify(x: Any) -> List[Any]:
-    """Make sure sure value is a list."""
-    if x is None:
-        return []
-    if isinstance(x, (list, tuple, set)):
-        return list(x)
-    return [x]
-
-
-def process_annotation(annotation: Any) -> Tuple[Any, bool]:
-    types = simplify_dtype(annotation)
-    if isinstance(types, list):
-        nullable = False
-        if type(None) in types:
-            nullable = True
-            types = [x for x in types if x is not type(None)]
-        if len(types) > 1:
-            raise ValueError(f'multiple types not supported: {annotation}')
-        return types[0], nullable
-    return types, True
 
+ParameterType = Union[
+    str,
+    Callable[..., str],
+    List[Union[str, Callable[..., str]]],
+    Type[Any],
+]
 
-def process_types(params: Any) -> Any:
-    if params is None:
-        return params, []
-
-    elif isinstance(params, (list, tuple)):
-        params = list(params)
-        for i, item in enumerate(params):
-            if params[i] in python_type_map:
-                params[i] = python_type_map[params[i]]()
-            elif callable(item):
-                params[i] = item()
-        for item in params:
-            if not isinstance(item, str):
-                raise TypeError(f'unrecognized type for parameter: {item}')
-        return params, []
-
-    elif isinstance(params, dict):
-        names = []
-        params = dict(params)
-        for k, v in list(params.items()):
-            names.append(k)
-            if params[k] in python_type_map:
-                params[k] = python_type_map[params[k]]()
-            elif callable(v):
-                params[k] = v()
-        for item in params.values():
-            if not isinstance(item, str):
-                raise TypeError(f'unrecognized type for parameter: {item}')
-        return params, names
-
-    elif dataclasses.is_dataclass(params):
-        names = []
-        out = []
-        for item in dataclasses.fields(params):
-            typ, nullable = process_annotation(item.type)
-            sql_type = process_types(typ)[0]
-            if not nullable:
-                sql_type = sql_type.replace('NULL', 'NOT NULL')
-            out.append(sql_type)
-            names.append(item.name)
-        return out, names
-
-    elif has_pydantic and inspect.isclass(params) \
-            and issubclass(params, pydantic.BaseModel):
-        names = []
-        out = []
-        for name, item in params.model_fields.items():
-            typ, nullable = process_annotation(item.annotation)
-            sql_type = process_types(typ)[0]
-            if not nullable:
-                sql_type = sql_type.replace('NULL', 'NOT NULL')
-            out.append(sql_type)
-            names.append(name)
-        return out, names
-
-    elif params in python_type_map:
-        return python_type_map[params](), []
+ReturnType = ParameterType
 
-    elif callable(params):
-        return params(), []
 
-    elif isinstance(params, str):
-        return params, []
+def expand_types(args: Any) -> Optional[Union[List[str], Type[Any]]]:
+    """Expand the types for the function arguments / return values."""
+    if args is None:
+        return None
 
-    raise TypeError(f'unrecognized data type for args: {params}')
+    # SQL string
+    if isinstance(args, str):
+        return [args]
 
+    # General way of accepting pydantic.BaseModel, NamedTuple, TypedDict
+    elif inspect.isclass(args):
+        return args
 
-ParameterType = Union[
-    str,
-    List[str],
-    Dict[str, str],
-    'pydantic.BaseModel',
-    type,
-]
+    # Callable that returns a SQL string
+    elif callable(args):
+        out = args()
+        if not isinstance(out, str):
+            raise TypeError(f'unrecognized type for parameter: {args}')
+        return [out]
 
-ReturnType = Union[
-    str,
-    List[DataType],
-    List[type],
-    'pydantic.BaseModel',
-    type,
-]
+    # List of SQL strings or callables
+    else:
+        new_args = []
+        for arg in args:
+            if isinstance(arg, str):
+                new_args.append(arg)
+            elif callable(arg):
+                new_args.append(arg())
+            else:
+                raise TypeError(f'unrecognized type for parameter: {arg}')
+        return new_args
 
 
 def _func(
@@ -145,40 +59,18 @@ def _func(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-    data_format: Optional[str] = None,
     include_masks: bool = False,
     function_type: str = 'udf',
-    output_fields: Optional[List[str]] = None,
 ) -> Callable[..., Any]:
     """Generic wrapper for UDF and TVF decorators."""
-    args, _ = process_types(args)
-    returns, fields = process_types(returns)
-
-    if not output_fields and fields:
-        output_fields = fields
-
-    if isinstance(returns, list) \
-            and isinstance(output_fields, list) \
-            and len(output_fields) != len(returns):
-        raise ValueError(
-            'The number of output fields must match the number of return types',
-        )
-
-    if include_masks and data_format == 'python':
-        raise RuntimeError(
-            'include_masks is only valid when using '
-            'vectors for input parameters',
-        )
 
     _singlestoredb_attrs = {  # type: ignore
         k: v for k, v in dict(
             name=name,
-            args=args,
-            returns=returns,
-            data_format=data_format,
+            args=expand_types(args),
+            returns=expand_types(returns),
             include_masks=include_masks,
             function_type=function_type,
-            output_fields=output_fields or None,
         ).items() if v is not None
     }
 
@@ -207,7 +99,6 @@ def udf(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-    data_format: Optional[str] = None,
     include_masks: bool = False,
 ) -> Callable[..., Any]:
     """
@@ -219,7 +110,7 @@ def udf(
         The UDF to apply parameters to
     name : str, optional
         The name to use for the UDF in the database
-    args : str | Callable | List[str | Callable] | Dict[str, str | Callable], optional
+    args : str | Callable | List[str | Callable], optional
         Specifies the data types of the function arguments. Typically,
         the function data types are derived from the function parameter
         annotations. These annotations can be overridden. If the function
@@ -235,8 +126,6 @@ def udf(
     returns : str, optional
         Specifies the return data type of the function. If not specified,
         the type annotation from the function is used.
-    data_format : str, optional
-        The data format of each parameter: python, pandas, arrow, polars
     include_masks : bool, optional
         Should boolean masks be included with each input parameter to indicate
         which elements are NULL? This is only used when a input parameters are
@@ -252,27 +141,18 @@ def udf(
         name=name,
         args=args,
         returns=returns,
-        data_format=data_format,
         include_masks=include_masks,
         function_type='udf',
     )
 
 
-udf.pandas = functools.partial(udf, data_format='pandas')  # type: ignore
-udf.polars = functools.partial(udf, data_format='polars')  # type: ignore
-udf.arrow = functools.partial(udf, data_format='arrow')  # type: ignore
-udf.numpy = functools.partial(udf, data_format='numpy')  # type: ignore
-
-
 def tvf(
     func: Optional[Callable[..., Any]] = None,
     *,
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-    data_format: Optional[str] = None,
     include_masks: bool = False,
-    output_fields: Optional[List[str]] = None,
 ) -> Callable[..., Any]:
     """
     Apply attributes to a TVF.
@@ -283,7 +163,7 @@ def tvf(
         The TVF to apply parameters to
     name : str, optional
         The name to use for the TVF in the database
-    args : str | Callable | List[str | Callable] | Dict[str, str | Callable], optional
+    args : str | Callable | List[str | Callable], optional
         Specifies the data types of the function arguments. Typically,
         the function data types are derived from the function parameter
         annotations. These annotations can be overridden. If the function
@@ -299,15 +179,10 @@ def tvf(
     returns : str, optional
         Specifies the return data type of the function. If not specified,
         the type annotation from the function is used.
-    data_format : str, optional
-        The data format of each parameter: python, pandas, arrow, polars
     include_masks : bool, optional
         Should boolean masks be included with each input parameter to indicate
         which elements are NULL? This is only used when a input parameters are
         configured to a vector type (numpy, pandas, polars, arrow).
-    output_fields : List[str], optional
-        The names of the output fields for the TVF. If not specified, the
-        names are generated.
 
     Returns
     -------
@@ -319,14 +194,6 @@ def tvf(
         name=name,
         args=args,
         returns=returns,
-        data_format=data_format,
         include_masks=include_masks,
         function_type='tvf',
-        output_fields=output_fields,
     )
-
-
-tvf.pandas = functools.partial(tvf, data_format='pandas')  # type: ignore
-tvf.polars = functools.partial(tvf, data_format='polars')  # type: ignore
-tvf.arrow = functools.partial(tvf, data_format='arrow')  # type: ignore
-tvf.numpy = functools.partial(tvf, data_format='numpy')  # type: ignore