1use std::collections::HashMap;
11use std::fs;
12use std::fs::File;
13use std::io::{BufReader, BufWriter, Write};
14use std::path::Path;
15
16use byteorder::{LittleEndian, ReadBytesExt};
17use fs2::FileExt;
18use memmap2::{Mmap, MmapMut};
19use ndarray::{Array1, Array2, ArrayView1, ArrayView2};
20
21use crate::error::{Error, Result};
22
23struct FileLockGuard {
26 _file: File,
27}
28
29impl FileLockGuard {
30 fn acquire(lock_path: &Path) -> Result<Self> {
34 let file = std::fs::OpenOptions::new()
35 .read(true)
36 .write(true)
37 .create(true)
38 .truncate(false)
39 .open(lock_path)
40 .map_err(|e| {
41 Error::IndexLoad(format!("Failed to open lock file {:?}: {}", lock_path, e))
42 })?;
43
44 file.lock_exclusive().map_err(|e| {
45 Error::IndexLoad(format!("Failed to acquire lock on {:?}: {}", lock_path, e))
46 })?;
47
48 Ok(Self { _file: file })
49 }
50}
51
52impl Drop for FileLockGuard {
53 fn drop(&mut self) {
54 let _ = self._file.unlock();
56 }
57}
58
59pub struct MmapArray2F32 {
63 _mmap: Mmap,
64 shape: (usize, usize),
65 data_offset: usize,
66}
67
68impl MmapArray2F32 {
69 pub fn from_raw_file(path: &Path) -> Result<Self> {
76 let file = File::open(path)
77 .map_err(|e| Error::IndexLoad(format!("Failed to open file {:?}: {}", path, e)))?;
78
79 let mmap = unsafe {
80 Mmap::map(&file)
81 .map_err(|e| Error::IndexLoad(format!("Failed to mmap file {:?}: {}", path, e)))?
82 };
83
84 if mmap.len() < 16 {
85 return Err(Error::IndexLoad("File too small for header".into()));
86 }
87
88 let mut cursor = std::io::Cursor::new(&mmap[..16]);
90 let nrows = cursor
91 .read_i64::<LittleEndian>()
92 .map_err(|e| Error::IndexLoad(format!("Failed to read nrows: {}", e)))?
93 as usize;
94 let ncols = cursor
95 .read_i64::<LittleEndian>()
96 .map_err(|e| Error::IndexLoad(format!("Failed to read ncols: {}", e)))?
97 as usize;
98
99 let expected_size = 16 + nrows * ncols * 4;
100 if mmap.len() < expected_size {
101 return Err(Error::IndexLoad(format!(
102 "File size {} too small for shape ({}, {})",
103 mmap.len(),
104 nrows,
105 ncols
106 )));
107 }
108
109 Ok(Self {
110 _mmap: mmap,
111 shape: (nrows, ncols),
112 data_offset: 16,
113 })
114 }
115
116 pub fn shape(&self) -> (usize, usize) {
118 self.shape
119 }
120
121 pub fn nrows(&self) -> usize {
123 self.shape.0
124 }
125
126 pub fn ncols(&self) -> usize {
128 self.shape.1
129 }
130
131 pub fn row(&self, idx: usize) -> ArrayView1<'_, f32> {
133 let start = self.data_offset + idx * self.shape.1 * 4;
134 let bytes = &self._mmap[start..start + self.shape.1 * 4];
135
136 let data =
138 unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const f32, self.shape.1) };
139
140 ArrayView1::from_shape(self.shape.1, data).unwrap()
141 }
142
143 pub fn load_rows(&self, start: usize, end: usize) -> Array2<f32> {
145 let nrows = end - start;
146 let byte_start = self.data_offset + start * self.shape.1 * 4;
147 let byte_end = self.data_offset + end * self.shape.1 * 4;
148 let bytes = &self._mmap[byte_start..byte_end];
149
150 let data = unsafe {
152 std::slice::from_raw_parts(bytes.as_ptr() as *const f32, nrows * self.shape.1)
153 };
154
155 Array2::from_shape_vec((nrows, self.shape.1), data.to_vec()).unwrap()
156 }
157
158 pub fn to_owned(&self) -> Array2<f32> {
160 self.load_rows(0, self.shape.0)
161 }
162}
163
164pub struct MmapArray2U8 {
166 _mmap: Mmap,
167 shape: (usize, usize),
168 data_offset: usize,
169}
170
171impl MmapArray2U8 {
172 pub fn from_raw_file(path: &Path) -> Result<Self> {
174 let file = File::open(path)
175 .map_err(|e| Error::IndexLoad(format!("Failed to open file {:?}: {}", path, e)))?;
176
177 let mmap = unsafe {
178 Mmap::map(&file)
179 .map_err(|e| Error::IndexLoad(format!("Failed to mmap file {:?}: {}", path, e)))?
180 };
181
182 if mmap.len() < 16 {
183 return Err(Error::IndexLoad("File too small for header".into()));
184 }
185
186 let mut cursor = std::io::Cursor::new(&mmap[..16]);
187 let nrows = cursor
188 .read_i64::<LittleEndian>()
189 .map_err(|e| Error::IndexLoad(format!("Failed to read nrows: {}", e)))?
190 as usize;
191 let ncols = cursor
192 .read_i64::<LittleEndian>()
193 .map_err(|e| Error::IndexLoad(format!("Failed to read ncols: {}", e)))?
194 as usize;
195
196 let expected_size = 16 + nrows * ncols;
197 if mmap.len() < expected_size {
198 return Err(Error::IndexLoad(format!(
199 "File size {} too small for shape ({}, {})",
200 mmap.len(),
201 nrows,
202 ncols
203 )));
204 }
205
206 Ok(Self {
207 _mmap: mmap,
208 shape: (nrows, ncols),
209 data_offset: 16,
210 })
211 }
212
213 pub fn shape(&self) -> (usize, usize) {
215 self.shape
216 }
217
218 pub fn view(&self) -> ArrayView2<'_, u8> {
220 let bytes = &self._mmap[self.data_offset..self.data_offset + self.shape.0 * self.shape.1];
221 ArrayView2::from_shape(self.shape, bytes).unwrap()
222 }
223
224 pub fn load_rows(&self, start: usize, end: usize) -> Array2<u8> {
226 let nrows = end - start;
227 let byte_start = self.data_offset + start * self.shape.1;
228 let byte_end = self.data_offset + end * self.shape.1;
229 let bytes = &self._mmap[byte_start..byte_end];
230
231 Array2::from_shape_vec((nrows, self.shape.1), bytes.to_vec()).unwrap()
232 }
233
234 pub fn to_owned(&self) -> Array2<u8> {
236 self.load_rows(0, self.shape.0)
237 }
238}
239
240pub struct MmapArray1I64 {
242 _mmap: Mmap,
243 len: usize,
244 data_offset: usize,
245}
246
247impl MmapArray1I64 {
248 pub fn from_raw_file(path: &Path) -> Result<Self> {
250 let file = File::open(path)
251 .map_err(|e| Error::IndexLoad(format!("Failed to open file {:?}: {}", path, e)))?;
252
253 let mmap = unsafe {
254 Mmap::map(&file)
255 .map_err(|e| Error::IndexLoad(format!("Failed to mmap file {:?}: {}", path, e)))?
256 };
257
258 if mmap.len() < 8 {
259 return Err(Error::IndexLoad("File too small for header".into()));
260 }
261
262 let mut cursor = std::io::Cursor::new(&mmap[..8]);
263 let len = cursor
264 .read_i64::<LittleEndian>()
265 .map_err(|e| Error::IndexLoad(format!("Failed to read length: {}", e)))?
266 as usize;
267
268 let expected_size = 8 + len * 8;
269 if mmap.len() < expected_size {
270 return Err(Error::IndexLoad(format!(
271 "File size {} too small for length {}",
272 mmap.len(),
273 len
274 )));
275 }
276
277 Ok(Self {
278 _mmap: mmap,
279 len,
280 data_offset: 8,
281 })
282 }
283
284 pub fn len(&self) -> usize {
286 self.len
287 }
288
289 pub fn is_empty(&self) -> bool {
291 self.len == 0
292 }
293
294 pub fn get(&self, idx: usize) -> i64 {
296 let start = self.data_offset + idx * 8;
297 let bytes = &self._mmap[start..start + 8];
298 i64::from_le_bytes(bytes.try_into().unwrap())
299 }
300
301 pub fn to_owned(&self) -> Array1<i64> {
303 let bytes = &self._mmap[self.data_offset..self.data_offset + self.len * 8];
304
305 let data = unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const i64, self.len) };
307
308 Array1::from_vec(data.to_vec())
309 }
310}
311
312pub fn write_array2_f32(array: &Array2<f32>, path: &Path) -> Result<()> {
314 use std::io::Write;
315
316 let file = File::create(path)
317 .map_err(|e| Error::IndexLoad(format!("Failed to create file {:?}: {}", path, e)))?;
318 let mut writer = std::io::BufWriter::new(file);
319
320 let nrows = array.nrows() as i64;
321 let ncols = array.ncols() as i64;
322
323 writer
324 .write_all(&nrows.to_le_bytes())
325 .map_err(|e| Error::IndexLoad(format!("Failed to write nrows: {}", e)))?;
326 writer
327 .write_all(&ncols.to_le_bytes())
328 .map_err(|e| Error::IndexLoad(format!("Failed to write ncols: {}", e)))?;
329
330 for val in array.iter() {
331 writer
332 .write_all(&val.to_le_bytes())
333 .map_err(|e| Error::IndexLoad(format!("Failed to write data: {}", e)))?;
334 }
335
336 writer
337 .flush()
338 .map_err(|e| Error::IndexLoad(format!("Failed to flush: {}", e)))?;
339
340 Ok(())
341}
342
343pub fn write_array2_u8(array: &Array2<u8>, path: &Path) -> Result<()> {
345 use std::io::Write;
346
347 let file = File::create(path)
348 .map_err(|e| Error::IndexLoad(format!("Failed to create file {:?}: {}", path, e)))?;
349 let mut writer = std::io::BufWriter::new(file);
350
351 let nrows = array.nrows() as i64;
352 let ncols = array.ncols() as i64;
353
354 writer
355 .write_all(&nrows.to_le_bytes())
356 .map_err(|e| Error::IndexLoad(format!("Failed to write nrows: {}", e)))?;
357 writer
358 .write_all(&ncols.to_le_bytes())
359 .map_err(|e| Error::IndexLoad(format!("Failed to write ncols: {}", e)))?;
360
361 for row in array.rows() {
362 writer
363 .write_all(row.as_slice().unwrap())
364 .map_err(|e| Error::IndexLoad(format!("Failed to write data: {}", e)))?;
365 }
366
367 writer
368 .flush()
369 .map_err(|e| Error::IndexLoad(format!("Failed to flush: {}", e)))?;
370
371 Ok(())
372}
373
374pub fn write_array1_i64(array: &Array1<i64>, path: &Path) -> Result<()> {
376 use std::io::Write;
377
378 let file = File::create(path)
379 .map_err(|e| Error::IndexLoad(format!("Failed to create file {:?}: {}", path, e)))?;
380 let mut writer = std::io::BufWriter::new(file);
381
382 let len = array.len() as i64;
383
384 writer
385 .write_all(&len.to_le_bytes())
386 .map_err(|e| Error::IndexLoad(format!("Failed to write length: {}", e)))?;
387
388 for val in array.iter() {
389 writer
390 .write_all(&val.to_le_bytes())
391 .map_err(|e| Error::IndexLoad(format!("Failed to write data: {}", e)))?;
392 }
393
394 writer
395 .flush()
396 .map_err(|e| Error::IndexLoad(format!("Failed to flush: {}", e)))?;
397
398 Ok(())
399}
400
401const NPY_MAGIC: &[u8] = b"\x93NUMPY";
407
408fn parse_dtype_from_header(header: &str) -> Result<String> {
410 let descr_start = header
412 .find("'descr':")
413 .ok_or_else(|| Error::IndexLoad("No descr in NPY header".into()))?;
414
415 let after_descr = &header[descr_start + 8..];
416 let quote_start = after_descr
417 .find('\'')
418 .ok_or_else(|| Error::IndexLoad("No dtype quote in NPY header".into()))?;
419 let rest = &after_descr[quote_start + 1..];
420 let quote_end = rest
421 .find('\'')
422 .ok_or_else(|| Error::IndexLoad("Unclosed dtype quote in NPY header".into()))?;
423
424 Ok(rest[..quote_end].to_string())
425}
426
427pub fn detect_npy_dtype(path: &Path) -> Result<String> {
429 let file = File::open(path)
430 .map_err(|e| Error::IndexLoad(format!("Failed to open NPY file {:?}: {}", path, e)))?;
431
432 let mmap = unsafe {
433 Mmap::map(&file)
434 .map_err(|e| Error::IndexLoad(format!("Failed to mmap NPY file {:?}: {}", path, e)))?
435 };
436
437 if mmap.len() < 10 {
438 return Err(Error::IndexLoad(format!(
439 "NPY file {:?} too small: {} bytes",
440 path,
441 mmap.len()
442 )));
443 }
444
445 if &mmap[..6] != NPY_MAGIC {
447 return Err(Error::IndexLoad("Invalid NPY magic".into()));
448 }
449
450 let major_version = mmap[6];
451
452 let header_len = if major_version == 1 {
454 u16::from_le_bytes([mmap[8], mmap[9]]) as usize
455 } else if major_version == 2 {
456 if mmap.len() < 12 {
457 return Err(Error::IndexLoad("NPY v2 file too small".into()));
458 }
459 u32::from_le_bytes([mmap[8], mmap[9], mmap[10], mmap[11]]) as usize
460 } else {
461 return Err(Error::IndexLoad(format!(
462 "Unsupported NPY version: {}",
463 major_version
464 )));
465 };
466
467 let header_start = if major_version == 1 { 10 } else { 12 };
468 let header_end = header_start + header_len;
469
470 if mmap.len() < header_end {
471 return Err(Error::IndexLoad("NPY header exceeds file size".into()));
472 }
473
474 let header_str = std::str::from_utf8(&mmap[header_start..header_end])
475 .map_err(|e| Error::IndexLoad(format!("Invalid NPY header encoding: {}", e)))?;
476
477 parse_dtype_from_header(header_str)
478}
479
480pub fn convert_f16_to_f32_npy(path: &Path) -> Result<()> {
482 use half::f16;
483 use std::io::Read;
484
485 let mut file = File::open(path)
487 .map_err(|e| Error::IndexLoad(format!("Failed to open {:?}: {}", path, e)))?;
488 let mut data = Vec::new();
489 file.read_to_end(&mut data)
490 .map_err(|e| Error::IndexLoad(format!("Failed to read {:?}: {}", path, e)))?;
491
492 if data.len() < 10 || &data[..6] != NPY_MAGIC {
493 return Err(Error::IndexLoad("Invalid NPY file".into()));
494 }
495
496 let major_version = data[6];
497 let header_start = if major_version == 1 { 10 } else { 12 };
498 let header_len = if major_version == 1 {
499 u16::from_le_bytes([data[8], data[9]]) as usize
500 } else {
501 u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize
502 };
503 let header_end = header_start + header_len;
504
505 let header_str = std::str::from_utf8(&data[header_start..header_end])
507 .map_err(|e| Error::IndexLoad(format!("Invalid header: {}", e)))?;
508 let shape = parse_shape_from_header(header_str)?;
509
510 let total_elements: usize = shape.iter().product();
512 let f16_data = &data[header_end..header_end + total_elements * 2];
513
514 let mut f32_data = Vec::with_capacity(total_elements * 4);
516 for chunk in f16_data.chunks(2) {
517 let f16_val = f16::from_le_bytes([chunk[0], chunk[1]]);
518 let f32_val: f32 = f16_val.to_f32();
519 f32_data.extend_from_slice(&f32_val.to_le_bytes());
520 }
521
522 let file = File::create(path)
524 .map_err(|e| Error::IndexLoad(format!("Failed to create {:?}: {}", path, e)))?;
525 let mut writer = BufWriter::new(file);
526
527 if shape.len() == 1 {
528 write_npy_header_1d(&mut writer, shape[0], "<f4")?;
529 } else if shape.len() == 2 {
530 write_npy_header_2d(&mut writer, shape[0], shape[1], "<f4")?;
531 } else {
532 return Err(Error::IndexLoad("Unsupported shape dimensions".into()));
533 }
534
535 writer
536 .write_all(&f32_data)
537 .map_err(|e| Error::IndexLoad(format!("Failed to write data: {}", e)))?;
538 writer.flush()?;
539
540 Ok(())
541}
542
543pub fn convert_i64_to_i32_npy(path: &Path) -> Result<()> {
545 use std::io::Read;
546
547 let mut file = File::open(path)
549 .map_err(|e| Error::IndexLoad(format!("Failed to open {:?}: {}", path, e)))?;
550 let mut data = Vec::new();
551 file.read_to_end(&mut data)
552 .map_err(|e| Error::IndexLoad(format!("Failed to read {:?}: {}", path, e)))?;
553
554 if data.len() < 10 || &data[..6] != NPY_MAGIC {
555 return Err(Error::IndexLoad("Invalid NPY file".into()));
556 }
557
558 let major_version = data[6];
559 let header_start = if major_version == 1 { 10 } else { 12 };
560 let header_len = if major_version == 1 {
561 u16::from_le_bytes([data[8], data[9]]) as usize
562 } else {
563 u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize
564 };
565 let header_end = header_start + header_len;
566
567 let header_str = std::str::from_utf8(&data[header_start..header_end])
569 .map_err(|e| Error::IndexLoad(format!("Invalid header: {}", e)))?;
570 let shape = parse_shape_from_header(header_str)?;
571
572 if shape.len() != 1 {
573 return Err(Error::IndexLoad("Expected 1D array for i64->i32".into()));
574 }
575
576 let len = shape[0];
577 let i64_data = &data[header_end..header_end + len * 8];
578
579 let mut i32_data = Vec::with_capacity(len * 4);
581 for chunk in i64_data.chunks(8) {
582 let i64_val = i64::from_le_bytes(chunk.try_into().unwrap());
583 let i32_val = i64_val as i32;
584 i32_data.extend_from_slice(&i32_val.to_le_bytes());
585 }
586
587 let file = File::create(path)
589 .map_err(|e| Error::IndexLoad(format!("Failed to create {:?}: {}", path, e)))?;
590 let mut writer = BufWriter::new(file);
591
592 write_npy_header_1d(&mut writer, len, "<i4")?;
593
594 writer
595 .write_all(&i32_data)
596 .map_err(|e| Error::IndexLoad(format!("Failed to write data: {}", e)))?;
597 writer.flush()?;
598
599 Ok(())
600}
601
602pub fn normalize_u8_npy(path: &Path) -> Result<()> {
607 use std::io::Read;
608
609 let mut file = File::open(path)
611 .map_err(|e| Error::IndexLoad(format!("Failed to open {:?}: {}", path, e)))?;
612 let mut data = Vec::new();
613 file.read_to_end(&mut data)
614 .map_err(|e| Error::IndexLoad(format!("Failed to read {:?}: {}", path, e)))?;
615
616 if data.len() < 10 || &data[..6] != NPY_MAGIC {
617 return Err(Error::IndexLoad("Invalid NPY file".into()));
618 }
619
620 let major_version = data[6];
621 let header_start = if major_version == 1 { 10 } else { 12 };
622 let header_len = if major_version == 1 {
623 u16::from_le_bytes([data[8], data[9]]) as usize
624 } else {
625 u32::from_le_bytes([data[8], data[9], data[10], data[11]]) as usize
626 };
627 let header_end = header_start + header_len;
628
629 let header_str = std::str::from_utf8(&data[header_start..header_end])
631 .map_err(|e| Error::IndexLoad(format!("Invalid header: {}", e)))?;
632 let shape = parse_shape_from_header(header_str)?;
633
634 if shape.len() != 2 {
635 return Err(Error::IndexLoad(
636 "Expected 2D array for u8 normalization".into(),
637 ));
638 }
639
640 let nrows = shape[0];
641 let ncols = shape[1];
642 let u8_data = &data[header_end..header_end + nrows * ncols];
643
644 let new_file = File::create(path)
646 .map_err(|e| Error::IndexLoad(format!("Failed to create {:?}: {}", path, e)))?;
647 let mut writer = BufWriter::new(new_file);
648
649 write_npy_header_2d(&mut writer, nrows, ncols, "|u1")?;
650
651 writer
652 .write_all(u8_data)
653 .map_err(|e| Error::IndexLoad(format!("Failed to write data: {}", e)))?;
654 writer.flush()?;
655
656 Ok(())
657}
658
659fn parse_npy_header(path: &Path, mmap: &Mmap) -> Result<(Vec<usize>, usize, bool)> {
661 if mmap.len() < 10 {
662 return Err(Error::IndexLoad(format!(
663 "NPY file {:?} too small: {} bytes",
664 path,
665 mmap.len()
666 )));
667 }
668
669 if &mmap[..6] != NPY_MAGIC {
671 return Err(Error::IndexLoad("Invalid NPY magic".into()));
672 }
673
674 let major_version = mmap[6];
675 let _minor_version = mmap[7];
676
677 let header_len = if major_version == 1 {
679 u16::from_le_bytes([mmap[8], mmap[9]]) as usize
680 } else if major_version == 2 {
681 if mmap.len() < 12 {
682 return Err(Error::IndexLoad(format!(
683 "NPY v2 file {:?} too small: {} bytes",
684 path,
685 mmap.len()
686 )));
687 }
688 u32::from_le_bytes([mmap[8], mmap[9], mmap[10], mmap[11]]) as usize
689 } else {
690 return Err(Error::IndexLoad(format!(
691 "Unsupported NPY version: {}",
692 major_version
693 )));
694 };
695
696 let header_start = if major_version == 1 { 10 } else { 12 };
697 let header_end = header_start + header_len;
698
699 if mmap.len() < header_end {
700 return Err(Error::IndexLoad(format!(
701 "NPY header exceeds file size for {:?}: header_end={}, file_size={}",
702 path,
703 header_end,
704 mmap.len()
705 )));
706 }
707
708 let header_str = std::str::from_utf8(&mmap[header_start..header_end])
710 .map_err(|e| Error::IndexLoad(format!("Invalid NPY header encoding: {}", e)))?;
711
712 let shape = parse_shape_from_header(header_str)?;
714 let fortran_order = header_str.contains("'fortran_order': True");
715
716 Ok((shape, header_end, fortran_order))
717}
718
719fn parse_shape_from_header(header: &str) -> Result<Vec<usize>> {
721 let shape_start = header
723 .find("'shape':")
724 .ok_or_else(|| Error::IndexLoad("No shape in NPY header".into()))?;
725
726 let after_shape = &header[shape_start + 8..];
727 let paren_start = after_shape
728 .find('(')
729 .ok_or_else(|| Error::IndexLoad("No shape tuple in NPY header".into()))?;
730 let paren_end = after_shape
731 .find(')')
732 .ok_or_else(|| Error::IndexLoad("Unclosed shape tuple in NPY header".into()))?;
733
734 let shape_content = &after_shape[paren_start + 1..paren_end];
735
736 let mut shape = Vec::new();
738 for part in shape_content.split(',') {
739 let trimmed = part.trim();
740 if !trimmed.is_empty() {
741 let dim: usize = trimmed.parse().map_err(|e| {
742 Error::IndexLoad(format!("Invalid shape dimension '{}': {}", trimmed, e))
743 })?;
744 shape.push(dim);
745 }
746 }
747
748 Ok(shape)
749}
750
751pub struct MmapNpyArray1I64 {
755 _mmap: Mmap,
756 len: usize,
757 data_offset: usize,
758}
759
760impl MmapNpyArray1I64 {
761 pub fn empty() -> Self {
766 let mmap = MmapMut::map_anon(1)
767 .expect("failed to create anonymous mmap")
768 .make_read_only()
769 .expect("failed to make anonymous mmap read-only");
770 Self {
771 _mmap: mmap,
772 len: 0,
773 data_offset: 0,
774 }
775 }
776
777 pub fn from_npy_file(path: &Path) -> Result<Self> {
779 let file = File::open(path)
780 .map_err(|e| Error::IndexLoad(format!("Failed to open NPY file {:?}: {}", path, e)))?;
781
782 let mmap = unsafe {
783 Mmap::map(&file).map_err(|e| {
784 Error::IndexLoad(format!("Failed to mmap NPY file {:?}: {}", path, e))
785 })?
786 };
787
788 let (shape, data_offset, _fortran_order) = parse_npy_header(path, &mmap)?;
789
790 if shape.is_empty() {
791 return Err(Error::IndexLoad("Empty shape in NPY file".into()));
792 }
793
794 let len = shape[0];
795
796 let expected_size = data_offset + len * 8;
798 if mmap.len() < expected_size {
799 return Err(Error::IndexLoad(format!(
800 "NPY file size {} too small for {} elements",
801 mmap.len(),
802 len
803 )));
804 }
805
806 Ok(Self {
807 _mmap: mmap,
808 len,
809 data_offset,
810 })
811 }
812
813 pub fn len(&self) -> usize {
815 self.len
816 }
817
818 pub fn is_empty(&self) -> bool {
820 self.len == 0
821 }
822
823 pub fn slice(&self, start: usize, end: usize) -> Vec<i64> {
830 let count = end - start;
831 let mut result = Vec::with_capacity(count);
832
833 for i in start..end {
834 result.push(self.get(i));
835 }
836
837 result
838 }
839
840 pub fn get(&self, idx: usize) -> i64 {
842 let start = self.data_offset + idx * 8;
843 let bytes = &self._mmap[start..start + 8];
844 i64::from_le_bytes(bytes.try_into().unwrap())
845 }
846}
847
848pub struct MmapNpyArray2F32 {
854 _mmap: Mmap,
855 shape: (usize, usize),
856 data_offset: usize,
857}
858
859impl MmapNpyArray2F32 {
860 pub fn from_npy_file(path: &Path) -> Result<Self> {
862 let file = File::open(path)
863 .map_err(|e| Error::IndexLoad(format!("Failed to open NPY file {:?}: {}", path, e)))?;
864
865 let mmap = unsafe {
866 Mmap::map(&file).map_err(|e| {
867 Error::IndexLoad(format!("Failed to mmap NPY file {:?}: {}", path, e))
868 })?
869 };
870
871 let (shape_vec, data_offset, _fortran_order) = parse_npy_header(path, &mmap)?;
872
873 if shape_vec.len() != 2 {
874 return Err(Error::IndexLoad(format!(
875 "Expected 2D array, got {}D",
876 shape_vec.len()
877 )));
878 }
879
880 let shape = (shape_vec[0], shape_vec[1]);
881
882 let expected_size = data_offset + shape.0 * shape.1 * 4;
884 if mmap.len() < expected_size {
885 return Err(Error::IndexLoad(format!(
886 "NPY file size {} too small for shape {:?}",
887 mmap.len(),
888 shape
889 )));
890 }
891
892 Ok(Self {
893 _mmap: mmap,
894 shape,
895 data_offset,
896 })
897 }
898
899 pub fn shape(&self) -> (usize, usize) {
901 self.shape
902 }
903
904 pub fn nrows(&self) -> usize {
906 self.shape.0
907 }
908
909 pub fn ncols(&self) -> usize {
911 self.shape.1
912 }
913
914 pub fn view(&self) -> ArrayView2<'_, f32> {
918 let byte_start = self.data_offset;
919 let byte_end = self.data_offset + self.shape.0 * self.shape.1 * 4;
920 let bytes = &self._mmap[byte_start..byte_end];
921
922 let data = unsafe {
924 std::slice::from_raw_parts(bytes.as_ptr() as *const f32, self.shape.0 * self.shape.1)
925 };
926
927 ArrayView2::from_shape(self.shape, data).unwrap()
928 }
929
930 pub fn row(&self, idx: usize) -> ArrayView1<'_, f32> {
932 let byte_start = self.data_offset + idx * self.shape.1 * 4;
933 let bytes = &self._mmap[byte_start..byte_start + self.shape.1 * 4];
934
935 let data =
937 unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const f32, self.shape.1) };
938
939 ArrayView1::from_shape(self.shape.1, data).unwrap()
940 }
941
942 pub fn slice_rows(&self, start: usize, end: usize) -> ArrayView2<'_, f32> {
944 let nrows = end - start;
945 let byte_start = self.data_offset + start * self.shape.1 * 4;
946 let byte_end = self.data_offset + end * self.shape.1 * 4;
947 let bytes = &self._mmap[byte_start..byte_end];
948
949 let data = unsafe {
951 std::slice::from_raw_parts(bytes.as_ptr() as *const f32, nrows * self.shape.1)
952 };
953
954 ArrayView2::from_shape((nrows, self.shape.1), data).unwrap()
955 }
956
957 pub fn to_owned(&self) -> Array2<f32> {
961 self.view().to_owned()
962 }
963}
964
965pub struct MmapNpyArray2U8 {
969 _mmap: Mmap,
970 shape: (usize, usize),
971 data_offset: usize,
972}
973
974impl MmapNpyArray2U8 {
975 pub fn empty() -> Self {
980 let mmap = MmapMut::map_anon(1)
981 .expect("failed to create anonymous mmap")
982 .make_read_only()
983 .expect("failed to make anonymous mmap read-only");
984 Self {
985 _mmap: mmap,
986 shape: (0, 0),
987 data_offset: 0,
988 }
989 }
990
991 pub fn from_npy_file(path: &Path) -> Result<Self> {
993 let file = File::open(path)
994 .map_err(|e| Error::IndexLoad(format!("Failed to open NPY file {:?}: {}", path, e)))?;
995
996 let mmap = unsafe {
997 Mmap::map(&file).map_err(|e| {
998 Error::IndexLoad(format!("Failed to mmap NPY file {:?}: {}", path, e))
999 })?
1000 };
1001
1002 let (shape_vec, data_offset, _fortran_order) = parse_npy_header(path, &mmap)?;
1003
1004 if shape_vec.len() != 2 {
1005 return Err(Error::IndexLoad(format!(
1006 "Expected 2D array, got {}D",
1007 shape_vec.len()
1008 )));
1009 }
1010
1011 let shape = (shape_vec[0], shape_vec[1]);
1012
1013 let expected_size = data_offset + shape.0 * shape.1;
1015 if mmap.len() < expected_size {
1016 return Err(Error::IndexLoad(format!(
1017 "NPY file size {} too small for shape {:?}",
1018 mmap.len(),
1019 shape
1020 )));
1021 }
1022
1023 Ok(Self {
1024 _mmap: mmap,
1025 shape,
1026 data_offset,
1027 })
1028 }
1029
1030 pub fn shape(&self) -> (usize, usize) {
1032 self.shape
1033 }
1034
1035 pub fn nrows(&self) -> usize {
1037 self.shape.0
1038 }
1039
1040 pub fn ncols(&self) -> usize {
1042 self.shape.1
1043 }
1044
1045 pub fn slice_rows(&self, start: usize, end: usize) -> ArrayView2<'_, u8> {
1047 let nrows = end - start;
1048 let byte_start = self.data_offset + start * self.shape.1;
1049 let byte_end = self.data_offset + end * self.shape.1;
1050 let bytes = &self._mmap[byte_start..byte_end];
1051
1052 ArrayView2::from_shape((nrows, self.shape.1), bytes).unwrap()
1053 }
1054
1055 pub fn view(&self) -> ArrayView2<'_, u8> {
1057 self.slice_rows(0, self.shape.0)
1058 }
1059
1060 pub fn row(&self, idx: usize) -> &[u8] {
1062 let byte_start = self.data_offset + idx * self.shape.1;
1063 let byte_end = byte_start + self.shape.1;
1064 &self._mmap[byte_start..byte_end]
1065 }
1066}
1067
1068#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1074pub struct ChunkManifestEntry {
1075 pub rows: usize,
1076 pub mtime: f64,
1077}
1078
1079#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1081pub struct MergeManifest {
1082 pub chunks: HashMap<String, ChunkManifestEntry>,
1084 #[serde(default)]
1086 pub padding_rows: usize,
1087 #[serde(default)]
1089 pub num_chunks: usize,
1090 #[serde(default)]
1092 pub metadata_mtime: f64,
1093 #[serde(default)]
1095 pub total_rows: usize,
1096 #[serde(default)]
1098 pub ncols: usize,
1099}
1100
1101pub type ChunkManifest = HashMap<String, ChunkManifestEntry>;
1103
1104fn load_merge_manifest(manifest_path: &Path) -> Option<MergeManifest> {
1107 if manifest_path.exists() {
1108 if let Ok(file) = File::open(manifest_path) {
1109 let reader = BufReader::new(file);
1111 if let Ok(manifest) = serde_json::from_reader::<_, MergeManifest>(reader) {
1112 return Some(manifest);
1113 }
1114 if let Ok(file) = File::open(manifest_path) {
1116 if let Ok(chunks) =
1117 serde_json::from_reader::<_, ChunkManifest>(BufReader::new(file))
1118 {
1119 return Some(MergeManifest {
1121 chunks,
1122 padding_rows: 0,
1123 total_rows: 0,
1124 ncols: 0,
1125 num_chunks: 0,
1126 metadata_mtime: 0.0,
1127 });
1128 }
1129 }
1130 }
1131 }
1132 None
1133}
1134
1135fn save_merge_manifest(manifest_path: &Path, manifest: &MergeManifest) -> Result<()> {
1137 let temp_path = manifest_path.with_extension("manifest.json.tmp");
1138
1139 let file = File::create(&temp_path)
1141 .map_err(|e| Error::IndexLoad(format!("Failed to create temp manifest: {}", e)))?;
1142 let mut writer = BufWriter::new(file);
1143 serde_json::to_writer(&mut writer, manifest)
1144 .map_err(|e| Error::IndexLoad(format!("Failed to write manifest: {}", e)))?;
1145 writer
1146 .flush()
1147 .map_err(|e| Error::IndexLoad(format!("Failed to flush manifest: {}", e)))?;
1148
1149 writer
1151 .into_inner()
1152 .map_err(|e| Error::IndexLoad(format!("Failed to get inner file: {}", e)))?
1153 .sync_all()
1154 .map_err(|e| Error::IndexLoad(format!("Failed to sync manifest: {}", e)))?;
1155
1156 fs::rename(&temp_path, manifest_path)
1158 .map_err(|e| Error::IndexLoad(format!("Failed to rename manifest: {}", e)))?;
1159
1160 Ok(())
1161}
1162
1163fn get_mtime(path: &Path) -> Result<f64> {
1165 let metadata = fs::metadata(path)
1166 .map_err(|e| Error::IndexLoad(format!("Failed to get metadata for {:?}: {}", path, e)))?;
1167 let mtime = metadata
1168 .modified()
1169 .map_err(|e| Error::IndexLoad(format!("Failed to get mtime: {}", e)))?;
1170 let duration = mtime
1171 .duration_since(std::time::UNIX_EPOCH)
1172 .map_err(|e| Error::IndexLoad(format!("Invalid mtime: {}", e)))?;
1173 Ok(duration.as_secs_f64())
1174}
1175
1176fn npy_header_layout(header_dict: &str) -> (usize, usize) {
1178 let header_len = header_dict.len();
1179 let padding = (64 - ((10 + header_len) % 64)) % 64;
1180 let total = 10 + header_len + padding + 1; (padding, total)
1182}
1183
1184fn npy_header_dict_1d(len: usize, dtype: &str) -> String {
1185 format!(
1186 "{{'descr': '{}', 'fortran_order': False, 'shape': ({},), }}",
1187 dtype, len
1188 )
1189}
1190
1191fn npy_header_dict_2d(nrows: usize, ncols: usize, dtype: &str) -> String {
1192 format!(
1193 "{{'descr': '{}', 'fortran_order': False, 'shape': ({}, {}), }}",
1194 dtype, nrows, ncols
1195 )
1196}
1197
1198fn npy_header_size_1d(len: usize, dtype: &str) -> usize {
1200 let dict = npy_header_dict_1d(len, dtype);
1201 npy_header_layout(&dict).1
1202}
1203
1204fn npy_header_size_2d(nrows: usize, ncols: usize, dtype: &str) -> usize {
1206 let dict = npy_header_dict_2d(nrows, ncols, dtype);
1207 npy_header_layout(&dict).1
1208}
1209
1210fn write_npy_header(writer: &mut impl Write, header_dict: &str) -> Result<usize> {
1212 let (padding, total) = npy_header_layout(header_dict);
1213 let padded_header = format!("{}{}\n", header_dict, " ".repeat(padding));
1214
1215 writer
1217 .write_all(NPY_MAGIC)
1218 .map_err(|e| Error::IndexLoad(format!("Failed to write NPY magic: {}", e)))?;
1219 writer
1220 .write_all(&[1, 0])
1221 .map_err(|e| Error::IndexLoad(format!("Failed to write version: {}", e)))?;
1222
1223 let header_len_bytes = (padded_header.len() as u16).to_le_bytes();
1225 writer
1226 .write_all(&header_len_bytes)
1227 .map_err(|e| Error::IndexLoad(format!("Failed to write header len: {}", e)))?;
1228
1229 writer
1231 .write_all(padded_header.as_bytes())
1232 .map_err(|e| Error::IndexLoad(format!("Failed to write header: {}", e)))?;
1233
1234 Ok(total)
1235}
1236
1237fn write_npy_header_1d(writer: &mut impl Write, len: usize, dtype: &str) -> Result<usize> {
1239 write_npy_header(writer, &npy_header_dict_1d(len, dtype))
1240}
1241
1242fn write_npy_header_2d(
1244 writer: &mut impl Write,
1245 nrows: usize,
1246 ncols: usize,
1247 dtype: &str,
1248) -> Result<usize> {
1249 write_npy_header(writer, &npy_header_dict_2d(nrows, ncols, dtype))
1250}
1251
1252struct ChunkInfo {
1254 path: std::path::PathBuf,
1255 filename: String,
1256 rows: usize,
1257 mtime: f64,
1258}
1259
1260pub fn merge_codes_chunks(
1267 index_path: &Path,
1268 num_chunks: usize,
1269 padding_rows: usize,
1270) -> Result<std::path::PathBuf> {
1271 use ndarray_npy::ReadNpyExt;
1272
1273 let merged_path = index_path.join("merged_codes.npy");
1274 let manifest_path = index_path.join("merged_codes.manifest.json");
1275 let temp_path = index_path.join("merged_codes.npy.tmp");
1276 let lock_path = index_path.join("merged_codes.lock");
1277
1278 let metadata_json_path = index_path.join("metadata.json");
1281 let current_metadata_mtime = get_mtime(&metadata_json_path).unwrap_or(0.0);
1282 if let Some(ref manifest) = load_merge_manifest(&manifest_path) {
1283 let mtime_matches = manifest.metadata_mtime > 0.0
1284 && (manifest.metadata_mtime - current_metadata_mtime).abs() < 0.001;
1285 if manifest.num_chunks == num_chunks
1286 && manifest.padding_rows == padding_rows
1287 && manifest.chunks.len() == num_chunks
1288 && manifest.total_rows > 0
1289 && mtime_matches
1290 && merged_path.exists()
1291 {
1292 if let Ok(meta) = std::fs::metadata(&merged_path) {
1293 let expected_size = npy_header_size_1d(manifest.total_rows, "<i8")
1294 + manifest.total_rows * std::mem::size_of::<i64>();
1295 if meta.len() == expected_size as u64 {
1296 return Ok(merged_path);
1297 }
1298 }
1299 }
1300 }
1301
1302 let _lock = FileLockGuard::acquire(&lock_path)?;
1305
1306 let old_manifest = load_merge_manifest(&manifest_path);
1311
1312 let mut chunks: Vec<ChunkInfo> = Vec::new();
1314 let mut total_rows = 0usize;
1315 let mut chain_broken = false;
1316
1317 for i in 0..num_chunks {
1318 let filename = format!("{}.codes.npy", i);
1319 let path = index_path.join(&filename);
1320
1321 if path.exists() {
1322 let mtime = get_mtime(&path)?;
1323
1324 let file = File::open(&path)?;
1326 let arr: Array1<i64> = Array1::read_npy(file)?;
1327 let rows = arr.len();
1328
1329 if rows > 0 {
1330 total_rows += rows;
1331
1332 let is_clean = if let Some(ref manifest) = old_manifest {
1334 manifest
1335 .chunks
1336 .get(&filename)
1337 .is_some_and(|entry| entry.mtime == mtime && entry.rows == rows)
1338 } else {
1339 false
1340 };
1341
1342 if !is_clean {
1343 chain_broken = true;
1344 }
1345
1346 chunks.push(ChunkInfo {
1347 path,
1348 filename,
1349 rows,
1350 mtime,
1351 });
1352 }
1353 }
1354 }
1355
1356 if total_rows == 0 {
1357 return Err(Error::IndexLoad("No data to merge".into()));
1358 }
1359
1360 let final_rows = total_rows + padding_rows;
1361
1362 let padding_changed = old_manifest
1368 .as_ref()
1369 .map(|m| m.padding_rows != padding_rows)
1370 .unwrap_or(true);
1371 let total_rows_mismatch = old_manifest
1372 .as_ref()
1373 .map(|m| m.total_rows != final_rows)
1374 .unwrap_or(true);
1375
1376 let needs_full_rewrite =
1377 !merged_path.exists() || chain_broken || padding_changed || total_rows_mismatch;
1378
1379 if needs_full_rewrite {
1380 let file = File::create(&temp_path)
1382 .map_err(|e| Error::IndexLoad(format!("Failed to create temp merged file: {}", e)))?;
1383 let mut writer = BufWriter::new(file);
1384
1385 let header_size = write_npy_header_1d(&mut writer, final_rows, "<i8")?;
1387
1388 let mut written_rows = 0usize;
1390 for chunk in &chunks {
1391 let file = File::open(&chunk.path)?;
1392 let arr: Array1<i64> = Array1::read_npy(file)?;
1393 for &val in arr.iter() {
1394 writer.write_all(&val.to_le_bytes())?;
1395 }
1396 written_rows += arr.len();
1397 }
1398
1399 for _ in 0..padding_rows {
1401 writer.write_all(&0i64.to_le_bytes())?;
1402 }
1403 written_rows += padding_rows;
1404
1405 writer
1407 .flush()
1408 .map_err(|e| Error::IndexLoad(format!("Failed to flush merged file: {}", e)))?;
1409 let file = writer
1410 .into_inner()
1411 .map_err(|e| Error::IndexLoad(format!("Failed to get inner file: {}", e)))?;
1412 file.sync_all()
1413 .map_err(|e| Error::IndexLoad(format!("Failed to sync merged file to disk: {}", e)))?;
1414
1415 let expected_size = header_size + written_rows * 8;
1417 let actual_size = fs::metadata(&temp_path)
1418 .map_err(|e| Error::IndexLoad(format!("Failed to get temp file metadata: {}", e)))?
1419 .len() as usize;
1420
1421 if actual_size != expected_size {
1422 let _ = fs::remove_file(&temp_path);
1424 return Err(Error::IndexLoad(format!(
1425 "Merged codes file size mismatch: expected {} bytes, got {} bytes",
1426 expected_size, actual_size
1427 )));
1428 }
1429
1430 fs::rename(&temp_path, &merged_path)
1432 .map_err(|e| Error::IndexLoad(format!("Failed to rename merged file: {}", e)))?;
1433 } else {
1434 if merged_path.exists() {
1436 let file_size = fs::metadata(&merged_path)
1437 .map_err(|e| {
1438 Error::IndexLoad(format!("Failed to get merged file metadata: {}", e))
1439 })?
1440 .len() as usize;
1441
1442 let min_expected_size = 64 + final_rows * 8;
1444 if file_size < min_expected_size {
1445 let _ = fs::remove_file(&merged_path);
1447 let _ = fs::remove_file(&manifest_path);
1448 drop(_lock);
1450 return merge_codes_chunks(index_path, num_chunks, padding_rows);
1451 }
1452 }
1453 }
1454
1455 let mut chunk_map = HashMap::new();
1457 for chunk in &chunks {
1458 chunk_map.insert(
1459 chunk.filename.clone(),
1460 ChunkManifestEntry {
1461 rows: chunk.rows,
1462 mtime: chunk.mtime,
1463 },
1464 );
1465 }
1466 let new_manifest = MergeManifest {
1467 chunks: chunk_map,
1468 padding_rows,
1469 total_rows: final_rows,
1470 ncols: 0, num_chunks,
1472 metadata_mtime: current_metadata_mtime,
1473 };
1474 save_merge_manifest(&manifest_path, &new_manifest)?;
1475
1476 Ok(merged_path)
1477}
1478
1479pub fn merge_residuals_chunks(
1484 index_path: &Path,
1485 num_chunks: usize,
1486 padding_rows: usize,
1487) -> Result<std::path::PathBuf> {
1488 use ndarray_npy::ReadNpyExt;
1489
1490 let merged_path = index_path.join("merged_residuals.npy");
1491 let manifest_path = index_path.join("merged_residuals.manifest.json");
1492 let temp_path = index_path.join("merged_residuals.npy.tmp");
1493 let lock_path = index_path.join("merged_residuals.lock");
1494
1495 let metadata_json_path = index_path.join("metadata.json");
1498 let current_metadata_mtime = get_mtime(&metadata_json_path).unwrap_or(0.0);
1499 if let Some(ref manifest) = load_merge_manifest(&manifest_path) {
1500 if manifest.num_chunks == num_chunks
1501 && manifest.padding_rows == padding_rows
1502 && manifest.chunks.len() == num_chunks
1503 && manifest.total_rows > 0
1504 && manifest.ncols > 0
1505 && manifest.metadata_mtime > 0.0
1506 && (manifest.metadata_mtime - current_metadata_mtime).abs() < 0.001
1507 && merged_path.exists()
1508 {
1509 if let Ok(meta) = std::fs::metadata(&merged_path) {
1510 let expected_size = npy_header_size_2d(manifest.total_rows, manifest.ncols, "|u1")
1511 + manifest.total_rows * manifest.ncols;
1512 if meta.len() == expected_size as u64 {
1513 return Ok(merged_path);
1514 }
1515 }
1516 }
1517 }
1518
1519 let _lock = FileLockGuard::acquire(&lock_path)?;
1522
1523 let old_manifest = load_merge_manifest(&manifest_path);
1528
1529 let mut chunks: Vec<ChunkInfo> = Vec::new();
1531 let mut total_rows = 0usize;
1532 let mut ncols = 0usize;
1533 let mut chain_broken = false;
1534
1535 for i in 0..num_chunks {
1536 let filename = format!("{}.residuals.npy", i);
1537 let path = index_path.join(&filename);
1538
1539 if path.exists() {
1540 let mtime = get_mtime(&path)?;
1541
1542 let file = File::open(&path)?;
1544 let arr: Array2<u8> = Array2::read_npy(file)?;
1545 let rows = arr.nrows();
1546 ncols = arr.ncols();
1547
1548 if rows > 0 {
1549 total_rows += rows;
1550
1551 let is_clean = if let Some(ref manifest) = old_manifest {
1552 manifest
1553 .chunks
1554 .get(&filename)
1555 .is_some_and(|entry| entry.mtime == mtime && entry.rows == rows)
1556 } else {
1557 false
1558 };
1559
1560 if !is_clean {
1561 chain_broken = true;
1562 }
1563
1564 chunks.push(ChunkInfo {
1565 path,
1566 filename,
1567 rows,
1568 mtime,
1569 });
1570 }
1571 }
1572 }
1573
1574 if total_rows == 0 || ncols == 0 {
1575 return Err(Error::IndexLoad("No residual data to merge".into()));
1576 }
1577
1578 let final_rows = total_rows + padding_rows;
1579
1580 let padding_changed = old_manifest
1586 .as_ref()
1587 .map(|m| m.padding_rows != padding_rows)
1588 .unwrap_or(true);
1589 let total_rows_mismatch = old_manifest
1590 .as_ref()
1591 .map(|m| m.total_rows != final_rows)
1592 .unwrap_or(true);
1593 let ncols_mismatch = old_manifest
1594 .as_ref()
1595 .map(|m| m.ncols != ncols && m.ncols != 0)
1596 .unwrap_or(false);
1597
1598 let needs_full_rewrite = !merged_path.exists()
1599 || chain_broken
1600 || padding_changed
1601 || total_rows_mismatch
1602 || ncols_mismatch;
1603
1604 if needs_full_rewrite {
1605 let file = File::create(&temp_path)
1607 .map_err(|e| Error::IndexLoad(format!("Failed to create temp merged file: {}", e)))?;
1608 let mut writer = BufWriter::new(file);
1609
1610 let header_size = write_npy_header_2d(&mut writer, final_rows, ncols, "|u1")?;
1612
1613 let mut written_rows = 0usize;
1615 for chunk in &chunks {
1616 let file = File::open(&chunk.path)?;
1617 let arr: Array2<u8> = Array2::read_npy(file)?;
1618 for row in arr.rows() {
1619 writer.write_all(row.as_slice().unwrap())?;
1620 }
1621 written_rows += arr.nrows();
1622 }
1623
1624 let zero_row = vec![0u8; ncols];
1626 for _ in 0..padding_rows {
1627 writer.write_all(&zero_row)?;
1628 }
1629 written_rows += padding_rows;
1630
1631 writer
1633 .flush()
1634 .map_err(|e| Error::IndexLoad(format!("Failed to flush merged residuals: {}", e)))?;
1635 let file = writer
1636 .into_inner()
1637 .map_err(|e| Error::IndexLoad(format!("Failed to get inner file: {}", e)))?;
1638 file.sync_all().map_err(|e| {
1639 Error::IndexLoad(format!("Failed to sync merged residuals to disk: {}", e))
1640 })?;
1641
1642 let expected_size = header_size + written_rows * ncols;
1644 let actual_size = fs::metadata(&temp_path)
1645 .map_err(|e| Error::IndexLoad(format!("Failed to get temp file metadata: {}", e)))?
1646 .len() as usize;
1647
1648 if actual_size != expected_size {
1649 let _ = fs::remove_file(&temp_path);
1651 return Err(Error::IndexLoad(format!(
1652 "Merged residuals file size mismatch: expected {} bytes, got {} bytes",
1653 expected_size, actual_size
1654 )));
1655 }
1656
1657 fs::rename(&temp_path, &merged_path)
1659 .map_err(|e| Error::IndexLoad(format!("Failed to rename merged residuals: {}", e)))?;
1660 } else {
1661 if merged_path.exists() {
1663 let file_size = fs::metadata(&merged_path)
1664 .map_err(|e| {
1665 Error::IndexLoad(format!("Failed to get merged file metadata: {}", e))
1666 })?
1667 .len() as usize;
1668
1669 let min_expected_size = 64 + final_rows * ncols;
1671 if file_size < min_expected_size {
1672 let _ = fs::remove_file(&merged_path);
1674 let _ = fs::remove_file(&manifest_path);
1675 drop(_lock);
1677 return merge_residuals_chunks(index_path, num_chunks, padding_rows);
1678 }
1679 }
1680 }
1681
1682 let mut chunk_map = HashMap::new();
1684 for chunk in &chunks {
1685 chunk_map.insert(
1686 chunk.filename.clone(),
1687 ChunkManifestEntry {
1688 rows: chunk.rows,
1689 mtime: chunk.mtime,
1690 },
1691 );
1692 }
1693 let new_manifest = MergeManifest {
1694 chunks: chunk_map,
1695 padding_rows,
1696 total_rows: final_rows,
1697 ncols,
1698 num_chunks,
1699 metadata_mtime: current_metadata_mtime,
1700 };
1701 save_merge_manifest(&manifest_path, &new_manifest)?;
1702
1703 Ok(merged_path)
1704}
1705
1706pub fn clear_merged_files(index_path: &Path) -> Result<()> {
1715 let codes_lock_path = index_path.join("merged_codes.lock");
1719 let residuals_lock_path = index_path.join("merged_residuals.lock");
1720 let _codes_lock = FileLockGuard::acquire(&codes_lock_path)?;
1721 let _residuals_lock = FileLockGuard::acquire(&residuals_lock_path)?;
1722
1723 let files_to_remove = [
1724 "merged_codes.npy",
1725 "merged_codes.npy.tmp",
1726 "merged_codes.manifest.json",
1727 "merged_codes.manifest.json.tmp",
1728 "merged_residuals.npy",
1729 "merged_residuals.npy.tmp",
1730 "merged_residuals.manifest.json",
1731 "merged_residuals.manifest.json.tmp",
1732 ];
1733
1734 for filename in files_to_remove {
1735 let path = index_path.join(filename);
1736 if path.exists() {
1737 fs::remove_file(&path)
1738 .map_err(|e| Error::IndexLoad(format!("Failed to remove {}: {}", filename, e)))?;
1739 }
1740 }
1741
1742 Ok(())
1743}
1744
1745pub fn convert_fastplaid_to_nextplaid(index_path: &Path) -> Result<bool> {
1758 let mut converted = false;
1759
1760 let float_files = [
1762 "centroids.npy",
1763 "avg_residual.npy",
1764 "bucket_cutoffs.npy",
1765 "bucket_weights.npy",
1766 ];
1767
1768 for filename in float_files {
1769 let path = index_path.join(filename);
1770 if path.exists() {
1771 let dtype = detect_npy_dtype(&path)?;
1772 if dtype == "<f2" {
1773 eprintln!(" Converting {} from float16 to float32", filename);
1774 convert_f16_to_f32_npy(&path)?;
1775 converted = true;
1776 }
1777 }
1778 }
1779
1780 let ivf_lengths_path = index_path.join("ivf_lengths.npy");
1782 if ivf_lengths_path.exists() {
1783 let dtype = detect_npy_dtype(&ivf_lengths_path)?;
1784 if dtype == "<i8" {
1785 eprintln!(" Converting ivf_lengths.npy from int64 to int32");
1786 convert_i64_to_i32_npy(&ivf_lengths_path)?;
1787 converted = true;
1788 }
1789 }
1790
1791 for entry in fs::read_dir(index_path)? {
1794 let entry = entry?;
1795 let filename = entry.file_name().to_string_lossy().to_string();
1796 if filename.ends_with(".residuals.npy") {
1797 let path = entry.path();
1798 let dtype = detect_npy_dtype(&path)?;
1799 if dtype == "<u1" {
1800 eprintln!(
1801 " Normalizing {} dtype descriptor from <u1 to |u1",
1802 filename
1803 );
1804 normalize_u8_npy(&path)?;
1805 converted = true;
1806 }
1807 }
1808 }
1809
1810 Ok(converted)
1811}
1812
1813#[cfg(test)]
1814mod tests {
1815 use super::*;
1816 use std::io::Write;
1817 use tempfile::NamedTempFile;
1818
1819 #[test]
1820 fn test_mmap_array2_f32() {
1821 let mut file = NamedTempFile::new().unwrap();
1823
1824 file.write_all(&3i64.to_le_bytes()).unwrap();
1826 file.write_all(&2i64.to_le_bytes()).unwrap();
1827
1828 for val in [1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0] {
1830 file.write_all(&val.to_le_bytes()).unwrap();
1831 }
1832
1833 file.flush().unwrap();
1834
1835 let mmap = MmapArray2F32::from_raw_file(file.path()).unwrap();
1837 assert_eq!(mmap.shape(), (3, 2));
1838
1839 let row0 = mmap.row(0);
1840 assert_eq!(row0[0], 1.0);
1841 assert_eq!(row0[1], 2.0);
1842
1843 let owned = mmap.to_owned();
1844 assert_eq!(owned[[2, 0]], 5.0);
1845 assert_eq!(owned[[2, 1]], 6.0);
1846 }
1847
1848 #[test]
1849 fn test_mmap_array1_i64() {
1850 let mut file = NamedTempFile::new().unwrap();
1851
1852 file.write_all(&4i64.to_le_bytes()).unwrap();
1854
1855 for val in [10i64, 20, 30, 40] {
1857 file.write_all(&val.to_le_bytes()).unwrap();
1858 }
1859
1860 file.flush().unwrap();
1861
1862 let mmap = MmapArray1I64::from_raw_file(file.path()).unwrap();
1863 assert_eq!(mmap.len(), 4);
1864 assert_eq!(mmap.get(0), 10);
1865 assert_eq!(mmap.get(3), 40);
1866
1867 let owned = mmap.to_owned();
1868 assert_eq!(owned[1], 20);
1869 assert_eq!(owned[2], 30);
1870 }
1871
1872 #[test]
1873 fn test_write_read_roundtrip() {
1874 let file = NamedTempFile::new().unwrap();
1875 let path = file.path();
1876
1877 let array = Array2::from_shape_vec((2, 3), vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
1879
1880 write_array2_f32(&array, path).unwrap();
1882
1883 let mmap = MmapArray2F32::from_raw_file(path).unwrap();
1885 let loaded = mmap.to_owned();
1886
1887 assert_eq!(array, loaded);
1888 }
1889}