quasardb 3.14.2.dev4__cp39-cp39-win_amd64.whl → 3.14.2.dev6__cp39-cp39-win_amd64.whl

quasardb/numpy/__init__.py

@@ -29,11 +29,12 @@
 
 import logging
 import time
+import warnings
 
 import quasardb
 import quasardb.table_cache as table_cache
 
-logger = logging.getLogger('quasardb.numpy')
+logger = logging.getLogger("quasardb.numpy")
 
 
 class NumpyRequired(ImportError):
@@ -41,6 +42,7 @@ class NumpyRequired(ImportError):
     Exception raised when trying to use QuasarDB pandas integration, but
     pandas has not been installed.
     """
+
     pass
 
 
@@ -52,7 +54,7 @@ except ImportError as err:
     logger.exception(err)
     raise NumpyRequired(
         "The numpy library is required to handle numpy arrays formats"
-        ) from err
+    ) from err
 
 
 class IncompatibleDtypeError(TypeError):
@@ -68,13 +70,16 @@ class IncompatibleDtypeError(TypeError):
         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)
+        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())
@@ -82,29 +87,33 @@ class IncompatibleDtypeErrors(TypeError):
     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))
+        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]')]
+    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]")],
 }
 
 
@@ -141,9 +150,12 @@ def _coerce_dtype(dtype, columns):
         # Any columns not provided will have a 'None' dtype.
         dtype_ = [None] * len(columns)
 
-        for (k, dt) in dtype.items():
+        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, )
+                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
@@ -151,16 +163,22 @@ def _coerce_dtype(dtype, columns):
         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)))
+        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)))
+        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.
@@ -174,7 +192,12 @@ def _add_desired_dtypes(dtype, columns):
         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)
+            logger.debug(
+                "using default dtype '%s' for column '%s' with type %s",
+                dtype_,
+                cname,
+                ctype,
+            )
             dtype[i] = dtype_
 
     return dtype
@@ -196,8 +219,11 @@ def _is_all_masked(xs):
         # 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)))
+    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.
@@ -208,7 +234,7 @@ def _is_all_masked(xs):
 
 
 def dtypes_equal(lhs, rhs):
-    if lhs.kind == 'U' or lhs.kind == 'S':
+    if lhs.kind == "U" or lhs.kind == "S":
         # Unicode and string data has variable length encoding, which means their itemsize
         # can be anything.
         #
@@ -236,43 +262,60 @@ def _validate_dtypes(data, columns):
         (cname, ctype, provided_dtype, expected_dtype) = e
         return
 
-    for (data_, (cname, ctype)) in zip(data, columns):
+    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_))
+            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 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 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))
+        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))
+            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
@@ -291,16 +334,16 @@ def _clean_nulls(xs, dtype):
 
     assert ma.isMA(xs)
 
-    if xs.dtype is not np.dtype('object'):
+    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')
+        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')
+    elif dtype == np.dtype("datetime64[ns]"):
+        fill_value = np.datetime64("nat")
 
     mask = xs.mask
     xs_ = xs.filled(fill_value)
@@ -308,7 +351,6 @@ def _clean_nulls(xs, dtype):
     return ma.array(xs_, mask=mask)
 
 
-
 def _coerce_data(data, dtype):
     """
     Coerces each numpy array of `data` to the dtype present in `dtype`.
@@ -325,7 +367,12 @@ def _coerce_data(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(
+                "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_))
@@ -338,14 +385,20 @@ def _coerce_data(data, dtype):
                 # 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_)
+                    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.error(
+                        "An error occured while coercing input data type from dtype '%s' to dtype '%s': ",
+                        data_.dtype,
+                        dtype_,
+                    )
                     logger.exception(err)
                     raise err
 
@@ -358,6 +411,7 @@ def _coerce_data(data, dtype):
 
     return data
 
+
 def _probe_length(xs):
     """
     Returns the length of the first non-null array in `xs`, or None if all arrays
@@ -372,6 +426,7 @@ def _probe_length(xs):
 
     return None
 
+
 def _ensure_list(xs, cinfos):
     """
     If input data is a dict, ensures it's converted to a list with the correct
@@ -422,12 +477,42 @@ def _coerce_retries(retries) -> quasardb.RetryOptions:
     elif isinstance(retries, quasardb.RetryOptions):
         return retries
     else:
-        raise TypeError("retries should either be an integer or quasardb.RetryOptions, got: " + type(retries))
+        raise TypeError(
+            "retries should either be an integer or quasardb.RetryOptions, got: "
+            + type(retries)
+        )
+
+
+def _kwarg_deprecation_warning(
+    old_kwarg, old_value, new_kwargs, new_values, stacklevel
+):
+    new_declaration = ", ".join(
+        f"{new_kwarg}={new_value}"
+        for new_kwarg, new_value in zip(new_kwargs, new_values)
+    )
+    warnings.warn(
+        f"The argument '{old_kwarg}' <{type(old_value).__name__}> is deprecated and will be removed in a future version. "
+        f"Please use '{new_declaration}' instead.",
+        DeprecationWarning,
+        stacklevel=stacklevel + 1,
+    )
+
+
+def _type_check(var, var_name, target_type, raise_error=True, allow_none=True):
+    if allow_none and var is None:
+        return True
+    if not isinstance(var, target_type):
+        if raise_error:
+            raise quasardb.quasardb.InvalidArgumentError(
+                f"Invalid '{var_name}' type, expected: {target_type}, got: {type(var)}"
+            )
+        return False
+    return True
 
 
 def ensure_ma(xs, dtype=None):
     if isinstance(dtype, list):
-        assert(isinstance(xs, list) == True)
+        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
@@ -440,7 +525,7 @@ def ensure_ma(xs, dtype=None):
 
     logger.debug("coercing array with dtype: %s", xs.dtype)
 
-    if xs.dtype.kind in ['O', 'U', 'S']:
+    if xs.dtype.kind in ["O", "U", "S"]:
         logger.debug("Data is object-like, masking None values")
 
         mask = xs == None
@@ -450,21 +535,17 @@ def ensure_ma(xs, dtype=None):
         return ma.masked_invalid(xs, copy=False)
 
 
-def read_array(table=None,
-               column=None,
-               ranges=None):
+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
-    }
+    kwargs = {"column": column}
 
     if ranges is not None:
-        kwargs['ranges'] = ranges
+        kwargs["ranges"] = ranges
 
     read_with = {
         quasardb.ColumnType.Double: table.double_get_ranges,
@@ -482,12 +563,8 @@ def read_array(table=None,
 
 
 def write_array(
-        data=None,
-        index=None,
-        table=None,
-        column=None,
-        dtype=None,
-        infer_types=True):
+    data=None, index=None, table=None, column=None, dtype=None, infer_types=True
+):
     """
     Write a Numpy array to a single column.
 
@@ -527,9 +604,8 @@ def write_array(
     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)
+    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
@@ -564,28 +640,39 @@ def write_array(
         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)
+    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):
+    data,
+    cluster,
+    table=None,
+    *,
+    dtype=None,
+    index=None,
+    # TODO: Set the default push_mode after removing _async, fast and truncate
+    push_mode=None,
+    _async=False,
+    fast=False,
+    truncate=False,
+    truncate_range=None,
+    deduplicate=False,
+    deduplication_mode="drop",
+    infer_types=True,
+    writer=None,
+    write_through=True,
+    retries=3,
+    # We accept additional kwargs that will be passed through the writer.push() methods
+    **kwargs,
+):
     """
     Write multiple aligned numpy arrays to a table.
 
@@ -649,13 +736,32 @@ def write_arrays(
       Defaults to True. For production use cases where you want to avoid implicit conversions,
       we recommend setting this to False.
 
+    push_mode: optional quasardb.WriterPushMode
+      The mode used for inserting data. Can be either a string or a `WriterPushMode` enumeration item.
+      Available options:
+      * `Truncate`: Truncate (also referred to as upsert) the data in-place. Will detect time range
+        to truncate from the time range inside the dataframe.
+      * `Async`: 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.
+      * `Fast`: 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.
+      * `Transactional`: Ensures full transactional consistency.
+
+      Defaults to `Transactional`.
+
     truncate: optional bool
+      **DEPRECATED** – Use `push_mode=WriterPushMode.Truncate` instead.
       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.
 
+    truncate_range: optional tuple
+      Time range to truncate from the time range inside the dataframe.
+
     _async: optional bool
+      **DEPRECATED** – Use `push_mode=WriterPushMode.Async` instead.
       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.
@@ -663,6 +769,7 @@ def write_arrays(
       Defaults to False.
 
     fast: optional bool
+      **DEPRECATED** – Use `push_mode=WriterPushMode.Fast` instead.
       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.
 
@@ -691,34 +798,60 @@ def write_arrays(
 
     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)
-
+        data = [(table, data)]
+        table = None
+
+    _type_check(push_mode, "push_mode", target_type=quasardb.WriterPushMode)
+    deprecation_stacklevel = kwargs.pop("deprecation_stacklevel", 1) + 1
+
+    if isinstance(truncate, tuple):
+        # Especial case, truncate might be a tuple indicating the range.
+        _kwarg_deprecation_warning(
+            "truncate",
+            truncate,
+            ["push_mode", "truncate_range"],
+            [quasardb.WriterPushMode.Truncate, truncate],
+            deprecation_stacklevel,
+        )
+        truncate_range = truncate_range or truncate
+        truncate = True
+
+    kwarg_to_mode = {
+        # "kwarg": (kwarg_type, kwarg_push_mode, is_deprecated)
+        "fast": (bool, quasardb.WriterPushMode.Fast, True),
+        "_async": (bool, quasardb.WriterPushMode.Async, True),
+        "truncate": (bool, quasardb.WriterPushMode.Truncate, True),
+        "truncate_range": (tuple, quasardb.WriterPushMode.Truncate, False),
+    }
 
-    ret = []
+    for kwarg, info in kwarg_to_mode.items():
+        expected_type, mode, deprecated = info
+        kwarg_value = locals()[kwarg]
+        _type_check(kwarg_value, kwarg, target_type=expected_type)
+
+        if kwarg_value:
+            if push_mode and push_mode != mode:
+                raise quasardb.quasardb.InvalidArgumentError(
+                    f"Found '{kwarg}' in kwargs, but push mode is already set to {push_mode}"
+                )
+            push_mode = mode
+            if deprecated:
+                _kwarg_deprecation_warning(
+                    kwarg, kwarg_value, ["push_mode"], [mode], deprecation_stacklevel
+                )
+
+    if not push_mode:
+        push_mode = quasardb.WriterPushMode.Transactional
 
     # Create batch column info from dataframe
     if writer is None:
         writer = cluster.writer()
 
+    ret = []
     n_rows = 0
-
     push_data = quasardb.WriterData()
 
-    for (table, data_) in data:
+    for table, data_ in data:
         # Acquire reference to table if string is provided
         if isinstance(table, str):
             table = table_cache.lookup(table, cluster)
@@ -729,10 +862,15 @@ def write_arrays(
         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_
+        if index is None and isinstance(data_, dict) and "$timestamp" in data_:
+            # Create shallow copy of `data_` so that we don't modify the reference, i.e.
+            # delete keys.
+            #
+            # This ensures that the user can call the same function multiple times without
+            # side-effects.
+            data_ = data_.copy()
+            index_ = data_.pop("$timestamp")
+            assert "$timestamp" not in data_
         elif index is not None:
             index_ = index
         else:
@@ -751,7 +889,6 @@ def write_arrays(
         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.
@@ -778,37 +915,27 @@ def write_arrays(
     # 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
+    push_kwargs["deduplicate"] = deduplicate
+    push_kwargs["deduplication_mode"] = deduplication_mode
+    push_kwargs["write_through"] = write_through
+    push_kwargs["retries"] = retries
+    push_kwargs["push_mode"] = push_mode
+    if truncate_range:
+        push_kwargs["range"] = truncate_range
 
     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))
+    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([]))
+        return (np.array([], np.dtype("datetime64[ns]")), np.array([]))
 
     n = None
     for x in xs:
@@ -822,8 +949,8 @@ def _xform_query_results(xs, index, dict):
     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)
+        xs_ = [("$index", np.arange(n))] + xs
+        return _xform_query_results(xs_, "$index", dict)
 
     if isinstance(index, str):
         for i in range(len(xs)):
@@ -833,10 +960,18 @@ def _xform_query_results(xs, index, dict):
                 # 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))
+        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))
+        raise TypeError(
+            "Unable to resolve index column: unrecognized type {}: {}".format(
+                type(index), index
+            )
+        )
 
     idx = xs[index][1]
     del xs[index]
@@ -845,7 +980,9 @@ def _xform_query_results(xs, index, dict):
     # 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.")
+            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
@@ -860,10 +997,7 @@ def _xform_query_results(xs, index, dict):
     return (idx, xs_)
 
 
-def query(cluster,
-          query,
-          index=None,
-          dict=False):
+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: