Add null masks

Author: kesmit13

setup.cfg

@@ -25,9 +25,9 @@ install_requires =
     requests
     setuptools
     sqlparams
-    typing-extensions<=4.13.2
     wheel
     tomli>=1.1.0;python_version < '3.11'
+    typing-extensions<=4.13.2;python_version < '3.11'
 python_requires = >=3.8
 include_package_data = True
 tests_require =

singlestoredb/functions/__init__.py

@@ -1,4 +1,3 @@
 from .decorator import udf  # noqa: F401
-from .decorator import udf_with_null_masks  # noqa: F401
 from .typing import Masked  # noqa: F401
 from .typing import Table  # noqa: F401

singlestoredb/functions/decorator.py

@@ -1,6 +1,5 @@
 import functools
 import inspect
-import typing
 from typing import Any
 from typing import Callable
 from typing import List
@@ -10,7 +9,6 @@
 
 from . import utils
 from .dtypes import SQLString
-from .typing import Masked
 
 
 ParameterType = Union[
@@ -61,27 +59,6 @@ def is_valid_callable(obj: Any) -> bool:
     )
 
 
-def verify_mask(obj: Any) -> bool:
-    """Verify that the object is a tuple of two vector types."""
-    if not typing.get_origin(obj) is Masked:
-        raise TypeError(
-            f'expected a Masked type, but got {type(obj)}',
-        )
-    return True
-
-
-def verify_masks(obj: Callable[..., Any]) -> bool:
-    """Verify that the function parameters and return value are all masks."""
-    ann = utils.get_annotations(obj)
-    for name, value in ann.items():
-        if not verify_mask(value):
-            raise TypeError(
-                f'Expected a vector type for the parameter {name} '
-                f'in function {obj.__name__}, but got {value}',
-            )
-    return True
-
-
 def expand_types(args: Any) -> Optional[Union[List[str], Type[Any]]]:
     """Expand the types for the function arguments / return values."""
     if args is None:
@@ -123,7 +100,6 @@ def _func(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-    with_null_masks: bool = False,
 ) -> Callable[..., Any]:
     """Generic wrapper for UDF and TVF decorators."""
 
@@ -132,7 +108,6 @@ def _func(
             name=name,
             args=expand_types(args),
             returns=expand_types(returns),
-            with_null_masks=with_null_masks,
         ).items() if v is not None
     }
 
@@ -141,8 +116,6 @@ def _func(
     # in at that time.
     if func is None:
         def decorate(func: Callable[..., Any]) -> Callable[..., Any]:
-            if with_null_masks:
-                verify_masks(func)
 
             def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
                 return func(*args, **kwargs)  # type: ignore
@@ -153,9 +126,6 @@ def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
 
         return decorate
 
-    if with_null_masks:
-        verify_masks(func)
-
     def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
         return func(*args, **kwargs)  # type: ignore
 
@@ -180,54 +150,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], 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
-        takes a single type for all parameters, `args` can be set to a
-        SQL string describing all parameters. If the function takes more
-        than one parameter and all of the parameters are being manually
-        defined, a list of SQL strings may be used (one for each parameter).
-        A dictionary of SQL strings may be used to specify a parameter type
-        for a subset of parameters; the keys are the names of the
-        function parameters. Callables may also be used for datatypes. This
-        is primarily for using the functions in the ``dtypes`` module that
-        are associated with SQL types with all default options (e.g., ``dt.FLOAT``).
-    returns : str, optional
-        Specifies the return data type of the function. If not specified,
-        the type annotation from the function is used.
-
-    Returns
-    -------
-    Callable
-
-    """
-    return _func(
-        func=func,
-        name=name,
-        args=args,
-        returns=returns,
-        with_null_masks=False,
-    )
-
-
-def udf_with_null_masks(
-    func: Optional[Callable[..., Any]] = None,
-    *,
-    name: Optional[str] = None,
-    args: Optional[ParameterType] = None,
-    returns: Optional[ReturnType] = None,
-) -> Callable[..., Any]:
-    """
-    Define a user-defined function (UDF) with null masks.
-
-    Parameters
-    ----------
-    func : callable, optional
-        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], optional
+    args : str | Type | 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
@@ -240,9 +163,10 @@ def udf_with_null_masks(
         function parameters. Callables may also be used for datatypes. This
         is primarily for using the functions in the ``dtypes`` module that
         are associated with SQL types with all default options (e.g., ``dt.FLOAT``).
-    returns : str, optional
-        Specifies the return data type of the function. If not specified,
-        the type annotation from the function is used.
+    returns : str | Type | Callable | List[str | Callable] | Table, optional
+        Specifies the return data type of the function. This parameter
+        works the same way as `args`. If the function is a table-valued
+        function, the return type should be a `Table` object.
 
     Returns
     -------
@@ -254,5 +178,4 @@ def udf_with_null_masks(
         name=name,
         args=args,
         returns=returns,
-        with_null_masks=True,
     )

singlestoredb/functions/ext/asgi.py

@@ -26,6 +26,7 @@
 import asyncio
 import dataclasses
 import importlib.util
+import inspect
 import io
 import itertools
 import json
@@ -36,6 +37,7 @@
 import sys
 import tempfile
 import textwrap
+import typing
 import urllib
 import zipfile
 import zipimport
@@ -62,6 +64,7 @@
 from ...mysql.constants import FIELD_TYPE as ft
 from ..signature import get_signature
 from ..signature import signature_to_sql
+from ..typing import Masked
 
 try:
     import cloudpickle
@@ -207,6 +210,25 @@ def get_array_class(data_format: str) -> Callable[..., Any]:
     return array_cls
 
 
+def get_masked_params(func: Callable[..., Any]) -> List[bool]:
+    """
+    Get the list of masked parameters for the function.
+
+    Parameters
+    ----------
+    func : Callable
+        The function to call as the endpoint
+
+    Returns
+    -------
+    List[bool]
+        Boolean list of masked parameters
+
+    """
+    params = inspect.signature(func).parameters
+    return [typing.get_origin(x.annotation) is Masked for x in params.values()]
+
+
 def make_func(
     name: str,
     func: Callable[..., Any],
@@ -226,8 +248,6 @@ def make_func(
     (Callable, Dict[str, Any])
 
     """
-    attrs = getattr(func, '_singlestoredb_attrs', {})
-    with_null_masks = attrs.get('with_null_masks', False)
     info: Dict[str, Any] = {}
 
     sig = get_signature(func, func_name=name)
@@ -236,6 +256,8 @@ def make_func(
     args_data_format = sig.get('args_data_format', 'scalar')
     returns_data_format = sig.get('returns_data_format', 'scalar')
 
+    masks = get_masked_params(func)
+
     if function_type == 'tvf':
         # Scalar (Python) types
         if returns_data_format == 'scalar':
@@ -265,24 +287,21 @@ async def do_func(  # type: ignore
                 #        each result row, so we just have to use the same
                 #        row ID for all rows in the result.
 
-                # If `with_null_masks` is set, the function is expected to return
-                # a tuple of (data, mask) for each column.
-                if with_null_masks:
-                    out = func(*cols)
-                    assert isinstance(out, tuple)
-                    row_ids = array_cls([row_ids[0]] * len(out[0][0]))
-                    return row_ids, [out]
+                def build_tuple(x: Any) -> Any:
+                    return tuple(x) if isinstance(x, Masked) else (x, None)
 
                 # Call function on each column of data
                 if cols and cols[0]:
-                    res = get_dataframe_columns(func(*[x[0] for x in cols]))
+                    res = get_dataframe_columns(
+                        func(*[x if m else x[0] for x, m in zip(cols, masks)]),
+                    )
                 else:
                     res = get_dataframe_columns(func())
 
                 # Generate row IDs
                 row_ids = array_cls([row_ids[0]] * len(res[0]))
 
-                return row_ids, [(x, None) for x in res]
+                return row_ids, [build_tuple(x) for x in res]
 
     else:
         # Scalar (Python) types
@@ -305,22 +324,22 @@ async def do_func(  # type: ignore
                 '''Call function on given cols of data.'''
                 row_ids = array_cls(row_ids)
 
-                # If `with_null_masks` is set, the function is expected to return
-                # a tuple of (data, mask) for each column.`
-                if with_null_masks:
-                    out = func(*cols)
-                    assert isinstance(out, tuple)
-                    return row_ids, [out]
+                def build_tuple(x: Any) -> Any:
+                    return tuple(x) if isinstance(x, Masked) else (x, None)
 
                 # Call the function with `cols` as the function parameters
                 if cols and cols[0]:
-                    out = func(*[x[0] for x in cols])
+                    out = func(*[x if m else x[0] for x, m in zip(cols, masks)])
                 else:
                     out = func()
 
+                # Single masked value
+                if isinstance(out, Masked):
+                    return row_ids, [tuple(out)]
+
                 # Multiple return values
                 if isinstance(out, tuple):
-                    return row_ids, [(x, None) for x in out]
+                    return row_ids, [build_tuple(x) for x in out]
 
                 # Single return value
                 return row_ids, [(out, None)]