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, encode_payloads_memory};
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    /// Total admitted bytes, including any private owner allowance.
80    pub fn bytes(&self) -> usize {
81        self.owner.lock().unwrap().bytes()
82    }
83
84    /// Admits an absolute total high-water mark; this never shrinks the claim.
85    /// Concurrent incremental allocators must use `try_grow_by`.
86    #[tracing::instrument(name = "storage.arrow.memory.try_grow_to", level = "trace", skip_all, fields(bytes = bytes, outcome = "cancelled"))]
87    pub fn try_grow_to(&self, bytes: usize) -> Result<()> {
88        let result: Result<()> = self.try_grow(bytes).map_err(crate::capacity::Refusal::into_error);
89        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
90        result
91    }
92
93    pub(crate) fn try_grow(&self, bytes: usize) -> std::result::Result<(), crate::capacity::Refusal> {
94        let mut owner = self.owner.lock().unwrap();
95        let additional = bytes.saturating_sub(owner.bytes());
96        owner.try_grow(additional)
97    }
98
99    /// Atomically admits an additional allocation peak without waiting.
100    ///
101    /// Refusal leaves the claim unchanged. Use this for independent increments
102    /// through aliases; `bytes() + try_grow_to()` is not an atomic increment.
103    #[tracing::instrument(name = "storage.arrow.memory.try_grow_by", level = "trace", skip_all, fields(additional = additional, outcome = "cancelled"))]
104    pub fn try_grow_by(&self, additional: usize) -> Result<()> {
105        let result: Result<()> = self.owner.lock().unwrap().try_grow(additional).map_err(crate::capacity::Refusal::into_error);
106        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
107        result
108    }
109
110    pub(crate) async fn grow_to(&self, bytes: usize) -> std::result::Result<(), crate::capacity::Refusal> {
111        Reservation::grow_shared(&self.owner, bytes).await
112    }
113
114    pub(crate) fn shrink_to(&self, bytes: usize) -> Result<()> {
115        let mut held = self.owner.lock().unwrap();
116        if bytes > held.bytes() {
117            return Err(invalid("retained memory cannot grow after allocation"));
118        }
119        held.shrink_to(bytes);
120        Ok(())
121    }
122}
123
124impl std::fmt::Debug for RetainedMemory {
125    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126        formatter.write_str("RetainedMemory")
127    }
128}
129
130impl std::fmt::Debug for Retained {
131    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
132        let _owner = &self.0;
133        formatter.write_str("RetainedMemory")
134    }
135}
136
137impl ArrowSchema {
138    #[tracing::instrument(name = "storage.arrow.schema.try_new", level = "trace", skip_all, fields(outcome = "cancelled"))]
139    pub fn try_new(schema: Arc<Schema>) -> Result<Self> {
140        let result: Result<Self> = (|| {
141            validate_schema(&schema)?;
142            let ipc = Arc::new(ipc::encode(&schema, None)?);
143            Ok(Self { schema, ipc, retention: None })
144        })();
145        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
146        result
147    }
148
149    pub(crate) fn retain(mut self, reservation: RetainedMemory) -> Self {
150        self.retention = Some(Retained(reservation));
151        self
152    }
153
154    pub fn as_schema(&self) -> &Schema {
155        &self.schema
156    }
157
158    pub(crate) fn ipc_capacity(&self) -> usize {
159        self.ipc.capacity()
160    }
161
162    pub fn ipc_bytes(&self) -> &[u8] {
163        &self.ipc
164    }
165
166    #[tracing::instrument(name = "storage.arrow.schema.from_ipc", level = "trace", skip_all, fields(outcome = "cancelled"))]
167    pub fn from_ipc(bytes: Vec<u8>) -> Result<Self> {
168        let result: Result<Self> = (|| {
169            ipc::preflight(&bytes, false)?;
170            let (schema, _) = ipc::decode(&bytes, false)?;
171            Ok(Self { schema, ipc: Arc::new(bytes), retention: None })
172        })();
173        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
174        result
175    }
176}
177
178impl ArrowBatch {
179    #[tracing::instrument(name = "storage.arrow.batch.try_new", level = "trace", skip_all, fields(outcome = "cancelled"))]
180    pub fn try_new(batch: RecordBatch) -> Result<Self> {
181        let result: Result<Self> = (|| {
182            validate_schema(batch.schema().as_ref())?;
183            let encoded = ipc::encode(batch.schema().as_ref(), Some(&batch))?;
184            let resident = retention::resident_bytes(&batch)?;
185            Ok(Self { batch, ipc: Arc::new(encoded), resident, retention: None })
186        })();
187        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
188        result
189    }
190
191    pub(crate) fn retain(mut self, reservation: RetainedMemory) -> Result<Self> {
192        let owner =
193            Arc::new(retention::Owner { _reservation: reservation.clone(), _encoded: Arc::clone(&self.ipc), _schema: self.batch.schema() });
194        self.batch = retention::attach(&self.batch, owner)?;
195        self.retention = Some(Retained(reservation));
196        Ok(self)
197    }
198
199    /// Projects an admitted batch while preserving its original shared claim.
200    ///
201    /// `additional_peak` bounds new output allocations and conversion scratch.
202    /// Private output-owner storage is admitted separately, before the callback.
203    /// Immutable Arrow buffer clones and slices may share the original claim.
204    /// The callback must transfer payloads only into its successful value; errors
205    /// and external captures must not retain them. Nested owners retain their guards.
206    /// Unretained batches from `try_new` or `from_ipc` refuse without conversion.
207    #[tracing::instrument(name = "storage.arrow.batch.try_map", level = "trace", skip_all, fields(additional_peak = additional_peak, outcome = "cancelled"))]
208    pub fn try_map<U, E>(
209        self, additional_peak: usize, map: impl FnOnce(RecordBatch) -> std::result::Result<U, E>,
210    ) -> std::result::Result<crate::ReadBatch<U>, E>
211    where
212        E: From<Error>,
213    {
214        let result: std::result::Result<crate::ReadBatch<U>, E> = (|| {
215            let memory = self.retention.as_ref().ok_or_else(|| E::from(invalid("projection requires an admitted batch")))?.0.clone();
216            let additional = additional_peak
217                .checked_add(crate::ReadBatch::<U>::owner_bytes().map_err(E::from)?)
218                .ok_or_else(|| E::from(Error::ResourceExhausted { resource: "Arrow projection memory" }))?;
219            memory.try_grow_by(additional).map_err(E::from)?;
220            let Self { batch, .. } = self;
221            let value = map(batch)?;
222            Ok(crate::ReadBatch::retained(value, memory))
223        })();
224        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
225        result
226    }
227
228    pub fn as_record_batch(&self) -> &RecordBatch {
229        &self.batch
230    }
231
232    pub(crate) fn ipc_capacity(&self) -> usize {
233        self.ipc.capacity()
234    }
235
236    pub fn ipc_bytes(&self) -> &[u8] {
237        &self.ipc
238    }
239
240    pub(crate) fn encoding_memory_size(&self) -> Result<usize> {
241        self.ipc
242            .capacity()
243            .checked_add(retention::owner_bytes(&self.batch)?)
244            .ok_or_else(|| invalid("Arrow encoding ownership size overflow"))
245    }
246
247    #[tracing::instrument(name = "storage.arrow.batch.from_ipc", level = "trace", skip_all, fields(outcome = "cancelled"))]
248    pub fn from_ipc(bytes: Vec<u8>) -> Result<Self> {
249        let result: Result<Self> = (|| {
250            ipc::preflight(&bytes, true)?;
251            let (_, batch) = ipc::decode(&bytes, true)?;
252            let batch = batch.ok_or_else(|| invalid("missing Arrow record batch"))?;
253            let resident = retention::resident_bytes(&batch)?;
254            Ok(Self { batch, ipc: Arc::new(bytes), resident, retention: None })
255        })();
256        tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
257        result
258    }
259
260    pub fn validate_schema(&self, schema: &Schema) -> Result<()> {
261        if self.batch.schema().as_ref() != schema {
262            return Err(invalid("batch schema differs from announced schema"));
263        }
264        Ok(())
265    }
266}
267
268impl MemorySize for ArrowSchema {
269    fn memory_size(&self) -> usize {
270        std::mem::size_of::<Self>().saturating_add(schema_memory(&self.schema)).saturating_add(self.ipc.capacity())
271    }
272}
273
274impl From<ArrowSchema> for harmos_domain::model::storage::ArrowIpc {
275    fn from(schema: ArrowSchema) -> Self {
276        Self::from(schema.ipc)
277    }
278}
279
280impl From<ArrowBatch> for harmos_domain::model::storage::ArrowIpc {
281    fn from(batch: ArrowBatch) -> Self {
282        let ArrowBatch { batch, ipc, retention, .. } = batch;
283        drop(batch);
284        let encoded = Self::from(ipc);
285        drop(retention);
286        encoded
287    }
288}
289
290impl MemorySize for ArrowBatch {
291    fn memory_size(&self) -> usize {
292        std::mem::size_of::<Self>().saturating_add(self.resident).saturating_add(self.ipc.capacity())
293    }
294}
295
296/* === Public Functions === */
297
298pub fn validate_schema(schema: &Schema) -> Result<()> {
299    let mut fields = 0;
300    for field in schema.fields() {
301        if schema.fields().iter().filter(|other| other.name() == field.name()).count() != 1 {
302            return Err(invalid("duplicate Arrow field name"));
303        }
304        validate_field(field, 0, &mut fields)?;
305    }
306    Ok(())
307}
308
309#[tracing::instrument(name = "storage.arrow.decoding_memory", level = "trace", skip_all, fields(bytes = bytes.len(), batch = batch, outcome = "cancelled"))]
310pub fn decoding_memory(bytes: &[u8], batch: bool) -> Result<usize> {
311    let result: Result<usize> = ipc::preflight(bytes, batch);
312    tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
313    result
314}
315
316/* === Private Functions === */
317
318fn validate_field(field: &Field, depth: usize, count: &mut usize) -> Result<()> {
319    *count = count.checked_add(1).ok_or_else(|| invalid("Arrow field count overflow"))?;
320    if depth > MAX_DEPTH || *count > MAX_FIELDS {
321        return Err(invalid("Arrow schema exceeds structural bounds"));
322    }
323    match field.data_type() {
324        DataType::Null
325        | DataType::Boolean
326        | DataType::Int8
327        | DataType::Int16
328        | DataType::Int32
329        | DataType::Int64
330        | DataType::UInt8
331        | DataType::UInt16
332        | DataType::UInt32
333        | DataType::UInt64
334        | DataType::Float32
335        | DataType::Float64
336        | DataType::Utf8
337        | DataType::Binary => Ok(()),
338        DataType::Timestamp(TimeUnit::Nanosecond, zone) if zone.as_deref() == Some("UTC") => Ok(()),
339        DataType::List(child) => validate_field(child, depth + 1, count),
340        DataType::Struct(children) => {
341            for child in children {
342                if children.iter().filter(|other| other.name() == child.name()).count() != 1 {
343                    return Err(invalid("duplicate Arrow child field name"));
344                }
345                validate_field(child, depth + 1, count)?;
346            }
347            Ok(())
348        }
349        _ => Err(invalid("unsupported Arrow data type")),
350    }
351}
352
353pub(crate) fn schema_memory(schema: &Schema) -> usize {
354    std::mem::size_of::<Schema>()
355        .saturating_add(schema.fields().iter().map(|field| field.size()).fold(0usize, usize::saturating_add))
356        .saturating_add(
357            schema
358                .metadata()
359                .iter()
360                .map(|(key, value)| key.capacity().saturating_add(value.capacity()).saturating_add(64))
361                .fold(0usize, usize::saturating_add),
362        )
363}
364
365pub(crate) fn resident_bytes(batch: &RecordBatch) -> Result<usize> {
366    retention::resident_bytes(batch)
367}
368
369pub(crate) fn new_resident_bytes(batch: &RecordBatch, input: &[&RecordBatch]) -> Result<usize> {
370    retention::new_resident_bytes(batch, input)
371}
372
373pub(crate) fn retain_native(batch: RecordBatch, reservation: RetainedMemory) -> Result<RecordBatch> {
374    let owner = Arc::new(retention::Owner { _reservation: reservation, _encoded: Arc::new(Vec::new()), _schema: batch.schema() });
375    retention::attach(&batch, owner)
376}
377
378pub(crate) fn invalid(reason: &'static str) -> Error {
379    Error::InvalidInput { field: "storage Arrow", reason }
380}
381
382/* === Tests === */
383
384#[cfg(test)]
385mod tests {
386    use super::{ArrowBatch, ArrowSchema, decoding_memory};
387    use crate::MemorySize;
388    use arrow_array::{Int64Array, RecordBatch};
389    use arrow_schema::{DataType, Field, Schema};
390    use std::sync::Arc;
391
392    fn batch() -> RecordBatch {
393        let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false), Field::new("row", DataType::Int64, false)]));
394        RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![2, 1, 1, 2])), Arc::new(Int64Array::from(vec![0, 1, 2, 3]))])
395            .unwrap()
396    }
397
398    #[test]
399    fn ipc_codec_roundtrips() {
400        // Native Buffers And Self-Contained Wire Values Agree.
401        let original = batch();
402        let encoded = ArrowBatch::try_new(original.clone()).unwrap();
403        let bound = decoding_memory(encoded.ipc_bytes(), true).unwrap();
404        let bytes = encoded.ipc_bytes().to_vec();
405        let allocation = bytes.as_ptr();
406        let decoded = ArrowBatch::from_ipc(bytes).unwrap();
407        assert_eq!(decoded.ipc_bytes().as_ptr(), allocation);
408        assert_eq!(decoded.as_record_batch(), &original);
409        assert!(bound >= decoded.memory_size());
410        let header = ArrowSchema::try_new(original.schema()).unwrap();
411        let decoded_header = ArrowSchema::from_ipc(header.ipc_bytes().to_vec()).unwrap();
412        decoded.validate_schema(decoded_header.as_schema()).unwrap();
413        let output = harmos_domain::model::storage::ArrowIpc::from(decoded);
414        assert_eq!(output.bytes().as_ptr(), allocation);
415        assert_eq!(output.bytes(), encoded.ipc_bytes());
416        let shared = encoded.clone();
417        let output = harmos_domain::model::storage::ArrowIpc::from(shared);
418        assert_eq!(output.bytes(), encoded.ipc_bytes());
419        assert_ne!(output.bytes().as_ptr(), encoded.ipc_bytes().as_ptr());
420        assert!(ArrowSchema::from_ipc(encoded.ipc_bytes().to_vec()).is_err());
421        assert!(ArrowBatch::from_ipc(header.ipc_bytes().to_vec()).is_err());
422
423        // Malformed Framing Is Refused Before Arrow Allocates Arrays.
424        for length in [0, 1, 4, encoded.ipc_bytes().len() - 1] {
425            assert!(decoding_memory(&encoded.ipc_bytes()[..length], true).is_err());
426        }
427        let mut trailing = encoded.ipc_bytes().to_vec();
428        trailing.extend_from_slice(&[0, 0, 0, 0]);
429        assert!(decoding_memory(&trailing, true).is_err());
430
431        // Unsupported And Ambiguous Schemas Are Refused.
432        let unsupported = Schema::new(vec![Field::new("value", DataType::LargeUtf8, false)]);
433        assert!(ArrowSchema::try_new(Arc::new(unsupported)).is_err());
434        let duplicate = Schema::new(vec![Field::new("x", DataType::Int64, false), Field::new("x", DataType::Int64, false)]);
435        assert!(ArrowSchema::try_new(Arc::new(duplicate)).is_err());
436    }
437    fn admitted(bytes: usize) -> (crate::StorageManager, ArrowBatch) {
438        use crate::capacity::{Held, Pool, Reservation};
439        use crate::test_backend::{Data, Recovery, history, storage_settings};
440        let storage = crate::StorageManager::new(
441            crate::StorageBackends { primary: Arc::new(Data::<()>::new(history())), state: Arc::new(Recovery::new(history())) },
442            storage_settings(),
443        )
444        .unwrap();
445        {
446            let mut state = storage.inner.state.lock().unwrap();
447            state.ready = true;
448            state.status.running = true;
449        }
450        let memory = Reservation::try_new(&storage.inner, Pool::Read, &Held::default(), bytes).unwrap().retain();
451        let batch = ArrowBatch::try_new(batch()).unwrap().retain(memory).unwrap();
452        (storage, batch)
453    }
454
455    #[test]
456    fn projection_aliases_retain() {
457        let (storage, batch) = admitted(16_384);
458        let alias = batch.clone();
459        let slice = batch.as_record_batch().slice(1, 2);
460        let pointer = batch.as_record_batch().column(0).as_any().downcast_ref::<Int64Array>().unwrap().values().as_ptr();
461        let projected = batch.try_map(0, Ok::<_, crate::Error>).unwrap();
462        assert_eq!(projected.column(0).as_any().downcast_ref::<Int64Array>().unwrap().values().as_ptr(), pointer);
463        drop(projected);
464        drop(alias);
465        assert!(storage.status().read_bytes > 0);
466        assert_eq!(slice.column(0).as_any().downcast_ref::<Int64Array>().unwrap().value(0), 1);
467        drop(slice);
468        assert_eq!(storage.status().read_bytes, 0);
469    }
470
471    #[test]
472    fn projection_dto_retains() {
473        let (storage, batch) = admitted(16_384);
474        let projected = batch
475            .try_map(4 * std::mem::size_of::<i64>(), |batch| {
476                Ok::<_, crate::Error>(batch.column(0).as_any().downcast_ref::<Int64Array>().unwrap().values().to_vec())
477            })
478            .unwrap();
479        assert_eq!(projected.as_slice(), &[2, 1, 1, 2]);
480        assert!(storage.status().read_bytes > 16_384);
481        drop(projected);
482        assert_eq!(storage.status().read_bytes, 0);
483    }
484
485    #[test]
486    fn unretained_projection_refuses() {
487        let encoded = ArrowBatch::try_new(batch()).unwrap();
488        let decoded = ArrowBatch::from_ipc(encoded.ipc_bytes().to_vec()).unwrap();
489        for batch in [encoded, decoded] {
490            let called = std::cell::Cell::new(false);
491            let result = batch.try_map(0, |batch| {
492                called.set(true);
493                Ok::<_, crate::Error>(batch)
494            });
495            assert!(result.is_err());
496            assert!(!called.get());
497        }
498    }
499
500    #[test]
501    fn projection_budget_refuses() {
502        let bytes = crate::test_backend::storage_settings().budget.read_memory_bytes.get();
503        let (storage, batch) = admitted(bytes);
504        let called = std::cell::Cell::new(false);
505        let result = batch.try_map(0, |batch| {
506            called.set(true);
507            Ok::<_, crate::Error>(batch)
508        });
509        assert!(matches!(result, Err(crate::Error::ResourceExhausted { .. })));
510        assert!(!called.get());
511        assert_eq!(storage.status().read_bytes, 0);
512    }
513
514    #[test]
515    fn projection_failure_retains() {
516        let (storage, batch) = admitted(16_384);
517        let alias = batch.clone();
518        let result = batch.try_map(64, |_| -> crate::Result<()> {
519            Err(crate::Error::InvalidInput { field: "projection", reason: "injected conversion failure" })
520        });
521        assert!(result.is_err());
522        assert!(storage.status().read_bytes > 16_384);
523        drop(alias);
524        assert_eq!(storage.status().read_bytes, 0);
525    }
526
527    #[test]
528    fn projection_alias_growth() {
529        let (storage, batch) = admitted(16_384);
530        let values = std::thread::scope(|scope| {
531            let mut tasks = Vec::new();
532            for _ in 0..2 {
533                let batch = batch.clone();
534                tasks.push(scope.spawn(move || batch.try_map(0, |batch| Ok::<_, crate::Error>(batch.num_rows())).unwrap()));
535            }
536            tasks.into_iter().map(|task| task.join().unwrap()).collect::<Vec<_>>()
537        });
538        let scaffold = crate::ReadBatch::<usize>::owner_bytes().unwrap();
539        assert_eq!(storage.status().read_bytes, 16_384 + 2 * scaffold);
540        drop(batch);
541        assert!(values.iter().all(|value| **value == 4));
542        drop(values);
543        assert_eq!(storage.status().read_bytes, 0);
544    }
545}
harmos_storage/arrow.rs

harmos_storage/
arrow.rs
