quasardb 3.14.2.dev8__cp310-cp310-manylinux_2_26_x86_64.manylinux_2_28_x86_64.whl

quasardb/pandas/__init__.py

@@ -0,0 +1,696 @@
+# pylint: disable=C0103,C0111,C0302,R0903
+
+# Copyright (c) 2009-2024, quasardb SAS. All rights reserved.
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+#    * Redistributions of source code must retain the above copyright
+#      notice, this list of conditions and the following disclaimer.
+#    * Redistributions in binary form must reproduce the above copyright
+#      notice, this list of conditions and the following disclaimer in the
+#      documentation and/or other materials provided with the distribution.
+#    * Neither the name of quasardb nor the names of its contributors may
+#      be used to endorse or promote products derived from this software
+#      without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY QUASARDB AND CONTRIBUTORS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE REGENTS AND CONTRIBUTORS BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+from __future__ import annotations
+
+import logging
+import warnings
+from datetime import timedelta
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+
+import quasardb
+import quasardb.numpy as qdbnp
+import quasardb.table_cache as table_cache
+from quasardb.quasardb import Cluster, Table, Writer
+from quasardb.typing import DType, MaskedArrayAny, Range, RangeSet
+
+logger = logging.getLogger("quasardb.pandas")
+
+
+class PandasRequired(ImportError):
+    """
+    Exception raised when trying to use QuasarDB pandas integration, but
+    pandas has not been installed.
+    """
+
+    pass
+
+
+try:
+    import numpy as np
+    import numpy.ma as ma
+    import pandas as pd
+    from pandas.core.api import DataFrame, Series
+    from pandas.core.base import PandasObject  # type: ignore[attr-defined]
+
+except ImportError:
+    raise PandasRequired("The pandas library is required to handle pandas data formats")
+
+
+# Constant mapping of numpy dtype to QuasarDB column type
+# TODO(leon): support this natively in qdb C api ? we have everything we need
+#             to understand dtypes.
+_dtype_map: Dict[Any, quasardb.ColumnType] = {
+    np.dtype("int64"): quasardb.ColumnType.Int64,
+    np.dtype("int32"): quasardb.ColumnType.Int64,
+    np.dtype("float64"): quasardb.ColumnType.Double,
+    np.dtype("object"): quasardb.ColumnType.String,
+    np.dtype("M8[ns]"): quasardb.ColumnType.Timestamp,
+    np.dtype("datetime64[ns]"): quasardb.ColumnType.Timestamp,
+    "int64": quasardb.ColumnType.Int64,
+    "int32": quasardb.ColumnType.Int64,
+    "float32": quasardb.ColumnType.Double,
+    "float64": quasardb.ColumnType.Double,
+    "timestamp": quasardb.ColumnType.Timestamp,
+    "string": quasardb.ColumnType.String,
+    "bytes": quasardb.ColumnType.Blob,
+    "floating": quasardb.ColumnType.Double,
+    "integer": quasardb.ColumnType.Int64,
+    "bytes": quasardb.ColumnType.Blob,
+    "string": quasardb.ColumnType.String,
+    "datetime64": quasardb.ColumnType.Timestamp,
+}
+
+# Type hint for TableLike parameter
+TableLike = Union[str, Table]
+
+
+def read_series(
+    table: Table, col_name: str, ranges: Optional[RangeSet] = None
+) -> pd.Series:
+    """
+    Read a Pandas Timeseries from a single column.
+
+    Parameters:
+    -----------
+
+    table : quasardb.Timeseries
+      QuasarDB Timeseries table object, e.g. qdb_cluster.table('my_table')
+
+    col_name : str
+      Name of the column to read.
+
+    ranges : list
+      A list of ranges to read, represented as tuples of Numpy datetime64[ns] objects.
+    """
+    read_with = {
+        quasardb.ColumnType.Double: table.double_get_ranges,
+        quasardb.ColumnType.Blob: table.blob_get_ranges,
+        quasardb.ColumnType.String: table.string_get_ranges,
+        quasardb.ColumnType.Int64: table.int64_get_ranges,
+        quasardb.ColumnType.Timestamp: table.timestamp_get_ranges,
+        quasardb.ColumnType.Symbol: table.string_get_ranges,
+    }
+
+    kwargs: Dict[str, Any] = {"column": col_name}
+
+    if ranges is not None:
+        kwargs["ranges"] = ranges
+
+    # Dispatch based on column type
+    t = table.column_type_by_id(col_name)
+
+    logger.info(
+        "reading Series from column %s.%s with type %s", table.get_name(), col_name, t
+    )
+
+    res = (read_with[t])(**kwargs)
+
+    return pd.Series(res[1], index=res[0])
+
+
+def write_series(
+    series: pd.Series,
+    table: Table,
+    col_name: str,
+    infer_types: bool = True,
+    dtype: Optional[DType] = None,
+) -> None:
+    """
+    Writes a Pandas Timeseries to a single column.
+
+    Parameters:
+    -----------
+
+    series : pandas.Series
+      Pandas Series, with a numpy.datetime64[ns] as index. Underlying data will be attempted
+      to be transformed to appropriate QuasarDB type.
+
+    table : quasardb.Timeseries
+      QuasarDB Timeseries table object, e.g. qdb_cluster.table('my_table')
+
+    col_name : str
+      Column name to store data in.
+    """
+
+    logger.debug(
+        "write_series, table=%s, col_name=%s, infer_types=%s, dtype=%s",
+        table.get_name(),
+        col_name,
+        infer_types,
+        dtype,
+    )
+
+    data = None
+    index = None
+
+    data = ma.masked_array(series.to_numpy(copy=False), mask=series.isna())
+
+    if infer_types is True:
+        index = series.index.to_numpy("datetime64[ns]", copy=False)
+    else:
+        index = series.index.to_numpy(copy=False)
+
+    assert data is not None
+    assert index is not None
+
+    return qdbnp.write_array(
+        data=data,
+        index=index,
+        table=table,
+        column=col_name,
+        dtype=dtype,
+        infer_types=infer_types,
+    )
+
+
+def query(
+    cluster: Cluster,
+    query: str,
+    index: Optional[str] = None,
+    blobs: bool = False,
+    numpy: bool = True,
+) -> pd.DataFrame:
+    """
+    Execute *query* and return the result as a pandas DataFrame.
+
+    Parameters
+    ----------
+    cluster : quasardb.Cluster
+        Active connection to the QuasarDB cluster.
+
+    query : str
+        The query to execute.
+
+    index : str | None, default None
+        Column to use as index.  When None a synthetic index is created and
+        named “$index”.
+
+    blobs, numpy
+        DEPRECATED – no longer used.  Supplying a non-default value raises a
+        DeprecationWarning and the argument is ignored.
+    """
+    # ------------------------------------------------------------------ deprecations
+    if blobs is not False:
+        warnings.warn(
+            "`blobs` is deprecated and will be removed in a future version; "
+            "the argument is ignored.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+    if numpy is not True:
+        warnings.warn(
+            "`numpy` is deprecated and will be removed in a future version; "
+            "the argument is ignored.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+    # ------------------------------------------------------------------------------
+
+    logger.debug("querying and returning as DataFrame: %s", query)
+    index_vals, m = qdbnp.query(cluster, query, index=index, dict=True)
+
+    index_name = "$index" if index is None else index
+    index_obj = pd.Index(index_vals, name=index_name)
+
+    return pd.DataFrame(m, index=index_obj)
+
+
+def stream_dataframes(
+    conn: Cluster,
+    tables: List[TableLike],
+    *,
+    batch_size: Optional[int] = 2**16,
+    column_names: Optional[List[str]] = None,
+    ranges: Optional[RangeSet] = None,
+) -> Iterator[pd.DataFrame]:
+    """
+    Read a Pandas Dataframe from a QuasarDB Timeseries table. Returns a generator with dataframes of size `batch_size`, which is useful
+    when traversing a large dataset which does not fit into memory.
+
+    Accepts the same parameters as `stream_dataframes`.
+
+    Parameters:
+    -----------
+
+    conn : quasardb.Cluster
+      Connection to the QuasarDB database.
+
+    tables : list[str | quasardb.Table]
+      QuasarDB tables to stream, as a list of strings or quasardb table objects.
+
+    batch_size : int
+      The amount of rows to fetch in a single read operation. If unset, uses 2^16 (65536) rows
+      as batch size by default.
+
+    column_names : optional list
+      List of columns to read in dataframe. The timestamp column '$timestamp' is
+      always read.
+
+      Defaults to all columns.
+
+    ranges: optional list
+      A list of time ranges to read, represented as tuples of Numpy datetime64[ns] objects.
+      Defaults to the entire table.
+
+    """
+    # Sanitize batch_size
+    if batch_size is None:
+        batch_size = 2**16
+    elif not isinstance(batch_size, int):
+        raise TypeError(
+            "batch_size should be an integer, but got: {} with value {}".format(
+                type(batch_size), str(batch_size)
+            )
+        )
+
+    kwargs: Dict[str, Any] = {"batch_size": batch_size}
+
+    if column_names:
+        kwargs["column_names"] = column_names
+
+    if ranges:
+        kwargs["ranges"] = ranges
+
+    coerce_table_name_fn = lambda x: x if isinstance(x, str) else x.get_name()
+    kwargs["table_names"] = [coerce_table_name_fn(x) for x in tables]
+
+    with conn.reader(**kwargs) as reader:
+        for batch in reader:
+            # We always expect the timestamp column, and set this as the index
+            assert "$timestamp" in batch
+
+            idx = pd.Index(batch.pop("$timestamp"), copy=False, name="$timestamp")
+            df = pd.DataFrame(batch, index=idx)
+
+            yield df
+
+
+def stream_dataframe(
+    conn: Cluster,
+    table: TableLike,
+    *,
+    batch_size: Optional[int] = 2**16,
+    column_names: Optional[List[str]] = None,
+    ranges: Optional[RangeSet] = None,
+) -> Iterator[pd.DataFrame]:
+    """
+    Read a single table and return a stream of dataframes. This is a convenience function that wraps around
+    `stream_dataframes`.
+    """
+    # For backwards compatibility, we drop the `$table` column returned: this is not strictly
+    # necessary, but it also is somewhat reasonable to drop it when we're reading from a single
+    # table, which is the case here.
+    clean_df_fn = lambda df: df.drop(columns=["$table"])
+
+    return (
+        clean_df_fn(df)
+        for df in stream_dataframes(
+            conn,
+            [table],
+            batch_size=batch_size,
+            column_names=column_names,
+            ranges=ranges,
+        )
+    )
+
+
+def read_dataframe(
+    conn: Cluster,
+    table: TableLike,
+    *,
+    batch_size: Optional[int] = 2**16,
+    column_names: Optional[List[str]] = None,
+    ranges: Optional[RangeSet] = None,
+) -> pd.DataFrame:
+    """
+    Read a Pandas Dataframe from a QuasarDB Timeseries table. Wraps around stream_dataframes(), and
+    returns everything as a single dataframe. batch_size is always explicitly set to 0.
+
+
+    Parameters:
+    -----------
+
+    conn : quasardb.Cluster
+      Connection to the QuasarDB database.
+
+    table : str | quasardb.Table
+      QuasarDB table to stream, either as a string or a table object. When re-executing the same function
+      multiple times on the same tables, providing the table as an object has a performance benefit.
+
+    """
+
+    if batch_size is not None and batch_size != 0:
+        logger.warning(
+            "Providing a batch size with read_dataframe is unsupported, overriding batch_size to 65536."
+        )
+        logger.warning(
+            "If you wish to traverse the data in smaller batches, please use: stream_dataframe()."
+        )
+        batch_size = 2**16
+
+    # Note that this is *lazy*, dfs is a generator, not a list -- as such, dataframes will be
+    # fetched on-demand, which means that an error could occur in the middle of processing
+    # dataframes.
+    dfs = stream_dataframe(
+        conn, table, batch_size=batch_size, column_names=column_names, ranges=ranges
+    )
+
+    # if result of stream_dataframe is empty this could result in ValueError on pd.concat()
+    # as stream_dataframe is a generator there is no easy way to check for this condition without evaluation
+    # the most simple way is to catch the ValueError and return an empty DataFrame
+    try:
+        return pd.concat(dfs, copy=False)  #  type: ignore[call-overload]
+    except ValueError as e:
+        logger.error(
+            "Error while concatenating dataframes. This can happen if result set is empty. Returning empty dataframe. Error: %s",
+            e,
+        )
+        return pd.DataFrame()
+
+
+def _extract_columns(
+    df: pd.DataFrame, cinfos: List[Tuple[str, quasardb.ColumnType]]
+) -> Dict[str, MaskedArrayAny]:
+    """
+    Converts dataframe to a number of numpy arrays, one for each column.
+
+    Arrays will be indexed by relative offset, in the same order as the table's columns.
+    If a table column is not present in the dataframe, it it have a None entry.
+    If a dataframe column is not present in the table, it will be ommitted.
+    """
+    ret: Dict[str, MaskedArrayAny] = {}
+
+    # Grab all columns from the DataFrame in the order of table columns,
+    # put None if not present in df.
+    for i in range(len(cinfos)):
+        (cname, _) = cinfos[i]
+
+        if cname in df.columns:
+            arr = df[cname].array
+            ret[cname] = ma.masked_array(arr.to_numpy(copy=False), mask=arr.isna())
+
+    return ret
+
+
+def write_dataframes(
+    dfs: Union[
+        Dict[TableLike, pd.DataFrame],
+        List[tuple[TableLike, pd.DataFrame]],
+    ],
+    cluster: quasardb.Cluster,
+    *,
+    create: bool = False,
+    shard_size: Optional[timedelta] = None,
+    # numpy.write_arrays passthrough options
+    dtype: Optional[
+        Union[DType, Dict[str, Optional[DType]], List[Optional[DType]]]
+    ] = None,
+    push_mode: Optional[quasardb.WriterPushMode] = None,
+    _async: bool = False,
+    fast: bool = False,
+    truncate: Union[bool, Range] = False,
+    truncate_range: Optional[Range] = None,
+    deduplicate: Union[bool, str, List[str]] = False,
+    deduplication_mode: str = "drop",
+    infer_types: bool = True,
+    writer: Optional[Writer] = None,
+    write_through: bool = True,
+    retries: Union[int, quasardb.RetryOptions] = 3,
+    **kwargs: Any,
+) -> List[Table]:
+    """
+    Store dataframes into a table. Any additional parameters not documented here
+    are passed to numpy.write_arrays(). Please consult the pydoc of that function
+    for additional accepted parameters.
+
+    Parameters:
+    -----------
+
+    dfs: dict[str | quasardb.Table, pd.DataFrame] | list[tuple[str | quasardb.Table, pd.DataFrame]]
+      This can be either a dict that maps table (either objects or names) to a dataframe, or a list
+      of table<>dataframe tuples.
+
+    cluster: quasardb.Cluster
+      Active connection to the QuasarDB cluster
+
+    create: optional bool
+      Whether to create the table. Defaults to False.
+
+    shard_size: optional datetime.timedelta
+      The shard size of the timeseries you wish to create when `create` is True.
+    """
+
+    # If dfs is a dict, we convert it to a list of tuples.
+    if isinstance(dfs, dict):
+        dfs = list(dfs.items())
+
+    if shard_size is not None and create == False:
+        raise ValueError("Invalid argument: shard size provided while create is False")
+
+    # If the tables are provided as strings, we look them up.
+    dfs_ = []
+    for table, df in dfs:
+        if isinstance(table, str):
+            table = table_cache.lookup(table, cluster)
+
+        dfs_.append((table, df))
+
+    data_by_table = []
+
+    for table, df in dfs_:
+        logger.debug("quasardb.pandas.write_dataframe, create = %s", create)
+        assert isinstance(df, pd.DataFrame)
+
+        # Create table if requested
+        if create:
+            _create_table_from_df(df, table, shard_size)
+
+        cinfos = [(x.name, x.type) for x in table.list_columns()]
+
+        if not df.index.is_monotonic_increasing:
+            logger.warning(
+                "dataframe index is unsorted, resorting dataframe based on index"
+            )
+            df = df.sort_index().reindex()
+
+        # We pass everything else to our qdbnp.write_arrays function, as generally speaking
+        # it is (much) more sensible to deal with numpy arrays than Pandas dataframes:
+        # pandas has the bad habit of wanting to cast data to different types if your data
+        # is sparse, most notably forcing sparse integer arrays to floating points.
+
+        data = _extract_columns(df, cinfos)
+        data["$timestamp"] = ma.masked_array(
+            df.index.to_numpy(copy=False, dtype="datetime64[ns]")
+        )  # We cast to masked_array to enforce typing compliance
+
+        data_by_table.append((table, data))
+
+    kwargs["deprecation_stacklevel"] = kwargs.get("deprecation_stacklevel", 1) + 1
+    return qdbnp.write_arrays(
+        data_by_table,
+        cluster,
+        table=None,
+        index=None,
+        dtype=dtype,
+        push_mode=push_mode,
+        _async=_async,
+        fast=fast,
+        truncate=truncate,
+        truncate_range=truncate_range,
+        deduplicate=deduplicate,
+        deduplication_mode=deduplication_mode,
+        infer_types=infer_types,
+        writer=writer,
+        write_through=write_through,
+        retries=retries,
+        **kwargs,
+    )
+
+
+def write_dataframe(
+    df: pd.DataFrame,
+    cluster: quasardb.Cluster,
+    table: TableLike,
+    *,
+    create: bool = False,
+    shard_size: Optional[timedelta] = None,
+    # numpy.write_arrays passthrough options
+    dtype: Optional[
+        Union[DType, Dict[str, Optional[DType]], List[Optional[DType]]]
+    ] = None,
+    push_mode: Optional[quasardb.WriterPushMode] = None,
+    _async: bool = False,
+    fast: bool = False,
+    truncate: Union[bool, Range] = False,
+    truncate_range: Optional[Range] = None,
+    deduplicate: Union[bool, str, List[str]] = False,
+    deduplication_mode: str = "drop",
+    infer_types: bool = True,
+    writer: Optional[Writer] = None,
+    write_through: bool = True,
+    retries: Union[int, quasardb.RetryOptions] = 3,
+    **kwargs: Any,
+) -> List[Table]:
+    """
+    Store a single dataframe into a table. Takes the same arguments as `write_dataframes`, except only
+    a single df/table combination.
+    """
+    kwargs["deprecation_stacklevel"] = kwargs.get("deprecation_stacklevel", 1) + 1
+    return write_dataframes(
+        [(table, df)],
+        cluster,
+        create=create,
+        shard_size=shard_size,
+        dtype=dtype,
+        push_mode=push_mode,
+        _async=_async,
+        fast=fast,
+        truncate=truncate,
+        truncate_range=truncate_range,
+        deduplicate=deduplicate,
+        deduplication_mode=deduplication_mode,
+        infer_types=infer_types,
+        writer=writer,
+        write_through=write_through,
+        retries=retries,
+        **kwargs,
+    )
+
+
+def write_pinned_dataframe(
+    df: pd.DataFrame,
+    cluster: quasardb.Cluster,
+    table: TableLike,
+    *,
+    create: bool = False,
+    shard_size: Optional[timedelta] = None,
+    # numpy.write_arrays passthrough options
+    dtype: Optional[
+        Union[DType, Dict[str, Optional[DType]], List[Optional[DType]]]
+    ] = None,
+    push_mode: Optional[quasardb.WriterPushMode] = None,
+    _async: bool = False,
+    fast: bool = False,
+    truncate: Union[bool, Range] = False,
+    truncate_range: Optional[Range] = None,
+    deduplicate: Union[bool, str, List[str]] = False,
+    deduplication_mode: str = "drop",
+    infer_types: bool = True,
+    writer: Optional[Writer] = None,
+    write_through: bool = True,
+    retries: Union[int, quasardb.RetryOptions] = 3,
+    **kwargs: Any,
+) -> List[Table]:
+    """
+    Legacy wrapper around write_dataframe()
+    """
+    logger.warning(
+        "write_pinned_dataframe is deprecated and will be removed in a future release."
+    )
+    logger.warning("Please use write_dataframe directly instead")
+    kwargs["deprecation_stacklevel"] = 2
+    return write_dataframe(
+        df,
+        cluster,
+        table,
+        create=create,
+        shard_size=shard_size,
+        dtype=dtype,
+        push_mode=push_mode,
+        _async=_async,
+        fast=fast,
+        truncate=truncate,
+        truncate_range=truncate_range,
+        deduplicate=deduplicate,
+        deduplication_mode=deduplication_mode,
+        infer_types=infer_types,
+        writer=writer,
+        write_through=write_through,
+        retries=retries,
+        **kwargs,
+    )
+
+
+def _create_table_from_df(
+    df: pd.DataFrame, table: Table, shard_size: Optional[timedelta] = None
+) -> Table:
+    cols = list()
+
+    dtypes = _get_inferred_dtypes(df)
+
+    logger.info("got inferred dtypes: %s", dtypes)
+    for c in df.columns:
+        dt = dtypes[c]
+        ct = _dtype_to_column_type(df[c].dtype, dt)
+        logger.debug(
+            "probed pandas dtype %s to inferred dtype %s and map to quasardb column type %s",
+            df[c].dtype,
+            dt,
+            ct,
+        )
+        cols.append(quasardb.ColumnInfo(ct, c))
+
+    try:
+        if not shard_size:
+            table.create(cols)
+        else:
+            table.create(cols, shard_size)
+    except quasardb.AliasAlreadyExistsError:
+        # TODO(leon): warn? how?
+        pass
+
+    return table
+
+
+def _dtype_to_column_type(dt: Any, inferred: Any) -> quasardb.ColumnType:
+    res = _dtype_map.get(inferred, None)
+    if res is None:
+        res = _dtype_map.get(dt, None)
+
+    if res is None:
+        raise ValueError("Incompatible data type: ", dt)
+
+    return res
+
+
+def _get_inferred_dtypes(df: pd.DataFrame) -> Dict[str, str]:
+    dtypes = {}
+    for i in range(len(df.columns)):
+        c = df.columns[i]
+        dt = pd.api.types.infer_dtype(df[c].values)
+        logger.debug("Determined dtype of column %s to be %s", c, dt)
+        dtypes[c] = dt
+    return dtypes
+
+
+def _get_inferred_dtypes_indexed(df: pd.DataFrame) -> List[str]:
+    dtypes = _get_inferred_dtypes(df)
+    # Performance improvement: avoid a expensive dict lookups by indexing
+    # the column types by relative offset within the df.
+    return list(dtypes[c] for c in df.columns)