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

arrow_json/writer/
encoder.rs
