Skip to main content

lance_arrow/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4//! Extend Arrow Functionality
5//!
6//! To improve Arrow-RS ergonomic
7
8#![warn(clippy::undocumented_unsafe_blocks)]
9
10// lance-arrow reinterprets value bytes as native numeric types in
11// `FloatArray::as_slice` for `bf16` (rust/lance-arrow/src/bfloat16.rs), which
12// requires the host byte order to match the on-disk byte order Lance writes.
13// Lance writes little-endian; building on a big-endian target would silently
14// produce wrong numeric values.
15#[cfg(not(target_endian = "little"))]
16compile_error!("lance-arrow only supports little-endian targets");
17
18use std::sync::Arc;
19use std::{collections::HashMap, ptr::NonNull};
20
21use arrow_array::{
22    Array, ArrayRef, ArrowNumericType, FixedSizeBinaryArray, FixedSizeListArray, GenericListArray,
23    LargeListArray, ListArray, MapArray, OffsetSizeTrait, PrimitiveArray, RecordBatch, StructArray,
24    UInt8Array, UInt32Array, cast::AsArray,
25};
26use arrow_array::{
27    Float32Array, Float64Array, Int8Array, Int16Array, Int32Array, Int64Array, new_null_array,
28};
29use arrow_buffer::MutableBuffer;
30use arrow_data::ArrayDataBuilder;
31use arrow_schema::{ArrowError, DataType, Field, Fields, IntervalUnit, Schema, SortOptions};
32use arrow_select::{interleave::interleave, take::take};
33use rand::prelude::*;
34
35pub mod deepcopy;
36pub mod schema;
37pub use schema::*;
38pub mod bfloat16;
39pub mod floats;
40use crate::list::ListArrayExt;
41pub use floats::*;
42
43pub mod ipc;
44pub mod json;
45pub mod list;
46pub mod memory;
47pub mod scalar;
48pub mod stream;
49pub mod r#struct;
50
51/// Arrow extension metadata key for extension name
52pub const ARROW_EXT_NAME_KEY: &str = "ARROW:extension:name";
53
54/// Arrow extension metadata key for extension metadata
55pub const ARROW_EXT_META_KEY: &str = "ARROW:extension:metadata";
56
57/// Key used by lance to mark a field as a blob
58/// TODO: Use Arrow extension mechanism instead?
59pub const BLOB_META_KEY: &str = "lance-encoding:blob";
60/// Arrow extension type name for Lance blob v2 columns
61pub const BLOB_V2_EXT_NAME: &str = "lance.blob.v2";
62/// Metadata key for overriding the dedicated blob size threshold (in bytes)
63pub const BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY: &str =
64    "lance-encoding:blob-dedicated-size-threshold";
65/// Metadata key for overriding the inline blob size threshold (in bytes)
66pub const BLOB_INLINE_SIZE_THRESHOLD_META_KEY: &str = "lance-encoding:blob-inline-size-threshold";
67/// Metadata key for overriding the maximum size (in bytes) of a packed blob sidecar file
68pub const BLOB_PACK_FILE_SIZE_THRESHOLD_META_KEY: &str =
69    "lance-encoding:blob-pack-file-size-threshold";
70
71type Result<T> = std::result::Result<T, ArrowError>;
72
73pub trait DataTypeExt {
74    /// Returns true if the data type is binary-like, such as Utf8, Binary, or the large and/or view variants.
75    fn is_binary_like(&self) -> bool;
76
77    /// Returns true if the data type is a struct.
78    fn is_struct(&self) -> bool;
79
80    /// Check whether the given Arrow DataType is fixed stride.
81    ///
82    /// A fixed stride type has the same byte width for all array elements
83    /// This includes all PrimitiveType's Boolean, FixedSizeList, FixedSizeBinary, and Decimals
84    fn is_fixed_stride(&self) -> bool;
85
86    /// Returns true if the [DataType] is a dictionary type.
87    fn is_dictionary(&self) -> bool;
88
89    /// Returns the byte width of the data type
90    /// Panics if the data type is not fixed stride.
91    fn byte_width(&self) -> usize;
92
93    /// Returns the byte width of the data type, if it is fixed stride.
94    /// Returns None if the data type is not fixed stride.
95    fn byte_width_opt(&self) -> Option<usize>;
96}
97
98impl DataTypeExt for DataType {
99    fn is_binary_like(&self) -> bool {
100        use DataType::*;
101        matches!(
102            self,
103            Utf8 | Binary | LargeUtf8 | LargeBinary | Utf8View | BinaryView
104        )
105    }
106
107    fn is_struct(&self) -> bool {
108        matches!(self, Self::Struct(_))
109    }
110
111    fn is_fixed_stride(&self) -> bool {
112        use DataType::*;
113        matches!(
114            self,
115            Boolean
116                | UInt8
117                | UInt16
118                | UInt32
119                | UInt64
120                | Int8
121                | Int16
122                | Int32
123                | Int64
124                | Float16
125                | Float32
126                | Float64
127                | Decimal128(_, _)
128                | Decimal256(_, _)
129                | FixedSizeList(_, _)
130                | FixedSizeBinary(_)
131                | Duration(_)
132                | Timestamp(_, _)
133                | Date32
134                | Date64
135                | Time32(_)
136                | Time64(_)
137        )
138    }
139
140    fn is_dictionary(&self) -> bool {
141        matches!(self, Self::Dictionary(_, _))
142    }
143
144    fn byte_width_opt(&self) -> Option<usize> {
145        match self {
146            Self::Int8 => Some(1),
147            Self::Int16 => Some(2),
148            Self::Int32 => Some(4),
149            Self::Int64 => Some(8),
150            Self::UInt8 => Some(1),
151            Self::UInt16 => Some(2),
152            Self::UInt32 => Some(4),
153            Self::UInt64 => Some(8),
154            Self::Float16 => Some(2),
155            Self::Float32 => Some(4),
156            Self::Float64 => Some(8),
157            Self::Date32 => Some(4),
158            Self::Date64 => Some(8),
159            Self::Time32(_) => Some(4),
160            Self::Time64(_) => Some(8),
161            Self::Timestamp(_, _) => Some(8),
162            Self::Duration(_) => Some(8),
163            Self::Decimal128(_, _) => Some(16),
164            Self::Decimal256(_, _) => Some(32),
165            Self::Interval(unit) => match unit {
166                IntervalUnit::YearMonth => Some(4),
167                IntervalUnit::DayTime => Some(8),
168                IntervalUnit::MonthDayNano => Some(16),
169            },
170            Self::FixedSizeBinary(s) => Some(*s as usize),
171            Self::FixedSizeList(dt, s) => dt
172                .data_type()
173                .byte_width_opt()
174                .map(|width| width * *s as usize),
175            _ => None,
176        }
177    }
178
179    fn byte_width(&self) -> usize {
180        self.byte_width_opt()
181            .unwrap_or_else(|| panic!("Expecting fixed stride data type, found {:?}", self))
182    }
183}
184
185/// Create an [`GenericListArray`] from values and offsets.
186///
187/// ```
188/// use arrow_array::{Int32Array, Int64Array, ListArray};
189/// use arrow_array::types::Int64Type;
190/// use lance_arrow::try_new_generic_list_array;
191///
192/// let offsets = Int32Array::from_iter([0, 2, 7, 10]);
193/// let int_values = Int64Array::from_iter(0..10);
194/// let list_arr = try_new_generic_list_array(int_values, &offsets).unwrap();
195/// assert_eq!(list_arr,
196///     ListArray::from_iter_primitive::<Int64Type, _, _>(vec![
197///         Some(vec![Some(0), Some(1)]),
198///         Some(vec![Some(2), Some(3), Some(4), Some(5), Some(6)]),
199///         Some(vec![Some(7), Some(8), Some(9)]),
200/// ]))
201/// ```
202pub fn try_new_generic_list_array<T: Array, Offset: ArrowNumericType>(
203    values: T,
204    offsets: &PrimitiveArray<Offset>,
205) -> Result<GenericListArray<Offset::Native>>
206where
207    Offset::Native: OffsetSizeTrait,
208{
209    let data_type = if Offset::Native::IS_LARGE {
210        DataType::LargeList(Arc::new(Field::new(
211            "item",
212            values.data_type().clone(),
213            true,
214        )))
215    } else {
216        DataType::List(Arc::new(Field::new(
217            "item",
218            values.data_type().clone(),
219            true,
220        )))
221    };
222    let data = ArrayDataBuilder::new(data_type)
223        .len(offsets.len() - 1)
224        .add_buffer(offsets.into_data().buffers()[0].clone())
225        .add_child_data(values.into_data())
226        .build()?;
227
228    Ok(GenericListArray::from(data))
229}
230
231pub fn fixed_size_list_type(list_width: i32, inner_type: DataType) -> DataType {
232    DataType::FixedSizeList(Arc::new(Field::new("item", inner_type, true)), list_width)
233}
234
235pub trait FixedSizeListArrayExt {
236    /// Create an [`FixedSizeListArray`] from values and list size.
237    ///
238    /// ```
239    /// use arrow_array::{Int64Array, FixedSizeListArray};
240    /// use arrow_array::types::Int64Type;
241    /// use lance_arrow::FixedSizeListArrayExt;
242    ///
243    /// let int_values = Int64Array::from_iter(0..10);
244    /// let fixed_size_list_arr = FixedSizeListArray::try_new_from_values(int_values, 2).unwrap();
245    /// assert_eq!(fixed_size_list_arr,
246    ///     FixedSizeListArray::from_iter_primitive::<Int64Type, _, _>(vec![
247    ///         Some(vec![Some(0), Some(1)]),
248    ///         Some(vec![Some(2), Some(3)]),
249    ///         Some(vec![Some(4), Some(5)]),
250    ///         Some(vec![Some(6), Some(7)]),
251    ///         Some(vec![Some(8), Some(9)])
252    /// ], 2))
253    /// ```
254    fn try_new_from_values<T: Array + 'static>(
255        values: T,
256        list_size: i32,
257    ) -> Result<FixedSizeListArray>;
258
259    /// Sample `n` rows from the [FixedSizeListArray]
260    ///
261    /// ```
262    /// use arrow_array::{Int64Array, FixedSizeListArray, Array};
263    /// use lance_arrow::FixedSizeListArrayExt;
264    ///
265    /// let int_values = Int64Array::from_iter(0..256);
266    /// let fixed_size_list_arr = FixedSizeListArray::try_new_from_values(int_values, 16).unwrap();
267    /// let sampled = fixed_size_list_arr.sample(10).unwrap();
268    /// assert_eq!(sampled.len(), 10);
269    /// assert_eq!(sampled.value_length(), 16);
270    /// assert_eq!(sampled.values().len(), 160);
271    /// ```
272    fn sample(&self, n: usize) -> Result<FixedSizeListArray>;
273
274    /// Ensure the [FixedSizeListArray] of Float16, Float32, Float64,
275    /// Int8, Int16, Int32, Int64, UInt8, UInt32 type to its closest floating point type.
276    fn convert_to_floating_point(&self) -> Result<FixedSizeListArray>;
277}
278
279impl FixedSizeListArrayExt for FixedSizeListArray {
280    fn try_new_from_values<T: Array + 'static>(values: T, list_size: i32) -> Result<Self> {
281        let field = Arc::new(Field::new("item", values.data_type().clone(), true));
282        let values = Arc::new(values);
283
284        Self::try_new(field, list_size, values, None)
285    }
286
287    fn sample(&self, n: usize) -> Result<FixedSizeListArray> {
288        if n >= self.len() {
289            return Ok(self.clone());
290        }
291        let mut rng = SmallRng::from_os_rng();
292        let chosen = (0..self.len() as u32).choose_multiple(&mut rng, n);
293        take(self, &UInt32Array::from(chosen), None).map(|arr| arr.as_fixed_size_list().clone())
294    }
295
296    fn convert_to_floating_point(&self) -> Result<FixedSizeListArray> {
297        match self.data_type() {
298            DataType::FixedSizeList(field, size) => match field.data_type() {
299                DataType::Float16 | DataType::Float32 | DataType::Float64 => Ok(self.clone()),
300                DataType::Int8 => Ok(Self::new(
301                    Arc::new(arrow_schema::Field::new(
302                        field.name(),
303                        DataType::Float32,
304                        field.is_nullable(),
305                    )),
306                    *size,
307                    Arc::new(Float32Array::from_iter(
308                        self.values()
309                            .as_any()
310                            .downcast_ref::<Int8Array>()
311                            .ok_or(ArrowError::ParseError(
312                                "Fail to cast primitive array to Int8Type".to_string(),
313                            ))?
314                            .into_iter()
315                            .map(|x| x.map(|y| y as f32)),
316                    )),
317                    self.nulls().cloned(),
318                )),
319                DataType::Int16 => Ok(Self::new(
320                    Arc::new(arrow_schema::Field::new(
321                        field.name(),
322                        DataType::Float32,
323                        field.is_nullable(),
324                    )),
325                    *size,
326                    Arc::new(Float32Array::from_iter(
327                        self.values()
328                            .as_any()
329                            .downcast_ref::<Int16Array>()
330                            .ok_or(ArrowError::ParseError(
331                                "Fail to cast primitive array to Int16Type".to_string(),
332                            ))?
333                            .into_iter()
334                            .map(|x| x.map(|y| y as f32)),
335                    )),
336                    self.nulls().cloned(),
337                )),
338                DataType::Int32 => Ok(Self::new(
339                    Arc::new(arrow_schema::Field::new(
340                        field.name(),
341                        DataType::Float32,
342                        field.is_nullable(),
343                    )),
344                    *size,
345                    Arc::new(Float32Array::from_iter(
346                        self.values()
347                            .as_any()
348                            .downcast_ref::<Int32Array>()
349                            .ok_or(ArrowError::ParseError(
350                                "Fail to cast primitive array to Int32Type".to_string(),
351                            ))?
352                            .into_iter()
353                            .map(|x| x.map(|y| y as f32)),
354                    )),
355                    self.nulls().cloned(),
356                )),
357                DataType::Int64 => Ok(Self::new(
358                    Arc::new(arrow_schema::Field::new(
359                        field.name(),
360                        DataType::Float64,
361                        field.is_nullable(),
362                    )),
363                    *size,
364                    Arc::new(Float64Array::from_iter(
365                        self.values()
366                            .as_any()
367                            .downcast_ref::<Int64Array>()
368                            .ok_or(ArrowError::ParseError(
369                                "Fail to cast primitive array to Int64Type".to_string(),
370                            ))?
371                            .into_iter()
372                            .map(|x| x.map(|y| y as f64)),
373                    )),
374                    self.nulls().cloned(),
375                )),
376                DataType::UInt8 => Ok(Self::new(
377                    Arc::new(arrow_schema::Field::new(
378                        field.name(),
379                        DataType::Float64,
380                        field.is_nullable(),
381                    )),
382                    *size,
383                    Arc::new(Float64Array::from_iter(
384                        self.values()
385                            .as_any()
386                            .downcast_ref::<UInt8Array>()
387                            .ok_or(ArrowError::ParseError(
388                                "Fail to cast primitive array to UInt8Type".to_string(),
389                            ))?
390                            .into_iter()
391                            .map(|x| x.map(|y| y as f64)),
392                    )),
393                    self.nulls().cloned(),
394                )),
395                DataType::UInt32 => Ok(Self::new(
396                    Arc::new(arrow_schema::Field::new(
397                        field.name(),
398                        DataType::Float64,
399                        field.is_nullable(),
400                    )),
401                    *size,
402                    Arc::new(Float64Array::from_iter(
403                        self.values()
404                            .as_any()
405                            .downcast_ref::<UInt32Array>()
406                            .ok_or(ArrowError::ParseError(
407                                "Fail to cast primitive array to UInt32Type".to_string(),
408                            ))?
409                            .into_iter()
410                            .map(|x| x.map(|y| y as f64)),
411                    )),
412                    self.nulls().cloned(),
413                )),
414                data_type => Err(ArrowError::ParseError(format!(
415                    "Expect either floating type or integer got {:?}",
416                    data_type
417                ))),
418            },
419            data_type => Err(ArrowError::ParseError(format!(
420                "Expect either FixedSizeList got {:?}",
421                data_type
422            ))),
423        }
424    }
425}
426
427/// Force downcast of an [`Array`], such as an [`ArrayRef`], to
428/// [`FixedSizeListArray`], panic'ing on failure.
429pub fn as_fixed_size_list_array(arr: &dyn Array) -> &FixedSizeListArray {
430    arr.as_any().downcast_ref::<FixedSizeListArray>().unwrap()
431}
432
433pub trait FixedSizeBinaryArrayExt {
434    /// Create an [`FixedSizeBinaryArray`] from values and stride.
435    ///
436    /// ```
437    /// use arrow_array::{UInt8Array, FixedSizeBinaryArray};
438    /// use arrow_array::types::UInt8Type;
439    /// use lance_arrow::FixedSizeBinaryArrayExt;
440    ///
441    /// let int_values = UInt8Array::from_iter(0..10);
442    /// let fixed_size_list_arr = FixedSizeBinaryArray::try_new_from_values(&int_values, 2).unwrap();
443    /// assert_eq!(fixed_size_list_arr,
444    ///     FixedSizeBinaryArray::from(vec![
445    ///         Some(vec![0, 1].as_slice()),
446    ///         Some(vec![2, 3].as_slice()),
447    ///         Some(vec![4, 5].as_slice()),
448    ///         Some(vec![6, 7].as_slice()),
449    ///         Some(vec![8, 9].as_slice())
450    /// ]))
451    /// ```
452    fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result<FixedSizeBinaryArray>;
453}
454
455impl FixedSizeBinaryArrayExt for FixedSizeBinaryArray {
456    fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result<Self> {
457        let data_type = DataType::FixedSizeBinary(stride);
458        let data = ArrayDataBuilder::new(data_type)
459            .len(values.len() / stride as usize)
460            .add_buffer(values.into_data().buffers()[0].clone())
461            .build()?;
462        Ok(Self::from(data))
463    }
464}
465
466pub fn as_fixed_size_binary_array(arr: &dyn Array) -> &FixedSizeBinaryArray {
467    arr.as_any().downcast_ref::<FixedSizeBinaryArray>().unwrap()
468}
469
470pub fn iter_str_array(arr: &dyn Array) -> Box<dyn Iterator<Item = Option<&str>> + Send + '_> {
471    match arr.data_type() {
472        DataType::Utf8 => Box::new(arr.as_string::<i32>().iter()),
473        DataType::LargeUtf8 => Box::new(arr.as_string::<i64>().iter()),
474        DataType::Utf8View => Box::new(arr.as_string_view().iter()),
475        _ => panic!(
476            "Expecting Utf8, LargeUtf8 or Utf8View, found {:?}",
477            arr.data_type()
478        ),
479    }
480}
481
482pub fn iter_binary_array(
483    arr: &dyn Array,
484) -> Result<Box<dyn Iterator<Item = Option<&[u8]>> + Send + '_>> {
485    match arr.data_type() {
486        DataType::Binary => Ok(Box::new(arr.as_binary::<i32>().iter())),
487        DataType::LargeBinary => Ok(Box::new(arr.as_binary::<i64>().iter())),
488        DataType::BinaryView => Ok(Box::new(arr.as_binary_view().iter())),
489        DataType::FixedSizeBinary(_) => Ok(Box::new(arr.as_fixed_size_binary().iter())),
490        data_type => Err(ArrowError::InvalidArgumentError(format!(
491            "Expecting a binary type, found {data_type}"
492        ))),
493    }
494}
495
496/// Extends Arrow's [RecordBatch].
497pub trait RecordBatchExt {
498    /// Append a new column to this [`RecordBatch`] and returns a new RecordBatch.
499    ///
500    /// ```
501    /// use std::sync::Arc;
502    /// use arrow_array::{RecordBatch, Int32Array, StringArray};
503    /// use arrow_schema::{Schema, Field, DataType};
504    /// use lance_arrow::*;
505    ///
506    /// let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
507    /// let int_arr = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
508    /// let record_batch = RecordBatch::try_new(schema, vec![int_arr.clone()]).unwrap();
509    ///
510    /// let new_field = Field::new("s", DataType::Utf8, true);
511    /// let str_arr = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
512    /// let new_record_batch = record_batch.try_with_column(new_field, str_arr.clone()).unwrap();
513    ///
514    /// assert_eq!(
515    ///     new_record_batch,
516    ///     RecordBatch::try_new(
517    ///         Arc::new(Schema::new(
518    ///             vec![
519    ///                 Field::new("a", DataType::Int32, true),
520    ///                 Field::new("s", DataType::Utf8, true)
521    ///             ])
522    ///         ),
523    ///         vec![int_arr, str_arr],
524    ///     ).unwrap()
525    /// )
526    /// ```
527    fn try_with_column(&self, field: Field, arr: ArrayRef) -> Result<RecordBatch>;
528
529    /// Created a new RecordBatch with column at index.
530    fn try_with_column_at(&self, index: usize, field: Field, arr: ArrayRef) -> Result<RecordBatch>;
531
532    /// Creates a new [`RecordBatch`] from the provided  [`StructArray`].
533    ///
534    /// The fields on the [`StructArray`] need to match this [`RecordBatch`] schema
535    fn try_new_from_struct_array(&self, arr: StructArray) -> Result<RecordBatch>;
536
537    /// Merge with another [`RecordBatch`] and returns a new one.
538    ///
539    /// Fields are merged based on name.  First we iterate the left columns.  If a matching
540    /// name is found in the right then we merge the two columns.  If there is no match then
541    /// we add the left column to the output.
542    ///
543    /// To merge two columns we consider the type.  If both arrays are struct arrays we recurse.
544    /// Otherwise we use the left array.
545    ///
546    /// Afterwards we add all non-matching right columns to the output.
547    ///
548    /// Note: This method likely does not handle nested fields correctly and you may want to consider
549    /// using [`Self::merge_with_schema`] instead.
550    /// ```
551    /// use std::sync::Arc;
552    /// use arrow_array::*;
553    /// use arrow_schema::{Schema, Field, DataType};
554    /// use lance_arrow::*;
555    ///
556    /// let left_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
557    /// let int_arr = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
558    /// let left = RecordBatch::try_new(left_schema, vec![int_arr.clone()]).unwrap();
559    ///
560    /// let right_schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, true)]));
561    /// let str_arr = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
562    /// let right = RecordBatch::try_new(right_schema, vec![str_arr.clone()]).unwrap();
563    ///
564    /// let new_record_batch = left.merge(&right).unwrap();
565    ///
566    /// assert_eq!(
567    ///     new_record_batch,
568    ///     RecordBatch::try_new(
569    ///         Arc::new(Schema::new(
570    ///             vec![
571    ///                 Field::new("a", DataType::Int32, true),
572    ///                 Field::new("s", DataType::Utf8, true)
573    ///             ])
574    ///         ),
575    ///         vec![int_arr, str_arr],
576    ///     ).unwrap()
577    /// )
578    /// ```
579    ///
580    /// TODO: add merge nested fields support.
581    fn merge(&self, other: &RecordBatch) -> Result<RecordBatch>;
582
583    /// Create a batch by merging columns between two batches with a given schema.
584    ///
585    /// A reference schema is used to determine the proper ordering of nested fields.
586    ///
587    /// For each field in the reference schema we look for corresponding fields in
588    /// the left and right batches.  If a field is found in both batches we recursively merge
589    /// it.
590    ///
591    /// If a field is only in the left or right batch we take it as it is.
592    fn merge_with_schema(&self, other: &RecordBatch, schema: &Schema) -> Result<RecordBatch>;
593
594    /// Drop one column specified with the name and return the new [`RecordBatch`].
595    ///
596    /// If the named column does not exist, it returns a copy of this [`RecordBatch`].
597    fn drop_column(&self, name: &str) -> Result<RecordBatch>;
598
599    /// Replace a column (specified by name) and return the new [`RecordBatch`].
600    fn replace_column_by_name(&self, name: &str, column: Arc<dyn Array>) -> Result<RecordBatch>;
601
602    /// Replace a column schema (specified by name) and return the new [`RecordBatch`].
603    fn replace_column_schema_by_name(
604        &self,
605        name: &str,
606        new_data_type: DataType,
607        column: Arc<dyn Array>,
608    ) -> Result<RecordBatch>;
609
610    /// Rename a column at a given index.
611    fn rename_column(&self, index: usize, new_name: &str) -> Result<RecordBatch>;
612
613    /// Get (potentially nested) column by qualified name.
614    fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef>;
615
616    /// Project the schema over the [RecordBatch].
617    fn project_by_schema(&self, schema: &Schema) -> Result<RecordBatch>;
618
619    /// metadata of the schema.
620    fn metadata(&self) -> &HashMap<String, String>;
621
622    /// Add metadata to the schema.
623    fn add_metadata(&self, key: String, value: String) -> Result<RecordBatch> {
624        let mut metadata = self.metadata().clone();
625        metadata.insert(key, value);
626        self.with_metadata(metadata)
627    }
628
629    /// Replace the schema metadata with the provided one.
630    fn with_metadata(&self, metadata: HashMap<String, String>) -> Result<RecordBatch>;
631
632    /// Take selected rows from the [RecordBatch].
633    fn take(&self, indices: &UInt32Array) -> Result<RecordBatch>;
634
635    /// Create a new RecordBatch with compacted memory after slicing.
636    fn shrink_to_fit(&self) -> Result<RecordBatch>;
637
638    /// Helper method to sort the RecordBatch by a column
639    fn sort_by_column(&self, column: usize, options: Option<SortOptions>) -> Result<RecordBatch>;
640}
641
642impl RecordBatchExt for RecordBatch {
643    fn try_with_column(&self, field: Field, arr: ArrayRef) -> Result<Self> {
644        let new_schema = Arc::new(self.schema().as_ref().try_with_column(field)?);
645        let mut new_columns = self.columns().to_vec();
646        new_columns.push(arr);
647        Self::try_new(new_schema, new_columns)
648    }
649
650    fn try_with_column_at(&self, index: usize, field: Field, arr: ArrayRef) -> Result<Self> {
651        let new_schema = Arc::new(self.schema().as_ref().try_with_column_at(index, field)?);
652        let mut new_columns = self.columns().to_vec();
653        new_columns.insert(index, arr);
654        Self::try_new(new_schema, new_columns)
655    }
656
657    fn try_new_from_struct_array(&self, arr: StructArray) -> Result<Self> {
658        let schema = Arc::new(Schema::new_with_metadata(
659            arr.fields().to_vec(),
660            self.schema().metadata.clone(),
661        ));
662        let batch = Self::from(arr);
663        batch.with_schema(schema)
664    }
665
666    fn merge(&self, other: &Self) -> Result<Self> {
667        if self.num_rows() != other.num_rows() {
668            return Err(ArrowError::InvalidArgumentError(format!(
669                "Attempt to merge two RecordBatch with different sizes: {} != {}",
670                self.num_rows(),
671                other.num_rows()
672            )));
673        }
674        let left_struct_array: StructArray = self.clone().into();
675        let right_struct_array: StructArray = other.clone().into();
676        self.try_new_from_struct_array(merge(&left_struct_array, &right_struct_array))
677    }
678
679    fn merge_with_schema(&self, other: &RecordBatch, schema: &Schema) -> Result<RecordBatch> {
680        if self.num_rows() != other.num_rows() {
681            return Err(ArrowError::InvalidArgumentError(format!(
682                "Attempt to merge two RecordBatch with different sizes: {} != {}",
683                self.num_rows(),
684                other.num_rows()
685            )));
686        }
687        let left_struct_array: StructArray = self.clone().into();
688        let right_struct_array: StructArray = other.clone().into();
689        self.try_new_from_struct_array(merge_with_schema(
690            &left_struct_array,
691            &right_struct_array,
692            schema.fields(),
693        ))
694    }
695
696    fn drop_column(&self, name: &str) -> Result<Self> {
697        let mut fields = vec![];
698        let mut columns = vec![];
699        for i in 0..self.schema().fields.len() {
700            if self.schema().field(i).name() != name {
701                fields.push(self.schema().field(i).clone());
702                columns.push(self.column(i).clone());
703            }
704        }
705        Self::try_new(
706            Arc::new(Schema::new_with_metadata(
707                fields,
708                self.schema().metadata().clone(),
709            )),
710            columns,
711        )
712    }
713
714    fn rename_column(&self, index: usize, new_name: &str) -> Result<RecordBatch> {
715        let mut fields = self.schema().fields().to_vec();
716        if index >= fields.len() {
717            return Err(ArrowError::InvalidArgumentError(format!(
718                "Index out of bounds: {}",
719                index
720            )));
721        }
722        fields[index] = Arc::new(Field::new(
723            new_name,
724            fields[index].data_type().clone(),
725            fields[index].is_nullable(),
726        ));
727        Self::try_new(
728            Arc::new(Schema::new_with_metadata(
729                fields,
730                self.schema().metadata().clone(),
731            )),
732            self.columns().to_vec(),
733        )
734    }
735
736    fn replace_column_by_name(&self, name: &str, column: Arc<dyn Array>) -> Result<RecordBatch> {
737        let mut columns = self.columns().to_vec();
738        let field_i = self
739            .schema()
740            .fields()
741            .iter()
742            .position(|f| f.name() == name)
743            .ok_or_else(|| ArrowError::SchemaError(format!("Field {} does not exist", name)))?;
744        columns[field_i] = column;
745        Self::try_new(self.schema(), columns)
746    }
747
748    fn replace_column_schema_by_name(
749        &self,
750        name: &str,
751        new_data_type: DataType,
752        column: Arc<dyn Array>,
753    ) -> Result<RecordBatch> {
754        let fields = self
755            .schema()
756            .fields()
757            .iter()
758            .map(|x| {
759                if x.name() != name {
760                    x.clone()
761                } else {
762                    let new_field = Field::new(name, new_data_type.clone(), x.is_nullable());
763                    Arc::new(new_field)
764                }
765            })
766            .collect::<Vec<_>>();
767        let schema = Schema::new_with_metadata(fields, self.schema().metadata.clone());
768        let mut columns = self.columns().to_vec();
769        let field_i = self
770            .schema()
771            .fields()
772            .iter()
773            .position(|f| f.name() == name)
774            .ok_or_else(|| ArrowError::SchemaError(format!("Field {} does not exist", name)))?;
775        columns[field_i] = column;
776        Self::try_new(Arc::new(schema), columns)
777    }
778
779    fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef> {
780        let split = name.split('.').collect::<Vec<_>>();
781        if split.is_empty() {
782            return None;
783        }
784
785        self.column_by_name(split[0])
786            .and_then(|arr| get_sub_array(arr, &split[1..]))
787    }
788
789    fn project_by_schema(&self, schema: &Schema) -> Result<Self> {
790        let struct_array: StructArray = self.clone().into();
791        self.try_new_from_struct_array(project(&struct_array, schema.fields())?)
792    }
793
794    fn metadata(&self) -> &HashMap<String, String> {
795        self.schema_ref().metadata()
796    }
797
798    fn with_metadata(&self, metadata: HashMap<String, String>) -> Result<RecordBatch> {
799        let mut schema = self.schema_ref().as_ref().clone();
800        schema.metadata = metadata;
801        Self::try_new(schema.into(), self.columns().into())
802    }
803
804    fn take(&self, indices: &UInt32Array) -> Result<Self> {
805        let struct_array: StructArray = self.clone().into();
806        let taken = take(&struct_array, indices, None)?;
807        self.try_new_from_struct_array(taken.as_struct().clone())
808    }
809
810    fn shrink_to_fit(&self) -> Result<Self> {
811        // Deep copy the sliced record batch, instead of whole batch
812        crate::deepcopy::deep_copy_batch_sliced(self)
813    }
814
815    fn sort_by_column(&self, column: usize, options: Option<SortOptions>) -> Result<Self> {
816        if column >= self.num_columns() {
817            return Err(ArrowError::InvalidArgumentError(format!(
818                "Column index out of bounds: {}",
819                column
820            )));
821        }
822        let column = self.column(column);
823        let sorted = arrow_ord::sort::sort_to_indices(column, options, None)?;
824        self.take(&sorted)
825    }
826}
827
828/// Recursively projects an array to match the target field's structure.
829/// This handles reordering fields inside nested List<Struct> types.
830fn project_array(array: &ArrayRef, target_field: &Field) -> Result<ArrayRef> {
831    match target_field.data_type() {
832        DataType::Struct(subfields) => {
833            let struct_arr = array.as_struct();
834            let projected = project(struct_arr, subfields)?;
835            Ok(Arc::new(projected))
836        }
837        DataType::List(inner_field) => {
838            let list_arr: &ListArray = array.as_list();
839            let projected_values = project_array(list_arr.values(), inner_field.as_ref())?;
840            Ok(Arc::new(ListArray::new(
841                inner_field.clone(),
842                list_arr.offsets().clone(),
843                projected_values,
844                list_arr.nulls().cloned(),
845            )))
846        }
847        DataType::LargeList(inner_field) => {
848            let list_arr: &LargeListArray = array.as_list();
849            let projected_values = project_array(list_arr.values(), inner_field.as_ref())?;
850            Ok(Arc::new(LargeListArray::new(
851                inner_field.clone(),
852                list_arr.offsets().clone(),
853                projected_values,
854                list_arr.nulls().cloned(),
855            )))
856        }
857        DataType::FixedSizeList(inner_field, size) => {
858            let list_arr = array.as_fixed_size_list();
859            let projected_values = project_array(list_arr.values(), inner_field.as_ref())?;
860            Ok(Arc::new(FixedSizeListArray::new(
861                inner_field.clone(),
862                *size,
863                projected_values,
864                list_arr.nulls().cloned(),
865            )))
866        }
867        // A nullable entries field fails MapArray::try_new unconditionally,
868        // so a (schema-invalid) map declared that way keeps the clone
869        // fallthrough it always had rather than gaining a new error.
870        DataType::Map(entries_field, sorted) if !entries_field.is_nullable() => {
871            let map_arr = array.as_map();
872            let DataType::Struct(entry_fields) = entries_field.data_type() else {
873                return Err(ArrowError::SchemaError(format!(
874                    "Map entries field must be a struct, got {}",
875                    entries_field.data_type()
876                )));
877            };
878            let projected_entries = project(map_arr.entries(), entry_fields)?;
879            // try_new re-checks the entries invariants (a non-null entries
880            // struct, two entry columns, offset bounds); null keys are ruled
881            // out one level down, by the struct rebuild against the
882            // non-nullable key field.
883            Ok(Arc::new(MapArray::try_new(
884                entries_field.clone(),
885                map_arr.offsets().clone(),
886                projected_entries,
887                map_arr.nulls().cloned(),
888                *sorted,
889            )?))
890        }
891        _ => Ok(array.clone()),
892    }
893}
894
895fn project(struct_array: &StructArray, fields: &Fields) -> Result<StructArray> {
896    if fields.is_empty() {
897        return Ok(StructArray::new_empty_fields(
898            struct_array.len(),
899            struct_array.nulls().cloned(),
900        ));
901    }
902    let mut columns: Vec<ArrayRef> = vec![];
903    for field in fields.iter() {
904        if let Some(col) = struct_array.column_by_name(field.name()) {
905            let projected = project_array(col, field.as_ref())?;
906            columns.push(projected);
907        } else {
908            return Err(ArrowError::SchemaError(format!(
909                "field {} does not exist in the RecordBatch",
910                field.name()
911            )));
912        }
913    }
914    // Preserve the struct's validity when projecting
915    StructArray::try_new(fields.clone(), columns, struct_array.nulls().cloned())
916}
917
918fn lists_have_same_offsets_helper<T: OffsetSizeTrait>(left: &dyn Array, right: &dyn Array) -> bool {
919    let left_list: &GenericListArray<T> = left.as_list();
920    let right_list: &GenericListArray<T> = right.as_list();
921    left_list.offsets().inner() == right_list.offsets().inner()
922}
923
924fn merge_list_structs_helper<T: OffsetSizeTrait>(
925    left: &dyn Array,
926    right: &dyn Array,
927    items_field_name: impl Into<String>,
928    items_nullable: bool,
929) -> Arc<dyn Array> {
930    let left_list: &GenericListArray<T> = left.as_list();
931    let right_list: &GenericListArray<T> = right.as_list();
932    let left_struct = left_list.values();
933    let right_struct = right_list.values();
934    let left_struct_arr = left_struct.as_struct();
935    let right_struct_arr = right_struct.as_struct();
936    let merged_items = Arc::new(merge(left_struct_arr, right_struct_arr));
937    let items_field = Arc::new(Field::new(
938        items_field_name,
939        merged_items.data_type().clone(),
940        items_nullable,
941    ));
942    Arc::new(GenericListArray::<T>::new(
943        items_field,
944        left_list.offsets().clone(),
945        merged_items,
946        left_list.nulls().cloned(),
947    ))
948}
949
950fn merge_list_struct_null_helper<T: OffsetSizeTrait>(
951    left: &dyn Array,
952    right: &dyn Array,
953    not_null: &dyn Array,
954    items_field_name: impl Into<String>,
955) -> Arc<dyn Array> {
956    let left_list: &GenericListArray<T> = left.as_list::<T>();
957    let not_null_list = not_null.as_list::<T>();
958    let right_list = right.as_list::<T>();
959
960    let left_struct = left_list.values().as_struct();
961    let not_null_struct: &StructArray = not_null_list.values().as_struct();
962    let right_struct = right_list.values().as_struct();
963
964    let values_len = not_null_list.values().len();
965    let mut merged_fields =
966        Vec::with_capacity(not_null_struct.num_columns() + right_struct.num_columns());
967    let mut merged_columns =
968        Vec::with_capacity(not_null_struct.num_columns() + right_struct.num_columns());
969
970    for (_, field) in left_struct.columns().iter().zip(left_struct.fields()) {
971        merged_fields.push(field.clone());
972        if let Some(val) = not_null_struct.column_by_name(field.name()) {
973            merged_columns.push(val.clone());
974        } else {
975            merged_columns.push(new_null_array(field.data_type(), values_len))
976        }
977    }
978    for (_, field) in right_struct
979        .columns()
980        .iter()
981        .zip(right_struct.fields())
982        .filter(|(_, field)| left_struct.column_by_name(field.name()).is_none())
983    {
984        merged_fields.push(field.clone());
985        if let Some(val) = not_null_struct.column_by_name(field.name()) {
986            merged_columns.push(val.clone());
987        } else {
988            merged_columns.push(new_null_array(field.data_type(), values_len));
989        }
990    }
991
992    let merged_struct = Arc::new(StructArray::new(
993        Fields::from(merged_fields),
994        merged_columns,
995        not_null_struct.nulls().cloned(),
996    ));
997    let items_field = Arc::new(Field::new(
998        items_field_name,
999        merged_struct.data_type().clone(),
1000        true,
1001    ));
1002    Arc::new(GenericListArray::<T>::new(
1003        items_field,
1004        not_null_list.offsets().clone(),
1005        merged_struct,
1006        not_null_list.nulls().cloned(),
1007    ))
1008}
1009
1010fn merge_list_struct_null(
1011    left: &dyn Array,
1012    right: &dyn Array,
1013    not_null: &dyn Array,
1014) -> Arc<dyn Array> {
1015    match left.data_type() {
1016        DataType::List(left_field) => {
1017            merge_list_struct_null_helper::<i32>(left, right, not_null, left_field.name())
1018        }
1019        DataType::LargeList(left_field) => {
1020            merge_list_struct_null_helper::<i64>(left, right, not_null, left_field.name())
1021        }
1022        _ => unreachable!(),
1023    }
1024}
1025
1026fn merge_list_struct(left: &dyn Array, right: &dyn Array) -> Arc<dyn Array> {
1027    // Merging fields into a list<struct<...>> is tricky and can only succeed
1028    // in two ways.  First, if both lists have the same offsets.  Second, if
1029    // one of the lists is all-null
1030    if left.null_count() == left.len() {
1031        return merge_list_struct_null(left, right, right);
1032    } else if right.null_count() == right.len() {
1033        return merge_list_struct_null(left, right, left);
1034    }
1035    match (left.data_type(), right.data_type()) {
1036        (DataType::List(left_field), DataType::List(_)) => {
1037            if !lists_have_same_offsets_helper::<i32>(left, right) {
1038                panic!("Attempt to merge list struct arrays which do not have same offsets");
1039            }
1040            merge_list_structs_helper::<i32>(
1041                left,
1042                right,
1043                left_field.name(),
1044                left_field.is_nullable(),
1045            )
1046        }
1047        (DataType::LargeList(left_field), DataType::LargeList(_)) => {
1048            if !lists_have_same_offsets_helper::<i64>(left, right) {
1049                panic!("Attempt to merge list struct arrays which do not have same offsets");
1050            }
1051            merge_list_structs_helper::<i64>(
1052                left,
1053                right,
1054                left_field.name(),
1055                left_field.is_nullable(),
1056            )
1057        }
1058        _ => unreachable!(),
1059    }
1060}
1061
1062/// Helper function to merge validity buffers from two struct arrays.
1063///
1064/// A row is valid if it is valid in either input.
1065/// An absent validity buffer means all rows are valid, an all-null buffer acts as the identity for this merge.
1066fn merge_struct_validity(
1067    left_validity: Option<&arrow_buffer::NullBuffer>,
1068    right_validity: Option<&arrow_buffer::NullBuffer>,
1069) -> Option<arrow_buffer::NullBuffer> {
1070    match (left_validity, right_validity) {
1071        // Fast paths: no computation needed
1072        (None, _) | (_, None) => None,
1073        (Some(left), Some(right)) => {
1074            if left.null_count() == 0 || right.null_count() == 0 {
1075                return None;
1076            }
1077            if left.null_count() == left.len() {
1078                return Some(right.clone());
1079            }
1080            if right.null_count() == right.len() {
1081                return Some(left.clone());
1082            }
1083
1084            let left_buffer = left.inner();
1085            let right_buffer = right.inner();
1086
1087            // Perform bitwise OR directly on BooleanBuffers
1088            // This preserves the correct semantics: 1 = valid, 0 = null
1089            let merged_buffer = left_buffer | right_buffer;
1090
1091            Some(arrow_buffer::NullBuffer::from(merged_buffer))
1092        }
1093    }
1094}
1095
1096fn merge_list_child_values(
1097    child_field: &Field,
1098    left_values: ArrayRef,
1099    right_values: ArrayRef,
1100) -> ArrayRef {
1101    match child_field.data_type() {
1102        DataType::Struct(child_fields) => Arc::new(merge_with_schema(
1103            left_values.as_struct(),
1104            right_values.as_struct(),
1105            child_fields,
1106        )) as ArrayRef,
1107        DataType::List(grandchild) => {
1108            let left_list = left_values
1109                .as_any()
1110                .downcast_ref::<ListArray>()
1111                .expect("left list values should be ListArray");
1112            let right_list = right_values
1113                .as_any()
1114                .downcast_ref::<ListArray>()
1115                .expect("right list values should be ListArray");
1116            let merged_values = merge_list_child_values(
1117                grandchild.as_ref(),
1118                left_list.values().clone(),
1119                right_list.values().clone(),
1120            );
1121            let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1122            Arc::new(ListArray::new(
1123                grandchild.clone(),
1124                left_list.offsets().clone(),
1125                merged_values,
1126                merged_validity,
1127            )) as ArrayRef
1128        }
1129        DataType::LargeList(grandchild) => {
1130            let left_list = left_values
1131                .as_any()
1132                .downcast_ref::<LargeListArray>()
1133                .expect("left list values should be LargeListArray");
1134            let right_list = right_values
1135                .as_any()
1136                .downcast_ref::<LargeListArray>()
1137                .expect("right list values should be LargeListArray");
1138            let merged_values = merge_list_child_values(
1139                grandchild.as_ref(),
1140                left_list.values().clone(),
1141                right_list.values().clone(),
1142            );
1143            let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1144            Arc::new(LargeListArray::new(
1145                grandchild.clone(),
1146                left_list.offsets().clone(),
1147                merged_values,
1148                merged_validity,
1149            )) as ArrayRef
1150        }
1151        DataType::FixedSizeList(grandchild, list_size) => {
1152            let left_list = left_values
1153                .as_any()
1154                .downcast_ref::<FixedSizeListArray>()
1155                .expect("left list values should be FixedSizeListArray");
1156            let right_list = right_values
1157                .as_any()
1158                .downcast_ref::<FixedSizeListArray>()
1159                .expect("right list values should be FixedSizeListArray");
1160            let merged_values = merge_list_child_values(
1161                grandchild.as_ref(),
1162                left_list.values().clone(),
1163                right_list.values().clone(),
1164            );
1165            let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1166            Arc::new(FixedSizeListArray::new(
1167                grandchild.clone(),
1168                *list_size,
1169                merged_values,
1170                merged_validity,
1171            )) as ArrayRef
1172        }
1173        _ => left_values.clone(),
1174    }
1175}
1176
1177// Helper function to adjust child array validity based on parent struct validity
1178// When parent struct is null, propagates null to child array
1179// Optimized with fast paths and SIMD operations
1180fn adjust_child_validity(
1181    child: &ArrayRef,
1182    parent_validity: Option<&arrow_buffer::NullBuffer>,
1183) -> ArrayRef {
1184    // Fast path: no parent validity means no adjustment needed
1185    let parent_validity = match parent_validity {
1186        None => return child.clone(),
1187        Some(p) if p.null_count() == 0 => return child.clone(), // No nulls to propagate
1188        Some(p) => p,
1189    };
1190
1191    // Fast path: DataType::Null arrays are always entirely null by definition and cannot
1192    // carry an explicit null bitmap (Arrow rejects it). No adjustment is needed.
1193    if child.data_type() == &DataType::Null {
1194        return child.clone();
1195    }
1196
1197    let child_validity = child.nulls();
1198
1199    // Compute the new validity: child_validity AND parent_validity
1200    let new_validity = match child_validity {
1201        None => {
1202            // Fast path: child has no existing validity, just use parent's
1203            parent_validity.clone()
1204        }
1205        Some(child_nulls) => {
1206            let child_buffer = child_nulls.inner();
1207            let parent_buffer = parent_validity.inner();
1208
1209            // Perform bitwise AND directly on BooleanBuffers
1210            // This preserves the correct semantics: 1 = valid, 0 = null
1211            let merged_buffer = child_buffer & parent_buffer;
1212
1213            arrow_buffer::NullBuffer::from(merged_buffer)
1214        }
1215    };
1216
1217    // Create new array with adjusted validity
1218    arrow_array::make_array(
1219        arrow_data::ArrayData::try_new(
1220            child.data_type().clone(),
1221            child.len(),
1222            Some(new_validity.into_inner().into_inner()),
1223            child.offset(),
1224            child.to_data().buffers().to_vec(),
1225            child.to_data().child_data().to_vec(),
1226        )
1227        .unwrap(),
1228    )
1229}
1230
1231fn merge(left_struct_array: &StructArray, right_struct_array: &StructArray) -> StructArray {
1232    let mut fields: Vec<Field> = vec![];
1233    let mut columns: Vec<ArrayRef> = vec![];
1234    let right_fields = right_struct_array.fields();
1235    let right_columns = right_struct_array.columns();
1236
1237    // Get the validity buffers from both structs
1238    let left_validity = left_struct_array.nulls();
1239    let right_validity = right_struct_array.nulls();
1240
1241    // Compute merged validity
1242    let merged_validity = merge_struct_validity(left_validity, right_validity);
1243
1244    // iterate through the fields on the left hand side
1245    for (left_field, left_column) in left_struct_array
1246        .fields()
1247        .iter()
1248        .zip(left_struct_array.columns().iter())
1249    {
1250        match right_fields
1251            .iter()
1252            .position(|f| f.name() == left_field.name())
1253        {
1254            // if the field exists on the right hand side, merge them recursively if appropriate
1255            Some(right_index) => {
1256                let right_field = right_fields.get(right_index).unwrap();
1257                let right_column = right_columns.get(right_index).unwrap();
1258                // if both fields are struct, merge them recursively
1259                match (left_field.data_type(), right_field.data_type()) {
1260                    (DataType::Struct(_), DataType::Struct(_)) => {
1261                        let left_sub_array = left_column.as_struct();
1262                        let right_sub_array = right_column.as_struct();
1263                        let merged_sub_array = merge(left_sub_array, right_sub_array);
1264                        fields.push(Field::new(
1265                            left_field.name(),
1266                            merged_sub_array.data_type().clone(),
1267                            left_field.is_nullable(),
1268                        ));
1269                        columns.push(Arc::new(merged_sub_array) as ArrayRef);
1270                    }
1271                    (DataType::List(left_list), DataType::List(right_list))
1272                        if left_list.data_type().is_struct()
1273                            && right_list.data_type().is_struct() =>
1274                    {
1275                        // Identical inner types: nothing to merge, keep the left column.
1276                        if left_list.data_type() == right_list.data_type() {
1277                            fields.push(left_field.as_ref().clone());
1278                            columns.push(left_column.clone());
1279                        } else {
1280                            // The struct fields differ, so merge them structurally. merge_list_struct
1281                            // only succeeds when both lists share offsets or one side is all-null;
1282                            // it panics otherwise.
1283                            let merged_sub_array = merge_list_struct(&left_column, &right_column);
1284
1285                            fields.push(Field::new(
1286                                left_field.name(),
1287                                merged_sub_array.data_type().clone(),
1288                                left_field.is_nullable(),
1289                            ));
1290                            columns.push(merged_sub_array);
1291                        }
1292                    }
1293                    // otherwise, just use the field on the left hand side
1294                    _ => {
1295                        // TODO handle list-of-struct and other types
1296                        fields.push(left_field.as_ref().clone());
1297                        // Adjust the column validity: if left struct was null, propagate to child
1298                        let adjusted_column = adjust_child_validity(left_column, left_validity);
1299                        columns.push(adjusted_column);
1300                    }
1301                }
1302            }
1303            None => {
1304                fields.push(left_field.as_ref().clone());
1305                // Adjust the column validity: if left struct was null, propagate to child
1306                let adjusted_column = adjust_child_validity(left_column, left_validity);
1307                columns.push(adjusted_column);
1308            }
1309        }
1310    }
1311
1312    // now iterate through the fields on the right hand side
1313    right_fields
1314        .iter()
1315        .zip(right_columns.iter())
1316        .for_each(|(field, column)| {
1317            // add new columns on the right
1318            if !left_struct_array
1319                .fields()
1320                .iter()
1321                .any(|f| f.name() == field.name())
1322            {
1323                fields.push(field.as_ref().clone());
1324                // This field doesn't exist on the left
1325                // We use the right's column but need to adjust for struct validity
1326                let adjusted_column = adjust_child_validity(column, right_validity);
1327                columns.push(adjusted_column);
1328            }
1329        });
1330
1331    StructArray::try_new(Fields::from(fields), columns, merged_validity).unwrap()
1332}
1333
1334fn merge_with_schema(
1335    left_struct_array: &StructArray,
1336    right_struct_array: &StructArray,
1337    fields: &Fields,
1338) -> StructArray {
1339    // Helper function that returns true if both types are struct or both are non-struct
1340    fn same_type_kind(left: &DataType, right: &DataType) -> bool {
1341        match (left, right) {
1342            (DataType::Struct(_), DataType::Struct(_)) => true,
1343            (DataType::Struct(_), _) => false,
1344            (_, DataType::Struct(_)) => false,
1345            _ => true,
1346        }
1347    }
1348
1349    let mut output_fields: Vec<Field> = Vec::with_capacity(fields.len());
1350    let mut columns: Vec<ArrayRef> = Vec::with_capacity(fields.len());
1351
1352    let left_fields = left_struct_array.fields();
1353    let left_columns = left_struct_array.columns();
1354    let right_fields = right_struct_array.fields();
1355    let right_columns = right_struct_array.columns();
1356
1357    // Get the validity buffers from both structs
1358    let left_validity = left_struct_array.nulls();
1359    let right_validity = right_struct_array.nulls();
1360
1361    // Compute merged validity
1362    let merged_validity = merge_struct_validity(left_validity, right_validity);
1363
1364    for field in fields {
1365        let left_match_idx = left_fields.iter().position(|f| {
1366            f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1367        });
1368        let right_match_idx = right_fields.iter().position(|f| {
1369            f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1370        });
1371
1372        match (left_match_idx, right_match_idx) {
1373            (None, Some(right_idx)) => {
1374                output_fields.push(right_fields[right_idx].as_ref().clone());
1375                // Adjust validity if the right struct was null
1376                let adjusted_column =
1377                    adjust_child_validity(&right_columns[right_idx], right_validity);
1378                columns.push(adjusted_column);
1379            }
1380            (Some(left_idx), None) => {
1381                output_fields.push(left_fields[left_idx].as_ref().clone());
1382                // Adjust validity if the left struct was null
1383                let adjusted_column = adjust_child_validity(&left_columns[left_idx], left_validity);
1384                columns.push(adjusted_column);
1385            }
1386            (Some(left_idx), Some(right_idx)) => {
1387                match field.data_type() {
1388                    DataType::Struct(child_fields) => {
1389                        let left_sub_array = left_columns[left_idx].as_struct();
1390                        let right_sub_array = right_columns[right_idx].as_struct();
1391                        let merged_sub_array =
1392                            merge_with_schema(left_sub_array, right_sub_array, child_fields);
1393                        output_fields.push(Field::new(
1394                            field.name(),
1395                            merged_sub_array.data_type().clone(),
1396                            field.is_nullable(),
1397                        ));
1398                        columns.push(Arc::new(merged_sub_array) as ArrayRef);
1399                    }
1400                    DataType::List(child_field) => {
1401                        let left_list = left_columns[left_idx]
1402                            .as_any()
1403                            .downcast_ref::<ListArray>()
1404                            .unwrap();
1405                        let right_list = right_columns[right_idx]
1406                            .as_any()
1407                            .downcast_ref::<ListArray>()
1408                            .unwrap();
1409                        let merged_values = merge_list_child_values(
1410                            child_field.as_ref(),
1411                            left_list.trimmed_values(),
1412                            right_list.trimmed_values(),
1413                        );
1414                        let merged_validity =
1415                            merge_struct_validity(left_list.nulls(), right_list.nulls());
1416                        // `trimmed_values` starts at the first used value, so offsets
1417                        // must be shifted to match or `ListArray::new` panics when the
1418                        // input list was sliced (e.g. from a filtered batch).
1419                        let merged_list = ListArray::new(
1420                            child_field.clone(),
1421                            left_list.trimmed_offsets(),
1422                            merged_values,
1423                            merged_validity,
1424                        );
1425                        output_fields.push(field.as_ref().clone());
1426                        columns.push(Arc::new(merged_list) as ArrayRef);
1427                    }
1428                    DataType::LargeList(child_field) => {
1429                        let left_list = left_columns[left_idx]
1430                            .as_any()
1431                            .downcast_ref::<LargeListArray>()
1432                            .unwrap();
1433                        let right_list = right_columns[right_idx]
1434                            .as_any()
1435                            .downcast_ref::<LargeListArray>()
1436                            .unwrap();
1437                        let merged_values = merge_list_child_values(
1438                            child_field.as_ref(),
1439                            left_list.trimmed_values(),
1440                            right_list.trimmed_values(),
1441                        );
1442                        let merged_validity =
1443                            merge_struct_validity(left_list.nulls(), right_list.nulls());
1444                        let merged_list = LargeListArray::new(
1445                            child_field.clone(),
1446                            left_list.trimmed_offsets(),
1447                            merged_values,
1448                            merged_validity,
1449                        );
1450                        output_fields.push(field.as_ref().clone());
1451                        columns.push(Arc::new(merged_list) as ArrayRef);
1452                    }
1453                    DataType::FixedSizeList(child_field, list_size) => {
1454                        let left_list = left_columns[left_idx]
1455                            .as_any()
1456                            .downcast_ref::<FixedSizeListArray>()
1457                            .unwrap();
1458                        let right_list = right_columns[right_idx]
1459                            .as_any()
1460                            .downcast_ref::<FixedSizeListArray>()
1461                            .unwrap();
1462                        let merged_values = merge_list_child_values(
1463                            child_field.as_ref(),
1464                            left_list.values().clone(),
1465                            right_list.values().clone(),
1466                        );
1467                        let merged_validity =
1468                            merge_struct_validity(left_list.nulls(), right_list.nulls());
1469                        let merged_list = FixedSizeListArray::new(
1470                            child_field.clone(),
1471                            *list_size,
1472                            merged_values,
1473                            merged_validity,
1474                        );
1475                        output_fields.push(field.as_ref().clone());
1476                        columns.push(Arc::new(merged_list) as ArrayRef);
1477                    }
1478                    _ => {
1479                        output_fields.push(left_fields[left_idx].as_ref().clone());
1480                        // For fields that exist in both, use left but adjust validity
1481                        let adjusted_column =
1482                            adjust_child_validity(&left_columns[left_idx], left_validity);
1483                        columns.push(adjusted_column);
1484                    }
1485                }
1486            }
1487            (None, None) => {
1488                // The field will not be included in the output
1489            }
1490        }
1491    }
1492
1493    StructArray::try_new(Fields::from(output_fields), columns, merged_validity).unwrap()
1494}
1495
1496fn get_sub_array<'a>(array: &'a ArrayRef, components: &[&str]) -> Option<&'a ArrayRef> {
1497    if components.is_empty() {
1498        return Some(array);
1499    }
1500    if !matches!(array.data_type(), DataType::Struct(_)) {
1501        return None;
1502    }
1503    let struct_arr = array.as_struct();
1504    struct_arr
1505        .column_by_name(components[0])
1506        .and_then(|arr| get_sub_array(arr, &components[1..]))
1507}
1508
1509/// Interleave multiple RecordBatches into a single RecordBatch.
1510///
1511/// Behaves like [`arrow_select::interleave::interleave`], but for RecordBatches.
1512pub fn interleave_batches(
1513    batches: &[RecordBatch],
1514    indices: &[(usize, usize)],
1515) -> Result<RecordBatch> {
1516    let first_batch = batches.first().ok_or_else(|| {
1517        ArrowError::InvalidArgumentError("Cannot interleave zero RecordBatches".to_string())
1518    })?;
1519    let schema = first_batch.schema();
1520    let num_columns = first_batch.num_columns();
1521    let mut columns = Vec::with_capacity(num_columns);
1522    let mut chunks = Vec::with_capacity(batches.len());
1523
1524    for i in 0..num_columns {
1525        for batch in batches {
1526            chunks.push(batch.column(i).as_ref());
1527        }
1528        let new_column = interleave(&chunks, indices)?;
1529        columns.push(new_column);
1530        chunks.clear();
1531    }
1532
1533    RecordBatch::try_new(schema, columns)
1534}
1535
1536pub trait BufferExt {
1537    /// Create an `arrow_buffer::Buffer`` from a `bytes::Bytes` object
1538    ///
1539    /// The alignment must be specified (as `bytes_per_value`) since we want to make
1540    /// sure we can safely reinterpret the buffer.
1541    ///
1542    /// If the buffer is properly aligned this will be zero-copy.  If not, a copy
1543    /// will be made and an owned buffer returned.
1544    ///
1545    /// If `bytes_per_value` is not a power of two, then we assume the buffer is
1546    /// never going to be reinterpreted into another type and we can safely
1547    /// ignore the alignment.
1548    ///
1549    /// Yes, the method name is odd.  It's because there is already a `from_bytes`
1550    /// which converts from `arrow_buffer::bytes::Bytes` (not `bytes::Bytes`)
1551    fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self;
1552
1553    /// Allocates a new properly aligned arrow buffer and copies `bytes` into it
1554    ///
1555    /// `size_bytes` can be larger than `bytes` and, if so, the trailing bytes will
1556    /// be zeroed out.
1557    ///
1558    /// # Panics
1559    ///
1560    /// Panics if `size_bytes` is less than `bytes.len()`
1561    fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self;
1562}
1563
1564fn is_pwr_two(n: u64) -> bool {
1565    n & (n - 1) == 0
1566}
1567
1568impl BufferExt for arrow_buffer::Buffer {
1569    fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self {
1570        if is_pwr_two(bytes_per_value) && bytes.as_ptr().align_offset(bytes_per_value as usize) != 0
1571        {
1572            // The original buffer is not aligned, cannot zero-copy
1573            let size_bytes = bytes.len();
1574            Self::copy_bytes_bytes(bytes, size_bytes)
1575        } else {
1576            // The original buffer is aligned, can zero-copy
1577            // SAFETY: the alignment is correct we can make this conversion
1578            unsafe {
1579                Self::from_custom_allocation(
1580                    NonNull::new(bytes.as_ptr() as _).expect("should be a valid pointer"),
1581                    bytes.len(),
1582                    Arc::new(bytes),
1583                )
1584            }
1585        }
1586    }
1587
1588    fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self {
1589        assert!(size_bytes >= bytes.len());
1590        let mut buf = MutableBuffer::with_capacity(size_bytes);
1591        let to_fill = size_bytes - bytes.len();
1592        buf.extend(bytes);
1593        buf.extend(std::iter::repeat_n(0_u8, to_fill));
1594
1595        // FIX for issue #4512: Shrink buffer to actual size before converting to immutable
1596        // This reduces memory overhead from capacity over-allocation
1597        buf.shrink_to_fit();
1598
1599        Self::from(buf)
1600    }
1601}
1602
1603#[cfg(test)]
1604mod tests {
1605    use super::*;
1606    use arrow_array::{
1607        BinaryArray, BinaryViewArray, FixedSizeBinaryArray, Float32Array, Int32Array,
1608        LargeBinaryArray, ListArray, NullArray, StringArray, StructArray, new_empty_array,
1609        new_null_array,
1610    };
1611    use arrow_buffer::OffsetBuffer;
1612
1613    #[test]
1614    fn test_convert_to_floating_point_preserves_inner_nulls() {
1615        // A FixedSizeList<Int8> with a null inner element must convert to a
1616        // FixedSizeList<Float32> with the null kept in place. Dropping it would
1617        // shorten the values array and shift every later element (and, when the
1618        // remaining count is not a multiple of the list size, panic).
1619        let values = Int8Array::from(vec![Some(1), None, Some(3), Some(4)]);
1620        let fsl = FixedSizeListArray::new(
1621            Arc::new(Field::new("item", DataType::Int8, true)),
1622            2,
1623            Arc::new(values),
1624            None,
1625        );
1626
1627        let converted = fsl.convert_to_floating_point().unwrap();
1628
1629        assert_eq!(converted.len(), 2);
1630        let conv_values = converted
1631            .values()
1632            .as_any()
1633            .downcast_ref::<Float32Array>()
1634            .unwrap();
1635        assert_eq!(conv_values.len(), 4);
1636        assert_eq!(conv_values.value(0), 1.0);
1637        assert!(conv_values.is_null(1));
1638        assert_eq!(conv_values.value(2), 3.0);
1639        assert_eq!(conv_values.value(3), 4.0);
1640    }
1641
1642    #[test]
1643    fn test_convert_to_floating_point_preserves_inner_nulls_f64_arm() {
1644        // The Float64-producing arms (Int64/UInt8/UInt32) share the same fix as the
1645        // Float32 arms; cover one representative (UInt8 -> Float64) so both branch
1646        // families are exercised.
1647        let values = UInt8Array::from(vec![Some(10u8), None, Some(30), Some(40)]);
1648        let fsl = FixedSizeListArray::new(
1649            Arc::new(Field::new("item", DataType::UInt8, true)),
1650            2,
1651            Arc::new(values),
1652            None,
1653        );
1654
1655        let converted = fsl.convert_to_floating_point().unwrap();
1656
1657        assert_eq!(converted.len(), 2);
1658        let conv_values = converted
1659            .values()
1660            .as_any()
1661            .downcast_ref::<Float64Array>()
1662            .unwrap();
1663        assert_eq!(conv_values.len(), 4);
1664        assert_eq!(conv_values.value(0), 10.0);
1665        assert!(conv_values.is_null(1));
1666        assert_eq!(conv_values.value(2), 30.0);
1667        assert_eq!(conv_values.value(3), 40.0);
1668    }
1669
1670    #[test]
1671    fn test_merge_recursive() {
1672        let a_array = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
1673        let e_array = Int32Array::from(vec![Some(4), Some(5), Some(6)]);
1674        let c_array = Int32Array::from(vec![Some(7), Some(8), Some(9)]);
1675        let d_array = StringArray::from(vec![Some("a"), Some("b"), Some("c")]);
1676
1677        let left_schema = Schema::new(vec![
1678            Field::new("a", DataType::Int32, true),
1679            Field::new(
1680                "b",
1681                DataType::Struct(vec![Field::new("c", DataType::Int32, true)].into()),
1682                true,
1683            ),
1684        ]);
1685        let left_batch = RecordBatch::try_new(
1686            Arc::new(left_schema),
1687            vec![
1688                Arc::new(a_array.clone()),
1689                Arc::new(StructArray::from(vec![(
1690                    Arc::new(Field::new("c", DataType::Int32, true)),
1691                    Arc::new(c_array.clone()) as ArrayRef,
1692                )])),
1693            ],
1694        )
1695        .unwrap();
1696
1697        let right_schema = Schema::new(vec![
1698            Field::new("e", DataType::Int32, true),
1699            Field::new(
1700                "b",
1701                DataType::Struct(vec![Field::new("d", DataType::Utf8, true)].into()),
1702                true,
1703            ),
1704        ]);
1705        let right_batch = RecordBatch::try_new(
1706            Arc::new(right_schema),
1707            vec![
1708                Arc::new(e_array.clone()),
1709                Arc::new(StructArray::from(vec![(
1710                    Arc::new(Field::new("d", DataType::Utf8, true)),
1711                    Arc::new(d_array.clone()) as ArrayRef,
1712                )])) as ArrayRef,
1713            ],
1714        )
1715        .unwrap();
1716
1717        let merged_schema = Schema::new(vec![
1718            Field::new("a", DataType::Int32, true),
1719            Field::new(
1720                "b",
1721                DataType::Struct(
1722                    vec![
1723                        Field::new("c", DataType::Int32, true),
1724                        Field::new("d", DataType::Utf8, true),
1725                    ]
1726                    .into(),
1727                ),
1728                true,
1729            ),
1730            Field::new("e", DataType::Int32, true),
1731        ]);
1732        let merged_batch = RecordBatch::try_new(
1733            Arc::new(merged_schema),
1734            vec![
1735                Arc::new(a_array) as ArrayRef,
1736                Arc::new(StructArray::from(vec![
1737                    (
1738                        Arc::new(Field::new("c", DataType::Int32, true)),
1739                        Arc::new(c_array) as ArrayRef,
1740                    ),
1741                    (
1742                        Arc::new(Field::new("d", DataType::Utf8, true)),
1743                        Arc::new(d_array) as ArrayRef,
1744                    ),
1745                ])) as ArrayRef,
1746                Arc::new(e_array) as ArrayRef,
1747            ],
1748        )
1749        .unwrap();
1750
1751        let result = left_batch.merge(&right_batch).unwrap();
1752        assert_eq!(result, merged_batch);
1753    }
1754
1755    #[test]
1756    fn test_merge_with_schema() {
1757        fn test_batch(names: &[&str], types: &[DataType]) -> (Schema, RecordBatch) {
1758            let fields: Fields = names
1759                .iter()
1760                .zip(types)
1761                .map(|(name, ty)| Field::new(name.to_string(), ty.clone(), false))
1762                .collect();
1763            let schema = Schema::new(vec![Field::new(
1764                "struct",
1765                DataType::Struct(fields.clone()),
1766                false,
1767            )]);
1768            let children = types.iter().map(new_empty_array).collect::<Vec<_>>();
1769            let batch = RecordBatch::try_new(
1770                Arc::new(schema.clone()),
1771                vec![Arc::new(StructArray::new(fields, children, None)) as ArrayRef],
1772            );
1773            (schema, batch.unwrap())
1774        }
1775
1776        let (_, left_batch) = test_batch(&["a", "b"], &[DataType::Int32, DataType::Int64]);
1777        let (_, right_batch) = test_batch(&["c", "b"], &[DataType::Int32, DataType::Int64]);
1778        let (output_schema, _) = test_batch(
1779            &["b", "a", "c"],
1780            &[DataType::Int64, DataType::Int32, DataType::Int32],
1781        );
1782
1783        // If we use merge_with_schema the schema is respected
1784        let merged = left_batch
1785            .merge_with_schema(&right_batch, &output_schema)
1786            .unwrap();
1787        assert_eq!(merged.schema().as_ref(), &output_schema);
1788
1789        // If we use merge we get first-come first-serve based on the left batch
1790        let (naive_schema, _) = test_batch(
1791            &["a", "b", "c"],
1792            &[DataType::Int32, DataType::Int64, DataType::Int32],
1793        );
1794        let merged = left_batch.merge(&right_batch).unwrap();
1795        assert_eq!(merged.schema().as_ref(), &naive_schema);
1796    }
1797
1798    #[test]
1799    fn test_merge_list_struct() {
1800        let x_field = Arc::new(Field::new("x", DataType::Int32, true));
1801        let y_field = Arc::new(Field::new("y", DataType::Int32, true));
1802        let x_struct_field = Arc::new(Field::new(
1803            "item",
1804            DataType::Struct(Fields::from(vec![x_field.clone()])),
1805            true,
1806        ));
1807        let y_struct_field = Arc::new(Field::new(
1808            "item",
1809            DataType::Struct(Fields::from(vec![y_field.clone()])),
1810            true,
1811        ));
1812        let both_struct_field = Arc::new(Field::new(
1813            "item",
1814            DataType::Struct(Fields::from(vec![x_field.clone(), y_field.clone()])),
1815            true,
1816        ));
1817        let left_schema = Schema::new(vec![Field::new(
1818            "list_struct",
1819            DataType::List(x_struct_field.clone()),
1820            true,
1821        )]);
1822        let right_schema = Schema::new(vec![Field::new(
1823            "list_struct",
1824            DataType::List(y_struct_field.clone()),
1825            true,
1826        )]);
1827        let both_schema = Schema::new(vec![Field::new(
1828            "list_struct",
1829            DataType::List(both_struct_field.clone()),
1830            true,
1831        )]);
1832
1833        let x = Arc::new(Int32Array::from(vec![1]));
1834        let y = Arc::new(Int32Array::from(vec![2]));
1835        let x_struct = Arc::new(StructArray::new(
1836            Fields::from(vec![x_field.clone()]),
1837            vec![x.clone()],
1838            None,
1839        ));
1840        let y_struct = Arc::new(StructArray::new(
1841            Fields::from(vec![y_field.clone()]),
1842            vec![y.clone()],
1843            None,
1844        ));
1845        let both_struct = Arc::new(StructArray::new(
1846            Fields::from(vec![x_field.clone(), y_field.clone()]),
1847            vec![x.clone(), y],
1848            None,
1849        ));
1850        let both_null_struct = Arc::new(StructArray::new(
1851            Fields::from(vec![x_field, y_field]),
1852            vec![x, Arc::new(new_null_array(&DataType::Int32, 1))],
1853            None,
1854        ));
1855        let offsets = OffsetBuffer::from_lengths([1]);
1856        let x_s_list = ListArray::new(x_struct_field, offsets.clone(), x_struct, None);
1857        let y_s_list = ListArray::new(y_struct_field, offsets.clone(), y_struct, None);
1858        let both_list = ListArray::new(
1859            both_struct_field.clone(),
1860            offsets.clone(),
1861            both_struct,
1862            None,
1863        );
1864        let both_null_list = ListArray::new(both_struct_field, offsets, both_null_struct, None);
1865        let x_batch =
1866            RecordBatch::try_new(Arc::new(left_schema), vec![Arc::new(x_s_list)]).unwrap();
1867        let y_batch = RecordBatch::try_new(
1868            Arc::new(right_schema.clone()),
1869            vec![Arc::new(y_s_list.clone())],
1870        )
1871        .unwrap();
1872        let merged = x_batch.merge(&y_batch).unwrap();
1873        let expected =
1874            RecordBatch::try_new(Arc::new(both_schema.clone()), vec![Arc::new(both_list)]).unwrap();
1875        assert_eq!(merged, expected);
1876
1877        let y_null_list = new_null_array(y_s_list.data_type(), 1);
1878        let y_null_batch =
1879            RecordBatch::try_new(Arc::new(right_schema), vec![Arc::new(y_null_list.clone())])
1880                .unwrap();
1881        let expected =
1882            RecordBatch::try_new(Arc::new(both_schema), vec![Arc::new(both_null_list)]).unwrap();
1883        let merged = x_batch.merge(&y_null_batch).unwrap();
1884        assert_eq!(merged, expected);
1885    }
1886
1887    #[test]
1888    fn test_merge_list_struct_identical_schema() {
1889        // Merging two batches whose `List<Struct>` columns have identical types
1890        // should yield a single column equal to the input (there is nothing to
1891        // merge), not a struct with the field pushed twice.
1892        let x_field = Arc::new(Field::new("x", DataType::Int32, true));
1893        let item_field = Arc::new(Field::new(
1894            "item",
1895            DataType::Struct(Fields::from(vec![x_field.clone()])),
1896            true,
1897        ));
1898        let schema = Arc::new(Schema::new(vec![Field::new(
1899            "list_struct",
1900            DataType::List(item_field.clone()),
1901            true,
1902        )]));
1903
1904        let build_list = |values: Vec<i32>| {
1905            let len = values.len();
1906            let item_struct = Arc::new(StructArray::new(
1907                Fields::from(vec![x_field.clone()]),
1908                vec![Arc::new(Int32Array::from(values))],
1909                None,
1910            ));
1911            ListArray::new(
1912                item_field.clone(),
1913                OffsetBuffer::from_lengths([len]),
1914                item_struct,
1915                None,
1916            )
1917        };
1918
1919        // Distinct values so the equality assertion proves the left column is kept,
1920        // not that some column with the same values happens to be present.
1921        let left =
1922            RecordBatch::try_new(schema.clone(), vec![Arc::new(build_list(vec![1, 2]))]).unwrap();
1923        let right = RecordBatch::try_new(schema, vec![Arc::new(build_list(vec![3, 4]))]).unwrap();
1924
1925        let merged = left.merge(&right).unwrap();
1926
1927        // Exactly one column, equal to the left input: the field must not be duplicated.
1928        assert_eq!(merged.num_columns(), 1);
1929        assert_eq!(merged, left);
1930    }
1931
1932    #[test]
1933    fn test_merge_nested_list_struct_identical_schema() {
1934        // A `List<Struct>` nested inside a struct reaches the identical-type
1935        // short-circuit through the recursive `merge` call. The recursion must
1936        // keep the field exactly once, equal to the left input.
1937        let x_field = Arc::new(Field::new("x", DataType::Int32, true));
1938        let item_field = Arc::new(Field::new(
1939            "item",
1940            DataType::Struct(Fields::from(vec![x_field.clone()])),
1941            true,
1942        ));
1943        let companies_field = Arc::new(Field::new(
1944            "companies",
1945            DataType::List(item_field.clone()),
1946            true,
1947        ));
1948        let schema = Arc::new(Schema::new(vec![Field::new(
1949            "outer",
1950            DataType::Struct(Fields::from(vec![companies_field.clone()])),
1951            true,
1952        )]));
1953
1954        let build_outer = |x: i32| {
1955            // One list row: [{x}].
1956            let item_struct = Arc::new(StructArray::new(
1957                Fields::from(vec![x_field.clone()]),
1958                vec![Arc::new(Int32Array::from(vec![x]))],
1959                None,
1960            ));
1961            let companies = Arc::new(ListArray::new(
1962                item_field.clone(),
1963                OffsetBuffer::from_lengths([1]),
1964                item_struct,
1965                None,
1966            ));
1967            StructArray::new(
1968                Fields::from(vec![companies_field.clone()]),
1969                vec![companies],
1970                None,
1971            )
1972        };
1973
1974        // Distinct values so equality proves the left column is kept.
1975        let left = RecordBatch::try_new(schema.clone(), vec![Arc::new(build_outer(10))]).unwrap();
1976        let right = RecordBatch::try_new(schema, vec![Arc::new(build_outer(20))]).unwrap();
1977
1978        let merged = left.merge(&right).unwrap();
1979
1980        // The recursive identical-type merge keeps exactly one `companies` field,
1981        // with no double-push inside the nested struct.
1982        assert_eq!(merged.column(0).as_struct().num_columns(), 1);
1983        assert_eq!(merged, left);
1984    }
1985
1986    #[test]
1987    fn test_byte_width_opt() {
1988        assert_eq!(DataType::Int32.byte_width_opt(), Some(4));
1989        assert_eq!(DataType::Int64.byte_width_opt(), Some(8));
1990        assert_eq!(DataType::Float32.byte_width_opt(), Some(4));
1991        assert_eq!(DataType::Float64.byte_width_opt(), Some(8));
1992        assert_eq!(DataType::Utf8.byte_width_opt(), None);
1993        assert_eq!(DataType::Binary.byte_width_opt(), None);
1994        assert_eq!(
1995            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))).byte_width_opt(),
1996            None
1997        );
1998        assert_eq!(
1999            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 3)
2000                .byte_width_opt(),
2001            Some(12)
2002        );
2003        assert_eq!(
2004            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 4)
2005                .byte_width_opt(),
2006            Some(16)
2007        );
2008        assert_eq!(
2009            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Utf8, true)), 5)
2010                .byte_width_opt(),
2011            None
2012        );
2013    }
2014
2015    #[test]
2016    fn test_take_record_batch() {
2017        let schema = Arc::new(Schema::new(vec![
2018            Field::new("a", DataType::Int32, true),
2019            Field::new("b", DataType::Utf8, true),
2020        ]));
2021        let batch = RecordBatch::try_new(
2022            schema.clone(),
2023            vec![
2024                Arc::new(Int32Array::from_iter_values(0..20)),
2025                Arc::new(StringArray::from_iter_values(
2026                    (0..20).map(|i| format!("str-{}", i)),
2027                )),
2028            ],
2029        )
2030        .unwrap();
2031        let taken = batch.take(&(vec![1_u32, 5_u32, 10_u32].into())).unwrap();
2032        assert_eq!(
2033            taken,
2034            RecordBatch::try_new(
2035                schema,
2036                vec![
2037                    Arc::new(Int32Array::from(vec![1, 5, 10])),
2038                    Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])),
2039                ],
2040            )
2041            .unwrap()
2042        )
2043    }
2044
2045    #[test]
2046    fn test_schema_project_by_schema() {
2047        let metadata = [("key".to_string(), "value".to_string())];
2048        let schema = Arc::new(
2049            Schema::new(vec![
2050                Field::new("a", DataType::Int32, true),
2051                Field::new("b", DataType::Utf8, true),
2052            ])
2053            .with_metadata(metadata.clone().into()),
2054        );
2055        let batch = RecordBatch::try_new(
2056            schema,
2057            vec![
2058                Arc::new(Int32Array::from_iter_values(0..20)),
2059                Arc::new(StringArray::from_iter_values(
2060                    (0..20).map(|i| format!("str-{}", i)),
2061                )),
2062            ],
2063        )
2064        .unwrap();
2065
2066        // Empty schema
2067        let empty_schema = Schema::empty();
2068        let empty_projected = batch.project_by_schema(&empty_schema).unwrap();
2069        let expected_schema = empty_schema.with_metadata(metadata.clone().into());
2070        assert_eq!(
2071            empty_projected,
2072            RecordBatch::from(StructArray::new_empty_fields(batch.num_rows(), None))
2073                .with_schema(Arc::new(expected_schema))
2074                .unwrap()
2075        );
2076
2077        // Re-ordered schema
2078        let reordered_schema = Schema::new(vec![
2079            Field::new("b", DataType::Utf8, true),
2080            Field::new("a", DataType::Int32, true),
2081        ]);
2082        let reordered_projected = batch.project_by_schema(&reordered_schema).unwrap();
2083        let expected_schema = Arc::new(reordered_schema.with_metadata(metadata.clone().into()));
2084        assert_eq!(
2085            reordered_projected,
2086            RecordBatch::try_new(
2087                expected_schema,
2088                vec![
2089                    Arc::new(StringArray::from_iter_values(
2090                        (0..20).map(|i| format!("str-{}", i)),
2091                    )),
2092                    Arc::new(Int32Array::from_iter_values(0..20)),
2093                ],
2094            )
2095            .unwrap()
2096        );
2097
2098        // Sub schema
2099        let sub_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
2100        let sub_projected = batch.project_by_schema(&sub_schema).unwrap();
2101        let expected_schema = Arc::new(sub_schema.with_metadata(metadata.into()));
2102        assert_eq!(
2103            sub_projected,
2104            RecordBatch::try_new(
2105                expected_schema,
2106                vec![Arc::new(Int32Array::from_iter_values(0..20))],
2107            )
2108            .unwrap()
2109        );
2110    }
2111
2112    #[test]
2113    fn test_project_rebuilds_sliced_map() {
2114        // A sliced MapArray keeps its full entries array behind sliced
2115        // offsets and validity; the Map projection arm must rebuild it
2116        // without renormalizing either.
2117        let entry_fields = Fields::from(vec![
2118            Field::new("key", DataType::Utf8, false),
2119            Field::new("value", DataType::Int32, true),
2120        ]);
2121        let entries_field = Arc::new(Field::new(
2122            "entries",
2123            DataType::Struct(entry_fields.clone()),
2124            false,
2125        ));
2126        let entries = StructArray::new(
2127            entry_fields,
2128            vec![
2129                Arc::new(StringArray::from(vec!["k0", "k1", "k2"])) as ArrayRef,
2130                Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef,
2131            ],
2132            None,
2133        );
2134        let map = MapArray::new(
2135            entries_field.clone(),
2136            OffsetBuffer::new(vec![0, 1, 1, 3].into()),
2137            entries,
2138            Some(arrow_buffer::NullBuffer::from(vec![true, false, true])),
2139            false,
2140        );
2141        let schema = Arc::new(Schema::new(vec![Field::new(
2142            "m",
2143            DataType::Map(entries_field, false),
2144            true,
2145        )]));
2146        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(map) as ArrayRef]).unwrap();
2147
2148        // Rows 1..3: a null slot and a two-entry slot, offsets not zero-based.
2149        let projected = batch
2150            .slice(1, 2)
2151            .project_by_schema(schema.as_ref())
2152            .unwrap();
2153        let map = projected.column(0).as_map();
2154        assert!(map.is_null(0));
2155        assert!(map.is_valid(1));
2156        assert_eq!(map.value_length(1), 2);
2157        assert_eq!(
2158            map.value(1)
2159                .column(1)
2160                .as_primitive::<arrow_array::types::Int32Type>()
2161                .values(),
2162            &[2, 3]
2163        );
2164    }
2165
2166    #[test]
2167    fn test_project_preserves_struct_validity() {
2168        // Test that projecting a struct array preserves its validity (fix for issue #4385)
2169        let fields = Fields::from(vec![
2170            Field::new("id", DataType::Int32, false),
2171            Field::new("value", DataType::Float32, true),
2172        ]);
2173
2174        // Create a struct array with validity
2175        let id_array = Int32Array::from(vec![1, 2, 3]);
2176        let value_array = Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]);
2177        let struct_array = StructArray::new(
2178            fields.clone(),
2179            vec![
2180                Arc::new(id_array) as ArrayRef,
2181                Arc::new(value_array) as ArrayRef,
2182            ],
2183            Some(vec![true, false, true].into()), // Second struct is null
2184        );
2185
2186        // Project the struct array
2187        let projected = project(&struct_array, &fields).unwrap();
2188
2189        // Verify the validity is preserved
2190        assert_eq!(projected.null_count(), 1);
2191        assert!(!projected.is_null(0));
2192        assert!(projected.is_null(1));
2193        assert!(!projected.is_null(2));
2194    }
2195
2196    #[test]
2197    fn test_merge_struct_with_different_validity() {
2198        // Test case from Weston's review comment
2199        // File 1 has height field with some nulls
2200        let height_array = Int32Array::from(vec![Some(500), None, Some(600), None]);
2201        let left_fields = Fields::from(vec![Field::new("height", DataType::Int32, true)]);
2202        let left_struct = StructArray::new(
2203            left_fields,
2204            vec![Arc::new(height_array) as ArrayRef],
2205            Some(vec![true, false, true, false].into()), // Rows 2 and 4 are null structs
2206        );
2207
2208        // File 2 has width field with some nulls
2209        let width_array = Int32Array::from(vec![Some(300), Some(200), None, None]);
2210        let right_fields = Fields::from(vec![Field::new("width", DataType::Int32, true)]);
2211        let right_struct = StructArray::new(
2212            right_fields,
2213            vec![Arc::new(width_array) as ArrayRef],
2214            Some(vec![true, true, false, false].into()), // Rows 3 and 4 are null structs
2215        );
2216
2217        // Merge the two structs
2218        let merged = merge(&left_struct, &right_struct);
2219
2220        // Expected:
2221        // Row 1: both non-null -> {width: 300, height: 500}
2222        // Row 2: left null, right non-null -> {width: 200, height: null}
2223        // Row 3: left non-null, right null -> {width: null, height: 600}
2224        // Row 4: both null -> null struct
2225
2226        assert_eq!(merged.null_count(), 1); // Only row 4 is null
2227        assert!(!merged.is_null(0));
2228        assert!(!merged.is_null(1));
2229        assert!(!merged.is_null(2));
2230        assert!(merged.is_null(3));
2231
2232        // Check field values
2233        let height_col = merged.column_by_name("height").unwrap();
2234        let height_values = height_col.as_any().downcast_ref::<Int32Array>().unwrap();
2235        assert_eq!(height_values.value(0), 500);
2236        assert!(height_values.is_null(1)); // height is null when left struct was null
2237        assert_eq!(height_values.value(2), 600);
2238
2239        let width_col = merged.column_by_name("width").unwrap();
2240        let width_values = width_col.as_any().downcast_ref::<Int32Array>().unwrap();
2241        assert_eq!(width_values.value(0), 300);
2242        assert_eq!(width_values.value(1), 200);
2243        assert!(width_values.is_null(2)); // width is null when right struct was null
2244
2245        // An all-null validity buffer is data, not a placeholder meaning "this side has no
2246        // validity": merging two of them keeps the rows null.
2247        let all_null_left = StructArray::new(
2248            Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2249            vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef],
2250            Some(vec![false, false].into()),
2251        );
2252        let all_null_right = StructArray::new(
2253            Fields::from(vec![Field::new("width", DataType::Int32, true)]),
2254            vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef],
2255            Some(vec![false, false].into()),
2256        );
2257
2258        let merged = merge(&all_null_left, &all_null_right);
2259        assert_eq!(merged.null_count(), 2);
2260
2261        // An all-null side is the identity of the merge, so the other side decides each row.
2262        let partial_left = StructArray::new(
2263            Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2264            vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef],
2265            Some(vec![true, false].into()),
2266        );
2267        let merged = merge(&partial_left, &all_null_right);
2268        assert!(!merged.is_null(0));
2269        assert!(merged.is_null(1));
2270
2271        // A missing validity buffer means all rows are valid, which absorbs an all-null side.
2272        let all_valid_left = StructArray::new(
2273            Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2274            vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef],
2275            None,
2276        );
2277        let merged = merge(&all_valid_left, &all_null_right);
2278        assert_eq!(merged.null_count(), 0);
2279
2280        // An explicit all-valid buffer has the same semantics as a missing buffer.
2281        let all_valid: arrow_buffer::NullBuffer = vec![true, true].into();
2282        let partial: arrow_buffer::NullBuffer = vec![true, false].into();
2283        assert!(merge_struct_validity(Some(&all_valid), Some(&partial)).is_none());
2284        assert!(merge_struct_validity(Some(&partial), Some(&all_valid)).is_none());
2285    }
2286
2287    #[test]
2288    fn test_merge_null_typed_column_with_parent_validity() {
2289        // Reproduces ENT-990: panic in adjust_child_validity when a Null-typed column
2290        // exists on one side and the parent struct has null rows.
2291        // Arrow's Null type has no null bitmap, so passing one to ArrayData::try_new panics.
2292        let left_struct = StructArray::new(
2293            Fields::from(vec![Field::new("a", DataType::Int32, true)]),
2294            vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef],
2295            Some(vec![true, false].into()),
2296        );
2297        let right_struct = StructArray::new(
2298            Fields::from(vec![Field::new("b", DataType::Null, true)]),
2299            vec![Arc::new(NullArray::new(2)) as ArrayRef],
2300            Some(vec![true, false].into()),
2301        );
2302
2303        // Previously panicked: "Arrays of type Null cannot contain a null bitmask"
2304        let merged = merge(&left_struct, &right_struct);
2305        assert_eq!(merged.len(), 2);
2306        let b_col = merged.column_by_name("b").unwrap();
2307        // DataType::Null implies all-null by definition; no null bitmap is stored.
2308        assert_eq!(b_col.data_type(), &DataType::Null);
2309        assert_eq!(b_col.len(), 2);
2310    }
2311
2312    #[test]
2313    fn test_merge_with_schema_with_nullable_struct_list_schema_mismatch() {
2314        // left_list setup
2315        let left_company_id = Arc::new(Int32Array::from(vec![None, None]));
2316        let left_count = Arc::new(Int32Array::from(vec![None, None]));
2317        let left_struct = Arc::new(StructArray::new(
2318            Fields::from(vec![
2319                Field::new("company_id", DataType::Int32, true),
2320                Field::new("count", DataType::Int32, true),
2321            ]),
2322            vec![left_company_id, left_count],
2323            None,
2324        ));
2325        let left_list = Arc::new(ListArray::new(
2326            Arc::new(Field::new(
2327                "item",
2328                DataType::Struct(left_struct.fields().clone()),
2329                true,
2330            )),
2331            OffsetBuffer::from_lengths([2]),
2332            left_struct,
2333            None,
2334        ));
2335
2336        // Right List Setup
2337        let right_company_name = Arc::new(StringArray::from(vec!["Google", "Microsoft"]));
2338        let right_struct = Arc::new(StructArray::new(
2339            Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2340            vec![right_company_name],
2341            None,
2342        ));
2343        let right_list = Arc::new(ListArray::new(
2344            Arc::new(Field::new(
2345                "item",
2346                DataType::Struct(right_struct.fields().clone()),
2347                true,
2348            )),
2349            OffsetBuffer::from_lengths([2]),
2350            right_struct,
2351            None,
2352        ));
2353
2354        let target_fields = Fields::from(vec![Field::new(
2355            "companies",
2356            DataType::List(Arc::new(Field::new(
2357                "item",
2358                DataType::Struct(Fields::from(vec![
2359                    Field::new("company_id", DataType::Int32, true),
2360                    Field::new("company_name", DataType::Utf8, true),
2361                    Field::new("count", DataType::Int32, true),
2362                ])),
2363                true,
2364            ))),
2365            true,
2366        )]);
2367
2368        let left_batch = RecordBatch::try_new(
2369            Arc::new(Schema::new(vec![Field::new(
2370                "companies",
2371                left_list.data_type().clone(),
2372                true,
2373            )])),
2374            vec![left_list as ArrayRef],
2375        )
2376        .unwrap();
2377
2378        let right_batch = RecordBatch::try_new(
2379            Arc::new(Schema::new(vec![Field::new(
2380                "companies",
2381                right_list.data_type().clone(),
2382                true,
2383            )])),
2384            vec![right_list as ArrayRef],
2385        )
2386        .unwrap();
2387
2388        let merged = left_batch
2389            .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2390            .unwrap();
2391
2392        // Verify the merged structure
2393        let merged_list = merged
2394            .column_by_name("companies")
2395            .unwrap()
2396            .as_any()
2397            .downcast_ref::<ListArray>()
2398            .unwrap();
2399        let merged_struct = merged_list.values().as_struct();
2400
2401        // Should have all 3 fields
2402        assert_eq!(merged_struct.num_columns(), 3);
2403        assert!(merged_struct.column_by_name("company_id").is_some());
2404        assert!(merged_struct.column_by_name("company_name").is_some());
2405        assert!(merged_struct.column_by_name("count").is_some());
2406
2407        // Verify values
2408        let company_id = merged_struct
2409            .column_by_name("company_id")
2410            .unwrap()
2411            .as_any()
2412            .downcast_ref::<Int32Array>()
2413            .unwrap();
2414        assert!(company_id.is_null(0));
2415        assert!(company_id.is_null(1));
2416
2417        let company_name = merged_struct
2418            .column_by_name("company_name")
2419            .unwrap()
2420            .as_any()
2421            .downcast_ref::<StringArray>()
2422            .unwrap();
2423        assert_eq!(company_name.value(0), "Google");
2424        assert_eq!(company_name.value(1), "Microsoft");
2425
2426        let count = merged_struct
2427            .column_by_name("count")
2428            .unwrap()
2429            .as_any()
2430            .downcast_ref::<Int32Array>()
2431            .unwrap();
2432        assert!(count.is_null(0));
2433        assert!(count.is_null(1));
2434    }
2435
2436    #[test]
2437    fn test_merge_struct_lists() {
2438        test_merge_struct_lists_generic::<i32>();
2439    }
2440
2441    #[test]
2442    fn test_merge_struct_large_lists() {
2443        test_merge_struct_lists_generic::<i64>();
2444    }
2445
2446    fn test_merge_struct_lists_generic<O: OffsetSizeTrait>() {
2447        // left_list setup
2448        let left_company_id = Arc::new(Int32Array::from(vec![
2449            Some(1),
2450            Some(2),
2451            Some(3),
2452            Some(4),
2453            Some(5),
2454            Some(6),
2455            Some(7),
2456            Some(8),
2457            Some(9),
2458            Some(10),
2459            Some(11),
2460            Some(12),
2461            Some(13),
2462            Some(14),
2463            Some(15),
2464            Some(16),
2465            Some(17),
2466            Some(18),
2467            Some(19),
2468            Some(20),
2469        ]));
2470        let left_count = Arc::new(Int32Array::from(vec![
2471            Some(10),
2472            Some(20),
2473            Some(30),
2474            Some(40),
2475            Some(50),
2476            Some(60),
2477            Some(70),
2478            Some(80),
2479            Some(90),
2480            Some(100),
2481            Some(110),
2482            Some(120),
2483            Some(130),
2484            Some(140),
2485            Some(150),
2486            Some(160),
2487            Some(170),
2488            Some(180),
2489            Some(190),
2490            Some(200),
2491        ]));
2492        let left_struct = Arc::new(StructArray::new(
2493            Fields::from(vec![
2494                Field::new("company_id", DataType::Int32, true),
2495                Field::new("count", DataType::Int32, true),
2496            ]),
2497            vec![left_company_id, left_count],
2498            None,
2499        ));
2500
2501        let left_list = Arc::new(GenericListArray::<O>::new(
2502            Arc::new(Field::new(
2503                "item",
2504                DataType::Struct(left_struct.fields().clone()),
2505                true,
2506            )),
2507            OffsetBuffer::from_lengths([3, 1]),
2508            left_struct.clone(),
2509            None,
2510        ));
2511
2512        let left_list_struct = Arc::new(StructArray::new(
2513            Fields::from(vec![Field::new(
2514                "companies",
2515                if O::IS_LARGE {
2516                    DataType::LargeList(Arc::new(Field::new(
2517                        "item",
2518                        DataType::Struct(left_struct.fields().clone()),
2519                        true,
2520                    )))
2521                } else {
2522                    DataType::List(Arc::new(Field::new(
2523                        "item",
2524                        DataType::Struct(left_struct.fields().clone()),
2525                        true,
2526                    )))
2527                },
2528                true,
2529            )]),
2530            vec![left_list as ArrayRef],
2531            None,
2532        ));
2533
2534        // right_list setup
2535        let right_company_name = Arc::new(StringArray::from(vec![
2536            "Google",
2537            "Microsoft",
2538            "Apple",
2539            "Facebook",
2540        ]));
2541        let right_struct = Arc::new(StructArray::new(
2542            Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2543            vec![right_company_name],
2544            None,
2545        ));
2546        let right_list = Arc::new(GenericListArray::<O>::new(
2547            Arc::new(Field::new(
2548                "item",
2549                DataType::Struct(right_struct.fields().clone()),
2550                true,
2551            )),
2552            OffsetBuffer::from_lengths([3, 1]),
2553            right_struct.clone(),
2554            None,
2555        ));
2556
2557        let right_list_struct = Arc::new(StructArray::new(
2558            Fields::from(vec![Field::new(
2559                "companies",
2560                if O::IS_LARGE {
2561                    DataType::LargeList(Arc::new(Field::new(
2562                        "item",
2563                        DataType::Struct(right_struct.fields().clone()),
2564                        true,
2565                    )))
2566                } else {
2567                    DataType::List(Arc::new(Field::new(
2568                        "item",
2569                        DataType::Struct(right_struct.fields().clone()),
2570                        true,
2571                    )))
2572                },
2573                true,
2574            )]),
2575            vec![right_list as ArrayRef],
2576            None,
2577        ));
2578
2579        // prepare schema
2580        let target_fields = Fields::from(vec![Field::new(
2581            "companies",
2582            if O::IS_LARGE {
2583                DataType::LargeList(Arc::new(Field::new(
2584                    "item",
2585                    DataType::Struct(Fields::from(vec![
2586                        Field::new("company_id", DataType::Int32, true),
2587                        Field::new("company_name", DataType::Utf8, true),
2588                        Field::new("count", DataType::Int32, true),
2589                    ])),
2590                    true,
2591                )))
2592            } else {
2593                DataType::List(Arc::new(Field::new(
2594                    "item",
2595                    DataType::Struct(Fields::from(vec![
2596                        Field::new("company_id", DataType::Int32, true),
2597                        Field::new("company_name", DataType::Utf8, true),
2598                        Field::new("count", DataType::Int32, true),
2599                    ])),
2600                    true,
2601                )))
2602            },
2603            true,
2604        )]);
2605
2606        // merge left_list and right_list
2607        let merged_array = merge_with_schema(&left_list_struct, &right_list_struct, &target_fields);
2608        assert_eq!(merged_array.len(), 2);
2609    }
2610
2611    #[test]
2612    fn test_merge_with_schema_sliced_list_struct() {
2613        test_merge_with_schema_sliced_list_struct_generic::<i32>();
2614    }
2615
2616    #[test]
2617    fn test_merge_with_schema_sliced_large_list_struct() {
2618        test_merge_with_schema_sliced_list_struct_generic::<i64>();
2619    }
2620
2621    // Regression for #6580: merge_with_schema panicked when the left list was a
2622    // sliced view whose offsets did not start at zero (common after a filtered
2623    // scan). Cloning those offsets alongside `trimmed_values` produced offsets
2624    // larger than the trimmed child, panicking in `(Large)ListArray::new`.
2625    fn test_merge_with_schema_sliced_list_struct_generic<O: OffsetSizeTrait>() {
2626        let make_list_dtype = |item_field: Arc<Field>| {
2627            if O::IS_LARGE {
2628                DataType::LargeList(item_field)
2629            } else {
2630                DataType::List(item_field)
2631            }
2632        };
2633
2634        // Build a List<Struct> with two rows of 5 items each, then slice away
2635        // the first row so the remaining list's offsets start at 5, not 0.
2636        let struct_fields_a = Fields::from(vec![Field::new("a", DataType::Int32, true)]);
2637        let left_values = Arc::new(StructArray::new(
2638            struct_fields_a.clone(),
2639            vec![Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef],
2640            None,
2641        ));
2642        let full_list = GenericListArray::<O>::new(
2643            Arc::new(Field::new("item", DataType::Struct(struct_fields_a), true)),
2644            OffsetBuffer::<O>::from_lengths([5, 5]),
2645            left_values,
2646            None,
2647        );
2648        let sliced_left = full_list.slice(1, 1);
2649        assert_eq!(sliced_left.offsets()[0].as_usize(), 5);
2650        assert_eq!(sliced_left.offsets()[1].as_usize(), 10);
2651
2652        let struct_fields_b = Fields::from(vec![Field::new("b", DataType::Int32, true)]);
2653        let right_values = Arc::new(StructArray::new(
2654            struct_fields_b.clone(),
2655            vec![Arc::new(Int32Array::from_iter_values(100..105)) as ArrayRef],
2656            None,
2657        ));
2658        let right_list = GenericListArray::<O>::new(
2659            Arc::new(Field::new("item", DataType::Struct(struct_fields_b), true)),
2660            OffsetBuffer::<O>::from_lengths([5]),
2661            right_values,
2662            None,
2663        );
2664
2665        let target_item_field = Arc::new(Field::new(
2666            "item",
2667            DataType::Struct(Fields::from(vec![
2668                Field::new("a", DataType::Int32, true),
2669                Field::new("b", DataType::Int32, true),
2670            ])),
2671            true,
2672        ));
2673        let target_fields = Fields::from(vec![Field::new(
2674            "items",
2675            make_list_dtype(target_item_field),
2676            true,
2677        )]);
2678
2679        let left_batch = RecordBatch::try_new(
2680            Arc::new(Schema::new(vec![Field::new(
2681                "items",
2682                sliced_left.data_type().clone(),
2683                true,
2684            )])),
2685            vec![Arc::new(sliced_left) as ArrayRef],
2686        )
2687        .unwrap();
2688        let right_batch = RecordBatch::try_new(
2689            Arc::new(Schema::new(vec![Field::new(
2690                "items",
2691                right_list.data_type().clone(),
2692                true,
2693            )])),
2694            vec![Arc::new(right_list) as ArrayRef],
2695        )
2696        .unwrap();
2697
2698        let merged = left_batch
2699            .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2700            .unwrap();
2701
2702        let merged_list = merged
2703            .column_by_name("items")
2704            .unwrap()
2705            .as_any()
2706            .downcast_ref::<GenericListArray<O>>()
2707            .unwrap();
2708        assert_eq!(merged_list.len(), 1);
2709        assert_eq!(merged_list.value_length(0).as_usize(), 5);
2710        let merged_struct = merged_list.values().as_struct();
2711        assert_eq!(merged_struct.num_columns(), 2);
2712        let a = merged_struct
2713            .column_by_name("a")
2714            .unwrap()
2715            .as_any()
2716            .downcast_ref::<Int32Array>()
2717            .unwrap();
2718        // After shifting offsets to zero, values 5..10 should be first.
2719        let a_vals: Vec<i32> = a.iter().map(|v| v.unwrap()).collect();
2720        assert_eq!(a_vals, vec![5, 6, 7, 8, 9]);
2721    }
2722
2723    #[test]
2724    fn test_project_by_schema_list_struct_reorder() {
2725        // Test that project_by_schema correctly reorders fields inside List<Struct>
2726        // This is a regression test for issue #5702
2727
2728        // Source schema with inner struct fields in order: c, b, a
2729        let source_inner_struct = DataType::Struct(Fields::from(vec![
2730            Field::new("c", DataType::Utf8, true),
2731            Field::new("b", DataType::Utf8, true),
2732            Field::new("a", DataType::Utf8, true),
2733        ]));
2734        let source_schema = Arc::new(Schema::new(vec![
2735            Field::new("id", DataType::Int32, false),
2736            Field::new(
2737                "data",
2738                DataType::List(Arc::new(Field::new(
2739                    "item",
2740                    source_inner_struct.clone(),
2741                    true,
2742                ))),
2743                true,
2744            ),
2745        ]));
2746
2747        // Create source data with c, b, a order
2748        let c_array = StringArray::from(vec!["c1", "c2"]);
2749        let b_array = StringArray::from(vec!["b1", "b2"]);
2750        let a_array = StringArray::from(vec!["a1", "a2"]);
2751        let inner_struct = StructArray::from(vec![
2752            (
2753                Arc::new(Field::new("c", DataType::Utf8, true)),
2754                Arc::new(c_array) as ArrayRef,
2755            ),
2756            (
2757                Arc::new(Field::new("b", DataType::Utf8, true)),
2758                Arc::new(b_array) as ArrayRef,
2759            ),
2760            (
2761                Arc::new(Field::new("a", DataType::Utf8, true)),
2762                Arc::new(a_array) as ArrayRef,
2763            ),
2764        ]);
2765
2766        let list_array = ListArray::new(
2767            Arc::new(Field::new("item", source_inner_struct, true)),
2768            OffsetBuffer::from_lengths([1, 1]),
2769            Arc::new(inner_struct),
2770            None,
2771        );
2772
2773        let batch = RecordBatch::try_new(
2774            source_schema,
2775            vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(list_array)],
2776        )
2777        .unwrap();
2778
2779        // Target schema with inner struct fields in order: a, b, c
2780        let target_inner_struct = DataType::Struct(Fields::from(vec![
2781            Field::new("a", DataType::Utf8, true),
2782            Field::new("b", DataType::Utf8, true),
2783            Field::new("c", DataType::Utf8, true),
2784        ]));
2785        let target_schema = Schema::new(vec![
2786            Field::new("id", DataType::Int32, false),
2787            Field::new(
2788                "data",
2789                DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
2790                true,
2791            ),
2792        ]);
2793
2794        // Project should reorder the inner struct fields
2795        let projected = batch.project_by_schema(&target_schema).unwrap();
2796
2797        // Verify the schema is correct
2798        assert_eq!(projected.schema().as_ref(), &target_schema);
2799
2800        // Verify the data is correct by checking inner struct field order
2801        let projected_list = projected.column(1).as_list::<i32>();
2802        let projected_struct = projected_list.values().as_struct();
2803
2804        // Fields should now be in order: a, b, c
2805        assert_eq!(
2806            projected_struct.column_by_name("a").unwrap().as_ref(),
2807            &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2808        );
2809        assert_eq!(
2810            projected_struct.column_by_name("b").unwrap().as_ref(),
2811            &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2812        );
2813        assert_eq!(
2814            projected_struct.column_by_name("c").unwrap().as_ref(),
2815            &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2816        );
2817
2818        // Also verify positional access matches expected order (a=0, b=1, c=2)
2819        assert_eq!(
2820            projected_struct.column(0).as_ref(),
2821            &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2822        );
2823        assert_eq!(
2824            projected_struct.column(1).as_ref(),
2825            &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2826        );
2827        assert_eq!(
2828            projected_struct.column(2).as_ref(),
2829            &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2830        );
2831    }
2832
2833    #[test]
2834    fn test_project_by_schema_nested_list_struct() {
2835        // Test deeply nested List<Struct<List<Struct>>> projection
2836        let inner_struct = DataType::Struct(Fields::from(vec![
2837            Field::new("y", DataType::Int32, true),
2838            Field::new("x", DataType::Int32, true),
2839        ]));
2840        let source_schema = Arc::new(Schema::new(vec![Field::new(
2841            "outer",
2842            DataType::List(Arc::new(Field::new(
2843                "item",
2844                DataType::Struct(Fields::from(vec![
2845                    Field::new("b", DataType::Utf8, true),
2846                    Field::new(
2847                        "inner_list",
2848                        DataType::List(Arc::new(Field::new("item", inner_struct.clone(), true))),
2849                        true,
2850                    ),
2851                    Field::new("a", DataType::Utf8, true),
2852                ])),
2853                true,
2854            ))),
2855            true,
2856        )]));
2857
2858        // Create deeply nested data
2859        let y_array = Int32Array::from(vec![1, 2]);
2860        let x_array = Int32Array::from(vec![3, 4]);
2861        let innermost_struct = StructArray::from(vec![
2862            (
2863                Arc::new(Field::new("y", DataType::Int32, true)),
2864                Arc::new(y_array) as ArrayRef,
2865            ),
2866            (
2867                Arc::new(Field::new("x", DataType::Int32, true)),
2868                Arc::new(x_array) as ArrayRef,
2869            ),
2870        ]);
2871        let inner_list = ListArray::new(
2872            Arc::new(Field::new("item", inner_struct.clone(), true)),
2873            OffsetBuffer::from_lengths([2]),
2874            Arc::new(innermost_struct),
2875            None,
2876        );
2877
2878        let b_array = StringArray::from(vec!["b1"]);
2879        let a_array = StringArray::from(vec!["a1"]);
2880        let middle_struct = StructArray::from(vec![
2881            (
2882                Arc::new(Field::new("b", DataType::Utf8, true)),
2883                Arc::new(b_array) as ArrayRef,
2884            ),
2885            (
2886                Arc::new(Field::new(
2887                    "inner_list",
2888                    DataType::List(Arc::new(Field::new("item", inner_struct, true))),
2889                    true,
2890                )),
2891                Arc::new(inner_list) as ArrayRef,
2892            ),
2893            (
2894                Arc::new(Field::new("a", DataType::Utf8, true)),
2895                Arc::new(a_array) as ArrayRef,
2896            ),
2897        ]);
2898
2899        let outer_list = ListArray::new(
2900            Arc::new(Field::new("item", middle_struct.data_type().clone(), true)),
2901            OffsetBuffer::from_lengths([1]),
2902            Arc::new(middle_struct),
2903            None,
2904        );
2905
2906        let batch =
2907            RecordBatch::try_new(source_schema, vec![Arc::new(outer_list) as ArrayRef]).unwrap();
2908
2909        // Target schema with reordered fields at all levels
2910        let target_inner_struct = DataType::Struct(Fields::from(vec![
2911            Field::new("x", DataType::Int32, true), // x before y now
2912            Field::new("y", DataType::Int32, true),
2913        ]));
2914        let target_schema = Schema::new(vec![Field::new(
2915            "outer",
2916            DataType::List(Arc::new(Field::new(
2917                "item",
2918                DataType::Struct(Fields::from(vec![
2919                    Field::new("a", DataType::Utf8, true), // a before b now
2920                    Field::new(
2921                        "inner_list",
2922                        DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
2923                        true,
2924                    ),
2925                    Field::new("b", DataType::Utf8, true),
2926                ])),
2927                true,
2928            ))),
2929            true,
2930        )]);
2931
2932        let projected = batch.project_by_schema(&target_schema).unwrap();
2933
2934        // Verify schema
2935        assert_eq!(projected.schema().as_ref(), &target_schema);
2936
2937        // Verify deeply nested data is reordered correctly
2938        let outer_list = projected.column(0).as_list::<i32>();
2939        let middle_struct = outer_list.values().as_struct();
2940
2941        // Middle struct should have a first, then inner_list, then b
2942        assert_eq!(
2943            middle_struct.column(0).as_ref(),
2944            &StringArray::from(vec!["a1"]) as &dyn Array
2945        );
2946        assert_eq!(
2947            middle_struct.column(2).as_ref(),
2948            &StringArray::from(vec!["b1"]) as &dyn Array
2949        );
2950
2951        // Inner list's struct should have x first, then y
2952        let inner_list = middle_struct.column(1).as_list::<i32>();
2953        let innermost_struct = inner_list.values().as_struct();
2954        assert_eq!(
2955            innermost_struct.column(0).as_ref(),
2956            &Int32Array::from(vec![3, 4]) as &dyn Array
2957        );
2958        assert_eq!(
2959            innermost_struct.column(1).as_ref(),
2960            &Int32Array::from(vec![1, 2]) as &dyn Array
2961        );
2962    }
2963
2964    #[test]
2965    fn test_iter_binary_array_accepts_binary_variants() {
2966        let binary = BinaryArray::from(vec![b"a".as_slice(), b"bc"]);
2967        assert_eq!(
2968            iter_binary_array(&binary).unwrap().collect::<Vec<_>>(),
2969            vec![Some(b"a".as_slice()), Some(b"bc".as_slice())]
2970        );
2971
2972        let large_binary = LargeBinaryArray::from(vec![b"x".as_slice(), b"yz"]);
2973        assert_eq!(
2974            iter_binary_array(&large_binary)
2975                .unwrap()
2976                .collect::<Vec<_>>(),
2977            vec![Some(b"x".as_slice()), Some(b"yz".as_slice())]
2978        );
2979
2980        let binary_view = BinaryViewArray::from(vec![b"1".as_slice(), b"23"]);
2981        assert_eq!(
2982            iter_binary_array(&binary_view).unwrap().collect::<Vec<_>>(),
2983            vec![Some(b"1".as_slice()), Some(b"23".as_slice())]
2984        );
2985
2986        let fixed_size = FixedSizeBinaryArray::from(vec![b"abcd", b"efgh"]);
2987        assert_eq!(
2988            iter_binary_array(&fixed_size).unwrap().collect::<Vec<_>>(),
2989            vec![Some(b"abcd".as_slice()), Some(b"efgh".as_slice())]
2990        );
2991    }
2992
2993    #[test]
2994    fn test_iter_binary_array_rejects_non_binary() {
2995        let int_array = Int32Array::from(vec![1, 2, 3]);
2996        let Err(error) = iter_binary_array(&int_array) else {
2997            panic!("expected an error for non-binary array");
2998        };
2999        assert!(error.to_string().contains("Expecting a binary type"));
3000    }
3001}