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