Add support for multiple return value types and return value field names

Author: kesmit13

setup.cfg

@@ -81,7 +81,7 @@ exclude =
     docs/*
     resources/*
     licenses/*
-max-complexity = 35
+max-complexity = 45
 max-line-length = 90
 per-file-ignores =
     singlestoredb/__init__.py:F401

singlestoredb/functions/decorator.py

@@ -1,13 +1,28 @@
+import datetime
 import functools
+import string
 from typing import Any
 from typing import Callable
 from typing import Dict
 from typing import List
 from typing import Optional
 from typing import Union
 
+from . import dtypes
 from .dtypes import DataType
 
+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."""
@@ -23,30 +38,37 @@ def _func(
     *,
     name: Optional[str] = None,
     args: Optional[Union[DataType, List[DataType], Dict[str, DataType]]] = None,
-    returns: Optional[str] = None,
+    returns: Optional[Union[str, List[DataType], List[type]]] = 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."""
     if args is None:
         pass
     elif isinstance(args, (list, tuple)):
         args = list(args)
         for i, item in enumerate(args):
-            if callable(item):
+            if args[i] in python_type_map:
+                args[i] = python_type_map[args[i]]()
+            elif callable(item):
                 args[i] = item()
         for item in args:
             if not isinstance(item, str):
                 raise TypeError(f'unrecognized type for parameter: {item}')
     elif isinstance(args, dict):
         args = dict(args)
         for k, v in list(args.items()):
-            if callable(v):
+            if args[k] in python_type_map:
+                args[k] = python_type_map[args[k]]()
+            elif callable(v):
                 args[k] = v()
         for item in args.values():
             if not isinstance(item, str):
                 raise TypeError(f'unrecognized type for parameter: {item}')
+    elif args in python_type_map:
+        args = python_type_map[args]()
     elif callable(args):
         args = args()
     elif isinstance(args, str):
@@ -56,16 +78,47 @@ def _func(
 
     if returns is None:
         pass
+    elif isinstance(returns, (list, tuple)):
+        returns = list(returns)
+        for i, item in enumerate(returns):
+            if item in python_type_map:
+                returns[i] = python_type_map[item]()
+            elif callable(item):
+                returns[i] = item()
+        for item in returns:
+            if not isinstance(item, str):
+                raise TypeError(f'unrecognized return type: {item}')
+    elif returns in python_type_map:
+        returns = python_type_map[returns]()
     elif callable(returns):
         returns = returns()
     elif isinstance(returns, str):
         returns = returns
     else:
         raise TypeError(f'unrecognized return type: {returns}')
 
-    if returns is not None and not isinstance(returns, str):
+    if returns is None:
+        pass
+    elif isinstance(returns, list):
+        for item in returns:
+            if not isinstance(item, str):
+                raise TypeError(f'unrecognized return type: {item}')
+    elif not isinstance(returns, str):
         raise TypeError(f'unrecognized return type: {returns}')
 
+    if not output_fields:
+        if isinstance(returns, list):
+            output_fields = []
+            for i, _ in enumerate(returns):
+                output_fields.append(string.ascii_letters[i])
+        else:
+            output_fields = [string.ascii_letters[0]]
+
+    if isinstance(returns, 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 '
@@ -80,6 +133,7 @@ def _func(
             data_format=data_format,
             include_masks=include_masks,
             function_type=function_type,
+            output_fields=output_fields,
         ).items() if v is not None
     }
 
@@ -107,7 +161,7 @@ def udf(
     *,
     name: Optional[str] = None,
     args: Optional[Union[DataType, List[DataType], Dict[str, DataType]]] = None,
-    returns: Optional[str] = None,
+    returns: Optional[Union[str, List[DataType], List[type]]] = None,
     data_format: Optional[str] = None,
     include_masks: bool = False,
 ) -> Callable[..., Any]:
@@ -170,9 +224,10 @@ def tvf(
     *,
     name: Optional[str] = None,
     args: Optional[Union[DataType, List[DataType], Dict[str, DataType]]] = None,
-    returns: Optional[str] = None,
+    returns: Optional[Union[str, List[DataType], List[type]]] = None,
     data_format: Optional[str] = None,
     include_masks: bool = False,
+    output_fields: Optional[List[str]] = None,
 ) -> Callable[..., Any]:
     """
     Apply attributes to a TVF.
@@ -205,6 +260,9 @@ def tvf(
         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
     -------
@@ -219,6 +277,7 @@ def tvf(
         data_format=data_format,
         include_masks=include_masks,
         function_type='tvf',
+        output_fields=output_fields,
     )
 
 

singlestoredb/functions/ext/asgi.py

@@ -190,7 +190,7 @@ async def do_func(  # type: ignore
                     assert isinstance(out, tuple)
                     return row_ids, [out]
 
-                out_ids, out = [], []
+                out = []
                 res = func(*[x[0] for x in cols])
                 rtype = str(type(res)).lower()
 
@@ -209,20 +209,20 @@ async def do_func(  # type: ignore
                 # NOTE: There is no way to determine which row ID belongs to
                 #        each result row, so we just have to use the same
                 #        row ID for all rows in the result.
-                if data_format == 'numpy':
-                    import numpy as np
-                    out_ids = np.array([row_ids[0]] * len(out[0][0]))
-                elif data_format == 'polars':
+                if data_format == 'polars':
                     import polars as pl
-                    out_ids = pl.Series([row_ids[0]] * len(out[0][0]))
+                    array_cls = pl.Series
                 elif data_format == 'arrow':
                     import pyarrow as pa
-                    out_ids = pa.array([row_ids[0]] * len(out[0][0]))
+                    array_cls = pa.array
                 elif data_format == 'pandas':
                     import pandas as pd
-                    out_ids = pd.Series([row_ids[0]] * len(out[0][0]))
+                    array_cls = pd.Series
+                else:
+                    import numpy as np
+                    array_cls = np.array
 
-                return out_ids, out
+                return array_cls([row_ids[0]] * len(out[0][0])), out
 
     else:
         if data_format == 'python':

singlestoredb/functions/signature.py

@@ -312,7 +312,10 @@ def classify_dtype(dtype: Any) -> str:
     if is_int:
         return int_type_map.get(name, 'int64')
 
-    raise TypeError(f'unsupported type annotation: {dtype}')
+    raise TypeError(
+        f'unsupported type annotation: {dtype}; '
+        'use `args`/`returns` on the @udf/@tvf decotator to specify the data type',
+    )
 
 
 def collapse_dtypes(dtypes: Union[str, List[str]]) -> str:
@@ -449,6 +452,7 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
 
     args_overrides = attrs.get('args', None)
     returns_overrides = attrs.get('returns', None)
+    output_fields = attrs.get('output_fields', None)
 
     spec_diff = set(arg_names).difference(set(annotations.keys()))
 
@@ -499,6 +503,22 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
     if isinstance(returns_overrides, str):
         sql = returns_overrides
         out_type = sql_to_dtype(sql)
+    elif isinstance(returns_overrides, list):
+        sqls = []
+        out_types = []
+        for i, item in enumerate(returns_overrides):
+            if not isinstance(item, str):
+                raise TypeError(f'unrecognized type for return value: {item}')
+            if output_fields:
+                sqls.append(f'`{output_fields[i]}` {item}')
+            else:
+                sqls.append(f'{string.ascii_letters[i]} {item}')
+            out_types.append(sql_to_dtype(item))
+        if function_type == 'tvf':
+            sql = 'TABLE({})'.format(', '.join(sqls))
+        else:
+            sql = 'RECORD({})'.format(', '.join(sqls))
+        out_type = 'tuple[{}]'.format(','.join(out_types))
     elif returns_overrides is not None and not isinstance(returns_overrides, str):
         raise TypeError(f'unrecognized type for return value: {returns_overrides}')
     else:

singlestoredb/tests/test_udf.py

@@ -391,16 +391,6 @@ def foo(x: int, y: float, z: str) -> int: ...
             '`y` DOUBLE NOT NULL, ' \
             '`z` CHAR(30) NULL) RETURNS SMALLINT NOT NULL'
 
-        # Override parameter with incorrect type
-        with self.assertRaises(TypeError):
-            @udf(args=dict(x=int))
-            def foo(x: int, y: float, z: str) -> int: ...
-
-        # Override return value with incorrect type
-        with self.assertRaises(TypeError):
-            @udf(returns=int)
-            def foo(x: int, y: float, z: str) -> int: ...
-
         # Change function name
         @udf(name='hello_world')
         def foo(x: int) -> int: ...