/Users/andrewlamb/Software/arrow-rs/arrow-data/src/transform/variable_size.rs
Line | Count | Source |
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 | | |
18 | | use crate::ArrayData; |
19 | | use arrow_buffer::{ArrowNativeType, MutableBuffer}; |
20 | | use num::traits::AsPrimitive; |
21 | | use num::{CheckedAdd, Integer}; |
22 | | |
23 | | use super::{ |
24 | | Extend, _MutableArrayData, |
25 | | utils::{extend_offsets, get_last_offset}, |
26 | | }; |
27 | | |
28 | | #[inline] |
29 | 9 | fn extend_offset_values<T: ArrowNativeType + AsPrimitive<usize>>( |
30 | 9 | buffer: &mut MutableBuffer, |
31 | 9 | offsets: &[T], |
32 | 9 | values: &[u8], |
33 | 9 | start: usize, |
34 | 9 | len: usize, |
35 | 9 | ) { |
36 | 9 | let start_values = offsets[start].as_(); |
37 | 9 | let end_values: usize = offsets[start + len].as_(); |
38 | 9 | let new_values = &values[start_values..end_values]; |
39 | 9 | buffer.extend_from_slice(new_values); |
40 | 9 | } |
41 | | |
42 | 9 | pub(super) fn build_extend<T: ArrowNativeType + Integer + CheckedAdd + AsPrimitive<usize>>( |
43 | 9 | array: &ArrayData, |
44 | 9 | ) -> Extend<'_> { |
45 | 9 | let offsets = array.buffer::<T>(0); |
46 | 9 | let values = array.buffers()[1].as_slice(); |
47 | 9 | Box::new( |
48 | 9 | move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| { |
49 | 9 | let offset_buffer = &mut mutable.buffer1; |
50 | 9 | let values_buffer = &mut mutable.buffer2; |
51 | | |
52 | | // this is safe due to how offset is built. See details on `get_last_offset` |
53 | 9 | let last_offset = unsafe { get_last_offset(offset_buffer) }; |
54 | | |
55 | 9 | extend_offsets::<T>(offset_buffer, last_offset, &offsets[start..start + len + 1]); |
56 | | // values |
57 | 9 | extend_offset_values::<T>(values_buffer, offsets, values, start, len); |
58 | 9 | }, |
59 | | ) |
60 | 9 | } |
61 | | |
62 | 0 | pub(super) fn extend_nulls<T: ArrowNativeType>(mutable: &mut _MutableArrayData, len: usize) { |
63 | 0 | let offset_buffer = &mut mutable.buffer1; |
64 | | |
65 | | // this is safe due to how offset is built. See details on `get_last_offset` |
66 | 0 | let last_offset: T = unsafe { get_last_offset(offset_buffer) }; |
67 | | |
68 | 0 | (0..len).for_each(|_| offset_buffer.push(last_offset)) |
69 | 0 | } |