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