arrow_data/transform/
mod.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
18//! Low-level array data abstractions.
19//!
20//! Provides utilities for creating, manipulating, and converting Arrow arrays
21//! made of primitive types, strings, and nested types.
22
23use super::{ArrayData, ArrayDataBuilder, ByteView, data::new_buffers};
24use crate::bit_mask::set_bits;
25use arrow_buffer::buffer::{BooleanBuffer, NullBuffer};
26use arrow_buffer::{ArrowNativeType, Buffer, IntervalMonthDayNano, MutableBuffer, bit_util, i256};
27use arrow_schema::{ArrowError, DataType, IntervalUnit, UnionMode};
28use half::f16;
29use num_integer::Integer;
30use std::mem;
31
32mod boolean;
33mod fixed_binary;
34mod fixed_size_list;
35mod list;
36mod list_view;
37mod null;
38mod primitive;
39mod run;
40mod structure;
41mod union;
42mod utils;
43mod variable_size;
44
45type ExtendNullBits<'a> = Box<dyn Fn(&mut _MutableArrayData, usize, usize) + 'a>;
46// function that extends `[start..start+len]` to the mutable array.
47// this is dynamic because different data_types influence how buffers and children are extended.
48type Extend<'a> =
49    Box<dyn Fn(&mut _MutableArrayData, usize, usize, usize) -> Result<(), ArrowError> + 'a>;
50
51type ExtendNulls = Box<dyn Fn(&mut _MutableArrayData, usize) -> Result<(), ArrowError>>;
52
53/// A mutable [ArrayData] that knows how to freeze itself into an [ArrayData].
54/// This is just a data container.
55#[derive(Debug)]
56struct _MutableArrayData<'a> {
57    pub data_type: DataType,
58    pub null_count: usize,
59
60    pub len: usize,
61    pub null_buffer: Option<MutableBuffer>,
62
63    // arrow specification only allows up to 3 buffers (2 ignoring the nulls above).
64    // Thus, we place them in the stack to avoid bound checks and greater data locality.
65    pub buffer1: MutableBuffer,
66    pub buffer2: MutableBuffer,
67    pub child_data: Vec<MutableArrayData<'a>>,
68}
69
70impl _MutableArrayData<'_> {
71    fn null_buffer(&mut self) -> &mut MutableBuffer {
72        self.null_buffer
73            .as_mut()
74            .expect("MutableArrayData not nullable")
75    }
76}
77
78fn build_extend_null_bits(array: &ArrayData, use_nulls: bool) -> ExtendNullBits<'_> {
79    if let Some(nulls) = array.nulls() {
80        let bytes = nulls.validity();
81        Box::new(move |mutable, start, len| {
82            let mutable_len = mutable.len;
83            let out = mutable.null_buffer();
84            utils::resize_for_bits(out, mutable_len + len);
85            mutable.null_count += set_bits(
86                out.as_slice_mut(),
87                bytes,
88                mutable_len,
89                nulls.offset() + start,
90                len,
91            );
92        })
93    } else if use_nulls {
94        Box::new(|mutable, _, len| {
95            let mutable_len = mutable.len;
96            let out = mutable.null_buffer();
97            utils::resize_for_bits(out, mutable_len + len);
98            let write_data = out.as_slice_mut();
99            (0..len).for_each(|i| {
100                bit_util::set_bit(write_data, mutable_len + i);
101            });
102        })
103    } else {
104        Box::new(|_, _, _| {})
105    }
106}
107
108/// Efficiently create an [ArrayData] from one or more existing [ArrayData]s by
109/// copying chunks.
110///
111/// The main use case of this struct is to perform unary operations to arrays of
112/// arbitrary types, such as `filter` and `take`.
113///
114/// # Example
115/// ```
116/// use arrow_buffer::Buffer;
117/// use arrow_data::ArrayData;
118/// use arrow_data::transform::MutableArrayData;
119/// use arrow_schema::DataType;
120/// fn i32_array(values: &[i32]) -> ArrayData {
121///   ArrayData::try_new(DataType::Int32, values.len(), None, 0, vec![Buffer::from_slice_ref(values)], vec![]).unwrap()
122/// }
123/// let arr1  = i32_array(&[1, 2, 3, 4, 5]);
124/// let arr2  = i32_array(&[6, 7, 8, 9, 10]);
125/// // Create a mutable array for copying values from arr1 and arr2, with a capacity for 6 elements
126/// let capacity = 3 * std::mem::size_of::<i32>();
127/// let mut mutable = MutableArrayData::new(vec![&arr1, &arr2], false, 10);
128/// // Copy the first 3 elements from arr1
129/// mutable.extend(0, 0, 3);
130/// // Copy the last 3 elements from arr2
131/// mutable.extend(1, 2, 5);
132/// // Complete the MutableArrayData into a new ArrayData
133/// let frozen = mutable.freeze();
134/// assert_eq!(frozen, i32_array(&[1, 2, 3, 8, 9, 10]));
135/// ```
136pub struct MutableArrayData<'a> {
137    /// Input arrays: the data being read FROM.
138    ///
139    /// Note this is "dead code" because all actual references to the arrays are
140    /// stored in closures for extending values and nulls.
141    #[allow(dead_code)]
142    arrays: Vec<&'a ArrayData>,
143
144    /// In progress output array: The data being written TO
145    ///
146    /// Note these fields are in a separate struct, [_MutableArrayData], as they
147    /// cannot be in [MutableArrayData] itself due to mutability invariants (interior
148    /// mutability): [MutableArrayData] contains a function that can only mutate
149    /// [_MutableArrayData], not [MutableArrayData] itself
150    data: _MutableArrayData<'a>,
151
152    /// The child data of the `Array` in Dictionary arrays.
153    ///
154    /// This is not stored in `_MutableArrayData` because these values are
155    /// constant and only needed at the end, when freezing [_MutableArrayData].
156    dictionary: Option<ArrayData>,
157
158    /// Variadic data buffers referenced by views.
159    ///
160    /// Note this this is not stored in `_MutableArrayData` because these values
161    /// are constant and only needed at the end, when freezing
162    /// [_MutableArrayData]
163    variadic_data_buffers: Vec<Buffer>,
164
165    /// function used to extend output array with values from input arrays.
166    ///
167    /// This function's lifetime is bound to the input arrays because it reads
168    /// values from them.
169    extend_values: Vec<Extend<'a>>,
170
171    /// function used to extend the output array with nulls from input arrays.
172    ///
173    /// This function's lifetime is bound to the input arrays because it reads
174    /// nulls from it.
175    extend_null_bits: Vec<ExtendNullBits<'a>>,
176
177    /// function used to extend the output array with null elements.
178    ///
179    /// This function is independent of the arrays and therefore has no lifetime.
180    extend_nulls: ExtendNulls,
181}
182
183impl std::fmt::Debug for MutableArrayData<'_> {
184    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
185        // ignores the closures.
186        f.debug_struct("MutableArrayData")
187            .field("data", &self.data)
188            .finish()
189    }
190}
191
192/// Builds an extend that adds `offset` to the source primitive
193/// Additionally validates that `max` fits into the
194/// the underlying primitive returning None if not
195fn build_extend_dictionary(array: &ArrayData, offset: usize, max: usize) -> Option<Extend<'_>> {
196    macro_rules! validate_and_build {
197        ($dt: ty) => {{
198            // `max` is the merged dictionary length; the largest key index is
199            // `max - 1`, so the key type only needs to hold `max - 1` (e.g. 256
200            // values use keys 0..=255, which fit in u8).
201            let _: $dt = max.saturating_sub(1).try_into().ok()?;
202            let offset: $dt = offset.try_into().ok()?;
203            Some(primitive::build_extend_with_offset(array, offset))
204        }};
205    }
206    match array.data_type() {
207        DataType::Dictionary(child_data_type, _) => match child_data_type.as_ref() {
208            DataType::UInt8 => validate_and_build!(u8),
209            DataType::UInt16 => validate_and_build!(u16),
210            DataType::UInt32 => validate_and_build!(u32),
211            DataType::UInt64 => validate_and_build!(u64),
212            DataType::Int8 => validate_and_build!(i8),
213            DataType::Int16 => validate_and_build!(i16),
214            DataType::Int32 => validate_and_build!(i32),
215            DataType::Int64 => validate_and_build!(i64),
216            _ => unreachable!(),
217        },
218        _ => None,
219    }
220}
221
222/// Builds an extend that adds `buffer_offset` to any buffer indices encountered
223fn build_extend_view(array: &ArrayData, buffer_offset: u32) -> Extend<'_> {
224    let views = array.buffer::<u128>(0);
225    Box::new(
226        move |mutable: &mut _MutableArrayData, _, start: usize, len: usize| {
227            mutable
228                .buffer1
229                .extend(views[start..start + len].iter().map(|v| {
230                    let len = *v as u32;
231                    if len <= 12 {
232                        return *v; // Stored inline
233                    }
234                    let mut view = ByteView::from(*v);
235                    view.buffer_index += buffer_offset;
236                    view.into()
237                }));
238            Ok(())
239        },
240    )
241}
242
243fn build_extend(array: &ArrayData) -> Extend<'_> {
244    match array.data_type() {
245        DataType::Null => null::build_extend(array),
246        DataType::Boolean => boolean::build_extend(array),
247        DataType::UInt8 => primitive::build_extend::<u8>(array),
248        DataType::UInt16 => primitive::build_extend::<u16>(array),
249        DataType::UInt32 => primitive::build_extend::<u32>(array),
250        DataType::UInt64 => primitive::build_extend::<u64>(array),
251        DataType::Int8 => primitive::build_extend::<i8>(array),
252        DataType::Int16 => primitive::build_extend::<i16>(array),
253        DataType::Int32 => primitive::build_extend::<i32>(array),
254        DataType::Int64 => primitive::build_extend::<i64>(array),
255        DataType::Float32 => primitive::build_extend::<f32>(array),
256        DataType::Float64 => primitive::build_extend::<f64>(array),
257        DataType::Date32 | DataType::Time32(_) | DataType::Interval(IntervalUnit::YearMonth) => {
258            primitive::build_extend::<i32>(array)
259        }
260        DataType::Date64
261        | DataType::Time64(_)
262        | DataType::Timestamp(_, _)
263        | DataType::Duration(_)
264        | DataType::Interval(IntervalUnit::DayTime) => primitive::build_extend::<i64>(array),
265        DataType::Interval(IntervalUnit::MonthDayNano) => {
266            primitive::build_extend::<IntervalMonthDayNano>(array)
267        }
268        DataType::Decimal32(_, _) => primitive::build_extend::<i32>(array),
269        DataType::Decimal64(_, _) => primitive::build_extend::<i64>(array),
270        DataType::Decimal128(_, _) => primitive::build_extend::<i128>(array),
271        DataType::Decimal256(_, _) => primitive::build_extend::<i256>(array),
272        DataType::Utf8 | DataType::Binary => variable_size::build_extend::<i32>(array),
273        DataType::LargeUtf8 | DataType::LargeBinary => variable_size::build_extend::<i64>(array),
274        DataType::BinaryView | DataType::Utf8View => unreachable!("should use build_extend_view"),
275        DataType::Map(_, _) | DataType::List(_) => list::build_extend::<i32>(array),
276        DataType::LargeList(_) => list::build_extend::<i64>(array),
277        DataType::ListView(_) => list_view::build_extend::<i32>(array),
278        DataType::LargeListView(_) => list_view::build_extend::<i64>(array),
279        DataType::Dictionary(_, _) => unreachable!("should use build_extend_dictionary"),
280        DataType::Struct(_) => structure::build_extend(array),
281        DataType::FixedSizeBinary(_) => fixed_binary::build_extend(array),
282        DataType::Float16 => primitive::build_extend::<f16>(array),
283        DataType::FixedSizeList(_, _) => fixed_size_list::build_extend(array),
284        DataType::Union(_, mode) => match mode {
285            UnionMode::Sparse => union::build_extend_sparse(array),
286            UnionMode::Dense => union::build_extend_dense(array),
287        },
288        DataType::RunEndEncoded(_, _) => run::build_extend(array),
289    }
290}
291
292fn build_extend_nulls(data_type: &DataType) -> ExtendNulls {
293    Box::new(match data_type {
294        DataType::Null => null::extend_nulls,
295        DataType::Boolean => boolean::extend_nulls,
296        DataType::UInt8 => primitive::extend_nulls::<u8>,
297        DataType::UInt16 => primitive::extend_nulls::<u16>,
298        DataType::UInt32 => primitive::extend_nulls::<u32>,
299        DataType::UInt64 => primitive::extend_nulls::<u64>,
300        DataType::Int8 => primitive::extend_nulls::<i8>,
301        DataType::Int16 => primitive::extend_nulls::<i16>,
302        DataType::Int32 => primitive::extend_nulls::<i32>,
303        DataType::Int64 => primitive::extend_nulls::<i64>,
304        DataType::Float32 => primitive::extend_nulls::<f32>,
305        DataType::Float64 => primitive::extend_nulls::<f64>,
306        DataType::Date32 | DataType::Time32(_) | DataType::Interval(IntervalUnit::YearMonth) => {
307            primitive::extend_nulls::<i32>
308        }
309        DataType::Date64
310        | DataType::Time64(_)
311        | DataType::Timestamp(_, _)
312        | DataType::Duration(_)
313        | DataType::Interval(IntervalUnit::DayTime) => primitive::extend_nulls::<i64>,
314        DataType::Interval(IntervalUnit::MonthDayNano) => {
315            primitive::extend_nulls::<IntervalMonthDayNano>
316        }
317        DataType::Decimal32(_, _) => primitive::extend_nulls::<i32>,
318        DataType::Decimal64(_, _) => primitive::extend_nulls::<i64>,
319        DataType::Decimal128(_, _) => primitive::extend_nulls::<i128>,
320        DataType::Decimal256(_, _) => primitive::extend_nulls::<i256>,
321        DataType::Utf8 | DataType::Binary => variable_size::extend_nulls::<i32>,
322        DataType::LargeUtf8 | DataType::LargeBinary => variable_size::extend_nulls::<i64>,
323        DataType::BinaryView | DataType::Utf8View => primitive::extend_nulls::<u128>,
324        DataType::Map(_, _) | DataType::List(_) => list::extend_nulls::<i32>,
325        DataType::LargeList(_) => list::extend_nulls::<i64>,
326        DataType::ListView(_) => list_view::extend_nulls::<i32>,
327        DataType::LargeListView(_) => list_view::extend_nulls::<i64>,
328        DataType::Dictionary(child_data_type, _) => match child_data_type.as_ref() {
329            DataType::UInt8 => primitive::extend_nulls::<u8>,
330            DataType::UInt16 => primitive::extend_nulls::<u16>,
331            DataType::UInt32 => primitive::extend_nulls::<u32>,
332            DataType::UInt64 => primitive::extend_nulls::<u64>,
333            DataType::Int8 => primitive::extend_nulls::<i8>,
334            DataType::Int16 => primitive::extend_nulls::<i16>,
335            DataType::Int32 => primitive::extend_nulls::<i32>,
336            DataType::Int64 => primitive::extend_nulls::<i64>,
337            _ => unreachable!(),
338        },
339        DataType::Struct(_) => structure::extend_nulls,
340        DataType::FixedSizeBinary(_) => fixed_binary::extend_nulls,
341        DataType::Float16 => primitive::extend_nulls::<f16>,
342        DataType::FixedSizeList(_, _) => fixed_size_list::extend_nulls,
343        DataType::Union(_, mode) => match mode {
344            UnionMode::Sparse => union::extend_nulls_sparse,
345            UnionMode::Dense => union::extend_nulls_dense,
346        },
347        DataType::RunEndEncoded(_, _) => run::extend_nulls,
348    })
349}
350
351fn preallocate_offset_and_binary_buffer<Offset: ArrowNativeType + Integer>(
352    capacity: usize,
353    binary_size: usize,
354) -> [MutableBuffer; 2] {
355    // offsets
356    let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<Offset>());
357    // safety: `unsafe` code assumes that this buffer is initialized with one element
358    buffer.push(Offset::zero());
359
360    [
361        buffer,
362        MutableBuffer::new(binary_size * mem::size_of::<u8>()),
363    ]
364}
365
366/// Define capacities to pre-allocate for child data or data buffers.
367#[derive(Debug, Clone)]
368pub enum Capacities {
369    /// Binary, Utf8 and LargeUtf8 data types
370    ///
371    /// Defines
372    /// * the capacity of the array offsets
373    /// * the capacity of the binary/ str buffer
374    Binary(usize, Option<usize>),
375    /// List and LargeList data types
376    ///
377    /// Defines
378    /// * the capacity of the array offsets
379    /// * the capacity of the child data
380    List(usize, Option<Box<Capacities>>),
381    /// Struct type
382    ///
383    /// Defines
384    /// * the capacity of the array
385    /// * the capacities of the fields
386    Struct(usize, Option<Vec<Capacities>>),
387    /// Dictionary type
388    ///
389    /// Defines
390    /// * the capacity of the array/keys
391    /// * the capacity of the values
392    Dictionary(usize, Option<Box<Capacities>>),
393    /// Don't preallocate inner buffers and rely on array growth strategy
394    Array(usize),
395}
396
397impl<'a> MutableArrayData<'a> {
398    /// Returns a new [MutableArrayData] with capacity to `capacity` slots and
399    /// specialized to create an [ArrayData] from multiple `arrays`.
400    ///
401    /// # Arguments
402    /// * `arrays` - the source arrays to copy from
403    /// * `use_nulls` - a flag indicating whether the caller intends to call `extend_nulls`.
404    ///   Note: null-handling is enabled automatically if any source array contains nulls.
405    /// * `capacity` - the preallocated capacity of the output array, in slots (number of elements)
406    ///
407    /// if `use_nulls` is `false` and no source arrays contains nulls, calling
408    /// [MutableArrayData::extend_nulls] or [MutableArrayData::try_extend_nulls] will panic.
409    pub fn new(arrays: Vec<&'a ArrayData>, use_nulls: bool, capacity: usize) -> Self {
410        Self::with_capacities(arrays, use_nulls, Capacities::Array(capacity))
411    }
412
413    /// Similar to [MutableArrayData::new], but lets users define the
414    /// preallocated capacities of the array with more granularity.
415    ///
416    /// See [MutableArrayData::new] for more information on the arguments.
417    ///
418    /// # Panics
419    ///
420    /// This function panics if the given `capacities` don't match the data type
421    /// of `arrays`. Or when a [Capacities] variant is not yet supported.
422    pub fn with_capacities(
423        arrays: Vec<&'a ArrayData>,
424        use_nulls: bool,
425        capacities: Capacities,
426    ) -> Self {
427        let data_type = arrays[0].data_type();
428
429        for a in arrays.iter().skip(1) {
430            assert_eq!(
431                data_type,
432                a.data_type(),
433                "Arrays with inconsistent types passed to MutableArrayData"
434            )
435        }
436
437        // if any of the arrays has nulls, insertions from any array requires setting bits
438        // as there is at least one array with nulls.
439        let use_nulls = use_nulls | arrays.iter().any(|array| array.null_count() > 0);
440
441        let mut array_capacity;
442
443        let [buffer1, buffer2] = match (data_type, &capacities) {
444            (
445                DataType::LargeUtf8 | DataType::LargeBinary,
446                Capacities::Binary(capacity, Some(value_cap)),
447            ) => {
448                array_capacity = *capacity;
449                preallocate_offset_and_binary_buffer::<i64>(*capacity, *value_cap)
450            }
451            (DataType::Utf8 | DataType::Binary, Capacities::Binary(capacity, Some(value_cap))) => {
452                array_capacity = *capacity;
453                preallocate_offset_and_binary_buffer::<i32>(*capacity, *value_cap)
454            }
455            (_, Capacities::Array(capacity)) => {
456                array_capacity = *capacity;
457                new_buffers(data_type, *capacity)
458            }
459            (
460                DataType::List(_)
461                | DataType::LargeList(_)
462                | DataType::ListView(_)
463                | DataType::LargeListView(_)
464                | DataType::FixedSizeList(_, _),
465                Capacities::List(capacity, _),
466            ) => {
467                array_capacity = *capacity;
468                new_buffers(data_type, *capacity)
469            }
470            _ => panic!("Capacities: {capacities:?} not yet supported"),
471        };
472
473        let child_data = match &data_type {
474            DataType::Decimal32(_, _)
475            | DataType::Decimal64(_, _)
476            | DataType::Decimal128(_, _)
477            | DataType::Decimal256(_, _)
478            | DataType::Null
479            | DataType::Boolean
480            | DataType::UInt8
481            | DataType::UInt16
482            | DataType::UInt32
483            | DataType::UInt64
484            | DataType::Int8
485            | DataType::Int16
486            | DataType::Int32
487            | DataType::Int64
488            | DataType::Float16
489            | DataType::Float32
490            | DataType::Float64
491            | DataType::Date32
492            | DataType::Date64
493            | DataType::Time32(_)
494            | DataType::Time64(_)
495            | DataType::Duration(_)
496            | DataType::Timestamp(_, _)
497            | DataType::Utf8
498            | DataType::Binary
499            | DataType::LargeUtf8
500            | DataType::LargeBinary
501            | DataType::BinaryView
502            | DataType::Utf8View
503            | DataType::Interval(_)
504            | DataType::FixedSizeBinary(_) => vec![],
505            DataType::Map(_, _)
506            | DataType::List(_)
507            | DataType::LargeList(_)
508            | DataType::ListView(_)
509            | DataType::LargeListView(_) => {
510                let children = arrays
511                    .iter()
512                    .map(|array| &array.child_data()[0])
513                    .collect::<Vec<_>>();
514
515                let capacities =
516                    if let Capacities::List(capacity, ref child_capacities) = capacities {
517                        child_capacities
518                            .clone()
519                            .map(|c| *c)
520                            .unwrap_or(Capacities::Array(capacity))
521                    } else {
522                        Capacities::Array(array_capacity)
523                    };
524
525                vec![MutableArrayData::with_capacities(
526                    children, use_nulls, capacities,
527                )]
528            }
529            // the dictionary type just appends keys and clones the values.
530            DataType::Dictionary(_, _) => vec![],
531            DataType::Struct(fields) => match capacities {
532                Capacities::Struct(capacity, Some(ref child_capacities)) => {
533                    array_capacity = capacity;
534                    (0..fields.len())
535                        .zip(child_capacities)
536                        .map(|(i, child_cap)| {
537                            let child_arrays = arrays
538                                .iter()
539                                .map(|array| &array.child_data()[i])
540                                .collect::<Vec<_>>();
541                            MutableArrayData::with_capacities(
542                                child_arrays,
543                                use_nulls,
544                                child_cap.clone(),
545                            )
546                        })
547                        .collect::<Vec<_>>()
548                }
549                Capacities::Struct(capacity, None) => {
550                    array_capacity = capacity;
551                    (0..fields.len())
552                        .map(|i| {
553                            let child_arrays = arrays
554                                .iter()
555                                .map(|array| &array.child_data()[i])
556                                .collect::<Vec<_>>();
557                            MutableArrayData::new(child_arrays, use_nulls, capacity)
558                        })
559                        .collect::<Vec<_>>()
560                }
561                _ => (0..fields.len())
562                    .map(|i| {
563                        let child_arrays = arrays
564                            .iter()
565                            .map(|array| &array.child_data()[i])
566                            .collect::<Vec<_>>();
567                        MutableArrayData::new(child_arrays, use_nulls, array_capacity)
568                    })
569                    .collect::<Vec<_>>(),
570            },
571            DataType::RunEndEncoded(_, _) => {
572                let run_ends_child = arrays
573                    .iter()
574                    .map(|array| &array.child_data()[0])
575                    .collect::<Vec<_>>();
576                let value_child = arrays
577                    .iter()
578                    .map(|array| &array.child_data()[1])
579                    .collect::<Vec<_>>();
580                vec![
581                    MutableArrayData::new(run_ends_child, false, array_capacity),
582                    MutableArrayData::new(value_child, use_nulls, array_capacity),
583                ]
584            }
585            DataType::FixedSizeList(_, size) => {
586                let children = arrays
587                    .iter()
588                    .map(|array| &array.child_data()[0])
589                    .collect::<Vec<_>>();
590                let capacities =
591                    if let Capacities::List(capacity, ref child_capacities) = capacities {
592                        child_capacities
593                            .clone()
594                            .map(|c| *c)
595                            .unwrap_or(Capacities::Array(capacity * *size as usize))
596                    } else {
597                        Capacities::Array(array_capacity * *size as usize)
598                    };
599                vec![MutableArrayData::with_capacities(
600                    children, use_nulls, capacities,
601                )]
602            }
603            DataType::Union(fields, _) => (0..fields.len())
604                .map(|i| {
605                    let child_arrays = arrays
606                        .iter()
607                        .map(|array| &array.child_data()[i])
608                        .collect::<Vec<_>>();
609                    MutableArrayData::new(child_arrays, use_nulls, array_capacity)
610                })
611                .collect::<Vec<_>>(),
612        };
613
614        // Get the dictionary if any, and if it is a concatenation of multiple
615        let (dictionary, dict_concat) = match &data_type {
616            DataType::Dictionary(_, _) => {
617                // If more than one dictionary, concatenate dictionaries together
618                let dict_concat = !arrays
619                    .windows(2)
620                    .all(|a| a[0].child_data()[0].ptr_eq(&a[1].child_data()[0]));
621
622                match dict_concat {
623                    false => (Some(arrays[0].child_data()[0].clone()), false),
624                    true => {
625                        if let Capacities::Dictionary(_, _) = capacities {
626                            panic!("dictionary capacity not yet supported")
627                        }
628                        let dictionaries: Vec<_> =
629                            arrays.iter().map(|array| &array.child_data()[0]).collect();
630                        let lengths: Vec<_> = dictionaries
631                            .iter()
632                            .map(|dictionary| dictionary.len())
633                            .collect();
634                        let capacity = lengths.iter().sum();
635
636                        let mut mutable = MutableArrayData::new(dictionaries, false, capacity);
637
638                        for (i, len) in lengths.iter().enumerate() {
639                            mutable.try_extend(i, 0, *len).expect(
640                                "extend failed while building dictionary; \
641                                 this is a bug in MutableArrayData",
642                            )
643                        }
644
645                        (Some(mutable.freeze()), true)
646                    }
647                }
648            }
649            _ => (None, false),
650        };
651
652        let variadic_data_buffers = match &data_type {
653            DataType::BinaryView | DataType::Utf8View => arrays
654                .iter()
655                .flat_map(|x| x.buffers().iter().skip(1))
656                .map(Buffer::clone)
657                .collect(),
658            _ => vec![],
659        };
660
661        let extend_nulls = build_extend_nulls(data_type);
662
663        let extend_null_bits = arrays
664            .iter()
665            .map(|array| build_extend_null_bits(array, use_nulls))
666            .collect();
667
668        let null_buffer = use_nulls.then(|| {
669            let null_bytes = bit_util::ceil(array_capacity, 8);
670            MutableBuffer::from_len_zeroed(null_bytes)
671        });
672
673        let extend_values = match &data_type {
674            DataType::Dictionary(_, _) => {
675                let mut next_offset = 0;
676                let extend_values: Result<Vec<_>, _> = arrays
677                    .iter()
678                    .map(|array| {
679                        let offset = next_offset;
680                        let dict_len = array.child_data()[0].len();
681
682                        if dict_concat {
683                            next_offset += dict_len;
684                        }
685
686                        build_extend_dictionary(array, offset, offset + dict_len)
687                            .ok_or(ArrowError::DictionaryKeyOverflowError)
688                    })
689                    .collect();
690
691                extend_values.expect("MutableArrayData::new is infallible")
692            }
693            DataType::BinaryView | DataType::Utf8View => {
694                let mut next_offset = 0u32;
695                arrays
696                    .iter()
697                    .map(|arr| {
698                        let num_data_buffers = (arr.buffers().len() - 1) as u32;
699                        let offset = next_offset;
700                        next_offset = next_offset
701                            .checked_add(num_data_buffers)
702                            .expect("view buffer index overflow");
703                        build_extend_view(arr, offset)
704                    })
705                    .collect()
706            }
707            _ => arrays.iter().map(|array| build_extend(array)).collect(),
708        };
709
710        let data = _MutableArrayData {
711            data_type: data_type.clone(),
712            len: 0,
713            null_count: 0,
714            null_buffer,
715            buffer1,
716            buffer2,
717            child_data,
718        };
719        Self {
720            arrays,
721            data,
722            dictionary,
723            variadic_data_buffers,
724            extend_values,
725            extend_null_bits,
726            extend_nulls,
727        }
728    }
729
730    /// Extends the in progress array with a region of the input arrays, returning an error on
731    /// overflow.
732    ///
733    /// # Arguments
734    /// * `index` - the index of array that you want to copy values from
735    /// * `start` - the start index of the chunk (inclusive)
736    /// * `end` - the end index of the chunk (exclusive)
737    ///
738    /// # Errors
739    /// Returns an error if offset arithmetic overflows the underlying integer type.
740    ///
741    /// # Panic
742    /// This function panics if there is an invalid index,
743    /// i.e. `index` >= the number of source arrays
744    /// or `end` > the length of the `index`th array
745    pub fn try_extend(&mut self, index: usize, start: usize, end: usize) -> Result<(), ArrowError> {
746        let len = end - start;
747        (self.extend_null_bits[index])(&mut self.data, start, len);
748        // Snapshot buffer lengths before attempting the extend so we can roll
749        // back to a consistent state if it fails.
750        let buf1_len = self.data.buffer1.len();
751        let buf2_len = self.data.buffer2.len();
752        if let Err(e) = (self.extend_values[index])(&mut self.data, index, start, len) {
753            // Restore buffers to their pre-call lengths so the array remains
754            // in a valid state for the caller to inspect or retry.
755            self.data.buffer1.truncate(buf1_len);
756            self.data.buffer2.truncate(buf2_len);
757            return Err(e);
758        }
759        self.data.len += len;
760        Ok(())
761    }
762
763    /// Extends the in progress array with a region of the input arrays.
764    ///
765    /// # Panic
766    /// This function panics if there is an invalid index,
767    /// i.e. `index` >= the number of source arrays,
768    /// `end` > the length of the `index`th array,
769    /// or the offset type overflows (e.g. more than 2 GiB in a `StringArray`).
770    #[deprecated(
771        since = "59.0.0",
772        note = "Use `try_extend` which returns an error on overflow instead of panicking"
773    )]
774    pub fn extend(&mut self, index: usize, start: usize, end: usize) {
775        self.try_extend(index, start, end)
776            .expect("extend failed due to offset overflow")
777    }
778
779    /// Extends the in progress array with null elements, ignoring the input arrays, returning an
780    /// error on overflow.
781    ///
782    /// Prefer this over [`extend_nulls`](Self::extend_nulls) to handle cases where the run-end
783    /// counter overflows (relevant for `RunEndEncoded` arrays).
784    ///
785    /// # Panics
786    ///
787    /// Panics if [`MutableArrayData`] not created with `use_nulls` or nullable source arrays
788    pub fn try_extend_nulls(&mut self, len: usize) -> Result<(), ArrowError> {
789        self.data.len += len;
790        let bit_len = bit_util::ceil(self.data.len, 8);
791        let nulls = self.data.null_buffer();
792        nulls.resize(bit_len, 0);
793        self.data.null_count += len;
794        (self.extend_nulls)(&mut self.data, len)?;
795        Ok(())
796    }
797
798    /// Extends the in progress array with null elements, ignoring the input arrays.
799    ///
800    /// # Panics
801    ///
802    /// Panics if [`MutableArrayData`] not created with `use_nulls` or nullable source arrays,
803    /// or if the run-end counter overflows for `RunEndEncoded` arrays.
804    #[deprecated(
805        since = "59.0.0",
806        note = "Use `try_extend_nulls` which returns an error on overflow instead of panicking"
807    )]
808    pub fn extend_nulls(&mut self, len: usize) {
809        self.try_extend_nulls(len)
810            .expect("extend_nulls failed due to overflow")
811    }
812
813    /// Returns the current length
814    #[inline]
815    pub fn len(&self) -> usize {
816        self.data.len
817    }
818
819    /// Returns true if len is 0
820    #[inline]
821    pub fn is_empty(&self) -> bool {
822        self.data.len == 0
823    }
824
825    /// Returns the current null count
826    #[inline]
827    pub fn null_count(&self) -> usize {
828        self.data.null_count
829    }
830
831    /// Creates a [ArrayData] from the in progress array, consuming `self`.
832    pub fn freeze(self) -> ArrayData {
833        unsafe { self.into_builder().build_unchecked() }
834    }
835
836    /// Consume self and returns the in progress array as [`ArrayDataBuilder`].
837    ///
838    /// This is useful for extending the default behavior of MutableArrayData.
839    pub fn into_builder(self) -> ArrayDataBuilder {
840        let data = self.data;
841
842        let buffers = match data.data_type {
843            DataType::Null
844            | DataType::Struct(_)
845            | DataType::FixedSizeList(_, _)
846            | DataType::RunEndEncoded(_, _) => {
847                vec![]
848            }
849            DataType::BinaryView | DataType::Utf8View => {
850                let mut b = self.variadic_data_buffers;
851                b.insert(0, data.buffer1.into());
852                b
853            }
854            DataType::Utf8
855            | DataType::Binary
856            | DataType::LargeUtf8
857            | DataType::LargeBinary
858            | DataType::ListView(_)
859            | DataType::LargeListView(_) => {
860                vec![data.buffer1.into(), data.buffer2.into()]
861            }
862            DataType::Union(_, mode) => {
863                match mode {
864                    // Based on Union's DataTypeLayout
865                    UnionMode::Sparse => vec![data.buffer1.into()],
866                    UnionMode::Dense => vec![data.buffer1.into(), data.buffer2.into()],
867                }
868            }
869            _ => vec![data.buffer1.into()],
870        };
871
872        let child_data = match data.data_type {
873            DataType::Dictionary(_, _) => vec![self.dictionary.unwrap()],
874            _ => data.child_data.into_iter().map(|x| x.freeze()).collect(),
875        };
876
877        let nulls = match data.data_type {
878            // RunEndEncoded, Null, and Union arrays cannot have top-level null bitmasks
879            DataType::RunEndEncoded(_, _) | DataType::Null | DataType::Union(_, _) => None,
880            _ => data
881                .null_buffer
882                .map(|nulls| {
883                    let bools = BooleanBuffer::new(nulls.into(), 0, data.len);
884                    unsafe { NullBuffer::new_unchecked(bools, data.null_count) }
885                })
886                .filter(|n| n.null_count() > 0),
887        };
888
889        ArrayDataBuilder::new(data.data_type)
890            .offset(0)
891            .len(data.len)
892            .nulls(nulls)
893            .buffers(buffers)
894            .child_data(child_data)
895    }
896}
897
898// See arrow/tests/array_transform.rs for tests of transform functionality
899
900#[cfg(test)]
901mod test {
902    use super::*;
903    use arrow_schema::Field;
904    use std::sync::Arc;
905
906    #[test]
907    fn test_list_append_with_capacities() {
908        let array = ArrayData::new_empty(&DataType::List(Arc::new(Field::new(
909            "element",
910            DataType::Int64,
911            false,
912        ))));
913
914        let mutable = MutableArrayData::with_capacities(
915            vec![&array],
916            false,
917            Capacities::List(6, Some(Box::new(Capacities::Array(17)))),
918        );
919
920        // capacities are rounded up to multiples of 64 by MutableBuffer
921        assert_eq!(mutable.data.buffer1.capacity(), 64);
922        assert_eq!(mutable.data.child_data[0].data.buffer1.capacity(), 192);
923    }
924}
arrow_data/transform/mod.rs

arrow_data/transform/
mod.rs
