laminar_core/state/
partition_key_schema.rs

1//! Persisted schema identity for partitioning ABI v1.
2//!
3//! This is deliberately narrower than a physical Arrow or state-artifact
4//! schema. SQL aliases are not routing identity, dictionary keys are hydrated,
5//! and field metadata is rejected by [`PartitionKeySchemaV1::try_new`].
6
7use arrow_schema::{DataType, FieldRef, IntervalUnit, TimeUnit};
8use sha2::{Digest, Sha256};
9
10use super::partition_key::{validate_key_type, PartitionKeyCodecError};
11
12/// Version of the partition-key schema descriptor byte format.
13pub const PARTITION_KEY_SCHEMA_VERSION: u16 = 1;
14/// Hard allocation bound for one partition-key schema descriptor.
15///
16/// The widest currently admitted schema is below 69 KiB: 256 nullable timestamp
17/// slots with 256-byte timezones plus framing. The 128-KiB ceiling leaves bounded
18/// encoder headroom but does not admit new types or wider parameters.
19pub const MAX_PARTITION_KEY_SCHEMA_BYTES: usize = 128 * 1024;
20
21const MAGIC: &[u8; 8] = b"LDBPKS\0\0";
22
23/// Persistable routing identity for an ordered partition-key schema.
24///
25/// The descriptor is opaque. Its bytes are a persisted ABI and may only change
26/// with an explicit partition-key schema version decision and new goldens.
27#[derive(Clone, Debug, Eq, PartialEq)]
28pub struct PartitionKeySchemaV1 {
29    bytes: Box<[u8]>,
30    sha256: [u8; 32],
31}
32
33impl PartitionKeySchemaV1 {
34    /// Validate and describe one ordered, resolved partition-key schema.
35    ///
36    /// Dictionary index widths are physical transport details and recursively
37    /// hydrate to the value type. Field names are SQL aliases and are omitted;
38    /// order and nullability remain identity.
39    ///
40    /// # Errors
41    /// Returns an error for an empty, over-wide, metadata-bearing, unsupported,
42    /// over-nested, or over-sized routing schema.
43    pub fn try_new(fields: &[FieldRef]) -> Result<Self, PartitionKeyCodecError> {
44        super::partition_key::validate_key_fields(fields)?;
45
46        let mut output = DescriptorWriter::new()?;
47        output.write_len(fields.len())?;
48        for (index, field) in fields.iter().enumerate() {
49            output.write_bool(field.is_nullable())?;
50            write_data_type(&mut output, field.data_type(), index)?;
51        }
52
53        let sha256 = Sha256::digest(&output.bytes).into();
54        Ok(Self {
55            bytes: output.bytes.into_boxed_slice(),
56            sha256,
57        })
58    }
59
60    /// Canonical, versioned routing-schema bytes.
61    #[must_use]
62    pub fn as_bytes(&self) -> &[u8] {
63        &self.bytes
64    }
65
66    /// SHA-256 of [`Self::as_bytes`].
67    #[must_use]
68    pub const fn sha256(&self) -> [u8; 32] {
69        self.sha256
70    }
71}
72
73struct DescriptorWriter {
74    bytes: Vec<u8>,
75}
76
77impl DescriptorWriter {
78    fn new() -> Result<Self, PartitionKeyCodecError> {
79        let mut writer = Self { bytes: Vec::new() };
80        writer.write_raw(MAGIC)?;
81        writer.write_raw(&PARTITION_KEY_SCHEMA_VERSION.to_be_bytes())?;
82        Ok(writer)
83    }
84
85    fn write_u8(&mut self, value: u8) -> Result<(), PartitionKeyCodecError> {
86        self.write_raw(&[value])
87    }
88
89    fn write_bool(&mut self, value: bool) -> Result<(), PartitionKeyCodecError> {
90        self.write_u8(u8::from(value))
91    }
92
93    fn write_len(&mut self, value: usize) -> Result<(), PartitionKeyCodecError> {
94        #[allow(clippy::cast_possible_truncation)]
95        self.write_raw(&(value as u64).to_be_bytes())
96    }
97
98    fn write_bytes(&mut self, value: &[u8]) -> Result<(), PartitionKeyCodecError> {
99        self.write_len(value.len())?;
100        self.write_raw(value)
101    }
102
103    fn write_raw(&mut self, value: &[u8]) -> Result<(), PartitionKeyCodecError> {
104        let next_len = self.bytes.len().checked_add(value.len()).ok_or(
105            PartitionKeyCodecError::KeySchemaDescriptorTooLarge {
106                limit: MAX_PARTITION_KEY_SCHEMA_BYTES,
107            },
108        )?;
109        if next_len > MAX_PARTITION_KEY_SCHEMA_BYTES {
110            return Err(PartitionKeyCodecError::KeySchemaDescriptorTooLarge {
111                limit: MAX_PARTITION_KEY_SCHEMA_BYTES,
112            });
113        }
114        if next_len > self.bytes.capacity() {
115            let target_capacity = self
116                .bytes
117                .capacity()
118                .max(256)
119                .saturating_mul(2)
120                .max(next_len)
121                .min(MAX_PARTITION_KEY_SCHEMA_BYTES);
122            self.bytes
123                .try_reserve_exact(target_capacity - self.bytes.len())
124                .map_err(|_| PartitionKeyCodecError::KeySchemaDescriptorTooLarge {
125                    limit: MAX_PARTITION_KEY_SCHEMA_BYTES,
126                })?;
127        }
128        if self.bytes.capacity() > MAX_PARTITION_KEY_SCHEMA_BYTES {
129            return Err(PartitionKeyCodecError::KeySchemaDescriptorTooLarge {
130                limit: MAX_PARTITION_KEY_SCHEMA_BYTES,
131            });
132        }
133        self.bytes.extend_from_slice(value);
134        Ok(())
135    }
136}
137
138// These tags are LaminarDB's persisted ABI, not Arrow enum discriminants. The
139// exhaustive match makes an Arrow upgrade stop here until the routing policy is
140// reviewed. `validate_key_type` is shared with the row codec and runs first.
141fn write_data_type(
142    output: &mut DescriptorWriter,
143    data_type: &DataType,
144    index: usize,
145) -> Result<(), PartitionKeyCodecError> {
146    validate_key_type(data_type, index, 0)?;
147    match data_type {
148        DataType::Null => output.write_u8(0),
149        DataType::Boolean => output.write_u8(1),
150        DataType::Int8 => output.write_u8(2),
151        DataType::Int16 => output.write_u8(3),
152        DataType::Int32 => output.write_u8(4),
153        DataType::Int64 => output.write_u8(5),
154        DataType::UInt8 => output.write_u8(6),
155        DataType::UInt16 => output.write_u8(7),
156        DataType::UInt32 => output.write_u8(8),
157        DataType::UInt64 => output.write_u8(9),
158        DataType::Timestamp(unit, timezone) => {
159            output.write_u8(10)?;
160            write_time_unit(output, unit)?;
161            match timezone {
162                None => output.write_u8(0),
163                Some(timezone) => {
164                    output.write_u8(1)?;
165                    output.write_bytes(timezone.as_bytes())
166                }
167            }
168        }
169        DataType::Date32 => output.write_u8(11),
170        DataType::Date64 => output.write_u8(12),
171        DataType::Time32(unit) => {
172            output.write_u8(13)?;
173            write_time_unit(output, unit)
174        }
175        DataType::Time64(unit) => {
176            output.write_u8(14)?;
177            write_time_unit(output, unit)
178        }
179        DataType::Duration(unit) => {
180            output.write_u8(15)?;
181            write_time_unit(output, unit)
182        }
183        DataType::Interval(unit) => {
184            output.write_u8(16)?;
185            write_interval_unit(output, unit)
186        }
187        DataType::Binary => output.write_u8(17),
188        DataType::FixedSizeBinary(width) => {
189            output.write_u8(18)?;
190            output.write_raw(&width.to_be_bytes())
191        }
192        DataType::LargeBinary => output.write_u8(19),
193        DataType::BinaryView => output.write_u8(20),
194        DataType::Utf8 => output.write_u8(21),
195        DataType::LargeUtf8 => output.write_u8(22),
196        DataType::Utf8View => output.write_u8(23),
197        DataType::Dictionary(_, values) => write_data_type(output, values, index),
198        DataType::Decimal32(precision, scale) => {
199            output.write_u8(24)?;
200            output.write_u8(*precision)?;
201            output.write_raw(&scale.to_be_bytes())
202        }
203        DataType::Decimal64(precision, scale) => {
204            output.write_u8(25)?;
205            output.write_u8(*precision)?;
206            output.write_raw(&scale.to_be_bytes())
207        }
208        DataType::Decimal128(precision, scale) => {
209            output.write_u8(26)?;
210            output.write_u8(*precision)?;
211            output.write_raw(&scale.to_be_bytes())
212        }
213        DataType::Decimal256(precision, scale) => {
214            output.write_u8(27)?;
215            output.write_u8(*precision)?;
216            output.write_raw(&scale.to_be_bytes())
217        }
218        DataType::Float16
219        | DataType::Float32
220        | DataType::Float64
221        | DataType::List(_)
222        | DataType::ListView(_)
223        | DataType::FixedSizeList(_, _)
224        | DataType::LargeList(_)
225        | DataType::LargeListView(_)
226        | DataType::Struct(_)
227        | DataType::Union(_, _)
228        | DataType::Map(_, _)
229        | DataType::RunEndEncoded(_, _) => Err(PartitionKeyCodecError::UnsupportedKeyTypeFamily {
230            index,
231            family: "non-scalar",
232        }),
233    }
234}
235
236fn write_time_unit(
237    output: &mut DescriptorWriter,
238    unit: &TimeUnit,
239) -> Result<(), PartitionKeyCodecError> {
240    output.write_u8(match unit {
241        TimeUnit::Second => 0,
242        TimeUnit::Millisecond => 1,
243        TimeUnit::Microsecond => 2,
244        TimeUnit::Nanosecond => 3,
245    })
246}
247
248fn write_interval_unit(
249    output: &mut DescriptorWriter,
250    unit: &IntervalUnit,
251) -> Result<(), PartitionKeyCodecError> {
252    output.write_u8(match unit {
253        IntervalUnit::YearMonth => 0,
254        IntervalUnit::DayTime => 1,
255        IntervalUnit::MonthDayNano => 2,
256    })
257}
laminar_core/state/partition_key_schema.rs

laminar_core/state/
partition_key_schema.rs
