laminar_connectors/serde/raw/
mod.rs

1//! Raw bytes pass-through serialization.
2//!
3//! Treats each record as one opaque value in a single column named `value`: `Utf8` (validated
4//! UTF-8, the default) or `Binary`/`LargeBinary` (bytes preserved exactly).
5
6use std::sync::Arc;
7
8use arrow_array::builder::{BinaryBuilder, LargeBinaryBuilder, StringBuilder};
9use arrow_array::cast::AsArray;
10use arrow_array::{ArrayRef, RecordBatch};
11use arrow_schema::{DataType, Field, Schema, SchemaRef};
12
13use super::{Format, RecordDeserializer, RecordSerializer};
14use crate::error::SerdeError;
15
16/// The default schema for raw bytes: a single `Utf8` column named "value".
17pub(crate) fn raw_schema() -> SchemaRef {
18    raw_schema_of(DataType::Utf8)
19}
20
21fn raw_schema_of(data_type: DataType) -> SchemaRef {
22    Arc::new(Schema::new(vec![Field::new("value", data_type, false)]))
23}
24
25#[cfg(any(
26    test,
27    feature = "kafka",
28    feature = "nats",
29    feature = "rabbitmq",
30    feature = "ibmmq",
31    feature = "files"
32))]
33/// The raw protocol schema selected by a declared schema: a single `value` column of `Binary` or
34/// `LargeBinary` keeps payload bytes exactly; anything else selects the UTF-8 default.
35pub(crate) fn raw_protocol_for(declared: Option<&SchemaRef>) -> SchemaRef {
36    match declared.map(|schema| schema.fields().as_ref()) {
37        Some([field]) if is_binary(field.data_type()) => raw_schema_of(field.data_type().clone()),
38        _ => raw_schema(),
39    }
40}
41
42fn is_binary(data_type: &DataType) -> bool {
43    matches!(data_type, DataType::Binary | DataType::LargeBinary)
44}
45
46fn binary_value_type(schema: &SchemaRef) -> Option<&DataType> {
47    match schema.fields().as_ref() {
48        [field] if is_binary(field.data_type()) => Some(field.data_type()),
49        _ => None,
50    }
51}
52
53fn build_binary(records: &[&[u8]], data_type: &DataType) -> Result<ArrayRef, SerdeError> {
54    let total = records
55        .iter()
56        .try_fold(0_usize, |total, record| total.checked_add(record.len()))
57        .ok_or_else(|| SerdeError::MalformedInput("raw batch size overflows usize".into()))?;
58    if *data_type == DataType::LargeBinary {
59        let mut builder = LargeBinaryBuilder::with_capacity(records.len(), total);
60        for record in records {
61            builder.append_value(record);
62        }
63        return Ok(Arc::new(builder.finish()));
64    }
65    if i32::try_from(total).is_err() {
66        return Err(SerdeError::MalformedInput(format!(
67            "raw batch holds {total} bytes, beyond the Binary offset limit; declare LargeBinary"
68        )));
69    }
70    let mut builder = BinaryBuilder::with_capacity(records.len(), total);
71    for record in records {
72        builder.append_value(record);
73    }
74    Ok(Arc::new(builder.finish()))
75}
76
77/// Raw bytes deserializer.
78///
79/// Stores each record as one `value`: exact bytes when the schema declares a binary `value`,
80/// otherwise validated UTF-8 text.
81#[derive(Debug, Clone)]
82pub struct RawBytesDeserializer {
83    _private: (),
84}
85
86impl RawBytesDeserializer {
87    /// Creates a new raw bytes deserializer.
88    #[must_use]
89    pub fn new() -> Self {
90        Self { _private: () }
91    }
92}
93
94impl Default for RawBytesDeserializer {
95    fn default() -> Self {
96        Self::new()
97    }
98}
99
100impl RecordDeserializer for RawBytesDeserializer {
101    fn deserialize(&self, data: &[u8], schema: &SchemaRef) -> Result<RecordBatch, SerdeError> {
102        if binary_value_type(schema).is_some() {
103            return self.deserialize_batch(&[data], schema);
104        }
105        let s = std::str::from_utf8(data)
106            .map_err(|e| SerdeError::MalformedInput(format!("invalid UTF-8: {e}")))?;
107        let mut builder = StringBuilder::with_capacity(1, data.len());
108        builder.append_value(s);
109
110        let schema = raw_schema();
111        RecordBatch::try_new(schema, vec![Arc::new(builder.finish())])
112            .map_err(|e| SerdeError::MalformedInput(format!("failed to create batch: {e}")))
113    }
114
115    fn deserialize_batch(
116        &self,
117        records: &[&[u8]],
118        schema: &SchemaRef,
119    ) -> Result<RecordBatch, SerdeError> {
120        if let Some(data_type) = binary_value_type(schema) {
121            let column = build_binary(records, data_type)?;
122            return RecordBatch::try_new(Arc::clone(schema), vec![column])
123                .map_err(|e| SerdeError::MalformedInput(format!("failed to create batch: {e}")));
124        }
125        let total_bytes: usize = records.iter().map(|r| r.len()).sum();
126        let mut builder = StringBuilder::with_capacity(records.len(), total_bytes);
127
128        for data in records {
129            let s = std::str::from_utf8(data)
130                .map_err(|e| SerdeError::MalformedInput(format!("invalid UTF-8: {e}")))?;
131            builder.append_value(s);
132        }
133
134        let schema = raw_schema();
135        RecordBatch::try_new(schema, vec![Arc::new(builder.finish())])
136            .map_err(|e| SerdeError::MalformedInput(format!("failed to create batch: {e}")))
137    }
138
139    fn format(&self) -> Format {
140        Format::Raw
141    }
142}
143
144/// Raw bytes serializer.
145///
146/// Emits the first column of each row: binary values verbatim, text as its UTF-8 bytes.
147#[derive(Debug, Clone)]
148pub struct RawBytesSerializer {
149    _private: (),
150}
151
152impl RawBytesSerializer {
153    /// Creates a new raw bytes serializer.
154    #[must_use]
155    pub fn new() -> Self {
156        Self { _private: () }
157    }
158}
159
160impl Default for RawBytesSerializer {
161    fn default() -> Self {
162        Self::new()
163    }
164}
165
166impl RecordSerializer for RawBytesSerializer {
167    fn serialize(&self, batch: &RecordBatch) -> Result<Vec<Vec<u8>>, SerdeError> {
168        if batch.num_columns() == 0 {
169            return Ok(Vec::new());
170        }
171
172        let column = batch.column(0);
173        match column.data_type() {
174            DataType::Utf8 => Ok(owned(
175                column
176                    .as_string::<i32>()
177                    .iter()
178                    .map(|v| v.map(str::as_bytes)),
179            )),
180            DataType::Binary => Ok(owned(column.as_binary::<i32>().iter())),
181            DataType::LargeBinary => Ok(owned(column.as_binary::<i64>().iter())),
182            _ => Err(SerdeError::UnsupportedFormat(
183                "raw serializer expects a Utf8, Binary, or LargeBinary column".into(),
184            )),
185        }
186    }
187
188    fn format(&self) -> Format {
189        Format::Raw
190    }
191}
192
193fn owned<'a>(values: impl Iterator<Item = Option<&'a [u8]>>) -> Vec<Vec<u8>> {
194    values
195        .map(|value| value.map_or_else(Vec::new, <[u8]>::to_vec))
196        .collect()
197}
198
199#[cfg(test)]
200mod tests;
laminar_connectors/serde/raw/mod.rs

laminar_connectors/serde/raw/
mod.rs
