singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/utils/results.py

@@ -0,0 +1,585 @@
+#!/usr/bin/env python
+"""SingleStoreDB package utilities."""
+import collections
+import warnings
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import NamedTuple
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+from .dtypes import NUMPY_TYPE_MAP
+from .dtypes import POLARS_TYPE_MAP
+from .dtypes import PYARROW_TYPE_MAP
+
+UNSIGNED_FLAG = 32
+BINARY_FLAG = 128
+
+try:
+    has_numpy = True
+    import numpy as np
+except ImportError:
+    has_numpy = False
+
+try:
+    has_pandas = True
+    import pandas as pd
+except ImportError:
+    has_pandas = False
+
+try:
+    has_polars = True
+    import polars as pl
+except ImportError:
+    has_polars = False
+
+try:
+    has_pyarrow = True
+    import pyarrow as pa
+except ImportError:
+    has_pyarrow = False
+
+DBAPIResult = Union[List[Tuple[Any, ...]], Tuple[Any, ...]]
+OneResult = Union[
+    Tuple[Any, ...], Dict[str, Any],
+    'np.ndarray', 'pd.DataFrame', 'pl.DataFrame', 'pa.Table',
+]
+ManyResult = Union[
+    List[Tuple[Any, ...]], List[Dict[str, Any]],
+    'np.ndarray', 'pd.DataFrame', 'pl.DataFrame', 'pa.Table',
+]
+Result = Union[OneResult, ManyResult]
+
+
+class Description(NamedTuple):
+    """Column definition."""
+    name: str
+    type_code: int
+    display_size: Optional[int]
+    internal_size: Optional[int]
+    precision: Optional[int]
+    scale: Optional[int]
+    null_ok: Optional[bool]
+    flags: Optional[int]
+    charset: Optional[int]
+
+
+if has_numpy:
+    # If an int column is nullable, we need to use floats rather than
+    # ints for numpy and pandas.
+    NUMPY_TYPE_MAP_CAST_FLOAT = NUMPY_TYPE_MAP.copy()
+    NUMPY_TYPE_MAP_CAST_FLOAT.update({
+        1: np.float32,  # Tiny
+        -1: np.float32,  # Unsigned Tiny
+        2: np.float32,  # Short
+        -2: np.float32,  # Unsigned Short
+        3: np.float64,  # Long
+        -3: np.float64,  # Unsigned Long
+        8: np.float64,  # LongLong
+        -8: np.float64,  # Unsigned LongLong
+        9: np.float64,  # Int24
+        -9: np.float64,  # Unsigned Int24
+        13: np.float64,  # Year
+    })
+
+if has_polars:
+    # Remap date/times to strings; let polars do the parsing
+    POLARS_TYPE_MAP = POLARS_TYPE_MAP.copy()
+    POLARS_TYPE_MAP.update({
+        7: pl.Utf8,
+        10: pl.Utf8,
+        12: pl.Utf8,
+    })
+
+
+INT_TYPES = set([1, 2, 3, 8, 9])
+CHAR_TYPES = set([15, 249, 250, 251, 252, 253, 254])
+DECIMAL_TYPES = set([0, 246])
+
+
+def signed(desc: Description) -> int:
+    if ((desc.flags or 0) & UNSIGNED_FLAG and desc.type_code in INT_TYPES) or \
+            (desc.charset == 63 and desc.type_code in CHAR_TYPES):
+        return -desc.type_code
+    return desc.type_code
+
+
+def _description_to_numpy_schema(desc: List[Description]) -> Dict[str, Any]:
+    """Convert description to numpy array schema info."""
+    if has_numpy:
+        return dict(
+            dtype=[
+                (
+                    x.name,
+                    NUMPY_TYPE_MAP_CAST_FLOAT[signed(x)]
+                    if x.null_ok else NUMPY_TYPE_MAP[signed(x)],
+                )
+                for x in desc
+            ],
+        )
+    return {}
+
+
+def _description_to_pandas_schema(desc: List[Description]) -> Dict[str, Any]:
+    """Convert description to pandas DataFrame schema info."""
+    if has_pandas:
+        return dict(columns=[x.name for x in desc])
+    return {}
+
+
+def _decimalize_polars(desc: Description) -> 'pl.Decimal':
+    return pl.Decimal(desc.precision or 10, desc.scale or 0)
+
+
+def _description_to_polars_schema(desc: List[Description]) -> Dict[str, Any]:
+    """Convert description to polars DataFrame schema info."""
+    if has_polars:
+        with_columns = {}
+        for x in desc:
+            if x.type_code in [7, 12]:
+                if x.scale == 6:
+                    with_columns[x.name] = pl.col(x.name).str.to_datetime(
+                        '%Y-%m-%d %H:%M:%S.%6f', time_unit='us',
+                    )
+                else:
+                    with_columns[x.name] = pl.col(x.name).str.to_datetime(
+                        '%Y-%m-%d %H:%M:%S', time_unit='us',
+                    )
+            elif x.type_code == 10:
+                with_columns[x.name] = pl.col(x.name).str.to_date('%Y-%m-%d')
+
+        return dict(
+            schema=dict(
+                schema=[
+                    (
+                        x.name, _decimalize_polars(x)
+                        if x.type_code in DECIMAL_TYPES else POLARS_TYPE_MAP[signed(x)],
+                    )
+                    for x in desc
+                ],
+            ),
+            with_columns=with_columns,
+        )
+    return {}
+
+
+def _decimalize_arrow(desc: Description) -> 'pa.Decimal128':
+    return pa.decimal128(desc.precision or 10, desc.scale or 0)
+
+
+def _description_to_arrow_schema(desc: List[Description]) -> Dict[str, Any]:
+    """Convert description to Arrow Table schema info."""
+    if has_pyarrow:
+        return dict(
+            schema=pa.schema([
+                (
+                    x.name, _decimalize_arrow(x)
+                    if x.type_code in DECIMAL_TYPES else PYARROW_TYPE_MAP[signed(x)],
+                )
+                for x in desc
+            ]),
+        )
+    return {}
+
+
+def results_to_numpy(
+    desc: List[Description],
+    res: Optional[DBAPIResult],
+    single: Optional[bool] = False,
+    schema: Optional[Dict[str, Any]] = None,
+) -> Optional[Result]:
+    """
+    Convert results to numpy.
+
+    Parameters
+    ----------
+    desc : list of Descriptions
+        The column metadata
+    res : tuple or list of tuples
+        The query results
+    single : bool, optional
+        Is this a single result (i.e., from `fetchone`)?
+    schema : Dict[str, Any], optional
+        Cached schema for current output format
+
+    Returns
+    -------
+    numpy.array
+        If `numpy` is available
+    tuple or list of tuples
+        If `numpy` is not available
+
+    """
+    if not res:
+        return res
+    if has_numpy:
+        schema = _description_to_numpy_schema(desc) if schema is None else schema
+        if single:
+            return np.array([res], **schema)
+        return np.array(list(res), **schema)
+    warnings.warn(
+        'numpy is not available; unable to convert to array',
+        RuntimeWarning,
+    )
+    return res
+
+
+def results_to_pandas(
+    desc: List[Description],
+    res: Optional[DBAPIResult],
+    single: Optional[bool] = False,
+    schema: Optional[Dict[str, Any]] = None,
+) -> Optional[Result]:
+    """
+    Convert results to pandas.
+
+    Parameters
+    ----------
+    desc : list of Descriptions
+        The column metadata
+    res : tuple or list of tuples
+        The query results
+    single : bool, optional
+        Is this a single result (i.e., from `fetchone`)?
+    schema : Dict[str, Any], optional
+        Cached schema for current output format
+
+    Returns
+    -------
+    DataFrame
+        If `pandas` is available
+    tuple or list of tuples
+        If `pandas` is not available
+
+    """
+    if not res:
+        return res
+    if has_pandas:
+        schema = _description_to_pandas_schema(desc) if schema is None else schema
+        return pd.DataFrame(results_to_numpy(desc, res, single=single, schema=schema))
+    warnings.warn(
+        'pandas is not available; unable to convert to DataFrame',
+        RuntimeWarning,
+    )
+    return res
+
+
+def results_to_polars(
+    desc: List[Description],
+    res: Optional[DBAPIResult],
+    single: Optional[bool] = False,
+    schema: Optional[Dict[str, Any]] = None,
+) -> Optional[Result]:
+    """
+    Convert results to polars.
+
+    Parameters
+    ----------
+    desc : list of Descriptions
+        The column metadata
+    res : tuple or list of tuples
+        The query results
+    single : bool, optional
+        Is this a single result (i.e., from `fetchone`)?
+    schema : Dict[str, Any], optional
+        Cached schema for current output format
+
+    Returns
+    -------
+    DataFrame
+        If `polars` is available
+    tuple or list of tuples
+        If `polars` is not available
+
+    """
+    if not res:
+        return res
+    if has_polars:
+        schema = _description_to_polars_schema(desc) if schema is None else schema
+        if single:
+            out = pl.DataFrame([res], orient='row', **schema.get('schema', {}))
+        else:
+            out = pl.DataFrame(res, orient='row', **schema.get('schema', {}))
+        with_columns = schema.get('with_columns')
+        if with_columns:
+            return out.with_columns(**with_columns)
+        return out
+    warnings.warn(
+        'polars is not available; unable to convert to DataFrame',
+        RuntimeWarning,
+    )
+    return res
+
+
+def results_to_arrow(
+    desc: List[Description],
+    res: Optional[DBAPIResult],
+    single: Optional[bool] = False,
+    schema: Optional[Dict[str, Any]] = None,
+) -> Optional[Result]:
+    """
+    Convert results to Arrow.
+
+    Parameters
+    ----------
+    desc : list of Descriptions
+        The column metadata
+    res : tuple or list of tuples
+        The query results
+    single : bool, optional
+        Is this a single result (i.e., from `fetchone`)?
+    schema : Dict[str, Any], optional
+        Cached schema for current output format
+
+    Returns
+    -------
+    Table
+        If `pyarrow` is available
+    tuple or list of tuples
+        If `pyarrow` is not available
+
+    """
+    if not res:
+        return res
+    if has_pyarrow:
+        names = [x[0] for x in desc]
+        schema = _description_to_arrow_schema(desc) if schema is None else schema
+        if single:
+            if isinstance(res, dict):
+                return pa.Table.from_pylist([res], **schema)
+            else:
+                return pa.Table.from_pylist([dict(zip(names, res))], **schema)
+        if isinstance(res[0], dict):
+            return pa.Table.from_pylist(res, **schema)
+        else:
+            return pa.Table.from_pylist([dict(zip(names, x)) for x in res], **schema)
+    warnings.warn(
+        'pyarrow is not available; unable to convert to Table',
+        RuntimeWarning,
+    )
+    return res
+
+
+def results_to_namedtuple(
+    desc: List[Description],
+    res: Optional[DBAPIResult],
+    single: Optional[bool] = False,
+    schema: Optional[Dict[str, Any]] = None,
+) -> Optional[Result]:
+    """
+    Convert results to namedtuples.
+
+    Parameters
+    ----------
+    desc : list of Descriptions
+        The column metadata
+    res : tuple or list of tuples
+        The query results
+    single : bool, optional
+        Is this a single result (i.e., from `fetchone`)?
+    schema : Dict[str, Any], optional
+        Cached schema for current output format
+
+    Returns
+    -------
+    namedtuple
+        If single is True
+    list of namedtuples
+        If single is False
+
+    """
+    if not res:
+        return res
+    tup = collections.namedtuple(  # type: ignore
+        'Row', list(
+            [x[0] for x in desc],
+        ), rename=True,
+    )
+    if single:
+        return tup(*res)
+    return [tup(*x) for x in res]
+
+
+def results_to_dict(
+    desc: List[Description],
+    res: Optional[DBAPIResult],
+    single: Optional[bool] = False,
+    schema: Optional[Dict[str, Any]] = None,
+) -> Optional[Result]:
+    """
+    Convert results to dicts.
+
+    Parameters
+    ----------
+    desc : list of Descriptions
+        The column metadata
+    res : tuple or list of tuples
+        The query results
+    single : bool, optional
+        Is this a single result (i.e., from `fetchone`)?
+    schema : Dict[str, Any], optional
+        Cached schema for current output format
+
+    Returns
+    -------
+    dict
+        If single is True
+    list of dicts
+        If single is False
+
+    """
+    if not res:
+        return res
+    names = [x[0] for x in desc]
+    if single:
+        return dict(zip(names, res))
+    return [dict(zip(names, x)) for x in res]
+
+
+def results_to_tuple(
+    desc: List[Description],
+    res: Optional[DBAPIResult],
+    single: Optional[bool] = False,
+    schema: Optional[Dict[str, Any]] = None,
+) -> Optional[Result]:
+    """
+    Convert results to tuples.
+
+    Parameters
+    ----------
+    desc : list of Descriptions
+        The column metadata
+    res : tuple or list of tuples
+        The query results
+    single : bool, optional
+        Is this a single result (i.e., from `fetchone`)?
+    schema : Dict[str, Any], optional
+        Cached schema for current output format
+
+    Returns
+    -------
+    tuple
+        If single is True
+    list of tuples
+        If single is False
+
+    """
+    if not res:
+        return res
+    if single:
+        if type(res) is tuple:
+            return res
+        return tuple(res)
+    if type(res[0]) is tuple:
+        return list(res)
+    return [tuple(x) for x in res]
+
+
+def _no_schema(desc: List[Description]) -> Optional[Dict[str, Any]]:
+    return {}
+
+
+_converters: Dict[
+    str, Callable[
+        [
+            List[Description], Optional[DBAPIResult],
+            Optional[bool], Optional[Dict[str, Any]],
+        ],
+        Optional[Result],
+    ],
+] = {
+    'tuple': results_to_tuple,
+    'tuples': results_to_tuple,
+    'namedtuple': results_to_namedtuple,
+    'namedtuples': results_to_namedtuple,
+    'dict': results_to_dict,
+    'dicts': results_to_dict,
+    'numpy': results_to_numpy,
+    'pandas': results_to_pandas,
+    'polars': results_to_polars,
+    'arrow': results_to_arrow,
+    'pyarrow': results_to_arrow,
+}
+
+_schema_converters: Dict[
+    str, Callable[
+        [List[Description]], Optional[Dict[str, Any]],
+    ],
+] = {
+    'tuple': _no_schema,
+    'tuples': _no_schema,
+    'namedtuple': _no_schema,
+    'namedtuples': _no_schema,
+    'dict': _no_schema,
+    'dicts': _no_schema,
+    'structsequence': _no_schema,
+    'structsequences': _no_schema,
+    'numpy': _description_to_numpy_schema,
+    'pandas': _description_to_numpy_schema,
+    'polars': _description_to_polars_schema,
+    'arrow': _description_to_arrow_schema,
+    'pyarrow': _description_to_arrow_schema,
+}
+
+
+def format_results(
+    format: str,
+    desc: List[Description],
+    res: Optional[DBAPIResult],
+    single: bool = False,
+    schema: Optional[Dict[str, Any]] = None,
+) -> Optional[Result]:
+    """
+    Convert results to format specified in the package options.
+
+    Parameters
+    ----------
+    format : str
+        Name of the format type
+    desc : list of Descriptions
+        The column metadata
+    res : tuple or list of tuples
+        The query results
+    single : bool, optional
+        Is this a single result (i.e., from `fetchone`)?
+    schema : Dict[str, Any], optional
+        Cached schema for current output format
+
+    Returns
+    -------
+    list of (named)tuples, list of dicts or DataFrame
+        If single is False
+    (named)tuple, dict, or DataFrame
+        If single is True
+
+    """
+    return _converters[format](desc, res, single, schema)
+
+
+def get_schema(
+    format: str,
+    desc: List[Description],
+) -> Dict[str, Any]:
+    """
+    Convert a DB-API description to a format schema.
+
+    Parameters
+    ----------
+    format : str
+        Name of the format type
+    desc : list of Descriptions
+        The column metadata
+
+    Returns
+    -------
+    Dict[str, Any]
+        A dictionary of function parameters containing schema information
+        for the given format type
+
+    """
+    if format in _schema_converters:
+        return _schema_converters[format](desc) or {}
+    return {}