<jupyter_start><jupyter_text>Advanced indexing<jupyter_code>import sys
sys.path.insert(0, '..')
import zarr
import numpy as np
np.random.seed(42)
import cProfile
zarr.__version__<jupyter_output><empty_output><jupyter_text>Functionality and API  Indexing a 1D array with a Boolean (mask) arraySupported via ``get/set_mask_selection()`` and ``.vindex[]``. Also supported via ``get/set_orthogonal_selection()`` and ``.oindex[]``.<jupyter_code>a = np.arange(10)
za = zarr.array(a, chunks=2)
ix = [False,  True,  False,  True, False, True, False,  True,  False,  True]
# get items
za.vindex[ix]
# get items
za.oindex[ix]
# set items
za.vindex[ix] = a[ix] * 10
za[:]
# set items
za.oindex[ix] = a[ix] * 100
za[:]
# if using .oindex, indexing array can be any array-like, e.g., Zarr array
zix = zarr.array(ix, chunks=2)
za = zarr.array(a, chunks=2)
za.oindex[zix]  # will not load all zix into memory<jupyter_output><empty_output><jupyter_text>Indexing a 1D array with a 1D integer (coordinate) arraySupported via ``get/set_coordinate_selection()`` and ``.vindex[]``. Also supported via ``get/set_orthogonal_selection()`` and ``.oindex[]``.<jupyter_code>a = np.arange(10)
za = zarr.array(a, chunks=2)
ix = [1, 3, 5, 7, 9]
# get items
za.vindex[ix]
# get items
za.oindex[ix]
# set items
za.vindex[ix] = a[ix] * 10
za[:]
# set items
za.oindex[ix] = a[ix] * 100
za[:]<jupyter_output><empty_output><jupyter_text>Indexing a 1D array with a multi-dimensional integer (coordinate) arraySupported via ``get/set_coordinate_selection()`` and ``.vindex[]``.<jupyter_code>a = np.arange(10)
za = zarr.array(a, chunks=2)
ix = np.array([[1, 3, 5], [2, 4, 6]])
# get items
za.vindex[ix]
# set items
za.vindex[ix] = a[ix] * 10
za[:]<jupyter_output><empty_output><jupyter_text>Slicing a 1D array with step > 1Slices with step > 1 are supported via ``get/set_basic_selection()``, ``get/set_orthogonal_selection()``, ``__getitem__`` and ``.oindex[]``. Negative steps are not supported.<jupyter_code>a = np.arange(10)
za = zarr.array(a, chunks=2)
# get items
za[1::2]
# set items
za.oindex[1::2] = a[1::2] * 10
za[:]<jupyter_output><empty_output><jupyter_text>Orthogonal (outer) indexing of multi-dimensional arraysOrthogonal (a.k.a. outer) indexing is supported with either Boolean or integer arrays, in combination with integers and slices. This functionality is provided via the ``get/set_orthogonal_selection()`` methods. For convenience, this functionality is also available via the ``.oindex[]`` property.<jupyter_code>a = np.arange(15).reshape(5, 3)
za = zarr.array(a, chunks=(3, 2))
za[:]
# orthogonal indexing with Boolean arrays
ix0 = [False, True, False, True, False]
ix1 = [True, False, True]
za.get_orthogonal_selection((ix0, ix1))
# alternative API
za.oindex[ix0, ix1]
# orthogonal indexing with integer arrays
ix0 = [1, 3]
ix1 = [0, 2]
za.get_orthogonal_selection((ix0, ix1))
# alternative API
za.oindex[ix0, ix1]
# combine with slice
za.oindex[[1,  3], :]
# combine with slice
za.oindex[:, [0, 2]]
# set items via Boolean selection
ix0 = [False, True, False, True, False]
ix1 = [True, False, True]
selection = ix0, ix1
value = 42
za.set_orthogonal_selection(selection, value)
za[:]
# alternative API
za.oindex[ix0, ix1] = 44
za[:]
# set items via integer selection
ix0 = [1, 3]
ix1 = [0, 2]
selection = ix0, ix1
value = 46
za.set_orthogonal_selection(selection, value)
za[:]
# alternative API
za.oindex[ix0, ix1] = 48
za[:]<jupyter_output><empty_output><jupyter_text>Coordinate indexing of multi-dimensional arraysSelecting arbitrary points from a multi-dimensional array by indexing with integer (coordinate) arrays is supported. This functionality is provided via the ``get/set_coordinate_selection()`` methods. For convenience, this functionality is also available via the ``.vindex[]`` property.<jupyter_code>a = np.arange(15).reshape(5, 3)
za = zarr.array(a, chunks=(3, 2))
za[:]
# get items
ix0 = [1, 3]
ix1 = [0, 2]
za.get_coordinate_selection((ix0, ix1))
# alternative API
za.vindex[ix0, ix1]
# set items
za.set_coordinate_selection((ix0, ix1), 42)
za[:]
# alternative API
za.vindex[ix0, ix1] = 44
za[:]<jupyter_output><empty_output><jupyter_text>Mask indexing of multi-dimensional arraysSelecting arbitrary points from a multi-dimensional array by a Boolean array is supported. This functionality is provided via the ``get/set_mask_selection()`` methods. For convenience, this functionality is also available via the ``.vindex[]`` property.<jupyter_code>a = np.arange(15).reshape(5, 3)
za = zarr.array(a, chunks=(3, 2))
za[:]
ix = np.zeros_like(a, dtype=bool)
ix[1, 0] = True
ix[3, 2] = True
za.get_mask_selection(ix)
za.vindex[ix]
za.set_mask_selection(ix, 42)
za[:]
za.vindex[ix] = 44
za[:]<jupyter_output><empty_output><jupyter_text>Selecting fields from arrays with a structured dtypeAll ``get/set_selection_...()`` methods support a ``fields`` argument which allows retrieving/replacing data for a specific field or fields. Also h5py-like API is supported where fields can be provided within ``__getitem__``, ``.oindex[]`` and ``.vindex[]``.<jupyter_code>a = np.array([(b'aaa', 1, 4.2),
              (b'bbb', 2, 8.4),
              (b'ccc', 3, 12.6)], 
             dtype=[('foo', 'S3'), ('bar', 'i4'), ('baz', 'f8')])
za = zarr.array(a, chunks=2, fill_value=None)
za[:]
za['foo']
za['foo', 'baz']
za[:2, 'foo']
za[:2, 'foo', 'baz']
za.oindex[[0, 2], 'foo']
za.vindex[[0, 2], 'foo']
za['bar'] = 42
za[:]
za[:2, 'bar'] = 84
za[:]<jupyter_output><empty_output><jupyter_text>Note that this API differs from numpy when selecting multiple fields. E.g.:<jupyter_code>a['foo', 'baz']
a[['foo', 'baz']]
za['foo', 'baz']
za[['foo', 'baz']]<jupyter_output><empty_output><jupyter_text>1D Benchmarking<jupyter_code>c = np.arange(100000000)
c.nbytes
%time zc = zarr.array(c)
zc.info
%timeit c.copy()
%timeit zc[:]<jupyter_output>254 ms ± 942 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>bool dense selection<jupyter_code># relatively dense selection - 10%
ix_dense_bool = np.random.binomial(1, 0.1, size=c.shape[0]).astype(bool)
np.count_nonzero(ix_dense_bool)
%timeit c[ix_dense_bool]
%timeit zc.oindex[ix_dense_bool]
%timeit zc.vindex[ix_dense_bool]
import tempfile
import cProfile
import pstats

def profile(statement, sort='time', restrictions=(7,)):
    with tempfile.NamedTemporaryFile() as f:
        cProfile.run(statement, filename=f.name)
        pstats.Stats(f.name).sort_stats(sort).print_stats(*restrictions)

profile('zc.oindex[ix_dense_bool]')<jupyter_output>Wed Nov  8 17:17:48 2017    /tmp/tmpruua2rs_

         98386 function calls in 0.483 seconds

   Ordered by: internal time
   List reduced from 83 to 7 due to restriction <7>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     1025    0.197    0.000    0.197    0.000 {method 'nonzero' of 'numpy.ndarray' objects}
     1024    0.149    0.000    0.159    0.000 ../zarr/core.py:1028(_decode_chunk)
     1024    0.044    0.000    0.231    0.000 ../zarr/core.py:849(_chunk_getitem)
     1024    0.009    0.000    0.009    0.000 {built-in method numpy.core.multiarray.count_nonzero}
     1025    0.007    0.000    0.238    0.000 ../zarr/indexing.py:541(__iter__)
     1024    0.006    0.000    0.207    0.000 /home/aliman/pyenv/zarr_20171023/lib/python3.6/site-packages/numpy/lib/index_tricks.py:26(ix_)
     2048    0.005    0.000    0.005    0.000 ../zarr/core.py:337(<genexpr>)<jupyter_text>Method ``nonzero`` is being called internally within numpy to convert bool to int selections, no way to avoid.<jupyter_code>profile('zc.vindex[ix_dense_bool]')<jupyter_output>Wed Nov  8 17:18:06 2017    /tmp/tmp7_bautep

         52382 function calls in 0.592 seconds

   Ordered by: internal time
   List reduced from 88 to 7 due to restriction <7>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        2    0.219    0.110    0.219    0.110 {method 'nonzero' of 'numpy.ndarray' objects}
     1024    0.096    0.000    0.101    0.000 ../zarr/core.py:1028(_decode_chunk)
        2    0.094    0.047    0.094    0.047 ../zarr/indexing.py:630(<genexpr>)
     1024    0.044    0.000    0.167    0.000 ../zarr/core.py:849(_chunk_getitem)
        1    0.029    0.029    0.029    0.029 {built-in method numpy.core.multiarray.ravel_multi_index}
        1    0.023    0.023    0.023    0.023 {built-in method numpy.core.multiarray.bincount}
        1    0.021    0.021    0.181    0.181 ../zarr/indexing.py:603(__init__)<jupyter_text>``.vindex[]`` is a bit slower, possibly because internally it converts to a coordinate array first.  int dense selection<jupyter_code>ix_dense_int = np.random.choice(c.shape[0], size=c.shape[0]//10, replace=True)
ix_dense_int_sorted = ix_dense_int.copy()
ix_dense_int_sorted.sort()
len(ix_dense_int)
%timeit c[ix_dense_int_sorted]
%timeit zc.oindex[ix_dense_int_sorted]
%timeit zc.vindex[ix_dense_int_sorted]
%timeit c[ix_dense_int]
%timeit zc.oindex[ix_dense_int]
%timeit zc.vindex[ix_dense_int]
profile('zc.oindex[ix_dense_int_sorted]')
profile('zc.vindex[ix_dense_int_sorted]')
profile('zc.oindex[ix_dense_int]')
profile('zc.vindex[ix_dense_int]')<jupyter_output>Wed Nov  8 17:19:22 2017    /tmp/tmpbskhj8de

         50320 function calls in 1.730 seconds

   Ordered by: internal time
   List reduced from 86 to 7 due to restriction <7>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    1.116    1.116    1.116    1.116 {method 'argsort' of 'numpy.ndarray' objects}
     1024    0.133    0.000    0.275    0.000 ../zarr/core.py:849(_chunk_getitem)
        2    0.121    0.060    0.121    0.060 ../zarr/indexing.py:654(<genexpr>)
     1024    0.113    0.000    0.119    0.000 ../zarr/core.py:1028(_decode_chunk)
        2    0.100    0.050    0.100    0.050 ../zarr/indexing.py:630(<genexpr>)
        1    0.030    0.030    0.030    0.030 {built-in method numpy.core.multiarray.ravel_multi_index}
        1    0.024    0.024    1.427    1.427 ../zarr/indexing.py:603(__init__)<jupyter_text>When indices are not sorted, zarr needs to partially sort them so the occur in chunk order, so we only have to visit each chunk once. This sorting dominates the processing time and is unavoidable AFAIK.  bool sparse selection<jupyter_code># relatively sparse selection
ix_sparse_bool = np.random.binomial(1, 0.0001, size=c.shape[0]).astype(bool)
np.count_nonzero(ix_sparse_bool)
%timeit c[ix_sparse_bool]
%timeit zc.oindex[ix_sparse_bool]
%timeit zc.vindex[ix_sparse_bool]
profile('zc.oindex[ix_sparse_bool]')
profile('zc.vindex[ix_sparse_bool]')<jupyter_output>Wed Nov  8 17:20:09 2017    /tmp/tmphsko8nvh

         52382 function calls in 0.160 seconds

   Ordered by: internal time
   List reduced from 88 to 7 due to restriction <7>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     1024    0.093    0.000    0.098    0.000 ../zarr/core.py:1028(_decode_chunk)
        2    0.017    0.008    0.017    0.008 {method 'nonzero' of 'numpy.ndarray' objects}
     1025    0.008    0.000    0.014    0.000 ../zarr/indexing.py:674(__iter__)
     1024    0.006    0.000    0.127    0.000 ../zarr/core.py:849(_chunk_getitem)
     2048    0.004    0.000    0.004    0.000 ../zarr/indexing.py:695(<genexpr>)
     2054    0.003    0.000    0.003    0.000 ../zarr/core.py:337(<genexpr>)
     1024    0.002    0.000    0.005    0.000 /home/aliman/pyenv/zarr_20171023/lib/python3.6/site-packages/numpy/core/arrayprint.py:381(wrapper)<jupyter_text>int sparse selection<jupyter_code>ix_sparse_int = np.random.choice(c.shape[0], size=c.shape[0]//10000, replace=True)
ix_sparse_int_sorted = ix_sparse_int.copy()
ix_sparse_int_sorted.sort()
len(ix_sparse_int)
%timeit c[ix_sparse_int_sorted]
%timeit c[ix_sparse_int]
%timeit zc.oindex[ix_sparse_int_sorted]
%timeit zc.vindex[ix_sparse_int_sorted]
%timeit zc.oindex[ix_sparse_int]
%timeit zc.vindex[ix_sparse_int]
profile('zc.oindex[ix_sparse_int]')
profile('zc.vindex[ix_sparse_int]')<jupyter_output>Wed Nov  8 17:21:19 2017    /tmp/tmpdwju98kn

         50320 function calls in 0.167 seconds

   Ordered by: internal time
   List reduced from 86 to 7 due to restriction <7>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     1024    0.105    0.000    0.111    0.000 ../zarr/core.py:1028(_decode_chunk)
     1025    0.009    0.000    0.017    0.000 ../zarr/indexing.py:674(__iter__)
     1024    0.006    0.000    0.142    0.000 ../zarr/core.py:849(_chunk_getitem)
     2048    0.005    0.000    0.005    0.000 ../zarr/indexing.py:695(<genexpr>)
     2054    0.004    0.000    0.004    0.000 ../zarr/core.py:337(<genexpr>)
        1    0.003    0.003    0.162    0.162 ../zarr/core.py:591(_get_selection)
     1027    0.003    0.000    0.003    0.000 {method 'reshape' of 'numpy.ndarray' objects}<jupyter_text>For sparse selections, processing time is dominated by decompression, so we can't do any better.  sparse bool selection as zarr array<jupyter_code>zix_sparse_bool = zarr.array(ix_sparse_bool)
zix_sparse_bool.info
%timeit zc.oindex[zix_sparse_bool]<jupyter_output>387 ms ± 5.47 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>slice with step<jupyter_code>%timeit np.array(c[::2])
%timeit zc[::2]
%timeit zc[::10]
%timeit zc[::100]
%timeit zc[::1000]
profile('zc[::2]')<jupyter_output>Wed Nov  8 17:22:44 2017    /tmp/tmpg9dxqcpg

         49193 function calls in 0.211 seconds

   Ordered by: internal time
   List reduced from 55 to 7 due to restriction <7>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
     1024    0.104    0.000    0.110    0.000 ../zarr/core.py:1028(_decode_chunk)
     1024    0.067    0.000    0.195    0.000 ../zarr/core.py:849(_chunk_getitem)
     1025    0.005    0.000    0.013    0.000 ../zarr/indexing.py:278(__iter__)
     2048    0.004    0.000    0.004    0.000 ../zarr/core.py:337(<genexpr>)
     2050    0.003    0.000    0.003    0.000 ../zarr/indexing.py:90(ceildiv)
     1025    0.003    0.000    0.006    0.000 ../zarr/indexing.py:109(__iter__)
     1024    0.003    0.000    0.003    0.000 {method 'reshape' of 'numpy.ndarray' objects}<jupyter_text>2D Benchmarking<jupyter_code>c.shape
d = c.reshape(-1, 1000)
d.shape
zd = zarr.array(d)
zd.info<jupyter_output><empty_output><jupyter_text>bool orthogonal selection<jupyter_code>ix0 = np.random.binomial(1, 0.5, size=d.shape[0]).astype(bool)
ix1 = np.random.binomial(1, 0.5, size=d.shape[1]).astype(bool)
%timeit d[np.ix_(ix0, ix1)]
%timeit zd.oindex[ix0, ix1]<jupyter_output>373 ms ± 5.45 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>int orthogonal selection<jupyter_code>ix0 = np.random.choice(d.shape[0], size=int(d.shape[0] * .5), replace=True)
ix1 = np.random.choice(d.shape[1], size=int(d.shape[1] * .5), replace=True)
%timeit d[np.ix_(ix0, ix1)]
%timeit zd.oindex[ix0, ix1]<jupyter_output>566 ms ± 12.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)<jupyter_text>coordinate (point) selection<jupyter_code>n = int(d.size * .1)
ix0 = np.random.choice(d.shape[0], size=n, replace=True)
ix1 = np.random.choice(d.shape[1], size=n, replace=True)
n
%timeit d[ix0, ix1]
%timeit zd.vindex[ix0, ix1]
profile('zd.vindex[ix0, ix1]')<jupyter_output>Wed Nov  8 17:24:31 2017    /tmp/tmp7c68z70p

         62673 function calls in 2.065 seconds

   Ordered by: internal time
   List reduced from 88 to 7 due to restriction <7>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    1.112    1.112    1.112    1.112 {method 'argsort' of 'numpy.ndarray' objects}
        3    0.244    0.081    0.244    0.081 ../zarr/indexing.py:654(<genexpr>)
        3    0.193    0.064    0.193    0.064 ../zarr/indexing.py:630(<genexpr>)
     1024    0.170    0.000    0.350    0.000 ../zarr/core.py:849(_chunk_getitem)
     1024    0.142    0.000    0.151    0.000 ../zarr/core.py:1028(_decode_chunk)
        1    0.044    0.044    0.044    0.044 {built-in method numpy.core.multiarray.ravel_multi_index}
        1    0.043    0.043    1.676    1.676 ../zarr/indexing.py:603(__init__)<jupyter_text>Points need to be partially sorted so all points in the same chunk are grouped and processed together. This requires ``argsort`` which dominates time.  h5py comparisonN.B., not really fair because using slower compressor, but for interest...<jupyter_code>import h5py
import tempfile
h5f = h5py.File(tempfile.mktemp(), driver='core', backing_store=False)
hc = h5f.create_dataset('c', data=c, compression='gzip', compression_opts=1, chunks=zc.chunks, shuffle=True)
hc
%time hc[:]
%time hc[ix_sparse_bool]
# # this is pathological, takes minutes 
# %time hc[ix_dense_bool]
# this is pretty slow
%time hc[::1000]<jupyter_output>CPU times: user 38.3 s, sys: 136 ms, total: 38.4 s
Wall time: 38.1 s
{'fraction_non_alphanumeric': 0.1375231495310353, 'fraction_numerical': 0.09307604994324631, 'mean_word_length': 3.7194812517620526, 'pattern_counts': {'":': 0, '<': 99, '<?xml version=': 0, '>': 101, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 1, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '23', 'dir': 'jupyter-structured-clean-dedup', 'id': '224153', 'license': 'MIT', 'n_tokens_mistral': 7435, 'n_tokens_neox': 6401, 'n_words': 1557, 'path': 'notebooks/advanced_indexing.ipynb', 'repo_name': 'hmaarrfk/zarr-python', 'size': '16739'}