pyoctree | Python Octree Implementation | GPU library

I am building a distributed spatial index in Python with a wrapped C++ extension. I am trying to use Dask (for the first time) instead of MPI. The current implementation works, but has several problems.

1. The critical part is the large overhead compared to serial execution. I would expect an almost linear speed-up from the algorithm.

2. Persist does not seem to do what I expect, as timings of the first query are much longer than in the second.

3. for my eyes the code looks unidiomatic, but as said, I am new to Dask. Is there a better way of doing this?

4. I need to to do some indexing into the delayed objects when using 2d chunks. That does not occur with 1d chunks and feels weird.

The algorithm builds sorts list of particles and then builds an octree, where nodes reference contiguous blocks of particles by in the sorted array by a pair of (start, end) indices. Queries take a boundary box and search overlapping nodes in the octree collect particles actually in the boundary box from the resulting candidates. Both building and queries are purely serial.

Parallelisation is done by randomly splitting up the particles, building tree on each subset. Queries are broadcast across all sub-indices and results are concatenated back together. Dask felt like a natural fit here. I use persist to generate the index once and keep it around as I expect many queries per index.

I tried map_blocks but this only seems to work for array to array transformations. Further, various permutations of persist/compute have been tried.