arrow_json/reader/
map_array.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.
17
18use std::sync::Arc;
19
20use arrow_array::{ArrayRef, MapArray, StructArray};
21use arrow_buffer::{ArrowNativeType, NullBufferBuilder, OffsetBuffer, ScalarBuffer};
22use arrow_schema::{ArrowError, DataType, FieldRef, Fields};
23
24use crate::reader::tape::{Tape, TapeElement};
25use crate::reader::{ArrayDecoder, DecoderContext};
26
27pub struct MapArrayDecoder {
28    entries_field: FieldRef,
29    key_value_fields: Fields,
30    ordered: bool,
31    keys: Box<dyn ArrayDecoder>,
32    values: Box<dyn ArrayDecoder>,
33    ignore_type_conflicts: bool,
34    is_nullable: bool,
35}
36
37impl MapArrayDecoder {
38    pub fn new(
39        ctx: &DecoderContext,
40        data_type: &DataType,
41        is_nullable: bool,
42    ) -> Result<Self, ArrowError> {
43        let (entries_field, ordered) = match data_type {
44            DataType::Map(_, true) => {
45                return Err(ArrowError::NotYetImplemented(
46                    "Decoding MapArray with sorted fields".to_string(),
47                ));
48            }
49            DataType::Map(f, ordered) => (f.clone(), *ordered),
50            _ => unreachable!(),
51        };
52
53        let key_value_fields = match entries_field.data_type() {
54            DataType::Struct(fields) if fields.len() == 2 => fields.clone(),
55            d => {
56                return Err(ArrowError::InvalidArgumentError(format!(
57                    "MapArray must contain struct with two fields, got {d}"
58                )));
59            }
60        };
61
62        let keys = ctx.make_decoder(&key_value_fields[0], key_value_fields[0].is_nullable())?;
63        let values = ctx.make_decoder(&key_value_fields[1], key_value_fields[1].is_nullable())?;
64
65        Ok(Self {
66            entries_field,
67            key_value_fields,
68            ordered,
69            keys,
70            values,
71            ignore_type_conflicts: ctx.ignore_type_conflicts(),
72            is_nullable,
73        })
74    }
75}
76
77impl ArrayDecoder for MapArrayDecoder {
78    fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef, ArrowError> {
79        let mut offsets = Vec::with_capacity(pos.len() + 1);
80        offsets.push(0);
81
82        let mut key_pos = Vec::with_capacity(pos.len());
83        let mut value_pos = Vec::with_capacity(pos.len());
84
85        let mut nulls = self.is_nullable.then(|| NullBufferBuilder::new(pos.len()));
86
87        for p in pos.iter().copied() {
88            let end_idx = match (tape.get(p), nulls.as_mut()) {
89                (TapeElement::StartObject(end_idx), None) => end_idx,
90                (TapeElement::StartObject(end_idx), Some(nulls)) => {
91                    nulls.append_non_null();
92                    end_idx
93                }
94                (TapeElement::Null, Some(nulls)) => {
95                    nulls.append_null();
96                    p + 1
97                }
98                (_, Some(nulls)) if self.ignore_type_conflicts => {
99                    nulls.append_null();
100                    p + 1
101                }
102                _ => return Err(tape.error(p, "{")),
103            };
104
105            let mut cur_idx = p + 1;
106            while cur_idx < end_idx {
107                let key = cur_idx;
108                let value = tape.next(key, "map key")?;
109                cur_idx = tape.next(value, "map value")?;
110
111                key_pos.push(key);
112                value_pos.push(value);
113            }
114
115            let offset = i32::from_usize(key_pos.len()).ok_or_else(|| {
116                ArrowError::JsonError("offset overflow decoding MapArray".to_string())
117            })?;
118            offsets.push(offset)
119        }
120
121        assert_eq!(key_pos.len(), value_pos.len());
122
123        let key_array = self.keys.decode(tape, &key_pos)?;
124        let value_array = self.values.decode(tape, &value_pos)?;
125
126        let entries = StructArray::try_new_with_length(
127            self.key_value_fields.clone(),
128            vec![key_array, value_array],
129            None,
130            key_pos.len(),
131        )?;
132
133        let nulls = nulls.as_mut().and_then(|x| x.finish());
134        // SAFETY: offsets are built monotonically starting from 0
135        let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets)) };
136
137        let array = MapArray::try_new(
138            self.entries_field.clone(),
139            offsets,
140            entries,
141            nulls,
142            self.ordered,
143        )?;
144        Ok(Arc::new(array))
145    }
146}
arrow_json/reader/map_array.rs

arrow_json/reader/
map_array.rs
