ailake_file/
reader.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2use crate::footer::Precision;
3use ailake_core::{AilakeError, AilakeResult, Centroid, VectorMetric};
4use ailake_index::{AnyIndex, HnswIndex, IvfPqSerializer, MmapLoader};
5use ailake_parquet::ParquetVectorReader;
6use arrow_array::RecordBatch;
7use bytes::Bytes;
8
9use crate::footer::{AilakeHeader, DistanceMetric, FLAG_INDEX_IVF_PQ, HEADER_SIZE};
10
11pub struct AilakeFileReader {
12    bytes: Bytes,
13    vector_column: String,
14    #[allow(dead_code)]
15    dim: u32,
16}
17
18impl AilakeFileReader {
19    pub fn new(bytes: Bytes, vector_column: &str, dim: u32) -> Self {
20        Self {
21            bytes,
22            vector_column: vector_column.to_string(),
23            dim,
24        }
25    }
26
27    /// Returns the absolute byte offset of the primary AILK section.
28    /// Reads `ailake.footer_offset` from the Parquet footer key-value metadata.
29    pub fn ailk_offset(&self) -> AilakeResult<u64> {
30        let reader = ParquetVectorReader::new(self.bytes.clone(), &self.vector_column);
31        let val = reader
32            .kv_metadata("ailake.footer_offset")?
33            .ok_or(AilakeError::NotAnAilakeFile)?;
34        val.parse::<u64>().map_err(|_| AilakeError::NotAnAilakeFile)
35    }
36
37    /// Returns the absolute byte offset of the AILK section for a named vector column.
38    ///
39    /// For additional columns tries `ailake.{column}.footer_offset` first,
40    /// then falls back to `ailake.footer_offset` (primary / single-column files).
41    pub fn ailk_offset_for_column(&self, column: &str) -> AilakeResult<u64> {
42        let reader = ParquetVectorReader::new(self.bytes.clone(), column);
43        let col_key = format!("ailake.{column}.footer_offset");
44        if let Some(val) = reader.kv_metadata(&col_key)? {
45            return val.parse::<u64>().map_err(|_| AilakeError::NotAnAilakeFile);
46        }
47        let val = reader
48            .kv_metadata("ailake.footer_offset")?
49            .ok_or(AilakeError::NotAnAilakeFile)?;
50        val.parse::<u64>().map_err(|_| AilakeError::NotAnAilakeFile)
51    }
52
53    /// Returns true if the file contains an embedded AILK section.
54    pub fn is_ailake_file(&self) -> bool {
55        self.ailk_offset().is_ok()
56    }
57
58    /// Parse the 64-byte AI-Lake header from the embedded AILK section.
59    pub fn read_header(&self) -> AilakeResult<AilakeHeader> {
60        self.read_header_at_offset(self.ailk_offset()?)
61    }
62
63    /// Parse the 64-byte AI-Lake header for a named vector column.
64    ///
65    /// Uses `ailake.{column}.footer_offset` for extra columns and falls back to
66    /// `ailake.footer_offset` for the primary column (single-column files).
67    pub fn read_header_for_column(&self, column: &str) -> AilakeResult<AilakeHeader> {
68        self.read_header_at_offset(self.ailk_offset_for_column(column)?)
69    }
70
71    fn read_header_at_offset(&self, offset: u64) -> AilakeResult<AilakeHeader> {
72        let offset = offset as usize;
73        if offset + HEADER_SIZE > self.bytes.len() {
74            return Err(AilakeError::NotAnAilakeFile);
75        }
76        let header_bytes: &[u8; HEADER_SIZE] = self.bytes[offset..offset + HEADER_SIZE]
77            .try_into()
78            .map_err(|_| AilakeError::NotAnAilakeFile)?;
79        AilakeHeader::from_bytes(header_bytes)
80    }
81
82    /// Read centroid + radius from the AILK section.
83    pub fn get_centroid(&self) -> AilakeResult<Centroid> {
84        let ailk_start = self.ailk_offset()? as usize;
85        let header = self.read_header()?;
86        let centroid_start = ailk_start + header.centroid_offset as usize;
87        let centroid_end = centroid_start + header.centroid_len as usize;
88
89        if centroid_end > self.bytes.len() {
90            return Err(AilakeError::NotAnAilakeFile);
91        }
92
93        let centroid_data = &self.bytes[centroid_start..centroid_end];
94        let dim = header.dim as usize;
95        let expected_len = dim * 4 + 4;
96        if centroid_data.len() != expected_len {
97            return Err(AilakeError::InvalidCentroidLength {
98                expected_dim: header.dim,
99                actual: centroid_data.len(),
100            });
101        }
102
103        let values: Vec<f32> = centroid_data[..dim * 4]
104            .chunks_exact(4)
105            .map(|b| f32::from_le_bytes(b.try_into().unwrap()))
106            .collect();
107        let radius = f32::from_le_bytes(centroid_data[dim * 4..].try_into().unwrap());
108        let metric = distance_metric_to_vector_metric(header.distance_metric);
109
110        Ok(Centroid {
111            values,
112            radius,
113            metric,
114        })
115    }
116
117    /// Load the HNSW index from the primary AILK section.
118    pub fn load_index(&self) -> AilakeResult<HnswIndex> {
119        self.load_index_for_column(&self.vector_column.clone())
120    }
121
122    /// Load the HNSW index for a specific vector column.
123    ///
124    /// Works for both single-column files (falls back to primary AILK) and
125    /// multi-column files written with `AilakeFileWriter::write_multi`.
126    pub fn load_index_for_column(&self, column: &str) -> AilakeResult<HnswIndex> {
127        let ailk_start = self.ailk_offset_for_column(column)? as usize;
128
129        if ailk_start + HEADER_SIZE > self.bytes.len() {
130            return Err(AilakeError::NotAnAilakeFile);
131        }
132        let header_bytes: &[u8; HEADER_SIZE] = self.bytes[ailk_start..ailk_start + HEADER_SIZE]
133            .try_into()
134            .map_err(|_| AilakeError::NotAnAilakeFile)?;
135        let header = AilakeHeader::from_bytes(header_bytes)?;
136
137        let hnsw_start = ailk_start + header.hnsw_offset as usize;
138        let hnsw_end = hnsw_start + header.hnsw_len as usize;
139
140        if hnsw_end > self.bytes.len() {
141            return Err(AilakeError::NotAnAilakeFile);
142        }
143        let mut idx = MmapLoader::from_bytes(&self.bytes[hnsw_start..hnsw_end])?;
144        if header.precision == Precision::F16 {
145            idx.quantize_to_f16();
146        }
147        Ok(idx)
148    }
149
150    /// Load primary index as `AnyIndex`, dispatching on header flags.
151    pub fn load_any_index(&self) -> AilakeResult<AnyIndex> {
152        self.load_any_index_for_column(&self.vector_column.clone())
153    }
154
155    /// Load index for a specific vector column as `AnyIndex`.
156    pub fn load_any_index_for_column(&self, column: &str) -> AilakeResult<AnyIndex> {
157        let ailk_start = self.ailk_offset_for_column(column)? as usize;
158
159        if ailk_start + HEADER_SIZE > self.bytes.len() {
160            return Err(AilakeError::NotAnAilakeFile);
161        }
162        let header_bytes: &[u8; HEADER_SIZE] = self.bytes[ailk_start..ailk_start + HEADER_SIZE]
163            .try_into()
164            .map_err(|_| AilakeError::NotAnAilakeFile)?;
165        let header = AilakeHeader::from_bytes(header_bytes)?;
166
167        let index_start = ailk_start + header.hnsw_offset as usize;
168        let index_end = index_start + header.hnsw_len as usize;
169
170        if index_end > self.bytes.len() {
171            return Err(AilakeError::NotAnAilakeFile);
172        }
173        let index_bytes = &self.bytes[index_start..index_end];
174
175        if header.flags & FLAG_INDEX_IVF_PQ != 0 {
176            let idx = IvfPqSerializer::from_bytes(index_bytes)?;
177            Ok(AnyIndex::IvfPq(idx))
178        } else {
179            let mut idx = MmapLoader::from_bytes(index_bytes)?;
180            if header.precision == Precision::F16 {
181                idx.quantize_to_f16();
182            }
183            Ok(AnyIndex::Hnsw(idx))
184        }
185    }
186
187    /// Read the Parquet section (tabular data + decoded embeddings).
188    /// The full file is valid Parquet; the AILK section is invisible to standard readers.
189    pub fn read_parquet(&self) -> AilakeResult<(RecordBatch, Vec<Vec<f32>>)> {
190        let reader = ParquetVectorReader::new(self.bytes.clone(), &self.vector_column);
191        reader.read_all()
192    }
193
194    /// Verify the positional invariant: Parquet record_count == HNSW node_count.
195    pub fn verify_integrity(&self) -> AilakeResult<()> {
196        let header = self.read_header()?;
197        let index = self.load_index()?;
198        let reader = ParquetVectorReader::new(self.bytes.clone(), &self.vector_column);
199        let parquet_count = reader.record_count()?;
200
201        if parquet_count != index.node_count() {
202            return Err(AilakeError::RowCountMismatch {
203                parquet: parquet_count,
204                hnsw: index.node_count(),
205            });
206        }
207        if parquet_count != header.record_count {
208            return Err(AilakeError::RowCountMismatch {
209                parquet: parquet_count,
210                hnsw: header.record_count,
211            });
212        }
213        Ok(())
214    }
215}
216
217fn distance_metric_to_vector_metric(dm: DistanceMetric) -> VectorMetric {
218    match dm {
219        DistanceMetric::Cosine => VectorMetric::Cosine,
220        DistanceMetric::Euclidean => VectorMetric::Euclidean,
221        DistanceMetric::DotProduct => VectorMetric::DotProduct,
222        DistanceMetric::NormalizedCosine => VectorMetric::NormalizedCosine,
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229    use crate::writer::AilakeFileWriter;
230    use ailake_core::{VectorMetric, VectorPrecision, VectorStoragePolicy};
231    use arrow_array::{Int32Array, RecordBatch};
232    use arrow_schema::{DataType, Field, Schema};
233    use std::sync::Arc;
234
235    fn make_policy(dim: u32) -> VectorStoragePolicy {
236        VectorStoragePolicy {
237            column_name: "embedding".to_string(),
238            dim,
239            metric: VectorMetric::Cosine,
240            precision: VectorPrecision::F16,
241            pq: None,
242            keep_raw_for_reranking: true,
243            pre_normalize: false,
244            hnsw_m: None,
245            hnsw_ef_construction: None,
246            ivf_residual: false,
247            embedding_model: None,
248            modality: None,
249        }
250    }
251
252    fn write_file(rows: usize, dim: u32) -> Bytes {
253        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
254        let ids: Vec<i32> = (0..rows as i32).collect();
255        let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(ids))]).unwrap();
256        let embs: Vec<Vec<f32>> = (0..rows)
257            .map(|i| {
258                let mut v = vec![0.0f32; dim as usize];
259                v[i % dim as usize] = 1.0;
260                v
261            })
262            .collect();
263        AilakeFileWriter::new(make_policy(dim))
264            .write(&batch, &embs)
265            .unwrap()
266    }
267
268    #[test]
269    fn is_ailake_file() {
270        let file = write_file(3, 4);
271        let reader = AilakeFileReader::new(file, "embedding", 4);
272        assert!(reader.is_ailake_file());
273    }
274
275    #[test]
276    fn integrity_check_passes() {
277        let file = write_file(10, 8);
278        let reader = AilakeFileReader::new(file, "embedding", 8);
279        reader.verify_integrity().unwrap();
280    }
281
282    #[test]
283    fn centroid_has_correct_dim() {
284        let file = write_file(5, 4);
285        let reader = AilakeFileReader::new(file, "embedding", 4);
286        let centroid = reader.get_centroid().unwrap();
287        assert_eq!(centroid.values.len(), 4);
288    }
289
290    #[test]
291    fn search_finds_nearest() {
292        let dim = 4u32;
293        let file = write_file(4, dim);
294        let reader = AilakeFileReader::new(file, "embedding", dim);
295        let index = reader.load_index().unwrap();
296        let query = vec![1.0f32, 0.0, 0.0, 0.0];
297        let results = index.search(&query, 1, 50);
298        assert_eq!(results.len(), 1);
299        assert_eq!(results[0].0, ailake_core::RowId::new(0));
300    }
301
302    #[test]
303    fn parquet_read_returns_tabular_data() {
304        let file = write_file(3, 4);
305        let reader = AilakeFileReader::new(file, "embedding", 4);
306        let (batch, embs) = reader.read_parquet().unwrap();
307        assert_eq!(batch.num_rows(), 3);
308        assert_eq!(embs.len(), 3);
309    }
310}
ailake_file/reader.rs

ailake_file/
reader.rs
