Skip to main content

next_plaid/
mmap.rs

1//! Memory-mapped file support for efficient large index loading.
2//!
3//! This module provides utilities for loading large arrays from disk using
4//! memory-mapped files, avoiding the need to load entire arrays into RAM.
5//!
6//! Two formats are supported:
7//! - Custom raw binary format (legacy): 8-byte header with shape, then raw data
8//! - NPY format: Standard NumPy format with header, used for index files
9
10use 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
23/// RAII guard for file-based locking to coordinate concurrent processes.
24/// The lock is released when this guard is dropped.
25struct FileLockGuard {
26    _file: File,
27}
28
29impl FileLockGuard {
30    /// Acquire an exclusive lock on the given lock file path.
31    /// Creates the lock file if it doesn't exist.
32    /// Blocks until the lock is acquired.
33    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        // Lock is automatically released when file is closed
55        let _ = self._file.unlock();
56    }
57}
58
59/// A memory-mapped array of f32 values.
60///
61/// This struct provides zero-copy access to large arrays stored on disk.
62pub struct MmapArray2F32 {
63    _mmap: Mmap,
64    shape: (usize, usize),
65    data_offset: usize,
66}
67
68impl MmapArray2F32 {
69    /// Load a 2D f32 array from a raw binary file.
70    ///
71    /// The file format is:
72    /// - 8 bytes: nrows (i64 little-endian)
73    /// - 8 bytes: ncols (i64 little-endian)
74    /// - nrows * ncols * 4 bytes: f32 data (little-endian)
75    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        // Read shape from header
89        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    /// Get the shape of the array.
117    pub fn shape(&self) -> (usize, usize) {
118        self.shape
119    }
120
121    /// Get the number of rows.
122    pub fn nrows(&self) -> usize {
123        self.shape.0
124    }
125
126    /// Get the number of columns.
127    pub fn ncols(&self) -> usize {
128        self.shape.1
129    }
130
131    /// Get a view of a row.
132    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        // Safety: We've verified the bounds and alignment
137        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    /// Load a range of rows into an owned Array2.
144    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        // Safety: We've verified the bounds
151        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    /// Convert to an owned Array2 (loads all data into memory).
159    pub fn to_owned(&self) -> Array2<f32> {
160        self.load_rows(0, self.shape.0)
161    }
162}
163
164/// A memory-mapped array of u8 values.
165pub struct MmapArray2U8 {
166    _mmap: Mmap,
167    shape: (usize, usize),
168    data_offset: usize,
169}
170
171impl MmapArray2U8 {
172    /// Load a 2D u8 array from a raw binary file.
173    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    /// Get the shape of the array.
214    pub fn shape(&self) -> (usize, usize) {
215        self.shape
216    }
217
218    /// Get a view of the data as ArrayView2.
219    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    /// Load a range of rows into an owned Array2.
225    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    /// Convert to an owned Array2.
235    pub fn to_owned(&self) -> Array2<u8> {
236        self.load_rows(0, self.shape.0)
237    }
238}
239
240/// A memory-mapped array of i64 values.
241pub struct MmapArray1I64 {
242    _mmap: Mmap,
243    len: usize,
244    data_offset: usize,
245}
246
247impl MmapArray1I64 {
248    /// Load a 1D i64 array from a raw binary file.
249    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    /// Get the length of the array.
285    pub fn len(&self) -> usize {
286        self.len
287    }
288
289    /// Returns true if the array is empty.
290    pub fn is_empty(&self) -> bool {
291        self.len == 0
292    }
293
294    /// Get a value at an index.
295    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    /// Convert to an owned Array1.
302    pub fn to_owned(&self) -> Array1<i64> {
303        let bytes = &self._mmap[self.data_offset..self.data_offset + self.len * 8];
304
305        // Safety: We've verified the bounds
306        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
312/// Write an `Array2<f32>` to a raw binary file format.
313pub 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
343/// Write an `Array2<u8>` to a raw binary file format.
344pub 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
374/// Write an `Array1<i64>` to a raw binary file format.
375pub 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
401// ============================================================================
402// NPY Format Memory-Mapped Arrays
403// ============================================================================
404
405/// NPY file magic bytes
406const NPY_MAGIC: &[u8] = b"\x93NUMPY";
407
408/// Parse dtype from NPY header string (e.g., "<f2" for float16, "<f4" for float32)
409fn parse_dtype_from_header(header: &str) -> Result<String> {
410    // Find 'descr': '...'
411    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
427/// Detect NPY file dtype without loading the entire file
428pub 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    // Check magic
446    if &mmap[..6] != NPY_MAGIC {
447        return Err(Error::IndexLoad("Invalid NPY magic".into()));
448    }
449
450    let major_version = mmap[6];
451
452    // Read header length
453    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
480/// Convert a float16 NPY file to float32 in place
481pub fn convert_f16_to_f32_npy(path: &Path) -> Result<()> {
482    use half::f16;
483    use std::io::Read;
484
485    // Read the entire file
486    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    // Parse header to get shape
506    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    // Calculate total elements
511    let total_elements: usize = shape.iter().product();
512    let f16_data = &data[header_end..header_end + total_elements * 2];
513
514    // Convert f16 to f32
515    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    // Write new file with f32 dtype
523    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
543/// Convert an int64 NPY file to int32 in place
544pub fn convert_i64_to_i32_npy(path: &Path) -> Result<()> {
545    use std::io::Read;
546
547    // Read the entire file
548    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    // Parse header to get shape
568    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    // Convert i64 to i32
580    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    // Write new file with i32 dtype
588    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
602/// Re-save a u8 NPY file to ensure dtype descriptor is "|u1" (platform-independent)
603///
604/// Note: We can't use ndarray_npy::ReadNpyExt here because it doesn't accept "<u1"
605/// descriptor, so we manually read the raw data and resave with "|u1".
606pub fn normalize_u8_npy(path: &Path) -> Result<()> {
607    use std::io::Read;
608
609    // Read the entire file
610    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    // Parse header to get shape
630    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    // Re-write with explicit "|u1" dtype
645    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
659/// Parse NPY header and return (shape, data_offset, is_fortran_order)
660fn 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    // Check magic
670    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    // Read header length
678    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    // Parse header dict (simplified Python dict parsing)
709    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    // Extract shape from header like: {'descr': '<i8', 'fortran_order': False, 'shape': (12345,), }
713    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
719/// Parse shape tuple from NPY header string
720fn parse_shape_from_header(header: &str) -> Result<Vec<usize>> {
721    // Find 'shape': (...)
722    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    // Parse comma-separated numbers
737    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
751/// Memory-mapped NPY array for i64 values (used for codes).
752///
753/// This struct provides zero-copy access to 1D i64 arrays stored in NPY format.
754pub struct MmapNpyArray1I64 {
755    _mmap: Mmap,
756    len: usize,
757    data_offset: usize,
758}
759
760impl MmapNpyArray1I64 {
761    /// Create an empty instance backed by an anonymous mmap (no file).
762    ///
763    /// Used to release file-backed mmap handles before file operations on Windows,
764    /// where deleting or renaming a memory-mapped file causes OS error 1224.
765    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    /// Load a 1D i64 array from an NPY file.
778    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        // Verify file size
797        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    /// Get the length of the array.
814    pub fn len(&self) -> usize {
815        self.len
816    }
817
818    /// Returns true if the array is empty.
819    pub fn is_empty(&self) -> bool {
820        self.len == 0
821    }
822
823    /// Get a slice of the data as &[i64].
824    ///
825    /// Returns a `Vec<i64>` instead of &[i64] to handle unaligned data safely.
826    ///
827    /// # Safety
828    /// The caller must ensure start <= end <= len.
829    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    /// Get a value at an index.
841    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
848/// Memory-mapped NPY array for f32 values (used for centroids).
849///
850/// This struct provides zero-copy access to 2D f32 arrays stored in NPY format.
851/// Unlike loading into an owned `Array2<f32>`, this approach lets the OS manage
852/// paging, reducing resident memory usage for large centroid matrices.
853pub struct MmapNpyArray2F32 {
854    _mmap: Mmap,
855    shape: (usize, usize),
856    data_offset: usize,
857}
858
859impl MmapNpyArray2F32 {
860    /// Load a 2D f32 array from an NPY file.
861    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        // Verify file size (f32 = 4 bytes)
883        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    /// Get the shape of the array.
900    pub fn shape(&self) -> (usize, usize) {
901        self.shape
902    }
903
904    /// Get the number of rows.
905    pub fn nrows(&self) -> usize {
906        self.shape.0
907    }
908
909    /// Get the number of columns.
910    pub fn ncols(&self) -> usize {
911        self.shape.1
912    }
913
914    /// Get a view of the entire array as ArrayView2.
915    ///
916    /// This provides zero-copy access to the memory-mapped data.
917    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        // Safety: We've verified bounds and f32 is 4-byte aligned in NPY format
923        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    /// Get a view of a single row.
931    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        // Safety: We've verified bounds and alignment
936        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    /// Get a view of rows [start..end] as ArrayView2.
943    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        // Safety: We've verified bounds
950        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    /// Convert to an owned Array2 (loads all data into memory).
958    ///
959    /// Use this only when you need an owned copy; prefer `view()` for read-only access.
960    pub fn to_owned(&self) -> Array2<f32> {
961        self.view().to_owned()
962    }
963}
964
965/// Memory-mapped NPY array for u8 values (used for residuals).
966///
967/// This struct provides zero-copy access to 2D u8 arrays stored in NPY format.
968pub struct MmapNpyArray2U8 {
969    _mmap: Mmap,
970    shape: (usize, usize),
971    data_offset: usize,
972}
973
974impl MmapNpyArray2U8 {
975    /// Create an empty instance backed by an anonymous mmap (no file).
976    ///
977    /// Used to release file-backed mmap handles before file operations on Windows,
978    /// where deleting or renaming a memory-mapped file causes OS error 1224.
979    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    /// Load a 2D u8 array from an NPY file.
992    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        // Verify file size
1014        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    /// Get the shape of the array.
1031    pub fn shape(&self) -> (usize, usize) {
1032        self.shape
1033    }
1034
1035    /// Get the number of rows.
1036    pub fn nrows(&self) -> usize {
1037        self.shape.0
1038    }
1039
1040    /// Get the number of columns.
1041    pub fn ncols(&self) -> usize {
1042        self.shape.1
1043    }
1044
1045    /// Get a view of rows [start..end] as ArrayView2.
1046    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    /// Get a view of the entire array.
1056    pub fn view(&self) -> ArrayView2<'_, u8> {
1057        self.slice_rows(0, self.shape.0)
1058    }
1059
1060    /// Get a single row as a slice.
1061    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// ============================================================================
1069// Merged File Creation
1070// ============================================================================
1071
1072/// Manifest entry for tracking chunk files
1073#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1074pub struct ChunkManifestEntry {
1075    pub rows: usize,
1076    pub mtime: f64,
1077}
1078
1079/// Manifest for merged files, including metadata about the merge
1080#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
1081pub struct MergeManifest {
1082    /// Chunk information
1083    pub chunks: HashMap<String, ChunkManifestEntry>,
1084    /// Number of padding rows used in the merge
1085    #[serde(default)]
1086    pub padding_rows: usize,
1087    /// Number of chunks expected (for fast-path validation)
1088    #[serde(default)]
1089    pub num_chunks: usize,
1090    /// Mtime of metadata.json at merge time (detects any index modifications)
1091    #[serde(default)]
1092    pub metadata_mtime: f64,
1093    /// Total rows in the merged file (including padding)
1094    #[serde(default)]
1095    pub total_rows: usize,
1096    /// Number of columns (for 2D arrays like residuals)
1097    #[serde(default)]
1098    pub ncols: usize,
1099}
1100
1101/// Legacy manifest type (for backwards compatibility during migration)
1102pub type ChunkManifest = HashMap<String, ChunkManifestEntry>;
1103
1104/// Load manifest from disk if it exists
1105/// Handles both new MergeManifest format and legacy ChunkManifest format
1106fn load_merge_manifest(manifest_path: &Path) -> Option<MergeManifest> {
1107    if manifest_path.exists() {
1108        if let Ok(file) = File::open(manifest_path) {
1109            // Try to load as new format first
1110            let reader = BufReader::new(file);
1111            if let Ok(manifest) = serde_json::from_reader::<_, MergeManifest>(reader) {
1112                return Some(manifest);
1113            }
1114            // Try legacy format
1115            if let Ok(file) = File::open(manifest_path) {
1116                if let Ok(chunks) =
1117                    serde_json::from_reader::<_, ChunkManifest>(BufReader::new(file))
1118                {
1119                    // Convert legacy format - missing padding info means we need to regenerate
1120                    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
1135/// Save manifest to disk atomically (write to temp file, then rename)
1136fn save_merge_manifest(manifest_path: &Path, manifest: &MergeManifest) -> Result<()> {
1137    let temp_path = manifest_path.with_extension("manifest.json.tmp");
1138
1139    // Write to temp file
1140    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    // Sync to disk
1150    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    // Atomic rename
1157    fs::rename(&temp_path, manifest_path)
1158        .map_err(|e| Error::IndexLoad(format!("Failed to rename manifest: {}", e)))?;
1159
1160    Ok(())
1161}
1162
1163/// Get file modification time as f64 seconds since epoch
1164fn 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
1176/// Build the NPY header dict string and compute the total header size (magic + version + len + padded dict).
1177fn 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; // +1 for the trailing newline
1181    (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
1198/// Compute the NPY header size for a 1D array (without writing).
1199fn 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
1204/// Compute the NPY header size for a 2D array (without writing).
1205fn 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
1210/// Write an NPY header (shared implementation for 1D and 2D).
1211fn 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    // Write magic + version (v1.0)
1216    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    // Write header length (2 bytes for v1.0)
1224    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    // Write header
1230    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
1237/// Write NPY header for a 1D array
1238fn 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
1242/// Write NPY header for a 2D array
1243fn 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
1252/// Information about a chunk file for merging
1253struct ChunkInfo {
1254    path: std::path::PathBuf,
1255    filename: String,
1256    rows: usize,
1257    mtime: f64,
1258}
1259
1260/// Merge chunked codes NPY files into a single merged file.
1261///
1262/// Uses incremental persistence with manifest tracking to skip unchanged chunks.
1263/// Uses atomic writes to prevent corruption from interrupted writes.
1264/// Uses file-based locking to coordinate concurrent processes.
1265/// Returns the path to the merged file.
1266pub 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    // Fast path: if manifest exists with matching params, metadata.json hasn't changed,
1279    // and merged file exists with correct size, skip chunk scanning entirely.
1280    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    // Acquire exclusive lock to prevent concurrent merge operations.
1303    // This is critical for multi-process scenarios (e.g., multiple API workers).
1304    let _lock = FileLockGuard::acquire(&lock_path)?;
1305
1306    // After acquiring the lock, re-check if merge is still needed.
1307    // Another process might have completed the merge while we were waiting.
1308
1309    // Load previous manifest (re-read after acquiring lock)
1310    let old_manifest = load_merge_manifest(&manifest_path);
1311
1312    // Scan chunks and detect changes
1313    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            // Get shape by reading header only
1325            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                // Check if this chunk changed
1333                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    // Check if we need to rewrite:
1363    // 1. Merged file doesn't exist
1364    // 2. Chunks have changed
1365    // 3. Padding has changed (stored in manifest)
1366    // 4. Total rows don't match (safety check)
1367    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        // Write to temp file first (atomic write pattern)
1381        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        // Write header
1386        let header_size = write_npy_header_1d(&mut writer, final_rows, "<i8")?;
1387
1388        // Write chunk data
1389        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        // Write padding zeros
1400        for _ in 0..padding_rows {
1401            writer.write_all(&0i64.to_le_bytes())?;
1402        }
1403        written_rows += padding_rows;
1404
1405        // Flush and sync to disk
1406        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        // Verify file size before renaming
1416        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            // Clean up temp file and return error
1423            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        // Atomic rename (overwrites existing file)
1431        fs::rename(&temp_path, &merged_path)
1432            .map_err(|e| Error::IndexLoad(format!("Failed to rename merged file: {}", e)))?;
1433    } else {
1434        // Validate existing merged file before using it
1435        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            // NPY header is at least 64 bytes, data is final_rows * 8 bytes
1443            let min_expected_size = 64 + final_rows * 8;
1444            if file_size < min_expected_size {
1445                // File is corrupted, force regeneration by recursing with empty manifest
1446                let _ = fs::remove_file(&merged_path);
1447                let _ = fs::remove_file(&manifest_path);
1448                // Lock is held, so we can safely drop it before recursing
1449                drop(_lock);
1450                return merge_codes_chunks(index_path, num_chunks, padding_rows);
1451            }
1452        }
1453    }
1454
1455    // Build and save manifest with full metadata
1456    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, // Not used for 1D codes array
1471        num_chunks,
1472        metadata_mtime: current_metadata_mtime,
1473    };
1474    save_merge_manifest(&manifest_path, &new_manifest)?;
1475
1476    Ok(merged_path)
1477}
1478
1479/// Merge chunked residuals NPY files into a single merged file.
1480///
1481/// Uses atomic writes to prevent corruption from interrupted writes.
1482/// Uses file-based locking to coordinate concurrent processes.
1483pub 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    // Fast path: if manifest exists with matching params, metadata.json hasn't changed,
1496    // and merged file has correct size, skip chunk scanning entirely.
1497    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    // Acquire exclusive lock to prevent concurrent merge operations.
1520    // This is critical for multi-process scenarios (e.g., multiple API workers).
1521    let _lock = FileLockGuard::acquire(&lock_path)?;
1522
1523    // After acquiring the lock, re-check if merge is still needed.
1524    // Another process might have completed the merge while we were waiting.
1525
1526    // Load previous manifest (re-read after acquiring lock)
1527    let old_manifest = load_merge_manifest(&manifest_path);
1528
1529    // Scan chunks and detect changes
1530    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            // Get shape by reading header
1543            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    // Check if we need to rewrite:
1581    // 1. Merged file doesn't exist
1582    // 2. Chunks have changed
1583    // 3. Padding has changed
1584    // 4. Total rows or ncols don't match
1585    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        // Write to temp file first (atomic write pattern)
1606        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        // Write header
1611        let header_size = write_npy_header_2d(&mut writer, final_rows, ncols, "|u1")?;
1612
1613        // Write chunk data
1614        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        // Write padding zeros
1625        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        // Flush and sync to disk
1632        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        // Verify file size before renaming
1643        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            // Clean up temp file and return error
1650            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        // Atomic rename
1658        fs::rename(&temp_path, &merged_path)
1659            .map_err(|e| Error::IndexLoad(format!("Failed to rename merged residuals: {}", e)))?;
1660    } else {
1661        // Validate existing merged file before using it
1662        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            // NPY header is at least 64 bytes, data is final_rows * ncols bytes
1670            let min_expected_size = 64 + final_rows * ncols;
1671            if file_size < min_expected_size {
1672                // File is corrupted, force regeneration
1673                let _ = fs::remove_file(&merged_path);
1674                let _ = fs::remove_file(&manifest_path);
1675                // Lock is held, so we can safely drop it before recursing
1676                drop(_lock);
1677                return merge_residuals_chunks(index_path, num_chunks, padding_rows);
1678            }
1679        }
1680    }
1681
1682    // Build and save manifest with full metadata
1683    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
1706/// Clear merged files and manifests to force regeneration on next load.
1707///
1708/// This should be called after index updates to ensure the merged files
1709/// are regenerated with the latest data. The function silently ignores
1710/// missing files.
1711///
1712/// Acquires file locks to prevent racing with concurrent merge operations
1713/// in multi-process deployments.
1714pub fn clear_merged_files(index_path: &Path) -> Result<()> {
1715    // Acquire locks to prevent racing with ongoing merge operations.
1716    // This is important in multi-process scenarios where one process might
1717    // be loading (merging) while another is updating (clearing).
1718    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
1745// ============================================================================
1746// Fast-PLAID Compatibility Conversion
1747// ============================================================================
1748
1749/// Convert a fast-plaid index to next-plaid compatible format.
1750///
1751/// This function detects and converts:
1752/// - float16 → float32 for centroids, avg_residual, bucket_cutoffs, bucket_weights
1753/// - int64 → int32 for ivf_lengths
1754/// - `<u1` → `|u1` for residuals
1755///
1756/// Returns true if any conversion was performed, false if already compatible.
1757pub fn convert_fastplaid_to_nextplaid(index_path: &Path) -> Result<bool> {
1758    let mut converted = false;
1759
1760    // Float files to convert from f16 to f32
1761    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    // Convert ivf_lengths from i64 to i32
1781    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    // Normalize residual files to use "|u1" descriptor
1792    // fast-plaid uses "<u1" which ndarray_npy doesn't accept
1793    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        // Create a test file
1822        let mut file = NamedTempFile::new().unwrap();
1823
1824        // Write header (3 rows, 2 cols)
1825        file.write_all(&3i64.to_le_bytes()).unwrap();
1826        file.write_all(&2i64.to_le_bytes()).unwrap();
1827
1828        // Write data
1829        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        // Load and verify
1836        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        // Write header (4 elements)
1853        file.write_all(&4i64.to_le_bytes()).unwrap();
1854
1855        // Write data
1856        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        // Create test array
1878        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
1881        write_array2_f32(&array, path).unwrap();
1882
1883        // Read back
1884        let mmap = MmapArray2F32::from_raw_file(path).unwrap();
1885        let loaded = mmap.to_owned();
1886
1887        assert_eq!(array, loaded);
1888    }
1889}