laminar_core/state/partition_key/
mod.rs

1//! Canonical typed partition-key encoding for partitioning ABI v1.
2//!
3//! The encoded bytes are Arrow row-format bytes produced with ascending,
4//! nulls-first sort fields. They are schema-bound rather than self-describing.
5//! The separately constructed [`PartitionKeySchemaV1`] freezes the persisted
6//! schema descriptor. It is not constructed by routing and adds no work to the
7//! record-processing hot path.
8//! Changing the admitted type set, row encoding, hash, or vnode mapping requires
9//! a partitioning ABI decision.
10
11use std::num::NonZeroU32;
12
13use arrow_array::types::{
14    validate_decimal_precision_and_scale, Decimal128Type, Decimal256Type, Decimal32Type,
15    Decimal64Type, DecimalType,
16};
17use arrow_array::ArrayRef;
18use arrow_row::{RowConverter, Rows, SortField};
19use arrow_schema::{DataType, FieldRef, TimeUnit};
20
21pub use super::partition_key_schema::PartitionKeySchemaV1;
22use super::vnode::key_hash;
23
24// These are structural ABI safety ceilings, not deployment SLOs. A numerical
25// workload profile may admit less but cannot raise them without an ABI review.
26/// Maximum number of columns retained by one ABI-v1 row-converter plan.
27pub(crate) const MAX_PARTITION_KEY_COLUMNS: usize = 256;
28/// Maximum recursive dictionary wrappers inspected while hydrating one key type.
29pub(crate) const MAX_PARTITION_KEY_NESTING: usize = 32;
30/// Maximum timezone parameter bytes copied into one routing descriptor.
31pub(crate) const MAX_PARTITION_KEY_TIMEZONE_BYTES: usize = 256;
32
33/// Failure to construct the ABI-v1 typed-key encoder.
34#[derive(Debug, thiserror::Error)]
35pub enum PartitionKeyCodecError {
36    /// A keyed partitioning codec requires at least one logical key column.
37    #[error("partition key schema is empty")]
38    EmptyKeySchema,
39    /// The composite key is too wide for bounded plan-time construction.
40    #[error("partition key has {count} columns; hard limit is {limit}")]
41    TooManyKeyColumns {
42        /// Requested key-column count.
43        count: usize,
44        /// Hard key-column limit.
45        limit: usize,
46    },
47    /// Arrow extension metadata has no routing semantics in partitioning ABI v1.
48    #[error("partition key column {index} has metadata unsupported by ABI v1")]
49    UnsupportedKeyMetadata {
50        /// Zero-based position in the composite key.
51        index: usize,
52    },
53    /// Recursive dictionary hydration is bounded before Arrow row construction.
54    #[error("partition key column {index} exceeds dictionary nesting depth {limit}")]
55    KeyTypeNestingTooDeep {
56        /// Zero-based position in the composite key.
57        index: usize,
58        /// Hard dictionary nesting limit.
59        limit: usize,
60    },
61    /// A variable type parameter crossed its partition ABI resource bound.
62    #[error(
63        "partition key column {index} {parameter} occupies {bytes} bytes; hard limit is {limit}"
64    )]
65    KeyTypeParameterTooLarge {
66        /// Zero-based position in the composite key.
67        index: usize,
68        /// Bounded type parameter.
69        parameter: &'static str,
70        /// Observed parameter bytes.
71        bytes: usize,
72        /// Hard parameter limit.
73        limit: usize,
74    },
75    /// The resolved key type has no frozen ABI-v1 equality and encoding contract.
76    #[error("partition key column {index} has unsupported ABI-v1 type {data_type}")]
77    UnsupportedKeyType {
78        /// Zero-based position in the composite key.
79        index: usize,
80        /// Rejected Arrow type.
81        data_type: DataType,
82    },
83    /// A recursive Arrow family is rejected without cloning attacker-shaped schema trees.
84    #[error("partition key column {index} has unsupported ABI-v1 type family {family}")]
85    UnsupportedKeyTypeFamily {
86        /// Zero-based position in the composite key.
87        index: usize,
88        /// Stable rejected family label.
89        family: &'static str,
90    },
91    /// Arrow rejected the otherwise admitted row layout.
92    #[error("partition key Arrow row encoding: {0}")]
93    Arrow(#[from] arrow_schema::ArrowError),
94    /// The opaque routing-schema descriptor crossed its hard allocation bound.
95    #[error("partition key schema descriptor exceeds {limit} bytes")]
96    KeySchemaDescriptorTooLarge {
97        /// Hard encoded descriptor limit.
98        limit: usize,
99    },
100}
101
102/// Vectorized encoder for the exact typed-key representation covered by
103/// [`super::PARTITIONING_ABI_VERSION`] 2.
104///
105/// Floats, nested values, and run-end encoding remain fail-closed. Dictionary
106/// indices are physical representation: supported dictionaries encode their
107/// hydrated scalar values.
108#[derive(Debug)]
109pub struct PartitionKeyCodecV1 {
110    converter: RowConverter,
111}
112
113/// Internal one-pass builder used where key-column lookup and validation must
114/// preserve their original request order.
115pub(crate) struct PartitionKeyCodecV1Builder {
116    fields: Vec<SortField>,
117}
118
119impl PartitionKeyCodecV1Builder {
120    pub(crate) fn with_capacity(capacity: usize) -> Self {
121        Self {
122            fields: Vec::with_capacity(capacity.min(MAX_PARTITION_KEY_COLUMNS)),
123        }
124    }
125
126    pub(crate) fn push(&mut self, data_type: DataType) -> Result<(), PartitionKeyCodecError> {
127        let index = self.fields.len();
128        if index >= MAX_PARTITION_KEY_COLUMNS {
129            return Err(PartitionKeyCodecError::TooManyKeyColumns {
130                count: index + 1,
131                limit: MAX_PARTITION_KEY_COLUMNS,
132            });
133        }
134        validate_key_type(&data_type, index, 0)?;
135        self.fields.push(SortField::new(data_type));
136        Ok(())
137    }
138
139    pub(crate) fn finish(self) -> Result<PartitionKeyCodecV1, PartitionKeyCodecError> {
140        if self.fields.is_empty() {
141            return Err(PartitionKeyCodecError::EmptyKeySchema);
142        }
143        Ok(PartitionKeyCodecV1 {
144            converter: RowConverter::new(self.fields)?,
145        })
146    }
147}
148
149impl PartitionKeyCodecV1 {
150    /// Construct a codec for one resolved, ordered composite-key schema.
151    ///
152    /// # Errors
153    /// Returns the first unsupported key type or an Arrow row-layout error.
154    pub fn try_new(
155        data_types: impl IntoIterator<Item = DataType>,
156    ) -> Result<Self, PartitionKeyCodecError> {
157        let data_types = data_types.into_iter();
158        let mut builder = PartitionKeyCodecV1Builder::with_capacity(data_types.size_hint().0);
159        for data_type in data_types {
160            builder.push(data_type)?;
161        }
162        builder.finish()
163    }
164
165    /// Encode all key columns in one vectorized Arrow-row pass.
166    ///
167    /// The returned rows borrow no input buffers and can be hashed or copied
168    /// into operator state. Column types and arity must match construction.
169    ///
170    /// # Errors
171    /// Returns an Arrow error for a mismatched schema or column length.
172    pub fn encode_columns(&self, columns: &[ArrayRef]) -> Result<Rows, arrow_schema::ArrowError> {
173        self.converter.convert_columns(columns)
174    }
175
176    /// Full ABI-v1 hash of one encoded key.
177    #[must_use]
178    pub fn hash_encoded(encoded: &[u8]) -> u64 {
179        key_hash(encoded)
180    }
181
182    /// Map one encoded key into a nonempty vnode space.
183    #[must_use]
184    pub fn vnode_for_encoded(encoded: &[u8], vnode_count: NonZeroU32) -> u32 {
185        #[allow(clippy::cast_possible_truncation)]
186        {
187            (Self::hash_encoded(encoded) % u64::from(vnode_count.get())) as u32
188        }
189    }
190}
191
192pub(crate) fn validate_key_fields(fields: &[FieldRef]) -> Result<(), PartitionKeyCodecError> {
193    if fields.is_empty() {
194        return Err(PartitionKeyCodecError::EmptyKeySchema);
195    }
196    if fields.len() > MAX_PARTITION_KEY_COLUMNS {
197        return Err(PartitionKeyCodecError::TooManyKeyColumns {
198            count: fields.len(),
199            limit: MAX_PARTITION_KEY_COLUMNS,
200        });
201    }
202    for (index, field) in fields.iter().enumerate() {
203        if !field.metadata().is_empty() {
204            return Err(PartitionKeyCodecError::UnsupportedKeyMetadata { index });
205        }
206        validate_key_type(field.data_type(), index, 0)?;
207    }
208    Ok(())
209}
210
211pub(crate) fn validate_key_type(
212    data_type: &DataType,
213    index: usize,
214    depth: usize,
215) -> Result<(), PartitionKeyCodecError> {
216    if depth > MAX_PARTITION_KEY_NESTING {
217        return Err(PartitionKeyCodecError::KeyTypeNestingTooDeep {
218            index,
219            limit: MAX_PARTITION_KEY_NESTING,
220        });
221    }
222
223    match data_type {
224        DataType::Null
225        | DataType::Boolean
226        | DataType::Int8
227        | DataType::Int16
228        | DataType::Int32
229        | DataType::Int64
230        | DataType::UInt8
231        | DataType::UInt16
232        | DataType::UInt32
233        | DataType::UInt64
234        | DataType::Date32
235        | DataType::Date64
236        | DataType::Duration(_)
237        | DataType::Interval(_)
238        | DataType::Binary
239        | DataType::LargeBinary
240        | DataType::BinaryView
241        | DataType::Utf8
242        | DataType::LargeUtf8
243        | DataType::Utf8View => Ok(()),
244        DataType::Timestamp(_, timezone) => {
245            let bytes = timezone.as_ref().map_or(0, |value| value.len());
246            if bytes <= MAX_PARTITION_KEY_TIMEZONE_BYTES {
247                Ok(())
248            } else {
249                Err(PartitionKeyCodecError::KeyTypeParameterTooLarge {
250                    index,
251                    parameter: "timezone",
252                    bytes,
253                    limit: MAX_PARTITION_KEY_TIMEZONE_BYTES,
254                })
255            }
256        }
257        DataType::FixedSizeBinary(length) if *length >= 0 => Ok(()),
258        DataType::Time32(unit) => {
259            if matches!(unit, TimeUnit::Second | TimeUnit::Millisecond) {
260                Ok(())
261            } else {
262                Err(PartitionKeyCodecError::UnsupportedKeyType {
263                    index,
264                    data_type: data_type.clone(),
265                })
266            }
267        }
268        DataType::Time64(unit) => {
269            if matches!(unit, TimeUnit::Microsecond | TimeUnit::Nanosecond) {
270                Ok(())
271            } else {
272                Err(PartitionKeyCodecError::UnsupportedKeyType {
273                    index,
274                    data_type: data_type.clone(),
275                })
276            }
277        }
278        DataType::Decimal32(precision, scale)
279            if valid_decimal::<Decimal32Type>(*precision, *scale) =>
280        {
281            Ok(())
282        }
283        DataType::Decimal64(precision, scale)
284            if valid_decimal::<Decimal64Type>(*precision, *scale) =>
285        {
286            Ok(())
287        }
288        DataType::Decimal128(precision, scale)
289            if valid_decimal::<Decimal128Type>(*precision, *scale) =>
290        {
291            Ok(())
292        }
293        DataType::Decimal256(precision, scale)
294            if valid_decimal::<Decimal256Type>(*precision, *scale) =>
295        {
296            Ok(())
297        }
298        DataType::Dictionary(indices, values)
299            if matches!(
300                indices.as_ref(),
301                DataType::Int8
302                    | DataType::Int16
303                    | DataType::Int32
304                    | DataType::Int64
305                    | DataType::UInt8
306                    | DataType::UInt16
307                    | DataType::UInt32
308                    | DataType::UInt64
309            ) =>
310        {
311            validate_key_type(values, index, depth + 1)
312        }
313        DataType::Float16
314        | DataType::Float32
315        | DataType::Float64
316        | DataType::FixedSizeBinary(_)
317        | DataType::Decimal32(_, _)
318        | DataType::Decimal64(_, _)
319        | DataType::Decimal128(_, _)
320        | DataType::Decimal256(_, _) => Err(PartitionKeyCodecError::UnsupportedKeyType {
321            index,
322            data_type: data_type.clone(),
323        }),
324        DataType::List(_) => unsupported_family(index, "list"),
325        DataType::ListView(_) => unsupported_family(index, "list-view"),
326        DataType::FixedSizeList(_, _) => unsupported_family(index, "fixed-size-list"),
327        DataType::LargeList(_) => unsupported_family(index, "large-list"),
328        DataType::LargeListView(_) => unsupported_family(index, "large-list-view"),
329        DataType::Struct(_) => unsupported_family(index, "struct"),
330        DataType::Union(_, _) => unsupported_family(index, "union"),
331        DataType::Map(_, _) => unsupported_family(index, "map"),
332        DataType::RunEndEncoded(_, _) => unsupported_family(index, "run-end-encoded"),
333        DataType::Dictionary(_, _) => unsupported_family(index, "dictionary-index"),
334    }
335}
336
337fn unsupported_family(index: usize, family: &'static str) -> Result<(), PartitionKeyCodecError> {
338    Err(PartitionKeyCodecError::UnsupportedKeyTypeFamily { index, family })
339}
340
341fn valid_decimal<T: DecimalType>(precision: u8, scale: i8) -> bool {
342    validate_decimal_precision_and_scale::<T>(precision, scale).is_ok()
343}
344
345#[cfg(test)]
346mod tests;
laminar_core/state/partition_key/mod.rs

laminar_core/state/partition_key/
mod.rs
