arrow_json/reader/
list_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::marker::PhantomData;
19use std::sync::Arc;
20
21use arrow_array::{
22    ArrayRef, FixedSizeListArray, GenericListArray, GenericListViewArray, OffsetSizeTrait,
23};
24use arrow_buffer::{NullBufferBuilder, OffsetBuffer, ScalarBuffer};
25use arrow_schema::{ArrowError, DataType, FieldRef};
26
27use crate::reader::tape::{Tape, TapeElement};
28use crate::reader::{ArrayDecoder, DecoderContext};
29
30pub type ListArrayDecoder<O> = ListLikeArrayDecoder<O, false>;
31pub type ListViewArrayDecoder<O> = ListLikeArrayDecoder<O, true>;
32
33pub struct ListLikeArrayDecoder<O, const IS_VIEW: bool> {
34    field: FieldRef,
35    decoder: Box<dyn ArrayDecoder>,
36    phantom: PhantomData<O>,
37    ignore_type_conflicts: bool,
38    is_nullable: bool,
39}
40
41impl<O: OffsetSizeTrait, const IS_VIEW: bool> ListLikeArrayDecoder<O, IS_VIEW> {
42    pub fn new(
43        ctx: &DecoderContext,
44        data_type: &DataType,
45        is_nullable: bool,
46    ) -> Result<Self, ArrowError> {
47        let field = match (IS_VIEW, data_type) {
48            (false, DataType::List(f)) if !O::IS_LARGE => f,
49            (false, DataType::LargeList(f)) if O::IS_LARGE => f,
50            (true, DataType::ListView(f)) if !O::IS_LARGE => f,
51            (true, DataType::LargeListView(f)) if O::IS_LARGE => f,
52            _ => unreachable!(),
53        };
54        let decoder = ctx.make_decoder(field.data_type(), field.is_nullable())?;
55
56        Ok(Self {
57            field: field.clone(),
58            decoder,
59            phantom: Default::default(),
60            ignore_type_conflicts: ctx.ignore_type_conflicts(),
61            is_nullable,
62        })
63    }
64}
65
66impl<O: OffsetSizeTrait, const IS_VIEW: bool> ArrayDecoder for ListLikeArrayDecoder<O, IS_VIEW> {
67    fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef, ArrowError> {
68        let mut child_pos = Vec::with_capacity(pos.len());
69        let mut offsets = Vec::with_capacity(pos.len() + 1);
70        offsets.push(O::from_usize(0).unwrap());
71
72        let mut nulls = self.is_nullable.then(|| NullBufferBuilder::new(pos.len()));
73
74        for p in pos {
75            let end_idx = match (tape.get(*p), nulls.as_mut()) {
76                (TapeElement::StartList(end_idx), None) => end_idx,
77                (TapeElement::StartList(end_idx), Some(nulls)) => {
78                    nulls.append_non_null();
79                    end_idx
80                }
81                (TapeElement::Null, Some(nulls)) => {
82                    nulls.append_null();
83                    *p + 1
84                }
85                (_, Some(nulls)) if self.ignore_type_conflicts => {
86                    nulls.append_null();
87                    *p + 1
88                }
89                _ => return Err(tape.error(*p, "[")),
90            };
91
92            let mut cur_idx = *p + 1;
93            while cur_idx < end_idx {
94                child_pos.push(cur_idx);
95
96                // Advance to next field
97                cur_idx = tape.next(cur_idx, "list value")?;
98            }
99
100            let offset = O::from_usize(child_pos.len()).ok_or_else(|| {
101                ArrowError::JsonError(format!("offset overflow decoding {}ListArray", O::PREFIX))
102            })?;
103            offsets.push(offset);
104        }
105
106        let values = self.decoder.decode(tape, &child_pos)?;
107        let nulls = nulls.as_mut().and_then(|x| x.finish());
108
109        if IS_VIEW {
110            let mut sizes = Vec::with_capacity(offsets.len() - 1);
111            for i in 1..offsets.len() {
112                sizes.push(offsets[i] - offsets[i - 1]);
113            }
114            offsets.pop();
115            let array = GenericListViewArray::<O>::try_new(
116                self.field.clone(),
117                ScalarBuffer::from(offsets),
118                ScalarBuffer::from(sizes),
119                values,
120                nulls,
121            )?;
122            Ok(Arc::new(array))
123        } else {
124            // SAFETY: offsets are built monotonically starting from 0
125            let offsets = unsafe { OffsetBuffer::<O>::new_unchecked(ScalarBuffer::from(offsets)) };
126
127            let array = GenericListArray::<O>::try_new(self.field.clone(), offsets, values, nulls)?;
128            Ok(Arc::new(array))
129        }
130    }
131}
132
133pub struct FixedSizeListArrayDecoder {
134    field: FieldRef,
135    size: i32,
136    decoder: Box<dyn ArrayDecoder>,
137    ignore_type_conflicts: bool,
138    is_nullable: bool,
139}
140
141impl FixedSizeListArrayDecoder {
142    pub fn new(
143        ctx: &DecoderContext,
144        data_type: &DataType,
145        is_nullable: bool,
146    ) -> Result<Self, ArrowError> {
147        let (field, size) = match data_type {
148            DataType::FixedSizeList(f, s) => (f, *s),
149            _ => unreachable!(),
150        };
151        let decoder = ctx.make_decoder(field.data_type(), field.is_nullable())?;
152
153        Ok(Self {
154            field: field.clone(),
155            size,
156            decoder,
157            ignore_type_conflicts: ctx.ignore_type_conflicts(),
158            is_nullable,
159        })
160    }
161}
162
163impl ArrayDecoder for FixedSizeListArrayDecoder {
164    fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef, ArrowError> {
165        let expected = self.size as usize;
166        let mut child_pos = Vec::with_capacity(pos.len() * expected);
167
168        let mut nulls = self.is_nullable.then(|| NullBufferBuilder::new(pos.len()));
169
170        for p in pos {
171            let end_idx = match (tape.get(*p), nulls.as_mut()) {
172                (TapeElement::StartList(end_idx), None) => end_idx,
173                (TapeElement::StartList(end_idx), Some(nulls)) => {
174                    nulls.append_non_null();
175                    end_idx
176                }
177                (TapeElement::Null, Some(nulls)) => {
178                    nulls.append_null();
179                    child_pos.resize(child_pos.len() + expected, 0);
180                    continue;
181                }
182                (_, Some(nulls)) if self.ignore_type_conflicts => {
183                    nulls.append_null();
184                    child_pos.resize(child_pos.len() + expected, 0);
185                    continue;
186                }
187                _ => return Err(tape.error(*p, "[")),
188            };
189
190            let child_start = child_pos.len();
191            let mut cur_idx = *p + 1;
192            while cur_idx < end_idx {
193                child_pos.push(cur_idx);
194                cur_idx = tape.next(cur_idx, "fixed-size list value")?;
195            }
196
197            let actual = child_pos.len() - child_start;
198            if actual != expected {
199                return Err(ArrowError::JsonError(format!(
200                    "Incorrect number of elements for FixedSizeList, \
201                     expected {expected} but got {actual}"
202                )));
203            }
204        }
205
206        let values = self.decoder.decode(tape, &child_pos)?;
207        let nulls = nulls.as_mut().and_then(|x| x.finish());
208
209        let array = FixedSizeListArray::try_new_with_length(
210            self.field.clone(),
211            self.size,
212            values,
213            nulls,
214            pos.len(),
215        )?;
216        Ok(Arc::new(array))
217    }
218}
arrow_json/reader/list_array.rs

arrow_json/reader/
list_array.rs
