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#[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#[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#[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}