/Users/andrewlamb/Software/arrow-rs/arrow-data/src/transform/union.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 super::{_MutableArrayData, Extend}; |
19 | | use crate::ArrayData; |
20 | | |
21 | 2 | pub(super) fn build_extend_sparse(array: &ArrayData) -> Extend<'_> { |
22 | 2 | let type_ids = array.buffer::<i8>(0); |
23 | | |
24 | 2 | Box::new( |
25 | 2 | move |mutable: &mut _MutableArrayData, index: usize, start: usize, len: usize| { |
26 | | // extends type_ids |
27 | 2 | mutable |
28 | 2 | .buffer1 |
29 | 2 | .extend_from_slice(&type_ids[start..start + len]); |
30 | | |
31 | 2 | mutable |
32 | 2 | .child_data |
33 | 2 | .iter_mut() |
34 | 2 | .for_each(|child| child.extend(index, start, start + len)) |
35 | 2 | }, |
36 | | ) |
37 | 2 | } |
38 | | |
39 | 6 | pub(super) fn build_extend_dense(array: &ArrayData) -> Extend<'_> { |
40 | 6 | let type_ids = array.buffer::<i8>(0); |
41 | 6 | let offsets = array.buffer::<i32>(1); |
42 | 6 | let arrow_schema::DataType::Union(src_fields, _) = array.data_type() else { |
43 | 0 | unreachable!(); |
44 | | }; |
45 | | |
46 | 6 | Box::new( |
47 | 8 | move |mutable: &mut _MutableArrayData, index: usize, start: usize, len: usize| { |
48 | | // extends type_ids |
49 | 8 | mutable |
50 | 8 | .buffer1 |
51 | 8 | .extend_from_slice(&type_ids[start..start + len]); |
52 | | |
53 | 11 | (8 start8 ..start + len).for_each8 (|i| { |
54 | 11 | let type_id = type_ids[i]; |
55 | 11 | let child_index = src_fields |
56 | 11 | .iter() |
57 | 14 | .position11 (|(r, _)| r == type_id) |
58 | 11 | .expect("invalid union type ID"); |
59 | 11 | let src_offset = offsets[i] as usize; |
60 | 11 | let child_data = &mut mutable.child_data[child_index]; |
61 | 11 | let dst_offset = child_data.len(); |
62 | | |
63 | | // Extend offsets |
64 | 11 | mutable.buffer2.push(dst_offset as i32); |
65 | 11 | mutable.child_data[child_index].extend(index, src_offset, src_offset + 1) |
66 | 11 | }) |
67 | 8 | }, |
68 | | ) |
69 | 6 | } |
70 | | |
71 | 0 | pub(super) fn extend_nulls_dense(_mutable: &mut _MutableArrayData, _len: usize) { |
72 | 0 | panic!("cannot call extend_nulls on UnionArray as cannot infer type"); |
73 | | } |
74 | | |
75 | 0 | pub(super) fn extend_nulls_sparse(_mutable: &mut _MutableArrayData, _len: usize) { |
76 | 0 | panic!("cannot call extend_nulls on UnionArray as cannot infer type"); |
77 | | } |