arrow_json/writer/
encoder.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17use std::io::Write;
18use std::sync::Arc;
19
20use crate::StructMode;
21use arrow_array::cast::AsArray;
22use arrow_array::types::*;
23use arrow_array::*;
24use arrow_buffer::{ArrowNativeType, NullBuffer, OffsetBuffer, ScalarBuffer};
25use arrow_cast::display::{ArrayFormatter, FormatOptions};
26use arrow_schema::{ArrowError, DataType, FieldRef};
27use half::f16;
28use lexical_core::FormattedSize;
29use serde_core::Serializer;
30
31/// Configuration options for the JSON encoder.
32#[derive(Debug, Clone, Default)]
33pub struct EncoderOptions {
34    /// Whether to include nulls in the output or elide them.
35    explicit_nulls: bool,
36    /// Whether to encode structs as JSON objects or JSON arrays of their values.
37    struct_mode: StructMode,
38    /// An optional hook for customizing encoding behavior.
39    encoder_factory: Option<Arc<dyn EncoderFactory>>,
40    /// Optional date format for date arrays
41    date_format: Option<String>,
42    /// Optional datetime format for datetime arrays
43    datetime_format: Option<String>,
44    /// Optional timestamp format for timestamp arrays
45    timestamp_format: Option<String>,
46    /// Optional timestamp format for timestamp with timezone arrays
47    timestamp_tz_format: Option<String>,
48    /// Optional time format for time arrays
49    time_format: Option<String>,
50}
51
52impl EncoderOptions {
53    /// Set whether to include nulls in the output or elide them.
54    pub fn with_explicit_nulls(mut self, explicit_nulls: bool) -> Self {
55        self.explicit_nulls = explicit_nulls;
56        self
57    }
58
59    /// Set whether to encode structs as JSON objects or JSON arrays of their values.
60    pub fn with_struct_mode(mut self, struct_mode: StructMode) -> Self {
61        self.struct_mode = struct_mode;
62        self
63    }
64
65    /// Set an optional hook for customizing encoding behavior.
66    pub fn with_encoder_factory(mut self, encoder_factory: Arc<dyn EncoderFactory>) -> Self {
67        self.encoder_factory = Some(encoder_factory);
68        self
69    }
70
71    /// Get whether to include nulls in the output or elide them.
72    pub fn explicit_nulls(&self) -> bool {
73        self.explicit_nulls
74    }
75
76    /// Get whether to encode structs as JSON objects or JSON arrays of their values.
77    pub fn struct_mode(&self) -> StructMode {
78        self.struct_mode
79    }
80
81    /// Get the optional hook for customizing encoding behavior.
82    pub fn encoder_factory(&self) -> Option<&Arc<dyn EncoderFactory>> {
83        self.encoder_factory.as_ref()
84    }
85
86    /// Set the JSON file's date format
87    pub fn with_date_format(mut self, format: String) -> Self {
88        self.date_format = Some(format);
89        self
90    }
91
92    /// Get the JSON file's date format if set, defaults to RFC3339
93    pub fn date_format(&self) -> Option<&str> {
94        self.date_format.as_deref()
95    }
96
97    /// Set the JSON file's datetime format
98    pub fn with_datetime_format(mut self, format: String) -> Self {
99        self.datetime_format = Some(format);
100        self
101    }
102
103    /// Get the JSON file's datetime format if set, defaults to RFC3339
104    pub fn datetime_format(&self) -> Option<&str> {
105        self.datetime_format.as_deref()
106    }
107
108    /// Set the JSON file's time format
109    pub fn with_time_format(mut self, format: String) -> Self {
110        self.time_format = Some(format);
111        self
112    }
113
114    /// Get the JSON file's datetime time if set, defaults to RFC3339
115    pub fn time_format(&self) -> Option<&str> {
116        self.time_format.as_deref()
117    }
118
119    /// Set the JSON file's timestamp format
120    pub fn with_timestamp_format(mut self, format: String) -> Self {
121        self.timestamp_format = Some(format);
122        self
123    }
124
125    /// Get the JSON file's timestamp format if set, defaults to RFC3339
126    pub fn timestamp_format(&self) -> Option<&str> {
127        self.timestamp_format.as_deref()
128    }
129
130    /// Set the JSON file's timestamp tz format
131    pub fn with_timestamp_tz_format(mut self, tz_format: String) -> Self {
132        self.timestamp_tz_format = Some(tz_format);
133        self
134    }
135
136    /// Get the JSON file's timestamp tz format if set, defaults to RFC3339
137    pub fn timestamp_tz_format(&self) -> Option<&str> {
138        self.timestamp_tz_format.as_deref()
139    }
140}
141
142/// Creates custom encoders for specific data types when writing JSON data.
143///
144/// This trait allows customizing JSON encoding for specific data types,
145/// or adding new encoders for unsupported or custom data types.
146///
147/// You can register an implementation of this trait using
148/// [`WriterBuilder::with_encoder_factory`].
149///
150/// [`WriterBuilder::with_encoder_factory`]: crate::writer::WriterBuilder::with_encoder_factory
151///
152/// # Example: Encode a `BinaryArray` as an array of integers
153///
154/// This example encodes each `BinaryArray` value as an array of integers. For
155/// example, the bytes `b"abc"` would be encoded as `[97, 98, 99]`. See the
156/// example on [`DecoderFactory`] for how to decode this back into a
157/// `BinaryArray`.
158///
159/// [`DecoderFactory`]: crate::DecoderFactory
160///
161/// ```
162/// use std::io::Write;
163/// use arrow_array::{ArrayAccessor, Array, BinaryArray, Float64Array, RecordBatch};
164/// use arrow_array::cast::AsArray;
165/// use arrow_schema::{DataType, Field, Schema, FieldRef};
166/// use arrow_json::{writer::{WriterBuilder, JsonArray, NullableEncoder}, StructMode};
167/// use arrow_json::{Encoder, EncoderFactory, EncoderOptions};
168/// use arrow_schema::ArrowError;
169/// use std::sync::Arc;
170/// use serde_json::json;
171/// use serde_json::Value;
172///
173/// /// Encoder for `BinaryArray` that encodes the bytes as an array of integers.
174/// /// For example, the bytes `b"abc"` would be encoded as `[97, 98, 99]`.
175/// struct IntArrayBinaryEncoder<B> {
176///     array: B,
177/// }
178///
179/// impl<'a, B> Encoder for IntArrayBinaryEncoder<B>
180/// where
181///     B: ArrayAccessor<Item = &'a [u8]>,
182/// {
183///     fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
184///         out.push(b'[');
185///         let child = self.array.value(idx);
186///         for (idx, byte) in child.iter().enumerate() {
187///             write!(out, "{byte}").unwrap();
188///             if idx < child.len() - 1 {
189///                 out.push(b',');
190///             }
191///         }
192///         out.push(b']');
193///     }
194/// }
195///
196/// /// Factory that creates an `IntArrayBinaryEncoder` for `BinaryArray` types.
197/// #[derive(Debug)]
198/// struct IntArrayBinaryEncoderFactory;
199///
200/// impl EncoderFactory for IntArrayBinaryEncoderFactory {
201///     fn make_default_encoder<'a>(
202///         &self,
203///         _field: &'a FieldRef,
204///         array: &'a dyn Array,
205///         _options: &'a EncoderOptions,
206///     ) -> Result<Option<NullableEncoder<'a>>, ArrowError> {
207///         match array.data_type() {
208///             DataType::Binary => {
209///                 let array = array.as_binary::<i32>();
210///                 let encoder = IntArrayBinaryEncoder { array };
211///                 let array_encoder = Box::new(encoder) as Box<dyn Encoder + 'a>;
212///                 let nulls = array.nulls().cloned();
213///                 Ok(Some(NullableEncoder::new(array_encoder, nulls)))
214///             }
215///             _ => Ok(None),
216///         }
217///     }
218/// }
219///
220/// // The input has two columns:
221/// // bytes: [b"a", null, b"b"]
222/// // float: [1.0, 2.3, null]
223/// let binary_array = BinaryArray::from_iter([Some(b"a".as_slice()), None, Some(b"b".as_slice())]);
224/// let float_array = Float64Array::from(vec![Some(1.0), Some(2.3), None]);
225/// let fields = vec![
226///     Field::new("bytes", DataType::Binary, true),
227///     Field::new("float", DataType::Float64, true),
228/// ];
229/// let batch = RecordBatch::try_new(
230///     Arc::new(Schema::new(fields)),
231///     vec![
232///         Arc::new(binary_array) as Arc<dyn Array>,
233///         Arc::new(float_array) as Arc<dyn Array>,
234///     ],
235/// )
236/// .unwrap();
237///
238/// // write the record batch to JSON using the custom encoder factory
239/// // and then reparse into serde_json::Value to verify the output
240/// let json_value: Value = {
241///     let mut buf = Vec::new();
242///     let mut writer = WriterBuilder::new()
243///         .with_encoder_factory(Arc::new(IntArrayBinaryEncoderFactory))
244///         .build::<_, JsonArray>(&mut buf);
245///     writer.write_batches(&[&batch]).unwrap();
246///     writer.finish().unwrap();
247///     serde_json::from_slice(&buf).unwrap()
248/// };
249///
250/// let expected = json!([
251///     {"bytes": [97], "float": 1.0},
252///     {"float": 2.3},
253///     {"bytes": [98]},
254/// ]);
255///
256/// assert_eq!(json_value, expected);
257/// ```
258pub trait EncoderFactory: std::fmt::Debug + Send + Sync {
259    /// Make an encoder that overrides the default encoder for a specific field and array or provides an encoder for a custom data type.
260    /// This can be used to override how e.g. binary data is encoded so that it is an encoded string or an array of integers.
261    ///
262    /// Note that the type of the field may not match the type of the array: for dictionary arrays unless the top-level dictionary is handled this
263    /// will be called again for the keys and values of the dictionary, at which point the field type will still be the outer dictionary type but the
264    /// array will have a different type.
265    /// For example, `field` might have the type `Dictionary(i32, Utf8)` but `array` will be `Utf8`.
266    fn make_default_encoder<'a>(
267        &self,
268        _field: &'a FieldRef,
269        _array: &'a dyn Array,
270        _options: &'a EncoderOptions,
271    ) -> Result<Option<NullableEncoder<'a>>, ArrowError> {
272        Ok(None)
273    }
274}
275
276/// An encoder + a null buffer.
277/// This is packaged together into a wrapper struct to minimize dynamic dispatch for null checks.
278pub struct NullableEncoder<'a> {
279    encoder: Box<dyn Encoder + 'a>,
280    nulls: Option<NullBuffer>,
281}
282
283impl<'a> NullableEncoder<'a> {
284    /// Create a new encoder with a null buffer.
285    #[inline]
286    pub fn new(encoder: Box<dyn Encoder + 'a>, nulls: Option<NullBuffer>) -> Self {
287        Self { encoder, nulls }
288    }
289
290    /// Encode the value at index `idx` to `out`.
291    #[inline]
292    pub fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
293        self.encoder.encode(idx, out)
294    }
295
296    /// Returns whether the value at index `idx` is null.
297    #[inline]
298    pub fn is_null(&self, idx: usize) -> bool {
299        match self.nulls {
300            Some(ref nulls) => nulls.is_null(idx),
301            None => false,
302        }
303    }
304
305    /// Returns whether the encoder has any nulls.
306    #[inline]
307    pub fn has_nulls(&self) -> bool {
308        match self.nulls {
309            Some(ref nulls) => nulls.null_count() > 0,
310            None => false,
311        }
312    }
313}
314
315impl Encoder for NullableEncoder<'_> {
316    #[inline]
317    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
318        self.encoder.encode(idx, out)
319    }
320}
321
322/// A trait to format array values as JSON values
323///
324/// Nullability is handled by the caller to allow encoding nulls implicitly, i.e. `{}` instead of `{"a": null}`
325pub trait Encoder {
326    /// Encode the non-null value at index `idx` to `out`.
327    ///
328    /// The behaviour is unspecified if `idx` corresponds to a null index.
329    fn encode(&mut self, idx: usize, out: &mut Vec<u8>);
330}
331
332/// Creates an encoder for the given array and field.
333///
334/// This first calls the EncoderFactory if one is provided, and then falls back to the default encoders.
335pub fn make_encoder<'a>(
336    field: &'a FieldRef,
337    array: &'a dyn Array,
338    options: &'a EncoderOptions,
339) -> Result<NullableEncoder<'a>, ArrowError> {
340    macro_rules! primitive_helper {
341        ($t:ty) => {{
342            let array = array.as_primitive::<$t>();
343            let nulls = array.nulls().cloned();
344            NullableEncoder::new(Box::new(PrimitiveEncoder::new(array)), nulls)
345        }};
346    }
347
348    if let Some(factory) = options.encoder_factory()
349        && let Some(encoder) = factory.make_default_encoder(field, array, options)?
350    {
351        return Ok(encoder);
352    }
353
354    let nulls = array.nulls().cloned();
355    let encoder = downcast_integer! {
356        array.data_type() => (primitive_helper),
357        DataType::Float16 => primitive_helper!(Float16Type),
358        DataType::Float32 => primitive_helper!(Float32Type),
359        DataType::Float64 => primitive_helper!(Float64Type),
360        DataType::Boolean => {
361            let array = array.as_boolean();
362            NullableEncoder::new(Box::new(BooleanEncoder(array)), array.nulls().cloned())
363        }
364        DataType::Null => NullableEncoder::new(Box::new(NullEncoder), array.logical_nulls()),
365        DataType::Utf8 => {
366            let array = array.as_string::<i32>();
367            NullableEncoder::new(Box::new(StringEncoder(array)), array.nulls().cloned())
368        }
369        DataType::LargeUtf8 => {
370            let array = array.as_string::<i64>();
371            NullableEncoder::new(Box::new(StringEncoder(array)), array.nulls().cloned())
372        }
373        DataType::Utf8View => {
374            let array = array.as_string_view();
375            NullableEncoder::new(Box::new(StringViewEncoder(array)), array.nulls().cloned())
376        }
377        DataType::BinaryView => {
378            let array = array.as_binary_view();
379            NullableEncoder::new(Box::new(BinaryViewEncoder(array)), array.nulls().cloned())
380        }
381        DataType::List(_) => {
382            let array = array.as_list::<i32>();
383            NullableEncoder::new(Box::new(ListLikeEncoder::try_new(field, array, options)?), array.nulls().cloned())
384        }
385        DataType::LargeList(_) => {
386            let array = array.as_list::<i64>();
387            NullableEncoder::new(Box::new(ListLikeEncoder::try_new(field, array, options)?), array.nulls().cloned())
388        }
389        DataType::ListView(_) => {
390            let array = array.as_list_view::<i32>();
391            NullableEncoder::new(Box::new(ListLikeEncoder::try_new(field, array, options)?), array.nulls().cloned())
392        }
393        DataType::LargeListView(_) => {
394            let array = array.as_list_view::<i64>();
395            NullableEncoder::new(Box::new(ListLikeEncoder::try_new(field, array, options)?), array.nulls().cloned())
396        }
397        DataType::FixedSizeList(_, _) => {
398            let array = array.as_fixed_size_list();
399            NullableEncoder::new(Box::new(ListLikeEncoder::try_new(field, array, options)?), array.nulls().cloned())
400        }
401
402        DataType::Dictionary(_, _) => downcast_dictionary_array! {
403            array => {
404                NullableEncoder::new(Box::new(DictionaryEncoder::try_new(field, array, options)?), array.nulls().cloned())
405            },
406            _ => unreachable!()
407        }
408
409        DataType::RunEndEncoded(_, _) => downcast_run_array! {
410            array => {
411                NullableEncoder::new(
412                    Box::new(RunEndEncodedEncoder::try_new(field, array, options)?),
413                    array.logical_nulls(),
414                )
415            },
416            _ => unreachable!()
417        }
418
419        DataType::Map(_, _) => {
420            let array = array.as_map();
421            NullableEncoder::new(Box::new(MapEncoder::try_new(field, array, options)?), array.nulls().cloned())
422        }
423
424        DataType::FixedSizeBinary(_) => {
425            let array = array.as_fixed_size_binary();
426            NullableEncoder::new(Box::new(BinaryEncoder::new(array)) as _, array.nulls().cloned())
427        }
428
429        DataType::Binary => {
430            let array: &BinaryArray = array.as_binary();
431            NullableEncoder::new(Box::new(BinaryEncoder::new(array)), array.nulls().cloned())
432        }
433
434        DataType::LargeBinary => {
435            let array: &LargeBinaryArray = array.as_binary();
436            NullableEncoder::new(Box::new(BinaryEncoder::new(array)), array.nulls().cloned())
437        }
438
439        DataType::Struct(fields) => {
440            let array = array.as_struct();
441            let encoders = fields.iter().zip(array.columns()).map(|(field, array)| {
442                let encoder = make_encoder(field, array, options)?;
443
444                // For typical ASCII names, this will be the exact length (includes 2x quotes, 1x colon).
445                let mut field_name = Vec::with_capacity(field.name().len() + 3);
446                encode_string(field.name(), &mut field_name);
447                field_name.push(b':');
448
449                Ok(FieldEncoder {
450                    field_name,
451                    encoder,
452                })
453            }).collect::<Result<Vec<_>, ArrowError>>()?;
454
455            let encoder = StructArrayEncoder{
456                encoders,
457                explicit_nulls: options.explicit_nulls(),
458                struct_mode: options.struct_mode(),
459            };
460            let nulls = array.nulls().cloned();
461            NullableEncoder::new(Box::new(encoder) as Box<dyn Encoder + 'a>, nulls)
462        }
463        DataType::Decimal32(_, _) | DataType::Decimal64(_, _) | DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => {
464            let options = FormatOptions::new().with_display_error(true);
465            let formatter = JsonArrayFormatter::new(ArrayFormatter::try_new(array, &options)?);
466            NullableEncoder::new(Box::new(RawArrayFormatter(formatter)) as Box<dyn Encoder + 'a>, nulls)
467        }
468        d => match d.is_temporal() {
469            true => {
470                // Note: the implementation of Encoder for ArrayFormatter assumes it does not produce
471                // characters that would need to be escaped within a JSON string, e.g. `'"'`.
472                // If support for user-provided format specifications is added, this assumption
473                // may need to be revisited
474                let fops = FormatOptions::new().with_display_error(true)
475                .with_date_format(options.date_format.as_deref())
476                .with_datetime_format(options.datetime_format.as_deref())
477                .with_timestamp_format(options.timestamp_format.as_deref())
478                .with_timestamp_tz_format(options.timestamp_tz_format.as_deref())
479                .with_time_format(options.time_format.as_deref());
480
481                let formatter = ArrayFormatter::try_new(array, &fops)?;
482                let formatter = JsonArrayFormatter::new(formatter);
483                NullableEncoder::new(Box::new(formatter) as Box<dyn Encoder + 'a>, nulls)
484            }
485            false => return Err(ArrowError::JsonError(format!(
486                "Unsupported data type for JSON encoding: {d:?}",
487            )))
488        }
489    };
490
491    Ok(encoder)
492}
493
494fn encode_string(s: &str, out: &mut Vec<u8>) {
495    let mut serializer = serde_json::Serializer::new(out);
496    serializer.serialize_str(s).unwrap();
497}
498
499fn encode_binary(bytes: &[u8], out: &mut Vec<u8>) {
500    out.push(b'"');
501    for byte in bytes {
502        write!(out, "{byte:02x}").unwrap();
503    }
504    out.push(b'"');
505}
506
507struct FieldEncoder<'a> {
508    field_name: Vec<u8>,
509    encoder: NullableEncoder<'a>,
510}
511
512impl FieldEncoder<'_> {
513    #[inline]
514    fn is_null(&self, idx: usize) -> bool {
515        self.encoder.is_null(idx)
516    }
517}
518
519struct StructArrayEncoder<'a> {
520    encoders: Vec<FieldEncoder<'a>>,
521    explicit_nulls: bool,
522    struct_mode: StructMode,
523}
524
525impl Encoder for StructArrayEncoder<'_> {
526    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
527        match self.struct_mode {
528            StructMode::ObjectOnly => out.push(b'{'),
529            StructMode::ListOnly => out.push(b'['),
530        }
531        let mut is_first = true;
532        // Nulls can only be dropped in explicit mode
533        let drop_nulls = (self.struct_mode == StructMode::ObjectOnly) && !self.explicit_nulls;
534
535        for field_encoder in &mut self.encoders {
536            let is_null = field_encoder.is_null(idx);
537            if is_null && drop_nulls {
538                continue;
539            }
540
541            if !is_first {
542                out.push(b',');
543            }
544            is_first = false;
545
546            if self.struct_mode == StructMode::ObjectOnly {
547                out.extend_from_slice(&field_encoder.field_name);
548            }
549
550            if is_null {
551                out.extend_from_slice(b"null");
552            } else {
553                field_encoder.encoder.encode(idx, out);
554            }
555        }
556        match self.struct_mode {
557            StructMode::ObjectOnly => out.push(b'}'),
558            StructMode::ListOnly => out.push(b']'),
559        }
560    }
561}
562
563trait PrimitiveEncode: ArrowNativeType {
564    type Buffer;
565
566    // Workaround https://github.com/rust-lang/rust/issues/61415
567    fn init_buffer() -> Self::Buffer;
568
569    /// Encode the primitive value as bytes, returning a reference to that slice.
570    ///
571    /// `buf` is temporary space that may be used
572    fn encode(self, buf: &mut Self::Buffer) -> &[u8];
573}
574
575macro_rules! integer_encode {
576    ($($t:ty),*) => {
577        $(
578            impl PrimitiveEncode for $t {
579                type Buffer = [u8; Self::FORMATTED_SIZE];
580
581                fn init_buffer() -> Self::Buffer {
582                    [0; Self::FORMATTED_SIZE]
583                }
584
585                fn encode(self, buf: &mut Self::Buffer) -> &[u8] {
586                    lexical_core::write(self, buf)
587                }
588            }
589        )*
590    };
591}
592integer_encode!(i8, i16, i32, i64, u8, u16, u32, u64);
593
594macro_rules! float_encode {
595    ($($t:ty),*) => {
596        $(
597            impl PrimitiveEncode for $t {
598                type Buffer = [u8; Self::FORMATTED_SIZE];
599
600                fn init_buffer() -> Self::Buffer {
601                    [0; Self::FORMATTED_SIZE]
602                }
603
604                fn encode(self, buf: &mut Self::Buffer) -> &[u8] {
605                    if self.is_infinite() || self.is_nan() {
606                        b"null"
607                    } else {
608                        lexical_core::write(self, buf)
609                    }
610                }
611            }
612        )*
613    };
614}
615float_encode!(f32, f64);
616
617impl PrimitiveEncode for f16 {
618    type Buffer = <f32 as PrimitiveEncode>::Buffer;
619
620    fn init_buffer() -> Self::Buffer {
621        f32::init_buffer()
622    }
623
624    fn encode(self, buf: &mut Self::Buffer) -> &[u8] {
625        self.to_f32().encode(buf)
626    }
627}
628
629struct PrimitiveEncoder<N: PrimitiveEncode> {
630    values: ScalarBuffer<N>,
631    buffer: N::Buffer,
632}
633
634impl<N: PrimitiveEncode> PrimitiveEncoder<N> {
635    fn new<P: ArrowPrimitiveType<Native = N>>(array: &PrimitiveArray<P>) -> Self {
636        Self {
637            values: array.values().clone(),
638            buffer: N::init_buffer(),
639        }
640    }
641}
642
643impl<N: PrimitiveEncode> Encoder for PrimitiveEncoder<N> {
644    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
645        out.extend_from_slice(self.values[idx].encode(&mut self.buffer));
646    }
647}
648
649struct BooleanEncoder<'a>(&'a BooleanArray);
650
651impl Encoder for BooleanEncoder<'_> {
652    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
653        match self.0.value(idx) {
654            true => out.extend_from_slice(b"true"),
655            false => out.extend_from_slice(b"false"),
656        }
657    }
658}
659
660struct StringEncoder<'a, O: OffsetSizeTrait>(&'a GenericStringArray<O>);
661
662impl<O: OffsetSizeTrait> Encoder for StringEncoder<'_, O> {
663    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
664        encode_string(self.0.value(idx), out);
665    }
666}
667
668struct StringViewEncoder<'a>(&'a StringViewArray);
669
670impl Encoder for StringViewEncoder<'_> {
671    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
672        encode_string(self.0.value(idx), out);
673    }
674}
675
676struct BinaryViewEncoder<'a>(&'a BinaryViewArray);
677
678impl Encoder for BinaryViewEncoder<'_> {
679    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
680        encode_binary(self.0.value(idx), out);
681    }
682}
683
684struct ListLikeEncoder<'a, L: ListLikeArray> {
685    list_array: &'a L,
686    encoder: NullableEncoder<'a>,
687}
688
689impl<'a, L: ListLikeArray> ListLikeEncoder<'a, L> {
690    fn try_new(
691        field: &'a FieldRef,
692        array: &'a L,
693        options: &'a EncoderOptions,
694    ) -> Result<Self, ArrowError> {
695        let encoder = make_encoder(field, array.values().as_ref(), options)?;
696        Ok(Self {
697            list_array: array,
698            encoder,
699        })
700    }
701}
702
703impl<L: ListLikeArray> Encoder for ListLikeEncoder<'_, L> {
704    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
705        let range = self.list_array.element_range(idx);
706        let start = range.start;
707        let end = range.end;
708        out.push(b'[');
709        if self.encoder.has_nulls() {
710            for idx in start..end {
711                if idx != start {
712                    out.push(b',')
713                }
714                if self.encoder.is_null(idx) {
715                    out.extend_from_slice(b"null");
716                } else {
717                    self.encoder.encode(idx, out);
718                }
719            }
720        } else {
721            for idx in start..end {
722                if idx != start {
723                    out.push(b',')
724                }
725                self.encoder.encode(idx, out);
726            }
727        }
728        out.push(b']');
729    }
730}
731
732struct DictionaryEncoder<'a, K: ArrowDictionaryKeyType> {
733    keys: ScalarBuffer<K::Native>,
734    encoder: NullableEncoder<'a>,
735}
736
737impl<'a, K: ArrowDictionaryKeyType> DictionaryEncoder<'a, K> {
738    fn try_new(
739        field: &'a FieldRef,
740        array: &'a DictionaryArray<K>,
741        options: &'a EncoderOptions,
742    ) -> Result<Self, ArrowError> {
743        let encoder = make_encoder(field, array.values().as_ref(), options)?;
744
745        Ok(Self {
746            keys: array.keys().values().clone(),
747            encoder,
748        })
749    }
750}
751
752impl<K: ArrowDictionaryKeyType> Encoder for DictionaryEncoder<'_, K> {
753    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
754        self.encoder.encode(self.keys[idx].as_usize(), out)
755    }
756}
757
758struct RunEndEncodedEncoder<'a, R: RunEndIndexType> {
759    run_array: &'a RunArray<R>,
760    encoder: NullableEncoder<'a>,
761}
762
763impl<'a, R: RunEndIndexType> RunEndEncodedEncoder<'a, R> {
764    fn try_new(
765        field: &'a FieldRef,
766        array: &'a RunArray<R>,
767        options: &'a EncoderOptions,
768    ) -> Result<Self, ArrowError> {
769        let encoder = make_encoder(field, array.values().as_ref(), options)?;
770        Ok(Self {
771            run_array: array,
772            encoder,
773        })
774    }
775}
776
777impl<R: RunEndIndexType> Encoder for RunEndEncodedEncoder<'_, R> {
778    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
779        let physical_idx = self.run_array.get_physical_index(idx);
780        self.encoder.encode(physical_idx, out)
781    }
782}
783
784/// A newtype wrapper around [`ArrayFormatter`] to keep our usage of it private and not implement `Encoder` for the public type
785struct JsonArrayFormatter<'a> {
786    formatter: ArrayFormatter<'a>,
787}
788
789impl<'a> JsonArrayFormatter<'a> {
790    fn new(formatter: ArrayFormatter<'a>) -> Self {
791        Self { formatter }
792    }
793}
794
795impl Encoder for JsonArrayFormatter<'_> {
796    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
797        out.push(b'"');
798        // Should be infallible
799        // Note: We are making an assumption that the formatter does not produce characters that require escaping
800        let _ = write!(out, "{}", self.formatter.value(idx));
801        out.push(b'"')
802    }
803}
804
805/// A newtype wrapper around [`JsonArrayFormatter`] that skips surrounding the value with `"`
806struct RawArrayFormatter<'a>(JsonArrayFormatter<'a>);
807
808impl Encoder for RawArrayFormatter<'_> {
809    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
810        let _ = write!(out, "{}", self.0.formatter.value(idx));
811    }
812}
813
814struct NullEncoder;
815
816impl Encoder for NullEncoder {
817    fn encode(&mut self, _idx: usize, _out: &mut Vec<u8>) {
818        unreachable!()
819    }
820}
821
822struct MapEncoder<'a> {
823    offsets: OffsetBuffer<i32>,
824    keys: NullableEncoder<'a>,
825    values: NullableEncoder<'a>,
826    explicit_nulls: bool,
827}
828
829impl<'a> MapEncoder<'a> {
830    fn try_new(
831        field: &'a FieldRef,
832        array: &'a MapArray,
833        options: &'a EncoderOptions,
834    ) -> Result<Self, ArrowError> {
835        let values = array.values();
836        let keys = array.keys();
837
838        if !matches!(
839            keys.data_type(),
840            DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
841        ) {
842            return Err(ArrowError::JsonError(format!(
843                "Only UTF8 keys supported by JSON MapArray Writer: got {:?}",
844                keys.data_type()
845            )));
846        }
847
848        let keys = make_encoder(field, keys, options)?;
849        let values = make_encoder(field, values, options)?;
850
851        // We sanity check nulls as these are currently not enforced by MapArray (#1697)
852        if keys.has_nulls() {
853            return Err(ArrowError::InvalidArgumentError(
854                "Encountered nulls in MapArray keys".to_string(),
855            ));
856        }
857
858        if array.entries().nulls().is_some_and(|x| x.null_count() != 0) {
859            return Err(ArrowError::InvalidArgumentError(
860                "Encountered nulls in MapArray entries".to_string(),
861            ));
862        }
863
864        Ok(Self {
865            offsets: array.offsets().clone(),
866            keys,
867            values,
868            explicit_nulls: options.explicit_nulls(),
869        })
870    }
871}
872
873impl Encoder for MapEncoder<'_> {
874    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
875        let end = self.offsets[idx + 1].as_usize();
876        let start = self.offsets[idx].as_usize();
877
878        let mut is_first = true;
879
880        out.push(b'{');
881
882        for idx in start..end {
883            let is_null = self.values.is_null(idx);
884            if is_null && !self.explicit_nulls {
885                continue;
886            }
887
888            if !is_first {
889                out.push(b',');
890            }
891            is_first = false;
892
893            self.keys.encode(idx, out);
894            out.push(b':');
895
896            if is_null {
897                out.extend_from_slice(b"null");
898            } else {
899                self.values.encode(idx, out);
900            }
901        }
902        out.push(b'}');
903    }
904}
905
906/// New-type wrapper for encoding the binary types in arrow: `Binary`, `LargeBinary`
907/// and `FixedSizeBinary` as hex strings in JSON.
908struct BinaryEncoder<B>(B);
909
910impl<'a, B> BinaryEncoder<B>
911where
912    B: ArrayAccessor<Item = &'a [u8]>,
913{
914    fn new(array: B) -> Self {
915        Self(array)
916    }
917}
918
919impl<'a, B> Encoder for BinaryEncoder<B>
920where
921    B: ArrayAccessor<Item = &'a [u8]>,
922{
923    fn encode(&mut self, idx: usize, out: &mut Vec<u8>) {
924        out.push(b'"');
925        for byte in self.0.value(idx) {
926            // this write is infallible
927            write!(out, "{byte:02x}").unwrap();
928        }
929        out.push(b'"');
930    }
931}
arrow_json/writer/encoder.rs

arrow_json/writer/
encoder.rs
