Improve ndindex performance (PyInf#12655)

[ArvidJB]

We're using ndindex as part of versioned-hdf5. Unfortunately in some common circumstances it can be pretty slow. Consider the following code:

>>> # THIS CELL WAS AUTO-GENERATED BY PYFLYBY
>>> import atexit
>>> import h5py
>>> import numpy as np
>>> import shutil
>>> import tempfile
>>> from versioned_hdf5 import VersionedHDF5File
>>> # END AUTO-GENERATED BLOCK


>>> d = tempfile.mkdtemp()
>>> atexit.register(shutil.rmtree, d)

[PYFLYBY] import atexit
[PYFLYBY] import shutil
[PYFLYBY] import tempfile

<function shutil.rmtree(path, ignore_errors=False, onerror=None, *, dir_fd=None)>

>>> def r0(d):
...     with h5py.File(f'{d}/data.h5', mode="w") as f:
...         vf = VersionedHDF5File(f)
...         with vf.stage_version("r0") as sv:
...             sv.create_dataset('values', data=np.arange(100_000_000), chunks=(1_000,), maxshape=(None,))

[PYFLYBY] from versioned_hdf5 import VersionedHDF5File
[PYFLYBY] import h5py
[PYFLYBY] import numpy as np


>>> def r1(d):
...     with h5py.File(f'{d}/data.h5', mode="r+") as f:
...         vf = VersionedHDF5File(f)
...         with vf.stage_version("r1") as sv:
...             values = sv['values']
...             # resizing creates an InMemoryDataset and populates data_dict
...             values.resize((110_000_000,))
...             # reading from InMemoryDataset is now slow
...             _ = values[500:90_000_000]


>>> %load_ext pyinstrument


>>> r0(d)
>>> %pyinstrument r1(d)


  _     ._   __/__   _ _  _  _ _/_   Recorded: 14:20:41  Samples:  30357
 /_//_/// /_\ / //_// / //_'/ //     Duration: 42.410    CPU time: 41.642
/   _/                      v4.6.1

Program: /usr/local/python/python-3.11/std/lib64/python3.11/site-packages/ipykernel_launcher.py -f /u/bessen/.local/share/jupyter/runtime/kernel-f2241356-7014-4cb5-99a3-98f8ade0d4d7.json

42.409 <module>  ../../../tmp/ipykernel_3610251/2296172108.py:1
`- 42.407 r1  ../../../tmp/ipykernel_3610251/2453629480.py:1
   |- 15.246 DatasetWrapper.__getitem__  versioned_hdf5/wrappers.py:1267
   |  `- 15.246 InMemoryDataset.__getitem__  versioned_hdf5/wrappers.py:757
   |     `- 15.246 InMemoryDataset.get_index  versioned_hdf5/wrappers.py:655
   |        |- 9.144 [self]  versioned_hdf5/wrappers.py
   |        |- 4.383 InMemoryDatasetID._read_chunk  versioned_hdf5/wrappers.py:1385
   |        |  `- 4.275 Dataset.__getitem__  h5py/_hl/dataset.py:749
   |        |     |- 1.916 Reader.read  <built-in>
   |        |     |- 1.704 Dataset._fast_reader  h5py/_hl/dataset.py:527
   |        |     |  `- 1.689 [self]  h5py/_hl/dataset.py
   |        |     `- 0.604 [self]  h5py/_hl/dataset.py
   |        `- 0.878 where  <__array_function__ internals>:177
   |              [2 frames hidden]  <__array_function__ internals>, <buil...
   |- 13.388 _GeneratorContextManager.__exit__  contextlib.py:141
   |  `- 13.388 VersionedHDF5File.stage_version  versioned_hdf5/api.py:267
   |     `- 13.371 commit_version  versioned_hdf5/versions.py:71
   |        |- 8.712 create_virtual_dataset  versioned_hdf5/backend.py:429
   |        |  |- 2.525 [self]  versioned_hdf5/backend.py
   |        |  |- 2.068 Group.create_virtual_dataset  h5py/_hl/group.py:188
   |        |  |  `- 1.646 VirtualLayout.make_dataset  h5py/_hl/vds.py:228
   |        |  |- 1.932 select  h5py/_hl/selections.py:19
   |        |  |  `- 1.452 [self]  h5py/_hl/selections.py
   |        |  |- 0.728 Dataset.shape  h5py/_hl/dataset.py:467
   |        |  |  `- 0.699 [self]  h5py/_hl/dataset.py
   |        |  `- 0.530 <dictcomp>  versioned_hdf5/backend.py:437
   |        |     `- 0.525 [self]  versioned_hdf5/backend.py
   |        |- 3.961 write_dataset_chunks  versioned_hdf5/backend.py:346
   |        |  |- 2.349 Hashtable.hash  versioned_hdf5/hashtable.py:116
   |        |  |  `- 1.822 openssl_sha256  <built-in>
   |        |  `- 0.659 Hashtable.__contains__  <frozen _collections_abc>:778
   |        |     `- 0.609 Hashtable.__getitem__  versioned_hdf5/hashtable.py:205
   |        |        `- 0.595 [self]  versioned_hdf5/hashtable.py
   |        `- 0.688 [self]  versioned_hdf5/versions.py
   |- 9.595 InMemoryDataset.resize  versioned_hdf5/wrappers.py:609
   |  |- 5.824 InMemoryDatasetID.data_dict  versioned_hdf5/wrappers.py:1298
   |  |  |- 2.915 <dictcomp>  versioned_hdf5/wrappers.py:1327
   |  |  |  `- 2.680 spaceid_to_slice  versioned_hdf5/slicetools.py:6
   |  |  |     `- 2.317 [self]  versioned_hdf5/slicetools.py
   |  |  |- 1.683 <listcomp>  versioned_hdf5/wrappers.py:1317
   |  |  |  `- 1.627 [self]  versioned_hdf5/wrappers.py
   |  |  `- 0.830 [self]  versioned_hdf5/wrappers.py
   |  |- 1.712 [self]  versioned_hdf5/wrappers.py
   |  `- 1.710 InMemoryDataset.get_index  versioned_hdf5/wrappers.py:655
   |     `- 1.710 InMemoryDataset.__getitem__  h5py/_hl/dataset.py:749
   |        `- 1.706 Reader.read  <built-in>
   |- 3.185 _GeneratorContextManager.__enter__  contextlib.py:132
   |  `- 3.185 VersionedHDF5File.stage_version  versioned_hdf5/api.py:267
   |     `- 3.185 create_version_group  versioned_hdf5/versions.py:22
   |        `- 3.184 Group.visititems  h5py/_hl/group.py:635
   |           |- 2.613 proxy  h5py/_hl/group.py:660
   |           |  |- 1.673 _get  versioned_hdf5/versions.py:57
   |           |  |  `- 1.275 InMemoryGroup.__setitem__  versioned_hdf5/wrappers.py:110
   |           |  |     `- 1.275 InMemoryGroup._add_to_data  versioned_hdf5/wrappers.py:114
   |           |  |        `- 1.275 InMemoryDataset.__init__  versioned_hdf5/wrappers.py:503
   |           |  |           `- 0.592 KeysViewHDF5.__iter__  <frozen _collections_abc>:835
   |           |  |              `- 0.592 AttributeManager.__iter__  h5py/_hl/attrs.py:257
   |           |  |                 `- 0.441 [self]  h5py/_hl/attrs.py
   |           |  `- 0.791 Group.__getitem__  h5py/_hl/group.py:348
   |           |     `- 0.544 [self]  h5py/_hl/group.py
   |           `- 0.571 [self]  h5py/_hl/group.py
   `- 0.985 File.__exit__  h5py/_hl/files.py:601
      `- 0.985 File.close  h5py/_hl/files.py:576



>>> %load_ext line_profiler


>>> from versioned_hdf5.wrappers import InMemoryDataset
>>> r0(d)
>>> %lprun -f InMemoryDataset.get_index.__wrapped__ r1(d)

Timer unit: 1e-09 s

Total time: 14.9595 s
File: /codemill/bessen/ndindex_venv/lib64/python3.11/site-packages/versioned_hdf5/wrappers.py
Function: get_index at line 655

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
   655                                               @with_phil
   656                                               def get_index(
   657                                                   self,
   658                                                   args: Union[slice, Slice, Tuple, tuple, h5r.RegionReference],
   659                                                   new_dtype: Optional[str] = None,
   660                                                   can_read_direct: Optional[bool] = None,
   661                                               ) -> np.ndarray:
   662                                                   """Read a slice from the HDF5 dataset given by the index.
   663                                           
   664                                                   Takes slices and recarray-style field names (more than one is
   665                                                   allowed!) in any order.  Obeys basic NumPy rules, including
   666                                                   broadcasting.
   667                                           
   668                                                   Parameters
   669                                                   ----------
   670                                                   args : Union[slice, Slice, Tuple, tuple, h5r.RegionReference]
   671                                                       Index to read from the Dataset
   672                                                   new_dtype : Optional[str]
   673                                                       Dtype of the returned array
   674                                                   can_read_direct : Optional[bool]
   675                                                       True if we can read directly from the underlying hdf5 Dataset, False otherwise.
   676                                                       This should be the value of the InMemoryDatasetID instance's ``can_read_direct``
   677                                                       property for this Dataset.
   678                                           
   679                                                       If None, ``self.id.can_read_direct`` is evaluated first to determine if data can
   680                                                       be read directly from the underlying dataset.
   681                                           
   682                                                   Returns
   683                                                   -------
   684                                                   np.ndarray
   685                                                       Array containing data from this dataset from the requested index
   686                                                   """
   687                                                   # This boilerplate code is based on h5py.Dataset.__getitem__
   688         2       8062.0   4031.0      0.0          args = args if isinstance(args, tuple) else (args,)
   689                                           
   690                                                   # Sort field names from the rest of the args.
   691         2      10665.0   5332.5      0.0          names = tuple(x for x in args if isinstance(x, str))
   692                                           
   693         2       2500.0   1250.0      0.0          if names:
   694                                                       # Read a subset of the fields in this structured dtype
   695                                                       if len(names) == 1:
   696                                                           names = names[0]  # Read with simpler dtype of this field
   697                                                       args = tuple(x for x in args if not isinstance(x, str))
   698                                                       return self.fields(names, _prior_dtype=new_dtype)[args]
   699                                           
   700         2       1745.0    872.5      0.0          if new_dtype is None:
   701         2      20228.0  10114.0      0.0              new_dtype = self.dtype
   702         2      73856.0  36928.0      0.0          mtype = h5t.py_create(new_dtype)
   703                                           
   704                                                   # === Special-case region references ====
   705                                           
   706         2       5677.0   2838.5      0.0          if len(args) == 1 and isinstance(args[0], h5r.RegionReference):
   707                                                       obj = h5r.dereference(args[0], self.id)
   708                                                       if obj != self.id:
   709                                                           raise ValueError("Region reference must point to this dataset")
   710                                           
   711                                                       sid = h5r.get_region(args[0], self.id)
   712                                                       mshape = guess_shape(sid)
   713                                                       if mshape is None:
   714                                                           # 0D with no data (NULL or deselected SCALAR)
   715                                                           return Empty(new_dtype)
   716                                                       out = np.empty(mshape, dtype=new_dtype)
   717                                                       if out.size == 0:
   718                                                           return out
   719                                           
   720                                                       sid_out = h5s.create_simple(mshape)
   721                                                       sid_out.select_all()
   722                                                       self.id.read(sid_out, sid, out, mtype)
   723                                                       return out
   724                                           
   725                                                   # === END CODE FROM h5py.Dataset.__getitem__ ===
   726                                           
   727         2     328533.0 164266.5      0.0          idx = ndindex(args).expand(self.shape)
   728                                           
   729         2       1098.0    549.0      0.0          if can_read_direct is None:
   730         1   48158088.0    5e+07      0.3              can_read_direct = self.id.can_read_direct
   731                                           
   732         2       1628.0    814.0      0.0          if can_read_direct:
   733         1 1411348097.0    1e+09      9.4              return super().__getitem__(idx.raw)
   734                                           
   735         1    1542927.0    2e+06      0.0          arr = np.ndarray(idx.newshape(self.shape), new_dtype, order="C")
   736                                           
   737     90001  994078595.0  11045.2      6.6          for chunk in self.chunks.as_subchunks(idx, self.shape):
   738     90000  367273224.0   4080.8      2.5              if chunk not in self.id.data_dict:
   739                                                           self.id.data_dict[chunk] = np.broadcast_to(
   740                                                               self.fillvalue, chunk.newshape(self.shape)
   741                                                           )
   742     90000  295761742.0   3286.2      2.0              elif isinstance(self.id.data_dict[chunk], (slice, Slice, tuple, Tuple)):
   743    360000  439296532.0   1220.3      2.9                  raw_idx = Tuple(
   744     90000  226675754.0   2518.6      1.5                      self.id.data_dict[chunk],
   745     90000  109581410.0   1217.6      0.7                      *[slice(0, len(i)) for i in chunk.args[1:]],
   746     90000   99604868.0   1106.7      0.7                  ).raw
   747     90000 4136405229.0  45960.1     27.7                  self.id.data_dict[chunk] = self.id._read_chunk(raw_idx)
   748                                           
   749     90000  269383258.0   2993.1      1.8              if self.id.data_dict[chunk].size != 0:
   750     90000 3196593424.0  35517.7     21.4                  arr_idx = chunk.as_subindex(idx)
   751     90000 2472868596.0  27476.3     16.5                  index = idx.as_subindex(chunk)
   752     90000  890518785.0   9894.7      6.0                  arr[arr_idx.raw] = self.id.data_dict[chunk][index.raw]
   753                                           
   754                                                   # Return arr as a scalar if it is shape () (matching h5py)
   755         1       1263.0   1263.0      0.0          return arr[()]

In particular note that in InMemoryDataset.get_index we spend 9.144 seconds in "[self]" and 0.878 seconds in "where":

   |     `- 15.246 InMemoryDataset.get_index  versioned_hdf5/wrappers.py:655
   |        |- 9.144 [self]  versioned_hdf5/wrappers.py
   |        |- 4.383 InMemoryDatasetID._read_chunk  versioned_hdf5/wrappers.py:1385
   |        |  `- 4.275 Dataset.__getitem__  h5py/_hl/dataset.py:749
   |        |     |- 1.916 Reader.read  <built-in>
   |        |     |- 1.704 Dataset._fast_reader  h5py/_hl/dataset.py:527
   |        |     |  `- 1.689 [self]  h5py/_hl/dataset.py
   |        |     `- 0.604 [self]  h5py/_hl/dataset.py
   |        `- 0.878 where  <__array_function__ internals>:177
   |              [2 frames hidden]  <__array_function__ internals>, <buil...

These timings come from ndindex but don't show up properly because it is compiled as Cython extension. If I disable cythonization you can see where it's spending it's time (and it's also even slower):

>>> r0(d)
>>> %pyinstrument r1(d)


  _     ._   __/__   _ _  _  _ _/_   Recorded: 14:24:32  Samples:  55759
 /_//_/// /_\ / //_// / //_'/ //     Duration: 66.820    CPU time: 66.438
/   _/                      v4.6.1

Program: /usr/local/python/python-3.11/std/lib64/python3.11/site-packages/ipykernel_launcher.py -f /u/bessen/.local/share/jupyter/runtime/kernel-f2241356-7014-4cb5-99a3-98f8ade0d4d7.json

66.820 <module>  ../../../tmp/ipykernel_3619864/2296172108.py:1
`- 66.817 r1  ../../../tmp/ipykernel_3619864/2453629480.py:1
   |- 32.529 DatasetWrapper.__getitem__  versioned_hdf5/wrappers.py:1267
   |  `- 32.529 InMemoryDataset.__getitem__  versioned_hdf5/wrappers.py:757
   |     `- 32.529 InMemoryDataset.get_index  versioned_hdf5/wrappers.py:655
   |        |- 19.605 Tuple.as_subindex  ndindex/tuple.py:627
   |        |  |- 8.551 Slice.as_subindex  ndindex/slice.py:501
   |        |  |  |- 3.618 Slice.reduce  ndindex/slice.py:212
   |        |  |  |  |- 2.811 Slice.__init__  ndindex/ndindex.py:159
   |        |  |  |  |  `- 2.539 Slice._typecheck  ndindex/slice.py:62
   |        |  |  |  |     `- 1.768 operator_index  ndindex/ndindex.py:681
   |        |  |  |  |        `- 1.345 [self]  ndindex/ndindex.py
   |        |  |  |  `- 0.807 [self]  ndindex/slice.py
   |        |  |  |- 2.226 subindex_slice  ndindex/subindex_helpers.py:83
   |        |  |  |  `- 1.114 where  ndindex/subindex_helpers.py:46
   |        |  |  |     `- 0.722 where  <__array_function__ internals>:177
   |        |  |  |- 1.144 [self]  ndindex/slice.py
   |        |  |  `- 0.997 Slice.__init__  ndindex/ndindex.py:159
   |        |  |     `- 0.911 Slice._typecheck  ndindex/slice.py:62
   |        |  |- 4.300 Tuple.reduce  ndindex/tuple.py:185
   |        |  |  |- 1.656 [self]  ndindex/tuple.py
   |        |  |  `- 1.513 Slice.reduce  ndindex/slice.py:212
   |        |  |     `- 1.107 Slice.__init__  ndindex/ndindex.py:159
   |        |  |        `- 0.972 Slice._typecheck  ndindex/slice.py:62
   |        |  |           `- 0.697 operator_index  ndindex/ndindex.py:681
   |        |  |- 3.529 Tuple.__init__  ndindex/ndindex.py:159
   |        |  |  `- 3.261 Tuple._typecheck  ndindex/tuple.py:49
   |        |  |     `- 1.929 [self]  ndindex/tuple.py
   |        |  `- 1.650 [self]  ndindex/tuple.py
   |        |- 4.108 InMemoryDatasetID._read_chunk  versioned_hdf5/wrappers.py:1385
   |        |  `- 4.002 Dataset.__getitem__  h5py/_hl/dataset.py:749
   |        |     |- 1.817 Reader.read  <built-in>
   |        |     `- 1.655 Dataset._fast_reader  h5py/_hl/dataset.py:527
   |        |        `- 1.637 [self]  h5py/_hl/dataset.py
   |        |- 2.265 ChunkSize.as_subchunks  ndindex/chunking.py:143
   |        |  `- 2.230 _indices  ndindex/chunking.py:288
   |        |     |- 0.909 <listcomp>  ndindex/chunking.py:292
   |        |     `- 0.871 Tuple.__init__  ndindex/ndindex.py:159
   |        |        `- 0.801 Tuple._typecheck  ndindex/tuple.py:49
   |        |- 2.179 [self]  versioned_hdf5/wrappers.py
   |        |- 1.487 Tuple.__hash__  ndindex/tuple.py:114
   |        |  `- 1.071 Slice.__hash__  ndindex/slice.py:85
   |        |- 1.057 Tuple.__eq__  ndindex/tuple.py:107
   |        `- 0.932 Tuple.__init__  ndindex/ndindex.py:159
   |           `- 0.860 Tuple._typecheck  ndindex/tuple.py:49
   |- 16.648 _GeneratorContextManager.__exit__  contextlib.py:141
   |  `- 16.648 VersionedHDF5File.stage_version  versioned_hdf5/api.py:267
   |     `- 16.633 commit_version  versioned_hdf5/versions.py:71
   |        |- 10.969 create_virtual_dataset  versioned_hdf5/backend.py:429
   |        |  |- 1.989 [self]  versioned_hdf5/backend.py
   |        |  |- 1.897 Group.create_virtual_dataset  h5py/_hl/group.py:188
   |        |  |  `- 1.467 VirtualLayout.make_dataset  h5py/_hl/vds.py:228
   |        |  |- 1.794 select  h5py/_hl/selections.py:19
   |        |  |  `- 1.361 [self]  h5py/_hl/selections.py
   |        |  |- 1.289 <dictcomp>  versioned_hdf5/backend.py:437
   |        |  |  `- 0.964 Slice.reduce  ndindex/slice.py:212
   |        |  |     `- 0.816 Slice.__init__  ndindex/ndindex.py:159
   |        |  |        `- 0.762 Slice._typecheck  ndindex/slice.py:62
   |        |  |- 1.111 Tuple.__init__  ndindex/ndindex.py:159
   |        |  |  `- 1.017 Tuple._typecheck  ndindex/tuple.py:49
   |        |  `- 0.846 Dataset.shape  h5py/_hl/dataset.py:467
   |        |     `- 0.799 [self]  h5py/_hl/dataset.py
   |        `- 4.984 write_dataset_chunks  versioned_hdf5/backend.py:346
   |           |- 2.393 Hashtable.hash  versioned_hdf5/hashtable.py:116
   |           |  `- 1.827 openssl_sha256  <built-in>
   |           `- 0.953 Hashtable.__contains__  <frozen _collections_abc>:778
   |              `- 0.903 Hashtable.__getitem__  versioned_hdf5/hashtable.py:205
   |                 `- 0.673 Slice.__init__  ndindex/ndindex.py:159
   |- 13.589 InMemoryDataset.resize  versioned_hdf5/wrappers.py:609
   |  |- 7.899 InMemoryDatasetID.data_dict  versioned_hdf5/wrappers.py:1298
   |  |  |- 4.633 <dictcomp>  versioned_hdf5/wrappers.py:1327
   |  |  |  `- 4.225 spaceid_to_slice  versioned_hdf5/slicetools.py:6
   |  |  |     |- 1.730 Tuple.__init__  ndindex/ndindex.py:159
   |  |  |     |  `- 1.607 Tuple._typecheck  ndindex/tuple.py:49
   |  |  |     |     `- 0.932 [self]  ndindex/tuple.py
   |  |  |     |- 1.228 [self]  versioned_hdf5/slicetools.py
   |  |  |     `- 1.193 hyperslab_to_slice  versioned_hdf5/slicetools.py:33
   |  |  |        `- 1.112 Slice.__init__  ndindex/ndindex.py:159
   |  |  |           `- 0.682 Slice._typecheck  ndindex/slice.py:62
   |  |  `- 1.464 <listcomp>  versioned_hdf5/wrappers.py:1317
   |  |     `- 1.405 [self]  versioned_hdf5/wrappers.py
   |  |- 2.367 ChunkSize.as_subchunks  ndindex/chunking.py:143
   |  |  `- 2.327 _indices  ndindex/chunking.py:288
   |  |     |- 1.155 <listcomp>  ndindex/chunking.py:292
   |  |     |  `- 0.884 Slice.__init__  ndindex/ndindex.py:159
   |  |     |     `- 0.710 Slice._typecheck  ndindex/slice.py:62
   |  |     `- 0.856 Tuple.__init__  ndindex/ndindex.py:159
   |  |        `- 0.788 Tuple._typecheck  ndindex/tuple.py:49
   |  |- 1.429 InMemoryDataset.get_index  versioned_hdf5/wrappers.py:655
   |  |  `- 1.429 InMemoryDataset.__getitem__  h5py/_hl/dataset.py:749
   |  |     `- 1.422 Reader.read  <built-in>
   |  `- 0.899 [self]  versioned_hdf5/wrappers.py
   |- 3.085 _GeneratorContextManager.__enter__  contextlib.py:132
   |  `- 3.085 VersionedHDF5File.stage_version  versioned_hdf5/api.py:267
   |     `- 3.085 create_version_group  versioned_hdf5/versions.py:22
   |        `- 3.084 Group.visititems  h5py/_hl/group.py:635
   |           `- 2.533 proxy  h5py/_hl/group.py:660
   |              |- 1.618 _get  versioned_hdf5/versions.py:57
   |              |  `- 1.225 InMemoryGroup.__setitem__  versioned_hdf5/wrappers.py:110
   |              |     `- 1.225 InMemoryGroup._add_to_data  versioned_hdf5/wrappers.py:114
   |              |        `- 1.225 InMemoryDataset.__init__  versioned_hdf5/wrappers.py:503
   |              `- 0.767 Group.__getitem__  h5py/_hl/group.py:348
   `- 0.958 File.__exit__  h5py/_hl/files.py:601
      `- 0.958 File.close  h5py/_hl/files.py:576