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