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

ailake_file/
reader.rs
