diff --git a/core/src/ivfflat_io.rs b/core/src/ivfflat_io.rs index 1409caa..25d24d7 100644 --- a/core/src/ivfflat_io.rs +++ b/core/src/ivfflat_io.rs @@ -38,14 +38,30 @@ pub const IVFFLAT_HEADER_SIZE: usize = 64; const FLAG_DELTA_IDS: u32 = 1 << 0; const REQUIRED_FLAGS: u32 = FLAG_DELTA_IDS; const SUPPORTED_FLAGS: u32 = REQUIRED_FLAGS; +const IVFFLAT_WRITE_BUFFER_SIZE: usize = 64 * 1024 * 1024; // Raw-vector scan cost scales with both rows and dimension. Below this amount, // Rayon scheduling and list-local heap merging outweigh the saved CPU time. const PARALLEL_FLAT_SCAN_MIN_COMPONENTS: usize = 1024 * 1024; pub fn write_ivfflat_index(index: &IVFFlatIndex, out: &mut dyn SeekWrite) -> io::Result<()> { + write_ivfflat_index_with_buffer_limit(index, out, IVFFLAT_WRITE_BUFFER_SIZE) +} + +fn write_ivfflat_index_with_buffer_limit( + index: &IVFFlatIndex, + out: &mut dyn SeekWrite, + buffer_limit: usize, +) -> io::Result<()> { let d = index.d; let nlist = index.nlist; validate_index_shape(index)?; + let bytes_per_vector = d.checked_mul(size_of::()).ok_or_else(|| { + io::Error::new( + io::ErrorKind::InvalidInput, + "IVF-FLAT bytes per vector overflow", + ) + })?; + let mut write_buffer = Vec::new(); let d_i32 = usize_to_i32(d, "dimension")?; let nlist_i32 = usize_to_i32(nlist, "nlist")?; let total_vectors = index.ids.iter().try_fold(0i64, |sum, ids| { @@ -85,7 +101,12 @@ pub fn write_ivfflat_index(index: &IVFFlatIndex, out: &mut dyn SeekWrite) -> io: write_u32_le(out, FLAG_DELTA_IDS)?; out.write_all(&[0u8; 32])?; - write_f32_slice(out, &index.quantizer_centroids)?; + write_f32_slice( + out, + &index.quantizer_centroids, + &mut write_buffer, + buffer_limit, + )?; let offset_table_size = nlist.checked_mul(16).ok_or_else(|| { io::Error::new( @@ -114,15 +135,7 @@ pub fn write_ivfflat_index(index: &IVFFlatIndex, out: &mut dyn SeekWrite) -> io: if count > 0 { let id_bytes_len = sorted_lists[list_id].1.len(); list_id_bytes_lens[list_id] = usize_to_i32(id_bytes_len, "delta ID section")?; - let vector_bytes = checked_list_bytes( - count, - d.checked_mul(4).ok_or_else(|| { - io::Error::new( - io::ErrorKind::InvalidInput, - "IVF-FLAT bytes per vector overflow", - ) - })?, - )?; + let vector_bytes = checked_list_bytes(count, bytes_per_vector)?; let list_bytes = 12usize .checked_add(id_bytes_len) .and_then(|len| len.checked_add(vector_bytes)) @@ -150,11 +163,15 @@ pub fn write_ivfflat_index(index: &IVFFlatIndex, out: &mut dyn SeekWrite) -> io: write_i64_le(out, index.ids[list_id][order[0]])?; write_i32_le(out, id_bytes.len() as i32)?; out.write_all(&id_bytes)?; - let mut sorted_vectors = Vec::with_capacity(order.len() * d); - for idx in order { - sorted_vectors.extend_from_slice(&index.vectors[list_id][idx * d..(idx + 1) * d]); - } - write_f32_slice(out, &sorted_vectors)?; + let vectors = &index.vectors[list_id]; + write_f32_iter( + out, + order + .into_iter() + .flat_map(|idx| vectors[idx * d..(idx + 1) * d].iter()), + &mut write_buffer, + buffer_limit, + )?; } Ok(()) @@ -1047,9 +1064,41 @@ fn write_i64_le(out: &mut dyn SeekWrite, v: i64) -> io::Result<()> { out.write_all(&v.to_le_bytes()) } -fn write_f32_slice(out: &mut dyn SeekWrite, data: &[f32]) -> io::Result<()> { - let bytes: Vec = data.iter().flat_map(|f| f.to_le_bytes()).collect(); - out.write_all(&bytes) +fn write_f32_slice( + out: &mut dyn SeekWrite, + data: &[f32], + buffer: &mut Vec, + buffer_limit: usize, +) -> io::Result<()> { + write_f32_iter(out, data.iter(), buffer, buffer_limit) +} + +fn write_f32_iter<'a>( + out: &mut dyn SeekWrite, + data: impl Iterator, + buffer: &mut Vec, + buffer_limit: usize, +) -> io::Result<()> { + let buffer_limit = buffer_limit.max(1); + buffer.clear(); + for value in data { + let bytes = value.to_le_bytes(); + let mut offset = 0; + while offset < bytes.len() { + let len = (buffer_limit - buffer.len()).min(bytes.len() - offset); + buffer.extend_from_slice(&bytes[offset..offset + len]); + offset += len; + if buffer.len() == buffer_limit { + out.write_all(buffer)?; + buffer.clear(); + } + } + } + if !buffer.is_empty() { + out.write_all(buffer)?; + buffer.clear(); + } + Ok(()) } fn validate_positive_i32(val: i32, field: &str) -> io::Result { @@ -1326,6 +1375,68 @@ mod tests { bytes } + struct MaxWriteWriter { + bytes: Vec, + max_write: usize, + } + + impl SeekWrite for MaxWriteWriter { + fn write_all(&mut self, buf: &[u8]) -> io::Result<()> { + self.max_write = self.max_write.max(buf.len()); + self.bytes.extend_from_slice(buf); + Ok(()) + } + + fn pos(&self) -> u64 { + self.bytes.len() as u64 + } + } + + #[test] + fn ivfflat_chunked_writer_preserves_format_and_bounds() { + const TEST_BUDGET: usize = 64; + + let mut index = IVFFlatIndex::new(3, 1, MetricType::L2); + index.quantizer_centroids = vec![1.0, 2.0, 3.0]; + index.ids[0] = vec![50, 10, 40, 20, 30, 60, 0]; + index.vectors[0] = index.ids[0] + .iter() + .flat_map(|&id| [id as f32, id as f32 + 0.25, id as f32 + 0.5]) + .collect(); + + let expected_bytes = serialized_flat_index(&index); + let mut chunked = MaxWriteWriter { + bytes: Vec::new(), + max_write: 0, + }; + write_ivfflat_index_with_buffer_limit(&index, &mut chunked, TEST_BUDGET).unwrap(); + + assert_eq!(chunked.bytes, expected_bytes); + assert!(chunked.max_write <= TEST_BUDGET); + + let mut reader = IVFFlatIndexReader::open(Cursor::new(chunked.bytes)).unwrap(); + let (ids, vectors) = reader.read_inverted_list(0).unwrap(); + assert_eq!(ids, vec![0, 10, 20, 30, 40, 50, 60]); + assert_eq!( + vectors, + ids.iter() + .flat_map(|&id| [id as f32, id as f32 + 0.25, id as f32 + 0.5]) + .collect::>() + ); + + let wide_dimension = TEST_BUDGET / size_of::() + 1; + let mut wide_index = IVFFlatIndex::new(wide_dimension, 1, MetricType::L2); + wide_index.quantizer_centroids = vec![0.0; wide_dimension]; + wide_index.ids[0] = vec![1]; + wide_index.vectors[0] = vec![1.0; wide_dimension]; + let mut wide = MaxWriteWriter { + bytes: Vec::new(), + max_write: 0, + }; + write_ivfflat_index_with_buffer_limit(&wide_index, &mut wide, TEST_BUDGET).unwrap(); + assert!(wide.max_write <= TEST_BUDGET); + } + #[test] fn ivfflat_streamed_list_reader_matches_full_payload() { let index = balanced_flat_index(8, 1, 257);