singlestoredb 1.10.0__cp38-abi3-win32.whl → 1.12.0__cp38-abi3-win32.whl

singlestoredb/__init__.py

@@ -13,7 +13,7 @@ Examples
 
 """
 
-__version__ = '1.10.0'
+__version__ = '1.12.0'
 
 from typing import Any
 

singlestoredb/config.py

@@ -201,6 +201,12 @@ register_option(
     environ='SINGLESTOREDB_BUFFERED',
 )
 
+register_option(
+    'parse_json', 'bool', check_bool, True,
+    'Parse JSON values into Python objects?',
+    environ='SINGLESTOREDB_PARSE_JSON',
+)
+
 register_option(
     'connect_timeout', 'int', check_int, 10,
     'The timeout for connecting to the database in seconds. '

singlestoredb/connection.py

@@ -541,10 +541,16 @@ class Cursor(metaclass=abc.ABCMeta):
 
         Examples
         --------
+        Query with no parameters
+
         >>> cur.execute('select * from mytable')
 
+        Query with positional parameters
+
         >>> cur.execute('select * from mytable where id < %s', [100])
 
+        Query with named parameters
+
         >>> cur.execute('select * from mytable where id < %(max)s', dict(max=100))
 
         Returns
@@ -1317,6 +1323,7 @@ def connect(
     track_env: Optional[bool] = None,
     enable_extended_data_types: Optional[bool] = None,
     vector_data_format: Optional[str] = None,
+    parse_json: Optional[bool] = None,
 ) -> Connection:
     """
     Return a SingleStoreDB connection.

singlestoredb/converters.py

@@ -594,7 +594,7 @@ def float32_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
     if has_numpy:
         return numpy.frombuffer(x, dtype=numpy.float32)
 
-    return struct.unpack(f'<{len(x)/4}f', x)
+    return struct.unpack(f'<{len(x)//4}f', x)
 
 
 def float64_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
@@ -650,7 +650,7 @@ def float64_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
     if has_numpy:
         return numpy.frombuffer(x, dtype=numpy.float64)
 
-    return struct.unpack(f'<{len(x)/8}d', x)
+    return struct.unpack(f'<{len(x)//8}d', x)
 
 
 def int8_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
@@ -762,7 +762,7 @@ def int16_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
     if has_numpy:
         return numpy.frombuffer(x, dtype=numpy.int16)
 
-    return struct.unpack(f'<{len(x)/2}h', x)
+    return struct.unpack(f'<{len(x)//2}h', x)
 
 
 def int32_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
@@ -818,7 +818,7 @@ def int32_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
     if has_numpy:
         return numpy.frombuffer(x, dtype=numpy.int32)
 
-    return struct.unpack(f'<{len(x)/4}l', x)
+    return struct.unpack(f'<{len(x)//4}l', x)
 
 
 def int64_vector_json_or_none(x: Optional[str]) -> Optional[Any]:
@@ -875,7 +875,7 @@ def int64_vector_or_none(x: Optional[bytes]) -> Optional[Any]:
     if has_numpy:
         return numpy.frombuffer(x, dtype=numpy.int64)
 
-    return struct.unpack(f'<{len(x)/8}l', x)
+    return struct.unpack(f'<{len(x)//8}l', x)
 
 
 def bson_or_none(x: Optional[bytes]) -> Optional[Any]:

singlestoredb/functions/__init__.py

@@ -1 +1,2 @@
+from .decorator import tvf  # noqa: F401
 from .decorator import udf  # noqa: F401

singlestoredb/functions/decorator.py

@@ -1,12 +1,36 @@
+import dataclasses
+import datetime
 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 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]:
@@ -18,12 +42,170 @@ def listify(x: Any) -> List[Any]:
     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
+
+
+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](), []
+
+    elif callable(params):
+        return params(), []
+
+    elif isinstance(params, str):
+        return params, []
+
+    raise TypeError(f'unrecognized data type for args: {params}')
+
+
+def _func(
+    func: Optional[Callable[..., Any]] = None,
+    *,
+    name: Optional[str] = None,
+    args: Optional[
+        Union[
+            DataType,
+            List[DataType],
+            Dict[str, DataType],
+            'pydantic.BaseModel',
+            type,
+        ]
+    ] = None,
+    returns: Optional[
+        Union[
+            str,
+            List[DataType],
+            List[type],
+            'pydantic.BaseModel',
+            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."""
+    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,
+            include_masks=include_masks,
+            function_type=function_type,
+            output_fields=output_fields or None,
+        ).items() if v is not None
+    }
+
+    # No func was specified, this is an uncalled decorator that will get
+    # called later, so the wrapper much be created with the func passed
+    # in at that time.
+    if func is None:
+        def decorate(func: Callable[..., Any]) -> Callable[..., Any]:
+            def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
+                return func(*args, **kwargs)  # type: ignore
+            wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+            return functools.wraps(func)(wrapper)
+        return decorate
+
+    def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
+        return func(*args, **kwargs)  # type: ignore
+
+    wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+
+    return functools.wraps(func)(wrapper)
+
+
 def udf(
     func: Optional[Callable[..., Any]] = None,
     *,
     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]:
@@ -64,79 +246,86 @@ def udf(
     Callable
 
     """
-    if args is None:
-        pass
-    elif isinstance(args, (list, tuple)):
-        args = list(args)
-        for i, item in enumerate(args):
-            if 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):
-                args[k] = v()
-        for item in args.values():
-            if not isinstance(item, str):
-                raise TypeError(f'unrecognized type for parameter: {item}')
-    elif callable(args):
-        args = args()
-    elif isinstance(args, str):
-        args = args
-    else:
-        raise TypeError(f'unrecognized data type for args: {args}')
-
-    if returns is None:
-        pass
-    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):
-        raise TypeError(f'unrecognized return type: {returns}')
+    return _func(
+        func=func,
+        name=name,
+        args=args,
+        returns=returns,
+        data_format=data_format,
+        include_masks=include_masks,
+        function_type='udf',
+    )
 
-    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,
-            include_masks=include_masks,
-        ).items() if v is not None
-    }
+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
 
-    # No func was specified, this is an uncalled decorator that will get
-    # called later, so the wrapper much be created with the func passed
-    # in at that time.
-    if func is None:
-        def decorate(func: Callable[..., Any]) -> Callable[..., Any]:
-            def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
-                return func(*args, **kwargs)  # type: ignore
-            wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
-            return functools.wraps(func)(wrapper)
-        return decorate
 
-    def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
-        return func(*args, **kwargs)  # type: ignore
+def tvf(
+    func: Optional[Callable[..., Any]] = None,
+    *,
+    name: Optional[str] = None,
+    args: Optional[Union[DataType, List[DataType], Dict[str, DataType]]] = 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.
 
-    wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+    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] | Dict[str, 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.
+    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.
 
-    return functools.wraps(func)(wrapper)
+    Returns
+    -------
+    Callable
 
+    """
+    return _func(
+        func=func,
+        name=name,
+        args=args,
+        returns=returns,
+        data_format=data_format,
+        include_masks=include_masks,
+        function_type='tvf',
+        output_fields=output_fields,
+    )
 
-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
+
+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

singlestoredb/functions/ext/asgi.py

@@ -24,6 +24,7 @@ Example
 """
 import argparse
 import asyncio
+import dataclasses
 import importlib.util
 import io
 import itertools
@@ -136,6 +137,14 @@ def get_func_names(funcs: str) -> List[Tuple[str, str]]:
     return out
 
 
+def as_tuple(x: Any) -> Any:
+    if hasattr(x, 'model_fields'):
+        return tuple(x.model_fields.values())
+    if dataclasses.is_dataclass(x):
+        return dataclasses.astuple(x)
+    return x
+
+
 def make_func(
     name: str,
     func: Callable[..., Any],
@@ -158,32 +167,104 @@ def make_func(
     attrs = getattr(func, '_singlestoredb_attrs', {})
     data_format = attrs.get('data_format') or 'python'
     include_masks = attrs.get('include_masks', False)
+    function_type = attrs.get('function_type', 'udf').lower()
     info: Dict[str, Any] = {}
 
-    if data_format == 'python':
-        async def do_func(
-            row_ids: Sequence[int],
-            rows: Sequence[Sequence[Any]],
-        ) -> Tuple[
-            Sequence[int],
-            List[Tuple[Any]],
-        ]:
-            '''Call function on given rows of data.'''
-            return row_ids, list(zip(func_map(func, rows)))
+    if function_type == 'tvf':
+        if data_format == 'python':
+            async def do_func(
+                row_ids: Sequence[int],
+                rows: Sequence[Sequence[Any]],
+            ) -> Tuple[
+                Sequence[int],
+                List[Tuple[Any]],
+            ]:
+                '''Call function on given rows of data.'''
+                out_ids: List[int] = []
+                out = []
+                for i, res in zip(row_ids, func_map(func, rows)):
+                    out.extend(as_tuple(res))
+                    out_ids.extend([row_ids[i]] * (len(out)-len(out_ids)))
+                return out_ids, out
+
+        else:
+            # Vector formats use the same function wrapper
+            async def do_func(  # type: ignore
+                row_ids: Sequence[int],
+                cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
+            ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
+                '''Call function on given cols of data.'''
+                if include_masks:
+                    out = func(*cols)
+                    assert isinstance(out, tuple)
+                    return row_ids, [out]
+
+                out = []
+                res = func(*[x[0] for x in cols])
+                rtype = str(type(res)).lower()
+
+                # Map tables / dataframes to a list of columns
+                if 'dataframe' in rtype:
+                    res = [res[x] for x in res.columns]
+                elif 'table' in rtype:
+                    res = res.columns
+
+                for vec in res:
+                    # C extension only supports Python objects as strings
+                    if data_format == 'numpy' and str(vec.dtype)[:2] in ['<U', '<S']:
+                        vec = vec.astype(object)
+                    out.append((vec, None))
+
+                # 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 == 'polars':
+                    import polars as pl
+                    array_cls = pl.Series
+                elif data_format == 'arrow':
+                    import pyarrow as pa
+                    array_cls = pa.array
+                elif data_format == 'pandas':
+                    import pandas as pd
+                    array_cls = pd.Series
+                else:
+                    import numpy as np
+                    array_cls = np.array
+
+                return array_cls([row_ids[0]] * len(out[0][0])), out
 
     else:
-        # Vector formats use the same function wrapper
-        async def do_func(  # type: ignore
-            row_ids: Sequence[int],
-            cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
-        ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
-            '''Call function on given cols of data.'''
-            # TODO: only supports a single return value
-            if include_masks:
-                out = func(*cols)
-                assert isinstance(out, tuple)
-                return row_ids, [out]
-            return row_ids, [(func(*[x[0] for x in cols]), None)]
+        if data_format == 'python':
+            async def do_func(
+                row_ids: Sequence[int],
+                rows: Sequence[Sequence[Any]],
+            ) -> Tuple[
+                Sequence[int],
+                List[Tuple[Any]],
+            ]:
+                '''Call function on given rows of data.'''
+                return row_ids, [as_tuple(x) for x in zip(func_map(func, rows))]
+
+        else:
+            # Vector formats use the same function wrapper
+            async def do_func(  # type: ignore
+                row_ids: Sequence[int],
+                cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
+            ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
+                '''Call function on given cols of data.'''
+                if include_masks:
+                    out = func(*cols)
+                    assert isinstance(out, tuple)
+                    return row_ids, [out]
+
+                out = func(*[x[0] for x in cols])
+
+                # Multiple return values
+                if isinstance(out, tuple):
+                    return row_ids, [(x, None) for x in out]
+
+                # Single return value
+                return row_ids, [(out, None)]
 
     do_func.__name__ = name
     do_func.__doc__ = func.__doc__
@@ -196,6 +277,9 @@ def make_func(
     # Set data format
     info['data_format'] = data_format
 
+    # Set function type
+    info['function_type'] = function_type
+
     # Setup argument types for rowdat_1 parser
     colspec = []
     for x in sig['args']:
@@ -205,11 +289,21 @@ def make_func(
         colspec.append((x['name'], rowdat_1_type_map[dtype]))
     info['colspec'] = colspec
 
+    def parse_return_type(s: str) -> List[str]:
+        if s.startswith('tuple['):
+            return s[6:-1].split(',')
+        if s.startswith('array[tuple['):
+            return s[12:-2].split(',')
+        return [s]
+
     # Setup return type
-    dtype = sig['returns']['dtype'].replace('?', '')
-    if dtype not in rowdat_1_type_map:
-        raise TypeError(f'no data type mapping for {dtype}')
-    info['returns'] = [rowdat_1_type_map[dtype]]
+    returns = []
+    for x in parse_return_type(sig['returns']['dtype']):
+        dtype = x.replace('?', '')
+        if dtype not in rowdat_1_type_map:
+            raise TypeError(f'no data type mapping for {dtype}')
+        returns.append(rowdat_1_type_map[dtype])
+    info['returns'] = returns
 
     return do_func, info
 
@@ -233,7 +327,7 @@ class Application(object):
             * Function aliases : <pkg1>.[<func1@alias1,func2@alias2,...]
             * Multiple packages : <pkg1>.<func1>:<pkg2>.<func2>
     app_mode : str, optional
-        The mode of operation for the application: remote or collocated
+        The mode of operation for the application: remote, managed, or collocated
     url : str, optional
         The URL of the function API
     data_format : str, optional