quasardb 3.14.2.dev4__cp311-cp311-macosx_14_0_x86_64.whl

quasardb/numpy/__init__.py

@@ -0,0 +1,901 @@
+# pylint: disable=C0103,C0111,C0302,R0903
+
+# Copyright (c) 2009-2022, 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.
+#
+
+import logging
+import time
+
+import quasardb
+import quasardb.table_cache as table_cache
+
+logger = logging.getLogger('quasardb.numpy')
+
+
+class NumpyRequired(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
+
+except ImportError as err:
+    logger.exception(err)
+    raise NumpyRequired(
+        "The numpy library is required to handle numpy arrays formats"
+        ) from err
+
+
+class IncompatibleDtypeError(TypeError):
+    """
+    Exception raised when a provided dtype is not the expected dtype.
+    """
+
+    def __init__(self, cname=None, ctype=None, expected=None, provided=None):
+        self.cname = cname
+        self.ctype = ctype
+        self.expected = expected
+        self.provided = provided
+        super().__init__(self.msg())
+
+    def msg(self):
+        return "Data for column '{}' with type '{}' was provided in dtype '{}' but need '{}'.".format(self.cname, self.ctype, self.provided, self.expected)
+
+
+class IncompatibleDtypeErrors(TypeError):
+    """
+    Wraps multiple dtype errors
+    """
+    def __init__(self, xs):
+        self.xs = xs
+        super().__init__(self.msg())
+
+    def msg(self):
+        return "\n".join(x.msg() for x in self.xs)
+
+class InvalidDataCardinalityError(ValueError):
+    """
+    Raised when the provided data arrays doesn't match the table's columns.
+    """
+    def __init__(self, data, cinfos):
+        self.data = data
+        self.cinfos = cinfos
+        super().__init__(self.msg())
+
+    def msg(self):
+        return "Provided data array length '{}' exceeds amount of table columns '{}', unable to map data to columns".format(len(self.data), len(self.cinfos))
+
+
+# Based on QuasarDB column types, which dtype do we accept?
+# First entry will always be the 'preferred' dtype, other ones
+# those that we can natively convert in native code.
+_ctype_to_dtype = {
+    quasardb.ColumnType.String: [np.dtype('U')],
+    quasardb.ColumnType.Symbol: [np.dtype('U')],
+    quasardb.ColumnType.Int64: [np.dtype('i8'), np.dtype('i4'), np.dtype('i2')],
+    quasardb.ColumnType.Double: [np.dtype('f8'), np.dtype('f4')],
+    quasardb.ColumnType.Blob: [np.dtype('S'), np.dtype('O')],
+    quasardb.ColumnType.Timestamp: [np.dtype('datetime64[ns]')]
+}
+
+
+def _best_dtype_for_ctype(ctype: quasardb.quasardb.ColumnType):
+    """
+    Returns the 'best' DType for a certain column type. For example, for blobs, even
+    though we accept py::bytes, prefer bytestrings (as they are faster to read in c++).
+    """
+    possible_dtypes = _ctype_to_dtype[ctype]
+    assert len(possible_dtypes) > 0
+
+    # By convention, the first entry is the preferred one
+    return possible_dtypes[0]
+
+
+def _coerce_dtype(dtype, columns):
+    if dtype is None:
+        dtype = [None] * len(columns)
+
+    if isinstance(dtype, np.dtype):
+        dtype = [dtype]
+
+    if type(dtype) is dict:
+
+        # Conveniently look up column index by label
+        offsets = {}
+        for i in range(len(columns)):
+            (cname, ctype) = columns[i]
+            offsets[cname] = i
+
+        # Now convert the provided dtype dict to a list that matches
+        # the relative offset within the table.
+        #
+        # Any columns not provided will have a 'None' dtype.
+        dtype_ = [None] * len(columns)
+
+        for (k, dt) in dtype.items():
+            if not k in offsets:
+                logger.warn("Forced dtype provided for column '%s' = %s, but that column is not found in the table. Skipping...", k, )
+
+            i = offsets[k]
+            dtype_[i] = dt
+
+        dtype = dtype_
+
+    if type(dtype) is not list:
+        raise ValueError("Forced dtype argument provided, but the argument has an incompatible type. Expected: list-like or dict-like, got: {}".format(type(dtype)))
+
+    if len(dtype) is not len(columns):
+        raise ValueError("Expected exactly one dtype for each column, but %d dtypes were provided for %d columns".format(len(dtype), len(columns)))
+
+    return dtype
+
+
+
+
+def _add_desired_dtypes(dtype, columns):
+    """
+    When infer_types=True, this function sets the 'desired' dtype for each of the columns.
+    `dtype` is expected to be the output of `_coerce_dtype`, that is, a list-like with an
+    entry for each column.
+    """
+    assert len(dtype) == len(columns)
+
+    for i in range(len(dtype)):
+        # No dtype explicitly provided by the user, otherwise we don't touch it
+        if dtype[i] is None:
+            (cname, ctype) = columns[i]
+            dtype_ = _best_dtype_for_ctype(ctype)
+            logger.debug("using default dtype '%s' for column '%s' with type %s", dtype_, cname, ctype)
+            dtype[i] = dtype_
+
+    return dtype
+
+
+def _is_all_masked(xs):
+    if ma.isMA(xs):
+        return ma.size(xs) == ma.count_masked(xs)
+
+    if np.size(xs) == 0:
+        # Empty lists are considered "masked"
+        return True
+
+    if xs.dtype == np.object_ and xs[0] is None:
+        # Very likely that all is None at this point, we can try a more "expensive"
+        # probing function
+        #
+        # We'll defer to the Python `all` function, as Numpy doesn't really have great
+        # built-ins for object arrays
+        return all(x is None for x in xs)
+
+
+    logger.debug("{} is not a masked array, not convertible to requested type... ".format(type(xs)))
+
+    # This array is *not* a masked array, it's *not* convertible to the type we want,
+    # and it's *not* an object array.
+    #
+    # The best we can do at this point is emit a warning and defer to the (expensive)
+    # python-based function as well.
+    return all(x is None for x in xs)
+
+
+def dtypes_equal(lhs, rhs):
+    if lhs.kind == 'U' or lhs.kind == 'S':
+        # Unicode and string data has variable length encoding, which means their itemsize
+        # can be anything.
+        #
+        # We only care about dtype kind in this case.
+        return lhs.kind == rhs.kind
+
+    return lhs == rhs
+
+
+def _dtype_found(needle, haystack):
+    """
+    Returns True if one of the dtypes in `haystack` matches that of `needle`.
+    """
+    for x in haystack:
+        if dtypes_equal(needle, x) is True:
+            return True
+
+    return False
+
+
+def _validate_dtypes(data, columns):
+    errors = list()
+
+    def _error_to_msg(e):
+        (cname, ctype, provided_dtype, expected_dtype) = e
+        return
+
+    for (data_, (cname, ctype)) in zip(data, columns):
+        expected_ = _ctype_to_dtype[ctype]
+
+        logger.debug("data_.dtype = %s, expected_ = %s", data_.dtype, expected_)
+
+        if _dtype_found(data_.dtype, expected_) == False:
+            errors.append(IncompatibleDtypeError(cname=cname, ctype=ctype, provided=data_.dtype, expected=expected_))
+
+    if len(errors) > 0:
+        raise IncompatibleDtypeErrors(errors)
+
+def _coerce_deduplicate(deduplicate, deduplication_mode, columns):
+    """
+    Throws an error when 'deduplicate' options are incorrect.
+    """
+    cnames = [cname for (cname, ctype) in columns]
+
+    if deduplication_mode not in ['drop', 'upsert']:
+        raise RuntimeError("deduplication_mode should be one of ['drop', 'upsert'], got: {}".format(deduplication_mode))
+
+    if isinstance(deduplicate, bool):
+        return deduplicate
+
+    # Special value of $timestamp, hardcoded
+    if isinstance(deduplicate, str) and deduplicate == '$timestamp':
+        deduplicate = ['$timestamp']
+    cnames.append('$timestamp')
+
+    if not isinstance(deduplicate, list):
+        raise TypeError("drop_duplicates should be either a bool or a list, got: " + type(deduplicate))
+
+    for column_name in deduplicate:
+        if not column_name in cnames:
+            raise RuntimeError("Provided deduplication column name '{}' not found in table columns.".format(column_name))
+
+    return deduplicate
+
+def _clean_nulls(xs, dtype):
+    """
+    Numpy's masked arrays have a downside that in case they're not able to convert a (masked!) value to
+    the desired dtype, they raise an error. So, for example, if I have a masked array of objects that
+    look like this
+
+    xs: [1.234 <pd.NA> 5.678]
+    mask: [1 0 1]
+
+    even though pd.NA is not "visible", because it cannot be converted to a float(), the operation will
+    fail!
+
+    This function fixes this by replacing the null values with an acceptable value that can always be
+    converted to the desired dtype.
+    """
+
+    assert ma.isMA(xs)
+
+    if xs.dtype is not np.dtype('object'):
+        return xs
+
+    fill_value = None
+    if dtype == np.float64 or dtype == np.float32 or dtype == np.float16:
+        fill_value = float('nan')
+    elif dtype == np.int64 or dtype == np.int32 or dtype == np.int16:
+        fill_value = -1
+    elif dtype == np.dtype('datetime64[ns]'):
+        fill_value = np.datetime64('nat')
+
+    mask = xs.mask
+    xs_ = xs.filled(fill_value)
+
+    return ma.array(xs_, mask=mask)
+
+
+
+def _coerce_data(data, dtype):
+    """
+    Coerces each numpy array of `data` to the dtype present in `dtype`.
+    """
+
+    assert len(data) == len(dtype)
+
+    for i in range(len(data)):
+        dtype_ = dtype[i]
+        data_ = data[i]
+
+        if dtype_ is not None and dtypes_equal(data_.dtype, dtype_) == False:
+            data_ = _clean_nulls(data_, dtype_)
+
+            assert ma.isMA(data_)
+
+            logger.debug("data for column with offset %d was provided in dtype '%s', but need '%s': converting data...", i, data_.dtype, dtype_)
+
+            logger.debug("dtype of data[%d] before: %s", i, data_.dtype)
+            logger.debug("type of data[%d] after: %s", i, type(data_))
+            logger.debug("size of data[%d] after: %s", i, ma.size(data_))
+            logger.debug("data of data[%d] after: %s", i, data_)
+
+            try:
+                data[i] = data_.astype(dtype_)
+            except TypeError as err:
+                # One 'bug' is that, if everything is masked, the underlying data type can be
+                # pretty much anything.
+                if _is_all_masked(data_):
+                    logger.debug("array completely empty, re-initializing to empty array of '%s'", dtype_)
+                    data[i] = ma.masked_all(ma.size(data_),
+                                            dtype=dtype_)
+
+                # Another 'bug' is that when the input data is objects, we may have null-like values (like pd.NA)
+                # that cannot easily be converted to, say, float.
+                else:
+                    logger.error("An error occured while coercing input data type from dtype '%s' to dtype '%s': ", data_.dtype, dtype_)
+                    logger.exception(err)
+                    raise err
+
+            assert data[i].dtype.kind == dtype_.kind
+
+            logger.debug("type of data[%d] after: %s", i, type(data[i]))
+            logger.debug("size of data[%d] after: %s", i, ma.size(data[i]))
+            logger.debug("data of data[%d] after: %s", i, data[i])
+            assert ma.size(data[i]) == ma.size(data_)
+
+    return data
+
+def _probe_length(xs):
+    """
+    Returns the length of the first non-null array in `xs`, or None if all arrays
+    are null.
+    """
+    if isinstance(xs, dict):
+        return _probe_length(xs.values())
+
+    for x in xs:
+        if x is not None:
+            return x.size
+
+    return None
+
+def _ensure_list(xs, cinfos):
+    """
+    If input data is a dict, ensures it's converted to a list with the correct
+    offsets.
+    """
+    if isinstance(xs, list):
+        return xs
+
+    if isinstance(xs, np.ndarray):
+        ret = []
+        for x in xs:
+            ret.append(x)
+
+        return ret
+
+    # As we only accept list-likes or dicts as input data, it *must* be a dict at this
+    # point
+    assert isinstance(xs, dict)
+
+    logger.debug("data was provided as dict, coercing to list")
+
+    # As we may have non-existing keys, we would like to initialize those as a masked
+    # array with all elements masked. In those cases, though, we need to know the size
+    # of the array.
+    n = _probe_length(xs)
+
+    ret = list()
+
+    for i in range(len(cinfos)):
+        (cname, ctype) = cinfos[i]
+
+        xs_ = None
+        if cname in xs:
+            xs_ = xs[cname]
+        else:
+            xs_ = ma.masked_all(n, dtype=_best_dtype_for_ctype(ctype))
+
+        ret.append(xs_)
+
+    return ret
+
+
+def _coerce_retries(retries) -> quasardb.RetryOptions:
+    if retries is None:
+        return quasardb.RetryOptions()
+    elif isinstance(retries, int):
+        return quasardb.RetryOptions(retries=retries)
+    elif isinstance(retries, quasardb.RetryOptions):
+        return retries
+    else:
+        raise TypeError("retries should either be an integer or quasardb.RetryOptions, got: " + type(retries))
+
+
+def ensure_ma(xs, dtype=None):
+    if isinstance(dtype, list):
+        assert(isinstance(xs, list) == True)
+        return [ensure_ma(xs_, dtype_) for (xs_, dtype_) in zip(xs, dtype)]
+
+    # Don't bother if we're already a masked array
+    if ma.isMA(xs):
+        return xs
+
+    if not isinstance(xs, np.ndarray):
+        logger.debug("Provided data is not a numpy array: %s", type(xs))
+        xs = np.array(xs, dtype=dtype)
+
+    logger.debug("coercing array with dtype: %s", xs.dtype)
+
+    if xs.dtype.kind in ['O', 'U', 'S']:
+        logger.debug("Data is object-like, masking None values")
+
+        mask = xs == None
+        return ma.masked_array(data=xs, mask=mask)
+    else:
+        logger.debug("Automatically masking invalid numbers")
+        return ma.masked_invalid(xs, copy=False)
+
+
+def read_array(table=None,
+               column=None,
+               ranges=None):
+    if table is None:
+        raise RuntimeError("A table is required.")
+
+    if column is None:
+        raise RuntimeError("A column is required.")
+
+    kwargs = {
+        'column': column
+    }
+
+    if ranges is not None:
+        kwargs['ranges'] = ranges
+
+    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.Symbol: table.string_get_ranges,
+        quasardb.ColumnType.Int64: table.int64_get_ranges,
+        quasardb.ColumnType.Timestamp: table.timestamp_get_ranges,
+    }
+
+    ctype = table.column_type_by_id(column)
+
+    fn = read_with[ctype]
+    return fn(**kwargs)
+
+
+def write_array(
+        data=None,
+        index=None,
+        table=None,
+        column=None,
+        dtype=None,
+        infer_types=True):
+    """
+    Write a Numpy array to a single column.
+
+    Parameters:
+    -----------
+
+    data: np.array
+      Numpy array with a dtype that is compatible with the column's type.
+
+    index: np.array
+      Numpy array with a datetime64[ns] dtype that will be used as the
+      $timestamp axis for the data to be stored.
+
+    dtype: optional np.dtype
+      If provided, ensures the data array is converted to this dtype before
+      insertion.
+
+    infer_types: optional bool
+      If true, when necessary will attempt to convert the data and index array
+      to the best type for the column. For example, if you provide float64 data
+      while the column's type is int64, it will automatically convert the data.
+
+      Defaults to True. For production use cases where you want to avoid implicit
+      conversions, we recommend always setting this to False.
+
+    """
+
+    if table is None:
+        raise RuntimeError("A table is required.")
+
+    if column is None:
+        raise RuntimeError("A column is required.")
+
+    if data is None:
+        raise RuntimeError("A data numpy array is required.")
+
+    if index is None:
+        raise RuntimeError("An index numpy timestamp array is required.")
+
+
+    data = ensure_ma(data, dtype=dtype)
+    ctype =  table.column_type_by_id(column)
+
+    # We try to reuse some of the other functions, which assume array-like
+    # shapes for column info and data. It's a bit hackish, but actually works
+    # well.
+    #
+    # We should probably generalize this block of code with the same found in
+    # write_arrays().
+
+    cinfos = [(column, ctype)]
+    dtype_ = [dtype]
+
+    dtype = _coerce_dtype(dtype_, cinfos)
+
+    if infer_types is True:
+        dtype = _add_desired_dtypes(dtype, cinfos)
+
+    # data_ = an array of [data]
+    data_ = [data]
+    data_ = _coerce_data(data_, dtype)
+    _validate_dtypes(data_, cinfos)
+
+    # No functions that assume array-of-data anymore, let's put it back
+    data = data_[0]
+
+    # Dispatch to the correct function
+    write_with = {
+        quasardb.ColumnType.Double: table.double_insert,
+        quasardb.ColumnType.Blob: table.blob_insert,
+        quasardb.ColumnType.String: table.string_insert,
+        quasardb.ColumnType.Symbol: table.string_insert,
+        quasardb.ColumnType.Int64: table.int64_insert,
+        quasardb.ColumnType.Timestamp: table.timestamp_insert,
+    }
+
+    logger.info("Writing array (%d rows of dtype %s) to columns %s.%s (type %s)", len(data), data.dtype, table.get_name(), column, ctype)
+    write_with[ctype](column, index, data)
+
+def write_arrays(
+        data,
+        cluster,
+        table = None,
+        *,
+        dtype = None,
+        index = None,
+        _async = False,
+        fast = False,
+        truncate = False,
+        deduplicate = False,
+        deduplication_mode = 'drop',
+        infer_types = True,
+        writer = None,
+        write_through = False,
+        retries = 3,
+
+        # We accept additional kwargs that will be passed through the writer.push() methods
+        **kwargs):
+    """
+    Write multiple aligned numpy arrays to a table.
+
+    Parameters:
+    -----------
+
+    data: Iterable of np.array, or dict-like of str:np.array
+      Numpy arrays to write into the database. Can either be a list of numpy arrays,
+      in which case they are expected to be in the same order as table.list_columns(),
+      and an array is provided for each of the columns. If `index` is None, the first
+      array will be assumed to be an index with dtype `datetime64[ns]`.
+
+      Alternatively, a dict of key/values may be provided, where the key is expected
+      to be a table column label, and the value is expected to be a np.array. If present,
+      a column with label '$timestamp' will be used as the index.
+
+      In all cases, all numpy arrays are expected to be of exactly the same length as the
+      index.
+
+    cluster: quasardb.Cluster
+      Active connection to the QuasarDB cluster
+
+    table: quasardb.Table or str
+      Either a string or a reference to a QuasarDB Timeseries table object.
+      For example, 'my_table' or cluster.table('my_table') are both valid values.
+
+      Defaults to False.
+
+    index: optional np.array with dtype datetime64[ns]
+      Optionally explicitly provide an array as the $timestamp index. If not provided,
+      the first array provided to `data` will be used as the index.
+
+    dtype: optional dtype, list of dtype, or dict of dtype
+      Optional data type to force. If a single dtype, will force that dtype to all
+      columns. If list-like, will map dtypes to dataframe columns by their offset.
+      If dict-like, will map dtypes to dataframe columns by their label.
+
+      If a dtype for a column is provided in this argument, and infer_types is also
+      True, this argument takes precedence.
+
+    deduplicate: bool or list[str]
+      Enables server-side deduplication of data when it is written into the table.
+      When True, automatically deduplicates rows when all values of a row are identical.
+      When a list of strings is provided, deduplicates only based on the values of
+      these columns.
+
+      Defaults to False.
+
+    deduplication_mode: 'drop' or 'upsert'
+      When `deduplicate` is enabled, decides how deduplication is performed. 'drop' means
+      any newly written duplicates are dropped, where 'upsert' means that the previously
+      written data is updated to reflect the new data.
+
+      Defaults to 'drop'.
+
+    infer_types: optional bool
+      If true, will attemp to convert types from Python to QuasarDB natives types if
+      the provided dataframe has incompatible types. For example, a dataframe with integers
+      will automatically convert these to doubles if the QuasarDB table expects it.
+
+      Defaults to True. For production use cases where you want to avoid implicit conversions,
+      we recommend setting this to False.
+
+    truncate: optional bool
+      Truncate (also referred to as upsert) the data in-place. Will detect time range to truncate
+      from the time range inside the dataframe.
+
+      Defaults to False.
+
+    _async: optional bool
+      If true, uses asynchronous insertion API where commits are buffered server-side and
+      acknowledged before they are written to disk. If you insert to the same table from
+      multiple processes, setting this to True may improve performance.
+
+      Defaults to False.
+
+    fast: optional bool
+      Whether to use 'fast push'. If you incrementally add small batches of data to table,
+      you may see better performance if you set this to True.
+
+      Defaults to False.
+
+    write_through: optional bool
+      If True, data is not cached after write.
+      By default is False, in which case caching is left at the discretion of the server.
+
+    writer: optional quasardb.Writer
+      Allows you to explicitly provide a Writer to use, which is expected to be
+      initialized with the `table`.
+
+      Reuse of the Writer allows for some performance improvements.
+
+    retries: optional int or quasardb.RetryOptions
+      Number of times to retry in case of a push failure. This is useful in case of async
+      pipeline failures, or when doing transactional inserts that may occasionally cause
+      transaction conflicts.
+
+      Retries with exponential backoff, starts at 3 seconds, and doubles every retry attempt.
+
+      Alternatively, a quasardb.RetryOptions object can be passed to more carefully fine-tune
+      retry behavior.
+    """
+
+    if table:
+        logger.debug("table explicitly provided, assuming single-table write")
+        return write_arrays([(table, data)],
+                            cluster,
+                            table=None,
+                            dtype=dtype,
+                            index=index,
+                            _async=_async,
+                            fast=fast,
+                            truncate=truncate,
+                            deduplicate=deduplicate,
+                            deduplication_mode=deduplication_mode,
+                            infer_types=infer_types,
+                            write_through=write_through,
+                            writer=writer,
+                            retries=retries,
+                            **kwargs)
+
+
+    ret = []
+
+    # Create batch column info from dataframe
+    if writer is None:
+        writer = cluster.writer()
+
+    n_rows = 0
+
+    push_data = quasardb.WriterData()
+
+    for (table, data_) in data:
+        # Acquire reference to table if string is provided
+        if isinstance(table, str):
+            table = table_cache.lookup(table, cluster)
+
+        cinfos = [(x.name, x.type) for x in table.list_columns()]
+        dtype = _coerce_dtype(dtype, cinfos)
+
+        assert type(dtype) is list
+        assert len(dtype) is len(cinfos)
+
+
+        if index is None and isinstance(data_, dict) and '$timestamp' in data_:
+            index_ = data_.pop('$timestamp')
+            assert '$timestamp' not in data_
+        elif index is not None:
+            index_ = index
+        else:
+            raise RuntimeError("Invalid index: no index provided.")
+
+        assert index_ is not None
+
+        if infer_types is True:
+            dtype = _add_desired_dtypes(dtype, cinfos)
+
+        data_ = _ensure_list(data_, cinfos)
+
+        if len(data_) != len(cinfos):
+            raise InvalidDataCardinalityError(data_, cinfos)
+
+        data_ = ensure_ma(data_, dtype=dtype)
+        data_ = _coerce_data(data_, dtype)
+
+
+        # Just some additional friendly information about incorrect dtypes, we'd
+        # prefer to have this information thrown from Python instead of native
+        # code as it generally makes for somewhat better error context.
+        _validate_dtypes(data_, cinfos)
+
+        deduplicate = _coerce_deduplicate(deduplicate, deduplication_mode, cinfos)
+
+        # Sanity check
+        assert len(data_) == len(cinfos)
+
+        for i in range(len(data_)):
+            assert len(data_[i]) == len(index_)
+
+        push_data.append(table, index_, data_)
+
+        n_rows += len(index_)
+        ret.append(table)
+
+    retries = _coerce_retries(retries)
+
+    # By default, we push all additional kwargs to the writer.push() function. This allows transparent propagation
+    # arguments.
+    #
+    # The initial use case was that so we can add additional parameters for test mocks, e.g. `mock_failures` so that
+    # we can validate the retry functionality.
+    push_kwargs = kwargs
+    push_kwargs['deduplicate'] = deduplicate
+    push_kwargs['deduplication_mode'] = deduplication_mode
+    push_kwargs['write_through'] = write_through
+    push_kwargs['retries'] = retries
+
+    logger.debug("pushing %d rows", n_rows)
+    start = time.time()
+
+    if fast is True:
+        push_kwargs['push_mode'] = quasardb.WriterPushMode.Fast
+    elif truncate is True:
+        push_kwargs['push_mode'] = quasardb.WriterPushMode.Truncate
+    elif isinstance(truncate, tuple):
+        push_kwargs['push_mode'] = quasardb.WriterPushMode.Truncate
+        push_kwargs['range'] = truncate
+    elif _async is True:
+        push_kwargs['push_mode'] = quasardb.WriterPushMode.Async
+    else:
+        push_kwargs['push_mode'] = quasardb.WriterPushMode.Transactional
+
+    writer.push(push_data, **push_kwargs)
+
+    logger.debug("pushed %d rows in %s seconds",
+                n_rows, (time.time() - start))
+
+    return ret
+
+
+def _xform_query_results(xs, index, dict):
+    if len(xs) == 0:
+        return (np.array([], np.dtype('datetime64[ns]')), np.array([]))
+
+    n = None
+    for x in xs:
+        assert isinstance(x, tuple)
+
+        if n is None:
+            n = x[1].size
+        else:
+            assert x[1].size == n
+
+    if index is None:
+        # Generate a range, put it in the front of the result list,
+        # recurse and tell the function to use that index.
+        xs_ = [('$index', np.arange(n))] + xs
+        return _xform_query_results(xs_, '$index', dict)
+
+    if isinstance(index, str):
+        for i in range(len(xs)):
+            (cname, _) = xs[i]
+            if cname == index:
+                # Now we know that this column has offset `i`,
+                # recurse with that offset
+                return _xform_query_results(xs, i, dict)
+
+        raise KeyError("Unable to resolve index column: column not found in results: {}".format(index))
+
+    if not isinstance(index, int):
+        raise TypeError("Unable to resolve index column: unrecognized type {}: {}".format(type(index), index))
+
+    idx = xs[index][1]
+    del xs[index]
+
+    # Our index *may* be a masked array, but if it is, there should be no
+    # masked items: we cannot not have an index for a certain row.
+    if ma.isMA(idx):
+        if ma.count_masked(idx) > 0:
+            raise ValueError("Invalid index: null values detected. An index is never allowed to have null values.")
+
+        assert isinstance(idx.data, np.ndarray)
+        idx = idx.data
+
+    xs_ = None
+
+    if dict is True:
+        xs_ = {x[0]: x[1] for x in xs}
+    else:
+        xs_ = [x[1] for x in xs]
+
+    return (idx, xs_)
+
+
+def query(cluster,
+          query,
+          index=None,
+          dict=False):
+    """
+    Execute a query and return the results as numpy arrays. The shape of the return value
+    is always:
+
+      tuple[index, dict | list[np.array]]
+
+
+    If `dict` is True, constructs a dict[str, np.array] where the key is the column name.
+    Otherwise, it returns a list of all the individual data arrays.
+
+
+
+    Parameters:
+    -----------
+
+    cluster : quasardb.Cluster
+      Active connection to the QuasarDB cluster
+
+    query : str
+      The query to execute.
+
+    index : optional[str | int]
+      If provided, resolves column and uses that as the index. If string (e.g. `$timestamp`), uses
+      that column as the index. If int (e.g. `1`), looks up the column based on that offset.
+
+    dict : bool
+      If true, returns data arrays as a dict, otherwise a list of np.arrays.
+      Defaults to False.
+
+    """
+
+    m = {}
+    xs = cluster.query_numpy(query)
+
+    return _xform_query_results(xs, index, dict)