Repository: pydata/sparse

Issue:
Indexing is O(n) when it can be made O(log n)
If we call `COO.sum_duplicates()` beforehand, it may be possible, with some trickery, to make indexing `O(log n)`. I'm still looking into how this can be done, and how to generalize it to multiple indices, but I know it's possible.

One concern I have is that I do not want it to become `O(n log n)` for multiple indices, so this needs to be thought out carefully.

Discussion:
codecov-io: # [Codecov](https://codecov.io/gh/pydata/sparse/pull/128?src=pr&el=h1) Report
> Merging [#128](https://codecov.io/gh/pydata/sparse/pull/128?src=pr&el=desc) into [master](https://codecov.io/gh/pydata/sparse/commit/250d7c6effb5438614400a004d9469f234e4f5a0?src=pr&el=desc) will **decrease** coverage by `2.38%`.
> The diff coverage is `62.62%`.

[![Impacted file tree graph](https://codecov.io/gh/pydata/sparse/pull/128/graphs/tree.svg?token=H212u0Uxxw&src=pr&width=650&height=150)](https://codecov.io/gh/pydata/sparse/pull/128?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #128      +/-   ##
==========================================
- Coverage   92.03%   89.65%   -2.39%     
==========================================
  Files          13       12       -1     
  Lines        1256     1295      +39     
==========================================
+ Hits         1156     1161       +5     
- Misses        100      134      +34
```

| Flag | Coverage Δ | |
|---|---|---|
| #python27 | `88.95% <62.62%> (-2.33%)` | :arrow_down: |
| #python36 | `89.34% <62.62%> (-2.38%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pydata/sparse/pull/128?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [benchmarks/benchmark\_coo.py](https://codecov.io/gh/pydata/sparse/pull/128/diff?src=pr&el=tree#diff-YmVuY2htYXJrcy9iZW5jaG1hcmtfY29vLnB5) | `0% <0%> (ø)` | :arrow_up: |
| [sparse/coo/core.py](https://codecov.io/gh/pydata/sparse/pull/128/diff?src=pr&el=tree#diff-c3BhcnNlL2Nvby9jb3JlLnB5) | `87.34% <54.16%> (-6.37%)` | :arrow_down: |
| [sparse/slicing.py](https://codecov.io/gh/pydata/sparse/pull/128/diff?src=pr&el=tree#diff-c3BhcnNlL3NsaWNpbmcucHk=) | `98.18% <95.83%> (+1.69%)` | :arrow_up: |
| [sparse/\_\_init\_\_.py](https://codecov.io/gh/pydata/sparse/pull/128/diff?src=pr&el=tree#diff-c3BhcnNlL19faW5pdF9fLnB5) | | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pydata/sparse/pull/128?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pydata/sparse/pull/128?src=pr&el=footer). Last update [250d7c6...0fdfdc1](https://codecov.io/gh/pydata/sparse/pull/128?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

hameerabbasi: I will note, however, that this indexing is geared towards integer or tuple of integer. Large slices WILL make this slow as molasses. Maybe a fallback if the slices are too large?
hameerabbasi: Latest benchmarks:
```
In[2]: import sparse
In[3]: x = sparse.random((1000, 1000, 1000), density=0.01)
In[4]: %timeit x[:500]
78.2 ms ± 613 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
In[5]: %timeit x[:500, :500]
223 ms ± 6.73 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In[6]: %timeit x[:500, :500, :500]
230 ms ± 6.07 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In[7]: %timeit x[5]
157 µs ± 423 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
```

Faster in all cases than old. Falls back to simple matching based on a "guesstimate" of which one will be faster.
hameerabbasi: cc @mrocklin Ready for review. :-)
mrocklin: Cool.  I'm looking forward to it.   I took a brief glance at it and noticed that it was fairly involved, including many new routines with complex logic.  This will probably take some time to review properly and so I may not get to this today.  

If you're inclined I suspect that you could make reviewing (and thus reading by future devs) easier by including small and simple doctests within each utility function that show what it does with a minimal example.  
hameerabbasi: I had already partially done that, but I made it more complete.
hameerabbasi: A couple of replies.
hameerabbasi: I added a bit of comments to the main function.
hameerabbasi: Rebased onto infrastructure fix commits.
hameerabbasi: @mrocklin I'm done working on this. Do you have time to review in the near future?
hameerabbasi: I think your idea of moving this to a separate module and adding a module-level docstring is probably best. I'll request you to hold off reviewing further until I make that change (it'll make things easier). I'll hopefully do this tomorrow at some point.

The algorithm itself probably isn't that difficult, but I can see the need for terminology being clarified.
mrocklin: Thanks.  In general I'm excited about the performance gains here.  Another option would just be to merge and move ahead without review.  Ideally in the future we'll have more people who have time to both write and review code.
hameerabbasi: I've clarified everything, hopefully. :-) Feel free to review and/or merge.
mrocklin: It seems like a lot of the logic here is based around performance considerations.  Do we want to add benchmarks for these various cases?
hameerabbasi: > It seems like a lot of the logic here is based around performance considerations. Do we want to add benchmarks for these various cases?

Seems like an excellent idea!
hameerabbasi: I noticed there were a few cases we supported but didn't actually test for. That isn't additional in this PR, but it was nice to add tests for those cases as well.
mrocklin: > I'm fairly certain how it will be done, but I haven't mapped it out in code yet. However, it requires list[list[int]] on Numba's end.

How confident are we that this will be implemented soon and, if implemented, that it will be fast enough?  If neither of those are true then do we have a backup plan?
hameerabbasi: > How confident are we that this will be implemented soon and, if implemented, that it will be fast enough? If neither of those are true then do we have a backup plan?

numba/numba#2560 is the issue we're looking for. It has a milestone of 0.38 (RC), which is due out April 11. Unless it's pushed to the next release, we should be okay.

There may be a way with flat data structures as well, but it's slightly more complicated. May be more performant, though. I'm shooting for next weekend. It requires some Numba trickery in any case (I think that trickery should be possible in `nopython` mode and if so, it'll be fast, roughly equivalent to `n` integer indexes where `n` is the size of the advanced index).

If nothing works out we can simply loop over the relevant indices in Python.

Do you need it urgently in your work?
mrocklin: No, I just want to make sure that we don't over-engineer into a solution that can not be extended.  I suspect that accepting lists and arrays of boolean and integer values will be important for XArray support, which is a nice use case to target.
hameerabbasi: Good thinking! But don't worry, I'm fairly sure that with integer indexing and some transposing and reshaping, we'll be good to go.
mrocklin: I still haven't gone through all of the logic here, but I've probably gone as far as I'm likely to.  In general things seem fine to me.  +1

Thanks for slogging through this @hameerabbasi .  I hope you're enjoying yourself :)

Patch:
diff --git a/.gitignore b/.gitignore
index e408e4a..71b3ae3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -45,6 +45,10 @@ htmlcov/
 nosetests.xml
 coverage.xml
 *,cover
+.pytest_cache/
+
+# Airspeed velocity
+.asv/
 
 # Translations
 *.mo
diff --git a/benchmarks/benchmark_coo.py b/benchmarks/benchmark_coo.py
index 5cb9f74..4bf760f 100644
--- a/benchmarks/benchmark_coo.py
+++ b/benchmarks/benchmark_coo.py
@@ -12,12 +12,17 @@ class ElemwiseSuite(object):
         self.x.sum_duplicates()
         self.y.sum_duplicates()
 
+        self.x + self.y  # Numba compilation
+
     def time_add(self):
         self.x + self.y
 
     def time_mul(self):
         self.x * self.y
 
+    def time_index(self):
+        self.x[5]
+
 
 class ElemwiseBroadcastingSuite(object):
     def setup(self):
@@ -33,3 +38,24 @@ class ElemwiseBroadcastingSuite(object):
 
     def time_mul(self):
         self.x * self.y
+
+
+class IndexingSuite(object):
+    def setup(self):
+        np.random.seed(0)
+        self.x = sparse.random((100, 100, 100), density=0.01)
+        self.x.sum_duplicates()
+
+        self.x[5]  # Numba compilation
+
+    def time_index_scalar(self):
+        self.x[5]
+
+    def time_index_slice(self):
+        self.x[:50]
+
+    def time_index_slice2(self):
+        self.x[:50, :50]
+
+    def time_index_slice3(self):
+        self.x[:50, :50, :50]
diff --git a/sparse/coo/core.py b/sparse/coo/core.py
index 6fbfdd9..cbc9612 100644
--- a/sparse/coo/core.py
+++ b/sparse/coo/core.py
@@ -1,4 +1,3 @@
-import numbers
 from collections import Iterable, defaultdict, deque
 
 import numpy as np
@@ -6,9 +5,9 @@ import scipy.sparse
 from numpy.lib.mixins import NDArrayOperatorsMixin
 
 from .common import dot
+from .indexing import getitem
 from .umath import elemwise, broadcast_to
 from ..compatibility import int, range
-from ..slicing import normalize_index
 from ..sparse_array import SparseArray
 from ..utils import _zero_of_dtype
 
@@ -487,93 +486,7 @@ class COO(SparseArray, NDArrayOperatorsMixin):
     def __sizeof__(self):
         return self.nbytes
 
-    def __getitem__(self, index):
-        if not isinstance(index, tuple):
-            if isinstance(index, str):
-                data = self.data[index]
-                idx = np.where(data)
-                coords = list(self.coords[:, idx[0]])
-                coords.extend(idx[1:])
-
-                return COO(coords, data[idx].flatten(),
-                           shape=self.shape + self.data.dtype[index].shape,
-                           has_duplicates=self.has_duplicates,
-                           sorted=self.sorted)
-            else:
-                index = (index,)
-
-        last_ellipsis = len(index) > 0 and index[-1] is Ellipsis
-        index = normalize_index(index, self.shape)
-        if len(index) != 0 and all(not isinstance(ind, Iterable) and ind == slice(None) for ind in index):
-            return self
-        mask = np.ones(self.nnz, dtype=np.bool)
-        for i, ind in enumerate([i for i in index if i is not None]):
-            if not isinstance(ind, Iterable) and ind == slice(None):
-                continue
-            mask &= _mask(self.coords[i], ind, self.shape[i])
-
-        n = mask.sum()
-        coords = []
-        shape = []
-        i = 0
-        for ind in index:
-            if isinstance(ind, numbers.Integral):
-                i += 1
-                continue
-            elif isinstance(ind, slice):
-                step = ind.step if ind.step is not None else 1
-                if step > 0:
-                    start = ind.start if ind.start is not None else 0
-                    start = max(start, 0)
-                    stop = ind.stop if ind.stop is not None else self.shape[i]
-                    stop = min(stop, self.shape[i])
-                    if start > stop:
-                        start = stop
-                    shape.append((stop - start + step - 1) // step)
-                else:
-                    start = ind.start or self.shape[i] - 1
-                    stop = ind.stop if ind.stop is not None else -1
-                    start = min(start, self.shape[i] - 1)
-                    stop = max(stop, -1)
-                    if start < stop:
-                        start = stop
-                    shape.append((start - stop - step - 1) // (-step))
-
-                dt = np.min_scalar_type(min(-(dim - 1) if dim != 0 else -1 for dim in shape))
-                coords.append((self.coords[i, mask].astype(dt) - start) // step)
-                i += 1
-            elif isinstance(ind, Iterable):
-                old = self.coords[i][mask]
-                new = np.empty(shape=old.shape, dtype=old.dtype)
-                for j, item in enumerate(ind):
-                    new[old == item] = j
-                coords.append(new)
-                shape.append(len(ind))
-                i += 1
-            elif ind is None:
-                coords.append(np.zeros(n))
-                shape.append(1)
-
-        for j in range(i, self.ndim):
-            coords.append(self.coords[j][mask])
-            shape.append(self.shape[j])
-
-        if coords:
-            coords = np.stack(coords, axis=0)
-        else:
-            if last_ellipsis:
-                coords = np.empty((0, np.sum(mask)), dtype=np.uint8)
-            else:
-                if np.sum(mask) != 0:
-                    return self.data[mask][0]
-                else:
-                    return _zero_of_dtype(self.dtype)[()]
-        shape = tuple(shape)
-        data = self.data[mask]
-
-        return COO(coords, data, shape=shape,
-                   has_duplicates=self.has_duplicates,
-                   sorted=self.sorted)
+    __getitem__ = getitem
 
     def __str__(self):
         return "<COO: shape=%s, dtype=%s, nnz=%d, sorted=%s, duplicates=%s>" % (
@@ -1572,28 +1485,6 @@ def _keepdims(original, new, axis):
     return new.reshape(shape)
 
 
-def _mask(coords, idx, shape):
-    if isinstance(idx, numbers.Integral):
-        return coords == idx
-    elif isinstance(idx, slice):
-        step = idx.step if idx.step is not None else 1
-        if step > 0:
-            start = idx.start if idx.start is not None else 0
-            stop = idx.stop if idx.stop is not None else shape
-            return (coords >= start) & (coords < stop) & \
-                   (coords % step == start % step)
-        else:
-            start = idx.start if idx.start is not None else (shape - 1)
-            stop = idx.stop if idx.stop is not None else -1
-            return (coords <= start) & (coords > stop) & \
-                   (coords % step == start % step)
-    elif isinstance(idx, Iterable):
-        mask = np.zeros(len(coords), dtype=np.bool)
-        for item in idx:
-            mask |= _mask(coords, item, shape)
-        return mask
-
-
 def _grouped_reduce(x, groups, method, **kwargs):
     """
     Performs a :code:`ufunc` grouped reduce.
diff --git a/sparse/coo/indexing.py b/sparse/coo/indexing.py
new file mode 100644
index 0000000..f172412
--- /dev/null
+++ b/sparse/coo/indexing.py
@@ -0,0 +1,452 @@
+from collections import Iterable
+from numbers import Integral
+
+import numba
+import numpy as np
+
+from ..compatibility import range, zip_longest
+from ..slicing import normalize_index
+from ..utils import _zero_of_dtype
+
+
+def getitem(x, index):
+    """
+    This function implements the indexing functionality for COO.
+
+    The overall algorithm has three steps:
+
+    1. Normalize the index to canonical form. Function: normalize_index
+    2. Get the mask, which is a list of integers corresponding to
+       the indices in coords/data for the output data. Function: _mask
+    3. Transform the coordinates to what they will be in the output.
+
+    Parameters
+    ----------
+    x : COO
+        The array to apply the indexing operation on.
+    index : {tuple, str}
+        The index into the array.
+    """
+    from .core import COO
+
+    x.sum_duplicates()
+
+    # If string, this is an index into an np.void
+    # Custom dtype.
+    if isinstance(index, str):
+        data = x.data[index]
+        idx = np.where(data)
+        coords = list(x.coords[:, idx[0]])
+        coords.extend(idx[1:])
+
+        return COO(coords, data[idx].flatten(),
+                   shape=x.shape + x.data.dtype[index].shape,
+                   has_duplicates=x.has_duplicates,
+                   sorted=x.sorted)
+
+    # Otherwise, convert into a tuple.
+    if not isinstance(index, tuple):
+        index = (index,)
+
+    # Check if the last index is an ellipsis.
+    last_ellipsis = len(index) > 0 and index[-1] is Ellipsis
+
+    # Normalize the index into canonical form.
+    index = normalize_index(index, x.shape)
+
+    # zip_longest so things like x[..., None] are picked up.
+    if len(index) != 0 and all(ind == slice(0, dim, 1) for ind, dim in zip_longest(index, x.shape)):
+        return x
+
+    # Get the mask
+    mask = _mask(x.coords, index, x.shape)
+
+    # Get the length of the mask
+    if isinstance(mask, slice):
+        n = len(range(mask.start, mask.stop, mask.step))
+    else:
+        n = len(mask)
+
+    coords = []
+    shape = []
+    i = 0
+    for ind in index:
+        # Nothing is added to shape or coords if the index is an integer.
+        if isinstance(ind, Integral):
+            i += 1
+            continue
+        # Add to the shape and transform the coords in the case of a slice.
+        elif isinstance(ind, slice):
+            shape.append(len(range(ind.start, ind.stop, ind.step)))
+            dt = np.min_scalar_type(min(-(dim - 1) if dim != 0 else -1 for dim in shape))
+            coords.append((x.coords[i, mask].astype(dt) - ind.start) // ind.step)
+            i += 1
+        elif isinstance(ind, Iterable):
+            raise NotImplementedError('Advanced indexing is not yet supported.')
+        # Add a dimension for None.
+        elif ind is None:
+            coords.append(np.zeros(n))
+            shape.append(1)
+
+    # Join all the transformed coords.
+    if coords:
+        coords = np.stack(coords, axis=0)
+    else:
+        # If index result is a scalar, return a 0-d COO or
+        # a scalar depending on whether the last index is an ellipsis.
+        if last_ellipsis:
+            coords = np.empty((0, n), dtype=np.uint8)
+        else:
+            if n != 0:
+                return x.data[mask][0]
+            else:
+                return _zero_of_dtype(x.dtype)[()]
+
+    shape = tuple(shape)
+    data = x.data[mask]
+
+    return COO(coords, data, shape=shape,
+               has_duplicates=x.has_duplicates,
+               sorted=x.sorted)
+
+
+def _mask(coords, indices, shape):
+    indices = _prune_indices(indices, shape)
+
+    ind_ar = np.empty((len(indices), 3), dtype=np.intp)
+
+    for i, idx in enumerate(indices):
+        if isinstance(idx, slice):
+            ind_ar[i] = [idx.start, idx.stop, idx.step]
+        else:  # idx is an integer
+            ind_ar[i] = [idx, idx + 1, 1]
+
+    mask, is_slice = _compute_mask(coords, ind_ar)
+
+    if is_slice:
+        return slice(mask[0], mask[1], 1)
+    else:
+        return mask
+
+
+def _prune_indices(indices, shape, prune_none=True):
+    """
+    Gets rid of the indices that do not contribute to the
+    overall mask, e.g. None and full slices.
+
+    Parameters
+    ----------
+    indices : tuple
+        The indices to the array.
+    shape : tuple[int]
+        The shape of the array.
+
+    Returns
+    -------
+    indices : tuple
+        The filtered indices.
+
+    Examples
+    --------
+    >>> _prune_indices((None, 5), (10,)) # None won't affect the mask
+    [5]
+    >>> _prune_indices((slice(0, 10, 1),), (10,)) # Full slices don't affect the mask
+    []
+    """
+    if prune_none:
+        indices = [idx for idx in indices if idx is not None]
+
+    i = 0
+    for idx, l in zip(indices[::-1], shape[::-1]):
+        if not isinstance(idx, slice):
+            break
+
+        if idx.start == 0 and idx.stop == l and idx.step == 1:
+            i += 1
+            continue
+
+        if idx.start == l - 1 and idx.stop == -1 and idx.step == -1:
+            i += 1
+            continue
+
+        break
+    if i != 0:
+        indices = indices[:-i]
+    return indices
+
+
+@numba.jit(nopython=True)
+def _compute_mask(coords, indices):  # pragma: no cover
+    """
+    Gets the mask for the coords given the indices in slice format.
+
+    Works with either start-stop ranges of matching indices into coords
+    called "pairs" (start-stop pairs) or filters the mask directly, based
+    on which is faster.
+
+    Exploits the structure in sorted coords, which is that for a constant
+    value of coords[i - 1], coords[i - 2] and so on, coords[i] is sorted.
+    Concretely, ``coords[i, coords[i - 1] == v1 & coords[i - 2] = v2, ...]``
+    is always sorted. It uses this sortedness to find sub-pairs for each
+    dimension given the previous, and so on. This is efficient for small
+    slices or ints, but not for large ones.
+
+    After it detects that working with pairs is rather inefficient (or after
+    going through each possible index), it constructs a filtered mask from the
+    start-stop pairs.
+
+    Parameters
+    ----------
+    coords : np.ndarray
+        The coordinates of the array.
+    indices : np.ndarray
+        The indices in the form of slices such that indices[:, 0] are starts,
+        indices[:, 1] are stops and indices[:, 2] are steps.
+
+    Returns
+    -------
+    mask : np.ndarray
+        The starts and stops in the mask.
+    is_slice : bool
+        Whether or not the array represents a continuous slice.
+
+    Examples
+    --------
+    Let's create some mock coords and indices
+
+    >>> import numpy as np
+    >>> coords = np.array([[0, 0, 1, 1, 2, 2]])
+    >>> indices = np.array([[0, 3, 2]])  # Equivalent to slice(0, 3, 2)
+
+    Now let's get the mask. Notice that the indices of ``0`` and ``2`` are matched.
+
+    >>> _compute_mask(coords, indices)
+    (array([0, 1, 4, 5]), False)
+
+    Now, let's try with a more "continuous" slice. Matches ``0`` and ``1``.
+
+    >>> indices = np.array([[0, 2, 1]])
+    >>> _compute_mask(coords, indices)
+    (array([0, 4]), True)
+
+    This is equivalent to mask being ``slice(0, 4, 1)``.
+    """
+    # Set the initial mask to be the entire range of coordinates.
+    starts = [0]
+    stops = [coords.shape[1]]
+    n_matches = coords.shape[1]
+
+    i = 0
+    while i < len(indices):
+        # Guesstimate whether working with pairs is more efficient or
+        # working with the mask directly.
+        # One side is the estimate of time taken for binary searches
+        # (n_searches * log(avg_length))
+        # The other is an estimated time of a linear filter for the mask.
+        n_pairs = len(starts)
+        n_current_slices = _get_slice_len(indices[i]) * n_pairs + 2
+        if n_current_slices * np.log(n_current_slices / max(n_pairs, 1)) > \
+                n_matches + n_pairs:
+            break
+
+        # For each of the pairs, search inside the coordinates for other
+        # matching sub-pairs.
+        # This gets the start-end coordinates in coords for each 'sub-array'
+        # Which would come out of indexing a single integer.
+        starts, stops, n_matches = _get_mask_pairs(starts, stops, coords[i], indices[i])
+
+        i += 1
+
+    # Combine adjacent pairs
+    starts, stops = _join_adjacent_pairs(starts, stops)
+
+    # If just one pair is left over, treat it as a slice.
+    if i == len(indices) and len(starts) == 1:
+        return np.array([starts[0], stops[0]]), True
+
+    # Convert start-stop pairs into mask, filtering by remaining
+    # coordinates.
+    mask = _filter_pairs(starts, stops, coords[i:], indices[i:])
+
+    return np.array(mask, dtype=np.intp), False
+
+
+@numba.jit(nopython=True)
+def _get_mask_pairs(starts_old, stops_old, c, idx):  # pragma: no cover
+    """
+    Gets the pairs for a following dimension given the pairs for
+    a dimension.
+
+    For each pair, it searches in the following dimension for
+    matching coords and returns those.
+
+    The total combined length of all pairs is returned to
+    help with the performance guesstimate.
+
+    Parameters
+    ----------
+    starts_old, stops_old : list[int]
+        The starts and stops from the previous index.
+    c : np.ndarray
+        The coords for this index's dimension.
+    idx : np.ndarray
+        The index in the form of a slice.
+        idx[0], idx[1], idx[2] = start, stop, step
+
+    Returns
+    -------
+    starts, stops: list
+        The starts and stops after applying the current index.
+    n_matches : int
+        The sum of elements in all ranges.
+
+    Examples
+    --------
+    >>> c = np.array([1, 2, 1, 2, 1, 1, 2, 2])
+    >>> starts_old = [4]
+    >>> stops_old = [8]
+    >>> idx = np.array([1, 2, 1])
+    >>> _get_mask_pairs(starts_old, stops_old, c, idx)
+    ([4], [6], 2)
+    """
+    starts = []
+    stops = []
+    n_matches = 0
+
+    for j in range(len(starts_old)):
+        # For each matching "integer" in the slice, search within the "sub-coords"
+        # Using binary search.
+        for p_match in range(idx[0], idx[1], idx[2]):
+            start = np.searchsorted(c[starts_old[j]:stops_old[j]], p_match) + starts_old[j]
+            stop = np.searchsorted(c[starts_old[j]:stops_old[j]], p_match + 1) + starts_old[j]
+
+            if start != stop:
+                starts.append(start)
+                stops.append(stop)
+                n_matches += stop - start
+
+    return starts, stops, n_matches
+
+
+@numba.jit(nopython=True)
+def _get_slice_len(idx):  # pragma: no cover
+    """
+    Get the number of elements in a slice.
+
+    Parameters
+    ----------
+    idx : np.ndarray
+        A (3,) shaped array containing start, stop, step
+
+    Returns
+    -------
+    n : int
+        The length of the slice.
+
+    Examples
+    --------
+    >>> idx = np.array([5, 15, 5])
+    >>> _get_slice_len(idx)
+    2
+    """
+    start, stop, step = idx[0], idx[1], idx[2]
+
+    if step > 0:
+        return (stop - start + step - 1) // step
+    else:
+        return (start - stop - step - 1) // (-step)
+
+
+@numba.jit(nopython=True)
+def _filter_pairs(starts, stops, coords, indices):  # pragma: no cover
+    """
+    Converts all the pairs into a single integer mask, additionally filtering
+    by the indices.
+
+    Parameters
+    ----------
+    starts, stops : list[int]
+        The starts and stops to convert into an array.
+    coords : np.ndarray
+        The coordinates to filter by.
+    indices : np.ndarray
+        The indices in the form of slices such that indices[:, 0] are starts,
+        indices[:, 1] are stops and indices[:, 2] are steps.
+
+    Returns
+    -------
+    mask : list
+        The output integer mask.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> starts = [2]
+    >>> stops = [7]
+    >>> coords = np.array([[0, 1, 2, 3, 4, 5, 6, 7]])
+    >>> indices = np.array([[2, 8, 2]]) # Start, stop, step pairs
+    >>> _filter_pairs(starts, stops, coords, indices)
+    [2, 4, 6]
+    """
+    mask = []
+
+    # For each pair,
+    for i in range(len(starts)):
+        # For each element match within the pair range
+        for j in range(starts[i], stops[i]):
+            match = True
+
+            # Check if it matches all indices
+            for k in range(len(indices)):
+                idx = indices[k]
+                elem = coords[k, j]
+
+                match &= ((elem - idx[0]) % idx[2] == 0 and
+                          ((idx[2] > 0 and idx[0] <= elem < idx[1])
+                           or (idx[2] < 0 and idx[0] >= elem > idx[1])))
+
+            # and append to the mask if so.
+            if match:
+                mask.append(j)
+
+    return mask
+
+
+@numba.jit(nopython=True)
+def _join_adjacent_pairs(starts_old, stops_old):  # pragma: no cover
+    """
+    Joins adjacent pairs into one. For example, 2-5 and 5-7
+    will reduce to 2-7 (a single pair). This may help in
+    returning a slice in the end which could be faster.
+
+    Parameters
+    ----------
+    starts_old, stops_old : list[int]
+        The input starts and stops
+
+    Returns
+    -------
+    starts, stops : list[int]
+        The reduced starts and stops.
+
+    Examples
+    --------
+    >>> starts = [2, 5]
+    >>> stops = [5, 7]
+    >>> _join_adjacent_pairs(starts, stops)
+    ([2], [7])
+    """
+    if len(starts_old) <= 1:
+        return starts_old, stops_old
+
+    starts = [starts_old[0]]
+    stops = []
+
+    for i in range(1, len(starts_old)):
+        if starts_old[i] != stops_old[i - 1]:
+            starts.append(starts_old[i])
+            stops.append(stops_old[i - 1])
+
+    stops.append(stops_old[-1])
+
+    return starts, stops
diff --git a/sparse/slicing.py b/sparse/slicing.py
index 962e792..56e5b25 100644
--- a/sparse/slicing.py
+++ b/sparse/slicing.py
@@ -3,6 +3,7 @@
 
 import math
 from numbers import Integral, Number
+
 import numpy as np
 
 
@@ -12,7 +13,7 @@ def normalize_index(idx, shape):
     2.  Adds full slices to end of index
     3.  Checks bounding conditions
     4.  Replaces numpy arrays with lists
-    5.  Posify's integers and lists
+    5.  Posify's slices integers and lists
     6.  Normalizes slices to canonical form
     Examples
     --------
@@ -23,9 +24,9 @@ def normalize_index(idx, shape):
     >>> normalize_index([-1], (10,))
     (array([9]),)
     >>> normalize_index(slice(-3, 10, 1), (10,))
-    (slice(7, None, None),)
+    (slice(7, 10, 1),)
     >>> normalize_index((Ellipsis, None), (10,))
-    (slice(None, None, None), None)
+    (slice(0, 10, 1), None)
     """
     if not isinstance(idx, tuple):
         idx = (idx,)
@@ -55,8 +56,9 @@ def normalize_index(idx, shape):
         if d is not None:
             check_index(i, d)
     idx = tuple(map(sanitize_index, idx))
-    idx = tuple(map(normalize_slice, idx, none_shape))
+    idx = tuple(map(replace_none, idx, none_shape))
     idx = posify_index(none_shape, idx)
+    idx = tuple(map(clip_slice, idx, none_shape))
     return idx
 
 
@@ -176,50 +178,6 @@ def _sanitize_index_element(ind):
     return int(ind)
 
 
-def normalize_slice(idx, dim):
-    """ Normalize slices to canonical form
-    Parameters
-    ----------
-    idx: slice or other index
-    dim: dimension length
-    Examples
-    --------
-    >>> normalize_slice(slice(0, 10, 1), 10)
-    slice(None, None, None)
-    """
-
-    if isinstance(idx, slice):
-        start, stop, step = idx.start, idx.stop, idx.step
-        if start is not None:
-            if start < 0 and not math.isnan(dim):
-                start = max(0, start + dim)
-            elif start > dim:
-                start = dim
-        if stop is not None:
-            if stop < 0 and not math.isnan(dim):
-                stop = max(0, stop + dim)
-            elif stop > dim:
-                stop = dim
-
-        step = 1 if step is None else step
-
-        if step > 0:
-            if start == 0:
-                start = None
-            if stop == dim:
-                stop = None
-        else:
-            if start == dim - 1:
-                start = None
-            if stop == -1:
-                stop = None
-
-        if step == 1:
-            step = None
-        return slice(start, stop, step)
-    return idx
-
-
 def posify_index(shape, ind):
     """ Flip negative indices around to positive ones
     >>> posify_index(10, 3)
@@ -243,5 +201,93 @@ def posify_index(shape, ind):
     if isinstance(ind, (np.ndarray, list)) and not math.isnan(shape):
         ind = np.asanyarray(ind)
         return np.where(ind < 0, ind + shape, ind)
+    if isinstance(ind, slice):
+        start, stop, step = ind.start, ind.stop, ind.step
+
+        if start < 0:
+            start += shape
+
+        if not (0 > stop >= step) and stop < 0:
+            stop += shape
+
+        return slice(start, stop, ind.step)
 
     return ind
+
+
+def clip_slice(idx, dim):
+    """
+    Clip slice to its effective size given the shape.
+
+    Parameters
+    ----------
+    idx : The index.
+    dim : The size along the corresponding dimension.
+
+    Returns
+    -------
+    idx : slice
+
+    Examples
+    --------
+    >>> clip_slice(slice(0, 20, 1), 10)
+    slice(0, 10, 1)
+    """
+    if not isinstance(idx, slice):
+        return idx
+
+    start, stop, step = idx.start, idx.stop, idx.step
+
+    if step > 0:
+        start = max(start, 0)
+        stop = min(stop, dim)
+
+        if start > stop:
+            start = stop
+    else:
+        start = min(start, dim - 1)
+        stop = max(stop, -1)
+
+        if start < stop:
+            start = stop
+
+    return slice(start, stop, step)
+
+
+def replace_none(idx, dim):
+    """
+    Normalize slices to canonical form, i.e.
+    replace ``None`` with the appropriate integers.
+
+    Parameters
+    ----------
+    idx: slice or other index
+    dim: dimension length
+
+    Examples
+    --------
+    >>> replace_none(slice(None, None, None), 10)
+    slice(0, 10, 1)
+    """
+    if not isinstance(idx, slice):
+        return idx
+
+    start, stop, step = idx.start, idx.stop, idx.step
+
+    if step is None:
+        step = 1
+
+    if step > 0:
+        if start is None:
+            start = 0
+
+        if stop is None:
+            stop = dim
+    else:
+        if start is None:
+            start = dim - 1
+
+        if stop is None:
+            stop = -1
+
+    return slice(start, stop, step)


Test patch:
diff --git a/sparse/tests/test_coo.py b/sparse/tests/test_coo.py
index d903683..6d83d93 100644
--- a/sparse/tests/test_coo.py
+++ b/sparse/tests/test_coo.py
@@ -1,14 +1,12 @@
-import pytest
-
 import operator
 
 import numpy as np
+import pytest
 import scipy.sparse
 import scipy.stats
 
-from sparse import COO
-
 import sparse
+from sparse import COO
 from sparse.utils import assert_eq, is_lexsorted, random_value_array
 
 
@@ -772,36 +770,27 @@ def test_gt():
 
 
 @pytest.mark.parametrize('index', [
+    # Integer
     0,
     1,
     -1,
+    (1, 1, 1),
+    # Pure slices
     (slice(0, 2),),
     (slice(None, 2), slice(None, 2)),
     (slice(1, None), slice(1, None)),
     (slice(None, None),),
+    (slice(None, None, -1),),
     (slice(None, 2, -1), slice(None, 2, -1)),
     (slice(1, None, 2), slice(1, None, 2)),
     (slice(None, None, 2),),
     (slice(None, 2, -1), slice(None, 2, -2)),
     (slice(1, None, 2), slice(1, None, 1)),
     (slice(None, None, -2),),
+    # Combinations
     (0, slice(0, 2),),
     (slice(0, 1), 0),
-    ([1, 0], 0),
-    (1, [0, 2]),
-    (0, [1, 0], 0),
-    (1, [2, 0], 0),
     (None, slice(1, 3), 0),
-    (Ellipsis, slice(1, 3)),
-    (1, Ellipsis, slice(1, 3)),
-    (slice(0, 1), Ellipsis),
-    (Ellipsis, None),
-    (None, Ellipsis),
-    (1, Ellipsis),
-    (1, Ellipsis, None),
-    (1, 1, 1),
-    (1, 1, 1, Ellipsis),
-    (Ellipsis, 1, None),
     (slice(0, 3), None, 0),
     (slice(1, 2), slice(2, 4)),
     (slice(1, 2), slice(None, None)),
@@ -815,10 +804,42 @@ def test_gt():
     (slice(2, 0, -1), slice(None, None), -1),
     (slice(-2, None, None),),
     (slice(-1, None, None), slice(-2, None, None)),
+    # With ellipsis
+    (Ellipsis, slice(1, 3)),
+    (1, Ellipsis, slice(1, 3)),
+    (slice(0, 1), Ellipsis),
+    (Ellipsis, None),
+    (None, Ellipsis),
+    (1, Ellipsis),
+    (1, Ellipsis, None),
+    (1, 1, 1, Ellipsis),
+    (Ellipsis, 1, None),
+    # Pathological - Slices larger than array
+    (slice(None, 1000)),
+    (slice(None), slice(None, 1000)),
+    (slice(None), slice(1000, -1000, -1)),
+    (slice(None), slice(1000, -1000, -50)),
+    # Pathological - Wrong ordering of start/stop
+    (slice(5, 0),),
+    (slice(0, 5, -1),),
+])
+def test_slicing(index):
+    s = sparse.random((2, 3, 4), density=0.5)
+    x = s.todense()
+
+    assert_eq(x[index], s[index])
+
+
+@pytest.mark.parametrize('index', [
+    ([1, 0], 0),
+    (1, [0, 2]),
+    (0, [1, 0], 0),
+    (1, [2, 0], 0),
     ([True, False], slice(1, None), slice(-2, None)),
     (slice(1, None), slice(-2, None), [True, False, True, False]),
 ])
-def test_slicing(index):
+@pytest.mark.xfail(reason='Advanced indexing is temporarily broken.')
+def test_advanced_indexing(index):
     s = sparse.random((2, 3, 4), density=0.5)
     x = s.todense()