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