singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/functions/typing/__init__.py

@@ -0,0 +1,41 @@
+from collections.abc import Iterable
+from typing import Any
+from typing import Tuple
+from typing import TypeVar
+
+try:
+    from typing import TypeVarTuple  # type: ignore
+    from typing import Unpack  # type: ignore
+except ImportError:
+    # Python 3.8 and earlier do not have TypeVarTuple
+    from typing_extensions import TypeVarTuple  # type: ignore
+    from typing_extensions import Unpack  # type: ignore
+
+
+T = TypeVar('T', bound=Iterable[Any])  # Generic type for iterable types
+
+#
+# Masked types are used for pairs of vectors where the first element is the
+# vector and the second element is a boolean mask indicating which elements
+# are NULL. The boolean mask is a vector of the same length as the first
+# element, where True indicates that the corresponding element in the first
+# element is NULL.
+#
+# This is needed for vector types that do not support NULL values, such as
+# numpy arrays and pandas Series.
+#
+
+
+class Masked(Tuple[T, T]):
+    def __new__(cls, *args: T) -> 'Masked[Tuple[T, T]]':  # type: ignore
+        return tuple.__new__(cls, (args[0], args[1]))  # type: ignore
+
+
+Ts = TypeVarTuple('Ts')
+
+
+class Table(Tuple[Unpack[Ts]]):
+    """Return type for a table valued function."""
+
+    def __new__(cls, *args: Unpack[Ts]) -> 'Table[Tuple[Unpack[Ts]]]':  # type: ignore
+        return tuple.__new__(cls, args)  # type: ignore

singlestoredb/functions/typing/numpy.py

@@ -0,0 +1,20 @@
+import numpy as np
+import numpy.typing as npt
+
+NDArray = npt.NDArray
+
+StringArray = StrArray = npt.NDArray[np.str_]
+BytesArray = npt.NDArray[np.bytes_]
+Float32Array = FloatArray = npt.NDArray[np.float32]
+Float64Array = DoubleArray = npt.NDArray[np.float64]
+IntArray = npt.NDArray[np.int_]
+Int8Array = npt.NDArray[np.int8]
+Int16Array = npt.NDArray[np.int16]
+Int32Array = npt.NDArray[np.int32]
+Int64Array = npt.NDArray[np.int64]
+UInt8Array = npt.NDArray[np.uint8]
+UInt16Array = npt.NDArray[np.uint16]
+UInt32Array = npt.NDArray[np.uint32]
+UInt64Array = npt.NDArray[np.uint64]
+DateTimeArray = npt.NDArray[np.datetime64]
+TimeDeltaArray = npt.NDArray[np.timedelta64]

singlestoredb/functions/typing/pandas.py

@@ -0,0 +1,2 @@
+from pandas import DataFrame  # noqa: F401
+from pandas import Series  # noqa: F401

singlestoredb/functions/typing/polars.py

@@ -0,0 +1,2 @@
+from polars import DataFrame  # noqa: F401
+from polars import Series  # noqa: F401

singlestoredb/functions/typing/pyarrow.py

@@ -0,0 +1,2 @@
+from pyarrow import Array  # noqa: F401
+from pyarrow import Table  # noqa: F401

singlestoredb/functions/utils.py

@@ -0,0 +1,421 @@
+import dataclasses
+import inspect
+import struct
+import sys
+import types
+import typing
+from collections.abc import Iterable
+from enum import Enum
+from typing import Any
+from typing import Dict
+from typing import Tuple
+from typing import Union
+
+from .typing import Masked
+
+if sys.version_info >= (3, 10):
+    _UNION_TYPES = {typing.Union, types.UnionType}
+else:
+    _UNION_TYPES = {typing.Union}
+
+
+is_dataclass = dataclasses.is_dataclass
+
+
+def is_masked(obj: Any) -> bool:
+    """Check if an object is a Masked type."""
+    origin = typing.get_origin(obj)
+    if origin is not None:
+        return origin is Masked or \
+            (inspect.isclass(origin) and issubclass(origin, Masked))
+    return False
+
+
+def is_union(x: Any) -> bool:
+    """Check if the object is a Union."""
+    return typing.get_origin(x) in _UNION_TYPES
+
+
+def get_annotations(obj: Any) -> Dict[str, Any]:
+    """Get the annotations of an object."""
+    return typing.get_type_hints(obj)
+
+
+def get_module(obj: Any) -> str:
+    """Get the module of an object."""
+    module = getattr(obj, '__module__', '').split('.')
+    if module:
+        return module[0]
+    return ''
+
+
+def get_type_name(obj: Any) -> str:
+    """Get the type name of an object."""
+    if hasattr(obj, '__name__'):
+        return obj.__name__
+    if hasattr(obj, '__class__'):
+        return obj.__class__.__name__
+    return ''
+
+
+def is_numpy(obj: Any) -> bool:
+    """Check if an object is a numpy array."""
+    if str(obj).startswith('numpy.ndarray['):
+        return True
+
+    if inspect.isclass(obj):
+        if get_module(obj) == 'numpy':
+            return get_type_name(obj) == 'ndarray'
+
+    origin = typing.get_origin(obj)
+    if get_module(origin) == 'numpy':
+        if get_type_name(origin) == 'ndarray':
+            return True
+
+    dtype = type(obj)
+    if get_module(dtype) == 'numpy':
+        return get_type_name(dtype) == 'ndarray'
+
+    return False
+
+
+def is_dataframe(obj: Any) -> bool:
+    """Check if an object is a DataFrame."""
+    # Cheating here a bit so we don't have to import pandas / polars / pyarrow:
+    # unless we absolutely need to
+    if get_module(obj) == 'pandas':
+        return get_type_name(obj) == 'DataFrame'
+    if get_module(obj) == 'polars':
+        return get_type_name(obj) == 'DataFrame'
+    if get_module(obj) == 'pyarrow':
+        return get_type_name(obj) == 'Table'
+    return False
+
+
+def is_vector(obj: Any, include_masks: bool = False) -> bool:
+    """Check if an object is a vector type."""
+    return is_pandas_series(obj) \
+        or is_polars_series(obj) \
+        or is_pyarrow_array(obj) \
+        or is_numpy(obj) \
+        or is_masked(obj)
+
+
+def get_data_format(obj: Any) -> str:
+    """Return the data format of the DataFrame / Table / vector."""
+    # Cheating here a bit so we don't have to import pandas / polars / pyarrow
+    # unless we absolutely need to
+    if get_module(obj) == 'pandas':
+        return 'pandas'
+    if get_module(obj) == 'polars':
+        return 'polars'
+    if get_module(obj) == 'pyarrow':
+        return 'arrow'
+    if get_module(obj) == 'numpy':
+        return 'numpy'
+    if isinstance(obj, list):
+        return 'list'
+    return 'scalar'
+
+
+def is_pandas_series(obj: Any) -> bool:
+    """Check if an object is a pandas Series."""
+    if is_union(obj):
+        obj = typing.get_args(obj)[0]
+    return (
+        get_module(obj) == 'pandas' and
+        get_type_name(obj) == 'Series'
+    )
+
+
+def is_polars_series(obj: Any) -> bool:
+    """Check if an object is a polars Series."""
+    if is_union(obj):
+        obj = typing.get_args(obj)[0]
+    return (
+        get_module(obj) == 'polars' and
+        get_type_name(obj) == 'Series'
+    )
+
+
+def is_pyarrow_array(obj: Any) -> bool:
+    """Check if an object is a pyarrow Array."""
+    if is_union(obj):
+        obj = typing.get_args(obj)[0]
+    return (
+        get_module(obj) == 'pyarrow' and
+        get_type_name(obj) == 'Array'
+    )
+
+
+def is_typeddict(obj: Any) -> bool:
+    """Check if an object is a TypedDict."""
+    if hasattr(typing, 'is_typeddict'):
+        return typing.is_typeddict(obj)  # noqa: TYP006
+    return False
+
+
+def is_namedtuple(obj: Any) -> bool:
+    """Check if an object is a named tuple."""
+    if inspect.isclass(obj):
+        return (
+                issubclass(obj, tuple) and
+                hasattr(obj, '_asdict') and
+                hasattr(obj, '_fields')
+        )
+    return (
+            isinstance(obj, tuple) and
+            hasattr(obj, '_asdict') and
+            hasattr(obj, '_fields')
+    )
+
+
+def is_pydantic(obj: Any) -> bool:
+    """Check if an object is a pydantic model."""
+    if not inspect.isclass(obj):
+        return False
+    # We don't want to import pydantic here, so we check if
+    # the class is a subclass
+    return bool([
+        x for x in inspect.getmro(obj)
+        if get_module(x) == 'pydantic'
+        and get_type_name(x) == 'BaseModel'
+    ])
+
+
+class VectorTypes(str, Enum):
+    """Enum for vector types."""
+    F16 = 'f16'
+    F32 = 'f32'
+    F64 = 'f64'
+    I8 = 'i8'
+    I16 = 'i16'
+    I32 = 'i32'
+    I64 = 'i64'
+
+
+def _vector_type_to_numpy_type(
+    vector_type: VectorTypes,
+) -> str:
+    """Convert a vector type to a numpy type."""
+    if vector_type == VectorTypes.F32:
+        return 'f4'
+    elif vector_type == VectorTypes.F64:
+        return 'f8'
+    elif vector_type == VectorTypes.I8:
+        return 'i1'
+    elif vector_type == VectorTypes.I16:
+        return 'i2'
+    elif vector_type == VectorTypes.I32:
+        return 'i4'
+    elif vector_type == VectorTypes.I64:
+        return 'i8'
+    raise ValueError(f'unsupported element type: {vector_type}')
+
+
+def _vector_type_to_struct_format(
+    vec: Any,
+    vector_type: VectorTypes,
+) -> str:
+    """Convert a vector type to a struct format string."""
+    n = len(vec)
+    if vector_type == VectorTypes.F32:
+        if isinstance(vec, (bytes, bytearray)):
+            n = n // 4
+        return f'<{n}f'
+    elif vector_type == VectorTypes.F64:
+        if isinstance(vec, (bytes, bytearray)):
+            n = n // 8
+        return f'<{n}d'
+    elif vector_type == VectorTypes.I8:
+        return f'<{n}b'
+    elif vector_type == VectorTypes.I16:
+        if isinstance(vec, (bytes, bytearray)):
+            n = n // 2
+        return f'<{n}h'
+    elif vector_type == VectorTypes.I32:
+        if isinstance(vec, (bytes, bytearray)):
+            n = n // 4
+        return f'<{n}i'
+    elif vector_type == VectorTypes.I64:
+        if isinstance(vec, (bytes, bytearray)):
+            n = n // 8
+        return f'<{n}q'
+    raise ValueError(f'unsupported element type: {vector_type}')
+
+
+def unpack_vector(
+    obj: Union[bytes, bytearray],
+    vec_type: VectorTypes = VectorTypes.F32,
+) -> Tuple[Any]:
+    """
+    Unpack a vector from bytes.
+
+    Parameters
+    ----------
+    obj : bytes or bytearray
+        The object to unpack.
+    vec_type : VectorTypes
+        The type of the elements in the vector.
+        Can be one of 'f32', 'f64', 'i8', 'i16', 'i32', or 'i64'.
+        Default is 'f32'.
+
+    Returns
+    -------
+    Tuple[Any]
+        The unpacked vector.
+
+    """
+    return struct.unpack(_vector_type_to_struct_format(obj, vec_type), obj)
+
+
+def pack_vector(
+    obj: Any,
+    vec_type: VectorTypes = VectorTypes.F32,
+) -> bytes:
+    """
+    Pack a vector into bytes.
+
+    Parameters
+    ----------
+    obj : Any
+        The object to pack.
+    vec_type : VectorTypes
+        The type of the elements in the vector.
+        Can be one of 'f32', 'f64', 'i8', 'i16', 'i32', or 'i64'.
+        Default is 'f32'.
+
+    Returns
+    -------
+    bytes
+        The packed vector.
+
+    """
+    if isinstance(obj, (list, tuple)):
+        return struct.pack(_vector_type_to_struct_format(obj, vec_type), *obj)
+
+    if is_numpy(obj):
+        return obj.tobytes()
+
+    if is_pandas_series(obj):
+        import pandas as pd
+        return pd.Series(obj).to_numpy().tobytes()
+
+    if is_polars_series(obj):
+        import polars as pl
+        return pl.Series(obj).to_numpy().tobytes()
+
+    if is_pyarrow_array(obj):
+        import pyarrow as pa
+        return pa.array(obj).to_numpy().tobytes()
+
+    raise ValueError(
+        f'unsupported object type: {type(obj)}',
+    )
+
+
+def unpack_vectors(
+    arr_of_vec: Any,
+    vec_type: VectorTypes = VectorTypes.F32,
+) -> Iterable[Any]:
+    """
+    Unpack a vector from an array of bytes.
+
+    Parameters
+    ----------
+    arr_of_vec : Iterable[Any]
+        The array of bytes to unpack.
+    vec_type : VectorTypes
+        The type of the elements in the vector.
+        Can be one of 'f32', 'f64', 'i8', 'i16', 'i32', or 'i64'.
+        Default is 'f32'.
+
+    Returns
+    -------
+    Iterable[Any]
+        The unpacked vector.
+
+    """
+    if isinstance(arr_of_vec, (list, tuple)):
+        return [unpack_vector(x, vec_type) for x in arr_of_vec]
+
+    import numpy as np
+
+    dtype = _vector_type_to_numpy_type(vec_type)
+
+    np_arr = np.array(
+        [np.frombuffer(x, dtype=dtype) for x in arr_of_vec],
+        dtype=dtype,
+    )
+
+    if is_numpy(arr_of_vec):
+        return np_arr
+
+    if is_pandas_series(arr_of_vec):
+        import pandas as pd
+        return pd.Series(np_arr)
+
+    if is_polars_series(arr_of_vec):
+        import polars as pl
+        return pl.Series(np_arr)
+
+    if is_pyarrow_array(arr_of_vec):
+        import pyarrow as pa
+        return pa.array(np_arr)
+
+    raise ValueError(
+        f'unsupported object type: {type(arr_of_vec)}',
+    )
+
+
+def pack_vectors(
+    arr_of_arr: Iterable[Any],
+    vec_type: VectorTypes = VectorTypes.F32,
+) -> Iterable[Any]:
+    """
+    Pack a vector into an array of bytes.
+
+    Parameters
+    ----------
+    arr_of_arr : Iterable[Any]
+        The array of bytes to pack.
+    vec_type : VectorTypes
+        The type of the elements in the vector.
+        Can be one of 'f32', 'f64', 'i8', 'i16', 'i32', or 'i64'.
+        Default is 'f32'.
+
+    Returns
+    -------
+    Iterable[Any]
+        The array of packed vectors.
+
+    """
+    if isinstance(arr_of_arr, (list, tuple)):
+        if not arr_of_arr:
+            return []
+        fmt = _vector_type_to_struct_format(arr_of_arr[0], vec_type)
+        return [struct.pack(fmt, x) for x in arr_of_arr]
+
+    import numpy as np
+
+    # Use object type because numpy truncates nulls at the end of fixed binary
+    np_arr = np.array([x.tobytes() for x in arr_of_arr], dtype=np.object_)
+
+    if is_numpy(arr_of_arr):
+        return np_arr
+
+    if is_pandas_series(arr_of_arr):
+        import pandas as pd
+        return pd.Series(np_arr)
+
+    if is_polars_series(arr_of_arr):
+        import polars as pl
+        return pl.Series(np_arr)
+
+    if is_pyarrow_array(arr_of_arr):
+        import pyarrow as pa
+        return pa.array(np_arr)
+
+    raise ValueError(
+        f'unsupported object type: {type(arr_of_arr)}',
+    )

singlestoredb/fusion/__init__.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python3
+import importlib
+import os
+
+from .registry import execute
+from .registry import get_handler
+
+# Load all files in handlers directory
+for f in os.listdir(os.path.join(os.path.dirname(__file__), 'handlers')):
+    if f.endswith('.py') and not f.startswith('_'):
+        importlib.import_module(f'singlestoredb.fusion.handlers.{f[:-3]}')