Skip to content

Commit 1385811

Browse files
committed
Micro optimize dataset.isel for speed on large datasets
This targets optimization for datasets with many "scalar" variables (that is variables without any dimensions). This can happen in the context where you have many pieces of small metadata that relate to various facts about an experimental condition. For example, we have about 80 of these in our datasets (and I want to incrase this number) Our datasets are quite large (On the order of 1TB uncompresed) so we often have one dimension that is in the 10's of thousands. However, it has become quite slow to index in the dataset. We therefore often "carefully slice out the matadata we need" prior to doing anything with our dataset, but that isn't quite possible with you want to orchestrate things with a parent application. These optimizations are likely "minor" but considering the results of the benchmark, I think they are quite worthwhile: * main (as of pydata#9001) - 2.5k its/s * With pydata#9002 - 4.2k its/s * With this Pull Request (on top of pydata#9002) -- 6.1k its/s Thanks for considering.
1 parent b1f3fea commit 1385811

File tree

1 file changed

+14
-4
lines changed

1 file changed

+14
-4
lines changed

xarray/core/dataset.py

+14-4
Original file line numberDiff line numberDiff line change
@@ -2981,20 +2981,30 @@ def isel(
29812981
coord_names = self._coord_names.copy()
29822982

29832983
indexes, index_variables = isel_indexes(self.xindexes, indexers)
2984+
all_keys = set(indexers.keys())
29842985

29852986
for name, var in self._variables.items():
29862987
# preserve variable order
29872988
if name in index_variables:
29882989
var = index_variables[name]
2989-
else:
2990-
var_indexers = {k: v for k, v in indexers.items() if k in var.dims}
2991-
if var_indexers:
2990+
dims.update(zip(var.dims, var.shape))
2991+
# Fastpath, skip all of this for variables with no dimensions
2992+
# Keep the result cached for future dictionary update
2993+
elif var_dims := var.dims:
2994+
# Large datasets with alot of metadata may have many scalars
2995+
# without any relevant dimensions for slicing.
2996+
# Pick those out quickly and avoid paying the cost below
2997+
# of resolving the var_indexers variables
2998+
if var_indexer_keys := all_keys.intersection(var_dims):
2999+
var_indexers = {k: indexers[k] for k in var_indexer_keys}
29923000
var = var.isel(var_indexers)
29933001
if drop and var.ndim == 0 and name in coord_names:
29943002
coord_names.remove(name)
29953003
continue
3004+
# Update our reference to `var_dims` after the call to isel
3005+
var_dims = var.dims
3006+
dims.update(zip(var_dims, var.shape))
29963007
variables[name] = var
2997-
dims.update(zip(var.dims, var.shape))
29983008

29993009
return self._construct_direct(
30003010
variables=variables,

0 commit comments

Comments
 (0)