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    /// Compile the installed Arrow writer against an empty array of each field type.
26    ///
27    /// # Errors
28    /// Rejects unsupported native encodings without serializing or publishing a record.
29    pub fn validate_schema(schema: &SchemaRef) -> SchemaResult<()> {
30        laminar_core::schema_binding::SchemaBinding::logical(
31            "json",
32            laminar_core::schema_binding::SchemaDirection::Sink,
33            laminar_core::schema_binding::SchemaOrigin::Query,
34            schema.as_ref().clone(),
35        )
36        .map_err(|error| SchemaError::Incompatible(error.to_string()))?;
37        let options =
38            EncoderOptions::default().with_encoder_factory(Arc::new(JsonbPassthroughFactory));
39        for field in schema.fields() {
40            let array = arrow_array::new_empty_array(field.data_type());
41            arrow_json::writer::make_encoder(field, array.as_ref(), &options).map_err(|error| {
42                SchemaError::Incompatible(format!(
43                    "JSON writer field '{}' cannot encode {}: {error}",
44                    field.name(),
45                    field.data_type()
46                ))
47            })?;
48        }
49        Ok(())
50    }
51
52    /// Creates a new JSON encoder for the given schema.
53    #[must_use]
54    pub fn new(schema: SchemaRef) -> Self {
55        Self { schema }
56    }
57}
58
59impl FormatEncoder for JsonEncoder {
60    fn input_schema(&self) -> SchemaRef {
61        self.schema.clone()
62    }
63
64    fn encode_batch(&self, batch: &RecordBatch) -> SchemaResult<Vec<Vec<u8>>> {
65        if batch.num_rows() == 0 {
66            return Ok(Vec::new());
67        }
68
69        let mut buf = Vec::new();
70        {
71            let mut writer = WriterBuilder::new()
72                .with_explicit_nulls(true)
73                .with_encoder_factory(Arc::new(JsonbPassthroughFactory))
74                .build::<_, LineDelimited>(&mut buf);
75            writer
76                .write(batch)
77                .map_err(|e| SchemaError::DecodeError(format!("JSON encode error: {e}")))?;
78            writer
79                .finish()
80                .map_err(|e| SchemaError::DecodeError(format!("JSON finish error: {e}")))?;
81        }
82
83        let output: Vec<Vec<u8>> = buf
84            .split(|&b| b == b'\n')
85            .filter(|line| !line.is_empty())
86            .map(<[u8]>::to_vec)
87            .collect();
88
89        Ok(output)
90    }
91
92    fn format_name(&self) -> &'static str {
93        "json"
94    }
95}
96
97/// Inlines `LargeBinary` as raw JSON when valid; falls through otherwise.
98#[derive(Debug)]
99struct JsonbPassthroughFactory;
100
101impl EncoderFactory for JsonbPassthroughFactory {
102    fn make_default_encoder<'a>(
103        &self,
104        _field: &'a FieldRef,
105        array: &'a dyn Array,
106        _options: &'a EncoderOptions,
107    ) -> Result<Option<NullableEncoder<'a>>, ArrowError> {
108        if *array.data_type() != DataType::LargeBinary {
109            return Ok(None);
110        }
111        let binary_array = array.as_binary::<i64>();
112        let nulls = binary_array.nulls().cloned();
113        let encoder = LargeBinaryJsonbEncoder {
114            array: binary_array,
115        };
116        Ok(Some(NullableEncoder::new(Box::new(encoder), nulls)))
117    }
118}
119
120struct LargeBinaryJsonbEncoder<'a> {
121    array: &'a arrow_array::LargeBinaryArray,
122}
123
124impl Encoder for LargeBinaryJsonbEncoder<'_> {
125    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
126        let bytes = self.array.value(idx);
127        if serde_json::from_slice::<serde_json::Value>(bytes).is_ok() {
128            out.extend_from_slice(bytes);
129        } else {
130            write!(out, "\"<binary:{} bytes>\"", bytes.len()).unwrap();
131        }
132    }
133}
134
135#[cfg(test)]
136mod tests;
laminar_connectors/schema/json/encoder/mod.rs

laminar_connectors/schema/json/encoder/
mod.rs
