harmos_storage/
arrow.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Storage Arrow Query Contract
11
12/* === Modules === */
13
14mod ipc;
15mod retention;
16
17/* === Re-Exports === */
18
19pub use super::writes::convert::encode_payloads;
20
21/* === Imports === */
22
23// Standard Library Imports
24use std::sync::Arc;
25// Third Party Imports
26use arrow_array::RecordBatch;
27use arrow_schema::{DataType, Field, Schema, TimeUnit};
28// Local Imports
29use crate::{Error, MemorySize, Result};
30
31/* === Constants === */
32
33pub(crate) const MAX_DEPTH: usize = 32;
34pub(crate) const MAX_FIELDS: usize = 4096;
35
36/* === Macros === */
37
38/* === Type Aliases === */
39
40/* === Structs === */
41
42#[derive(Clone, Debug)]
43pub struct ArrowSchema {
44    schema: Arc<Schema>,
45    ipc: Arc<Vec<u8>>,
46}
47
48#[derive(Clone, Debug)]
49pub struct ArrowBatch {
50    batch: RecordBatch,
51    ipc: Arc<Vec<u8>>,
52    resident: usize,
53}
54
55/* === Enums === */
56
57/* === Traits === */
58
59/* === Implementations === */
60
61impl ArrowSchema {
62    #[tracing::instrument(name = "storage.arrow.schema.try_new", level = "trace", skip_all, fields(outcome = "cancelled"))]
63    pub fn try_new(schema: Arc<Schema>) -> Result<Self> {
64        let result: Result<Self> = (|| {
65            validate_schema(&schema)?;
66            let ipc = Arc::new(ipc::encode(&schema, None)?);
67            Ok(Self { schema, ipc })
68        })();
69        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
70        result
71    }
72
73    pub fn as_schema(&self) -> &Schema {
74        &self.schema
75    }
76
77    pub fn ipc_capacity(&self) -> usize {
78        self.ipc.capacity()
79    }
80
81    pub fn ipc_bytes(&self) -> &[u8] {
82        &self.ipc
83    }
84
85    #[tracing::instrument(name = "storage.arrow.schema.from_ipc", level = "trace", skip_all, fields(outcome = "cancelled"))]
86    pub fn from_ipc(bytes: Vec<u8>) -> Result<Self> {
87        let result: Result<Self> = (|| {
88            ipc::preflight(&bytes, false)?;
89            let (schema, _) = ipc::decode(&bytes, false)?;
90            Ok(Self { schema, ipc: Arc::new(bytes) })
91        })();
92        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
93        result
94    }
95}
96
97impl ArrowBatch {
98    #[tracing::instrument(name = "storage.arrow.batch.try_new", level = "trace", skip_all, fields(outcome = "cancelled"))]
99    pub fn try_new(batch: RecordBatch) -> Result<Self> {
100        let result: Result<Self> = (|| {
101            validate_schema(batch.schema().as_ref())?;
102            let encoded = ipc::encode(batch.schema().as_ref(), Some(&batch))?;
103            let resident = retention::resident_bytes(&batch)?;
104            Ok(Self { batch, ipc: Arc::new(encoded), resident })
105        })();
106        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
107        result
108    }
109
110    pub fn try_map<U, E>(self, map: impl FnOnce(RecordBatch) -> std::result::Result<U, E>) -> std::result::Result<crate::ReadBatch<U>, E> {
111        Ok(crate::ReadBatch::new(map(self.batch)?))
112    }
113
114    pub fn as_record_batch(&self) -> &RecordBatch {
115        &self.batch
116    }
117
118    pub fn ipc_capacity(&self) -> usize {
119        self.ipc.capacity()
120    }
121
122    pub fn ipc_bytes(&self) -> &[u8] {
123        &self.ipc
124    }
125
126    #[tracing::instrument(name = "storage.arrow.batch.from_ipc", level = "trace", skip_all, fields(outcome = "cancelled"))]
127    pub fn from_ipc(bytes: Vec<u8>) -> Result<Self> {
128        let result: Result<Self> = (|| {
129            ipc::preflight(&bytes, true)?;
130            let (_, batch) = ipc::decode(&bytes, true)?;
131            let batch = batch.ok_or_else(|| invalid("missing Arrow record batch"))?;
132            let resident = retention::resident_bytes(&batch)?;
133            Ok(Self { batch, ipc: Arc::new(bytes), resident })
134        })();
135        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
136        result
137    }
138
139    pub fn validate_schema(&self, schema: &Schema) -> Result<()> {
140        if self.batch.schema().as_ref() != schema {
141            return Err(invalid("batch schema differs from announced schema"));
142        }
143        Ok(())
144    }
145}
146
147impl MemorySize for ArrowSchema {
148    fn memory_size(&self) -> usize {
149        std::mem::size_of::<Self>().saturating_add(schema_memory(&self.schema)).saturating_add(self.ipc.capacity())
150    }
151}
152
153impl From<ArrowSchema> for harmos_domain::model::storage::ArrowIpc {
154    fn from(schema: ArrowSchema) -> Self {
155        Self::from(schema.ipc)
156    }
157}
158
159impl From<ArrowBatch> for harmos_domain::model::storage::ArrowIpc {
160    fn from(batch: ArrowBatch) -> Self {
161        let ArrowBatch { batch, ipc, .. } = batch;
162        drop(batch);
163        Self::from(ipc)
164    }
165}
166
167impl MemorySize for ArrowBatch {
168    fn memory_size(&self) -> usize {
169        std::mem::size_of::<Self>().saturating_add(self.resident).saturating_add(self.ipc.capacity())
170    }
171}
172
173/* === Public Functions === */
174
175pub fn validate_schema(schema: &Schema) -> Result<()> {
176    let mut fields = 0;
177    for field in schema.fields() {
178        if schema.fields().iter().filter(|other| other.name() == field.name()).count() != 1 {
179            return Err(invalid("duplicate Arrow field name"));
180        }
181        validate_field(field, 0, &mut fields)?;
182    }
183    Ok(())
184}
185
186/* === Private Functions === */
187
188fn validate_field(field: &Field, depth: usize, count: &mut usize) -> Result<()> {
189    *count = count.checked_add(1).ok_or_else(|| invalid("Arrow field count overflow"))?;
190    if depth > MAX_DEPTH || *count > MAX_FIELDS {
191        return Err(invalid("Arrow schema exceeds structural bounds"));
192    }
193    match field.data_type() {
194        DataType::Null
195        | DataType::Boolean
196        | DataType::Int8
197        | DataType::Int16
198        | DataType::Int32
199        | DataType::Int64
200        | DataType::UInt8
201        | DataType::UInt16
202        | DataType::UInt32
203        | DataType::UInt64
204        | DataType::Float32
205        | DataType::Float64
206        | DataType::Utf8
207        | DataType::Binary => Ok(()),
208        DataType::Timestamp(TimeUnit::Nanosecond, zone) if zone.as_deref() == Some("UTC") => Ok(()),
209        DataType::List(child) => validate_field(child, depth + 1, count),
210        DataType::Struct(children) => {
211            for child in children {
212                if children.iter().filter(|other| other.name() == child.name()).count() != 1 {
213                    return Err(invalid("duplicate Arrow child field name"));
214                }
215                validate_field(child, depth + 1, count)?;
216            }
217            Ok(())
218        }
219        _ => Err(invalid("unsupported Arrow data type")),
220    }
221}
222
223pub(crate) fn schema_memory(schema: &Schema) -> usize {
224    std::mem::size_of::<Schema>()
225        .saturating_add(schema.fields().iter().map(|field| field.size()).fold(0usize, usize::saturating_add))
226        .saturating_add(
227            schema
228                .metadata()
229                .iter()
230                .map(|(key, value)| key.capacity().saturating_add(value.capacity()).saturating_add(64))
231                .fold(0usize, usize::saturating_add),
232        )
233}
234
235pub(crate) fn invalid(reason: &'static str) -> Error {
236    Error::InvalidInput { field: "storage Arrow", reason }
237}
238
239/* === Tests === */
240
241#[cfg(test)]
242mod tests {
243    use super::{ArrowBatch, ArrowSchema, ipc};
244    use arrow_array::{Int64Array, RecordBatch};
245    use arrow_schema::{DataType, Field, Schema};
246    use std::sync::Arc;
247
248    fn batch() -> RecordBatch {
249        let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false), Field::new("row", DataType::Int64, false)]));
250        RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![2, 1, 1, 2])), Arc::new(Int64Array::from(vec![0, 1, 2, 3]))])
251            .unwrap()
252    }
253
254    #[test]
255    fn ipc_codec_roundtrips() {
256        // Native Buffers And Self-Contained Wire Values Agree.
257        let original = batch();
258        let encoded = ArrowBatch::try_new(original.clone()).unwrap();
259        ipc::preflight(encoded.ipc_bytes(), true).unwrap();
260        let bytes = encoded.ipc_bytes().to_vec();
261        let allocation = bytes.as_ptr();
262        let decoded = ArrowBatch::from_ipc(bytes).unwrap();
263        assert_eq!(decoded.ipc_bytes().as_ptr(), allocation);
264        assert_eq!(decoded.as_record_batch(), &original);
265        let header = ArrowSchema::try_new(original.schema()).unwrap();
266        let decoded_header = ArrowSchema::from_ipc(header.ipc_bytes().to_vec()).unwrap();
267        decoded.validate_schema(decoded_header.as_schema()).unwrap();
268        let output = harmos_domain::model::storage::ArrowIpc::from(decoded);
269        assert_eq!(output.bytes().as_ptr(), allocation);
270        assert_eq!(output.bytes(), encoded.ipc_bytes());
271        let shared = encoded.clone();
272        let output = harmos_domain::model::storage::ArrowIpc::from(shared);
273        assert_eq!(output.bytes(), encoded.ipc_bytes());
274        assert_ne!(output.bytes().as_ptr(), encoded.ipc_bytes().as_ptr());
275        assert!(ArrowSchema::from_ipc(encoded.ipc_bytes().to_vec()).is_err());
276        assert!(ArrowBatch::from_ipc(header.ipc_bytes().to_vec()).is_err());
277
278        // Malformed Framing Is Refused Before Arrow Allocates Arrays.
279        for length in [0, 1, 4, encoded.ipc_bytes().len() - 1] {
280            assert!(ipc::preflight(&encoded.ipc_bytes()[..length], true).is_err());
281        }
282        let mut trailing = encoded.ipc_bytes().to_vec();
283        trailing.extend_from_slice(&[0, 0, 0, 0]);
284        assert!(ipc::preflight(&trailing, true).is_err());
285
286        // Unsupported And Ambiguous Schemas Are Refused.
287        let unsupported = Schema::new(vec![Field::new("value", DataType::LargeUtf8, false)]);
288        assert!(ArrowSchema::try_new(Arc::new(unsupported)).is_err());
289        let duplicate = Schema::new(vec![Field::new("x", DataType::Int64, false), Field::new("x", DataType::Int64, false)]);
290        assert!(ArrowSchema::try_new(Arc::new(duplicate)).is_err());
291    }
292}
harmos_storage/arrow.rs

harmos_storage/
arrow.rs
