perf: internal integer indexing

[selmanozleyen]

In this PR we use row indices for the fetching as much as possible instead of slices.

In the future if we bring zarrs native integer indexing the transition will be easier. Also according to @ilan-gold for bigger chunk sizes it has no performance impact and it gives a %20 boost for integer indexing.

[ilan-gold]

Is use_original_space used anywhere anymore? I am fine getting rid its use (in generating the indices), but then lets remove the arg because it's dead code

[selmanozleyen]

true, removed it

[ilan-gold]

Let's remove that arg!

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.33%. Comparing base (254ad41) to head (d6155ba).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #209      +/-   ##
==========================================
+ Coverage   93.18%   93.33%   +0.15%     
==========================================
  Files          14       14              
  Lines        1130     1111      -19     
==========================================
- Hits         1053     1037      -16     
+ Misses         77       74       -3     

Files with missing lines	Coverage Δ
src/annbatch/loader.py	93.70% <100.00%> (+0.56%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[selmanozleyen]

A lookup of the list-placement index of a dataset to the sorted row indices to fetch.

@ilan-gold this isn't correct right? Is it sorted row indices?

[ilan-gold]

Don't think anything is sorted. Let me check

[ilan-gold]

YEah, I mean you can double check, but I don't think we have any sorting enforcement. That's why constructing the batches to yield is completely the sampler's responsibility and relies on the sampler knowing its own fetch order

[selmanozleyen]

I can't access the cluster right now but my worry was creating integers from slices to then creating slices again from those integers. I want to just point out when the number of slices to loop over is small (which I thought was the goal for having slices) main should win right?

So I ran this local isolated benchmark: https://gist.github.com/selmanozleyen/199b73288ce23cd3928e0b5f00425924

The cutoff chunksize for main to win is like below:

preload_nchunks  cutoff
32               32-64
64               64-128
128              64-128
256              64-128
512              64-128
32768            64-128

When the number of datasets is 5, I found these cutoffs to be higher around 64-128. So the main should win when preload_nchunks is low and chunksize is high. Whenever the dataset number is higher main win gets harder.

But I'd consider this a branch win in general.

[ilan-gold]

I want to just point out when the number of slices to loop over is small (which I thought was the goal for having slices) main should win

Maybe, but running locally, it seems that the amount of time it takes on small numbers of slices once you get bigger chunks is 1/5th-1/10th regardless of the method. So it isn't really a bottleneck because the absolute time is much smaller regardless. For example, yes, the below case is 3x slower on this branch

❯ python tester.py --n-rows=100_000_000 --n-datasets=1 --preload-nchunks=32 --chunk-sizes=512 
n_rows=100000000 n_datasets=1 preload_nchunks=32 mean_nnz=20.0
chunk_size      rows/request    main_us branch_us       branch/main
512     16384   80.7    264.7   3.28

but compared to the amount of time the chunk_size=1 case takes in absolute numbers

❯ python tester.py --n-rows=20_000_000 --n-datasets=1 --preload-nchunks=4096 --chunk-sizes=1 
n_rows=20000000 n_datasets=1 preload_nchunks=4096 mean_nnz=20.0
chunk_size      rows/request    main_us branch_us       branch/main
1       4096    8516.1  5411.5  0.64

it's irrelevant. Also, your benchmark with row_runs bit can happen in somewhat in parallel to fetching because our _async_array._get_selection releases the GIL, but not sure how much impact that has.

Are you seeing this as well?

[selmanozleyen]

Thats very convincing! Thanks! And its kind of mind blowing