Skip to main content

arrow_data/
data.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//! Contains [`ArrayData`], a generic representation of Arrow array data which encapsulates
19//! common attributes and operations for Arrow array.
20
21use crate::bit_iterator::BitSliceIterator;
22use arrow_buffer::buffer::{BooleanBuffer, NullBuffer};
23use arrow_buffer::{
24    ArrowNativeType, Buffer, IntervalDayTime, IntervalMonthDayNano, MutableBuffer, bit_util, i256,
25};
26use arrow_schema::{ArrowError, DataType, UnionMode};
27use std::mem;
28use std::ops::Range;
29use std::sync::Arc;
30
31use crate::{equal, validate_binary_view, validate_string_view};
32
33#[inline]
34pub(crate) fn contains_nulls(
35    null_bit_buffer: Option<&NullBuffer>,
36    offset: usize,
37    len: usize,
38) -> bool {
39    match null_bit_buffer {
40        Some(buffer) => {
41            match BitSliceIterator::new(buffer.validity(), buffer.offset() + offset, len).next() {
42                Some((start, end)) => start != 0 || end != len,
43                None => len != 0, // No non-null values
44            }
45        }
46        None => false, // No null buffer
47    }
48}
49
50#[inline]
51pub(crate) fn count_nulls(
52    null_bit_buffer: Option<&NullBuffer>,
53    offset: usize,
54    len: usize,
55) -> usize {
56    if let Some(buf) = null_bit_buffer {
57        let buffer = buf.buffer();
58        len - buffer.count_set_bits_offset(offset + buf.offset(), len)
59    } else {
60        0
61    }
62}
63
64/// creates 2 [`MutableBuffer`]s with a given `capacity` (in slots).
65#[inline]
66pub(crate) fn new_buffers(data_type: &DataType, capacity: usize) -> [MutableBuffer; 2] {
67    let empty_buffer = MutableBuffer::new(0);
68    match data_type {
69        DataType::Null => [empty_buffer, MutableBuffer::new(0)],
70        DataType::Boolean => {
71            let bytes = bit_util::ceil(capacity, 8);
72            let buffer = MutableBuffer::new(bytes);
73            [buffer, empty_buffer]
74        }
75        DataType::UInt8
76        | DataType::UInt16
77        | DataType::UInt32
78        | DataType::UInt64
79        | DataType::Int8
80        | DataType::Int16
81        | DataType::Int32
82        | DataType::Int64
83        | DataType::Float16
84        | DataType::Float32
85        | DataType::Float64
86        | DataType::Decimal32(_, _)
87        | DataType::Decimal64(_, _)
88        | DataType::Decimal128(_, _)
89        | DataType::Decimal256(_, _)
90        | DataType::Date32
91        | DataType::Time32(_)
92        | DataType::Date64
93        | DataType::Time64(_)
94        | DataType::Duration(_)
95        | DataType::Timestamp(_, _)
96        | DataType::Interval(_) => [
97            MutableBuffer::new(capacity * data_type.primitive_width().unwrap()),
98            empty_buffer,
99        ],
100        DataType::Utf8 | DataType::Binary => {
101            let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i32>());
102            // safety: `unsafe` code assumes that this buffer is initialized with one element
103            buffer.push(0i32);
104            [buffer, MutableBuffer::new(capacity * mem::size_of::<u8>())]
105        }
106        DataType::LargeUtf8 | DataType::LargeBinary => {
107            let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i64>());
108            // safety: `unsafe` code assumes that this buffer is initialized with one element
109            buffer.push(0i64);
110            [buffer, MutableBuffer::new(capacity * mem::size_of::<u8>())]
111        }
112        DataType::BinaryView | DataType::Utf8View => [
113            MutableBuffer::new(capacity * mem::size_of::<u128>()),
114            empty_buffer,
115        ],
116        DataType::List(_) | DataType::Map(_, _) => {
117            // offset buffer always starts with a zero
118            let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i32>());
119            buffer.push(0i32);
120            [buffer, empty_buffer]
121        }
122        DataType::ListView(_) => [
123            MutableBuffer::new(capacity * mem::size_of::<i32>()),
124            MutableBuffer::new(capacity * mem::size_of::<i32>()),
125        ],
126        DataType::LargeList(_) => {
127            // offset buffer always starts with a zero
128            let mut buffer = MutableBuffer::new((1 + capacity) * mem::size_of::<i64>());
129            buffer.push(0i64);
130            [buffer, empty_buffer]
131        }
132        DataType::LargeListView(_) => [
133            MutableBuffer::new(capacity * mem::size_of::<i64>()),
134            MutableBuffer::new(capacity * mem::size_of::<i64>()),
135        ],
136        DataType::FixedSizeBinary(size) => {
137            if *size < 0 {
138                panic!("cannot construct buffers from FixedSizeBinary({size})");
139            }
140            [MutableBuffer::new(capacity * *size as usize), empty_buffer]
141        }
142        DataType::Dictionary(k, _) => [
143            MutableBuffer::new(capacity * k.primitive_width().unwrap()),
144            empty_buffer,
145        ],
146        DataType::FixedSizeList(_, _) | DataType::Struct(_) | DataType::RunEndEncoded(_, _) => {
147            [empty_buffer, MutableBuffer::new(0)]
148        }
149        DataType::Union(_, mode) => {
150            let type_ids = MutableBuffer::new(capacity * mem::size_of::<i8>());
151            match mode {
152                UnionMode::Sparse => [type_ids, empty_buffer],
153                UnionMode::Dense => {
154                    let offsets = MutableBuffer::new(capacity * mem::size_of::<i32>());
155                    [type_ids, offsets]
156                }
157            }
158        }
159    }
160}
161
162/// A generic representation of Arrow array data which encapsulates common attributes
163/// and operations for Arrow array.
164///
165/// Specific operations for different arrays types (e.g., primitive, list, struct)
166/// are implemented in `Array`.
167///
168/// # Memory Layout
169///
170/// `ArrayData` has references to one or more underlying data buffers
171/// and optional child ArrayData, depending on type as illustrated
172/// below. Bitmaps are not shown for simplicity but they are stored
173/// similarly to the buffers.
174///
175/// ```text
176///                        offset
177///                       points to
178/// ┌───────────────────┐ start of  ┌───────┐       Different
179/// │                   │   data    │       │     ArrayData may
180/// │ArrayData {        │           │....   │     also refers to
181/// │  data_type: ...   │   ─ ─ ─ ─▶│1234   │  ┌ ─  the same
182/// │  offset: ... ─ ─ ─│─ ┘        │4372   │      underlying
183/// │  len: ...    ─ ─ ─│─ ┐        │4888   │  │     buffer with different offset/len
184/// │  buffers: [       │           │5882   │◀─
185/// │    ...            │  │        │4323   │
186/// │  ]                │   ─ ─ ─ ─▶│4859   │
187/// │  child_data: [    │           │....   │
188/// │    ...            │           │       │
189/// │  ]                │           └───────┘
190/// │}                  │
191/// │                   │            Shared Buffer uses
192/// │               │   │            bytes::Bytes to hold
193/// └───────────────────┘            actual data values
194///           ┌ ─ ─ ┘
195///
196///           ▼
197/// ┌───────────────────┐
198/// │ArrayData {        │
199/// │  ...              │
200/// │}                  │
201/// │                   │
202/// └───────────────────┘
203///
204/// Child ArrayData may also have its own buffers and children
205/// ```
206
207#[derive(Debug, Clone)]
208pub struct ArrayData {
209    /// The data type
210    data_type: DataType,
211
212    /// The number of elements
213    len: usize,
214
215    /// The offset in number of items (not bytes).
216    ///
217    /// The offset applies to [`Self::child_data`] and [`Self::buffers`]. It
218    /// does NOT apply to [`Self::nulls`].
219    ///
220    /// See [`Self::offset()`] for details and diagrams.
221    offset: usize,
222
223    /// The buffers that store the actual data for this array, as defined
224    /// in the [Arrow Spec].
225    ///
226    /// Depending on the array types, [`Self::buffers`] can hold different
227    /// kinds of buffers (e.g., value buffer, value offset buffer) at different
228    /// positions.
229    ///
230    /// The buffer may be larger than needed.  Some items at the beginning may be skipped if
231    /// there is an `offset`.  Some items at the end may be skipped if the buffer is longer than
232    /// we need to satisfy `len`.
233    ///
234    /// [Arrow Spec](https://arrow.apache.org/docs/format/Columnar.html#physical-memory-layout)
235    buffers: Vec<Buffer>,
236
237    /// The child(ren) of this array.
238    ///
239    /// Only non-empty for nested types, such as `ListArray` and
240    /// `StructArray`.
241    ///
242    /// The first logical element in each child element begins at `offset`.
243    ///
244    /// If the child element also has an offset then these offsets are
245    /// cumulative.
246    ///
247    /// See [`Self::child_data()`] and [`Self::offset()`] for details.
248    child_data: Vec<ArrayData>,
249
250    /// The null bitmap.
251    ///
252    /// `None` indicates all values are non-null in this array.
253    ///
254    /// [`Self::offset()`] does not apply to the null bitmap. While the
255    /// BooleanBuffer may be sliced (have its own offset) internally, this
256    /// `NullBuffer` always represents exactly `len` elements.
257    nulls: Option<NullBuffer>,
258}
259
260/// A thread-safe, shared reference to the Arrow array data.
261pub type ArrayDataRef = Arc<ArrayData>;
262
263fn checked_len_plus_offset(
264    data_type: &DataType,
265    len: usize,
266    offset: usize,
267) -> Result<usize, ArrowError> {
268    len.checked_add(offset).ok_or_else(|| {
269        ArrowError::InvalidArgumentError(format!(
270            "Length {len} with offset {offset} overflows usize for {data_type}"
271        ))
272    })
273}
274
275/// Returns the first and last of `offsets`, and the bytes of `values` between
276/// them as a `str`. Returns `None` if the offsets are not valid indexes into
277/// `values` or those bytes are not valid UTF-8.
278fn utf8_span<'a, T: ArrowNativeType>(
279    offsets: &[T],
280    values: &'a [u8],
281) -> Option<(usize, usize, &'a str)> {
282    let first = offsets.first()?.to_usize()?;
283    let last = offsets.last()?.to_usize()?;
284    let values_str = std::str::from_utf8(values.get(first..last)?).ok()?;
285    Some((first, last, values_str))
286}
287
288impl ArrayData {
289    /// Create a new ArrayData instance;
290    ///
291    /// If `null_count` is not specified, the number of nulls in
292    /// null_bit_buffer is calculated.
293    ///
294    /// If the number of nulls is 0 then the null_bit_buffer
295    /// is set to `None`.
296    ///
297    /// # Safety
298    ///
299    /// The input values *must* form a valid Arrow array for
300    /// `data_type`, or undefined behavior can result.
301    ///
302    /// Note: This is a low level API and most users of the arrow
303    /// crate should create arrays using the methods in the `array`
304    /// module.
305    pub unsafe fn new_unchecked(
306        data_type: DataType,
307        len: usize,
308        null_count: Option<usize>,
309        null_bit_buffer: Option<Buffer>,
310        offset: usize,
311        buffers: Vec<Buffer>,
312        child_data: Vec<ArrayData>,
313    ) -> Self {
314        let builder = Self::inner_new_builder(
315            data_type,
316            len,
317            null_count,
318            null_bit_buffer,
319            offset,
320            buffers,
321            child_data,
322        );
323
324        // SAFETY: caller responsible for ensuring data is valid
325        unsafe { builder.build_unchecked() }
326    }
327
328    /// Create a new ArrayData, validating that the provided buffers form a valid
329    /// Arrow array of the specified data type.
330    ///
331    /// If the number of nulls in `null_bit_buffer` is 0 then the null_bit_buffer
332    /// is set to `None`.
333    ///
334    /// Internally this calls through to [`Self::validate_data`]
335    ///
336    /// Note: This is a low level API and most users of the arrow crate should create
337    /// arrays using the builders found in [arrow_array](https://docs.rs/arrow-array)
338    /// or [`ArrayDataBuilder`].
339    ///
340    /// See also [`Self::into_parts`] to recover the fields
341    pub fn try_new(
342        data_type: DataType,
343        len: usize,
344        null_bit_buffer: Option<Buffer>,
345        offset: usize,
346        buffers: Vec<Buffer>,
347        child_data: Vec<ArrayData>,
348    ) -> Result<Self, ArrowError> {
349        let builder = Self::inner_new_builder(
350            data_type,
351            len,
352            None,
353            null_bit_buffer,
354            offset,
355            buffers,
356            child_data,
357        );
358
359        assert!(!builder.skip_validation.get());
360
361        // As the data is not trusted, do a full validation of its contents
362        // We don't need to validate children as we can assume that the
363        // [`ArrayData`] in `child_data` have already been validated through
364        // a call to `ArrayData::try_new` or created using unsafe
365        builder.build()
366    }
367
368    fn inner_new_builder(
369        data_type: DataType,
370        len: usize,
371        null_count: Option<usize>,
372        null_bit_buffer: Option<Buffer>,
373        offset: usize,
374        buffers: Vec<Buffer>,
375        child_data: Vec<ArrayData>,
376    ) -> ArrayDataBuilder {
377        ArrayDataBuilder {
378            data_type,
379            len,
380            null_count,
381            null_bit_buffer,
382            nulls: None,
383            offset,
384            buffers,
385            child_data,
386            align_buffers: false,
387            skip_validation: UnsafeFlag::new(),
388        }
389    }
390
391    /// Return the constituent parts of this ArrayData
392    ///
393    /// This is the inverse of [`ArrayData::try_new`].
394    ///
395    /// Returns `(data_type, len, nulls, offset, buffers, child_data)`
396    pub fn into_parts(
397        self,
398    ) -> (
399        DataType,
400        usize,
401        Option<NullBuffer>,
402        usize,
403        Vec<Buffer>,
404        Vec<ArrayData>,
405    ) {
406        let Self {
407            data_type,
408            len,
409            nulls,
410            offset,
411            buffers,
412            child_data,
413        } = self;
414
415        (data_type, len, nulls, offset, buffers, child_data)
416    }
417
418    /// Returns a builder to construct a [`ArrayData`] instance of the same [`DataType`]
419    #[inline]
420    pub const fn builder(data_type: DataType) -> ArrayDataBuilder {
421        ArrayDataBuilder::new(data_type)
422    }
423
424    /// Returns a reference to the [`DataType`] of this [`ArrayData`]
425    #[inline]
426    pub const fn data_type(&self) -> &DataType {
427        &self.data_type
428    }
429
430    /// Returns the [`Buffer`] storing data for this [`ArrayData`]
431    pub fn buffers(&self) -> &[Buffer] {
432        &self.buffers
433    }
434
435    /// Returns a slice of children [`ArrayData`]. This will be non
436    /// empty for type such as lists and structs.
437    ///
438    /// Note: For nested types where the parent element `i` corresponds directly
439    /// to child element `i` (such as structs), both the parent's offset and
440    /// each child's own offset apply when locating child values — see
441    /// [`Self::offset`] for details.
442    pub fn child_data(&self) -> &[ArrayData] {
443        &self.child_data[..]
444    }
445
446    /// Returns whether the element at index `i` is null
447    #[inline]
448    pub fn is_null(&self, i: usize) -> bool {
449        match &self.nulls {
450            Some(v) => v.is_null(i),
451            None => false,
452        }
453    }
454
455    /// Returns a reference to the null buffer of this [`ArrayData`] if any
456    ///
457    /// Note: [`ArrayData::offset`] does NOT apply to the returned [`NullBuffer`]
458    #[inline]
459    pub fn nulls(&self) -> Option<&NullBuffer> {
460        self.nulls.as_ref()
461    }
462
463    /// Returns whether the element at index `i` is not null
464    #[inline]
465    pub fn is_valid(&self, i: usize) -> bool {
466        !self.is_null(i)
467    }
468
469    /// Returns the length (i.e., number of elements) of this [`ArrayData`].
470    #[inline]
471    pub const fn len(&self) -> usize {
472        self.len
473    }
474
475    /// Returns whether this [`ArrayData`] is empty
476    #[inline]
477    pub const fn is_empty(&self) -> bool {
478        self.len == 0
479    }
480
481    /// Returns the offset in elements of this [`ArrayData`]
482    ///
483    /// The offset applies to [`Self::buffers`] and [`Self::child_data`],
484    /// but does NOT apply to [`Self::nulls`], which always represents exactly
485    /// [`Self::len`] elements.
486    ///
487    /// # Offsets for Non-nested types
488    ///
489    /// For non-nested types, the offset skips leading elements in the buffers.
490    /// Logical element `i` is stored at physical position `offset + i`.
491    ///
492    /// For example, with `offset = 2` and `len = 3` the following array
493    /// represents elements `[C, D, E]`:
494    ///
495    /// ```text
496    ///                     offset: 2        len: 3
497    ///                   ◀───────────▶◀────────────────▶
498    ///                   ┌─────┬─────┬─────┬─────┬─────┬─────┐
499    ///    values buffer  │  A  │  B  │  C  │  D  │  E  │  F  │
500    ///                   └─────┴─────┴─────┴─────┴─────┴─────┘
501    ///    physical index    0     1     2     3     4     5
502    ///    logical index                 0     1     2
503    /// ```
504    ///
505    /// # Offsets for Struct types
506    ///
507    /// For [struct]s, logical element `i` of the parent corresponds directly to
508    /// element `i` of each child, with no indirection in between. Since a
509    /// struct has no buffers of its own, its offset applies to each child,
510    /// composing cumulatively with any child offset. Logical element `i` of the
511    /// struct corresponds to element `offset + i` of each child.
512    ///
513    /// For example, a struct with `offset = 2` and `len = 3` whose children
514    /// `c1` and `c2` themselves each have an offset of `1` represents the
515    /// elements `{c1: D, c2: d}`, `{c1: E, c2: e}`, `{c1: F, c2: f}`:
516    ///
517    /// ```text
518    ///                        struct offset: 2    len: 3
519    ///                          ◀───────────▶◀────────────────▶
520    ///    child offset: 1 ◀─────▶
521    ///                    ┌─────┬─────┬─────┬─────┬─────┬─────┬─────┐
522    ///    child c1        │  A  │  B  │  C  │  D  │  E  │  F  │  G  │
523    ///                    ├─────┼─────┼─────┼─────┼─────┼─────┼─────┤
524    ///    child c2        │  a  │  b  │  c  │  d  │  e  │  f  │  g  │
525    ///                    └─────┴─────┴─────┴─────┴─────┴─────┴─────┘
526    ///    physical index     0     1     2     3     4     5     6
527    ///    child index              0     1     2     3     4     5
528    ///    struct index                         0     1     2
529    /// ```
530    ///
531    /// [struct]: https://arrow.apache.org/docs/format/Columnar.html#struct-layout
532    #[inline]
533    pub const fn offset(&self) -> usize {
534        self.offset
535    }
536
537    /// Returns the total number of nulls in this array
538    #[inline]
539    pub fn null_count(&self) -> usize {
540        self.nulls
541            .as_ref()
542            .map(|x| x.null_count())
543            .unwrap_or_default()
544    }
545
546    /// Returns the total number of bytes of memory occupied by the
547    /// buffers owned by this [`ArrayData`] and all of its
548    /// children. (See also diagram on [`ArrayData`]).
549    ///
550    /// Note that this [`ArrayData`] may only refer to a subset of the
551    /// data in the underlying [`Buffer`]s (due to `offset` and
552    /// `length`), but the size returned includes the entire size of
553    /// the buffers.
554    ///
555    /// If multiple [`ArrayData`]s refer to the same underlying
556    /// [`Buffer`]s they will both report the same size.
557    pub fn get_buffer_memory_size(&self) -> usize {
558        let mut size = 0;
559        for buffer in &self.buffers {
560            size += buffer.capacity();
561        }
562        if let Some(bitmap) = &self.nulls {
563            size += bitmap.buffer().capacity()
564        }
565        for child in &self.child_data {
566            size += child.get_buffer_memory_size();
567        }
568        size
569    }
570
571    /// Returns the total number of the bytes of memory occupied by
572    /// the buffers by this slice of [`ArrayData`] (See also diagram on [`ArrayData`]).
573    ///
574    /// This is approximately the number of bytes if a new
575    /// [`ArrayData`] was formed by creating new [`Buffer`]s with
576    /// exactly the data needed. For variadic layouts, this includes the full
577    /// capacity of every variadic buffer retained by a zero-copy slice, without
578    /// inspecting which buffers or ranges are referenced by the slice.
579    ///
580    /// For example, a [`DataType::Int64`] with `100` elements,
581    /// [`Self::get_slice_memory_size`] would return `100 * 8 = 800`. If
582    /// the [`ArrayData`] was then [`Self::slice`]ed to refer to its
583    /// first `20` elements, then [`Self::get_slice_memory_size`] on the
584    /// sliced [`ArrayData`] would return `20 * 8 = 160`.
585    pub fn get_slice_memory_size(&self) -> Result<usize, ArrowError> {
586        let mut result: usize = 0;
587        let layout = layout(&self.data_type);
588
589        for spec in &layout.buffers {
590            match spec {
591                BufferSpec::FixedWidth { byte_width, .. } => {
592                    // Offset buffers contain len+1 elements: one boundary per element
593                    // plus a final boundary marking the end of the last element.
594                    let len = match self.data_type {
595                        DataType::Utf8
596                        | DataType::LargeUtf8
597                        | DataType::Binary
598                        | DataType::LargeBinary
599                        | DataType::List(_)
600                        | DataType::LargeList(_)
601                        | DataType::Map(_, _) => self.len + 1,
602                        _ => self.len,
603                    };
604                    let buffer_size = len.checked_mul(*byte_width).ok_or_else(|| {
605                        ArrowError::ComputeError(
606                            "Integer overflow computing buffer size".to_string(),
607                        )
608                    })?;
609                    result += buffer_size;
610                }
611                BufferSpec::VariableWidth => {
612                    let buffer_len = match self.data_type {
613                        DataType::Utf8 | DataType::Binary => {
614                            let offsets = self.typed_offsets::<i32>()?;
615                            (offsets[self.len] - offsets[0]) as usize
616                        }
617                        DataType::LargeUtf8 | DataType::LargeBinary => {
618                            let offsets = self.typed_offsets::<i64>()?;
619                            (offsets[self.len] - offsets[0]) as usize
620                        }
621                        _ => {
622                            return Err(ArrowError::NotYetImplemented(format!(
623                                "Invalid data type for VariableWidth buffer. Expected Utf8, LargeUtf8, Binary or LargeBinary. Got {}",
624                                self.data_type
625                            )));
626                        }
627                    };
628                    result += buffer_len;
629                }
630                BufferSpec::BitMap => {
631                    let buffer_size = bit_util::ceil(self.len, 8);
632                    result += buffer_size;
633                }
634                BufferSpec::AlwaysNull => {
635                    // Nothing to do
636                }
637            }
638        }
639
640        if layout.variadic {
641            // Slicing view arrays retains all variadic data buffers unchanged.
642            for buffer in self.buffers.iter().skip(layout.buffers.len()) {
643                result += buffer.capacity();
644            }
645        }
646
647        if self.nulls().is_some() {
648            result += bit_util::ceil(self.len, 8);
649        }
650
651        for child in &self.child_data {
652            result += child.get_slice_memory_size()?;
653        }
654        Ok(result)
655    }
656
657    /// Returns the total number of bytes of memory occupied
658    /// physically by this [`ArrayData`] and all its [`Buffer`]s and
659    /// children. (See also diagram on [`ArrayData`]).
660    ///
661    /// Equivalent to:
662    ///  `size_of_val(self)` +
663    ///  [`Self::get_buffer_memory_size`] +
664    ///  `size_of_val(child)` for all children
665    pub fn get_array_memory_size(&self) -> usize {
666        let mut size = mem::size_of_val(self);
667
668        // Calculate rest of the fields top down which contain actual data
669        for buffer in &self.buffers {
670            size += mem::size_of::<Buffer>();
671            size += buffer.capacity();
672        }
673        if let Some(nulls) = &self.nulls {
674            size += nulls.buffer().capacity();
675        }
676        for child in &self.child_data {
677            size += child.get_array_memory_size();
678        }
679
680        size
681    }
682
683    /// Creates a zero-copy slice of itself. This creates a new
684    /// [`ArrayData`] pointing at the same underlying [`Buffer`]s with a
685    /// different offset and len
686    ///
687    /// # Panics
688    ///
689    /// Panics if `offset + length` overflows or is greater than `self.len()`.
690    pub fn slice(&self, offset: usize, length: usize) -> ArrayData {
691        self.clone().sliced(offset, length)
692    }
693
694    /// Creates a zero-copy slice of this array, reusing its buffer and child allocations.
695    ///
696    /// This is equivalent to [`Self::slice`], but consumes `self`.
697    ///
698    /// # Panics
699    ///
700    /// Panics if `offset + length` overflows or is greater than `self.len()`.
701    pub fn sliced(mut self, offset: usize, length: usize) -> Self {
702        self.slice_in_place(offset, length);
703        self
704    }
705
706    fn slice_in_place(&mut self, offset: usize, length: usize) {
707        let end = offset
708            .checked_add(length)
709            .expect("offset + length overflow");
710        assert!(end <= self.len());
711
712        if let DataType::Struct(_) = self.data_type() {
713            // A struct has no buffers of its own, and reading child element `i`
714            // combines this array's offset with the child's own offset. Applying
715            // the slice to both would count it twice, so the cumulative offset
716            // goes to the children and this array's offset is reset to 0.
717            let child_offset = self.offset + offset;
718            for child in &mut self.child_data {
719                child.slice_in_place(child_offset, length);
720            }
721            self.offset = 0;
722        } else {
723            self.offset += offset;
724        }
725        self.len = length;
726        // `nulls` belongs to this array rather than to the children, so
727        // it is sliced by `offset` alone.
728        self.nulls = self.nulls.as_ref().map(|x| x.slice(offset, length));
729    }
730
731    /// Returns the `buffer` as a slice of type `T` starting at self.offset
732    ///
733    /// # Panics
734    /// This function panics if:
735    /// * the buffer is not byte-aligned with type T, or
736    /// * the datatype is `Boolean` (it corresponds to a bit-packed buffer where the offset is not applicable)
737    pub fn buffer<T: ArrowNativeType>(&self, buffer: usize) -> &[T] {
738        &self.buffers()[buffer].typed_data()[self.offset..]
739    }
740
741    /// Returns a new [`ArrayData`] valid for `data_type` containing `len` null values
742    ///
743    /// # Panics
744    /// This function panics if:
745    /// * the datatype `data_type` has incorrect layout
746    pub fn new_null(data_type: &DataType, len: usize) -> Self {
747        let bit_len = bit_util::ceil(len, 8);
748        let zeroed = |len: usize| Buffer::from(MutableBuffer::from_len_zeroed(len));
749
750        let (buffers, child_data, has_nulls) = match data_type.primitive_width() {
751            Some(width) => (vec![zeroed(width * len)], vec![], true),
752            None => match data_type {
753                DataType::Null => (vec![], vec![], false),
754                DataType::Boolean => (vec![zeroed(bit_len)], vec![], true),
755                DataType::Binary | DataType::Utf8 => {
756                    (vec![zeroed((len + 1) * 4), zeroed(0)], vec![], true)
757                }
758                DataType::BinaryView | DataType::Utf8View => (vec![zeroed(len * 16)], vec![], true),
759                DataType::LargeBinary | DataType::LargeUtf8 => {
760                    (vec![zeroed((len + 1) * 8), zeroed(0)], vec![], true)
761                }
762                DataType::FixedSizeBinary(i) => {
763                    if *i < 0 {
764                        panic!("cannot construct null data from FixedSizeBinary({i})");
765                    }
766                    (vec![zeroed(*i as usize * len)], vec![], true)
767                }
768                DataType::List(f) | DataType::Map(f, _) => (
769                    vec![zeroed((len + 1) * 4)],
770                    vec![ArrayData::new_empty(f.data_type())],
771                    true,
772                ),
773                DataType::LargeList(f) => (
774                    vec![zeroed((len + 1) * 8)],
775                    vec![ArrayData::new_empty(f.data_type())],
776                    true,
777                ),
778                DataType::ListView(f) => (
779                    vec![zeroed(len * 4), zeroed(len * 4)],
780                    vec![ArrayData::new_empty(f.data_type())],
781                    true,
782                ),
783                DataType::LargeListView(f) => (
784                    vec![zeroed(len * 8), zeroed(len * 8)],
785                    vec![ArrayData::new_empty(f.data_type())],
786                    true,
787                ),
788                DataType::FixedSizeList(f, list_len) => (
789                    vec![],
790                    vec![ArrayData::new_null(f.data_type(), *list_len as usize * len)],
791                    true,
792                ),
793                DataType::Struct(fields) => (
794                    vec![],
795                    fields
796                        .iter()
797                        .map(|f| Self::new_null(f.data_type(), len))
798                        .collect(),
799                    true,
800                ),
801                DataType::Dictionary(k, v) => (
802                    vec![zeroed(k.primitive_width().unwrap() * len)],
803                    vec![ArrayData::new_empty(v.as_ref())],
804                    true,
805                ),
806                DataType::Union(f, mode) => match f.iter().next() {
807                    // Every slot carries a type id naming one of the children,
808                    // so an empty union has nothing to put in a slot and can
809                    // only be the empty array.
810                    None => {
811                        assert_eq!(
812                            len, 0,
813                            "cannot construct null data from an empty union of length {len}, a slot has no type id to carry"
814                        );
815                        let buffers = match mode {
816                            UnionMode::Sparse => vec![zeroed(0)],
817                            UnionMode::Dense => vec![zeroed(0), zeroed(0)],
818                        };
819                        (buffers, vec![], false)
820                    }
821                    Some((id, _)) => {
822                        let ids = Buffer::from_iter(std::iter::repeat_n(id, len));
823                        let buffers = match mode {
824                            UnionMode::Sparse => vec![ids],
825                            UnionMode::Dense => {
826                                let end_offset = i32::from_usize(len).unwrap();
827                                vec![ids, Buffer::from_iter(0_i32..end_offset)]
828                            }
829                        };
830
831                        let children = f
832                            .iter()
833                            .enumerate()
834                            .map(|(idx, (_, f))| {
835                                if idx == 0 || *mode == UnionMode::Sparse {
836                                    Self::new_null(f.data_type(), len)
837                                } else {
838                                    Self::new_empty(f.data_type())
839                                }
840                            })
841                            .collect();
842
843                        (buffers, children, false)
844                    }
845                },
846                DataType::RunEndEncoded(r, v) => {
847                    if len == 0 {
848                        // For empty arrays, create zero-length child arrays.
849                        let runs = ArrayData::new_empty(r.data_type());
850                        let values = ArrayData::new_empty(v.data_type());
851                        (vec![], vec![runs, values], false)
852                    } else {
853                        let runs = match r.data_type() {
854                            DataType::Int16 => {
855                                let i = i16::from_usize(len).expect("run overflow");
856                                Buffer::from_slice_ref([i])
857                            }
858                            DataType::Int32 => {
859                                let i = i32::from_usize(len).expect("run overflow");
860                                Buffer::from_slice_ref([i])
861                            }
862                            DataType::Int64 => {
863                                let i = i64::from_usize(len).expect("run overflow");
864                                Buffer::from_slice_ref([i])
865                            }
866                            dt => unreachable!("Invalid run ends data type {dt}"),
867                        };
868
869                        let builder = ArrayData::builder(r.data_type().clone())
870                            .len(1)
871                            .buffers(vec![runs]);
872
873                        // SAFETY:
874                        // Valid by construction
875                        let runs = unsafe { builder.build_unchecked() };
876                        (
877                            vec![],
878                            vec![runs, ArrayData::new_null(v.data_type(), 1)],
879                            false,
880                        )
881                    }
882                }
883                // Handled by Some(width) branch above
884                DataType::Int8
885                | DataType::Int16
886                | DataType::Int32
887                | DataType::Int64
888                | DataType::UInt8
889                | DataType::UInt16
890                | DataType::UInt32
891                | DataType::UInt64
892                | DataType::Float16
893                | DataType::Float32
894                | DataType::Float64
895                | DataType::Timestamp(_, _)
896                | DataType::Date32
897                | DataType::Date64
898                | DataType::Time32(_)
899                | DataType::Time64(_)
900                | DataType::Duration(_)
901                | DataType::Interval(_)
902                | DataType::Decimal32(_, _)
903                | DataType::Decimal64(_, _)
904                | DataType::Decimal128(_, _)
905                | DataType::Decimal256(_, _) => unreachable!("{data_type}"),
906            },
907        };
908
909        let mut builder = ArrayDataBuilder::new(data_type.clone())
910            .len(len)
911            .buffers(buffers)
912            .child_data(child_data);
913
914        if has_nulls {
915            builder = builder.nulls(Some(NullBuffer::new_null(len)))
916        }
917
918        // SAFETY:
919        // Data valid by construction
920        unsafe { builder.build_unchecked() }
921    }
922
923    /// Returns a new empty [ArrayData] valid for `data_type`.
924    pub fn new_empty(data_type: &DataType) -> Self {
925        Self::new_null(data_type, 0)
926    }
927
928    /// Verifies that the buffers meet the minimum alignment requirements for the data type
929    ///
930    /// Buffers that are not adequately aligned will be copied to a new aligned allocation
931    ///
932    /// This can be useful for when interacting with data sent over IPC or FFI, that may
933    /// not meet the minimum alignment requirements
934    ///
935    /// This also aligns buffers of children data
936    pub fn align_buffers(&mut self) {
937        let layout = layout(&self.data_type);
938        for (buffer, spec) in self.buffers.iter_mut().zip(&layout.buffers) {
939            if let BufferSpec::FixedWidth { alignment, .. } = spec
940                && buffer.as_ptr().align_offset(*alignment) != 0
941            {
942                *buffer = Buffer::from_slice_ref(buffer.as_ref());
943            }
944        }
945        // align children data recursively
946        for data in &mut self.child_data {
947            data.align_buffers()
948        }
949    }
950
951    /// "cheap" validation of an `ArrayData`. Ensures buffers are
952    /// sufficiently sized to store `len` + `offset` total elements of
953    /// `data_type` and performs other inexpensive consistency checks.
954    ///
955    /// This check is "cheap" in the sense that it does not validate the
956    /// contents of the buffers (e.g. that all offsets for UTF8 arrays
957    /// are within the bounds of the values buffer).
958    ///
959    /// See [ArrayData::validate_data] to validate fully the offset content
960    /// and the validity of utf8 data
961    pub fn validate(&self) -> Result<(), ArrowError> {
962        // Need at least this much space in each buffer
963        let len_plus_offset = checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
964
965        // Check that the data layout conforms to the spec
966        let layout = layout(&self.data_type);
967
968        if !layout.can_contain_null_mask && self.nulls.is_some() {
969            return Err(ArrowError::InvalidArgumentError(format!(
970                "Arrays of type {:?} cannot contain a null bitmask",
971                self.data_type,
972            )));
973        }
974
975        // Check data buffers length for view types and other types
976        if self.buffers.len() < layout.buffers.len()
977            || (!layout.variadic && self.buffers.len() != layout.buffers.len())
978        {
979            return Err(ArrowError::InvalidArgumentError(format!(
980                "Expected {} buffers in array of type {:?}, got {}",
981                layout.buffers.len(),
982                self.data_type,
983                self.buffers.len(),
984            )));
985        }
986
987        for (i, (buffer, spec)) in self.buffers.iter().zip(layout.buffers.iter()).enumerate() {
988            match spec {
989                BufferSpec::FixedWidth {
990                    byte_width,
991                    alignment,
992                } => {
993                    let min_buffer_size = len_plus_offset.saturating_mul(*byte_width);
994
995                    if buffer.len() < min_buffer_size {
996                        return Err(ArrowError::InvalidArgumentError(format!(
997                            "Need at least {} bytes in buffers[{}] in array of type {:?}, but got {}",
998                            min_buffer_size,
999                            i,
1000                            self.data_type,
1001                            buffer.len()
1002                        )));
1003                    }
1004
1005                    let align_offset = buffer.as_ptr().align_offset(*alignment);
1006                    if align_offset != 0 {
1007                        return Err(ArrowError::InvalidArgumentError(format!(
1008                            "Misaligned buffers[{i}] in array of type {:?}, offset from expected alignment of {alignment} by {}",
1009                            self.data_type,
1010                            align_offset.min(alignment - align_offset)
1011                        )));
1012                    }
1013                }
1014                BufferSpec::VariableWidth => {
1015                    // not cheap to validate (need to look at the
1016                    // data). Partially checked in validate_offsets
1017                    // called below. Can check with `validate_full`
1018                }
1019                BufferSpec::BitMap => {
1020                    let min_buffer_size = bit_util::ceil(len_plus_offset, 8);
1021                    if buffer.len() < min_buffer_size {
1022                        return Err(ArrowError::InvalidArgumentError(format!(
1023                            "Need at least {} bytes for bitmap in buffers[{}] in array of type {:?}, but got {}",
1024                            min_buffer_size,
1025                            i,
1026                            self.data_type,
1027                            buffer.len()
1028                        )));
1029                    }
1030                }
1031                BufferSpec::AlwaysNull => {
1032                    // Nothing to validate
1033                }
1034            }
1035        }
1036
1037        // check null bit buffer size
1038        if let Some(nulls) = self.nulls() {
1039            if nulls.null_count() > self.len {
1040                return Err(ArrowError::InvalidArgumentError(format!(
1041                    "null_count {} for an array exceeds length of {} elements",
1042                    nulls.null_count(),
1043                    self.len
1044                )));
1045            }
1046
1047            if nulls.len() != self.len {
1048                return Err(ArrowError::InvalidArgumentError(format!(
1049                    "null buffer incorrect size. got {} expected {}",
1050                    nulls.len(),
1051                    self.len
1052                )));
1053            }
1054        }
1055
1056        self.validate_child_data()?;
1057
1058        // Additional Type specific checks
1059        match &self.data_type {
1060            DataType::Utf8 | DataType::Binary => {
1061                self.validate_offsets::<i32>(self.buffers[1].len())?;
1062            }
1063            DataType::LargeUtf8 | DataType::LargeBinary => {
1064                self.validate_offsets::<i64>(self.buffers[1].len())?;
1065            }
1066            DataType::Dictionary(key_type, _value_type) => {
1067                // At the moment, constructing a DictionaryArray will also check this
1068                if !DataType::is_dictionary_key_type(key_type) {
1069                    return Err(ArrowError::InvalidArgumentError(format!(
1070                        "Dictionary key type must be integer, but was {key_type}"
1071                    )));
1072                }
1073            }
1074            DataType::RunEndEncoded(run_ends_type, _) => {
1075                if run_ends_type.is_nullable() {
1076                    return Err(ArrowError::InvalidArgumentError(
1077                        "The nullable should be set to false for the field defining run_ends array.".to_string()
1078                    ));
1079                }
1080                if !DataType::is_run_ends_type(run_ends_type.data_type()) {
1081                    return Err(ArrowError::InvalidArgumentError(format!(
1082                        "RunArray run_ends types must be Int16, Int32 or Int64, but was {}",
1083                        run_ends_type.data_type()
1084                    )));
1085                }
1086            }
1087            DataType::Map(f, _) if f.is_nullable() => {
1088                return Err(ArrowError::InvalidArgumentError(
1089                    "The nullable should be set to false for the map entries field.".to_string(),
1090                ));
1091            }
1092            _ => {}
1093        }
1094
1095        Ok(())
1096    }
1097
1098    /// Returns a reference to the data in `buffer` as a typed slice
1099    /// (typically `&[i32]` or `&[i64]`) after validating. The
1100    /// returned slice is guaranteed to have at least `self.len + 1`
1101    /// entries.
1102    ///
1103    /// For an empty array, the `buffer` can also be empty.
1104    fn typed_offsets<T: ArrowNativeType + num_traits::Num>(&self) -> Result<&[T], ArrowError> {
1105        // An empty list-like array can have 0 offsets
1106        if self.len == 0 && self.buffer_at(0)?.is_empty() {
1107            return Ok(&[]);
1108        }
1109
1110        let len = checked_len_plus_offset(&self.data_type, self.len, 1)?;
1111
1112        self.typed_buffer(0, len)
1113    }
1114
1115    /// Returns a reference to the data in `buffers[idx]` as a typed slice after validating
1116    fn typed_buffer<T: ArrowNativeType + num_traits::Num>(
1117        &self,
1118        idx: usize,
1119        len: usize,
1120    ) -> Result<&[T], ArrowError> {
1121        let buffer = self.buffer_at(idx)?;
1122
1123        let required_elements = checked_len_plus_offset(&self.data_type, len, self.offset)?;
1124        let byte_width = mem::size_of::<T>();
1125        let required_len = required_elements.checked_mul(byte_width).ok_or_else(|| {
1126            ArrowError::InvalidArgumentError(format!(
1127                "Buffer {idx} of {} byte length overflow: {} elements of {} bytes exceeds usize",
1128                self.data_type, required_elements, byte_width
1129            ))
1130        })?;
1131
1132        if buffer.len() < required_len {
1133            return Err(ArrowError::InvalidArgumentError(format!(
1134                "Buffer {} of {} isn't large enough. Expected {} bytes got {}",
1135                idx,
1136                self.data_type,
1137                required_len,
1138                buffer.len()
1139            )));
1140        }
1141
1142        // Only the bytes in use are typed: the buffer may be longer than
1143        // `required_len` by a number of bytes that is not a multiple of
1144        // `byte_width`, which `Buffer::typed_data` would assert on.
1145        let bytes = &buffer.as_slice()[..required_len];
1146        // SAFETY: ArrowNativeType is trivially transmutable and sealed (as in
1147        // `Buffer::typed_data`); alignment is checked below
1148        let (prefix, values, suffix) = unsafe { bytes.align_to::<T>() };
1149        if !prefix.is_empty() || !suffix.is_empty() {
1150            return Err(ArrowError::InvalidArgumentError(format!(
1151                "Buffer {idx} of {} is not aligned to {} bytes",
1152                self.data_type,
1153                mem::align_of::<T>()
1154            )));
1155        }
1156        Ok(&values[self.offset..])
1157    }
1158
1159    /// Does a cheap sanity check that the `self.len` values in `buffer` are valid
1160    /// offsets (of type T) into some other buffer of `values_length` bytes long
1161    fn validate_offsets<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1162        &self,
1163        values_length: usize,
1164    ) -> Result<(), ArrowError> {
1165        // Justification: buffer size was validated above
1166        let offsets = self.typed_offsets::<T>()?;
1167        if offsets.is_empty() {
1168            return Ok(());
1169        }
1170
1171        let first_offset = offsets[0].to_usize().ok_or_else(|| {
1172            ArrowError::InvalidArgumentError(format!(
1173                "Error converting offset[0] ({}) to usize for {}",
1174                offsets[0], self.data_type
1175            ))
1176        })?;
1177
1178        let last_offset = offsets[self.len].to_usize().ok_or_else(|| {
1179            ArrowError::InvalidArgumentError(format!(
1180                "Error converting offset[{}] ({}) to usize for {}",
1181                self.len, offsets[self.len], self.data_type
1182            ))
1183        })?;
1184
1185        if first_offset > values_length {
1186            return Err(ArrowError::InvalidArgumentError(format!(
1187                "First offset {} of {} is larger than values length {}",
1188                first_offset, self.data_type, values_length,
1189            )));
1190        }
1191
1192        if last_offset > values_length {
1193            return Err(ArrowError::InvalidArgumentError(format!(
1194                "Last offset {} of {} is larger than values length {}",
1195                last_offset, self.data_type, values_length,
1196            )));
1197        }
1198
1199        if first_offset > last_offset {
1200            return Err(ArrowError::InvalidArgumentError(format!(
1201                "First offset {} in {} is smaller than last offset {}",
1202                first_offset, self.data_type, last_offset,
1203            )));
1204        }
1205
1206        Ok(())
1207    }
1208
1209    /// Does a cheap sanity check that the `self.len` values in `buffer` are valid
1210    /// offsets and sizes (of type T) into some other buffer of `values_length` bytes long
1211    fn validate_offsets_and_sizes<T: ArrowNativeType + num_traits::Num + std::fmt::Display>(
1212        &self,
1213        values_length: usize,
1214    ) -> Result<(), ArrowError> {
1215        let offsets: &[T] = self.typed_buffer(0, self.len)?;
1216        let sizes: &[T] = self.typed_buffer(1, self.len)?;
1217        if offsets.len() != sizes.len() {
1218            return Err(ArrowError::ComputeError(format!(
1219                "ListView offsets len {} does not match sizes len {}",
1220                offsets.len(),
1221                sizes.len()
1222            )));
1223        }
1224
1225        for i in 0..sizes.len() {
1226            let size = sizes[i].to_usize().ok_or_else(|| {
1227                ArrowError::InvalidArgumentError(format!(
1228                    "Error converting size[{}] ({}) to usize for {}",
1229                    i, sizes[i], self.data_type
1230                ))
1231            })?;
1232            let offset = offsets[i].to_usize().ok_or_else(|| {
1233                ArrowError::InvalidArgumentError(format!(
1234                    "Error converting offset[{}] ({}) to usize for {}",
1235                    i, offsets[i], self.data_type
1236                ))
1237            })?;
1238            if size
1239                .checked_add(offset)
1240                .expect("Offset and size have exceeded the usize boundary")
1241                > values_length
1242            {
1243                return Err(ArrowError::InvalidArgumentError(format!(
1244                    "Size {} at index {} is larger than the remaining values for {}",
1245                    size, i, self.data_type
1246                )));
1247            }
1248        }
1249        Ok(())
1250    }
1251
1252    /// Validates the layout of `child_data` ArrayData structures
1253    fn validate_child_data(&self) -> Result<(), ArrowError> {
1254        match &self.data_type {
1255            DataType::List(field) => {
1256                let values_data = self.get_single_valid_child_data(field.data_type())?;
1257                self.validate_offsets::<i32>(values_data.len)?;
1258            }
1259            DataType::LargeList(field) => {
1260                let values_data = self.get_single_valid_child_data(field.data_type())?;
1261                self.validate_offsets::<i64>(values_data.len)?;
1262            }
1263            DataType::Map(field, _) => {
1264                let DataType::Struct(entries_fields) = field.data_type() else {
1265                    return Err(ArrowError::InvalidArgumentError(format!(
1266                        "Map field should be a entries struct data type, got {:?} instead",
1267                        field.data_type()
1268                    )));
1269                };
1270                if entries_fields.len() != 2 {
1271                    return Err(ArrowError::InvalidArgumentError(format!(
1272                        "Map entries data type should be a struct containing 2 fields, got {} fields",
1273                        entries_fields.len()
1274                    )));
1275                }
1276
1277                // Key field
1278                if entries_fields[0].is_nullable() {
1279                    return Err(ArrowError::InvalidArgumentError(
1280                        "Map key field must not be nullable".to_string(),
1281                    ));
1282                }
1283                let values_data = self.get_single_valid_child_data(field.data_type())?;
1284                self.validate_offsets::<i32>(values_data.len)?;
1285            }
1286            DataType::ListView(field) => {
1287                let values_data = self.get_single_valid_child_data(field.data_type())?;
1288                self.validate_offsets_and_sizes::<i32>(values_data.len)?;
1289            }
1290            DataType::LargeListView(field) => {
1291                let values_data = self.get_single_valid_child_data(field.data_type())?;
1292                self.validate_offsets_and_sizes::<i64>(values_data.len)?;
1293            }
1294            DataType::FixedSizeList(field, list_size) => {
1295                let values_data = self.get_single_valid_child_data(field.data_type())?;
1296
1297                let list_size: usize = (*list_size).try_into().map_err(|_| {
1298                    ArrowError::InvalidArgumentError(format!(
1299                        "{} has a negative list_size {}",
1300                        self.data_type, list_size
1301                    ))
1302                })?;
1303
1304                let len_plus_offset =
1305                    checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1306                let expected_values_len = len_plus_offset.checked_mul(list_size).ok_or_else(|| {
1307                    ArrowError::InvalidArgumentError(format!(
1308                        "Integer overflow computing expected number of values in {}: length + offset ({}) multiplied by list size ({}) exceeds usize",
1309                        self.data_type, len_plus_offset, list_size
1310                    ))
1311                })?;
1312
1313                if values_data.len < expected_values_len {
1314                    return Err(ArrowError::InvalidArgumentError(format!(
1315                        "Values length {} is less than the length + offset ({}) multiplied by the value size ({}) for {}",
1316                        values_data.len, len_plus_offset, list_size, self.data_type
1317                    )));
1318                }
1319            }
1320            DataType::Struct(fields) => {
1321                self.validate_num_child_data(fields.len())?;
1322                let len_plus_offset =
1323                    checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1324                for (i, field) in fields.iter().enumerate() {
1325                    let field_data = self.get_valid_child_data(i, field.data_type())?;
1326
1327                    // Ensure child field has sufficient size
1328                    if field_data.len < len_plus_offset {
1329                        return Err(ArrowError::InvalidArgumentError(format!(
1330                            "{} child array #{} for field {} has length smaller than expected for struct array ({} < {})",
1331                            self.data_type,
1332                            i,
1333                            field.name(),
1334                            field_data.len,
1335                            len_plus_offset
1336                        )));
1337                    }
1338                }
1339            }
1340            DataType::RunEndEncoded(run_ends_field, values_field) => {
1341                // See the `RunArray` invariants for the rules enforced here and
1342                // in `check_run_ends`:
1343                // https://docs.rs/arrow/latest/arrow/array/struct.RunArray.html#offset-and-len
1344                self.validate_num_child_data(2)?;
1345                let run_ends_data = self.get_valid_child_data(0, run_ends_field.data_type())?;
1346                let values_data = self.get_valid_child_data(1, values_field.data_type())?;
1347                if run_ends_data.len != values_data.len {
1348                    return Err(ArrowError::InvalidArgumentError(format!(
1349                        "The run_ends array length should be the same as values array length. Run_ends array length is {}, values array length is {}",
1350                        run_ends_data.len, values_data.len
1351                    )));
1352                }
1353                if run_ends_data.nulls.is_some() {
1354                    return Err(ArrowError::InvalidArgumentError(
1355                        "Found null values in run_ends array. The run_ends array should not have null values.".to_string(),
1356                    ));
1357                }
1358            }
1359            DataType::Union(fields, mode) => {
1360                self.validate_num_child_data(fields.len())?;
1361
1362                for (i, (_, field)) in fields.iter().enumerate() {
1363                    let field_data = self.get_valid_child_data(i, field.data_type())?;
1364
1365                    if mode == &UnionMode::Sparse {
1366                        let len_plus_offset =
1367                            checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1368                        if field_data.len < len_plus_offset {
1369                            return Err(ArrowError::InvalidArgumentError(format!(
1370                                "Sparse union child array #{} has length smaller than expected for union array ({} < {})",
1371                                i, field_data.len, len_plus_offset
1372                            )));
1373                        }
1374                    }
1375                }
1376            }
1377            DataType::Dictionary(_key_type, value_type) => {
1378                self.get_single_valid_child_data(value_type)?;
1379            }
1380            _ => {
1381                // other types do not have child data
1382                if !self.child_data.is_empty() {
1383                    return Err(ArrowError::InvalidArgumentError(format!(
1384                        "Expected no child arrays for type {} but got {}",
1385                        self.data_type,
1386                        self.child_data.len()
1387                    )));
1388                }
1389            }
1390        }
1391        Ok(())
1392    }
1393
1394    /// Ensures that this array data has a single child_data with the
1395    /// expected type, and calls `validate()` on it. Returns a
1396    /// reference to that child_data
1397    fn get_single_valid_child_data(
1398        &self,
1399        expected_type: &DataType,
1400    ) -> Result<&ArrayData, ArrowError> {
1401        self.validate_num_child_data(1)?;
1402        self.get_valid_child_data(0, expected_type)
1403    }
1404
1405    /// Returns `buffers[idx]`, or an error if there is no such buffer.
1406    ///
1407    /// [`Self::validate_values`] can be called on its own, without the buffer counts
1408    /// having been checked by [`Self::validate`] first, so the index may be missing.
1409    fn buffer_at(&self, idx: usize) -> Result<&Buffer, ArrowError> {
1410        self.buffers.get(idx).ok_or_else(|| {
1411            ArrowError::InvalidArgumentError(format!(
1412                "{} should contain at least {} buffer(s), had {}",
1413                self.data_type,
1414                idx + 1,
1415                self.buffers.len()
1416            ))
1417        })
1418    }
1419
1420    /// Returns `child_data[idx]`, or an error if there is no such child.
1421    ///
1422    /// [`Self::validate_values`] can be called on its own, without the child counts
1423    /// having been checked by [`Self::validate`] first, so the index may be missing.
1424    fn child_at(&self, idx: usize) -> Result<&ArrayData, ArrowError> {
1425        self.child_data.get(idx).ok_or_else(|| {
1426            ArrowError::InvalidArgumentError(format!(
1427                "{} should contain at least {} child data array(s), had {}",
1428                self.data_type,
1429                idx + 1,
1430                self.child_data.len()
1431            ))
1432        })
1433    }
1434
1435    /// Returns `Err` if self.child_data does not have exactly `expected_len` elements
1436    fn validate_num_child_data(&self, expected_len: usize) -> Result<(), ArrowError> {
1437        if self.child_data.len() != expected_len {
1438            Err(ArrowError::InvalidArgumentError(format!(
1439                "Value data for {} should contain {} child data array(s), had {}",
1440                self.data_type,
1441                expected_len,
1442                self.child_data.len()
1443            )))
1444        } else {
1445            Ok(())
1446        }
1447    }
1448
1449    /// Ensures that `child_data[i]` has the expected type, calls
1450    /// `validate()` on it, and returns a reference to that child_data
1451    fn get_valid_child_data(
1452        &self,
1453        i: usize,
1454        expected_type: &DataType,
1455    ) -> Result<&ArrayData, ArrowError> {
1456        let values_data = self.child_data.get(i).ok_or_else(|| {
1457            ArrowError::InvalidArgumentError(format!(
1458                "{} did not have enough child arrays. Expected at least {} but had only {}",
1459                self.data_type,
1460                i + 1,
1461                self.child_data.len()
1462            ))
1463        })?;
1464
1465        if expected_type != &values_data.data_type {
1466            return Err(ArrowError::InvalidArgumentError(format!(
1467                "Child type mismatch for {}. Expected {} but child data had {}",
1468                self.data_type, expected_type, values_data.data_type
1469            )));
1470        }
1471
1472        values_data.validate()?;
1473        Ok(values_data)
1474    }
1475
1476    /// Validate that the data contained within this [`ArrayData`] is valid
1477    ///
1478    /// 1. Null count is correct
1479    /// 2. All offsets are valid
1480    /// 3. All String data is valid UTF-8
1481    /// 4. All dictionary offsets are valid
1482    ///
1483    /// Internally this calls:
1484    ///
1485    /// * [`Self::validate`]
1486    /// * [`Self::validate_nulls`]
1487    /// * [`Self::validate_values`]
1488    ///
1489    /// Note: this does not recurse into children, for a recursive variant
1490    /// see [`Self::validate_full`]
1491    pub fn validate_data(&self) -> Result<(), ArrowError> {
1492        self.validate()?;
1493
1494        self.validate_nulls()?;
1495        self.validate_values()?;
1496        Ok(())
1497    }
1498
1499    /// Performs a full recursive validation of this [`ArrayData`] and all its children
1500    ///
1501    /// This is equivalent to calling [`Self::validate_data`] on this [`ArrayData`]
1502    /// and all its children recursively
1503    pub fn validate_full(&self) -> Result<(), ArrowError> {
1504        self.validate_data()?;
1505        // validate all children recursively
1506        self.child_data
1507            .iter()
1508            .enumerate()
1509            .try_for_each(|(i, child_data)| {
1510                child_data.validate_full().map_err(|e| {
1511                    ArrowError::InvalidArgumentError(format!(
1512                        "{} child #{} invalid: {}",
1513                        self.data_type, i, e
1514                    ))
1515                })
1516            })?;
1517        Ok(())
1518    }
1519
1520    /// Validates the values stored within this [`ArrayData`] are valid
1521    /// without recursing into child [`ArrayData`]
1522    ///
1523    /// Does not (yet) check
1524    /// 1. Union type_ids are valid see [#85](https://github.com/apache/arrow-rs/issues/85)
1525    /// 2. the null count is correct and that any
1526    /// 3. nullability requirements of its children are correct
1527    ///
1528    /// [#85]: https://github.com/apache/arrow-rs/issues/85
1529    pub fn validate_nulls(&self) -> Result<(), ArrowError> {
1530        if let Some(nulls) = &self.nulls {
1531            let actual = nulls.len() - nulls.inner().count_set_bits();
1532            if actual != nulls.null_count() {
1533                return Err(ArrowError::InvalidArgumentError(format!(
1534                    "null_count value ({}) doesn't match actual number of nulls in array ({})",
1535                    nulls.null_count(),
1536                    actual
1537                )));
1538            }
1539        }
1540
1541        // In general non-nullable children should not contain nulls, however, for certain
1542        // types, such as StructArray and FixedSizeList, nulls in the parent take up
1543        // space in the child. As such we permit nulls in the children in the corresponding
1544        // positions for such types
1545        match &self.data_type {
1546            DataType::List(f) | DataType::LargeList(f) | DataType::Map(f, _) => {
1547                if !f.is_nullable() {
1548                    let child = &self.child_data[0];
1549                    self.validate_non_nullable(None, child, child.nulls())?
1550                }
1551            }
1552            DataType::FixedSizeList(field, len) => {
1553                let child = &self.child_data[0];
1554                if !field.is_nullable() {
1555                    let element_len = *len as usize;
1556                    let child_nulls = child.nulls().map(|nulls| {
1557                        nulls.slice(self.offset * element_len, self.len * element_len)
1558                    });
1559                    let expanded = self.nulls.as_ref().map(|nulls| nulls.expand(element_len));
1560                    self.validate_non_nullable(expanded.as_ref(), child, child_nulls.as_ref())?;
1561                }
1562            }
1563            DataType::Struct(fields) => {
1564                for (field, child) in fields.iter().zip(&self.child_data) {
1565                    if !field.is_nullable() {
1566                        let child_nulls = child
1567                            .nulls()
1568                            .map(|nulls| nulls.slice(self.offset, self.len));
1569                        self.validate_non_nullable(self.nulls(), child, child_nulls.as_ref())?
1570                    }
1571                }
1572            }
1573            _ => {}
1574        }
1575
1576        Ok(())
1577    }
1578
1579    /// Verifies that `child` contains no nulls not present in `mask`
1580    fn validate_non_nullable(
1581        &self,
1582        mask: Option<&NullBuffer>,
1583        child: &ArrayData,
1584        child_nulls: Option<&NullBuffer>,
1585    ) -> Result<(), ArrowError> {
1586        let Some(mask) = mask else {
1587            return match child_nulls.map(NullBuffer::null_count).unwrap_or_default() {
1588                0 => Ok(()),
1589                _ => Err(ArrowError::InvalidArgumentError(format!(
1590                    "non-nullable child of type {} contains nulls not present in parent {}",
1591                    child.data_type, self.data_type
1592                ))),
1593            };
1594        };
1595
1596        match child_nulls {
1597            Some(nulls) if !mask.contains(nulls) => Err(ArrowError::InvalidArgumentError(format!(
1598                "non-nullable child of type {} contains nulls not present in parent",
1599                child.data_type
1600            ))),
1601            _ => Ok(()),
1602        }
1603    }
1604
1605    /// Validates the values stored within this [`ArrayData`] are valid
1606    /// without recursing into child [`ArrayData`]
1607    ///
1608    /// Does not (yet) check
1609    /// 1. Union type_ids are valid see [#85](https://github.com/apache/arrow-rs/issues/85)
1610    pub fn validate_values(&self) -> Result<(), ArrowError> {
1611        match &self.data_type {
1612            DataType::Utf8 => self.validate_utf8::<i32>(),
1613            DataType::LargeUtf8 => self.validate_utf8::<i64>(),
1614            DataType::Binary => self.validate_offsets_full::<i32>(self.buffer_at(1)?.len()),
1615            DataType::LargeBinary => self.validate_offsets_full::<i64>(self.buffer_at(1)?.len()),
1616            DataType::BinaryView => {
1617                let views = self.typed_buffer::<u128>(0, self.len)?;
1618                validate_binary_view(views, &self.buffers[1..])
1619            }
1620            DataType::Utf8View => {
1621                let views = self.typed_buffer::<u128>(0, self.len)?;
1622                validate_string_view(views, &self.buffers[1..])
1623            }
1624            DataType::List(_) | DataType::Map(_, _) => {
1625                let child = self.child_at(0)?;
1626                self.validate_offsets_full::<i32>(child.len)
1627            }
1628            DataType::LargeList(_) => {
1629                let child = self.child_at(0)?;
1630                self.validate_offsets_full::<i64>(child.len)
1631            }
1632            DataType::Union(fields, mode) => {
1633                let mut child_lens: Vec<Option<usize>> = vec![None; 128];
1634                for (index, (type_id, _)) in fields.iter().enumerate() {
1635                    let child = self.child_at(index)?;
1636                    child_lens[type_id as usize] = Some(child.len());
1637                }
1638
1639                // Type id values must match one of the fields.
1640                let type_ids = self.typed_buffer::<i8>(0, self.len)?;
1641                for (index, &type_id) in type_ids.iter().enumerate() {
1642                    if type_id < 0 {
1643                        return Err(ArrowError::InvalidArgumentError(format!(
1644                            "Type id values cannot be negative, got {type_id} at index {index}"
1645                        )));
1646                    }
1647                    let type_id = type_id as usize;
1648                    if child_lens[type_id].is_none() {
1649                        return Err(ArrowError::InvalidArgumentError(format!(
1650                            "Type id values must match one of the field type ids, got {type_id} at index {index}"
1651                        )));
1652                    }
1653                }
1654
1655                // Check the value offsets are in bounds.
1656                if mode == &UnionMode::Dense {
1657                    let offsets = self.typed_buffer::<i32>(1, self.len)?;
1658
1659                    if type_ids
1660                        .iter()
1661                        .zip(offsets.iter())
1662                        .any(|(&type_id, &offset)| {
1663                            offset < 0 || offset as usize >= child_lens[type_id as usize].unwrap()
1664                        })
1665                    {
1666                        return Err(ArrowError::InvalidArgumentError(
1667                            "Offsets must be non-negative and within the length of the child array"
1668                                .to_string(),
1669                        ));
1670                    }
1671                }
1672
1673                Ok(())
1674            }
1675            DataType::Dictionary(key_type, _value_type) => {
1676                let dictionary_length = self.child_at(0)?.len;
1677                let dictionary_length = i64::try_from(dictionary_length).map_err(|_| {
1678                    ArrowError::InvalidArgumentError(format!(
1679                        "Dictionary of {dictionary_length} values is too long for an i64"
1680                    ))
1681                })?;
1682                let max_value = dictionary_length - 1;
1683                match key_type.as_ref() {
1684                    DataType::UInt8 => self.check_bounds::<u8>(max_value),
1685                    DataType::UInt16 => self.check_bounds::<u16>(max_value),
1686                    DataType::UInt32 => self.check_bounds::<u32>(max_value),
1687                    DataType::UInt64 => self.check_bounds::<u64>(max_value),
1688                    DataType::Int8 => self.check_bounds::<i8>(max_value),
1689                    DataType::Int16 => self.check_bounds::<i16>(max_value),
1690                    DataType::Int32 => self.check_bounds::<i32>(max_value),
1691                    DataType::Int64 => self.check_bounds::<i64>(max_value),
1692                    _ => Err(ArrowError::InvalidArgumentError(format!(
1693                        "Dictionary key type must be an integer, got {key_type}"
1694                    ))),
1695                }
1696            }
1697            DataType::RunEndEncoded(run_ends, _values) => {
1698                let len_plus_offset =
1699                    checked_len_plus_offset(&self.data_type, self.len, self.offset)?;
1700                let run_ends_data = self.child_at(0)?;
1701                match run_ends.data_type() {
1702                    DataType::Int16 => {
1703                        run_ends_data.check_run_ends::<i16>(self.len, len_plus_offset)
1704                    }
1705                    DataType::Int32 => {
1706                        run_ends_data.check_run_ends::<i32>(self.len, len_plus_offset)
1707                    }
1708                    DataType::Int64 => {
1709                        run_ends_data.check_run_ends::<i64>(self.len, len_plus_offset)
1710                    }
1711                    data_type => Err(ArrowError::InvalidArgumentError(format!(
1712                        "Run end type must be Int16, Int32 or Int64, got {data_type}"
1713                    ))),
1714                }
1715            }
1716            _ => {
1717                // No extra validation check required for other types
1718                Ok(())
1719            }
1720        }
1721    }
1722
1723    /// Calls the `validate(item_index, range)` function for each of
1724    /// the ranges specified in the arrow offsets buffer of type
1725    /// `T`. Also validates that each offset is smaller than
1726    /// `offset_limit`
1727    ///
1728    /// For an empty array, the offsets buffer can either be empty
1729    /// or contain a single `0`.
1730    ///
1731    /// For example, the offsets buffer contained `[1, 2, 4]`, this
1732    /// function would call `validate([1,2])`, and `validate([2,4])`
1733    fn validate_each_offset<T, V>(&self, offset_limit: usize, validate: V) -> Result<(), ArrowError>
1734    where
1735        T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1736        V: Fn(usize, Range<usize>) -> Result<(), ArrowError>,
1737    {
1738        self.typed_offsets::<T>()?
1739            .iter()
1740            .enumerate()
1741            .map(|(i, x)| {
1742                // check if the offset can be converted to usize
1743                let r = x.to_usize().ok_or_else(|| {
1744                    ArrowError::InvalidArgumentError(format!(
1745                        "Offset invariant failure: Could not convert offset {x} to usize at position {i}"))}
1746                    );
1747                // check if the offset exceeds the limit
1748                match r {
1749                    Ok(n) if n <= offset_limit => Ok((i, n)),
1750                    Ok(_) => Err(ArrowError::InvalidArgumentError(format!(
1751                        "Offset invariant failure: offset at position {i} out of bounds: {x} > {offset_limit}"))
1752                    ),
1753                    Err(e) => Err(e),
1754                }
1755            })
1756            .scan(0_usize, |start, end| {
1757                // check offsets are monotonically increasing
1758                match end {
1759                    Ok((i, end)) if *start <= end => {
1760                        let range = Some(Ok((i, *start..end)));
1761                        *start = end;
1762                        range
1763                    }
1764                    Ok((i, end)) => Some(Err(ArrowError::InvalidArgumentError(format!(
1765                        "Offset invariant failure: non-monotonic offset at slot {}: {} > {}",
1766                        i - 1, start, end))
1767                    )),
1768                    Err(err) => Some(Err(err)),
1769                }
1770            })
1771            .skip(1) // the first element is meaningless
1772            .try_for_each(|res: Result<(usize, Range<usize>), ArrowError>| {
1773                let (item_index, range) = res?;
1774                validate(item_index-1, range)
1775            })
1776    }
1777
1778    /// Ensures that all strings formed by the offsets in `buffers[0]`
1779    /// into `buffers[1]` are valid utf8 sequences. Bytes of `buffers[1]`
1780    /// outside the range of the offsets are not checked.
1781    fn validate_utf8<T>(&self) -> Result<(), ArrowError>
1782    where
1783        T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1784    {
1785        let values_buffer = &self.buffer_at(1)?.as_slice();
1786        let offsets = self.typed_offsets::<T>()?;
1787        if let Some((first, last, values_str)) = utf8_span(offsets, values_buffer) {
1788            self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1789                // `values_str` ends at the last offset. Checking each pair of offsets doesn't
1790                // show that `range.end <= last`: a later offset can still be smaller, which
1791                // `validate_each_offset` reports when it gets there.
1792                if range.end > last {
1793                    return Ok(());
1794                }
1795                let (start, end) = (range.start - first, range.end - first);
1796                if !values_str.is_char_boundary(start) || !values_str.is_char_boundary(end) {
1797                    return Err(ArrowError::InvalidArgumentError(format!(
1798                        "incomplete utf-8 byte sequence from index {string_index}"
1799                    )));
1800                }
1801                Ok(())
1802            })
1803        } else {
1804            // Find specific offset that failed utf8 validation
1805            self.validate_each_offset::<T, _>(values_buffer.len(), |string_index, range| {
1806                std::str::from_utf8(&values_buffer[range.clone()]).map_err(|e| {
1807                    ArrowError::InvalidArgumentError(format!(
1808                        "Invalid UTF8 sequence at string index {string_index} ({range:?}): {e}"
1809                    ))
1810                })?;
1811                Ok(())
1812            })
1813        }
1814    }
1815
1816    /// Ensures that all offsets in `buffers[0]` into `buffers[1]` are
1817    /// between `0` and `offset_limit`
1818    fn validate_offsets_full<T>(&self, offset_limit: usize) -> Result<(), ArrowError>
1819    where
1820        T: ArrowNativeType + TryInto<usize> + num_traits::Num + std::fmt::Display,
1821    {
1822        self.validate_each_offset::<T, _>(offset_limit, |_string_index, _range| {
1823            // No validation applied to each value, but the iteration
1824            // itself applies bounds checking to each range
1825            Ok(())
1826        })
1827    }
1828
1829    /// Validates that each value in self.buffers (typed as T)
1830    /// is within the range [0, max_value], inclusive
1831    fn check_bounds<T>(&self, max_value: i64) -> Result<(), ArrowError>
1832    where
1833        T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1834    {
1835        // `validate()` checks the buffer size too, but `validate_values()` can be called
1836        // on its own, so do not assume it has run.
1837        let indexes: &[T] = self.typed_buffer::<T>(0, self.len)?;
1838
1839        indexes.iter().enumerate().try_for_each(|(i, &dict_index)| {
1840            // Do not check the value is null (value can be arbitrary)
1841            if self.is_null(i) {
1842                return Ok(());
1843            }
1844            let dict_index: i64 = dict_index.try_into().map_err(|_| {
1845                ArrowError::InvalidArgumentError(format!(
1846                    "Value at position {i} out of bounds: {dict_index} (can not convert to i64)"
1847                ))
1848            })?;
1849
1850            if dict_index < 0 || dict_index > max_value {
1851                return Err(ArrowError::InvalidArgumentError(format!(
1852                    "Value at position {i} out of bounds: {dict_index} (should be in [0, {max_value}])"
1853                )));
1854            }
1855            Ok(())
1856        })
1857    }
1858
1859    /// Validates that each value in run_ends array is positive and strictly
1860    /// increasing, and that the last value covers the parent's `offset + len`
1861    /// (when the parent is non-empty), as required by the [`RunArray` invariants].
1862    ///
1863    /// [`RunArray` invariants]: https://docs.rs/arrow/latest/arrow/array/struct.RunArray.html#offset-and-len
1864    fn check_run_ends<T>(
1865        &self,
1866        parent_len: usize,
1867        parent_len_plus_offset: usize,
1868    ) -> Result<(), ArrowError>
1869    where
1870        T: ArrowNativeType + TryInto<i64> + num_traits::Num + std::fmt::Display,
1871    {
1872        let values = self.typed_buffer::<T>(0, self.len)?;
1873        let mut prev_value = 0_i64;
1874        values.iter().enumerate().try_for_each(|(ix, &inp_value)| {
1875            let value: i64 = inp_value.try_into().map_err(|_| {
1876                ArrowError::InvalidArgumentError(format!(
1877                    "Value at position {ix} out of bounds: {inp_value} (can not convert to i64)"
1878                ))
1879            })?;
1880            if value <= 0_i64 {
1881                return Err(ArrowError::InvalidArgumentError(format!(
1882                    "The values in run_ends array should be strictly positive. Found value {value} at index {ix} that does not match the criteria."
1883                )));
1884            }
1885            if ix > 0 && value <= prev_value {
1886                return Err(ArrowError::InvalidArgumentError(format!(
1887                    "The values in run_ends array should be strictly increasing. Found value {value} at index {ix} with previous value {prev_value} that does not match the criteria."
1888                )));
1889            }
1890
1891            prev_value = value;
1892            Ok(())
1893        })?;
1894
1895        if parent_len > 0 && prev_value.as_usize() < parent_len_plus_offset {
1896            return Err(ArrowError::InvalidArgumentError(format!(
1897                "The offset + length of array should be less or equal to last value in the run_ends array. The last value of run_ends array is {prev_value} and offset + length of array is {parent_len_plus_offset}."
1898            )));
1899        }
1900        Ok(())
1901    }
1902
1903    /// Returns true if this `ArrayData` is equal to `other`, using pointer comparisons
1904    /// to determine buffer equality. This is cheaper than `PartialEq::eq` but may
1905    /// return false when the arrays are logically equal
1906    pub fn ptr_eq(&self, other: &Self) -> bool {
1907        if self.offset != other.offset
1908            || self.len != other.len
1909            || self.data_type != other.data_type
1910            || self.buffers.len() != other.buffers.len()
1911            || self.child_data.len() != other.child_data.len()
1912        {
1913            return false;
1914        }
1915
1916        match (&self.nulls, &other.nulls) {
1917            (Some(a), Some(b)) if !a.inner().ptr_eq(b.inner()) => return false,
1918            (Some(_), None) | (None, Some(_)) => return false,
1919            _ => {}
1920        }
1921
1922        if !self
1923            .buffers
1924            .iter()
1925            .zip(other.buffers.iter())
1926            .all(|(a, b)| a.as_ptr() == b.as_ptr())
1927        {
1928            return false;
1929        }
1930
1931        self.child_data
1932            .iter()
1933            .zip(other.child_data.iter())
1934            .all(|(a, b)| a.ptr_eq(b))
1935    }
1936
1937    /// Converts this [`ArrayData`] into an [`ArrayDataBuilder`]
1938    pub fn into_builder(self) -> ArrayDataBuilder {
1939        self.into()
1940    }
1941
1942    /// Claim memory used by this ArrayData in the provided memory pool.
1943    ///
1944    /// This claims memory for:
1945    /// - All buffers in self.buffers
1946    /// - All child ArrayData recursively
1947    /// - The null buffer if present
1948    #[cfg(feature = "pool")]
1949    pub fn claim(&self, pool: &dyn arrow_buffer::MemoryPool) {
1950        // Claim all data buffers
1951        for buffer in &self.buffers {
1952            buffer.claim(pool);
1953        }
1954
1955        // Claim null buffer if present
1956        if let Some(nulls) = &self.nulls {
1957            nulls.claim(pool);
1958        }
1959
1960        // Recursively claim child data
1961        for child in &self.child_data {
1962            child.claim(pool);
1963        }
1964    }
1965}
1966
1967/// Return the expected [`DataTypeLayout`] Arrays of this data
1968/// type are expected to have
1969pub fn layout(data_type: &DataType) -> DataTypeLayout {
1970    // based on C/C++ implementation in
1971    // https://github.com/apache/arrow/blob/661c7d749150905a63dd3b52e0a04dac39030d95/cpp/src/arrow/type.h (and .cc)
1972    use arrow_schema::IntervalUnit::*;
1973
1974    match data_type {
1975        DataType::Null => DataTypeLayout {
1976            buffers: vec![],
1977            can_contain_null_mask: false,
1978            variadic: false,
1979        },
1980        DataType::Boolean => DataTypeLayout {
1981            buffers: vec![BufferSpec::BitMap],
1982            can_contain_null_mask: true,
1983            variadic: false,
1984        },
1985        DataType::Int8 => DataTypeLayout::new_fixed_width::<i8>(),
1986        DataType::Int16 => DataTypeLayout::new_fixed_width::<i16>(),
1987        DataType::Int32 => DataTypeLayout::new_fixed_width::<i32>(),
1988        DataType::Int64 => DataTypeLayout::new_fixed_width::<i64>(),
1989        DataType::UInt8 => DataTypeLayout::new_fixed_width::<u8>(),
1990        DataType::UInt16 => DataTypeLayout::new_fixed_width::<u16>(),
1991        DataType::UInt32 => DataTypeLayout::new_fixed_width::<u32>(),
1992        DataType::UInt64 => DataTypeLayout::new_fixed_width::<u64>(),
1993        DataType::Float16 => DataTypeLayout::new_fixed_width::<half::f16>(),
1994        DataType::Float32 => DataTypeLayout::new_fixed_width::<f32>(),
1995        DataType::Float64 => DataTypeLayout::new_fixed_width::<f64>(),
1996        DataType::Timestamp(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
1997        DataType::Date32 => DataTypeLayout::new_fixed_width::<i32>(),
1998        DataType::Date64 => DataTypeLayout::new_fixed_width::<i64>(),
1999        DataType::Time32(_) => DataTypeLayout::new_fixed_width::<i32>(),
2000        DataType::Time64(_) => DataTypeLayout::new_fixed_width::<i64>(),
2001        DataType::Interval(YearMonth) => DataTypeLayout::new_fixed_width::<i32>(),
2002        DataType::Interval(DayTime) => DataTypeLayout::new_fixed_width::<IntervalDayTime>(),
2003        DataType::Interval(MonthDayNano) => {
2004            DataTypeLayout::new_fixed_width::<IntervalMonthDayNano>()
2005        }
2006        DataType::Duration(_) => DataTypeLayout::new_fixed_width::<i64>(),
2007        DataType::Decimal32(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
2008        DataType::Decimal64(_, _) => DataTypeLayout::new_fixed_width::<i64>(),
2009        DataType::Decimal128(_, _) => DataTypeLayout::new_fixed_width::<i128>(),
2010        DataType::Decimal256(_, _) => DataTypeLayout::new_fixed_width::<i256>(),
2011        DataType::FixedSizeBinary(size) => {
2012            let spec = BufferSpec::FixedWidth {
2013                byte_width: (*size).try_into().unwrap(),
2014                alignment: mem::align_of::<u8>(),
2015            };
2016            DataTypeLayout {
2017                buffers: vec![spec],
2018                can_contain_null_mask: true,
2019                variadic: false,
2020            }
2021        }
2022        DataType::Binary => DataTypeLayout::new_binary::<i32>(),
2023        DataType::LargeBinary => DataTypeLayout::new_binary::<i64>(),
2024        DataType::Utf8 => DataTypeLayout::new_binary::<i32>(),
2025        DataType::LargeUtf8 => DataTypeLayout::new_binary::<i64>(),
2026        DataType::BinaryView | DataType::Utf8View => DataTypeLayout::new_view(),
2027        DataType::FixedSizeList(_, _) => DataTypeLayout::new_nullable_empty(), // all in child data
2028        DataType::List(_) => DataTypeLayout::new_fixed_width::<i32>(),
2029        DataType::ListView(_) => DataTypeLayout::new_list_view::<i32>(),
2030        DataType::LargeListView(_) => DataTypeLayout::new_list_view::<i64>(),
2031        DataType::LargeList(_) => DataTypeLayout::new_fixed_width::<i64>(),
2032        DataType::Map(_, _) => DataTypeLayout::new_fixed_width::<i32>(),
2033        DataType::Struct(_) => DataTypeLayout::new_nullable_empty(), // all in child data,
2034        DataType::RunEndEncoded(_, _) => DataTypeLayout::new_empty(), // all in child data,
2035        DataType::Union(_, mode) => {
2036            let type_ids = BufferSpec::FixedWidth {
2037                byte_width: mem::size_of::<i8>(),
2038                alignment: mem::align_of::<i8>(),
2039            };
2040
2041            DataTypeLayout {
2042                buffers: match mode {
2043                    UnionMode::Sparse => {
2044                        vec![type_ids]
2045                    }
2046                    UnionMode::Dense => {
2047                        vec![
2048                            type_ids,
2049                            BufferSpec::FixedWidth {
2050                                byte_width: mem::size_of::<i32>(),
2051                                alignment: mem::align_of::<i32>(),
2052                            },
2053                        ]
2054                    }
2055                },
2056                can_contain_null_mask: false,
2057                variadic: false,
2058            }
2059        }
2060        DataType::Dictionary(key_type, _value_type) => layout(key_type),
2061    }
2062}
2063
2064/// Layout specification for a data type
2065#[derive(Debug, PartialEq, Eq)]
2066// Note: Follows structure from C++: https://github.com/apache/arrow/blob/master/cpp/src/arrow/type.h#L91
2067pub struct DataTypeLayout {
2068    /// A vector of buffer layout specifications, one for each expected buffer
2069    pub buffers: Vec<BufferSpec>,
2070
2071    /// Can contain a null bitmask
2072    pub can_contain_null_mask: bool,
2073
2074    /// This field only applies to the view type [`DataType::BinaryView`] and [`DataType::Utf8View`]
2075    /// If `variadic` is true, the number of buffers expected is only lower-bounded by
2076    /// buffers.len(). Buffers that exceed the lower bound are legal.
2077    pub variadic: bool,
2078}
2079
2080impl DataTypeLayout {
2081    /// Describes a basic numeric array where each element has type `T`
2082    pub fn new_fixed_width<T>() -> Self {
2083        Self {
2084            buffers: vec![BufferSpec::FixedWidth {
2085                byte_width: mem::size_of::<T>(),
2086                alignment: mem::align_of::<T>(),
2087            }],
2088            can_contain_null_mask: true,
2089            variadic: false,
2090        }
2091    }
2092
2093    /// Describes arrays which have no data of their own
2094    /// but may still have a Null Bitmap (e.g. FixedSizeList)
2095    pub fn new_nullable_empty() -> Self {
2096        Self {
2097            buffers: vec![],
2098            can_contain_null_mask: true,
2099            variadic: false,
2100        }
2101    }
2102
2103    /// Describes arrays which have no data of their own
2104    /// (e.g. RunEndEncoded).
2105    pub fn new_empty() -> Self {
2106        Self {
2107            buffers: vec![],
2108            can_contain_null_mask: false,
2109            variadic: false,
2110        }
2111    }
2112
2113    /// Describes a basic numeric array where each element has a fixed
2114    /// with offset buffer of type `T`, followed by a
2115    /// variable width data buffer
2116    pub fn new_binary<T>() -> Self {
2117        Self {
2118            buffers: vec![
2119                // offsets
2120                BufferSpec::FixedWidth {
2121                    byte_width: mem::size_of::<T>(),
2122                    alignment: mem::align_of::<T>(),
2123                },
2124                // values
2125                BufferSpec::VariableWidth,
2126            ],
2127            can_contain_null_mask: true,
2128            variadic: false,
2129        }
2130    }
2131
2132    /// Describes a view type
2133    pub fn new_view() -> Self {
2134        Self {
2135            buffers: vec![BufferSpec::FixedWidth {
2136                byte_width: mem::size_of::<u128>(),
2137                alignment: mem::align_of::<u128>(),
2138            }],
2139            can_contain_null_mask: true,
2140            variadic: true,
2141        }
2142    }
2143
2144    /// Describes a list view type
2145    pub fn new_list_view<T>() -> Self {
2146        Self {
2147            buffers: vec![
2148                BufferSpec::FixedWidth {
2149                    byte_width: mem::size_of::<T>(),
2150                    alignment: mem::align_of::<T>(),
2151                },
2152                BufferSpec::FixedWidth {
2153                    byte_width: mem::size_of::<T>(),
2154                    alignment: mem::align_of::<T>(),
2155                },
2156            ],
2157            can_contain_null_mask: true,
2158            variadic: false,
2159        }
2160    }
2161}
2162
2163/// Layout specification for a single data type buffer
2164#[derive(Debug, PartialEq, Eq)]
2165pub enum BufferSpec {
2166    /// Each element is a fixed width primitive, with the given `byte_width` and `alignment`
2167    ///
2168    /// `alignment` is the alignment required by Rust for an array of the corresponding primitive,
2169    /// see [`Layout::array`](std::alloc::Layout::array) and [`std::mem::align_of`].
2170    ///
2171    /// Arrow-rs requires that all buffers have at least this alignment, to allow for
2172    /// [slice](std::slice) based APIs. Alignment in excess of this is not required to allow
2173    /// for array slicing and interoperability with `Vec`, which cannot be over-aligned.
2174    ///
2175    /// Note that these alignment requirements will vary between architectures
2176    FixedWidth {
2177        /// The width of each element in bytes
2178        byte_width: usize,
2179        /// The alignment required by Rust for an array of the corresponding primitive
2180        alignment: usize,
2181    },
2182    /// Variable width, such as string data for utf8 data
2183    VariableWidth,
2184    /// Buffer holds a bitmap.
2185    ///
2186    /// Note: Unlike the C++ implementation, the null/validity buffer
2187    /// is handled specially rather than as another of the buffers in
2188    /// the spec, so this variant is only used for the Boolean type.
2189    BitMap,
2190    /// Buffer is always null. Unused currently in Rust implementation,
2191    /// (used in C++ for Union type)
2192    AlwaysNull,
2193}
2194
2195impl PartialEq for ArrayData {
2196    fn eq(&self, other: &Self) -> bool {
2197        equal::equal(self, other)
2198    }
2199}
2200
2201/// A boolean flag that cannot be mutated outside of unsafe code.
2202///
2203/// Defaults to a value of false.
2204///
2205/// This structure is used to enforce safety in the [`ArrayDataBuilder`]
2206///
2207/// [`ArrayDataBuilder`]: super::ArrayDataBuilder
2208///
2209/// # Example
2210/// ```rust
2211/// use arrow_data::UnsafeFlag;
2212/// assert!(!UnsafeFlag::default().get()); // default is false
2213/// let mut flag = UnsafeFlag::new();
2214/// assert!(!flag.get()); // defaults to false
2215/// // can only set it to true in unsafe code
2216/// unsafe { flag.set(true) };
2217/// assert!(flag.get()); // now true
2218/// ```
2219#[derive(Debug, Clone)]
2220#[doc(hidden)]
2221pub struct UnsafeFlag(bool);
2222
2223impl UnsafeFlag {
2224    /// Creates a new `UnsafeFlag` with the value set to `false`.
2225    ///
2226    /// See examples on [`Self::new`]
2227    #[inline]
2228    pub const fn new() -> Self {
2229        Self(false)
2230    }
2231
2232    /// Sets the value of the flag to the given value
2233    ///
2234    /// Note this can purposely only be done in `unsafe` code
2235    ///
2236    /// # Safety
2237    ///
2238    /// If set, the flag will be set to the given value. There is nothing
2239    /// immediately unsafe about doing so, however, the flag can be used to
2240    /// subsequently bypass safety checks in the [`ArrayDataBuilder`].
2241    #[inline]
2242    pub unsafe fn set(&mut self, val: bool) {
2243        self.0 = val;
2244    }
2245
2246    /// Returns the value of the flag
2247    #[inline]
2248    pub fn get(&self) -> bool {
2249        self.0
2250    }
2251}
2252
2253// Manual impl to make it clear you can not construct unsafe with true
2254impl Default for UnsafeFlag {
2255    fn default() -> Self {
2256        Self::new()
2257    }
2258}
2259
2260/// Builder for [`ArrayData`] type
2261#[derive(Debug)]
2262pub struct ArrayDataBuilder {
2263    data_type: DataType,
2264    len: usize,
2265    null_count: Option<usize>,
2266    null_bit_buffer: Option<Buffer>,
2267    nulls: Option<NullBuffer>,
2268    offset: usize,
2269    buffers: Vec<Buffer>,
2270    child_data: Vec<ArrayData>,
2271    /// Should buffers be realigned (copying if necessary)?
2272    ///
2273    /// Defaults to false.
2274    align_buffers: bool,
2275    /// Should data validation be skipped for this [`ArrayData`]?
2276    ///
2277    /// Defaults to false.
2278    ///
2279    /// # Safety
2280    ///
2281    /// This flag can only be set to true using `unsafe` APIs. However, once true
2282    /// subsequent calls to `build()` may result in undefined behavior if the data
2283    /// is not valid.
2284    skip_validation: UnsafeFlag,
2285}
2286
2287impl ArrayDataBuilder {
2288    #[inline]
2289    /// Creates a new array data builder
2290    pub const fn new(data_type: DataType) -> Self {
2291        Self {
2292            data_type,
2293            len: 0,
2294            null_count: None,
2295            null_bit_buffer: None,
2296            nulls: None,
2297            offset: 0,
2298            buffers: vec![],
2299            child_data: vec![],
2300            align_buffers: false,
2301            skip_validation: UnsafeFlag::new(),
2302        }
2303    }
2304
2305    /// Creates a new array data builder from an existing one, changing the data type
2306    pub fn data_type(self, data_type: DataType) -> Self {
2307        Self { data_type, ..self }
2308    }
2309
2310    #[inline]
2311    /// Sets the length of the [ArrayData]
2312    pub const fn len(mut self, n: usize) -> Self {
2313        self.len = n;
2314        self
2315    }
2316
2317    /// Sets the null buffer of the [ArrayData]
2318    pub fn nulls(mut self, nulls: Option<NullBuffer>) -> Self {
2319        self.nulls = nulls;
2320        self.null_count = None;
2321        self.null_bit_buffer = None;
2322        self
2323    }
2324
2325    /// Sets the null count of the [ArrayData]
2326    pub fn null_count(mut self, null_count: usize) -> Self {
2327        self.null_count = Some(null_count);
2328        self
2329    }
2330
2331    /// Sets the `null_bit_buffer` of the [ArrayData]
2332    pub fn null_bit_buffer(mut self, buf: Option<Buffer>) -> Self {
2333        self.nulls = None;
2334        self.null_bit_buffer = buf;
2335        self
2336    }
2337
2338    /// Sets the offset of the [ArrayData]
2339    #[inline]
2340    pub const fn offset(mut self, n: usize) -> Self {
2341        self.offset = n;
2342        self
2343    }
2344
2345    /// Sets the buffers of the [ArrayData]
2346    pub fn buffers(mut self, v: Vec<Buffer>) -> Self {
2347        self.buffers = v;
2348        self
2349    }
2350
2351    /// Adds a single buffer to the [ArrayData]'s buffers
2352    pub fn add_buffer(mut self, b: Buffer) -> Self {
2353        self.buffers.push(b);
2354        self
2355    }
2356
2357    /// Adds multiple buffers to the [ArrayData]'s buffers
2358    pub fn add_buffers<I: IntoIterator<Item = Buffer>>(mut self, bs: I) -> Self {
2359        self.buffers.extend(bs);
2360        self
2361    }
2362
2363    /// Sets the child data of the [ArrayData]
2364    pub fn child_data(mut self, v: Vec<ArrayData>) -> Self {
2365        self.child_data = v;
2366        self
2367    }
2368
2369    /// Adds a single child data to the [ArrayData]'s child data
2370    pub fn add_child_data(mut self, r: ArrayData) -> Self {
2371        self.child_data.push(r);
2372        self
2373    }
2374
2375    /// Creates an array data, without any validation
2376    ///
2377    /// Note: This is shorthand for
2378    /// ```rust
2379    /// # #[expect(unsafe_op_in_unsafe_fn)]
2380    /// # let mut builder = arrow_data::ArrayDataBuilder::new(arrow_schema::DataType::Null);
2381    /// # let _ = unsafe {
2382    /// builder.skip_validation(true).build().unwrap()
2383    /// # };
2384    /// ```
2385    ///
2386    /// # Safety
2387    ///
2388    /// The same caveats as [`ArrayData::new_unchecked`]
2389    /// apply.
2390    pub unsafe fn build_unchecked(self) -> ArrayData {
2391        unsafe { self.skip_validation(true) }.build().unwrap()
2392    }
2393
2394    /// Creates an `ArrayData`, consuming `self`
2395    ///
2396    /// # Undefined behavior
2397    ///
2398    /// By default the underlying buffers are checked to ensure they are valid
2399    /// Arrow data. However, if the [`Self::skip_validation`] flag has been set
2400    /// to true (by the `unsafe` API) this validation is skipped. If the data is
2401    /// not valid, undefined behavior will result.
2402    pub fn build(self) -> Result<ArrayData, ArrowError> {
2403        let Self {
2404            data_type,
2405            len,
2406            null_count,
2407            null_bit_buffer,
2408            nulls,
2409            offset,
2410            buffers,
2411            child_data,
2412            align_buffers,
2413            skip_validation,
2414        } = self;
2415
2416        // SAFETY: `skip_validation` is only set to true using `unsafe` APIs.
2417        let validate = !skip_validation.get() || cfg!(feature = "force_validate");
2418        if validate && let Some(buffer) = null_bit_buffer.as_ref() {
2419            // Check before constructing the BooleanBuffer, which would otherwise panic.
2420            let len_plus_offset = checked_len_plus_offset(&data_type, len, offset)?;
2421            let needed_len = bit_util::ceil(len_plus_offset, 8);
2422            if buffer.len() < needed_len {
2423                return Err(ArrowError::InvalidArgumentError(format!(
2424                    "null_bit_buffer size too small. got {} needed {}",
2425                    buffer.len(),
2426                    needed_len
2427                )));
2428            }
2429        }
2430
2431        let nulls = nulls
2432            .or_else(|| {
2433                let buffer = null_bit_buffer?;
2434                let buffer = BooleanBuffer::new(buffer, offset, len);
2435                Some(match null_count {
2436                    Some(n) => {
2437                        // SAFETY: call to `data.validate_data()` below validates the null buffer is valid
2438                        unsafe { NullBuffer::new_unchecked(buffer, n) }
2439                    }
2440                    None => NullBuffer::new(buffer),
2441                })
2442            })
2443            .filter(|b| b.null_count() != 0);
2444
2445        let mut data = ArrayData {
2446            data_type,
2447            len,
2448            offset,
2449            buffers,
2450            child_data,
2451            nulls,
2452        };
2453
2454        if align_buffers {
2455            data.align_buffers();
2456        }
2457
2458        if validate {
2459            data.validate_data()?;
2460        }
2461        Ok(data)
2462    }
2463
2464    /// Ensure that all buffers are aligned, copying data if necessary
2465    ///
2466    /// Rust requires that arrays are aligned to their corresponding primitive,
2467    /// see [`Layout::array`](std::alloc::Layout::array) and [`std::mem::align_of`].
2468    ///
2469    /// [`ArrayData`] therefore requires that all buffers have at least this alignment,
2470    /// to allow for [slice](std::slice) based APIs. See [`BufferSpec::FixedWidth`].
2471    ///
2472    /// As this alignment is architecture specific, and not guaranteed by all arrow implementations,
2473    /// this flag is provided to automatically copy buffers to a new correctly aligned allocation
2474    /// when necessary, making it useful when interacting with buffers produced by other systems,
2475    /// e.g. IPC or FFI.
2476    ///
2477    /// If this flag is not enabled, `[Self::build`] return an error on encountering
2478    /// insufficiently aligned buffers.
2479    pub fn align_buffers(mut self, align_buffers: bool) -> Self {
2480        self.align_buffers = align_buffers;
2481        self
2482    }
2483
2484    /// Skips validation of the data.
2485    ///
2486    /// If this flag is enabled, `[Self::build`] will skip validation of the
2487    /// data
2488    ///
2489    /// If this flag is not enabled, `[Self::build`] will validate that all
2490    /// buffers are valid and will return an error if any data is invalid.
2491    /// Validation can be expensive.
2492    ///
2493    /// # Safety
2494    ///
2495    /// If validation is skipped, the buffers must form a valid Arrow array,
2496    /// otherwise undefined behavior will result
2497    pub unsafe fn skip_validation(mut self, skip_validation: bool) -> Self {
2498        unsafe {
2499            self.skip_validation.set(skip_validation);
2500        }
2501        self
2502    }
2503}
2504
2505impl From<ArrayData> for ArrayDataBuilder {
2506    fn from(d: ArrayData) -> Self {
2507        Self {
2508            data_type: d.data_type,
2509            len: d.len,
2510            offset: d.offset,
2511            buffers: d.buffers,
2512            child_data: d.child_data,
2513            nulls: d.nulls,
2514            null_bit_buffer: None,
2515            null_count: None,
2516            align_buffers: false,
2517            skip_validation: UnsafeFlag::new(),
2518        }
2519    }
2520}
2521
2522/// Get byte width of FixedSizeBinary size
2523/// # Panics:
2524/// - Panics if the `data_type` is not FixedSizeBinary
2525/// - Panics if byte width is negative
2526pub(crate) fn get_fixed_size_binary_width(data_type: &DataType) -> usize {
2527    match data_type {
2528        DataType::FixedSizeBinary(i) => {
2529            if *i < 0 {
2530                panic!("cannot compare FixedSizeBinary({})", *i);
2531            }
2532            *i as usize
2533        }
2534        _ => unreachable!(),
2535    }
2536}
2537
2538#[cfg(test)]
2539mod tests {
2540    use arrow_schema::UnionFields;
2541
2542    use super::*;
2543    use crate::ByteView;
2544    use crate::transform::MutableArrayData;
2545    use arrow_buffer::{OffsetBuffer, ScalarBuffer};
2546    use arrow_schema::{Field, Fields};
2547
2548    // See arrow/tests/array_data_validation.rs for test of array validation
2549
2550    /// returns a buffer initialized with some constant value for tests
2551    fn make_i32_buffer(n: usize) -> Buffer {
2552        Buffer::from_slice_ref(vec![42i32; n])
2553    }
2554
2555    /// returns a buffer initialized with some constant value for tests
2556    fn make_f32_buffer(n: usize) -> Buffer {
2557        Buffer::from_slice_ref(vec![42f32; n])
2558    }
2559
2560    #[test]
2561    fn test_builder() {
2562        // Buffer needs to be at least 25 long
2563        let v = (0..25).collect::<Vec<i32>>();
2564        let b1 = Buffer::from_slice_ref(&v);
2565        let arr_data = ArrayData::builder(DataType::Int32)
2566            .len(20)
2567            .offset(5)
2568            .add_buffer(b1)
2569            .null_bit_buffer(Some(Buffer::from([
2570                0b01011111, 0b10110101, 0b01100011, 0b00011110,
2571            ])))
2572            .build()
2573            .unwrap();
2574
2575        assert_eq!(20, arr_data.len());
2576        assert_eq!(10, arr_data.null_count());
2577        assert_eq!(5, arr_data.offset());
2578        assert_eq!(1, arr_data.buffers().len());
2579        assert_eq!(
2580            Buffer::from_slice_ref(&v).as_slice(),
2581            arr_data.buffers()[0].as_slice()
2582        );
2583    }
2584
2585    #[test]
2586    fn test_builder_with_child_data() {
2587        let child_arr_data = ArrayData::try_new(
2588            DataType::Int32,
2589            5,
2590            None,
2591            0,
2592            vec![Buffer::from_slice_ref([1i32, 2, 3, 4, 5])],
2593            vec![],
2594        )
2595        .unwrap();
2596
2597        let field = Arc::new(Field::new("x", DataType::Int32, true));
2598        let data_type = DataType::Struct(vec![field].into());
2599
2600        let arr_data = ArrayData::builder(data_type)
2601            .len(5)
2602            .offset(0)
2603            .add_child_data(child_arr_data.clone())
2604            .build()
2605            .unwrap();
2606
2607        assert_eq!(5, arr_data.len());
2608        assert_eq!(1, arr_data.child_data().len());
2609        assert_eq!(child_arr_data, arr_data.child_data()[0]);
2610    }
2611
2612    #[test]
2613    fn test_struct_validation_accounts_for_parent_offset() {
2614        let data_type =
2615            DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2616        let child = ArrayData::builder(DataType::Int32)
2617            .len(5)
2618            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2619            .build()
2620            .unwrap();
2621
2622        // The parent needs child elements 1..6, but the child only has five.
2623        let err = ArrayData::builder(data_type)
2624            .len(5)
2625            .offset(1)
2626            .add_child_data(child)
2627            .build()
2628            .unwrap_err()
2629            .to_string();
2630
2631        assert!(err.contains(
2632            "child array #0 for field x has length smaller than expected for struct array (5 < 6)"
2633        ));
2634    }
2635
2636    #[test]
2637    fn test_struct_non_nullable_child_nulls_account_for_parent_offset() {
2638        let build = |parent_nulls| {
2639            let child = ArrayData::builder(DataType::Int32)
2640                .len(5)
2641                .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2642                .nulls(Some(NullBuffer::new(BooleanBuffer::from(vec![
2643                    true, true, false, true, true,
2644                ]))))
2645                .build()
2646                .unwrap();
2647
2648            ArrayData::builder(DataType::Struct(Fields::from(vec![Field::new(
2649                "x",
2650                DataType::Int32,
2651                false,
2652            )])))
2653            .len(4)
2654            .offset(1)
2655            .nulls(Some(NullBuffer::new(BooleanBuffer::from(parent_nulls))))
2656            .add_child_data(child)
2657            .build()
2658        };
2659
2660        assert!(build(vec![true, false, true, true]).is_ok());
2661        assert!(build(vec![true, true, false, true]).is_err());
2662    }
2663
2664    #[test]
2665    fn test_fixed_size_list_non_nullable_child_nulls_account_for_parent_offset() {
2666        for child_offset in [0, 1] {
2667            let child = ArrayData::builder(DataType::Int32)
2668                .len(6)
2669                .offset(child_offset)
2670                .add_buffer(Buffer::from_slice_ref([0i32; 7]))
2671                .nulls(Some(NullBuffer::from(vec![
2672                    true, true, false, false, true, true,
2673                ])))
2674                .build()
2675                .unwrap();
2676            let build = |parent_nulls| {
2677                ArrayData::builder(DataType::FixedSizeList(
2678                    Arc::new(Field::new("item", DataType::Int32, false)),
2679                    2,
2680                ))
2681                .offset(1)
2682                .len(2)
2683                .nulls(Some(NullBuffer::from(parent_nulls)))
2684                .add_child_data(child.clone())
2685                .build()
2686            };
2687
2688            // The first visible list contains both child nulls.
2689            assert!(build(vec![false, true]).is_ok());
2690            assert!(build(vec![true, false]).is_err());
2691        }
2692    }
2693
2694    #[test]
2695    fn test_fixed_size_list_non_nullable_child_nulls_without_parent_nulls() {
2696        let child = ArrayData::builder(DataType::Int32)
2697            .len(6)
2698            .add_buffer(Buffer::from_slice_ref([0i32; 6]))
2699            .nulls(Some(NullBuffer::from(vec![
2700                false, false, true, true, false, false,
2701            ])))
2702            .build()
2703            .unwrap();
2704        let build = |offset, len, size| {
2705            ArrayData::builder(DataType::FixedSizeList(
2706                Arc::new(Field::new("item", DataType::Int32, false)),
2707                size,
2708            ))
2709            .offset(offset)
2710            .len(len)
2711            .add_child_data(child.clone())
2712            .build()
2713        };
2714
2715        // Nulls before and after the visible child range are irrelevant.
2716        assert!(build(1, 1, 2).is_ok());
2717        assert!(build(0, 1, 2).is_err());
2718        assert!(build(2, 1, 2).is_err());
2719        assert!(build(3, 0, 2).is_ok());
2720        assert!(build(1, 2, 0).is_ok());
2721    }
2722
2723    #[test]
2724    fn test_fixed_size_list_validation_accounts_for_parent_offset() {
2725        for nulls in [None, Some(NullBuffer::new_null(4))] {
2726            let child = ArrayData::builder(DataType::Int32)
2727                .len(4)
2728                .add_buffer(Buffer::from_slice_ref([0i32; 4]))
2729                .nulls(nulls)
2730                .build()
2731                .unwrap();
2732            let err = ArrayData::builder(DataType::FixedSizeList(
2733                Arc::new(Field::new("item", DataType::Int32, false)),
2734                2,
2735            ))
2736            .offset(1)
2737            .len(2)
2738            .add_child_data(child)
2739            .build()
2740            .unwrap_err()
2741            .to_string();
2742            assert!(err.contains("Values length 4 is less than"));
2743        }
2744    }
2745
2746    #[test]
2747    fn test_fixed_size_list_validation_child_len_mul_overflow() {
2748        // Empty child and no parent null buffer: huge logical lengths allocate nothing.
2749        let child = ArrayData::builder(DataType::Int32)
2750            .len(0)
2751            .add_buffer(Buffer::from_slice_ref([0i32; 0]))
2752            .build()
2753            .unwrap();
2754        let list_size = 2i32;
2755        let data_type = DataType::FixedSizeList(
2756            Arc::new(Field::new("item", DataType::Int32, true)),
2757            list_size,
2758        );
2759        // `(len + offset) * list_size` overflows; the addition itself does not.
2760        let len_plus_offset = (usize::MAX / 2) + 1;
2761        let overflow_cases = [(len_plus_offset, 0usize), (usize::MAX / 2, 1)];
2762
2763        for (len, offset) in overflow_cases {
2764            let err = ArrayData::builder(data_type.clone())
2765                .len(len)
2766                .offset(offset)
2767                .add_child_data(child.clone())
2768                .build()
2769                .unwrap_err();
2770
2771            match err {
2772                ArrowError::InvalidArgumentError(msg) => assert_eq!(
2773                    msg,
2774                    format!(
2775                        "Integer overflow computing expected number of values in {data_type}: length + offset ({len_plus_offset}) multiplied by list size ({list_size}) exceeds usize"
2776                    )
2777                ),
2778                other => panic!("unexpected error type {other}"),
2779            }
2780        }
2781
2782        // Same product is representable here, so validation reports a short child.
2783        let err = ArrayData::builder(data_type.clone())
2784            .len(usize::MAX / 2)
2785            .offset(0)
2786            .add_child_data(child)
2787            .build()
2788            .unwrap_err();
2789
2790        match err {
2791            ArrowError::InvalidArgumentError(msg) => assert_eq!(
2792                msg,
2793                format!(
2794                    "Values length 0 is less than the length + offset ({}) multiplied by the value size ({list_size}) for {data_type}",
2795                    usize::MAX / 2
2796                )
2797            ),
2798            other => panic!("unexpected error type {other}"),
2799        }
2800    }
2801
2802    #[test]
2803    fn test_struct_equal_accounts_for_parent_offset() {
2804        let data_type =
2805            DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2806
2807        let child1 = ArrayData::builder(DataType::Int32)
2808            .len(5)
2809            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2810            .build()
2811            .unwrap();
2812        let child2 = child1.slice(1, 4);
2813
2814        // data1 has offset at parent level; data2 has offset at child level.
2815        let data1 = ArrayData::builder(data_type.clone())
2816            .len(4)
2817            .offset(1)
2818            .add_child_data(child1)
2819            .build()
2820            .unwrap();
2821        let data2 = ArrayData::builder(data_type)
2822            .len(4)
2823            .add_child_data(child2)
2824            .build()
2825            .unwrap();
2826
2827        assert_eq!(data1, data2);
2828    }
2829
2830    #[test]
2831    fn test_extend_struct_accounts_for_parent_offset() {
2832        let data_type =
2833            DataType::Struct(Fields::from(vec![Field::new("x", DataType::Int32, false)]));
2834        let child = ArrayData::builder(DataType::Int32)
2835            .len(5)
2836            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4]))
2837            .build()
2838            .unwrap();
2839
2840        let data = ArrayData::builder(data_type)
2841            .len(4)
2842            .offset(1)
2843            .add_child_data(child)
2844            .build()
2845            .unwrap();
2846
2847        let mut mutable = MutableArrayData::new(vec![&data], false, data.len());
2848        mutable.try_extend(0, 0, data.len()).unwrap();
2849        let output = mutable.freeze();
2850
2851        assert_eq!(output.child_data()[0].buffer::<i32>(0), &[1, 2, 3, 4]);
2852    }
2853
2854    #[test]
2855    fn test_extend_fixed_size_list_accounts_for_parent_offset() {
2856        let field = Arc::new(Field::new("item", DataType::Int32, false));
2857        let child = ArrayData::builder(DataType::Int32)
2858            .len(6)
2859            .add_buffer(Buffer::from_slice_ref([0, 1, 2, 3, 4, 5]))
2860            .build()
2861            .unwrap();
2862
2863        let data = ArrayData::builder(DataType::FixedSizeList(field, 2))
2864            .len(2)
2865            .offset(1)
2866            .add_child_data(child)
2867            .build()
2868            .unwrap();
2869
2870        let mut mutable = MutableArrayData::new(vec![&data], false, data.len());
2871        mutable.try_extend(0, 0, data.len()).unwrap();
2872        let output = mutable.freeze();
2873
2874        assert_eq!(output.child_data()[0].buffer::<i32>(0), &[2, 3, 4, 5]);
2875    }
2876
2877    #[test]
2878    fn test_extend_sparse_union_accounts_for_parent_offset() {
2879        let fields = UnionFields::try_new(
2880            [0, 1],
2881            [
2882                Field::new("a", DataType::Int32, false),
2883                Field::new("b", DataType::Int32, false),
2884            ],
2885        )
2886        .unwrap();
2887        let child_a = ArrayData::builder(DataType::Int32)
2888            .len(3)
2889            .add_buffer(Buffer::from_slice_ref([10, 20, 30]))
2890            .build()
2891            .unwrap();
2892        let child_b = ArrayData::builder(DataType::Int32)
2893            .len(3)
2894            .add_buffer(Buffer::from_slice_ref([100, 200, 300]))
2895            .build()
2896            .unwrap();
2897
2898        // The children are left whole, as `ArrayData::slice` and C Data
2899        // Interface producers leave them for a sparse union
2900        let data = ArrayData::builder(DataType::Union(fields, UnionMode::Sparse))
2901            .len(2)
2902            .offset(1)
2903            .add_buffer(Buffer::from_slice_ref([0i8, 1, 0]))
2904            .add_child_data(child_a)
2905            .add_child_data(child_b)
2906            .build()
2907            .unwrap();
2908
2909        let mut mutable = MutableArrayData::new(vec![&data], false, data.len());
2910        mutable.try_extend(0, 0, data.len()).unwrap();
2911        let output = mutable.freeze();
2912
2913        assert_eq!(output.buffer::<i8>(0), &[1, 0]);
2914        assert_eq!(output.child_data()[0].buffer::<i32>(0), &[20, 30]);
2915        assert_eq!(output.child_data()[1].buffer::<i32>(0), &[200, 300]);
2916    }
2917
2918    #[test]
2919    fn test_null_count() {
2920        let mut bit_v: [u8; 2] = [0; 2];
2921        bit_util::set_bit(&mut bit_v, 0);
2922        bit_util::set_bit(&mut bit_v, 3);
2923        bit_util::set_bit(&mut bit_v, 10);
2924        let arr_data = ArrayData::builder(DataType::Int32)
2925            .len(16)
2926            .add_buffer(make_i32_buffer(16))
2927            .null_bit_buffer(Some(Buffer::from(bit_v)))
2928            .build()
2929            .unwrap();
2930        assert_eq!(13, arr_data.null_count());
2931
2932        // Test with offset
2933        let mut bit_v: [u8; 2] = [0; 2];
2934        bit_util::set_bit(&mut bit_v, 0);
2935        bit_util::set_bit(&mut bit_v, 3);
2936        bit_util::set_bit(&mut bit_v, 10);
2937        let arr_data = ArrayData::builder(DataType::Int32)
2938            .len(12)
2939            .offset(2)
2940            .add_buffer(make_i32_buffer(14)) // requires at least 14 bytes of space,
2941            .null_bit_buffer(Some(Buffer::from(bit_v)))
2942            .build()
2943            .unwrap();
2944        assert_eq!(10, arr_data.null_count());
2945    }
2946
2947    #[test]
2948    fn test_null_buffer_ref() {
2949        let mut bit_v: [u8; 2] = [0; 2];
2950        bit_util::set_bit(&mut bit_v, 0);
2951        bit_util::set_bit(&mut bit_v, 3);
2952        bit_util::set_bit(&mut bit_v, 10);
2953        let arr_data = ArrayData::builder(DataType::Int32)
2954            .len(16)
2955            .add_buffer(make_i32_buffer(16))
2956            .null_bit_buffer(Some(Buffer::from(bit_v)))
2957            .build()
2958            .unwrap();
2959        assert!(arr_data.nulls().is_some());
2960        assert_eq!(&bit_v, arr_data.nulls().unwrap().validity());
2961    }
2962
2963    #[test]
2964    fn test_slice() {
2965        let mut bit_v: [u8; 2] = [0; 2];
2966        bit_util::set_bit(&mut bit_v, 0);
2967        bit_util::set_bit(&mut bit_v, 3);
2968        bit_util::set_bit(&mut bit_v, 10);
2969        let data = ArrayData::builder(DataType::Int32)
2970            .len(16)
2971            .add_buffer(make_i32_buffer(16))
2972            .null_bit_buffer(Some(Buffer::from(bit_v)))
2973            .build()
2974            .unwrap();
2975        let new_data = data.slice(1, 15);
2976        assert_eq!(data.len() - 1, new_data.len());
2977        assert_eq!(1, new_data.offset());
2978        assert_eq!(data.null_count(), new_data.null_count());
2979
2980        // slice of a slice (removes one null)
2981        let new_data = new_data.slice(1, 14);
2982        assert_eq!(data.len() - 2, new_data.len());
2983        assert_eq!(2, new_data.offset());
2984        assert_eq!(data.null_count() - 1, new_data.null_count());
2985    }
2986
2987    #[test]
2988    #[should_panic(expected = "offset + length overflow")]
2989    fn test_slice_panics_on_offset_length_overflow() {
2990        let data = ArrayData::builder(DataType::Int32)
2991            .len(4)
2992            .add_buffer(make_i32_buffer(4))
2993            .build()
2994            .unwrap();
2995        let sliced = data.slice(1, 3);
2996
2997        sliced.slice(1, usize::MAX);
2998    }
2999
3000    #[test]
3001    fn test_sliced() {
3002        for nulls in [
3003            None,
3004            Some(NullBuffer::from(vec![true, false, true, false, true])),
3005        ] {
3006            let data = ArrayData::builder(DataType::Int32)
3007                .len(5)
3008                .offset(1)
3009                .add_buffer(make_i32_buffer(6))
3010                .nulls(nulls)
3011                .build()
3012                .unwrap();
3013
3014            for offset in 0..=data.len() {
3015                for length in 0..=data.len() - offset {
3016                    let owned = data.clone();
3017                    let buffers = owned.buffers().as_ptr();
3018                    let sliced = owned.sliced(offset, length);
3019                    assert_eq!(sliced, data.slice(offset, length));
3020                    assert_eq!(sliced.len(), length);
3021                    assert_eq!(sliced.offset(), 1 + offset);
3022                    assert_eq!(sliced.buffers().as_ptr(), buffers);
3023                    assert_eq!(sliced.buffers()[0].as_ptr(), data.buffers()[0].as_ptr());
3024                    if let Some(nulls) = data.nulls() {
3025                        assert_eq!(
3026                            sliced.nulls().unwrap().iter().collect::<Vec<_>>(),
3027                            nulls.iter().skip(offset).take(length).collect::<Vec<_>>()
3028                        );
3029                    }
3030                    sliced.validate_full().unwrap();
3031                }
3032            }
3033        }
3034
3035        let empty = ArrayData::new_empty(&DataType::Int32);
3036        assert_eq!(empty.clone().sliced(0, 0), empty);
3037    }
3038
3039    #[test]
3040    fn test_sliced_nested_struct() {
3041        let child = ArrayData::builder(DataType::Int32)
3042            .len(8)
3043            .offset(1)
3044            .add_buffer(make_i32_buffer(9))
3045            .nulls(Some(NullBuffer::from(vec![
3046                true, false, true, false, true, true, false, true,
3047            ])))
3048            .build()
3049            .unwrap();
3050        let inner = ArrayData::builder(DataType::Struct(Fields::from(vec![Field::new(
3051            "x",
3052            DataType::Int32,
3053            true,
3054        )])))
3055        .len(6)
3056        .offset(1)
3057        .nulls(Some(NullBuffer::from(vec![
3058            true, false, true, true, false, true,
3059        ])))
3060        .add_child_data(child)
3061        .build()
3062        .unwrap();
3063        let data = ArrayData::builder(DataType::Struct(Fields::from(vec![Field::new(
3064            "inner",
3065            inner.data_type().clone(),
3066            true,
3067        )])))
3068        .len(4)
3069        .offset(1)
3070        .nulls(Some(NullBuffer::from(vec![false, true, false, true])))
3071        .add_child_data(inner)
3072        .build()
3073        .unwrap();
3074
3075        for (offset, length) in [(0, 4), (1, 2), (4, 0), (0, 0)] {
3076            let owned = data.clone();
3077            let children = owned.child_data().as_ptr();
3078            let grandchildren = owned.child_data()[0].child_data().as_ptr();
3079            let buffers = owned.child_data()[0].child_data()[0].buffers().as_ptr();
3080            let sliced = owned.sliced(offset, length);
3081            let inner = &sliced.child_data()[0];
3082            let child = &inner.child_data()[0];
3083            assert_eq!(sliced, data.slice(offset, length));
3084            assert_eq!(sliced.offset(), 0);
3085            assert_eq!(inner.offset(), 0);
3086            assert_eq!(child.offset(), 3 + offset);
3087            assert_eq!(sliced.child_data().as_ptr(), children);
3088            assert_eq!(inner.child_data().as_ptr(), grandchildren);
3089            assert_eq!(child.buffers().as_ptr(), buffers);
3090            sliced.validate_full().unwrap();
3091        }
3092
3093        let sliced = data.sliced(1, 2).sliced(1, 1);
3094        assert_eq!(sliced.null_count(), 1);
3095        assert_eq!(sliced.child_data()[0].null_count(), 0);
3096        assert_eq!(sliced.child_data()[0].child_data()[0].offset(), 5);
3097        assert_eq!(sliced.child_data()[0].child_data()[0].null_count(), 0);
3098    }
3099
3100    #[test]
3101    fn test_sliced_list_preserves_children() {
3102        let child = ArrayData::builder(DataType::Int32)
3103            .len(6)
3104            .offset(1)
3105            .add_buffer(make_i32_buffer(7))
3106            .nulls(Some(NullBuffer::from(vec![
3107                true, false, true, false, true, true,
3108            ])))
3109            .build()
3110            .unwrap();
3111        let data = ArrayData::builder(DataType::List(Arc::new(Field::new(
3112            "item",
3113            DataType::Int32,
3114            true,
3115        ))))
3116        .len(4)
3117        .offset(1)
3118        .add_buffer(Buffer::from_slice_ref([0i32, 1, 2, 3, 4, 5]))
3119        .nulls(Some(NullBuffer::from(vec![false, true, false, true])))
3120        .add_child_data(child)
3121        .build()
3122        .unwrap();
3123        let buffers = data.buffers().as_ptr();
3124        let children = data.child_data().as_ptr();
3125        let child_buffers = data.child_data()[0].buffers().as_ptr();
3126        let expected = data.slice(1, 2);
3127        let sliced = data.sliced(1, 2);
3128        assert_eq!(sliced, expected);
3129        assert_eq!(sliced.offset(), 2);
3130        assert_eq!(sliced.child_data()[0].len(), 6);
3131        assert_eq!(sliced.child_data()[0].offset(), 1);
3132        assert_eq!(sliced.buffers().as_ptr(), buffers);
3133        assert_eq!(sliced.child_data().as_ptr(), children);
3134        assert_eq!(sliced.child_data()[0].buffers().as_ptr(), child_buffers);
3135        sliced.validate_full().unwrap();
3136    }
3137
3138    #[test]
3139    #[should_panic(expected = "offset + length overflow")]
3140    fn test_sliced_panics_on_offset_length_overflow() {
3141        ArrayData::new_null(&DataType::Int32, 4).sliced(1, usize::MAX);
3142    }
3143
3144    #[test]
3145    #[should_panic(expected = "end <= self.len()")]
3146    fn test_sliced_panics_on_offset_out_of_bounds() {
3147        ArrayData::new_null(&DataType::Int32, 4).sliced(5, 0);
3148    }
3149
3150    #[test]
3151    #[should_panic(expected = "end <= self.len()")]
3152    fn test_sliced_panics_on_length_out_of_bounds() {
3153        ArrayData::new_null(&DataType::Int32, 4).sliced(3, 2);
3154    }
3155
3156    #[test]
3157    fn test_typed_offsets_length_overflow() {
3158        let data = ArrayData {
3159            data_type: DataType::Binary,
3160            len: usize::MAX,
3161            offset: 0,
3162            buffers: vec![Buffer::from_slice_ref([0_i32])],
3163            child_data: vec![],
3164            nulls: None,
3165        };
3166        let err = data.typed_offsets::<i32>().unwrap_err();
3167
3168        assert_eq!(
3169            err.to_string(),
3170            format!(
3171                "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
3172                usize::MAX
3173            )
3174        );
3175    }
3176
3177    #[test]
3178    fn test_validate_typed_buffer_length_overflow() {
3179        let data = ArrayData {
3180            data_type: DataType::Binary,
3181            len: 0,
3182            offset: 2,
3183            buffers: vec![Buffer::from_slice_ref([0_i32])],
3184            child_data: vec![],
3185            nulls: None,
3186        };
3187        let err = data.typed_buffer::<i32>(0, usize::MAX).unwrap_err();
3188
3189        assert_eq!(
3190            err.to_string(),
3191            format!(
3192                "Invalid argument error: Length {} with offset 2 overflows usize for Binary",
3193                usize::MAX
3194            )
3195        );
3196    }
3197
3198    // Exercises ArrayData::try_new with len + offset overflowing
3199    fn try_new_binary_length_offset_overflow() -> Result<ArrayData, ArrowError> {
3200        ArrayData::try_new(
3201            DataType::Binary,
3202            usize::MAX,
3203            None,
3204            1,
3205            vec![
3206                Buffer::from_slice_ref([0_i32]),
3207                Buffer::from_iter(std::iter::empty::<u8>()),
3208            ],
3209            vec![],
3210        )
3211    }
3212
3213    #[cfg(not(feature = "force_validate"))]
3214    #[test]
3215    fn test_try_new_length_offset_overflow() {
3216        let err = try_new_binary_length_offset_overflow().unwrap_err();
3217
3218        assert_eq!(
3219            err.to_string(),
3220            format!(
3221                "Invalid argument error: Length {} with offset 1 overflows usize for Binary",
3222                usize::MAX
3223            )
3224        );
3225    }
3226
3227    #[cfg(feature = "force_validate")]
3228    #[test]
3229    #[should_panic(
3230        expected = "Length 18446744073709551615 with offset 1 overflows usize for Binary"
3231    )]
3232    fn test_try_new_length_offset_overflow_force_validate() {
3233        try_new_binary_length_offset_overflow().unwrap();
3234    }
3235
3236    #[test]
3237    fn test_equality() {
3238        let int_data = ArrayData::builder(DataType::Int32)
3239            .len(1)
3240            .add_buffer(make_i32_buffer(1))
3241            .build()
3242            .unwrap();
3243
3244        let float_data = ArrayData::builder(DataType::Float32)
3245            .len(1)
3246            .add_buffer(make_f32_buffer(1))
3247            .build()
3248            .unwrap();
3249        assert_ne!(int_data, float_data);
3250        assert!(!int_data.ptr_eq(&float_data));
3251        assert!(int_data.ptr_eq(&int_data));
3252
3253        let int_data_clone = int_data.clone();
3254        assert_eq!(int_data, int_data_clone);
3255        assert!(int_data.ptr_eq(&int_data_clone));
3256        assert!(int_data_clone.ptr_eq(&int_data));
3257
3258        let int_data_slice = int_data_clone.slice(1, 0);
3259        assert!(int_data_slice.ptr_eq(&int_data_slice));
3260        assert!(!int_data.ptr_eq(&int_data_slice));
3261        assert!(!int_data_slice.ptr_eq(&int_data));
3262
3263        let data_buffer = Buffer::from_slice_ref(b"abcdef");
3264        let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
3265        let string_data = ArrayData::try_new(
3266            DataType::Utf8,
3267            3,
3268            Some(Buffer::from_iter(vec![true, false, true])),
3269            0,
3270            vec![offsets_buffer, data_buffer],
3271            vec![],
3272        )
3273        .unwrap();
3274
3275        assert_ne!(float_data, string_data);
3276        assert!(!float_data.ptr_eq(&string_data));
3277
3278        assert!(string_data.ptr_eq(&string_data));
3279
3280        let string_data_cloned = string_data.clone();
3281        assert!(string_data_cloned.ptr_eq(&string_data));
3282        assert!(string_data.ptr_eq(&string_data_cloned));
3283
3284        let string_data_slice = string_data.slice(1, 2);
3285        assert!(string_data_slice.ptr_eq(&string_data_slice));
3286        assert!(!string_data_slice.ptr_eq(&string_data))
3287    }
3288
3289    #[test]
3290    fn test_slice_memory_size_view_payload_buffers() {
3291        for data_type in [DataType::Utf8View, DataType::BinaryView] {
3292            let inline_only = ArrayData::builder(data_type.clone())
3293                .len(2)
3294                .add_buffer(Buffer::from_vec(vec![0_u128; 2]))
3295                .build()
3296                .unwrap();
3297            assert_eq!(
3298                inline_only.get_slice_memory_size().unwrap(),
3299                2 * mem::size_of::<u128>()
3300            );
3301
3302            let mut first_payload = Vec::with_capacity(32);
3303            first_payload.extend_from_slice(b"first payload");
3304            let first_view =
3305                ByteView::new(first_payload.len().try_into().unwrap(), &first_payload[..4])
3306                    .as_u128();
3307            let first_payload = Buffer::from_vec(first_payload);
3308            assert!(first_payload.capacity() > first_payload.len());
3309            let first_payload_capacity = first_payload.capacity();
3310
3311            let mut second_payload = Vec::with_capacity(64);
3312            second_payload.extend_from_slice(b"second payload");
3313            let second_view = ByteView::new(
3314                second_payload.len().try_into().unwrap(),
3315                &second_payload[..4],
3316            )
3317            .with_buffer_index(1)
3318            .as_u128();
3319            let second_payload = Buffer::from_vec(second_payload);
3320            assert!(second_payload.capacity() > second_payload.len());
3321            let second_payload_capacity = second_payload.capacity();
3322
3323            let data = ArrayData::builder(data_type)
3324                .len(3)
3325                .add_buffer(Buffer::from_vec(vec![first_view, 0_u128, second_view]))
3326                .add_buffer(first_payload)
3327                .add_buffer(second_payload)
3328                .build()
3329                .unwrap();
3330            let sliced = data.slice(1, 1);
3331
3332            assert_eq!(
3333                sliced.get_slice_memory_size().unwrap(),
3334                mem::size_of::<u128>() + first_payload_capacity + second_payload_capacity
3335            );
3336        }
3337    }
3338
3339    #[test]
3340    fn test_slice_memory_size_utf8_offset_buffer_len_plus_one() {
3341        // 2-element array ["hello", "world"]: array len = 2, 10 bytes
3342        let data_buffer = Buffer::from_slice_ref(b"helloworld");
3343        // offsets need array_len+1 entries to mark the end of every string:
3344        //   [0, 5, 10] -> 3 i32s = 12 bytes
3345        let offsets_buffer = Buffer::from_slice_ref([0_i32, 5_i32, 10_i32]);
3346        let array = ArrayData::try_new(
3347            DataType::Utf8,
3348            2,
3349            None,
3350            0,
3351            vec![offsets_buffer, data_buffer],
3352            vec![],
3353        )
3354        .unwrap();
3355        assert_eq!(array.get_slice_memory_size().unwrap(), 22); // 12 + 10
3356    }
3357
3358    #[test]
3359    fn test_slice_memory_size_binary_offset_buffer_len_plus_one() {
3360        // 2-element array: array len = 2, not 3
3361        // values: 5 bytes
3362        let data_buffer = Buffer::from_slice_ref([0u8, 1, 2, 3, 4]);
3363        // offsets need array_len+1 entries to mark the end of every element:
3364        let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 5_i32]);
3365        let array = ArrayData::try_new(
3366            DataType::Binary,
3367            2,
3368            None,
3369            0,
3370            vec![offsets_buffer, data_buffer],
3371            vec![],
3372        )
3373        .unwrap();
3374        assert_eq!(array.get_slice_memory_size().unwrap(), 17); // 12 + 5
3375    }
3376
3377    #[test]
3378    fn test_slice_memory_size() {
3379        let mut bit_v: [u8; 2] = [0; 2];
3380        bit_util::set_bit(&mut bit_v, 0);
3381        bit_util::set_bit(&mut bit_v, 3);
3382        bit_util::set_bit(&mut bit_v, 10);
3383        let data = ArrayData::builder(DataType::Int32)
3384            .len(16)
3385            .add_buffer(make_i32_buffer(16))
3386            .null_bit_buffer(Some(Buffer::from(bit_v)))
3387            .build()
3388            .unwrap();
3389        let new_data = data.slice(1, 14);
3390        assert_eq!(
3391            data.get_slice_memory_size().unwrap() - 8,
3392            new_data.get_slice_memory_size().unwrap()
3393        );
3394        let data_buffer = Buffer::from_slice_ref(b"abcdef");
3395        let offsets_buffer = Buffer::from_slice_ref([0_i32, 2_i32, 2_i32, 5_i32]);
3396        let string_data = ArrayData::try_new(
3397            DataType::Utf8,
3398            3,
3399            Some(Buffer::from_iter(vec![true, false, true])),
3400            0,
3401            vec![offsets_buffer, data_buffer],
3402            vec![],
3403        )
3404        .unwrap();
3405        let string_data_slice = string_data.slice(1, 2);
3406        //4 bytes of offset and 2 bytes of data reduced by slicing.
3407        assert_eq!(
3408            string_data.get_slice_memory_size().unwrap() - 6,
3409            string_data_slice.get_slice_memory_size().unwrap()
3410        );
3411    }
3412
3413    #[test]
3414    fn test_builder_rejects_short_null_bit_buffer() {
3415        for (len, offset) in [(8000, 0), (8, 1)] {
3416            let err = ArrayData::builder(DataType::Int32)
3417                .len(len)
3418                .offset(offset)
3419                .add_buffer(make_i32_buffer(len + offset))
3420                .null_bit_buffer(Some(Buffer::from([0_u8])))
3421                .build()
3422                .unwrap_err();
3423            assert_eq!(
3424                err.to_string(),
3425                format!(
3426                    "Invalid argument error: null_bit_buffer size too small. got 1 needed {}",
3427                    bit_util::ceil(len + offset, 8)
3428                )
3429            );
3430        }
3431    }
3432
3433    #[test]
3434    fn test_builder_null_bit_buffer_length_overflow() {
3435        let err = ArrayData::builder(DataType::Int32)
3436            .len(usize::MAX)
3437            .offset(1)
3438            .null_bit_buffer(Some(Buffer::default()))
3439            .build()
3440            .unwrap_err();
3441        assert_eq!(
3442            err.to_string(),
3443            format!(
3444                "Invalid argument error: Length {} with offset 1 overflows usize for Int32",
3445                usize::MAX
3446            )
3447        );
3448    }
3449
3450    #[test]
3451    fn test_builder_accepts_valid_null_bit_buffer() {
3452        for (len, offset) in [(8, 0), (7, 1)] {
3453            let data = ArrayData::builder(DataType::Int32)
3454                .len(len)
3455                .offset(offset)
3456                .add_buffer(make_i32_buffer(len + offset))
3457                .null_bit_buffer(Some(Buffer::from([0_u8])))
3458                .build()
3459                .unwrap();
3460            assert_eq!(data.len(), len);
3461            assert_eq!(data.offset(), offset);
3462            assert_eq!(data.null_count(), len);
3463        }
3464    }
3465
3466    #[test]
3467    fn test_count_nulls() {
3468        let buffer = Buffer::from([0b00010110, 0b10011111]);
3469        let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 16));
3470        let count = count_nulls(Some(&buffer), 0, 16);
3471        assert_eq!(count, 7);
3472
3473        let count = count_nulls(Some(&buffer), 4, 8);
3474        assert_eq!(count, 3);
3475    }
3476
3477    #[test]
3478    fn test_contains_nulls() {
3479        let buffer: Buffer =
3480            MutableBuffer::from_iter([false, false, false, true, true, false]).into();
3481        let buffer = NullBuffer::new(BooleanBuffer::new(buffer, 0, 6));
3482        assert!(contains_nulls(Some(&buffer), 0, 6));
3483        assert!(contains_nulls(Some(&buffer), 0, 3));
3484        assert!(!contains_nulls(Some(&buffer), 3, 2));
3485        assert!(!contains_nulls(Some(&buffer), 0, 0));
3486    }
3487
3488    #[test]
3489    fn test_alignment() {
3490        let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
3491        let sliced = buffer.slice(1);
3492
3493        let mut data = ArrayData {
3494            data_type: DataType::Int32,
3495            len: 0,
3496            offset: 0,
3497            buffers: vec![buffer],
3498            child_data: vec![],
3499            nulls: None,
3500        };
3501        data.validate_full().unwrap();
3502
3503        // break alignment in data
3504        data.buffers[0] = sliced;
3505        let err = data.validate().unwrap_err();
3506
3507        assert_eq!(
3508            err.to_string(),
3509            "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
3510        );
3511
3512        data.align_buffers();
3513        data.validate_full().unwrap();
3514    }
3515
3516    #[test]
3517    fn test_alignment_struct() {
3518        let buffer = Buffer::from_vec(vec![1_i32, 2_i32, 3_i32]);
3519        let sliced = buffer.slice(1);
3520
3521        let child_data = ArrayData {
3522            data_type: DataType::Int32,
3523            len: 0,
3524            offset: 0,
3525            buffers: vec![buffer],
3526            child_data: vec![],
3527            nulls: None,
3528        };
3529
3530        let schema = DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, false)]));
3531        let mut data = ArrayData {
3532            data_type: schema,
3533            len: 0,
3534            offset: 0,
3535            buffers: vec![],
3536            child_data: vec![child_data],
3537            nulls: None,
3538        };
3539        data.validate_full().unwrap();
3540
3541        // break alignment in child data
3542        data.child_data[0].buffers[0] = sliced;
3543        let err = data.validate().unwrap_err();
3544
3545        assert_eq!(
3546            err.to_string(),
3547            "Invalid argument error: Misaligned buffers[0] in array of type Int32, offset from expected alignment of 4 by 1"
3548        );
3549
3550        data.align_buffers();
3551        data.validate_full().unwrap();
3552    }
3553
3554    #[test]
3555    fn test_null_view_types() {
3556        let array_len = 32;
3557        let array = ArrayData::new_null(&DataType::BinaryView, array_len);
3558        assert_eq!(array.len(), array_len);
3559        for i in 0..array.len() {
3560            assert!(array.is_null(i));
3561        }
3562
3563        let array = ArrayData::new_null(&DataType::Utf8View, array_len);
3564        assert_eq!(array.len(), array_len);
3565        for i in 0..array.len() {
3566            assert!(array.is_null(i));
3567        }
3568
3569        let array = ArrayData::new_null(
3570            &DataType::ListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
3571            array_len,
3572        );
3573        assert_eq!(array.len(), array_len);
3574        for i in 0..array.len() {
3575            assert!(array.is_null(i));
3576        }
3577
3578        let array = ArrayData::new_null(
3579            &DataType::LargeListView(Arc::new(Field::new_list_field(DataType::Int32, true))),
3580            array_len,
3581        );
3582        assert_eq!(array.len(), array_len);
3583        for i in 0..array.len() {
3584            assert!(array.is_null(i));
3585        }
3586    }
3587
3588    // Even when `force_validate` feature is on
3589    #[test]
3590    fn test_dont_panic_on_bad_input_when_using_try_new() {
3591        let empty_bytes = Buffer::default();
3592
3593        let array_data = ArrayData::try_new(
3594            DataType::Utf8,
3595            1, // len
3596            None,
3597            0,
3598            // the offsets says that we have 2 bytes but the buffer is empty
3599            vec![Buffer::from_vec(vec![0i32, 2i32]), empty_bytes],
3600            vec![],
3601        );
3602
3603        let res = array_data.expect_err("should get error");
3604
3605        assert_eq!(
3606            res.to_string(),
3607            "Invalid argument error: Last offset 2 of Utf8 is larger than values length 0"
3608        );
3609    }
3610
3611    /// Without `force_validate`, `build_unchecked` skips validation, so these can
3612    /// reach `validate_values` with a data type that has an invalid child type.
3613    #[test]
3614    #[cfg(not(feature = "force_validate"))]
3615    fn test_validate_values_rejects_a_non_integer_dictionary_key() {
3616        let values = valid_non_nullable_int32_array_data(2);
3617        let data_type = DataType::Dictionary(Box::new(DataType::Utf8), Box::new(DataType::Int32));
3618        let dictionary = unsafe {
3619            ArrayData::builder(data_type)
3620                .len(1)
3621                .add_child_data(values)
3622                .build_unchecked()
3623        };
3624
3625        let err = dictionary.validate_values().expect_err("should get error");
3626        assert_eq!(
3627            err.to_string(),
3628            "Invalid argument error: Dictionary key type must be an integer, got Utf8"
3629        );
3630    }
3631
3632    #[test]
3633    #[cfg(not(feature = "force_validate"))]
3634    fn test_validate_values_rejects_a_non_integer_run_end() {
3635        let data_type = DataType::RunEndEncoded(
3636            Arc::new(Field::new(
3637                Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
3638                DataType::Utf8,
3639                false,
3640            )),
3641            Arc::new(Field::new(
3642                Field::REE_VALUES_FIELD_DEFAULT_NAME,
3643                DataType::Int32,
3644                true,
3645            )),
3646        );
3647        let run_end_encoded = unsafe {
3648            ArrayData::builder(data_type)
3649                .len(1)
3650                .add_child_data(valid_non_nullable_int32_array_data(1))
3651                .add_child_data(valid_non_nullable_int32_array_data(1))
3652                .build_unchecked()
3653        };
3654
3655        let err = run_end_encoded
3656            .validate_values()
3657            .expect_err("should get error");
3658        assert_eq!(
3659            err.to_string(),
3660            "Invalid argument error: Run end type must be Int16, Int32 or Int64, got Utf8"
3661        );
3662    }
3663
3664    /// Returns a `RunEndEncoded<Int32, Float32>` data type and run ends / values
3665    /// children whose runs end at `[1, 3]`, i.e. a logical length of 3.
3666    fn run_array_test_children() -> (DataType, Vec<ArrayData>) {
3667        let dt = DataType::RunEndEncoded(
3668            Arc::new(Field::new(
3669                Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
3670                DataType::Int32,
3671                false,
3672            )),
3673            Arc::new(Field::new(
3674                Field::REE_VALUES_FIELD_DEFAULT_NAME,
3675                DataType::Float32,
3676                true,
3677            )),
3678        );
3679
3680        let runs = ArrayData::builder(DataType::Int32)
3681            .len(2)
3682            .add_buffer(vec![1_i32, 3].into())
3683            .build()
3684            .unwrap();
3685        let values = ArrayData::builder(DataType::Float32)
3686            .len(2)
3687            .add_buffer(vec![1.5_f32, 2.5].into())
3688            .build()
3689            .unwrap();
3690
3691        (dt, vec![runs, values])
3692    }
3693
3694    #[test]
3695    fn test_run_array_validates_last_run() {
3696        let (dt, children) = run_array_test_children();
3697
3698        let err = ArrayDataBuilder::new(dt)
3699            .len(4)
3700            .offset(0)
3701            .child_data(children)
3702            .build()
3703            .unwrap_err();
3704        assert_eq!(
3705            err.to_string(),
3706            "Invalid argument error: The offset + length of array should be less or equal to last value in the run_ends array. The last value of run_ends array is 3 and offset + length of array is 4."
3707        );
3708    }
3709
3710    #[test]
3711    fn test_run_array_empty_validates() {
3712        let (dt, children) = run_array_test_children();
3713
3714        let data = ArrayDataBuilder::new(dt)
3715            .len(0)
3716            .offset(4)
3717            .child_data(children)
3718            .build()
3719            .unwrap();
3720        assert_eq!(data.len(), 0);
3721    }
3722
3723    /// The parent's offset must be included when checking the last run end:
3724    /// offset 2 + length 2 = 4 exceeds the last run end of 3.
3725    #[test]
3726    fn test_run_array_validates_last_run_with_offset() {
3727        let (dt, children) = run_array_test_children();
3728
3729        let err = ArrayDataBuilder::new(dt)
3730            .len(2)
3731            .offset(2)
3732            .child_data(children)
3733            .build()
3734            .unwrap_err();
3735        assert_eq!(
3736            err.to_string(),
3737            "Invalid argument error: The offset + length of array should be less or equal to last value in the run_ends array. The last value of run_ends array is 3 and offset + length of array is 4."
3738        );
3739    }
3740
3741    /// A slice that stays within the last run end is valid:
3742    /// offset 1 + length 2 = 3 equals the last run end of 3.
3743    #[test]
3744    fn test_run_array_sliced_validates() {
3745        let (dt, children) = run_array_test_children();
3746
3747        let data = ArrayDataBuilder::new(dt)
3748            .len(2)
3749            .offset(1)
3750            .child_data(children)
3751            .build()
3752            .unwrap();
3753        assert_eq!(data.len(), 2);
3754        assert_eq!(data.offset(), 1);
3755    }
3756
3757    /// `validate_values` must report missing children rather than index out of bounds.
3758    #[test]
3759    #[cfg(not(feature = "force_validate"))]
3760    fn test_validate_values_rejects_missing_child_data() {
3761        let int32 = Box::new(DataType::Int32);
3762        let field = || Arc::new(Field::new("f", DataType::Int32, true));
3763        let data_types = [
3764            DataType::Dictionary(int32.clone(), int32.clone()),
3765            DataType::List(field()),
3766            DataType::LargeList(field()),
3767            DataType::RunEndEncoded(field(), field()),
3768        ];
3769
3770        for data_type in data_types {
3771            let data = unsafe {
3772                ArrayData::builder(data_type.clone())
3773                    .len(1)
3774                    .build_unchecked()
3775            };
3776            let err = data.validate_values().expect_err("should get error");
3777            assert_eq!(
3778                err.to_string(),
3779                format!(
3780                    "Invalid argument error: {data_type} should contain at least 1 child data array(s), had 0"
3781                )
3782            );
3783        }
3784    }
3785
3786    /// `validate_values` must report missing buffers rather than index out of bounds.
3787    #[test]
3788    #[cfg(not(feature = "force_validate"))]
3789    fn test_validate_values_rejects_missing_buffers() {
3790        // (data type, index of the first missing buffer)
3791        let cases = [
3792            (DataType::Utf8, 1),
3793            (DataType::LargeUtf8, 1),
3794            (DataType::Binary, 1),
3795            (DataType::LargeBinary, 1),
3796            (DataType::BinaryView, 0),
3797            (DataType::Utf8View, 0),
3798        ];
3799
3800        for (data_type, missing) in cases {
3801            let data = unsafe {
3802                ArrayData::builder(data_type.clone())
3803                    .len(1)
3804                    .build_unchecked()
3805            };
3806            let err = data.validate_values().expect_err("should get error");
3807            assert_eq!(
3808                err.to_string(),
3809                format!(
3810                    "Invalid argument error: {data_type} should contain at least {} buffer(s), had 0",
3811                    missing + 1
3812                )
3813            );
3814        }
3815    }
3816
3817    /// A dictionary whose keys buffer is too small must be reported, not asserted on.
3818    #[test]
3819    #[cfg(not(feature = "force_validate"))]
3820    fn test_validate_values_rejects_a_short_dictionary_keys_buffer() {
3821        let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Int32));
3822        let dictionary = unsafe {
3823            ArrayData::builder(data_type)
3824                .len(4)
3825                .add_buffer(Buffer::from_slice_ref([1_i32, 0]))
3826                .add_child_data(valid_non_nullable_int32_array_data(2))
3827                .build_unchecked()
3828        };
3829
3830        let err = dictionary.validate_values().expect_err("should get error");
3831        assert_eq!(
3832            err.to_string(),
3833            "Invalid argument error: Buffer 0 of Dictionary(Int32, Int32) isn't large enough. Expected 16 bytes got 8"
3834        );
3835    }
3836
3837    /// `validate` only requires a minimum buffer size, so an offsets buffer
3838    /// may carry trailing bytes that are not a whole number of offsets.
3839    #[test]
3840    fn test_validate_offsets_buffer_with_trailing_bytes() {
3841        let mut offsets = Buffer::from_slice_ref([0_i32, 2, 5]).as_slice().to_vec();
3842        offsets.extend_from_slice(&[0, 0]);
3843        let data = ArrayData::builder(DataType::Utf8)
3844            .len(2)
3845            .add_buffer(Buffer::from_slice_ref(&offsets))
3846            .add_buffer(Buffer::from_slice_ref(b"hello"))
3847            .build()
3848            .unwrap();
3849        data.validate_full().unwrap();
3850    }
3851
3852    #[test]
3853    fn should_fail_validation_when_having_map_field_type_is_not_struct() {
3854        let map_field = Field::new("key", DataType::Int32, false);
3855
3856        let map_field_data = valid_non_nullable_int32_array_data(2);
3857
3858        let results = test_both_builder_and_array_data(
3859            DataType::Map(map_field.into(), false),
3860            1,
3861            None,
3862            0,
3863            vec![
3864                OffsetBuffer::<i32>::from_lengths(vec![2])
3865                    .into_inner()
3866                    .into(),
3867            ],
3868            vec![map_field_data],
3869        );
3870
3871        for result in results {
3872            let array_data_err = result.expect_err("should fail for non struct field");
3873
3874            match array_data_err {
3875                ArrowError::InvalidArgumentError(msg) => {
3876                    assert_eq!(
3877                        msg,
3878                        "Map field should be a entries struct data type, got Int32 instead"
3879                    )
3880                }
3881                _ => panic!("unexpected error type {array_data_err}"),
3882            }
3883        }
3884    }
3885
3886    #[test]
3887    fn should_fail_validation_when_having_map_entries_only_have_1_field() {
3888        let struct_data_type = DataType::Struct(Fields::from(vec![Field::new(
3889            Field::MAP_KEY_FIELD_DEFAULT_NAME,
3890            DataType::Int32,
3891            false,
3892        )]));
3893
3894        let key_array_data = valid_non_nullable_int32_array_data(2);
3895
3896        let struct_data = {
3897            let builder = ArrayDataBuilder::new(struct_data_type.clone())
3898                .len(2)
3899                .nulls(None)
3900                .child_data(vec![key_array_data]);
3901
3902            builder.build().unwrap()
3903        };
3904
3905        let results = test_both_builder_and_array_data(
3906            DataType::Map(
3907                Field::new(
3908                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3909                    struct_data_type,
3910                    false,
3911                )
3912                .into(),
3913                false,
3914            ),
3915            1,
3916            None,
3917            0,
3918            vec![
3919                OffsetBuffer::<i32>::from_lengths(vec![2])
3920                    .into_inner()
3921                    .into(),
3922            ],
3923            vec![struct_data],
3924        );
3925
3926        for result in results {
3927            let array_data_err = result.expect_err("should fail for nullable key");
3928
3929            match array_data_err {
3930                ArrowError::InvalidArgumentError(msg) => {
3931                    assert_eq!(
3932                        msg,
3933                        "Map entries data type should be a struct containing 2 fields, got 1 fields"
3934                    )
3935                }
3936                _ => panic!("unexpected error type {array_data_err}"),
3937            }
3938        }
3939    }
3940
3941    #[test]
3942    fn should_fail_validation_when_having_map_entries_have_3_fields() {
3943        let struct_data_type = DataType::Struct(Fields::from(vec![
3944            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
3945            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
3946            Field::new("other", DataType::Int32, true),
3947        ]));
3948
3949        let key_array_data = valid_non_nullable_int32_array_data(2);
3950
3951        let values_array_data = valid_string_array_data(2);
3952
3953        let other_array_data = key_array_data.clone();
3954
3955        let struct_data = {
3956            let builder = ArrayDataBuilder::new(struct_data_type.clone())
3957                .len(2)
3958                .nulls(None)
3959                .child_data(vec![key_array_data, values_array_data, other_array_data]);
3960
3961            builder.build().unwrap()
3962        };
3963
3964        let results = test_both_builder_and_array_data(
3965            DataType::Map(
3966                Field::new(
3967                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
3968                    struct_data_type,
3969                    false,
3970                )
3971                .into(),
3972                false,
3973            ),
3974            1,
3975            None,
3976            0,
3977            vec![
3978                OffsetBuffer::<i32>::from_lengths(vec![2])
3979                    .into_inner()
3980                    .into(),
3981            ],
3982            vec![struct_data],
3983        );
3984
3985        for result in results {
3986            let array_data_err = result.expect_err("should fail for nullable key");
3987
3988            match array_data_err {
3989                ArrowError::InvalidArgumentError(msg) => {
3990                    assert_eq!(
3991                        msg,
3992                        "Map entries data type should be a struct containing 2 fields, got 3 fields"
3993                    )
3994                }
3995                _ => panic!("unexpected error type {array_data_err}"),
3996            }
3997        }
3998    }
3999
4000    #[test]
4001    fn should_fail_validation_when_having_nullable_map_keys() {
4002        let struct_data_type = DataType::Struct(Fields::from(vec![
4003            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, true),
4004            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
4005        ]));
4006
4007        let key_array_data = valid_non_nullable_int32_array_data(2);
4008        let values_array_data = valid_string_array_data(2);
4009
4010        let struct_data = {
4011            let builder = ArrayDataBuilder::new(struct_data_type.clone())
4012                .len(2)
4013                .nulls(None)
4014                .child_data(vec![key_array_data, values_array_data]);
4015
4016            builder.build().unwrap()
4017        };
4018
4019        let results = test_both_builder_and_array_data(
4020            DataType::Map(
4021                Field::new(
4022                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4023                    struct_data_type,
4024                    false,
4025                )
4026                .into(),
4027                false,
4028            ),
4029            1,
4030            None,
4031            0,
4032            vec![
4033                OffsetBuffer::<i32>::from_lengths(vec![2])
4034                    .into_inner()
4035                    .into(),
4036            ],
4037            vec![struct_data],
4038        );
4039
4040        for result in results {
4041            let array_data_err = result.expect_err("should fail for nullable key");
4042
4043            match array_data_err {
4044                ArrowError::InvalidArgumentError(msg) => {
4045                    assert_eq!(msg, "Map key field must not be nullable")
4046                }
4047                _ => panic!("unexpected error type {array_data_err}"),
4048            }
4049        }
4050    }
4051
4052    #[test]
4053    fn should_fail_validation_when_having_entries_is_nullable_for_map() {
4054        let struct_data_type = DataType::Struct(Fields::from(vec![
4055            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
4056            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
4057        ]));
4058
4059        let key_array_data = valid_non_nullable_int32_array_data(2);
4060
4061        let values_array_data = valid_string_array_data(2);
4062
4063        let struct_data = {
4064            let builder = ArrayDataBuilder::new(struct_data_type.clone())
4065                .len(2)
4066                .nulls(None)
4067                .child_data(vec![key_array_data, values_array_data]);
4068
4069            builder.build().unwrap()
4070        };
4071
4072        let results = test_both_builder_and_array_data(
4073            DataType::Map(
4074                Field::new(
4075                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4076                    struct_data_type,
4077                    true,
4078                )
4079                .into(),
4080                false,
4081            ),
4082            1,
4083            None,
4084            0,
4085            vec![
4086                OffsetBuffer::<i32>::from_lengths(vec![2])
4087                    .into_inner()
4088                    .into(),
4089            ],
4090            vec![struct_data],
4091        );
4092
4093        for result in results {
4094            let array_data_err = result.expect_err("should fail for nullable entries");
4095
4096            match array_data_err {
4097                ArrowError::InvalidArgumentError(msg) => assert_eq!(
4098                    msg,
4099                    "The nullable should be set to false for the map entries field."
4100                ),
4101                _ => panic!("unexpected error type {array_data_err}"),
4102            }
4103        }
4104    }
4105
4106    #[test]
4107    fn should_allow_to_create_map_from_data() {
4108        let struct_data_type = DataType::Struct(Fields::from(vec![
4109            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
4110            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
4111        ]));
4112
4113        let key_array_data = valid_non_nullable_int32_array_data(2);
4114        let values_array_data = valid_string_array_data(2);
4115
4116        let struct_data = {
4117            let builder = ArrayDataBuilder::new(struct_data_type.clone())
4118                .len(2)
4119                .nulls(None)
4120                .child_data(vec![key_array_data, values_array_data]);
4121
4122            builder.build().unwrap()
4123        };
4124
4125        let results = test_both_builder_and_array_data(
4126            DataType::Map(
4127                Field::new(
4128                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4129                    struct_data_type,
4130                    false,
4131                )
4132                .into(),
4133                false,
4134            ),
4135            1,
4136            None,
4137            0,
4138            vec![
4139                OffsetBuffer::<i32>::from_lengths(vec![2])
4140                    .into_inner()
4141                    .into(),
4142            ],
4143            vec![struct_data],
4144        );
4145
4146        for result in results {
4147            result.expect("should be able to create map ArrayData");
4148        }
4149    }
4150
4151    fn valid_string_array_data(length: usize) -> ArrayData {
4152        let offsets = OffsetBuffer::<i32>::from_lengths(vec![0; length])
4153            .into_inner()
4154            .into_inner();
4155        let empty_bytes = Buffer::default();
4156
4157        let builder = ArrayDataBuilder::new(DataType::Utf8)
4158            .len(length)
4159            .buffers(vec![offsets, empty_bytes])
4160            .nulls(None);
4161
4162        builder.build().unwrap()
4163    }
4164
4165    fn valid_non_nullable_int32_array_data(length: usize) -> ArrayData {
4166        let builder = ArrayDataBuilder::new(DataType::Int32)
4167            .len(length)
4168            .nulls(None)
4169            .buffers(vec![
4170                ScalarBuffer::<i32>::from(vec![1; length]).into_inner(),
4171            ]);
4172
4173        builder.build().unwrap()
4174    }
4175
4176    #[test]
4177    fn empty_and_null_map_array_should_pass_validation() {
4178        let dt = DataType::Map(
4179            Field::new(
4180                Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
4181                DataType::Struct(Fields::from(vec![
4182                    Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Int32, false),
4183                    Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
4184                ])),
4185                false,
4186            )
4187            .into(),
4188            false,
4189        );
4190
4191        ArrayData::new_empty(&dt).validate_full().unwrap();
4192        ArrayData::new_null(&dt, 1).validate_full().unwrap();
4193    }
4194
4195    #[test]
4196    fn null_buffer_offset_is_independent_of_data_offset() {
4197        // 100 values sliced down to the last 50, so the data has offset 50.
4198        let int_data = ArrayData::builder(DataType::UInt32)
4199            .offset(50)
4200            .len(50)
4201            .add_buffer(Buffer::from_vec(vec![0_u32; 100]))
4202            .build()
4203            .unwrap();
4204        int_data.validate().unwrap();
4205
4206        // A null buffer that happens to share the data's offset.
4207        let nulls = NullBuffer::new(BooleanBuffer::from(vec![false; 100]).slice(0, 50));
4208        let with_sliced_nulls = int_data
4209            .clone()
4210            .into_builder()
4211            .nulls(Some(nulls))
4212            .build()
4213            .unwrap();
4214        with_sliced_nulls.validate().unwrap();
4215
4216        // The same 50 nulls at offset 0. ArrayData::offset does not apply to the
4217        // null buffer, so this is just as valid and must not be rejected.
4218        let nulls = NullBuffer::new(BooleanBuffer::from(vec![false; 50]));
4219        let with_unsliced_nulls = int_data.into_builder().nulls(Some(nulls)).build().unwrap();
4220        with_unsliced_nulls.validate().unwrap();
4221        assert_eq!(with_unsliced_nulls.null_count(), 50);
4222    }
4223
4224    fn test_both_builder_and_array_data(
4225        data_type: DataType,
4226        len: usize,
4227        null_bit_buffer: Option<Buffer>,
4228        offset: usize,
4229        buffers: Vec<Buffer>,
4230        child_data: Vec<ArrayData>,
4231    ) -> [Result<ArrayData, ArrowError>; 2] {
4232        let from_builder_res = ArrayData::builder(data_type.clone())
4233            .len(len)
4234            .add_buffers(buffers.clone())
4235            .null_bit_buffer(null_bit_buffer.clone())
4236            .offset(offset)
4237            .child_data(child_data.clone())
4238            .build();
4239
4240        let from_try_new_res =
4241            ArrayData::try_new(data_type, len, null_bit_buffer, offset, buffers, child_data);
4242
4243        [from_builder_res, from_try_new_res]
4244    }
4245
4246    #[test]
4247    fn test_new_null_empty_union() {
4248        for mode in [UnionMode::Sparse, UnionMode::Dense] {
4249            let data_type = DataType::Union(UnionFields::empty(), mode);
4250            let data = ArrayData::new_null(&data_type, 0);
4251            data.validate_full()
4252                .expect("an empty union of length zero is valid");
4253            assert_eq!(data.len(), 0);
4254            assert!(data.child_data().is_empty());
4255        }
4256    }
4257
4258    #[test]
4259    #[should_panic(expected = "cannot construct null data from an empty union")]
4260    fn test_new_null_empty_union_with_slots() {
4261        let data_type = DataType::Union(UnionFields::empty(), UnionMode::Dense);
4262        let _ = ArrayData::new_null(&data_type, 1);
4263    }
4264
4265    #[test]
4266    fn test_union_validate_type_ids() {
4267        let dt = DataType::Union(
4268            UnionFields::try_new(
4269                vec![5],
4270                vec![Arc::new(Field::new("a", DataType::Int32, true))],
4271            )
4272            .unwrap(),
4273            UnionMode::Sparse,
4274        );
4275
4276        let child = ArrayData::builder(DataType::Int32)
4277            .len(2)
4278            .add_buffer(make_i32_buffer(2))
4279            .build()
4280            .unwrap();
4281
4282        let err = ArrayDataBuilder::new(dt)
4283            .len(2)
4284            .offset(0)
4285            .buffers(vec![vec![5_i8, 2].into()])
4286            .child_data(vec![child])
4287            .build()
4288            .unwrap_err();
4289        assert_eq!(
4290            err.to_string(),
4291            "Invalid argument error: Type id values must match one of the field type ids, got 2 at index 1"
4292        );
4293    }
4294
4295    #[test]
4296    fn test_union_validate_dense_offsets() {
4297        let dt = DataType::Union(
4298            UnionFields::try_new(
4299                vec![5],
4300                vec![Arc::new(Field::new("a", DataType::Int32, true))],
4301            )
4302            .unwrap(),
4303            UnionMode::Dense,
4304        );
4305
4306        let child = ArrayData::builder(DataType::Int32)
4307            .len(2)
4308            .add_buffer(make_i32_buffer(2))
4309            .build()
4310            .unwrap();
4311
4312        let err = ArrayDataBuilder::new(dt.clone())
4313            .len(2)
4314            // 2nd offset '10' is out of bounds
4315            .buffers(vec![vec![5_i8, 5].into(), vec![0_i32, 10].into()])
4316            .child_data(vec![child.clone()])
4317            .build()
4318            .unwrap_err();
4319        assert_eq!(
4320            err.to_string(),
4321            "Invalid argument error: Offsets must be non-negative and within the length of the child array"
4322        );
4323
4324        let err = ArrayDataBuilder::new(dt)
4325            .len(2)
4326            // 2nd offset is negative
4327            .buffers(vec![vec![5_i8, 5].into(), vec![0_i32, -1].into()])
4328            .child_data(vec![child])
4329            .build()
4330            .unwrap_err();
4331        assert_eq!(
4332            err.to_string(),
4333            "Invalid argument error: Offsets must be non-negative and within the length of the child array"
4334        );
4335    }
4336}