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, 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#[tracing::instrument(name = "storage.arrow.encode_payloads", level = "trace", skip_all, fields(rows = payload.len(), max_bytes = max_bytes.get(), outcome = "cancelled"))]
52pub fn encode_payloads(schema: &Schema, payload: &StructArray, max_bytes: NonZeroUsize) -> Result<Vec<Vec<u8>>> {
53    let result: Result<Vec<Vec<u8>>> = (|| {
54        let mut count = Counter { bytes: 0, limit: max_bytes.get() };
55        let mut rows = Vec::new();
56        rows.try_reserve_exact(payload.len()).map_err(|_| exhausted())?;
57        for row in 0..payload.len() {
58            let previous = count.bytes;
59            encode_message(schema, schema.root, payload, row, 0, &mut count)?;
60            let mut bytes = Vec::with_capacity(count.bytes - previous);
61            encode_message(schema, schema.root, payload, row, 0, &mut bytes)?;
62            rows.push(bytes);
63        }
64        Ok(rows)
65    })();
66    tracing::Span::current().record("outcome", if result.is_ok() { "completed" } else { "failed" });
67    result
68}
69
70pub fn records(scope: Option<&Identifier>, batch: &DecodedBatch, schema: &Schema) -> Result<Vec<RecordInput>> {
71    let data = batch.data.as_record_batch();
72    let metadata = struct_column(data.column_by_name("_storage").ok_or_else(invalid)?.as_ref())?;
73    let payload = struct_column(data.column_by_name("payload").ok_or_else(invalid)?.as_ref())?;
74    (0..data.num_rows())
75        .map(|row| {
76            if metadata.is_null(row) {
77                return Err(invalid());
78            }
79            let authored_scope = optional_id(metadata, "scope", row)?;
80            if scope.is_some_and(|expected| authored_scope.as_ref().is_some_and(|actual| actual != expected)) {
81                return Err(invalid());
82            }
83            let source_time = metadata
84                .column_by_name("source_time")
85                .filter(|column| !column.is_null(row))
86                .map(|column| {
87                    let nanos = typed::<TimestampNanosecondArray>(column.as_ref())?.value(row);
88                    if nanos >= 0 {
89                        UNIX_EPOCH.checked_add(Duration::from_nanos(nanos as u64))
90                    } else {
91                        UNIX_EPOCH.checked_sub(Duration::from_nanos(nanos.unsigned_abs()))
92                    }
93                    .ok_or_else(invalid)
94                })
95                .transpose()?;
96            let mut bytes = Vec::new();
97            encode_message(schema, schema.root, payload, row, 0, &mut bytes)?;
98            Ok(RecordInput {
99                schema: batch.schema.clone(),
100                scope: scope.cloned().or(authored_scope),
101                item: required_id(metadata, "item", row)?,
102                source: optional_id(metadata, "source", row)?,
103                source_time,
104                payload: bytes,
105            })
106        })
107        .collect()
108}
109
110pub fn streams(scope: Option<&Identifier>, batch: &DecodedBatch, schema: &Schema) -> Result<Vec<StreamInput>> {
111    let records = records(scope, batch, schema)?;
112    let metadata = struct_column(batch.data.as_record_batch().column_by_name("_storage").ok_or_else(invalid)?.as_ref())?;
113    let sequence = typed::<UInt64Array>(metadata.column_by_name("sequence").ok_or_else(invalid)?.as_ref())?;
114    records
115        .into_iter()
116        .enumerate()
117        .map(|(row, record)| {
118            if sequence.is_null(row) {
119                return Err(invalid());
120            }
121            if optional_id(metadata, "stream", row)?.is_some_and(|stream| stream != batch.destination) {
122                return Err(invalid());
123            }
124            Ok(StreamInput {
125                record,
126                stream: batch.destination.clone(),
127                sequence: sequence.value(row),
128            })
129        })
130        .collect()
131}
132
133pub fn objects(batch: &DecodedBatch, schema: &Schema, max_key_bytes: NonZeroUsize) -> Result<Vec<ObjectValue>> {
134    let data = batch.data.as_record_batch();
135    let metadata = struct_column(data.column_by_name("_storage").ok_or_else(invalid)?.as_ref())?;
136    let items = typed::<StringArray>(metadata.column_by_name("item").ok_or_else(invalid)?.as_ref())?;
137    let payload = struct_column(data.column_by_name("payload").ok_or_else(invalid)?.as_ref())?;
138    (0..data.num_rows())
139        .map(|row| {
140            if metadata.is_null(row) || items.is_null(row) {
141                return Err(invalid());
142            }
143            let mut bytes = Vec::new();
144            encode_message(schema, schema.root, payload, row, 0, &mut bytes)?;
145            Ok(ObjectValue {
146                schema: batch.schema.clone(),
147                item: ObjectKey::new(items.value(row), max_key_bytes)?,
148                payload: bytes,
149            })
150        })
151        .collect()
152}
153
154/* === Private Functions === */
155
156trait Sink {
157    fn append(&mut self, value: &[u8]) -> Result<()>;
158}
159struct Counter {
160    bytes: usize,
161    limit: usize,
162}
163impl Sink for Counter {
164    fn append(&mut self, value: &[u8]) -> Result<()> {
165        self.bytes = self.bytes.checked_add(value.len()).filter(|bytes| *bytes <= self.limit).ok_or_else(exhausted)?;
166        Ok(())
167    }
168}
169impl Sink for Vec<u8> {
170    fn append(&mut self, value: &[u8]) -> Result<()> {
171        self.extend_from_slice(value);
172        Ok(())
173    }
174}
175fn scalar<T: serde::Serialize>(value: T, output: &mut impl Sink) -> Result<()> {
176    let mut bytes = [0u8; 16];
177    let encoded = postcard::to_slice(&value, &mut bytes).map_err(|_| invalid())?;
178    output.append(encoded)
179}
180fn encode_message(
181    schema: &Schema, definition: usize, values: &StructArray, row: usize, depth: usize, output: &mut impl Sink,
182) -> Result<()> {
183    if depth > MAX_DEPTH || values.is_null(row) {
184        return Err(invalid());
185    }
186    let fields = schema.definitions.get(definition).ok_or_else(invalid)?;
187    if fields.iter().any(|field| field.cardinality == Cardinality::Variant) {
188        let mut selected = None;
189        for (index, field) in fields.iter().enumerate() {
190            if let Some(value) = values.column_by_name(&field.name)
191                && !value.is_null(row)
192            {
193                if selected.is_some() {
194                    return Err(invalid());
195                }
196                selected = Some((index, field, value));
197            }
198        }
199        let (index, field, value) = selected.ok_or_else(invalid)?;
200        let discriminant = typed::<UInt32Array>(values.column_by_name("__variant").ok_or_else(invalid)?.as_ref())?;
201        if discriminant.is_null(row) || discriminant.value(row) as usize != index {
202            return Err(invalid());
203        }
204        scalar(index as u32, output)?;
205        return encode_value(schema, field, value.as_ref(), row, depth + 1, output);
206    }
207    for field in fields {
208        let value = values.column_by_name(&field.name).ok_or_else(invalid)?;
209        match field.cardinality {
210            Cardinality::Optional => {
211                scalar(!value.is_null(row), output)?;
212                if value.is_null(row) {
213                    continue;
214                }
215            }
216            Cardinality::Repeated => {
217                let lists = typed::<ListArray>(value.as_ref())?;
218                if lists.is_null(row) {
219                    return Err(invalid());
220                }
221                let start = lists.value_offsets()[row] as usize;
222                let end = lists.value_offsets()[row + 1] as usize;
223                scalar((end - start) as u64, output)?;
224                for index in start..end {
225                    encode_value(schema, field, lists.values().as_ref(), index, depth + 1, output)?;
226                }
227                continue;
228            }
229            Cardinality::Required | Cardinality::Variant => {}
230        }
231        encode_value(schema, field, value.as_ref(), row, depth + 1, output)?;
232    }
233    Ok(())
234}
235fn encode_value(schema: &Schema, field: &Field, array: &dyn Array, row: usize, depth: usize, output: &mut impl Sink) -> Result<()> {
236    if array.is_null(row) {
237        return Err(invalid());
238    }
239    macro_rules! number {
240        ($array:ty) => {{
241            scalar(typed::<$array>(array)?.value(row), output)?;
242        }};
243    }
244    match field.kind {
245        Kind::Bool => number!(BooleanArray),
246        Kind::U32 => number!(UInt32Array),
247        Kind::U64 => number!(UInt64Array),
248        Kind::I32 => number!(Int32Array),
249        Kind::I64 => number!(Int64Array),
250        Kind::F32 => number!(Float32Array),
251        Kind::F64 => number!(Float64Array),
252        Kind::String => {
253            let value = typed::<StringArray>(array)?.value(row);
254            scalar(value.len() as u64, output)?;
255            output.append(value.as_bytes())?;
256        }
257        Kind::Bytes => {
258            let value = typed::<BinaryArray>(array)?.value(row);
259            scalar(value.len() as u64, output)?;
260            output.append(value)?;
261        }
262        Kind::Message { definition } => {
263            let values = struct_column(array)?;
264            encode_message(schema, definition, values, row, depth + 1, output)?;
265        }
266    }
267    Ok(())
268}
269
270fn typed<T: Array + 'static>(array: &dyn Array) -> Result<&T> {
271    array.as_any().downcast_ref().ok_or_else(invalid)
272}
273fn struct_column(array: &dyn Array) -> Result<&StructArray> {
274    typed(array)
275}
276fn optional_id(values: &StructArray, name: &str, row: usize) -> Result<Option<Identifier>> {
277    let Some(column) = values.column_by_name(name) else {
278        return Ok(None);
279    };
280    let strings = typed::<StringArray>(column.as_ref())?;
281    if strings.is_null(row) {
282        return Ok(None);
283    }
284    Identifier::new(strings.value(row)).map(Some).map_err(|_| invalid())
285}
286fn required_id(values: &StructArray, name: &str, row: usize) -> Result<Identifier> {
287    optional_id(values, name, row)?.ok_or_else(invalid)
288}
289fn exhausted() -> Error {
290    Error::ResourceExhausted { resource: "Arrow payload encoding bytes" }
291}
292fn invalid() -> Error {
293    Error::InvalidInput {
294        field: "write Arrow batch",
295        reason: "batch metadata or typed payload does not match its schema",
296    }
297}
298
299/* === Tests === */
300
301#[cfg(test)]
302mod tests {
303    use super::*;
304    use arrow_array::{
305        ArrayRef,
306        builder::{ListBuilder, UInt64Builder},
307    };
308    use arrow_schema::{DataType, Field as ArrowField};
309    use std::sync::Arc;
310
311    fn field(name: &str, tag: u32, kind: Kind, cardinality: Cardinality) -> Field {
312        Field {
313            name: name.to_owned().into(),
314            tag,
315            kind,
316            cardinality,
317        }
318    }
319
320    fn schema() -> Schema {
321        Schema {
322            root: 0,
323            definitions: vec![
324                vec![
325                    field("text", 1, Kind::String, Cardinality::Required),
326                    field("bytes", 2, Kind::Bytes, Cardinality::Required),
327                    field("values", 3, Kind::U64, Cardinality::Repeated),
328                    field("child", 4, Kind::Message { definition: 1 }, Cardinality::Required),
329                ],
330                vec![field("left", 1, Kind::U32, Cardinality::Variant), field("right", 2, Kind::String, Cardinality::Variant)],
331            ],
332        }
333    }
334
335    fn payload(conflict: bool) -> StructArray {
336        let mut lists = ListBuilder::new(UInt64Builder::new());
337        lists.values().append_value(1);
338        lists.values().append_value(300);
339        lists.append(true);
340        lists.append(true);
341        let lists = lists.finish();
342        let child = StructArray::from(vec![
343            (Arc::new(ArrowField::new("__variant", DataType::UInt32, false)), Arc::new(UInt32Array::from(vec![0, 1])) as ArrayRef),
344            (Arc::new(ArrowField::new("left", DataType::UInt32, true)), Arc::new(UInt32Array::from(vec![Some(7), None])) as ArrayRef),
345            (
346                Arc::new(ArrowField::new("right", DataType::Utf8, true)),
347                Arc::new(StringArray::from(vec![if conflict { Some("both") } else { None }, Some("chosen")])) as ArrayRef,
348            ),
349        ]);
350        StructArray::from(vec![
351            (Arc::new(ArrowField::new("text", DataType::Utf8, false)), Arc::new(StringArray::from(vec!["", "payload"])) as ArrayRef),
352            (
353                Arc::new(ArrowField::new("bytes", DataType::Binary, false)),
354                Arc::new(BinaryArray::from(vec![b"".as_slice(), &[0u8, 1, 255]])) as ArrayRef,
355            ),
356            (Arc::new(ArrowField::new("values", lists.data_type().clone(), false)), Arc::new(lists) as ArrayRef),
357            (Arc::new(ArrowField::new("child", child.data_type().clone(), false)), Arc::new(child) as ArrayRef),
358        ])
359    }
360
361    #[test]
362    fn payload_encoded_limits() {
363        let schema = schema();
364        let payload = payload(false);
365        {
366            let rows = encode_payloads(&schema, &payload, NonZeroUsize::new(1_000_000).unwrap()).unwrap();
367            let length = rows.iter().map(Vec::len).sum::<usize>();
368            assert_eq!(rows.len(), 2);
369            let exact = NonZeroUsize::new(length).unwrap();
370            assert_eq!(encode_payloads(&schema, &payload, exact).unwrap(), rows);
371            let smaller = NonZeroUsize::new(length - 1).unwrap();
372            assert!(matches!(encode_payloads(&schema, &payload, smaller), Err(Error::ResourceExhausted { .. })));
373        }
374    }
375
376    #[test]
377    fn payload_empty_and_malformed() {
378        let empty = Schema { root: 0, definitions: vec![vec![]] };
379        {
380            let maximum = NonZeroUsize::new(1_000_000).unwrap();
381            assert!(encode_payloads(&empty, &StructArray::new_empty_fields(0, None), maximum).unwrap().is_empty());
382            assert!(matches!(encode_payloads(&schema(), &payload(true), maximum), Err(Error::InvalidInput { .. })));
383            let oversized = StructArray::new_empty_fields(usize::MAX, None);
384            assert!(matches!(encode_payloads(&empty, &oversized, maximum), Err(Error::ResourceExhausted { .. })));
385        }
386    }
387}
harmos_storage/writes/convert.rs

harmos_storage/writes/
convert.rs
