From 9ac76329b94c86ad73e1d9343ffa51ef09dafd55 Mon Sep 17 00:00:00 2001 From: Nils Homer Date: Thu, 2 Apr 2026 14:33:13 -0700 Subject: [PATCH] perf(sort): reuse buffers in merge phase to reduce allocations Change GenericKeyedChunkReader::next_record and all ChunkSource wrappers to accept a &mut Vec buffer parameter instead of returning an owned Vec. This lets merge loops write directly from each heap entry's buffer and refill it in-place, avoiding per-record Vec allocations and the intermediate output_buffer that previously collected owned Vecs before flushing. --- src/lib/sort/raw.rs | 154 ++++++++++++++++++++------------------------ 1 file changed, 69 insertions(+), 85 deletions(-) diff --git a/src/lib/sort/raw.rs b/src/lib/sort/raw.rs index c097fd032..58329fca8 100644 --- a/src/lib/sort/raw.rs +++ b/src/lib/sort/raw.rs @@ -289,6 +289,9 @@ impl GenericKeyedChunkWriter { /// Auto-detects BGZF compression via magic bytes. pub struct GenericKeyedChunkReader { receiver: Receiver)>>, + /// Return channel for empty buffers — the consumer sends its old buffer + /// back so the producer can reuse the allocation instead of allocating. + buf_return: Sender>, _handle: JoinHandle<()>, } @@ -305,6 +308,7 @@ impl GenericKeyedChunkReader { /// Returns an error if the file cannot be opened. pub fn open(path: &Path, concurrency_limit: Option>) -> Result { let (tx, rx) = bounded(MERGE_PREFETCH_SIZE); + let (buf_tx, buf_rx) = bounded::>(MERGE_PREFETCH_SIZE); let path = path.to_path_buf(); let handle = thread::spawn(move || { @@ -336,13 +340,13 @@ impl GenericKeyedChunkReader { // Read using appropriate decoder if is_compressed { let bgzf_reader = BgzfReader::new(buf_reader); - Self::read_records(bgzf_reader, tx, concurrency_limit); + Self::read_records(bgzf_reader, tx, buf_rx, concurrency_limit); } else { - Self::read_records(buf_reader, tx, concurrency_limit); + Self::read_records(buf_reader, tx, buf_rx, concurrency_limit); } }); - Ok(Self { receiver: rx, _handle: handle }) + Ok(Self { receiver: rx, buf_return: buf_tx, _handle: handle }) } /// Read records from a reader and send them through the channel. @@ -355,6 +359,7 @@ impl GenericKeyedChunkReader { fn read_records( mut reader: R, tx: crossbeam_channel::Sender)>>, + buf_pool: crossbeam_channel::Receiver>, semaphore: Option>, ) { const BATCH_SIZE: usize = 64; @@ -391,7 +396,10 @@ impl GenericKeyedChunkReader { } let len = u32::from_le_bytes(len_buf) as usize; - let mut record = vec![0u8; len]; + // Try to reuse a buffer from the pool; allocate only if empty. + let mut record = buf_pool.try_recv().unwrap_or_default(); + record.clear(); + record.resize(len, 0); if reader.read_exact(&mut record).is_err() { log::error!("Error reading keyed chunk record"); error = true; @@ -420,10 +428,21 @@ impl GenericKeyedChunkReader { } } - /// Read the next keyed record from the prefetch buffer. - pub fn next_record(&mut self) -> Option<(K, Vec)> { + /// Read the next keyed record from the prefetch buffer into `buf`. + /// + /// On success the record bytes are swapped into `buf` and the sort key is + /// returned. The old contents of `buf` are returned to the producer thread + /// for reuse, avoiding per-record allocation on the disk path. + /// Returns `None` at EOF. + pub fn next_record(&mut self, buf: &mut Vec) -> Option { match self.receiver.recv() { - Ok(Some(entry)) => Some(entry), + Ok(Some((key, mut data))) => { + std::mem::swap(buf, &mut data); + // Return the old buffer to the producer for reuse. + // Ignore errors — producer may have finished. + let _ = self.buf_return.try_send(data); + Some(key) + } Ok(None) | Err(_) => None, } } @@ -639,7 +658,8 @@ impl RawExternalSorter { // Initialize heap with first record from each reader let mut heap: Vec> = Vec::with_capacity(readers.len()); for (reader_idx, reader) in readers.iter_mut().enumerate() { - if let Some((key, record)) = reader.next_record() { + let mut record = Vec::new(); + if let Some(key) = reader.next_record(&mut record) { heap.push(HeapEntry { key, record, reader_idx }); } } @@ -666,14 +686,11 @@ impl RawExternalSorter { // Merge loop while heap_size > 0 { let reader_idx = heap[0].reader_idx; - let key = std::mem::take(&mut heap[0].key); - let record = std::mem::take(&mut heap[0].record); - writer.write_record(&key, &record)?; + writer.write_record(&heap[0].key, &heap[0].record)?; - if let Some((next_key, next_record)) = readers[reader_idx].next_record() { + if let Some(next_key) = readers[reader_idx].next_record(&mut heap[0].record) { heap[0].key = next_key; - heap[0].record = next_record; heap_sift_down(&mut heap, 0, heap_size, <); } else { heap_size -= 1; @@ -1466,16 +1483,17 @@ impl RawExternalSorter { } impl KeyedChunkSource { - fn next_record(&mut self) -> Option<(TemplateKey, Vec)> { + /// Fill `buf` with the next record's bytes and return the sort key, + /// or `None` at EOF. + fn next_record(&mut self, buf: &mut Vec) -> Option { match self { - KeyedChunkSource::Disk(reader) => reader.next_record(), + KeyedChunkSource::Disk(reader) => reader.next_record(buf), KeyedChunkSource::Memory { records, idx } => { if *idx < records.len() { - // Use replace with zeroed key and empty vec to avoid Default requirement - let dummy = (TemplateKey::zeroed(), Vec::new()); - let entry = std::mem::replace(&mut records[*idx], dummy); + let (key, ref mut data) = records[*idx]; + std::mem::swap(buf, data); *idx += 1; - Some(entry) + Some(key) } else { None } @@ -1491,9 +1509,6 @@ impl RawExternalSorter { chunk_idx: usize, } - // Output buffer to reduce write syscalls - stores raw bytes directly - const OUTPUT_BUFFER_SIZE: usize = 2048; - let num_disk = chunk_files.len(); let num_memory = memory_chunks.iter().filter(|c| !c.is_empty()).count(); info!("Merging from {num_disk} files and {num_memory} in-memory blocks..."); @@ -1525,7 +1540,8 @@ impl RawExternalSorter { // Initialize heap with first record from each chunk let mut heap: Vec = Vec::with_capacity(sources.len()); for (chunk_idx, source) in sources.iter_mut().enumerate() { - if let Some((key, record)) = source.next_record() { + let mut record = Vec::new(); + if let Some(key) = source.next_record(&mut record) { heap.push(KeyedHeapEntry { key, record, chunk_idx }); } } @@ -1555,47 +1571,29 @@ impl RawExternalSorter { )?; let mut records_merged = 0u64; - let mut output_buffer: Vec> = Vec::with_capacity(OUTPUT_BUFFER_SIZE); // Merge loop using fixed-array heap with O(1) comparisons while heap_size > 0 { - // Get the minimum record (at heap[0]) let chunk_idx = heap[0].chunk_idx; - let record_bytes = std::mem::take(&mut heap[0].record); - // Buffer the raw bytes directly (no Record conversion) - output_buffer.push(record_bytes); + // Write directly from the heap entry's buffer (no intermediate Vec) + writer.write_raw_record(&heap[0].record)?; records_merged += 1; - // Flush buffer when full - if output_buffer.len() >= OUTPUT_BUFFER_SIZE { - for rec in output_buffer.drain(..) { - writer.write_raw_record(&rec)?; - } - } - - // Get next record from the same chunk - if let Some((key, next_record)) = sources[chunk_idx].next_record() { - // Replace top with new record and sift down + // Refill the heap entry's buffer from the same chunk source + if let Some(key) = sources[chunk_idx].next_record(&mut heap[0].record) { heap[0].key = key; - heap[0].record = next_record; heap_sift_down(&mut heap, 0, heap_size, <); } else { // Chunk exhausted - remove from heap heap_size -= 1; if heap_size > 0 { - // Move last element to top and sift down heap.swap(0, heap_size); heap_sift_down(&mut heap, 0, heap_size, <); } } } - // Flush remaining buffered records - for rec in output_buffer { - writer.write_raw_record(&rec)?; - } - writer.finish()?; info!("Merge complete: {records_merged} records merged"); @@ -1623,15 +1621,18 @@ impl RawExternalSorter { } impl GenericKeyedChunkSource { - fn next_record(&mut self) -> Option<(K, Vec)> { + /// Fill `buf` with the next record's bytes and return the sort key, + /// or `None` at EOF. + fn next_record(&mut self, buf: &mut Vec) -> Option { match self { - GenericKeyedChunkSource::Disk(reader) => reader.next_record(), + GenericKeyedChunkSource::Disk(reader) => reader.next_record(buf), GenericKeyedChunkSource::Memory { records, idx } => { if *idx < records.len() { - // Take ownership of the record - let entry = std::mem::take(&mut records[*idx]); + let (ref mut key, ref mut data) = records[*idx]; + std::mem::swap(buf, data); + let key = std::mem::take(key); *idx += 1; - Some(entry) + Some(key) } else { None } @@ -1647,9 +1648,6 @@ impl RawExternalSorter { chunk_idx: usize, } - // Output buffer to reduce write syscalls - const OUTPUT_BUFFER_SIZE: usize = 2048; - let num_disk = chunk_files.len(); let num_memory = memory_chunks.iter().filter(|c| !c.is_empty()).count(); info!("Merging from {num_disk} files and {num_memory} in-memory blocks..."); @@ -1682,7 +1680,8 @@ impl RawExternalSorter { // Initialize heap with first record from each chunk let mut heap: Vec> = Vec::with_capacity(sources.len()); for (chunk_idx, source) in sources.iter_mut().enumerate() { - if let Some((key, record)) = source.next_record() { + let mut record = Vec::new(); + if let Some(key) = source.next_record(&mut record) { heap.push(GenericKeyedHeapEntry { key, record, chunk_idx }); } } @@ -1711,47 +1710,29 @@ impl RawExternalSorter { )?; let mut records_merged = 0u64; - let mut output_buffer: Vec> = Vec::with_capacity(OUTPUT_BUFFER_SIZE); // Merge loop using fixed-array heap with key comparisons while heap_size > 0 { - // Get the minimum record (at heap[0]) let chunk_idx = heap[0].chunk_idx; - let record_bytes = std::mem::take(&mut heap[0].record); - // Buffer the raw bytes directly (no Record conversion) - output_buffer.push(record_bytes); + // Write directly from the heap entry's buffer (no intermediate Vec) + writer.write_raw_record(&heap[0].record)?; records_merged += 1; - // Flush buffer when full - if output_buffer.len() >= OUTPUT_BUFFER_SIZE { - for rec in output_buffer.drain(..) { - writer.write_raw_record(&rec)?; - } - } - - // Get next record from the same chunk - if let Some((key, next_record)) = sources[chunk_idx].next_record() { - // Replace top with new record and sift down + // Refill the heap entry's buffer from the same chunk source + if let Some(key) = sources[chunk_idx].next_record(&mut heap[0].record) { heap[0].key = key; - heap[0].record = next_record; heap_sift_down(&mut heap, 0, heap_size, <); } else { // Chunk exhausted - remove from heap heap_size -= 1; if heap_size > 0 { - // Move last element to top and sift down heap.swap(0, heap_size); heap_sift_down(&mut heap, 0, heap_size, <); } } } - // Flush remaining buffered records - for rec in output_buffer { - writer.write_raw_record(&rec)?; - } - writer.finish()?; info!("Merge complete: {records_merged} records merged"); @@ -1778,14 +1759,18 @@ impl RawExternalSorter { } impl KeyedSource { - fn next_record(&mut self) -> Option<(K, Vec)> { + /// Fill `buf` with the next record's bytes and return the sort key, + /// or `None` at EOF. + fn next_record(&mut self, buf: &mut Vec) -> Option { match self { - KeyedSource::Disk(reader) => reader.next_record(), + KeyedSource::Disk(reader) => reader.next_record(buf), KeyedSource::Memory { records, idx } => { if *idx < records.len() { - let entry = std::mem::take(&mut records[*idx]); + let (ref mut key, ref mut data) = records[*idx]; + std::mem::swap(buf, data); + let key = std::mem::take(key); *idx += 1; - Some(entry) + Some(key) } else { None } @@ -1828,7 +1813,8 @@ impl RawExternalSorter { // Initialize heap let mut heap: Vec> = Vec::with_capacity(sources.len()); for (chunk_idx, source) in sources.iter_mut().enumerate() { - if let Some((key, record)) = source.next_record() { + let mut record = Vec::new(); + if let Some(key) = source.next_record(&mut record) { heap.push(HeapEntry { key, record, chunk_idx }); } } @@ -1866,14 +1852,12 @@ impl RawExternalSorter { // for each record, so we write directly while heap_size > 0 { let chunk_idx = heap[0].chunk_idx; - let record_bytes = std::mem::take(&mut heap[0].record); - writer.write_raw_record(&record_bytes)?; + writer.write_raw_record(&heap[0].record)?; records_merged += 1; - if let Some((key, next_record)) = sources[chunk_idx].next_record() { + if let Some(key) = sources[chunk_idx].next_record(&mut heap[0].record) { heap[0].key = key; - heap[0].record = next_record; heap_sift_down(&mut heap, 0, heap_size, <); } else { heap_size -= 1;