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

arrow_json/reader/
run_end_array.rs
