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