From 1404974bfda5407c2645d683c0b0642616ad4a60 Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 13 Aug 2026 15:45:55 +0800 Subject: [PATCH 1/2] fix: bound IVF-Flat serialization memory --- core/src/ivfflat_io.rs | 129 +++++++++++++++++++++++++++++++++++------ 1 file changed, 112 insertions(+), 17 deletions(-) diff --git a/core/src/ivfflat_io.rs b/core/src/ivfflat_io.rs index 1409caa..8e4f2c4 100644 --- a/core/src/ivfflat_io.rs +++ b/core/src/ivfflat_io.rs @@ -38,14 +38,31 @@ 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 buffer_limit = buffer_limit.max(bytes_per_vector); + 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 +102,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 +136,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 +164,17 @@ 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]); + let rows_per_chunk = (buffer_limit / bytes_per_vector).max(1); + for chunk in order.chunks(rows_per_chunk) { + checked_list_bytes(chunk.len(), bytes_per_vector)?; + write_buffer.clear(); + for &idx in chunk { + for value in &index.vectors[list_id][idx * d..(idx + 1) * d] { + write_buffer.extend_from_slice(&value.to_le_bytes()); + } + } + out.write_all(&write_buffer)?; } - write_f32_slice(out, &sorted_vectors)?; } Ok(()) @@ -1047,9 +1067,22 @@ 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<()> { + let floats_per_chunk = (buffer_limit / size_of::()).max(1); + for chunk in data.chunks(floats_per_chunk) { + checked_list_bytes(chunk.len(), size_of::())?; + buffer.clear(); + for value in chunk { + buffer.extend_from_slice(&value.to_le_bytes()); + } + out.write_all(buffer)?; + } + Ok(()) } fn validate_positive_i32(val: i32, field: &str) -> io::Result { @@ -1326,6 +1359,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_eq!(wide.max_write, wide_dimension * size_of::()); + } + #[test] fn ivfflat_streamed_list_reader_matches_full_payload() { let index = balanced_flat_index(8, 1, 257); From 184573f0df3a52d2e11fc90c57d89cd1cdcd61ad Mon Sep 17 00:00:00 2001 From: yantian Date: Thu, 13 Aug 2026 17:15:18 +0800 Subject: [PATCH 2/2] fix: enforce IVF-Flat write buffer limit --- core/src/ivfflat_io.rs | 54 +++++++++++++++++++++++++++--------------- 1 file changed, 35 insertions(+), 19 deletions(-) diff --git a/core/src/ivfflat_io.rs b/core/src/ivfflat_io.rs index 8e4f2c4..25d24d7 100644 --- a/core/src/ivfflat_io.rs +++ b/core/src/ivfflat_io.rs @@ -61,7 +61,6 @@ fn write_ivfflat_index_with_buffer_limit( "IVF-FLAT bytes per vector overflow", ) })?; - let buffer_limit = buffer_limit.max(bytes_per_vector); let mut write_buffer = Vec::new(); let d_i32 = usize_to_i32(d, "dimension")?; let nlist_i32 = usize_to_i32(nlist, "nlist")?; @@ -164,17 +163,15 @@ fn write_ivfflat_index_with_buffer_limit( 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 rows_per_chunk = (buffer_limit / bytes_per_vector).max(1); - for chunk in order.chunks(rows_per_chunk) { - checked_list_bytes(chunk.len(), bytes_per_vector)?; - write_buffer.clear(); - for &idx in chunk { - for value in &index.vectors[list_id][idx * d..(idx + 1) * d] { - write_buffer.extend_from_slice(&value.to_le_bytes()); - } - } - out.write_all(&write_buffer)?; - } + 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(()) @@ -1073,14 +1070,33 @@ fn write_f32_slice( buffer: &mut Vec, buffer_limit: usize, ) -> io::Result<()> { - let floats_per_chunk = (buffer_limit / size_of::()).max(1); - for chunk in data.chunks(floats_per_chunk) { - checked_list_bytes(chunk.len(), size_of::())?; - buffer.clear(); - for value in chunk { - buffer.extend_from_slice(&value.to_le_bytes()); + 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(()) } @@ -1418,7 +1434,7 @@ mod tests { max_write: 0, }; write_ivfflat_index_with_buffer_limit(&wide_index, &mut wide, TEST_BUDGET).unwrap(); - assert_eq!(wide.max_write, wide_dimension * size_of::()); + assert!(wide.max_write <= TEST_BUDGET); } #[test]