perf!: remove shuffle buffer#5912
Conversation
This removes a buffer in the shuffler that accumulated batches for batched writes to temporary storage. This was configured with a public buffer_size parameter hence the breaking change. Previously, when we shuffled data we accumulated this many batches for each partition in memory and then flushed them all to disk at once. This may have been intended as an optimization in the original implementation of the shuffler, which supported external shuffling through arbitrary object storage. However, the shuffler was subsequently hardcoded to use local disk (where this kind of buffering serves no benefit) and even on remote object storage, we already have a layer of buffering in the storage writer. Instead of buffering batches, just write them directly to the FileWriter. This results in much more predictable memory usage and also faster index builds.
Here is the result on an index build over 100M 384d vectors. Note that this chart still shows a slow leak over the course of the shuffle. I assume that's unrelated and haven't looked at it yet. |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
westonpace
left a comment
westonpace
left a comment
There was a problem hiding this comment.
Nice work! The writers will already do their own buffering if they need to so I agree this extra layer of buffering is not needed.
| if !batches.is_empty() { | ||
| partition_sizes[part_id] += batches.iter().map(|b| b.num_rows()).sum::<usize>(); | ||
| futs.push(writer.write_batches(batches.iter())); | ||
| writers[part_id].write_batches(batches.iter()).await?; |
There was a problem hiding this comment.
We can do this in a follow-up but it might be nice to still do all the writes in parallel. E.g. keep the futs Vec. shuffled is a Vec and not any kind of stream / iterator so the data is all in memory already (I think the important point is getting rid of the if counter % self.buffer_size == 0)
let mut futs = vec![];
if !batches.is_empty() {
futs.push(writers[part_id].write_batches(batches.iter()));
}
try_join_all(futs).await?;
3 participants
This removes a buffer in the shuffler that accumulated batches for batched writes to temporary storage. This was configured with a public buffer_size parameter hence the breaking change.
Previously, when we shuffled data we accumulated this many batches for each partition in memory and then flushed them all to disk at once. This may have been intended as an optimization in the original implementation of the shuffler, which supported external shuffling through arbitrary object storage. However, the shuffler was subsequently hardcoded to use local disk (where this kind of buffering serves no benefit) and even on remote object storage, we already have a layer of buffering in the storage writer.
Instead of buffering batches, just write them directly to the FileWriter. This results in much more predictable memory usage and also faster index builds.