Use dask array native reduction

[fjetter]

This came up in dask/dask#11844 (comment)

I don't see a reason why you shouldn't be able to use the dask native da.reduction. This way, dask owns the tree reduction logic and it will work as intended once we move to an expression backend.

The only caveat with this is that the very final reduction / Session merge would happen locally. This should typically only be a very small number of sessions (one per dask.array). The only reason why this would be bad is if that merge would be expensive or would require IO but I believe this is just merging metadata, isn't it?

cc @dcherian

[dcherian]

The only reason why this would be bad is if that merge would be expensive or would require IO but I believe this is just merging metadata, isn't it?

Yes, we were being a bit extra around being scale-aware here. I figured if I wanted tree-reduce across arrays, I'd need to write out the graph, so I just did it for the arrays themselves too. I am surprised this just works, are these Session objects being automatically cast to object arrays of the right shape under the hood in _concatenate2? They must be. I had originally thought that I'd need to write custom combine and aggregate ...

[fjetter]

There is no _concatenate2 involved since we're disabling concatenate in the reduction. This way the chunks are just passed to merge_sessions as a nested list, e.g.

and there is no final concat step because the final reduction is also a merge_sessions which returns a singleton.

Sorry for the typing mess and thanks for cleaning up

[fjetter]

sorry for the linting problems. I didn't setup pre-commit because I struggled with the uv setup, etc. I think it should be working

[fjetter]

ok, rust ci is now failing with... is this related?

advisories FAILED, bans ok, licenses ok, sources ok
error: Recipe `check-deps` failed on line 34 with exit code 1
error: Recipe `pre-commit` failed on line 47 with exit code 1

I'll try to reproduce. Main doesn't seem to eb affected

[dcherian]

@fjetter the PR is fine, this is an unrelated error.

[dcherian]

Thanks @fjetter!