ailake_file/
footer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// Binary layout for the AI-Lake footer extension.
3// See docs/specs/FILE_FORMAT.md for field-by-field spec.
4
5use ailake_core::{AilakeError, AilakeResult, VectorMetric, VectorPrecision};
6
7pub const AILAKE_MAGIC: [u8; 4] = *b"AILK";
8pub const AILAKE_FORMAT_VERSION: u16 = 1;
9pub const TRAILER_SIZE: usize = 24;
10pub const HEADER_SIZE: usize = 64;
11
12/// `flags` bit 0 = 1: IVF-PQ index. Default (flags = 0): HNSW index.
13pub const FLAG_INDEX_IVF_PQ: u16 = 0x0001;
14
15/// Magic for the AILK_FTS section (separate from the vector AILK section).
16pub const AILK_FTS_MAGIC: [u8; 4] = *b"AFTS";
17/// AILK_FTS section header size: magic(4) + version(2) + reserved(2) + blob_len(8) = 16 bytes.
18pub const AILK_FTS_HEADER_SIZE: usize = 16;
19/// Parquet KV key storing the absolute byte offset of the AILK_FTS section.
20pub const KV_FTS_OFFSET: &str = "ailake.fts_offset";
21
22#[repr(u8)]
23#[derive(Debug, Clone, Copy, PartialEq)]
24pub enum Precision {
25    F32 = 0,
26    F16 = 1,
27    I8 = 2,
28    Binary = 3,
29}
30
31#[repr(u8)]
32#[derive(Debug, Clone, Copy, PartialEq)]
33pub enum DistanceMetric {
34    Cosine = 0,
35    Euclidean = 1,
36    DotProduct = 2,
37    NormalizedCosine = 3,
38}
39
40impl From<VectorPrecision> for Precision {
41    fn from(p: VectorPrecision) -> Self {
42        match p {
43            VectorPrecision::F32 => Precision::F32,
44            VectorPrecision::F16 => Precision::F16,
45            VectorPrecision::I8 => Precision::I8,
46            VectorPrecision::Binary => Precision::Binary,
47        }
48    }
49}
50
51impl From<VectorMetric> for DistanceMetric {
52    fn from(m: VectorMetric) -> Self {
53        match m {
54            VectorMetric::Cosine => DistanceMetric::Cosine,
55            VectorMetric::Euclidean => DistanceMetric::Euclidean,
56            VectorMetric::DotProduct => DistanceMetric::DotProduct,
57            VectorMetric::NormalizedCosine => DistanceMetric::NormalizedCosine,
58        }
59    }
60}
61
62impl TryFrom<u8> for Precision {
63    type Error = AilakeError;
64    fn try_from(v: u8) -> AilakeResult<Self> {
65        match v {
66            0 => Ok(Precision::F32),
67            1 => Ok(Precision::F16),
68            2 => Ok(Precision::I8),
69            3 => Ok(Precision::Binary),
70            _ => Err(AilakeError::InvalidArgument(format!(
71                "invalid precision byte: {v} (valid: 0=F32, 1=F16, 2=I8, 3=Binary)"
72            ))),
73        }
74    }
75}
76
77impl TryFrom<u8> for DistanceMetric {
78    type Error = AilakeError;
79    fn try_from(v: u8) -> AilakeResult<Self> {
80        match v {
81            0 => Ok(DistanceMetric::Cosine),
82            1 => Ok(DistanceMetric::Euclidean),
83            2 => Ok(DistanceMetric::DotProduct),
84            3 => Ok(DistanceMetric::NormalizedCosine),
85            _ => Err(AilakeError::InvalidArgument(format!(
86                "invalid distance metric byte: {v} (valid: 0=Cosine, 1=Euclidean, 2=DotProduct, 3=NormalizedCosine)"
87            ))),
88        }
89    }
90}
91
92/// 64-byte header at the start of the AI-Lake footer extension.
93#[derive(Debug, Clone)]
94pub struct AilakeHeader {
95    pub format_version: u16,
96    pub flags: u16,
97    pub dim: u32,
98    pub precision: Precision,
99    pub distance_metric: DistanceMetric,
100    pub record_count: u64,
101    pub centroid_offset: u64,
102    pub centroid_len: u64,
103    pub hnsw_offset: u64,
104    pub hnsw_len: u64,
105}
106
107impl AilakeHeader {
108    pub fn to_bytes(&self) -> [u8; HEADER_SIZE] {
109        let mut b = [0u8; HEADER_SIZE];
110        b[0..4].copy_from_slice(&AILAKE_MAGIC);
111        b[4..6].copy_from_slice(&self.format_version.to_le_bytes());
112        b[6..8].copy_from_slice(&self.flags.to_le_bytes());
113        b[8..12].copy_from_slice(&self.dim.to_le_bytes());
114        b[12] = self.precision as u8;
115        b[13] = self.distance_metric as u8;
116        // b[14..16] reserved = 0
117        b[16..24].copy_from_slice(&self.record_count.to_le_bytes());
118        b[24..32].copy_from_slice(&self.centroid_offset.to_le_bytes());
119        b[32..40].copy_from_slice(&self.centroid_len.to_le_bytes());
120        b[40..48].copy_from_slice(&self.hnsw_offset.to_le_bytes());
121        b[48..56].copy_from_slice(&self.hnsw_len.to_le_bytes());
122        // b[56..64] reserved = 0
123        b
124    }
125
126    pub fn from_bytes(b: &[u8; HEADER_SIZE]) -> AilakeResult<Self> {
127        if b[0..4] != AILAKE_MAGIC {
128            return Err(AilakeError::InvalidAilakeMagic([b[0], b[1], b[2], b[3]]));
129        }
130        let format_version = u16::from_le_bytes([b[4], b[5]]);
131        if format_version != AILAKE_FORMAT_VERSION {
132            return Err(AilakeError::UnsupportedFormatVersion(format_version));
133        }
134        Ok(AilakeHeader {
135            format_version,
136            flags: u16::from_le_bytes([b[6], b[7]]),
137            dim: u32::from_le_bytes([b[8], b[9], b[10], b[11]]),
138            precision: Precision::try_from(b[12])?,
139            distance_metric: DistanceMetric::try_from(b[13])?,
140            record_count: u64::from_le_bytes([
141                b[16], b[17], b[18], b[19], b[20], b[21], b[22], b[23],
142            ]),
143            centroid_offset: u64::from_le_bytes([
144                b[24], b[25], b[26], b[27], b[28], b[29], b[30], b[31],
145            ]),
146            centroid_len: u64::from_le_bytes([
147                b[32], b[33], b[34], b[35], b[36], b[37], b[38], b[39],
148            ]),
149            hnsw_offset: u64::from_le_bytes([
150                b[40], b[41], b[42], b[43], b[44], b[45], b[46], b[47],
151            ]),
152            hnsw_len: u64::from_le_bytes([b[48], b[49], b[50], b[51], b[52], b[53], b[54], b[55]]),
153        })
154    }
155}
156
157/// 24-byte trailer — the last bytes of every AI-Lake file.
158#[derive(Debug, Clone)]
159pub struct AilakeTrailer {
160    pub footer_offset: u64,
161    pub footer_len: u64,
162    pub format_version: u16,
163    pub flags: u16,
164}
165
166impl AilakeTrailer {
167    pub fn to_bytes(&self) -> [u8; TRAILER_SIZE] {
168        let mut b = [0u8; TRAILER_SIZE];
169        b[0..8].copy_from_slice(&self.footer_offset.to_le_bytes());
170        b[8..16].copy_from_slice(&self.footer_len.to_le_bytes());
171        b[16..18].copy_from_slice(&self.format_version.to_le_bytes());
172        b[18..20].copy_from_slice(&self.flags.to_le_bytes());
173        b[20..24].copy_from_slice(&AILAKE_MAGIC);
174        b
175    }
176
177    pub fn from_bytes(b: &[u8; TRAILER_SIZE]) -> AilakeResult<Self> {
178        if b[20..24] != AILAKE_MAGIC {
179            return Err(AilakeError::InvalidAilakeMagic([
180                b[20], b[21], b[22], b[23],
181            ]));
182        }
183        Ok(AilakeTrailer {
184            footer_offset: u64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]),
185            footer_len: u64::from_le_bytes([b[8], b[9], b[10], b[11], b[12], b[13], b[14], b[15]]),
186            format_version: u16::from_le_bytes([b[16], b[17]]),
187            flags: u16::from_le_bytes([b[18], b[19]]),
188        })
189    }
190}
191
192/// Returns the byte offset in `buf` where the Parquet footer thrift starts.
193///
194/// Parquet tail layout: `[...footer_thrift...][footer_len: u32 LE][PAR1: 4 bytes]`
195///
196/// Used by both the writer (to know where to splice AILK sections) and the reader
197/// (to locate the AILK trailer for KV-less bootstrap).
198pub fn parquet_footer_start(buf: &[u8]) -> AilakeResult<usize> {
199    let len = buf.len();
200    if len < 8 {
201        return Err(AilakeError::Parquet("file too small".into()));
202    }
203    if &buf[len - 4..] != b"PAR1" {
204        return Err(AilakeError::Parquet("missing PAR1 footer magic".into()));
205    }
206    let footer_thrift_len = u32::from_le_bytes(buf[len - 8..len - 4].try_into().unwrap()) as usize;
207    len.checked_sub(8 + footer_thrift_len)
208        .ok_or_else(|| AilakeError::Parquet("footer length overflow".into()))
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[test]
216    fn header_roundtrip() {
217        let h = AilakeHeader {
218            format_version: 1,
219            flags: 0,
220            dim: 1536,
221            precision: Precision::F16,
222            distance_metric: DistanceMetric::Cosine,
223            record_count: 50_000,
224            centroid_offset: 64,
225            centroid_len: 1536 * 4 + 4,
226            hnsw_offset: 64 + 1536 * 4 + 4,
227            hnsw_len: 4_194_304,
228        };
229        let bytes = h.to_bytes();
230        let h2 = AilakeHeader::from_bytes(&bytes).unwrap();
231        assert_eq!(h2.dim, 1536);
232        assert_eq!(h2.precision, Precision::F16);
233        assert_eq!(h2.distance_metric, DistanceMetric::Cosine);
234        assert_eq!(h2.record_count, 50_000);
235    }
236
237    #[test]
238    fn trailer_roundtrip() {
239        let t = AilakeTrailer {
240            footer_offset: 12_582_912,
241            footer_len: 4_194_304,
242            format_version: 1,
243            flags: 0,
244        };
245        let bytes = t.to_bytes();
246        let t2 = AilakeTrailer::from_bytes(&bytes).unwrap();
247        assert_eq!(t2.footer_offset, 12_582_912);
248        assert_eq!(&bytes[20..24], b"AILK");
249    }
250
251    #[test]
252    fn invalid_magic_rejected() {
253        let mut bytes = [0u8; HEADER_SIZE];
254        bytes[0..4].copy_from_slice(b"BLAH");
255        assert!(AilakeHeader::from_bytes(&bytes).is_err());
256    }
257}
ailake_file/footer.rs

ailake_file/
footer.rs
