feat: allow externally generated shuffle buffer and multiple of them #1801
Conversation
chebbyChefNEQ
force-pushed
the
rmeng/shuffle-only-index
branch
from
43a9fdf
to
3bea032
Compare
| /// # Safety | ||
| /// | ||
| /// user must ensure the buffers are valid. | ||
| pub unsafe fn set_unsorted_buffers(&mut self, unsorted_buffers: Vec<String>) { |
There was a problem hiding this comment.
Why is it a unsafe method?
Btw, Can we just hold a reference to &[String]?
There was a problem hiding this comment.
This is meant to note to user that if invalid buffers are passed here it can cause undefined behavior
There was a problem hiding this comment.
changed it to &[impl Into<String>]
| .expect("part id should exist"); | ||
|
|
||
| let mut stream = stream::iter(start..end) | ||
| .map(|i| reader.read_batch(i as i32, ReadBatchParams::RangeFull, &lance_schema)) |
There was a problem hiding this comment.
you can just do reader.read_batch(i as i32, .., &lance_schema) i think
| .map(|i| reader.read_batch(i as i32, ReadBatchParams::RangeFull, &lance_schema)) | ||
| .buffered(16); | ||
|
|
||
| while let Some(batch) = stream.next().await { |
There was a problem hiding this comment.
we need to update partition_sizes Fn closures can't do that easily because Stream::for_each captures multiple FnMut
| .expect("part id should exist"); | ||
|
|
||
| let mut stream = stream::iter(start..end) | ||
| .map(|i| reader.read_batch(i as i32, ReadBatchParams::RangeFull, &lance_schema)) |
There was a problem hiding this comment.
Also might consider putting read_batch into a tokio task. Otherwise progress on these IO calls will be blocked on the CPU-bound work happenining in your while let loop below.
| pq_codes.values()[i * num_sub_vectors..(i + 1) * num_sub_vectors] | ||
| .iter(), |
There was a problem hiding this comment.
Could be wrong, but I think you can make this a slice and it can do a contiguous mem copy instead of running through the iterator.
| pq_codes.values()[i * num_sub_vectors..(i + 1) * num_sub_vectors] | |
| .iter(), | |
| &pq_codes.values()[i * num_sub_vectors..(i + 1) * num_sub_vectors], |
|
Btw, delete |
chebbyChefNEQ
force-pushed
the
rmeng/shuffle-only-index
branch
2 times, most recently
from
70ebbe0
to
87af1c7
Compare
chebbyChefNEQ
force-pushed
the
rmeng/shuffle-only-index
branch
from
87af1c7
to
f3b7288
Compare
This PR does a few things to the IVF shuffler
set_unsorted_buffersmethod. This method allowes caller to set a list of buffers that is different from theunsorted.lancebuffer used by default.I will wire up python in the next PR to keep this smaller and more reviewable.