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