laminar_connectors/schema/json/encoder/
mod.rs

1//! JSON format encoder implementing [`FormatEncoder`].
2
3use std::io::Write as _;
4use std::sync::Arc;
5
6use arrow_array::cast::AsArray;
7use arrow_array::{Array, RecordBatch};
8use arrow_json::writer::{
9    Encoder, EncoderFactory, EncoderOptions, LineDelimited, NullableEncoder, WriterBuilder,
10};
11use arrow_schema::{ArrowError, DataType, FieldRef, SchemaRef};
12
13use crate::schema::error::{SchemaError, SchemaResult};
14use crate::schema::traits::FormatEncoder;
15
16/// Encodes Arrow `RecordBatch`es into one JSON object per row via
17/// `arrow_json::writer`. `LargeBinary` columns are inlined as JSON
18/// when valid (JSONB passthrough), matching `CollectExtra` semantics.
19#[derive(Debug)]
20pub struct JsonEncoder {
21    schema: SchemaRef,
22}
23
24impl JsonEncoder {
25    /// Creates a new JSON encoder for the given schema.
26    #[must_use]
27    pub fn new(schema: SchemaRef) -> Self {
28        Self { schema }
29    }
30}
31
32impl FormatEncoder for JsonEncoder {
33    fn input_schema(&self) -> SchemaRef {
34        self.schema.clone()
35    }
36
37    fn encode_batch(&self, batch: &RecordBatch) -> SchemaResult<Vec<Vec<u8>>> {
38        if batch.num_rows() == 0 {
39            return Ok(Vec::new());
40        }
41
42        let mut buf = Vec::new();
43        {
44            let mut writer = WriterBuilder::new()
45                .with_explicit_nulls(true)
46                .with_encoder_factory(Arc::new(JsonbPassthroughFactory))
47                .build::<_, LineDelimited>(&mut buf);
48            writer
49                .write(batch)
50                .map_err(|e| SchemaError::DecodeError(format!("JSON encode error: {e}")))?;
51            writer
52                .finish()
53                .map_err(|e| SchemaError::DecodeError(format!("JSON finish error: {e}")))?;
54        }
55
56        let output: Vec<Vec<u8>> = buf
57            .split(|&b| b == b'\n')
58            .filter(|line| !line.is_empty())
59            .map(<[u8]>::to_vec)
60            .collect();
61
62        Ok(output)
63    }
64
65    fn format_name(&self) -> &'static str {
66        "json"
67    }
68}
69
70/// Inlines `LargeBinary` as raw JSON when valid; falls through otherwise.
71#[derive(Debug)]
72struct JsonbPassthroughFactory;
73
74impl EncoderFactory for JsonbPassthroughFactory {
75    fn make_default_encoder<'a>(
76        &self,
77        _field: &'a FieldRef,
78        array: &'a dyn Array,
79        _options: &'a EncoderOptions,
80    ) -> Result<Option<NullableEncoder<'a>>, ArrowError> {
81        if *array.data_type() != DataType::LargeBinary {
82            return Ok(None);
83        }
84        let binary_array = array.as_binary::<i64>();
85        let nulls = binary_array.nulls().cloned();
86        let encoder = LargeBinaryJsonbEncoder {
87            array: binary_array,
88        };
89        Ok(Some(NullableEncoder::new(Box::new(encoder), nulls)))
90    }
91}
92
93struct LargeBinaryJsonbEncoder<'a> {
94    array: &'a arrow_array::LargeBinaryArray,
95}
96
97impl Encoder for LargeBinaryJsonbEncoder<'_> {
98    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
99        let bytes = self.array.value(idx);
100        if serde_json::from_slice::<serde_json::Value>(bytes).is_ok() {
101            out.extend_from_slice(bytes);
102        } else {
103            write!(out, "\"<binary:{} bytes>\"", bytes.len()).unwrap();
104        }
105    }
106}
107
108#[cfg(test)]
109mod tests;
laminar_connectors/schema/json/encoder/mod.rs

laminar_connectors/schema/json/encoder/
mod.rs
