|
19 | 19 | use crate::{Array, ArrayRef, make_array}; |
20 | 20 | use arrow_buffer::bit_chunk_iterator::{BitChunkIterator, BitChunks}; |
21 | 21 | use arrow_buffer::buffer::NullBuffer; |
22 | | -use arrow_buffer::{BooleanBuffer, MutableBuffer, ScalarBuffer}; |
| 22 | +use arrow_buffer::{BooleanBuffer, Buffer, MutableBuffer, ScalarBuffer}; |
23 | 23 | use arrow_data::{ArrayData, ArrayDataBuilder}; |
24 | 24 | use arrow_schema::{ArrowError, DataType, UnionFields, UnionMode}; |
25 | 25 | /// Contains the `UnionArray` type. |
@@ -680,32 +680,36 @@ impl UnionArray { |
680 | 680 |
|
681 | 681 | impl From<ArrayData> for UnionArray { |
682 | 682 | fn from(data: ArrayData) -> Self { |
683 | | - let (fields, mode) = match data.data_type() { |
684 | | - DataType::Union(fields, mode) => (fields, *mode), |
| 683 | + let (data_type, len, _nulls, offset, buffers, child_data) = data.into_parts(); |
| 684 | + |
| 685 | + let (fields, mode) = match &data_type { |
| 686 | + DataType::Union(fields, mode) => (fields, mode), |
685 | 687 | d => panic!("UnionArray expected ArrayData with type Union got {d}"), |
686 | 688 | }; |
| 689 | + |
687 | 690 | let (type_ids, offsets) = match mode { |
688 | | - UnionMode::Sparse => ( |
689 | | - ScalarBuffer::new(data.buffers()[0].clone(), data.offset(), data.len()), |
690 | | - None, |
691 | | - ), |
692 | | - UnionMode::Dense => ( |
693 | | - ScalarBuffer::new(data.buffers()[0].clone(), data.offset(), data.len()), |
694 | | - Some(ScalarBuffer::new( |
695 | | - data.buffers()[1].clone(), |
696 | | - data.offset(), |
697 | | - data.len(), |
698 | | - )), |
699 | | - ), |
| 691 | + UnionMode::Sparse => { |
| 692 | + let [buffer]: [Buffer; 1] = buffers.try_into().expect("1 buffer for type_ids"); |
| 693 | + (ScalarBuffer::new(buffer, offset, len), None) |
| 694 | + } |
| 695 | + UnionMode::Dense => { |
| 696 | + let [type_ids_buffer, offsets_buffer]: [Buffer; 2] = buffers |
| 697 | + .try_into() |
| 698 | + .expect("2 buffers for type_ids and offsets"); |
| 699 | + ( |
| 700 | + ScalarBuffer::new(type_ids_buffer, offset, len), |
| 701 | + Some(ScalarBuffer::new(offsets_buffer, offset, len)), |
| 702 | + ) |
| 703 | + } |
700 | 704 | }; |
701 | 705 |
|
702 | 706 | let max_id = fields.iter().map(|(i, _)| i).max().unwrap_or_default() as usize; |
703 | 707 | let mut boxed_fields = vec![None; max_id + 1]; |
704 | | - for (cd, (field_id, _)) in data.child_data().iter().zip(fields.iter()) { |
705 | | - boxed_fields[field_id as usize] = Some(make_array(cd.clone())); |
| 708 | + for (cd, (field_id, _)) in child_data.into_iter().zip(fields.iter()) { |
| 709 | + boxed_fields[field_id as usize] = Some(make_array(cd)); |
706 | 710 | } |
707 | 711 | Self { |
708 | | - data_type: data.data_type().clone(), |
| 712 | + data_type, |
709 | 713 | type_ids, |
710 | 714 | offsets, |
711 | 715 | fields: boxed_fields, |
|
0 commit comments