diff --git a/core/src/io.rs b/core/src/io.rs index 3a55823..618514a 100644 --- a/core/src/io.rs +++ b/core/src/io.rs @@ -39,6 +39,7 @@ pub const FLAG_DELTA_IDS: u32 = 1 << 2; pub const FLAG_TRANSPOSED_CODES: u32 = 1 << 3; const REQUIRED_FLAGS: u32 = FLAG_DELTA_IDS | FLAG_TRANSPOSED_CODES; const SUPPORTED_FLAGS: u32 = FLAG_HAS_OPQ | FLAG_BY_RESIDUAL | REQUIRED_FLAGS; +const IVFPQ_WRITE_BUFFER_SIZE: usize = 64 * 1024 * 1024; pub struct ReadRequest<'a> { pub pos: u64, @@ -220,9 +221,32 @@ 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 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 { @@ -278,6 +302,15 @@ fn checked_list_bytes(count: usize, bytes_per_entry: usize) -> io::Result /// Write a complete IVF-PQ index with delta-varint ID encoding. pub fn write_index(index: &IVFPQIndex, out: &mut dyn SeekWrite) -> io::Result<()> { + write_index_with_buffer_limit(index, out, IVFPQ_WRITE_BUFFER_SIZE) +} + +fn write_index_with_buffer_limit( + index: &IVFPQIndex, + out: &mut dyn SeekWrite, + buffer_limit: usize, +) -> io::Result<()> { + let mut write_buffer = Vec::new(); let d = index.d; let nlist = index.nlist; let m = index.pq.m; @@ -360,11 +393,16 @@ pub fn write_index(index: &IVFPQIndex, out: &mut dyn SeekWrite) -> io::Result<() out.write_all(&[0u8; 20])?; if let Some(ref opq) = index.opq { - write_f32_slice(out, &opq.rotation)?; + write_f32_slice(out, &opq.rotation, &mut write_buffer, buffer_limit)?; } - write_f32_slice(out, index.quantizer_centroids())?; - write_f32_slice(out, &index.pq.centroids)?; + write_f32_slice( + out, + index.quantizer_centroids(), + &mut write_buffer, + buffer_limit, + )?; + write_f32_slice(out, &index.pq.centroids, &mut write_buffer, buffer_limit)?; // Compute offsets for inverted lists // Delta-varint format per list: [base_id: i64][id_bytes_len: u32][id_bytes][codes] @@ -1354,6 +1392,62 @@ mod tests { assert!(encoded.len() < ids.len() * 8); } + #[test] + fn ivfpq_chunked_writer_preserves_format_and_bounds() { + // pq.centroids is m * ksub * dsub f32 (64 KiB here), far larger than any + // other section, so a small budget exercises chunked codebook writes. + const TEST_BUDGET: usize = 4096; + + 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 + } + } + + let (d, nlist, m, n) = (64usize, 4usize, 8usize, 512usize); + let mut index = IVFPQIndex::new(d, nlist, m, MetricType::L2, false); + let mut rng = rand::rngs::StdRng::seed_from_u64(7); + let data: Vec = (0..n * d).map(|_| rng.gen::()).collect(); + let ids: Vec = (0..n as i64).collect(); + index.train(&data, n); + index.add(&data, &ids, n); + + let mut expected = Vec::new(); + write_index(&index, &mut PosWriter::new(&mut expected)).unwrap(); + + let mut chunked = MaxWriteWriter { + bytes: Vec::new(), + max_write: 0, + }; + write_index_with_buffer_limit(&index, &mut chunked, TEST_BUDGET).unwrap(); + + assert_eq!( + chunked.bytes, expected, + "chunked output must be byte-identical" + ); + assert!( + chunked.max_write <= TEST_BUDGET, + "single write of {} bytes exceeds the {} byte budget", + chunked.max_write, + TEST_BUDGET + ); + + let mut cursor = Cursor::new(&chunked.bytes); + let reader = IVFPQIndexReader::open(&mut cursor).unwrap(); + assert_eq!(reader.total_vectors, n as i64); + } + #[test] fn test_write_read_roundtrip_delta_ids() { let d = 8;