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