vindex as outer indexer: memory and time performance

[ilan-gold]

Describe the issue:

Emulating outerindexing via vindex + np.ix_ appears to be much slower and more memory intensive (prohibitively so for very large arrays where for an 1,000,000x1,000,000 array, it tried allocating 1.5TB of memory) than twice indexing. I know this is basically stated in the docs, but maybe there is something to be done here? If not, feel free to close.

Minimal Complete Verifiable Example:

%load_ext memory_profiler

import dask.array as da
import numpy as np
import scipy as sp

chunksize = 100
size = 10_000
n_points = 5000

X = da.random.poisson(15, (size, size), chunks = (chunksize, chunksize))

index_0 = np.random.randint(0, X.shape[0], n_points)
index_0.sort()
index_1 = np.random.randint(0, X.shape[1], n_points)
index_1.sort()

print('vindex timing:')
%timeit X.vindex[np.ix_(index_0, index_1)].compute()
print('vindex memory usage:')
%memit X.vindex[np.ix_(index_0, index_1)]
print('double-index timing:')
%timeit X[index_0, :][:, index_1].compute()
print('double-index memory usage:')
%memit X[index_0, :][:, index_1]

Anything else we need to know?:

Environment:

• Dask version: 2024.3.1
• Python version: 3.12
• Operating System: mac
• Install method (conda, pip, source): pip