Skip to main content

laser_wire/
arrow.rs

1use crate::error::InvalidError;
2use crate::schema::SchemaFingerprint;
3use crate::validate::Validate;
4use serde::{Deserialize, Serialize};
5
6pub const ARROW_IPC_CONTRACT_VERSION: u32 = 1;
7pub const ARROW_IPC_MEDIA_TYPE: &str = "application/vnd.apache.arrow.stream";
8pub const MAX_ARROW_IPC_MESSAGE_BYTES: u64 = 8 * 1024 * 1024;
9pub const MAX_ARROW_IPC_FIELDS: u32 = 4_096;
10pub const MAX_ARROW_IPC_BATCHES: u32 = 64;
11pub const MAX_ARROW_IPC_ROWS: u64 = 1_000_000;
12pub const MAX_ARROW_IPC_DICTIONARIES: u32 = 4_096;
13pub const MAX_ARROW_DECIMAL_BITS: u16 = 128;
14
15/// Validated facts about one self-contained Arrow IPC stream carried by one
16/// Iggy message. The raw IPC stream remains the message body.
17#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
18pub struct ArrowIpcMessageMetadata {
19    pub contract_version: u32,
20    pub schema_fingerprint: SchemaFingerprint,
21    pub encoded_bytes: u64,
22    pub field_count: u32,
23    pub record_batch_count: u32,
24    pub row_count: u64,
25    pub dictionary_count: u32,
26}
27
28impl Validate for ArrowIpcMessageMetadata {
29    fn validate(&self) -> Result<(), InvalidError> {
30        if self.contract_version != ARROW_IPC_CONTRACT_VERSION {
31            return Err(InvalidError::new(format!(
32                "Arrow IPC contract version must be {ARROW_IPC_CONTRACT_VERSION}, got {}",
33                self.contract_version
34            )));
35        }
36        self.schema_fingerprint.validate()?;
37        if self.encoded_bytes == 0 || self.encoded_bytes > MAX_ARROW_IPC_MESSAGE_BYTES {
38            return Err(InvalidError::new(format!(
39                "Arrow IPC message bytes must be in 1..={MAX_ARROW_IPC_MESSAGE_BYTES}"
40            )));
41        }
42        if self.field_count == 0 || self.field_count > MAX_ARROW_IPC_FIELDS {
43            return Err(InvalidError::new(format!(
44                "Arrow IPC field count must be in 1..={MAX_ARROW_IPC_FIELDS}"
45            )));
46        }
47        if self.record_batch_count == 0 || self.record_batch_count > MAX_ARROW_IPC_BATCHES {
48            return Err(InvalidError::new(format!(
49                "Arrow IPC batch count must be in 1..={MAX_ARROW_IPC_BATCHES}"
50            )));
51        }
52        if self.row_count > MAX_ARROW_IPC_ROWS {
53            return Err(InvalidError::new(format!(
54                "Arrow IPC row count exceeds cap {MAX_ARROW_IPC_ROWS}"
55            )));
56        }
57        if self.dictionary_count > MAX_ARROW_IPC_DICTIONARIES {
58            return Err(InvalidError::new(format!(
59                "Arrow IPC dictionary count exceeds cap {MAX_ARROW_IPC_DICTIONARIES}"
60            )));
61        }
62        Ok(())
63    }
64}
65
66/// Frozen producer-side Arrow input policy. One message contains one complete
67/// stream with its schema, dictionaries, and record batches.
68#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
69pub struct ArrowIpcPolicy {
70    pub contract_version: u32,
71    pub stream_format_only: bool,
72    pub self_contained: bool,
73    pub dictionary_deltas: bool,
74    pub replacement_dictionaries: bool,
75    pub timestamp_unit: ArrowTimestampUnit,
76    pub max_decimal_bits: u16,
77    pub unions: bool,
78    pub extension_types: bool,
79}
80
81impl Default for ArrowIpcPolicy {
82    fn default() -> Self {
83        Self {
84            contract_version: ARROW_IPC_CONTRACT_VERSION,
85            stream_format_only: true,
86            self_contained: true,
87            dictionary_deltas: false,
88            replacement_dictionaries: false,
89            timestamp_unit: ArrowTimestampUnit::Microsecond,
90            max_decimal_bits: MAX_ARROW_DECIMAL_BITS,
91            unions: false,
92            extension_types: false,
93        }
94    }
95}
96
97impl Validate for ArrowIpcPolicy {
98    fn validate(&self) -> Result<(), InvalidError> {
99        if self.contract_version != ARROW_IPC_CONTRACT_VERSION {
100            return Err(InvalidError::new(format!(
101                "Arrow IPC policy version must be {ARROW_IPC_CONTRACT_VERSION}, got {}",
102                self.contract_version
103            )));
104        }
105        if !self.stream_format_only || !self.self_contained {
106            return Err(InvalidError::new(
107                "Arrow IPC input must be one self-contained stream per message",
108            ));
109        }
110        if self.dictionary_deltas || self.replacement_dictionaries {
111            return Err(InvalidError::new(
112                "Arrow IPC dictionary deltas and replacements are not supported",
113            ));
114        }
115        if self.max_decimal_bits != MAX_ARROW_DECIMAL_BITS {
116            return Err(InvalidError::new(format!(
117                "Arrow IPC decimal width must be {MAX_ARROW_DECIMAL_BITS} bits"
118            )));
119        }
120        if self.unions || self.extension_types {
121            return Err(InvalidError::new(
122                "Arrow IPC unions and extension types are not supported",
123            ));
124        }
125        Ok(())
126    }
127}
128
129#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
130#[serde(rename_all = "snake_case")]
131#[non_exhaustive]
132pub enum ArrowTimestampUnit {
133    Microsecond,
134}
135
136/// Stable reason returned when an Arrow stream cannot enter the logical schema gate.
137#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
138#[serde(rename_all = "snake_case")]
139#[non_exhaustive]
140pub enum ArrowIpcRejectionCode {
141    FileFormat,
142    MissingSchema,
143    MissingDictionary,
144    DictionaryDelta,
145    DictionaryReplacement,
146    Union,
147    ExtensionType,
148    TimestampUnit,
149    DecimalWidth,
150    SchemaFingerprint,
151    FieldLimit,
152    BatchLimit,
153    RowLimit,
154    ByteLimit,
155    MalformedStream,
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn given_the_frozen_policy_when_validated_then_should_accept_only_the_supported_arrow_subset() {
164        ArrowIpcPolicy::default()
165            .validate()
166            .expect("default policy is valid");
167        let policy = ArrowIpcPolicy {
168            dictionary_deltas: true,
169            ..ArrowIpcPolicy::default()
170        };
171        assert!(policy.validate().is_err());
172    }
173
174    #[test]
175    fn given_message_metadata_when_validated_then_should_enforce_every_admission_bound() {
176        let metadata = ArrowIpcMessageMetadata {
177            contract_version: ARROW_IPC_CONTRACT_VERSION,
178            schema_fingerprint: SchemaFingerprint::new([7; 32]),
179            encoded_bytes: 1024,
180            field_count: 4,
181            record_batch_count: 1,
182            row_count: 128,
183            dictionary_count: 0,
184        };
185        metadata.validate().expect("metadata is valid");
186        let mut oversized = metadata.clone();
187        oversized.encoded_bytes = MAX_ARROW_IPC_MESSAGE_BYTES + 1;
188        assert!(oversized.validate().is_err());
189    }
190}