harmos-storage 0.1.7

Harmos storage: durable journals, snapshots and Arrow-backed stream storage.
Documentation
/* ***
 *
 * Harmos
 *
 * SPDX-License-Identifier: Apache-2.0
 *
 * ***
 */

// Storage Arrow Query Contract

/* === Modules === */

/* === Re-Exports === */

/* === Imports === */

// Standard Library Imports
use std::collections::BTreeMap;
// Third Party Imports
use arrow_array::RecordBatch;
use arrow_data::ArrayData;
// Local Imports
use super::invalid;
use crate::Result;

/* === Constants === */

/* === Macros === */

/* === Type Aliases === */

/* === Structs === */

/* === Enums === */

/* === Traits === */

/* === Implementations === */

/* === Public Functions === */

pub(super) fn resident_bytes(batch: &RecordBatch) -> Result<usize> {
    let mut seen = BTreeMap::new();
    let mut bytes = super::schema_memory(batch.schema().as_ref());
    for column in batch.columns() {
        count(&column.to_data(), &mut seen, &mut bytes)?;
    }
    Ok(bytes)
}

/* === Private Functions === */

fn count(data: &ArrayData, seen: &mut BTreeMap<usize, usize>, bytes: &mut usize) -> Result<()> {
    *bytes = bytes.checked_add(std::mem::size_of::<ArrayData>()).ok_or_else(|| invalid("Arrow resident size overflow"))?;
    for buffer in data.buffers().iter().chain(data.nulls().map(|nulls| nulls.buffer())) {
        let base =
            (buffer.as_ptr() as usize).checked_sub(buffer.ptr_offset()).ok_or_else(|| invalid("Arrow allocation offset underflow"))?;
        // Include Accessible Spans When Custom Buffers Report Zero Capacity.
        let span = buffer.len().checked_add(buffer.ptr_offset()).ok_or_else(|| invalid("Arrow accessible allocation span overflow"))?;
        let capacity = buffer.capacity().max(span);
        let previous = seen.entry(base).or_insert(0);
        if capacity > *previous {
            *bytes = bytes.checked_add(capacity - *previous).ok_or_else(|| invalid("Arrow resident size overflow"))?;
            *previous = capacity;
        }
    }
    for child in data.child_data() {
        count(child, seen, bytes)?;
    }
    Ok(())
}

/* === Tests === */