UDF annotation work (#61)

* Refactor to use Table / Masked types

* Downgrade mypy for Python 3.8

* Fix annotations for older versions of Python

* Fix annotations for older versions of Python

* Fix annotations for older versions of Python

* Fix annotations for older versions of Python

* Fix annotations for older versions of Python

* Add null masks

* Fix difference in numpy detection

* Short circuit common valid types

* Add 3.13 checks

* Update autopep

* Add 3.13 to smoke tests

* Fix Table wrappers

* Fix masks in table results

* Vector utility functions

Author: kesmit13

.github/workflows/pre-commit.yml

@@ -14,6 +14,7 @@ jobs:
           - "3.10"
           - "3.11"
           - "3.12"
+          - "3.13"
 
     steps:
     - uses: actions/checkout@v3

.github/workflows/smoke-test.yml

@@ -54,6 +54,7 @@ jobs:
           - "3.10"
           - "3.11"
           - "3.12"
+          - "3.13"
         driver:
           - mysql
           - https

.pre-commit-config.yaml

@@ -21,7 +21,7 @@ repos:
         exclude: singlestoredb/clients/pymysqlsv/
         additional_dependencies: [flake8-typing-imports==1.12.0]
 -   repo: https://github.com/hhatto/autopep8
-    rev: v2.0.4
+    rev: v2.3.1
     hooks:
     -   id: autopep8
         args: [--diff]
@@ -40,7 +40,7 @@ repos:
     hooks:
     -   id: setup-cfg-fmt
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.6.1
+    rev: v1.14.1
     hooks:
     -   id: mypy
         additional_dependencies: [types-requests]

requirements.txt

@@ -5,4 +5,5 @@ requests
 setuptools
 sqlparams
 tomli>=1.1.0; python_version < '3.11'
+typing_extensions<=4.13.2
 wheel

setup.cfg

@@ -27,6 +27,7 @@ install_requires =
     sqlparams
     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,6 +1,12 @@
-from .decorator import tvf  # noqa: F401
-from .decorator import tvf_with_null_masks  # noqa: F401
 from .decorator import udf  # noqa: F401
-from .decorator import udf_with_null_masks  # noqa: F401
 from .typing import Masked  # noqa: F401
-from .typing import MaskedNDArray  # noqa: F401
+from .typing import Table  # noqa: F401
+from .utils import VectorTypes
+
+
+F32 = VectorTypes.F32
+F64 = VectorTypes.F64
+I8 = VectorTypes.I8
+I16 = VectorTypes.I16
+I32 = VectorTypes.I32
+I64 = VectorTypes.I64

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
@@ -60,40 +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 typing.get_origin(obj) is not tuple or len(typing.get_args(obj)) != 2:
-        raise TypeError(
-            f'Expected a tuple of two vector types, but got {type(obj)}',
-        )
-
-    args = typing.get_args(obj)
-
-    if not utils.is_vector(args[0]):
-        raise TypeError(
-            f'Expected a vector type for the first element, but got {args[0]}',
-        )
-
-    if not utils.is_vector(args[1]):
-        raise TypeError(
-            f'Expected a vector type for the second element, but got {args[1]}',
-        )
-
-    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:
@@ -135,8 +100,6 @@ def _func(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-    with_null_masks: bool = False,
-    function_type: str = 'udf',
 ) -> Callable[..., Any]:
     """Generic wrapper for UDF and TVF decorators."""
 
@@ -145,8 +108,6 @@ def _func(
             name=name,
             args=expand_types(args),
             returns=expand_types(returns),
-            with_null_masks=with_null_masks,
-            function_type=function_type,
         ).items() if v is not None
     }
 
@@ -155,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
@@ -167,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
 
@@ -194,151 +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,
-        function_type='udf',
-    )
-
-
-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
-        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=True,
-        function_type='udf',
-    )
-
-
-def tvf(
-    func: Optional[Callable[..., Any]] = None,
-    *,
-    name: Optional[str] = None,
-    args: Optional[ParameterType] = None,
-    returns: Optional[ReturnType] = None,
-) -> Callable[..., Any]:
-    """
-    Define a table-valued function (TVF).
-
-    Parameters
-    ----------
-    func : callable, optional
-        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], 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,
-        function_type='tvf',
-    )
-
-
-def tvf_with_null_masks(
-    func: Optional[Callable[..., Any]] = None,
-    *,
-    name: Optional[str] = None,
-    args: Optional[ParameterType] = None,
-    returns: Optional[ReturnType] = None,
-) -> Callable[..., Any]:
-    """
-    Define a table-valued function (TVF) using null masks.
-
-    Parameters
-    ----------
-    func : callable, optional
-        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], 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
@@ -351,9 +163,10 @@ def tvf_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
     -------
@@ -365,6 +178,4 @@ def tvf_with_null_masks(
         name=name,
         args=args,
         returns=returns,
-        with_null_masks=True,
-        function_type='tvf',
     )