harmos_storage/writes/
convert.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17use std::num::NonZeroUsize;
18use std::time::{Duration, UNIX_EPOCH};
19// Third Party Imports
20use arrow_array::{
21    Array, BinaryArray, BooleanArray, Float32Array, Float64Array, Int32Array, Int64Array, ListArray, StringArray, StructArray,
22    TimestampNanosecondArray, UInt32Array, UInt64Array,
23};
24// Local Imports
25use super::DecodedBatch;
26use crate::schema::{
27    Cardinality,
28    graph::{Field, Kind, Schema},
29};
30use crate::storage::{ObjectKey, ObjectValue, PayloadEncoding, RecordInput, StreamInput};
31use crate::{Error, Identifier, Result};
32
33/* === Constants === */
34
35const MAX_DEPTH: usize = 64;
36
37/* === Macros === */
38
39/* === Type Aliases === */
40
41/* === Structs === */
42
43/* === Enums === */
44
45/* === Traits === */
46
47/* === Implementations === */
48
49/* === Public Functions === */
50
51/// Bounds the owned row-vector allocation before encoding Arrow payloads.
52///
53/// Includes the outer Vec, its row slots and each exact encoded row capacity.
54/// Counting borrows Arrow values and uses only bounded stack scratch. The same
55/// schema/encoding validation and cumulative `max_bytes` limit apply as encoding.
56/// This does not include consumer codec decoding or its projected DTO allocations.
57#[tracing::instrument(name = "storage.arrow.encode_payloads_memory", level = "trace", skip_all, fields(rows = payload.len(), max_bytes = max_bytes.get(), outcome = "cancelled"))]
58pub fn encode_payloads_memory(schema: &Schema, payload: &StructArray, encoding: PayloadEncoding, max_bytes: NonZeroUsize) -> Result<usize> {
59    let result: Result<usize> = (|| {
60        let mut count = payload_counter(payload.len(), max_bytes)?;
61        for row in 0..payload.len() {
62            encode_message(schema, schema.root, payload, row, encoding, 0, &mut count)?;
63        }
64        Ok(count.bytes)
65    })();
66    tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
67    result
68}
69
70#[tracing::instrument(name = "storage.arrow.encode_payloads", level = "trace", skip_all, fields(rows = payload.len(), max_bytes = max_bytes.get(), outcome = "cancelled"))]
71pub fn encode_payloads(schema: &Schema, payload: &StructArray, encoding: PayloadEncoding, max_bytes: NonZeroUsize) -> Result<Vec<Vec<u8>>> {
72    let result: Result<Vec<Vec<u8>>> = (|| {
73        let mut count = payload_counter(payload.len(), max_bytes)?;
74        let mut rows = Vec::with_capacity(payload.len());
75        for row in 0..payload.len() {
76            let previous = count.bytes;
77            encode_message(schema, schema.root, payload, row, encoding, 0, &mut count)?;
78            let mut bytes = Vec::with_capacity(count.bytes - previous);
79            encode_message(schema, schema.root, payload, row, encoding, 0, &mut bytes)?;
80            rows.push(bytes);
81        }
82        Ok(rows)
83    })();
84    tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
85    result
86}
87
88pub fn records(scope: Option<&Identifier>, batch: &DecodedBatch, schema: &Schema) -> Result<Vec<RecordInput>> {
89    let data = batch.data.as_record_batch();
90    let metadata = struct_column(data.column_by_name("_storage").ok_or_else(invalid)?.as_ref())?;
91    let payload = struct_column(data.column_by_name("payload").ok_or_else(invalid)?.as_ref())?;
92    (0..data.num_rows())
93        .map(|row| {
94            if metadata.is_null(row) {
95                return Err(invalid());
96            }
97            let authored_scope = optional_id(metadata, "scope", row)?;
98            if scope.is_some_and(|expected| authored_scope.as_ref().is_some_and(|actual| actual != expected)) {
99                return Err(invalid());
100            }
101            let source_time = metadata
102                .column_by_name("source_time")
103                .filter(|column| !column.is_null(row))
104                .map(|column| {
105                    let nanos = typed::<TimestampNanosecondArray>(column.as_ref())?.value(row);
106                    if nanos >= 0 {
107                        UNIX_EPOCH.checked_add(Duration::from_nanos(nanos as u64))
108                    } else {
109                        UNIX_EPOCH.checked_sub(Duration::from_nanos(nanos.unsigned_abs()))
110                    }
111                    .ok_or_else(invalid)
112                })
113                .transpose()?;
114            let mut bytes = Vec::new();
115            encode_message(schema, schema.root, payload, row, batch.schema.payload_encoding, 0, &mut bytes)?;
116            Ok(RecordInput {
117                schema: batch.schema.clone(),
118                scope: scope.cloned().or(authored_scope),
119                item: required_id(metadata, "item", row)?,
120                source: optional_id(metadata, "source", row)?,
121                source_time,
122                payload: bytes,
123            })
124        })
125        .collect()
126}
127
128pub fn streams(scope: Option<&Identifier>, batch: &DecodedBatch, schema: &Schema) -> Result<Vec<StreamInput>> {
129    let records = records(scope, batch, schema)?;
130    let metadata = struct_column(batch.data.as_record_batch().column_by_name("_storage").ok_or_else(invalid)?.as_ref())?;
131    let sequence = typed::<UInt64Array>(metadata.column_by_name("sequence").ok_or_else(invalid)?.as_ref())?;
132    records
133        .into_iter()
134        .enumerate()
135        .map(|(row, record)| {
136            if sequence.is_null(row) {
137                return Err(invalid());
138            }
139            if optional_id(metadata, "stream", row)?.is_some_and(|stream| stream != batch.destination) {
140                return Err(invalid());
141            }
142            Ok(StreamInput { record, stream: batch.destination.clone(), sequence: sequence.value(row) })
143        })
144        .collect()
145}
146
147pub fn objects(batch: &DecodedBatch, schema: &Schema, max_key_bytes: NonZeroUsize) -> Result<Vec<ObjectValue>> {
148    let data = batch.data.as_record_batch();
149    let metadata = struct_column(data.column_by_name("_storage").ok_or_else(invalid)?.as_ref())?;
150    let items = typed::<StringArray>(metadata.column_by_name("item").ok_or_else(invalid)?.as_ref())?;
151    let payload = struct_column(data.column_by_name("payload").ok_or_else(invalid)?.as_ref())?;
152    (0..data.num_rows())
153        .map(|row| {
154            if metadata.is_null(row) || items.is_null(row) {
155                return Err(invalid());
156            }
157            let mut bytes = Vec::new();
158            encode_message(schema, schema.root, payload, row, batch.schema.payload_encoding, 0, &mut bytes)?;
159            Ok(ObjectValue { schema: batch.schema.clone(), item: ObjectKey::new(items.value(row), max_key_bytes)?, payload: bytes })
160        })
161        .collect()
162}
163
164/* === Private Functions === */
165
166fn payload_counter(rows: usize, max_bytes: NonZeroUsize) -> Result<Counter> {
167    let metadata = rows
168        .checked_mul(std::mem::size_of::<Vec<u8>>())
169        .and_then(|bytes| bytes.checked_add(std::mem::size_of::<Vec<Vec<u8>>>()))
170        .filter(|bytes| *bytes <= max_bytes.get())
171        .ok_or_else(exhausted)?;
172    Ok(Counter { bytes: metadata, limit: max_bytes.get() })
173}
174
175trait Sink {
176    fn append(&mut self, value: &[u8]) -> Result<()>;
177}
178struct Counter {
179    bytes: usize,
180    limit: usize,
181}
182impl Sink for Counter {
183    fn append(&mut self, value: &[u8]) -> Result<()> {
184        self.bytes = self.bytes.checked_add(value.len()).filter(|bytes| *bytes <= self.limit).ok_or_else(exhausted)?;
185        Ok(())
186    }
187}
188impl Sink for Vec<u8> {
189    fn append(&mut self, value: &[u8]) -> Result<()> {
190        self.extend_from_slice(value);
191        Ok(())
192    }
193}
194fn scalar<T: serde::Serialize>(value: T, output: &mut impl Sink) -> Result<()> {
195    let mut bytes = [0u8; 16];
196    let encoded = postcard::to_slice(&value, &mut bytes).map_err(|_| invalid())?;
197    output.append(encoded)
198}
199fn encode_message(
200    schema: &Schema, definition: usize, values: &StructArray, row: usize, encoding: PayloadEncoding, depth: usize, output: &mut impl Sink,
201) -> Result<()> {
202    if depth > MAX_DEPTH || values.is_null(row) {
203        return Err(invalid());
204    }
205    let fields = schema.definitions.get(definition).ok_or_else(invalid)?;
206    if fields.iter().any(|field| field.cardinality == Cardinality::Variant) {
207        let mut selected = None;
208        for (index, field) in fields.iter().enumerate() {
209            if let Some(value) = values.column_by_name(&field.name)
210                && !value.is_null(row)
211            {
212                if selected.is_some() {
213                    return Err(invalid());
214                }
215                selected = Some((index, field, value));
216            }
217        }
218        let (index, field, value) = selected.ok_or_else(invalid)?;
219        let discriminant = typed::<UInt32Array>(values.column_by_name("__variant").ok_or_else(invalid)?.as_ref())?;
220        if discriminant.is_null(row) || discriminant.value(row) as usize != index {
221            return Err(invalid());
222        }
223        if encoding == PayloadEncoding::PostcardV1 {
224            scalar(index as u32, output)?;
225        }
226        return encode_value(schema, field, value.as_ref(), row, encoding, depth + 1, output);
227    }
228    for field in fields {
229        let value = values.column_by_name(&field.name).ok_or_else(invalid)?;
230        match field.cardinality {
231            Cardinality::Optional => {
232                if encoding == PayloadEncoding::PostcardV1 {
233                    scalar(!value.is_null(row), output)?;
234                }
235                if value.is_null(row) {
236                    continue;
237                }
238            }
239            Cardinality::Repeated => {
240                let lists = typed::<ListArray>(value.as_ref())?;
241                if lists.is_null(row) {
242                    return Err(invalid());
243                }
244                let start = lists.value_offsets()[row] as usize;
245                let end = lists.value_offsets()[row + 1] as usize;
246                if encoding == PayloadEncoding::PostcardV1 {
247                    scalar((end - start) as u64, output)?;
248                }
249                for index in start..end {
250                    encode_value(schema, field, lists.values().as_ref(), index, encoding, depth + 1, output)?;
251                }
252                continue;
253            }
254            Cardinality::Required | Cardinality::Variant => {}
255        }
256        encode_value(schema, field, value.as_ref(), row, encoding, depth + 1, output)?;
257    }
258    Ok(())
259}
260fn encode_value(
261    schema: &Schema, field: &Field, array: &dyn Array, row: usize, encoding: PayloadEncoding, depth: usize, output: &mut impl Sink,
262) -> Result<()> {
263    if array.is_null(row) {
264        return Err(invalid());
265    }
266    let protobuf = encoding == PayloadEncoding::ProtobufV1;
267    macro_rules! number {
268        ($array:ty, $wire:expr, $encode:expr) => {{
269            let value = typed::<$array>(array)?.value(row);
270            if protobuf {
271                varint((u64::from(field.tag) << 3) | $wire, output)?;
272                ($encode)(value, output)?;
273            } else {
274                scalar(value, output)?;
275            }
276        }};
277    }
278    match field.kind {
279        Kind::Bool => number!(BooleanArray, 0, |value: bool, out: &mut _| varint(u64::from(value), out)),
280        Kind::U32 => number!(UInt32Array, 0, |value: u32, out: &mut _| varint(u64::from(value), out)),
281        Kind::U64 => number!(UInt64Array, 0, |value: u64, out: &mut _| varint(value, out)),
282        Kind::I32 => number!(Int32Array, 0, |value: i32, out: &mut _| varint(value as i64 as u64, out)),
283        Kind::I64 => number!(Int64Array, 0, |value: i64, out: &mut _| varint(value as u64, out)),
284        Kind::F32 => number!(Float32Array, 5, |value: f32, out: &mut _| Sink::append(out, &value.to_le_bytes())),
285        Kind::F64 => number!(Float64Array, 1, |value: f64, out: &mut _| Sink::append(out, &value.to_le_bytes())),
286        Kind::String => {
287            let value = typed::<StringArray>(array)?.value(row);
288            if protobuf {
289                header(field.tag, value.len(), output)?;
290            } else {
291                scalar(value.len() as u64, output)?;
292            }
293            output.append(value.as_bytes())?;
294        }
295        Kind::Bytes => {
296            let value = typed::<BinaryArray>(array)?.value(row);
297            if protobuf {
298                header(field.tag, value.len(), output)?;
299            } else {
300                scalar(value.len() as u64, output)?;
301            }
302            output.append(value)?;
303        }
304        Kind::Message { definition } => {
305            let values = struct_column(array)?;
306            if protobuf {
307                let mut count = Counter { bytes: 0, limit: usize::MAX };
308                encode_message(schema, definition, values, row, encoding, depth + 1, &mut count)?;
309                header(field.tag, count.bytes, output)?;
310            }
311            encode_message(schema, definition, values, row, encoding, depth + 1, output)?;
312        }
313    }
314    Ok(())
315}
316
317fn typed<T: Array + 'static>(array: &dyn Array) -> Result<&T> {
318    array.as_any().downcast_ref().ok_or_else(invalid)
319}
320fn struct_column(array: &dyn Array) -> Result<&StructArray> {
321    typed(array)
322}
323fn optional_id(values: &StructArray, name: &str, row: usize) -> Result<Option<Identifier>> {
324    let Some(column) = values.column_by_name(name) else {
325        return Ok(None);
326    };
327    let strings = typed::<StringArray>(column.as_ref())?;
328    if strings.is_null(row) {
329        return Ok(None);
330    }
331    Identifier::new(strings.value(row)).map(Some).map_err(|_| invalid())
332}
333fn required_id(values: &StructArray, name: &str, row: usize) -> Result<Identifier> {
334    optional_id(values, name, row)?.ok_or_else(invalid)
335}
336fn varint(mut value: u64, output: &mut impl Sink) -> Result<()> {
337    let mut bytes = [0u8; 10];
338    let mut count = 0;
339    while value >= 128 {
340        bytes[count] = (value as u8 & 127) | 128;
341        value >>= 7;
342        count += 1;
343    }
344    bytes[count] = value as u8;
345    output.append(&bytes[..=count])
346}
347fn header(tag: u32, length: usize, output: &mut impl Sink) -> Result<()> {
348    varint((u64::from(tag) << 3) | 2, output)?;
349    varint(length as u64, output)
350}
351fn exhausted() -> Error {
352    Error::ResourceExhausted { resource: "Arrow payload encoding bytes" }
353}
354fn invalid() -> Error {
355    Error::InvalidInput { field: "write Arrow batch", reason: "batch metadata or typed payload does not match its schema" }
356}
357
358/* === Tests === */
359
360#[cfg(test)]
361mod tests {
362    use super::*;
363    use arrow_array::{
364        ArrayRef,
365        builder::{ListBuilder, UInt64Builder},
366    };
367    use arrow_schema::{DataType, Field as ArrowField};
368    use std::sync::Arc;
369
370    fn field(name: &str, tag: u32, kind: Kind, cardinality: Cardinality) -> Field {
371        Field { name: name.to_owned().into(), tag, kind, cardinality }
372    }
373
374    fn schema() -> Schema {
375        Schema {
376            root: 0,
377            definitions: vec![
378                vec![
379                    field("text", 1, Kind::String, Cardinality::Required),
380                    field("bytes", 2, Kind::Bytes, Cardinality::Required),
381                    field("values", 3, Kind::U64, Cardinality::Repeated),
382                    field("child", 4, Kind::Message { definition: 1 }, Cardinality::Required),
383                ],
384                vec![field("left", 1, Kind::U32, Cardinality::Variant), field("right", 2, Kind::String, Cardinality::Variant)],
385            ],
386        }
387    }
388
389    fn payload(conflict: bool) -> StructArray {
390        let mut lists = ListBuilder::new(UInt64Builder::new());
391        lists.values().append_value(1);
392        lists.values().append_value(300);
393        lists.append(true);
394        lists.append(true);
395        let lists = lists.finish();
396        let child = StructArray::from(vec![
397            (Arc::new(ArrowField::new("__variant", DataType::UInt32, false)), Arc::new(UInt32Array::from(vec![0, 1])) as ArrayRef),
398            (Arc::new(ArrowField::new("left", DataType::UInt32, true)), Arc::new(UInt32Array::from(vec![Some(7), None])) as ArrayRef),
399            (
400                Arc::new(ArrowField::new("right", DataType::Utf8, true)),
401                Arc::new(StringArray::from(vec![if conflict { Some("both") } else { None }, Some("chosen")])) as ArrayRef,
402            ),
403        ]);
404        StructArray::from(vec![
405            (Arc::new(ArrowField::new("text", DataType::Utf8, false)), Arc::new(StringArray::from(vec!["", "payload"])) as ArrayRef),
406            (
407                Arc::new(ArrowField::new("bytes", DataType::Binary, false)),
408                Arc::new(BinaryArray::from(vec![b"".as_slice(), &[0u8, 1, 255]])) as ArrayRef,
409            ),
410            (Arc::new(ArrowField::new("values", lists.data_type().clone(), false)), Arc::new(lists) as ArrayRef),
411            (Arc::new(ArrowField::new("child", child.data_type().clone(), false)), Arc::new(child) as ArrayRef),
412        ])
413    }
414
415    fn resident(rows: &[Vec<u8>], capacity: usize) -> usize {
416        std::mem::size_of::<Vec<Vec<u8>>>() + capacity * std::mem::size_of::<Vec<u8>>() + rows.iter().map(Vec::capacity).sum::<usize>()
417    }
418    #[test]
419    fn payload_bound_capacities() {
420        let schema = schema();
421        let payload = payload(false);
422        for encoding in [PayloadEncoding::PostcardV1, PayloadEncoding::ProtobufV1] {
423            let maximum = NonZeroUsize::new(1_000_000).unwrap();
424            let bound = encode_payloads_memory(&schema, &payload, encoding, maximum).unwrap();
425            let exact = NonZeroUsize::new(bound).unwrap();
426            let rows = encode_payloads(&schema, &payload, encoding, exact).unwrap();
427            assert_eq!(rows.len(), 2);
428            assert!(resident(&rows, rows.capacity()) <= bound);
429            assert_eq!(encode_payloads_memory(&schema, &payload, encoding, exact).unwrap(), bound);
430            let smaller = NonZeroUsize::new(bound - 1).unwrap();
431            assert!(matches!(encode_payloads_memory(&schema, &payload, encoding, smaller), Err(Error::ResourceExhausted { .. })));
432            assert!(matches!(encode_payloads(&schema, &payload, encoding, smaller), Err(Error::ResourceExhausted { .. })));
433        }
434    }
435
436    #[test]
437    fn payload_bound_empty() {
438        let schema = Schema { root: 0, definitions: vec![vec![]] };
439        let payload = StructArray::new_empty_fields(0, None);
440        for encoding in [PayloadEncoding::PostcardV1, PayloadEncoding::ProtobufV1] {
441            let maximum = NonZeroUsize::new(1024).unwrap();
442            let bound = encode_payloads_memory(&schema, &payload, encoding, maximum).unwrap();
443            let rows = encode_payloads(&schema, &payload, encoding, maximum).unwrap();
444            assert!(rows.is_empty());
445            assert_eq!(bound, resident(&rows, rows.capacity()));
446        }
447    }
448
449    #[test]
450    fn payload_bound_malformed() {
451        let schema = schema();
452        let payload = payload(true);
453        for encoding in [PayloadEncoding::PostcardV1, PayloadEncoding::ProtobufV1] {
454            let maximum = NonZeroUsize::new(1_000_000).unwrap();
455            assert!(matches!(encode_payloads_memory(&schema, &payload, encoding, maximum), Err(Error::InvalidInput { .. })));
456            assert!(matches!(encode_payloads(&schema, &payload, encoding, maximum), Err(Error::InvalidInput { .. })));
457        }
458    }
459
460    #[test]
461    fn payload_bound_overflow() {
462        let schema = Schema { root: 0, definitions: vec![vec![]] };
463        // No Column Allocation Is Needed To Exercise The Row-Slot Overflow.
464        let payload = StructArray::new_empty_fields(usize::MAX, None);
465        let maximum = NonZeroUsize::new(usize::MAX).unwrap();
466        for encoding in [PayloadEncoding::PostcardV1, PayloadEncoding::ProtobufV1] {
467            assert!(matches!(encode_payloads_memory(&schema, &payload, encoding, maximum), Err(Error::ResourceExhausted { .. })));
468            assert!(matches!(encode_payloads(&schema, &payload, encoding, maximum), Err(Error::ResourceExhausted { .. })));
469        }
470    }
471}
harmos_storage/writes/convert.rs

harmos_storage/writes/
convert.rs
