array_format/
layout.rs

1//! Array layout definitions and array metadata.
2//!
3//! Each array stored in the file has an [`ArrayMeta`] entry in the footer
4//! describing its type, dimensions, and how its data is laid out across blocks.
5
6use rkyv::{Archive, Deserialize, Serialize};
7
8use crate::address::ChunkAddress;
9use crate::dtype::DType;
10
11/// A scalar value for a per-array key-value attribute.
12#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
13pub enum AttributeValue {
14    /// Boolean value.
15    Bool(bool),
16    /// Signed 8-bit integer.
17    Int8(i8),
18    /// Signed 16-bit integer.
19    Int16(i16),
20    /// Signed 32-bit integer.
21    Int32(i32),
22    /// Signed 64-bit integer.
23    Int64(i64),
24    /// Unsigned 8-bit integer.
25    UInt8(u8),
26    /// Unsigned 16-bit integer.
27    UInt16(u16),
28    /// Unsigned 32-bit integer.
29    UInt32(u32),
30    /// Unsigned 64-bit integer.
31    UInt64(u64),
32    /// 32-bit floating point.
33    Float32(f32),
34    /// 64-bit floating point.
35    Float64(f64),
36    /// UTF-8 string.
37    String(String),
38    /// Variable-length binary data.
39    Binary(Vec<u8>),
40    /// List of boolean values.
41    BoolList(Vec<bool>),
42    /// List of signed 8-bit integers.
43    Int8List(Vec<i8>),
44    /// List of signed 16-bit integers.
45    Int16List(Vec<i16>),
46    /// List of signed 32-bit integers.
47    Int32List(Vec<i32>),
48    /// List of signed 64-bit integers.
49    Int64List(Vec<i64>),
50    /// List of unsigned 8-bit integers.
51    UInt8List(Vec<u8>),
52    /// List of unsigned 16-bit integers.
53    UInt16List(Vec<u16>),
54    /// List of unsigned 32-bit integers.
55    UInt32List(Vec<u32>),
56    /// List of unsigned 64-bit integers.
57    UInt64List(Vec<u64>),
58    /// List of 32-bit floating point values.
59    Float32List(Vec<f32>),
60    /// List of 64-bit floating point values.
61    Float64List(Vec<f64>),
62    /// List of UTF-8 strings.
63    StringList(Vec<String>),
64    /// List of variable-length binary values.
65    BinaryList(Vec<Vec<u8>>),
66}
67
68impl PartialEq for AttributeValue {
69    fn eq(&self, other: &Self) -> bool {
70        match (self, other) {
71            (Self::Bool(a), Self::Bool(b)) => a == b,
72            (Self::Int8(a), Self::Int8(b)) => a == b,
73            (Self::Int16(a), Self::Int16(b)) => a == b,
74            (Self::Int32(a), Self::Int32(b)) => a == b,
75            (Self::Int64(a), Self::Int64(b)) => a == b,
76            (Self::UInt8(a), Self::UInt8(b)) => a == b,
77            (Self::UInt16(a), Self::UInt16(b)) => a == b,
78            (Self::UInt32(a), Self::UInt32(b)) => a == b,
79            (Self::UInt64(a), Self::UInt64(b)) => a == b,
80            (Self::Float32(a), Self::Float32(b)) => a.to_bits() == b.to_bits(),
81            (Self::Float64(a), Self::Float64(b)) => a.to_bits() == b.to_bits(),
82            (Self::String(a), Self::String(b)) => a == b,
83            (Self::Binary(a), Self::Binary(b)) => a == b,
84            (Self::BoolList(a), Self::BoolList(b)) => a == b,
85            (Self::Int8List(a), Self::Int8List(b)) => a == b,
86            (Self::Int16List(a), Self::Int16List(b)) => a == b,
87            (Self::Int32List(a), Self::Int32List(b)) => a == b,
88            (Self::Int64List(a), Self::Int64List(b)) => a == b,
89            (Self::UInt8List(a), Self::UInt8List(b)) => a == b,
90            (Self::UInt16List(a), Self::UInt16List(b)) => a == b,
91            (Self::UInt32List(a), Self::UInt32List(b)) => a == b,
92            (Self::UInt64List(a), Self::UInt64List(b)) => a == b,
93            // Compare floats by bit pattern so NaN == NaN, matching the scalar
94            // float variants and keeping value interning deduplication stable.
95            (Self::Float32List(a), Self::Float32List(b)) => {
96                a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
97            }
98            (Self::Float64List(a), Self::Float64List(b)) => {
99                a.len() == b.len() && a.iter().zip(b).all(|(x, y)| x.to_bits() == y.to_bits())
100            }
101            (Self::StringList(a), Self::StringList(b)) => a == b,
102            (Self::BinaryList(a), Self::BinaryList(b)) => a == b,
103            _ => false,
104        }
105    }
106}
107
108/// Controls the integer width used for attribute key/value dictionary indices.
109#[derive(Debug, Clone, Copy, PartialEq, Default, Archive, Serialize, Deserialize)]
110pub enum AttrIndexKind {
111    /// 16-bit dictionary indices (the default).
112    #[default]
113    U16,
114    /// 32-bit dictionary indices.
115    U32,
116    /// 64-bit dictionary indices.
117    U64,
118}
119
120/// Per-array attribute entries as pairs of dictionary indices.
121///
122/// The variant determines the storage width per index. All entries are sorted
123/// by key index to allow O(log n) binary search. Arrays in the same file may
124/// use different variants; each is self-describing.
125#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize)]
126pub enum Attributes {
127    /// `(key_index, value_index)` pairs stored as 16-bit indices.
128    U16(Vec<(u16, u16)>),
129    /// `(key_index, value_index)` pairs stored as 32-bit indices.
130    U32(Vec<(u32, u32)>),
131    /// `(key_index, value_index)` pairs stored as 64-bit indices.
132    U64(Vec<(u64, u64)>),
133}
134
135impl Attributes {
136    /// Returns an empty attribute set for the given index kind.
137    pub fn empty(kind: AttrIndexKind) -> Self {
138        match kind {
139            AttrIndexKind::U16 => Self::U16(Vec::new()),
140            AttrIndexKind::U32 => Self::U32(Vec::new()),
141            AttrIndexKind::U64 => Self::U64(Vec::new()),
142        }
143    }
144
145    /// Returns the value index for `key_idx`, or `None` if absent.
146    pub fn get(&self, key_idx: usize) -> Option<usize> {
147        match self {
148            Self::U16(v) => v
149                .binary_search_by_key(&(key_idx as u16), |(k, _)| *k)
150                .ok()
151                .map(|pos| v[pos].1 as usize),
152            Self::U32(v) => v
153                .binary_search_by_key(&(key_idx as u32), |(k, _)| *k)
154                .ok()
155                .map(|pos| v[pos].1 as usize),
156            Self::U64(v) => v
157                .binary_search_by_key(&(key_idx as u64), |(k, _)| *k)
158                .ok()
159                .map(|pos| v[pos].1 as usize),
160        }
161    }
162
163    /// Inserts or replaces the value index for `key_idx`, keeping entries sorted.
164    pub fn upsert(&mut self, key_idx: usize, val_idx: usize) {
165        match self {
166            Self::U16(v) => {
167                let (ki, vi) = (key_idx as u16, val_idx as u16);
168                match v.binary_search_by_key(&ki, |(k, _)| *k) {
169                    Ok(pos) => v[pos].1 = vi,
170                    Err(pos) => v.insert(pos, (ki, vi)),
171                }
172            }
173            Self::U32(v) => {
174                let (ki, vi) = (key_idx as u32, val_idx as u32);
175                match v.binary_search_by_key(&ki, |(k, _)| *k) {
176                    Ok(pos) => v[pos].1 = vi,
177                    Err(pos) => v.insert(pos, (ki, vi)),
178                }
179            }
180            Self::U64(v) => {
181                let (ki, vi) = (key_idx as u64, val_idx as u64);
182                match v.binary_search_by_key(&ki, |(k, _)| *k) {
183                    Ok(pos) => v[pos].1 = vi,
184                    Err(pos) => v.insert(pos, (ki, vi)),
185                }
186            }
187        }
188    }
189
190    /// Iterates over all `(key_index, value_index)` pairs as `usize`.
191    pub fn iter_entries(&self) -> Box<dyn Iterator<Item = (usize, usize)> + '_> {
192        match self {
193            Self::U16(v) => Box::new(v.iter().map(|&(k, v)| (k as usize, v as usize))),
194            Self::U32(v) => Box::new(v.iter().map(|&(k, v)| (k as usize, v as usize))),
195            Self::U64(v) => Box::new(v.iter().map(|&(k, v)| (k as usize, v as usize))),
196        }
197    }
198
199    /// Returns the maximum index value that fits in this variant.
200    pub fn max_index(&self) -> usize {
201        match self {
202            Self::U16(_) => u16::MAX as usize,
203            Self::U32(_) => u32::MAX as usize,
204            Self::U64(_) => usize::MAX,
205        }
206    }
207}
208
209/// Describes how an array's data is distributed across blocks.
210///
211/// Carries the array's global shape and dimension names so callers always
212/// know the full n-dimensional structure, independent of how the data is
213/// stored internally.
214#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize)]
215pub struct ArrayLayout {
216    /// Size of the array in each dimension (number of elements).
217    pub shape: Vec<u32>,
218    /// Name of each dimension (e.g. `["x", "y", "time"]`).
219    pub dimension_names: Vec<String>,
220    /// Storage strategy for the array data.
221    pub storage: StorageLayout,
222}
223
224/// A single entry in the chunk table: a coordinate vector paired with its
225/// storage address.
226///
227/// `coord` identifies the chunk within the array's grid (one index per
228/// dimension). `address` locates the chunk's bytes within a compressed block.
229#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize)]
230pub struct ChunkEntry {
231    /// Chunk coordinate within the array's grid (one index per dimension).
232    pub coord: Vec<u32>,
233    /// Location of the chunk's bytes within a compressed block.
234    pub address: ChunkAddress,
235}
236
237/// Describes the physical storage strategy for array data.
238///
239/// All arrays are stored as a grid of chunks. A "flat" array is simply one
240/// where `chunk_shape == shape`, yielding a single chunk at the origin.
241#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize)]
242pub struct StorageLayout {
243    /// The shape of each chunk in number of elements per dimension.
244    pub chunk_shape: Vec<u32>,
245    /// Mapping from chunk coordinates to their storage addresses.
246    ///
247    /// A layer sidecar carries only the chunks that changed; unchanged chunks
248    /// fall through to lower layers.
249    pub chunks: Vec<ChunkEntry>,
250}
251
252/// A scalar fill value for an array.
253///
254/// Represents the value used for unwritten or missing elements. Supports
255/// numeric types and strings; complex types (Binary, List, FixedSizeList)
256/// are not representable here.
257#[derive(Debug, Clone, Archive, Serialize, Deserialize)]
258pub enum FillValue {
259    /// Boolean fill value.
260    Bool(bool),
261    /// Signed integer fill value (for the signed integer dtypes).
262    Int(i64),
263    /// Unsigned integer fill value (for the unsigned integer dtypes).
264    UInt(u64),
265    /// Floating-point fill value (for `Float32`/`Float64`).
266    Float(f64),
267    /// String fill value.
268    String(String),
269    /// Fill value for [`DType::TimestampNs`] arrays — interpreted as `i64`
270    /// nanoseconds since the Unix epoch.
271    TimestampNs(i64),
272}
273
274impl PartialEq for FillValue {
275    fn eq(&self, other: &Self) -> bool {
276        match (self, other) {
277            (Self::Bool(a), Self::Bool(b)) => a == b,
278            (Self::Int(a), Self::Int(b)) => a == b,
279            (Self::UInt(a), Self::UInt(b)) => a == b,
280            // Compare by bit pattern so NaN == NaN.
281            (Self::Float(a), Self::Float(b)) => a.to_bits() == b.to_bits(),
282            (Self::String(a), Self::String(b)) => a == b,
283            (Self::TimestampNs(a), Self::TimestampNs(b)) => a == b,
284            _ => false,
285        }
286    }
287}
288
289/// Metadata for a single array stored in the file.
290///
291/// Stored in the footer's array table. Contains the schema, layout
292/// information, and a logical deletion flag.
293#[derive(Debug, Clone, PartialEq, Archive, Serialize, Deserialize)]
294pub struct ArrayMeta {
295    /// Unique name of the array within the file.
296    pub name: String,
297    /// Element data type.
298    pub dtype: DType,
299    /// Shape, dimension names, and physical chunk layout.
300    pub layout: ArrayLayout,
301    /// Fill value for unwritten or missing elements. `None` means zero/empty.
302    pub fill_value: Option<FillValue>,
303    /// Whether this array has been logically deleted.
304    ///
305    /// Deleted arrays are ignored during reads and removed during compaction.
306    pub deleted: bool,
307    /// User-defined key-value attributes.
308    ///
309    /// Each entry is `(key_index, value_index)` referencing `Footer::attr_keys`
310    /// and `Footer::attr_values`. Entries are sorted by key index for O(log n)
311    /// lookup. The `Attributes` variant controls the storage width per index.
312    pub attributes: Attributes,
313}
314
315impl ArrayLayout {
316    /// Looks up the [`ChunkAddress`] for a given chunk coordinate.
317    pub fn get_chunk(&self, coord: &[u32]) -> Option<&ChunkAddress> {
318        self.storage
319            .chunks
320            .iter()
321            .find(|e| e.coord.as_slice() == coord)
322            .map(|e| &e.address)
323    }
324
325    /// Returns all [`ChunkAddress`]es for this layout.
326    pub fn all_addresses(&self) -> Vec<&ChunkAddress> {
327        self.storage.chunks.iter().map(|e| &e.address).collect()
328    }
329}
330
331#[cfg(test)]
332mod tests {
333    use super::*;
334    use crate::address::BlockId;
335
336    fn sample_addr(block: u32, offset: u32, size: u32) -> ChunkAddress {
337        ChunkAddress {
338            block_id: BlockId(block),
339            offset,
340            size,
341        }
342    }
343
344    fn chunk(coord: Vec<u32>, block: u32, offset: u32, size: u32) -> ChunkEntry {
345        ChunkEntry {
346            coord,
347            address: sample_addr(block, offset, size),
348        }
349    }
350
351    #[test]
352    fn attribute_value_binary_and_list_equality() {
353        assert_eq!(
354            AttributeValue::Binary(vec![1, 2, 3]),
355            AttributeValue::Binary(vec![1, 2, 3])
356        );
357        assert_ne!(
358            AttributeValue::Binary(vec![1, 2, 3]),
359            AttributeValue::Binary(vec![1, 2])
360        );
361        assert_eq!(
362            AttributeValue::Int32List(vec![1, 2, 3]),
363            AttributeValue::Int32List(vec![1, 2, 3])
364        );
365        assert_eq!(
366            AttributeValue::StringList(vec!["a".into(), "b".into()]),
367            AttributeValue::StringList(vec!["a".into(), "b".into()])
368        );
369        assert_eq!(
370            AttributeValue::BinaryList(vec![vec![0], vec![1, 2]]),
371            AttributeValue::BinaryList(vec![vec![0], vec![1, 2]])
372        );
373        // Different variants are never equal.
374        assert_ne!(
375            AttributeValue::Int32List(vec![1]),
376            AttributeValue::Int64List(vec![1])
377        );
378    }
379
380    #[test]
381    fn attribute_value_float_list_compares_by_bits() {
382        // NaN == NaN by bit pattern, matching the scalar float variants so
383        // interning deduplication stays stable.
384        assert_eq!(
385            AttributeValue::Float64List(vec![f64::NAN, 1.0]),
386            AttributeValue::Float64List(vec![f64::NAN, 1.0])
387        );
388        // +0.0 and -0.0 have different bit patterns, so they are not equal.
389        assert_ne!(
390            AttributeValue::Float32List(vec![0.0]),
391            AttributeValue::Float32List(vec![-0.0])
392        );
393        assert_ne!(
394            AttributeValue::Float64List(vec![1.0, 2.0]),
395            AttributeValue::Float64List(vec![1.0])
396        );
397    }
398
399    #[test]
400    fn flat_layout_addresses() {
401        let layout = ArrayLayout {
402            shape: vec![100],
403            dimension_names: vec!["x".into()],
404            storage: StorageLayout {
405                chunk_shape: vec![100],
406                chunks: vec![chunk(vec![0], 0, 0, 400)],
407            },
408        };
409        assert_eq!(layout.all_addresses().len(), 1);
410        assert!(layout.get_chunk(&[0]).is_some());
411        assert!(layout.get_chunk(&[1]).is_none());
412    }
413
414    #[test]
415    fn chunked_layout_lookup() {
416        let layout = ArrayLayout {
417            shape: vec![128, 128],
418            dimension_names: vec!["x".into(), "y".into()],
419            storage: StorageLayout {
420                chunk_shape: vec![64, 64],
421                chunks: vec![
422                    chunk(vec![0, 0], 3, 0, 1000),
423                    chunk(vec![0, 1], 7, 2000, 1000),
424                ],
425            },
426        };
427        let addr = layout.get_chunk(&[0, 1]).unwrap();
428        assert_eq!(addr.block_id, BlockId(7));
429        assert_eq!(addr.offset, 2000);
430
431        assert!(layout.get_chunk(&[1, 0]).is_none());
432    }
433
434    #[test]
435    fn chunked_all_addresses() {
436        let layout = ArrayLayout {
437            shape: vec![30],
438            dimension_names: vec!["t".into()],
439            storage: StorageLayout {
440                chunk_shape: vec![10],
441                chunks: vec![
442                    chunk(vec![0], 0, 0, 500),
443                    chunk(vec![1], 0, 500, 500),
444                    chunk(vec![2], 1, 0, 500),
445                ],
446            },
447        };
448        assert_eq!(layout.all_addresses().len(), 3);
449    }
450}
array_format/layout.rs

array_format/
layout.rs
