singlestoredb 1.12.4__py3-none-any.whl → 1.13.1__py3-none-any.whl

singlestoredb/functions/typing.py

@@ -0,0 +1,41 @@
+from typing import Any
+from typing import Iterable
+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/utils.py

@@ -0,0 +1,421 @@
+import dataclasses
+import inspect
+import struct
+import sys
+import types
+import typing
+from enum import Enum
+from typing import Any
+from typing import Dict
+from typing import Iterable
+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/http/connection.py

@@ -648,7 +648,9 @@ class Cursor(connection.Cursor):
                         if 'UNSIGNED' in data_type:
                             flags = 32
                         if data_type.endswith('BLOB') or data_type.endswith('BINARY'):
-                            converter = functools.partial(b64decode_converter, converter)
+                            converter = functools.partial(
+                                b64decode_converter, converter,  # type: ignore
+                            )
                             charset = 63  # BINARY
                         if type_code == 0:  # DECIMAL
                             type_code = types.ColumnType.get_code('NEWDECIMAL')

singlestoredb/management/inference_api.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python
+"""SingleStoreDB Cloud Inference API."""
+import os
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from .utils import vars_to_str
+from singlestoredb.exceptions import ManagementError
+from singlestoredb.management.manager import Manager
+
+
+class InferenceAPIInfo(object):
+    """
+    Inference API definition.
+
+    This object is not directly instantiated. It is used in results
+    of API calls on the :class:`InferenceAPIManager`. See :meth:`InferenceAPIManager.get`.
+    """
+
+    service_id: str
+    model_name: str
+    name: str
+    connection_url: str
+    project_id: str
+
+    def __init__(
+        self,
+        service_id: str,
+        model_name: str,
+        name: str,
+        connection_url: str,
+        project_id: str,
+    ):
+        self.service_id = service_id
+        self.connection_url = connection_url
+        self.model_name = model_name
+        self.name = name
+        self.project_id = project_id
+
+    @classmethod
+    def from_dict(
+        cls,
+        obj: Dict[str, Any],
+    ) -> 'InferenceAPIInfo':
+        """
+        Construct a Inference API from a dictionary of values.
+
+        Parameters
+        ----------
+        obj : dict
+            Dictionary of values
+
+        Returns
+        -------
+        :class:`Job`
+
+        """
+        out = cls(
+            service_id=obj['serviceID'],
+            project_id=obj['projectID'],
+            model_name=obj['modelName'],
+            name=obj['name'],
+            connection_url=obj['connectionURL'],
+        )
+        return out
+
+    def __str__(self) -> str:
+        """Return string representation."""
+        return vars_to_str(self)
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return str(self)
+
+
+class InferenceAPIManager(object):
+    """
+    SingleStoreDB Inference APIs manager.
+
+    This class should be instantiated using :attr:`Organization.inference_apis`.
+
+    Parameters
+    ----------
+    manager : InferenceAPIManager, optional
+        The InferenceAPIManager the InferenceAPIManager belongs to
+
+    See Also
+    --------
+    :attr:`InferenceAPI`
+    """
+
+    def __init__(self, manager: Optional[Manager]):
+        self._manager = manager
+        self.project_id = os.environ.get('SINGLESTOREDB_PROJECT')
+
+    def get(self, model_name: str) -> InferenceAPIInfo:
+        if self._manager is None:
+            raise ManagementError(msg='Manager not initialized')
+        res = self._manager._get(f'inferenceapis/{self.project_id}/{model_name}').json()
+        return InferenceAPIInfo.from_dict(res)

singlestoredb/management/manager.py

@@ -62,6 +62,7 @@ class Manager(object):
         )
         if not new_access_token:
             raise ManagementError(msg='No management token was configured.')
+
         self._is_jwt = not access_token and new_access_token and is_jwt(new_access_token)
         self._sess = requests.Session()
         self._sess.headers.update({
@@ -70,10 +71,14 @@ class Manager(object):
             'Accept': 'application/json',
             'User-Agent': f'SingleStoreDB-Python/{client_version}',
         })
+
         self._base_url = urljoin(
-            base_url or type(self).default_base_url,
+            base_url
+            or config.get_option('management.base_url')
+            or type(self).default_base_url,
             version or type(self).default_version,
         ) + '/'
+
         self._params: Dict[str, str] = {}
         if organization_id:
             self._params['organizationID'] = organization_id

singlestoredb/management/organization.py

@@ -7,6 +7,7 @@ from typing import Optional
 from typing import Union
 
 from ..exceptions import ManagementError
+from .inference_api import InferenceAPIManager
 from .job import JobsManager
 from .manager import Manager
 from .utils import vars_to_str
@@ -207,3 +208,19 @@ class Organization(object):
         :class:`JobsManager`
         """
         return JobsManager(self._manager)
+
+    @property
+    def inference_apis(self) -> InferenceAPIManager:
+        """
+        Retrieve a SingleStoreDB inference api manager.
+
+        Parameters
+        ----------
+        manager : WorkspaceManager, optional
+            The WorkspaceManager the InferenceAPIManager belongs to
+
+        Returns
+        -------
+        :class:`InferenceAPIManager`
+        """
+        return InferenceAPIManager(self._manager)

singlestoredb/management/utils.py

@@ -30,7 +30,7 @@ JSONList = List[JSON]
 T = TypeVar('T')
 
 if sys.version_info < (3, 10):
-    PathLike = Union[str, os.PathLike]
+    PathLike = Union[str, os.PathLike]  # type: ignore
     PathLikeABC = os.PathLike
 else:
     PathLike = Union[str, os.PathLike[str]]
@@ -73,7 +73,7 @@ def ttl_property(ttl: datetime.timedelta) -> Callable[[Any], Any]:
     """Property with a time-to-live."""
     def wrapper(func: Callable[[Any], Any]) -> Any:
         out = TTLProperty(func, ttl=ttl)
-        return functools.wraps(func)(out)
+        return functools.wraps(func)(out)  # type: ignore
     return wrapper