arrow_json/reader/
run_end_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;
19
20use arrow_array::types::RunEndIndexType;
21use arrow_array::{Array, ArrayRef, PrimitiveArray, make_array, new_empty_array};
22use arrow_buffer::{ArrowNativeType, ScalarBuffer};
23use arrow_data::transform::MutableArrayData;
24use arrow_data::{ArrayData, ArrayDataBuilder};
25use arrow_schema::{ArrowError, DataType};
26
27use crate::reader::tape::Tape;
28use crate::reader::{ArrayDecoder, DecoderContext};
29
30pub struct RunEndEncodedArrayDecoder<R> {
31    data_type: DataType,
32    decoder: Box<dyn ArrayDecoder>,
33    phantom: PhantomData<R>,
34}
35
36impl<R: RunEndIndexType> RunEndEncodedArrayDecoder<R> {
37    pub fn new(
38        ctx: &DecoderContext,
39        data_type: &DataType,
40        is_nullable: bool,
41    ) -> Result<Self, ArrowError> {
42        let values_field = match data_type {
43            DataType::RunEndEncoded(_, v) => v,
44            _ => unreachable!(),
45        };
46        let decoder = ctx.make_decoder(
47            values_field.data_type(),
48            values_field.is_nullable() || is_nullable,
49        )?;
50
51        Ok(Self {
52            data_type: data_type.clone(),
53            decoder,
54            phantom: Default::default(),
55        })
56    }
57}
58
59impl<R: RunEndIndexType + Send> ArrayDecoder for RunEndEncodedArrayDecoder<R> {
60    fn decode(&mut self, tape: &Tape<'_>, pos: &[u32]) -> Result<ArrayRef, ArrowError> {
61        let len = pos.len();
62        if len == 0 {
63            return Ok(new_empty_array(&self.data_type));
64        }
65
66        let flat_data = self.decoder.decode(tape, pos)?.to_data();
67
68        let mut run_ends: Vec<R::Native> = Vec::new();
69        let mut mutable = MutableArrayData::new(vec![&flat_data], false, len);
70
71        let mut run_start = 0;
72        for i in 1..len {
73            if !same_run(&flat_data, run_start, i) {
74                let run_end = R::Native::from_usize(i).ok_or_else(|| {
75                    ArrowError::JsonError(format!(
76                        "Run end value {i} exceeds {:?} range",
77                        R::DATA_TYPE
78                    ))
79                })?;
80                run_ends.push(run_end);
81                mutable.extend(0, run_start, run_start + 1);
82                run_start = i;
83            }
84        }
85        let run_end = R::Native::from_usize(len).ok_or_else(|| {
86            ArrowError::JsonError(format!(
87                "Run end value {len} exceeds {:?} range",
88                R::DATA_TYPE
89            ))
90        })?;
91        run_ends.push(run_end);
92        mutable.extend(0, run_start, run_start + 1);
93
94        let values_data = mutable.freeze();
95        let run_ends_data =
96            PrimitiveArray::<R>::new(ScalarBuffer::from(run_ends), None).into_data();
97
98        let data = ArrayDataBuilder::new(self.data_type.clone())
99            .len(len)
100            .add_child_data(run_ends_data)
101            .add_child_data(values_data);
102
103        // Safety:
104        // run_ends are strictly increasing with the last value equal to len,
105        // and values has the same length as run_ends
106        Ok(make_array(unsafe { data.build_unchecked() }))
107    }
108}
109
110fn same_run(data: &ArrayData, i: usize, j: usize) -> bool {
111    let null_i = data.is_null(i);
112    let null_j = data.is_null(j);
113    if null_i != null_j {
114        return false;
115    }
116    if null_i {
117        return true;
118    }
119    data.slice(i, 1) == data.slice(j, 1)
120}
arrow_json/reader/run_end_array.rs

arrow_json/reader/
run_end_array.rs
