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
51pub fn encode_payloads(schema: &Schema, payload: &StructArray, encoding: PayloadEncoding, max_bytes: NonZeroUsize) -> Result<Vec<Vec<u8>>> {
52    let metadata = payload
53        .len()
54        .checked_mul(std::mem::size_of::<Vec<u8>>())
55        .and_then(|bytes| bytes.checked_add(std::mem::size_of::<Vec<Vec<u8>>>()))
56        .ok_or_else(exhausted)?;
57    let mut count = Counter { bytes: metadata, limit: max_bytes.get() };
58    if metadata > max_bytes.get() {
59        return Err(exhausted());
60    }
61    for row in 0..payload.len() {
62        encode_message(schema, schema.root, payload, row, encoding, 0, &mut count)?;
63    }
64    let mut rows = Vec::with_capacity(payload.len());
65    for row in 0..payload.len() {
66        let mut size = Counter { bytes: 0, limit: max_bytes.get() };
67        encode_message(schema, schema.root, payload, row, encoding, 0, &mut size)?;
68        let mut bytes = Vec::with_capacity(size.bytes);
69        encode_message(schema, schema.root, payload, row, encoding, 0, &mut bytes)?;
70        rows.push(bytes);
71    }
72    Ok(rows)
73}
74
75pub fn records(scope: Option<&Identifier>, batch: &DecodedBatch, schema: &Schema) -> Result<Vec<RecordInput>> {
76    let data = batch.data.as_record_batch();
77    let metadata = struct_column(data.column_by_name("_storage").ok_or_else(invalid)?.as_ref())?;
78    let payload = struct_column(data.column_by_name("payload").ok_or_else(invalid)?.as_ref())?;
79    (0..data.num_rows())
80        .map(|row| {
81            if metadata.is_null(row) {
82                return Err(invalid());
83            }
84            let authored_scope = optional_id(metadata, "scope", row)?;
85            if scope.is_some_and(|expected| authored_scope.as_ref().is_some_and(|actual| actual != expected)) {
86                return Err(invalid());
87            }
88            let source_time = metadata
89                .column_by_name("source_time")
90                .filter(|column| !column.is_null(row))
91                .map(|column| {
92                    let nanos = typed::<TimestampNanosecondArray>(column.as_ref())?.value(row);
93                    if nanos >= 0 {
94                        UNIX_EPOCH.checked_add(Duration::from_nanos(nanos as u64))
95                    } else {
96                        UNIX_EPOCH.checked_sub(Duration::from_nanos(nanos.unsigned_abs()))
97                    }
98                    .ok_or_else(invalid)
99                })
100                .transpose()?;
101            let mut bytes = Vec::new();
102            encode_message(schema, schema.root, payload, row, batch.schema.payload_encoding, 0, &mut bytes)?;
103            Ok(RecordInput {
104                schema: batch.schema.clone(),
105                scope: scope.cloned().or(authored_scope),
106                item: required_id(metadata, "item", row)?,
107                source: optional_id(metadata, "source", row)?,
108                source_time,
109                payload: bytes,
110            })
111        })
112        .collect()
113}
114
115pub fn streams(scope: Option<&Identifier>, batch: &DecodedBatch, schema: &Schema) -> Result<Vec<StreamInput>> {
116    let records = records(scope, batch, schema)?;
117    let metadata = struct_column(batch.data.as_record_batch().column_by_name("_storage").ok_or_else(invalid)?.as_ref())?;
118    let sequence = typed::<UInt64Array>(metadata.column_by_name("sequence").ok_or_else(invalid)?.as_ref())?;
119    records
120        .into_iter()
121        .enumerate()
122        .map(|(row, record)| {
123            if sequence.is_null(row) {
124                return Err(invalid());
125            }
126            if optional_id(metadata, "stream", row)?.is_some_and(|stream| stream != batch.destination) {
127                return Err(invalid());
128            }
129            Ok(StreamInput { record, stream: batch.destination.clone(), sequence: sequence.value(row) })
130        })
131        .collect()
132}
133
134pub fn objects(batch: &DecodedBatch, schema: &Schema, max_key_bytes: NonZeroUsize) -> Result<Vec<ObjectValue>> {
135    let data = batch.data.as_record_batch();
136    let metadata = struct_column(data.column_by_name("_storage").ok_or_else(invalid)?.as_ref())?;
137    let items = typed::<StringArray>(metadata.column_by_name("item").ok_or_else(invalid)?.as_ref())?;
138    let payload = struct_column(data.column_by_name("payload").ok_or_else(invalid)?.as_ref())?;
139    (0..data.num_rows())
140        .map(|row| {
141            if metadata.is_null(row) || items.is_null(row) {
142                return Err(invalid());
143            }
144            let mut bytes = Vec::new();
145            encode_message(schema, schema.root, payload, row, batch.schema.payload_encoding, 0, &mut bytes)?;
146            Ok(ObjectValue { schema: batch.schema.clone(), item: ObjectKey::new(items.value(row), max_key_bytes)?, payload: bytes })
147        })
148        .collect()
149}
150
151/* === Private Functions === */
152
153trait Sink {
154    fn append(&mut self, value: &[u8]) -> Result<()>;
155}
156struct Counter {
157    bytes: usize,
158    limit: usize,
159}
160impl Sink for Counter {
161    fn append(&mut self, value: &[u8]) -> Result<()> {
162        self.bytes = self.bytes.checked_add(value.len()).filter(|bytes| *bytes <= self.limit).ok_or_else(exhausted)?;
163        Ok(())
164    }
165}
166impl Sink for Vec<u8> {
167    fn append(&mut self, value: &[u8]) -> Result<()> {
168        self.extend_from_slice(value);
169        Ok(())
170    }
171}
172fn scalar<T: serde::Serialize>(value: T, output: &mut impl Sink) -> Result<()> {
173    let mut bytes = [0u8; 16];
174    let encoded = postcard::to_slice(&value, &mut bytes).map_err(|_| invalid())?;
175    output.append(encoded)
176}
177fn encode_message(
178    schema: &Schema, definition: usize, values: &StructArray, row: usize, encoding: PayloadEncoding, depth: usize, output: &mut impl Sink,
179) -> Result<()> {
180    if depth > MAX_DEPTH || values.is_null(row) {
181        return Err(invalid());
182    }
183    let fields = schema.definitions.get(definition).ok_or_else(invalid)?;
184    if fields.iter().any(|field| field.cardinality == Cardinality::Variant) {
185        let mut selected = None;
186        for (index, field) in fields.iter().enumerate() {
187            if let Some(value) = values.column_by_name(&field.name)
188                && !value.is_null(row)
189            {
190                if selected.is_some() {
191                    return Err(invalid());
192                }
193                selected = Some((index, field, value));
194            }
195        }
196        let (index, field, value) = selected.ok_or_else(invalid)?;
197        let discriminant = typed::<UInt32Array>(values.column_by_name("__variant").ok_or_else(invalid)?.as_ref())?;
198        if discriminant.is_null(row) || discriminant.value(row) as usize != index {
199            return Err(invalid());
200        }
201        if encoding == PayloadEncoding::PostcardV1 {
202            scalar(index as u32, output)?;
203        }
204        return encode_value(schema, field, value.as_ref(), row, encoding, depth + 1, output);
205    }
206    for field in fields {
207        let value = values.column_by_name(&field.name).ok_or_else(invalid)?;
208        match field.cardinality {
209            Cardinality::Optional => {
210                if encoding == PayloadEncoding::PostcardV1 {
211                    scalar(!value.is_null(row), output)?;
212                }
213                if value.is_null(row) {
214                    continue;
215                }
216            }
217            Cardinality::Repeated => {
218                let lists = typed::<ListArray>(value.as_ref())?;
219                if lists.is_null(row) {
220                    return Err(invalid());
221                }
222                let start = lists.value_offsets()[row] as usize;
223                let end = lists.value_offsets()[row + 1] as usize;
224                if encoding == PayloadEncoding::PostcardV1 {
225                    scalar((end - start) as u64, output)?;
226                }
227                for index in start..end {
228                    encode_value(schema, field, lists.values().as_ref(), index, encoding, depth + 1, output)?;
229                }
230                continue;
231            }
232            Cardinality::Required | Cardinality::Variant => {}
233        }
234        encode_value(schema, field, value.as_ref(), row, encoding, depth + 1, output)?;
235    }
236    Ok(())
237}
238fn encode_value(
239    schema: &Schema, field: &Field, array: &dyn Array, row: usize, encoding: PayloadEncoding, depth: usize, output: &mut impl Sink,
240) -> Result<()> {
241    if array.is_null(row) {
242        return Err(invalid());
243    }
244    let protobuf = encoding == PayloadEncoding::ProtobufV1;
245    macro_rules! number {
246        ($array:ty, $wire:expr, $encode:expr) => {{
247            let value = typed::<$array>(array)?.value(row);
248            if protobuf {
249                varint((u64::from(field.tag) << 3) | $wire, output)?;
250                ($encode)(value, output)?;
251            } else {
252                scalar(value, output)?;
253            }
254        }};
255    }
256    match field.kind {
257        Kind::Bool => number!(BooleanArray, 0, |value: bool, out: &mut _| varint(u64::from(value), out)),
258        Kind::U32 => number!(UInt32Array, 0, |value: u32, out: &mut _| varint(u64::from(value), out)),
259        Kind::U64 => number!(UInt64Array, 0, |value: u64, out: &mut _| varint(value, out)),
260        Kind::I32 => number!(Int32Array, 0, |value: i32, out: &mut _| varint(value as i64 as u64, out)),
261        Kind::I64 => number!(Int64Array, 0, |value: i64, out: &mut _| varint(value as u64, out)),
262        Kind::F32 => number!(Float32Array, 5, |value: f32, out: &mut _| Sink::append(out, &value.to_le_bytes())),
263        Kind::F64 => number!(Float64Array, 1, |value: f64, out: &mut _| Sink::append(out, &value.to_le_bytes())),
264        Kind::String => {
265            let value = typed::<StringArray>(array)?.value(row);
266            if protobuf {
267                header(field.tag, value.len(), output)?;
268            } else {
269                scalar(value.len() as u64, output)?;
270            }
271            output.append(value.as_bytes())?;
272        }
273        Kind::Bytes => {
274            let value = typed::<BinaryArray>(array)?.value(row);
275            if protobuf {
276                header(field.tag, value.len(), output)?;
277            } else {
278                scalar(value.len() as u64, output)?;
279            }
280            output.append(value)?;
281        }
282        Kind::Message { definition } => {
283            let values = struct_column(array)?;
284            if protobuf {
285                let mut count = Counter { bytes: 0, limit: usize::MAX };
286                encode_message(schema, definition, values, row, encoding, depth + 1, &mut count)?;
287                header(field.tag, count.bytes, output)?;
288            }
289            encode_message(schema, definition, values, row, encoding, depth + 1, output)?;
290        }
291    }
292    Ok(())
293}
294
295fn typed<T: Array + 'static>(array: &dyn Array) -> Result<&T> {
296    array.as_any().downcast_ref().ok_or_else(invalid)
297}
298fn struct_column(array: &dyn Array) -> Result<&StructArray> {
299    typed(array)
300}
301fn optional_id(values: &StructArray, name: &str, row: usize) -> Result<Option<Identifier>> {
302    let Some(column) = values.column_by_name(name) else {
303        return Ok(None);
304    };
305    let strings = typed::<StringArray>(column.as_ref())?;
306    if strings.is_null(row) {
307        return Ok(None);
308    }
309    Identifier::new(strings.value(row)).map(Some).map_err(|_| invalid())
310}
311fn required_id(values: &StructArray, name: &str, row: usize) -> Result<Identifier> {
312    optional_id(values, name, row)?.ok_or_else(invalid)
313}
314fn varint(mut value: u64, output: &mut impl Sink) -> Result<()> {
315    let mut bytes = [0u8; 10];
316    let mut count = 0;
317    while value >= 128 {
318        bytes[count] = (value as u8 & 127) | 128;
319        value >>= 7;
320        count += 1;
321    }
322    bytes[count] = value as u8;
323    output.append(&bytes[..=count])
324}
325fn header(tag: u32, length: usize, output: &mut impl Sink) -> Result<()> {
326    varint((u64::from(tag) << 3) | 2, output)?;
327    varint(length as u64, output)
328}
329fn exhausted() -> Error {
330    Error::ResourceExhausted { resource: "Arrow payload encoding bytes" }
331}
332fn invalid() -> Error {
333    Error::InvalidInput { field: "write Arrow batch", reason: "batch metadata or typed payload does not match its schema" }
334}
335
336/* === Tests === */
harmos_storage/writes/convert.rs

harmos_storage/writes/
convert.rs
