Skip to main content

lance_arrow/
lib.rs

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