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, Mutex};
25// Third Party Imports
26use arrow_array::RecordBatch;
27use arrow_schema::{DataType, Field, Schema, TimeUnit};
28// Local Imports
29use super::capacity::Reservation;
30use crate::{Error, MemorySize, Result};
31
32/* === Constants === */
33
34pub(crate) const MAX_DEPTH: usize = 32;
35pub(crate) const MAX_FIELDS: usize = 4096;
36
37/* === Macros === */
38
39/* === Type Aliases === */
40
41/* === Structs === */
42
43#[derive(Clone)]
44pub struct RetainedMemory {
45    // Arrow Allocation Requires Unwind-Safe Shared Ownership; The Mutex Owner
46    // Keeps The One Storage Reservation Type Behind That Contract.
47    owner: Arc<Mutex<Reservation>>,
48}
49
50#[derive(Clone, Debug)]
51pub struct ArrowSchema {
52    schema: Arc<Schema>,
53    ipc: Arc<Vec<u8>>,
54    retention: Option<Retained>,
55}
56
57#[derive(Clone, Debug)]
58pub struct ArrowBatch {
59    batch: RecordBatch,
60    ipc: Arc<Vec<u8>>,
61    resident: usize,
62    retention: Option<Retained>,
63}
64
65#[derive(Clone)]
66struct Retained(RetainedMemory);
67
68/* === Enums === */
69
70/* === Traits === */
71
72/* === Implementations === */
73
74impl RetainedMemory {
75    pub(super) fn new(owner: Arc<Mutex<Reservation>>) -> Self {
76        Self { owner }
77    }
78
79    pub(crate) fn shrink_to(&self, bytes: usize) -> Result<()> {
80        let mut held = self.owner.lock().unwrap();
81        if bytes > held.bytes() {
82            return Err(invalid("retained memory cannot grow after allocation"));
83        }
84        held.shrink_to(bytes);
85        Ok(())
86    }
87}
88
89impl std::fmt::Debug for Retained {
90    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91        let _owner = &self.0;
92        formatter.write_str("RetainedMemory")
93    }
94}
95
96impl ArrowSchema {
97    pub fn try_new(schema: Arc<Schema>) -> Result<Self> {
98        validate_schema(&schema)?;
99        let ipc = Arc::new(ipc::encode(&schema, None)?);
100        Ok(Self { schema, ipc, retention: None })
101    }
102
103    pub(crate) fn retain(mut self, reservation: RetainedMemory) -> Self {
104        self.retention = Some(Retained(reservation));
105        self
106    }
107
108    pub fn as_schema(&self) -> &Schema {
109        &self.schema
110    }
111
112    pub fn ipc_bytes(&self) -> &[u8] {
113        &self.ipc
114    }
115
116    pub fn from_ipc(bytes: &[u8]) -> Result<Self> {
117        ipc::preflight(bytes, false)?;
118        let (schema, _) = ipc::decode(bytes, false)?;
119        Ok(Self { schema, ipc: Arc::new(bytes.to_vec()), retention: None })
120    }
121}
122
123impl ArrowBatch {
124    pub fn try_new(batch: RecordBatch) -> Result<Self> {
125        validate_schema(batch.schema().as_ref())?;
126        let encoded = ipc::encode(batch.schema().as_ref(), Some(&batch))?;
127        let resident = retention::resident_bytes(&batch)?;
128        Ok(Self { batch, ipc: Arc::new(encoded), resident, retention: None })
129    }
130
131    pub(crate) fn retain(mut self, reservation: RetainedMemory) -> Result<Self> {
132        let owner =
133            Arc::new(retention::Owner { _reservation: reservation.clone(), _encoded: Arc::clone(&self.ipc), _schema: self.batch.schema() });
134        self.batch = retention::attach(&self.batch, owner)?;
135        self.retention = Some(Retained(reservation));
136        Ok(self)
137    }
138
139    pub fn as_record_batch(&self) -> &RecordBatch {
140        &self.batch
141    }
142
143    pub fn ipc_bytes(&self) -> &[u8] {
144        &self.ipc
145    }
146
147    pub(crate) fn encoding_memory_size(&self) -> Result<usize> {
148        self.ipc
149            .capacity()
150            .checked_add(retention::owner_bytes(&self.batch)?)
151            .ok_or_else(|| invalid("Arrow encoding ownership size overflow"))
152    }
153
154    pub fn from_ipc(bytes: &[u8]) -> Result<Self> {
155        ipc::preflight(bytes, true)?;
156        let (_, batch) = ipc::decode(bytes, true)?;
157        let batch = batch.ok_or_else(|| invalid("missing Arrow record batch"))?;
158        let resident = retention::resident_bytes(&batch)?;
159        Ok(Self { batch, ipc: Arc::new(bytes.to_vec()), resident, retention: None })
160    }
161
162    pub fn validate_schema(&self, schema: &Schema) -> Result<()> {
163        if self.batch.schema().as_ref() != schema {
164            return Err(invalid("batch schema differs from announced schema"));
165        }
166        Ok(())
167    }
168}
169
170impl MemorySize for ArrowSchema {
171    fn memory_size(&self) -> usize {
172        std::mem::size_of::<Self>().saturating_add(schema_memory(&self.schema)).saturating_add(self.ipc.capacity())
173    }
174}
175
176impl From<&ArrowBatch> for harmos_domain::model::storage::ArrowIpc {
177    fn from(batch: &ArrowBatch) -> Self {
178        Self::new(batch.ipc_bytes().to_vec())
179    }
180}
181
182impl From<ArrowBatch> for harmos_domain::model::storage::ArrowIpc {
183    fn from(batch: ArrowBatch) -> Self {
184        Self::from(&batch)
185    }
186}
187
188impl MemorySize for ArrowBatch {
189    fn memory_size(&self) -> usize {
190        std::mem::size_of::<Self>().saturating_add(self.resident).saturating_add(self.ipc.capacity())
191    }
192}
193
194/* === Public Functions === */
195
196pub fn validate_schema(schema: &Schema) -> Result<()> {
197    let mut fields = 0;
198    for field in schema.fields() {
199        if schema.fields().iter().filter(|other| other.name() == field.name()).count() != 1 {
200            return Err(invalid("duplicate Arrow field name"));
201        }
202        validate_field(field, 0, &mut fields)?;
203    }
204    Ok(())
205}
206
207pub fn decoding_memory(bytes: &[u8], batch: bool) -> Result<usize> {
208    ipc::preflight(bytes, batch)
209}
210
211/* === Private Functions === */
212
213fn validate_field(field: &Field, depth: usize, count: &mut usize) -> Result<()> {
214    *count = count.checked_add(1).ok_or_else(|| invalid("Arrow field count overflow"))?;
215    if depth > MAX_DEPTH || *count > MAX_FIELDS {
216        return Err(invalid("Arrow schema exceeds structural bounds"));
217    }
218    match field.data_type() {
219        DataType::Null
220        | DataType::Boolean
221        | DataType::Int8
222        | DataType::Int16
223        | DataType::Int32
224        | DataType::Int64
225        | DataType::UInt8
226        | DataType::UInt16
227        | DataType::UInt32
228        | DataType::UInt64
229        | DataType::Float32
230        | DataType::Float64
231        | DataType::Utf8
232        | DataType::Binary => Ok(()),
233        DataType::Timestamp(TimeUnit::Nanosecond, zone) if zone.as_deref() == Some("UTC") => Ok(()),
234        DataType::List(child) => validate_field(child, depth + 1, count),
235        DataType::Struct(children) => {
236            for child in children {
237                if children.iter().filter(|other| other.name() == child.name()).count() != 1 {
238                    return Err(invalid("duplicate Arrow child field name"));
239                }
240                validate_field(child, depth + 1, count)?;
241            }
242            Ok(())
243        }
244        _ => Err(invalid("unsupported Arrow data type")),
245    }
246}
247
248pub(crate) fn schema_memory(schema: &Schema) -> usize {
249    std::mem::size_of::<Schema>()
250        .saturating_add(schema.fields().iter().map(|field| field.size()).fold(0usize, usize::saturating_add))
251        .saturating_add(
252            schema
253                .metadata()
254                .iter()
255                .map(|(key, value)| key.capacity().saturating_add(value.capacity()).saturating_add(64))
256                .fold(0usize, usize::saturating_add),
257        )
258}
259
260pub(crate) fn resident_bytes(batch: &RecordBatch) -> Result<usize> {
261    retention::resident_bytes(batch)
262}
263
264pub(crate) fn new_resident_bytes(batch: &RecordBatch, input: &[&RecordBatch]) -> Result<usize> {
265    retention::new_resident_bytes(batch, input)
266}
267
268pub(crate) fn retain_native(batch: RecordBatch, reservation: RetainedMemory) -> Result<RecordBatch> {
269    let owner = Arc::new(retention::Owner { _reservation: reservation, _encoded: Arc::new(Vec::new()), _schema: batch.schema() });
270    retention::attach(&batch, owner)
271}
272
273pub(crate) fn invalid(reason: &'static str) -> Error {
274    Error::InvalidInput { field: "storage Arrow", reason }
275}
276
277/* === Tests === */
278
279#[cfg(test)]
280mod tests {
281    use super::{ArrowBatch, ArrowSchema, decoding_memory};
282    use crate::MemorySize;
283    use arrow_array::{Int64Array, RecordBatch};
284    use arrow_schema::{DataType, Field, Schema};
285    use std::sync::Arc;
286
287    fn batch() -> RecordBatch {
288        let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false), Field::new("row", DataType::Int64, false)]));
289        RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![2, 1, 1, 2])), Arc::new(Int64Array::from(vec![0, 1, 2, 3]))])
290            .unwrap()
291    }
292
293    #[test]
294    fn ipc_codec_roundtrips() {
295        // Native Buffers And Self-Contained Wire Values Agree.
296        let original = batch();
297        let encoded = ArrowBatch::try_new(original.clone()).unwrap();
298        let bound = decoding_memory(encoded.ipc_bytes(), true).unwrap();
299        let decoded = ArrowBatch::from_ipc(encoded.ipc_bytes()).unwrap();
300        assert_eq!(decoded.as_record_batch(), &original);
301        assert!(bound >= decoded.memory_size());
302        let header = ArrowSchema::try_new(original.schema()).unwrap();
303        let decoded_header = ArrowSchema::from_ipc(header.ipc_bytes()).unwrap();
304        decoded.validate_schema(decoded_header.as_schema()).unwrap();
305        assert!(ArrowSchema::from_ipc(encoded.ipc_bytes()).is_err());
306        assert!(ArrowBatch::from_ipc(header.ipc_bytes()).is_err());
307
308        // Malformed Framing Is Refused Before Arrow Allocates Arrays.
309        for length in [0, 1, 4, encoded.ipc_bytes().len() - 1] {
310            assert!(decoding_memory(&encoded.ipc_bytes()[..length], true).is_err());
311        }
312        let mut trailing = encoded.ipc_bytes().to_vec();
313        trailing.extend_from_slice(&[0, 0, 0, 0]);
314        assert!(decoding_memory(&trailing, true).is_err());
315
316        // Unsupported And Ambiguous Schemas Are Refused.
317        let unsupported = Schema::new(vec![Field::new("value", DataType::LargeUtf8, false)]);
318        assert!(ArrowSchema::try_new(Arc::new(unsupported)).is_err());
319        let duplicate = Schema::new(vec![Field::new("x", DataType::Int64, false), Field::new("x", DataType::Int64, false)]);
320        assert!(ArrowSchema::try_new(Arc::new(duplicate)).is_err());
321    }
322}
harmos_storage/arrow.rs

harmos_storage/
arrow.rs
