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 normalize validity buffers
1021/// Returns None for all-null validity (placeholder structs)
1022fn normalize_validity(
1023    validity: Option<&arrow_buffer::NullBuffer>,
1024) -> Option<&arrow_buffer::NullBuffer> {
1025    validity.filter(|v| v.null_count() != v.len())
1026}
1027
1028/// Helper function to merge validity buffers from two struct arrays
1029/// Returns None only if both arrays are null at the same position
1030///
1031/// Special handling for placeholder structs (all-null validity)
1032fn merge_struct_validity(
1033    left_validity: Option<&arrow_buffer::NullBuffer>,
1034    right_validity: Option<&arrow_buffer::NullBuffer>,
1035) -> Option<arrow_buffer::NullBuffer> {
1036    // Normalize both validity buffers (convert all-null to None)
1037    let left_normalized = normalize_validity(left_validity);
1038    let right_normalized = normalize_validity(right_validity);
1039
1040    match (left_normalized, right_normalized) {
1041        // Fast paths: no computation needed
1042        (None, None) => None,
1043        (Some(left), None) => Some(left.clone()),
1044        (None, Some(right)) => Some(right.clone()),
1045        (Some(left), Some(right)) => {
1046            // Fast path: if both have no nulls, can return either one
1047            if left.null_count() == 0 && right.null_count() == 0 {
1048                return Some(left.clone());
1049            }
1050
1051            let left_buffer = left.inner();
1052            let right_buffer = right.inner();
1053
1054            // Perform bitwise OR directly on BooleanBuffers
1055            // This preserves the correct semantics: 1 = valid, 0 = null
1056            let merged_buffer = left_buffer | right_buffer;
1057
1058            Some(arrow_buffer::NullBuffer::from(merged_buffer))
1059        }
1060    }
1061}
1062
1063fn merge_list_child_values(
1064    child_field: &Field,
1065    left_values: ArrayRef,
1066    right_values: ArrayRef,
1067) -> ArrayRef {
1068    match child_field.data_type() {
1069        DataType::Struct(child_fields) => Arc::new(merge_with_schema(
1070            left_values.as_struct(),
1071            right_values.as_struct(),
1072            child_fields,
1073        )) as ArrayRef,
1074        DataType::List(grandchild) => {
1075            let left_list = left_values
1076                .as_any()
1077                .downcast_ref::<ListArray>()
1078                .expect("left list values should be ListArray");
1079            let right_list = right_values
1080                .as_any()
1081                .downcast_ref::<ListArray>()
1082                .expect("right list values should be ListArray");
1083            let merged_values = merge_list_child_values(
1084                grandchild.as_ref(),
1085                left_list.values().clone(),
1086                right_list.values().clone(),
1087            );
1088            let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1089            Arc::new(ListArray::new(
1090                grandchild.clone(),
1091                left_list.offsets().clone(),
1092                merged_values,
1093                merged_validity,
1094            )) as ArrayRef
1095        }
1096        DataType::LargeList(grandchild) => {
1097            let left_list = left_values
1098                .as_any()
1099                .downcast_ref::<LargeListArray>()
1100                .expect("left list values should be LargeListArray");
1101            let right_list = right_values
1102                .as_any()
1103                .downcast_ref::<LargeListArray>()
1104                .expect("right list values should be LargeListArray");
1105            let merged_values = merge_list_child_values(
1106                grandchild.as_ref(),
1107                left_list.values().clone(),
1108                right_list.values().clone(),
1109            );
1110            let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1111            Arc::new(LargeListArray::new(
1112                grandchild.clone(),
1113                left_list.offsets().clone(),
1114                merged_values,
1115                merged_validity,
1116            )) as ArrayRef
1117        }
1118        DataType::FixedSizeList(grandchild, list_size) => {
1119            let left_list = left_values
1120                .as_any()
1121                .downcast_ref::<FixedSizeListArray>()
1122                .expect("left list values should be FixedSizeListArray");
1123            let right_list = right_values
1124                .as_any()
1125                .downcast_ref::<FixedSizeListArray>()
1126                .expect("right list values should be FixedSizeListArray");
1127            let merged_values = merge_list_child_values(
1128                grandchild.as_ref(),
1129                left_list.values().clone(),
1130                right_list.values().clone(),
1131            );
1132            let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1133            Arc::new(FixedSizeListArray::new(
1134                grandchild.clone(),
1135                *list_size,
1136                merged_values,
1137                merged_validity,
1138            )) as ArrayRef
1139        }
1140        _ => left_values.clone(),
1141    }
1142}
1143
1144// Helper function to adjust child array validity based on parent struct validity
1145// When parent struct is null, propagates null to child array
1146// Optimized with fast paths and SIMD operations
1147fn adjust_child_validity(
1148    child: &ArrayRef,
1149    parent_validity: Option<&arrow_buffer::NullBuffer>,
1150) -> ArrayRef {
1151    // Fast path: no parent validity means no adjustment needed
1152    let parent_validity = match parent_validity {
1153        None => return child.clone(),
1154        Some(p) if p.null_count() == 0 => return child.clone(), // No nulls to propagate
1155        Some(p) => p,
1156    };
1157
1158    // Fast path: DataType::Null arrays are always entirely null by definition and cannot
1159    // carry an explicit null bitmap (Arrow rejects it). No adjustment is needed.
1160    if child.data_type() == &DataType::Null {
1161        return child.clone();
1162    }
1163
1164    let child_validity = child.nulls();
1165
1166    // Compute the new validity: child_validity AND parent_validity
1167    let new_validity = match child_validity {
1168        None => {
1169            // Fast path: child has no existing validity, just use parent's
1170            parent_validity.clone()
1171        }
1172        Some(child_nulls) => {
1173            let child_buffer = child_nulls.inner();
1174            let parent_buffer = parent_validity.inner();
1175
1176            // Perform bitwise AND directly on BooleanBuffers
1177            // This preserves the correct semantics: 1 = valid, 0 = null
1178            let merged_buffer = child_buffer & parent_buffer;
1179
1180            arrow_buffer::NullBuffer::from(merged_buffer)
1181        }
1182    };
1183
1184    // Create new array with adjusted validity
1185    arrow_array::make_array(
1186        arrow_data::ArrayData::try_new(
1187            child.data_type().clone(),
1188            child.len(),
1189            Some(new_validity.into_inner().into_inner()),
1190            child.offset(),
1191            child.to_data().buffers().to_vec(),
1192            child.to_data().child_data().to_vec(),
1193        )
1194        .unwrap(),
1195    )
1196}
1197
1198fn merge(left_struct_array: &StructArray, right_struct_array: &StructArray) -> StructArray {
1199    let mut fields: Vec<Field> = vec![];
1200    let mut columns: Vec<ArrayRef> = vec![];
1201    let right_fields = right_struct_array.fields();
1202    let right_columns = right_struct_array.columns();
1203
1204    // Get the validity buffers from both structs
1205    let left_validity = left_struct_array.nulls();
1206    let right_validity = right_struct_array.nulls();
1207
1208    // Compute merged validity
1209    let merged_validity = merge_struct_validity(left_validity, right_validity);
1210
1211    // iterate through the fields on the left hand side
1212    for (left_field, left_column) in left_struct_array
1213        .fields()
1214        .iter()
1215        .zip(left_struct_array.columns().iter())
1216    {
1217        match right_fields
1218            .iter()
1219            .position(|f| f.name() == left_field.name())
1220        {
1221            // if the field exists on the right hand side, merge them recursively if appropriate
1222            Some(right_index) => {
1223                let right_field = right_fields.get(right_index).unwrap();
1224                let right_column = right_columns.get(right_index).unwrap();
1225                // if both fields are struct, merge them recursively
1226                match (left_field.data_type(), right_field.data_type()) {
1227                    (DataType::Struct(_), DataType::Struct(_)) => {
1228                        let left_sub_array = left_column.as_struct();
1229                        let right_sub_array = right_column.as_struct();
1230                        let merged_sub_array = merge(left_sub_array, right_sub_array);
1231                        fields.push(Field::new(
1232                            left_field.name(),
1233                            merged_sub_array.data_type().clone(),
1234                            left_field.is_nullable(),
1235                        ));
1236                        columns.push(Arc::new(merged_sub_array) as ArrayRef);
1237                    }
1238                    (DataType::List(left_list), DataType::List(right_list))
1239                        if left_list.data_type().is_struct()
1240                            && right_list.data_type().is_struct() =>
1241                    {
1242                        // If there is nothing to merge just use the left field
1243                        if left_list.data_type() == right_list.data_type() {
1244                            fields.push(left_field.as_ref().clone());
1245                            columns.push(left_column.clone());
1246                        }
1247                        // If we have two List<Struct> and they have different sets of fields then
1248                        // we can merge them if the offsets arrays are the same.  Otherwise, we
1249                        // have to consider it an error.
1250                        let merged_sub_array = merge_list_struct(&left_column, &right_column);
1251
1252                        fields.push(Field::new(
1253                            left_field.name(),
1254                            merged_sub_array.data_type().clone(),
1255                            left_field.is_nullable(),
1256                        ));
1257                        columns.push(merged_sub_array);
1258                    }
1259                    // otherwise, just use the field on the left hand side
1260                    _ => {
1261                        // TODO handle list-of-struct and other types
1262                        fields.push(left_field.as_ref().clone());
1263                        // Adjust the column validity: if left struct was null, propagate to child
1264                        let adjusted_column = adjust_child_validity(left_column, left_validity);
1265                        columns.push(adjusted_column);
1266                    }
1267                }
1268            }
1269            None => {
1270                fields.push(left_field.as_ref().clone());
1271                // Adjust the column validity: if left struct was null, propagate to child
1272                let adjusted_column = adjust_child_validity(left_column, left_validity);
1273                columns.push(adjusted_column);
1274            }
1275        }
1276    }
1277
1278    // now iterate through the fields on the right hand side
1279    right_fields
1280        .iter()
1281        .zip(right_columns.iter())
1282        .for_each(|(field, column)| {
1283            // add new columns on the right
1284            if !left_struct_array
1285                .fields()
1286                .iter()
1287                .any(|f| f.name() == field.name())
1288            {
1289                fields.push(field.as_ref().clone());
1290                // This field doesn't exist on the left
1291                // We use the right's column but need to adjust for struct validity
1292                let adjusted_column = adjust_child_validity(column, right_validity);
1293                columns.push(adjusted_column);
1294            }
1295        });
1296
1297    StructArray::try_new(Fields::from(fields), columns, merged_validity).unwrap()
1298}
1299
1300fn merge_with_schema(
1301    left_struct_array: &StructArray,
1302    right_struct_array: &StructArray,
1303    fields: &Fields,
1304) -> StructArray {
1305    // Helper function that returns true if both types are struct or both are non-struct
1306    fn same_type_kind(left: &DataType, right: &DataType) -> bool {
1307        match (left, right) {
1308            (DataType::Struct(_), DataType::Struct(_)) => true,
1309            (DataType::Struct(_), _) => false,
1310            (_, DataType::Struct(_)) => false,
1311            _ => true,
1312        }
1313    }
1314
1315    let mut output_fields: Vec<Field> = Vec::with_capacity(fields.len());
1316    let mut columns: Vec<ArrayRef> = Vec::with_capacity(fields.len());
1317
1318    let left_fields = left_struct_array.fields();
1319    let left_columns = left_struct_array.columns();
1320    let right_fields = right_struct_array.fields();
1321    let right_columns = right_struct_array.columns();
1322
1323    // Get the validity buffers from both structs
1324    let left_validity = left_struct_array.nulls();
1325    let right_validity = right_struct_array.nulls();
1326
1327    // Compute merged validity
1328    let merged_validity = merge_struct_validity(left_validity, right_validity);
1329
1330    for field in fields {
1331        let left_match_idx = left_fields.iter().position(|f| {
1332            f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1333        });
1334        let right_match_idx = right_fields.iter().position(|f| {
1335            f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1336        });
1337
1338        match (left_match_idx, right_match_idx) {
1339            (None, Some(right_idx)) => {
1340                output_fields.push(right_fields[right_idx].as_ref().clone());
1341                // Adjust validity if the right struct was null
1342                let adjusted_column =
1343                    adjust_child_validity(&right_columns[right_idx], right_validity);
1344                columns.push(adjusted_column);
1345            }
1346            (Some(left_idx), None) => {
1347                output_fields.push(left_fields[left_idx].as_ref().clone());
1348                // Adjust validity if the left struct was null
1349                let adjusted_column = adjust_child_validity(&left_columns[left_idx], left_validity);
1350                columns.push(adjusted_column);
1351            }
1352            (Some(left_idx), Some(right_idx)) => {
1353                match field.data_type() {
1354                    DataType::Struct(child_fields) => {
1355                        let left_sub_array = left_columns[left_idx].as_struct();
1356                        let right_sub_array = right_columns[right_idx].as_struct();
1357                        let merged_sub_array =
1358                            merge_with_schema(left_sub_array, right_sub_array, child_fields);
1359                        output_fields.push(Field::new(
1360                            field.name(),
1361                            merged_sub_array.data_type().clone(),
1362                            field.is_nullable(),
1363                        ));
1364                        columns.push(Arc::new(merged_sub_array) as ArrayRef);
1365                    }
1366                    DataType::List(child_field) => {
1367                        let left_list = left_columns[left_idx]
1368                            .as_any()
1369                            .downcast_ref::<ListArray>()
1370                            .unwrap();
1371                        let right_list = right_columns[right_idx]
1372                            .as_any()
1373                            .downcast_ref::<ListArray>()
1374                            .unwrap();
1375                        let merged_values = merge_list_child_values(
1376                            child_field.as_ref(),
1377                            left_list.trimmed_values(),
1378                            right_list.trimmed_values(),
1379                        );
1380                        let merged_validity =
1381                            merge_struct_validity(left_list.nulls(), right_list.nulls());
1382                        // `trimmed_values` starts at the first used value, so offsets
1383                        // must be shifted to match or `ListArray::new` panics when the
1384                        // input list was sliced (e.g. from a filtered batch).
1385                        let merged_list = ListArray::new(
1386                            child_field.clone(),
1387                            left_list.trimmed_offsets(),
1388                            merged_values,
1389                            merged_validity,
1390                        );
1391                        output_fields.push(field.as_ref().clone());
1392                        columns.push(Arc::new(merged_list) as ArrayRef);
1393                    }
1394                    DataType::LargeList(child_field) => {
1395                        let left_list = left_columns[left_idx]
1396                            .as_any()
1397                            .downcast_ref::<LargeListArray>()
1398                            .unwrap();
1399                        let right_list = right_columns[right_idx]
1400                            .as_any()
1401                            .downcast_ref::<LargeListArray>()
1402                            .unwrap();
1403                        let merged_values = merge_list_child_values(
1404                            child_field.as_ref(),
1405                            left_list.trimmed_values(),
1406                            right_list.trimmed_values(),
1407                        );
1408                        let merged_validity =
1409                            merge_struct_validity(left_list.nulls(), right_list.nulls());
1410                        let merged_list = LargeListArray::new(
1411                            child_field.clone(),
1412                            left_list.trimmed_offsets(),
1413                            merged_values,
1414                            merged_validity,
1415                        );
1416                        output_fields.push(field.as_ref().clone());
1417                        columns.push(Arc::new(merged_list) as ArrayRef);
1418                    }
1419                    DataType::FixedSizeList(child_field, list_size) => {
1420                        let left_list = left_columns[left_idx]
1421                            .as_any()
1422                            .downcast_ref::<FixedSizeListArray>()
1423                            .unwrap();
1424                        let right_list = right_columns[right_idx]
1425                            .as_any()
1426                            .downcast_ref::<FixedSizeListArray>()
1427                            .unwrap();
1428                        let merged_values = merge_list_child_values(
1429                            child_field.as_ref(),
1430                            left_list.values().clone(),
1431                            right_list.values().clone(),
1432                        );
1433                        let merged_validity =
1434                            merge_struct_validity(left_list.nulls(), right_list.nulls());
1435                        let merged_list = FixedSizeListArray::new(
1436                            child_field.clone(),
1437                            *list_size,
1438                            merged_values,
1439                            merged_validity,
1440                        );
1441                        output_fields.push(field.as_ref().clone());
1442                        columns.push(Arc::new(merged_list) as ArrayRef);
1443                    }
1444                    _ => {
1445                        output_fields.push(left_fields[left_idx].as_ref().clone());
1446                        // For fields that exist in both, use left but adjust validity
1447                        let adjusted_column =
1448                            adjust_child_validity(&left_columns[left_idx], left_validity);
1449                        columns.push(adjusted_column);
1450                    }
1451                }
1452            }
1453            (None, None) => {
1454                // The field will not be included in the output
1455            }
1456        }
1457    }
1458
1459    StructArray::try_new(Fields::from(output_fields), columns, merged_validity).unwrap()
1460}
1461
1462fn get_sub_array<'a>(array: &'a ArrayRef, components: &[&str]) -> Option<&'a ArrayRef> {
1463    if components.is_empty() {
1464        return Some(array);
1465    }
1466    if !matches!(array.data_type(), DataType::Struct(_)) {
1467        return None;
1468    }
1469    let struct_arr = array.as_struct();
1470    struct_arr
1471        .column_by_name(components[0])
1472        .and_then(|arr| get_sub_array(arr, &components[1..]))
1473}
1474
1475/// Interleave multiple RecordBatches into a single RecordBatch.
1476///
1477/// Behaves like [`arrow_select::interleave::interleave`], but for RecordBatches.
1478pub fn interleave_batches(
1479    batches: &[RecordBatch],
1480    indices: &[(usize, usize)],
1481) -> Result<RecordBatch> {
1482    let first_batch = batches.first().ok_or_else(|| {
1483        ArrowError::InvalidArgumentError("Cannot interleave zero RecordBatches".to_string())
1484    })?;
1485    let schema = first_batch.schema();
1486    let num_columns = first_batch.num_columns();
1487    let mut columns = Vec::with_capacity(num_columns);
1488    let mut chunks = Vec::with_capacity(batches.len());
1489
1490    for i in 0..num_columns {
1491        for batch in batches {
1492            chunks.push(batch.column(i).as_ref());
1493        }
1494        let new_column = interleave(&chunks, indices)?;
1495        columns.push(new_column);
1496        chunks.clear();
1497    }
1498
1499    RecordBatch::try_new(schema, columns)
1500}
1501
1502pub trait BufferExt {
1503    /// Create an `arrow_buffer::Buffer`` from a `bytes::Bytes` object
1504    ///
1505    /// The alignment must be specified (as `bytes_per_value`) since we want to make
1506    /// sure we can safely reinterpret the buffer.
1507    ///
1508    /// If the buffer is properly aligned this will be zero-copy.  If not, a copy
1509    /// will be made and an owned buffer returned.
1510    ///
1511    /// If `bytes_per_value` is not a power of two, then we assume the buffer is
1512    /// never going to be reinterpreted into another type and we can safely
1513    /// ignore the alignment.
1514    ///
1515    /// Yes, the method name is odd.  It's because there is already a `from_bytes`
1516    /// which converts from `arrow_buffer::bytes::Bytes` (not `bytes::Bytes`)
1517    fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self;
1518
1519    /// Allocates a new properly aligned arrow buffer and copies `bytes` into it
1520    ///
1521    /// `size_bytes` can be larger than `bytes` and, if so, the trailing bytes will
1522    /// be zeroed out.
1523    ///
1524    /// # Panics
1525    ///
1526    /// Panics if `size_bytes` is less than `bytes.len()`
1527    fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self;
1528}
1529
1530fn is_pwr_two(n: u64) -> bool {
1531    n & (n - 1) == 0
1532}
1533
1534impl BufferExt for arrow_buffer::Buffer {
1535    fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self {
1536        if is_pwr_two(bytes_per_value) && bytes.as_ptr().align_offset(bytes_per_value as usize) != 0
1537        {
1538            // The original buffer is not aligned, cannot zero-copy
1539            let size_bytes = bytes.len();
1540            Self::copy_bytes_bytes(bytes, size_bytes)
1541        } else {
1542            // The original buffer is aligned, can zero-copy
1543            // SAFETY: the alignment is correct we can make this conversion
1544            unsafe {
1545                Self::from_custom_allocation(
1546                    NonNull::new(bytes.as_ptr() as _).expect("should be a valid pointer"),
1547                    bytes.len(),
1548                    Arc::new(bytes),
1549                )
1550            }
1551        }
1552    }
1553
1554    fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self {
1555        assert!(size_bytes >= bytes.len());
1556        let mut buf = MutableBuffer::with_capacity(size_bytes);
1557        let to_fill = size_bytes - bytes.len();
1558        buf.extend(bytes);
1559        buf.extend(std::iter::repeat_n(0_u8, to_fill));
1560
1561        // FIX for issue #4512: Shrink buffer to actual size before converting to immutable
1562        // This reduces memory overhead from capacity over-allocation
1563        buf.shrink_to_fit();
1564
1565        Self::from(buf)
1566    }
1567}
1568
1569#[cfg(test)]
1570mod tests {
1571    use super::*;
1572    use arrow_array::{Float32Array, Int32Array, NullArray, StructArray};
1573    use arrow_array::{ListArray, StringArray, new_empty_array, new_null_array};
1574    use arrow_buffer::OffsetBuffer;
1575
1576    #[test]
1577    fn test_convert_to_floating_point_preserves_inner_nulls() {
1578        // A FixedSizeList<Int8> with a null inner element must convert to a
1579        // FixedSizeList<Float32> with the null kept in place. Dropping it would
1580        // shorten the values array and shift every later element (and, when the
1581        // remaining count is not a multiple of the list size, panic).
1582        let values = Int8Array::from(vec![Some(1), None, Some(3), Some(4)]);
1583        let fsl = FixedSizeListArray::new(
1584            Arc::new(Field::new("item", DataType::Int8, true)),
1585            2,
1586            Arc::new(values),
1587            None,
1588        );
1589
1590        let converted = fsl.convert_to_floating_point().unwrap();
1591
1592        assert_eq!(converted.len(), 2);
1593        let conv_values = converted
1594            .values()
1595            .as_any()
1596            .downcast_ref::<Float32Array>()
1597            .unwrap();
1598        assert_eq!(conv_values.len(), 4);
1599        assert_eq!(conv_values.value(0), 1.0);
1600        assert!(conv_values.is_null(1));
1601        assert_eq!(conv_values.value(2), 3.0);
1602        assert_eq!(conv_values.value(3), 4.0);
1603    }
1604
1605    #[test]
1606    fn test_convert_to_floating_point_preserves_inner_nulls_f64_arm() {
1607        // The Float64-producing arms (Int64/UInt8/UInt32) share the same fix as the
1608        // Float32 arms; cover one representative (UInt8 -> Float64) so both branch
1609        // families are exercised.
1610        let values = UInt8Array::from(vec![Some(10u8), None, Some(30), Some(40)]);
1611        let fsl = FixedSizeListArray::new(
1612            Arc::new(Field::new("item", DataType::UInt8, true)),
1613            2,
1614            Arc::new(values),
1615            None,
1616        );
1617
1618        let converted = fsl.convert_to_floating_point().unwrap();
1619
1620        assert_eq!(converted.len(), 2);
1621        let conv_values = converted
1622            .values()
1623            .as_any()
1624            .downcast_ref::<Float64Array>()
1625            .unwrap();
1626        assert_eq!(conv_values.len(), 4);
1627        assert_eq!(conv_values.value(0), 10.0);
1628        assert!(conv_values.is_null(1));
1629        assert_eq!(conv_values.value(2), 30.0);
1630        assert_eq!(conv_values.value(3), 40.0);
1631    }
1632
1633    #[test]
1634    fn test_merge_recursive() {
1635        let a_array = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
1636        let e_array = Int32Array::from(vec![Some(4), Some(5), Some(6)]);
1637        let c_array = Int32Array::from(vec![Some(7), Some(8), Some(9)]);
1638        let d_array = StringArray::from(vec![Some("a"), Some("b"), Some("c")]);
1639
1640        let left_schema = Schema::new(vec![
1641            Field::new("a", DataType::Int32, true),
1642            Field::new(
1643                "b",
1644                DataType::Struct(vec![Field::new("c", DataType::Int32, true)].into()),
1645                true,
1646            ),
1647        ]);
1648        let left_batch = RecordBatch::try_new(
1649            Arc::new(left_schema),
1650            vec![
1651                Arc::new(a_array.clone()),
1652                Arc::new(StructArray::from(vec![(
1653                    Arc::new(Field::new("c", DataType::Int32, true)),
1654                    Arc::new(c_array.clone()) as ArrayRef,
1655                )])),
1656            ],
1657        )
1658        .unwrap();
1659
1660        let right_schema = Schema::new(vec![
1661            Field::new("e", DataType::Int32, true),
1662            Field::new(
1663                "b",
1664                DataType::Struct(vec![Field::new("d", DataType::Utf8, true)].into()),
1665                true,
1666            ),
1667        ]);
1668        let right_batch = RecordBatch::try_new(
1669            Arc::new(right_schema),
1670            vec![
1671                Arc::new(e_array.clone()),
1672                Arc::new(StructArray::from(vec![(
1673                    Arc::new(Field::new("d", DataType::Utf8, true)),
1674                    Arc::new(d_array.clone()) as ArrayRef,
1675                )])) as ArrayRef,
1676            ],
1677        )
1678        .unwrap();
1679
1680        let merged_schema = Schema::new(vec![
1681            Field::new("a", DataType::Int32, true),
1682            Field::new(
1683                "b",
1684                DataType::Struct(
1685                    vec![
1686                        Field::new("c", DataType::Int32, true),
1687                        Field::new("d", DataType::Utf8, true),
1688                    ]
1689                    .into(),
1690                ),
1691                true,
1692            ),
1693            Field::new("e", DataType::Int32, true),
1694        ]);
1695        let merged_batch = RecordBatch::try_new(
1696            Arc::new(merged_schema),
1697            vec![
1698                Arc::new(a_array) as ArrayRef,
1699                Arc::new(StructArray::from(vec![
1700                    (
1701                        Arc::new(Field::new("c", DataType::Int32, true)),
1702                        Arc::new(c_array) as ArrayRef,
1703                    ),
1704                    (
1705                        Arc::new(Field::new("d", DataType::Utf8, true)),
1706                        Arc::new(d_array) as ArrayRef,
1707                    ),
1708                ])) as ArrayRef,
1709                Arc::new(e_array) as ArrayRef,
1710            ],
1711        )
1712        .unwrap();
1713
1714        let result = left_batch.merge(&right_batch).unwrap();
1715        assert_eq!(result, merged_batch);
1716    }
1717
1718    #[test]
1719    fn test_merge_with_schema() {
1720        fn test_batch(names: &[&str], types: &[DataType]) -> (Schema, RecordBatch) {
1721            let fields: Fields = names
1722                .iter()
1723                .zip(types)
1724                .map(|(name, ty)| Field::new(name.to_string(), ty.clone(), false))
1725                .collect();
1726            let schema = Schema::new(vec![Field::new(
1727                "struct",
1728                DataType::Struct(fields.clone()),
1729                false,
1730            )]);
1731            let children = types.iter().map(new_empty_array).collect::<Vec<_>>();
1732            let batch = RecordBatch::try_new(
1733                Arc::new(schema.clone()),
1734                vec![Arc::new(StructArray::new(fields, children, None)) as ArrayRef],
1735            );
1736            (schema, batch.unwrap())
1737        }
1738
1739        let (_, left_batch) = test_batch(&["a", "b"], &[DataType::Int32, DataType::Int64]);
1740        let (_, right_batch) = test_batch(&["c", "b"], &[DataType::Int32, DataType::Int64]);
1741        let (output_schema, _) = test_batch(
1742            &["b", "a", "c"],
1743            &[DataType::Int64, DataType::Int32, DataType::Int32],
1744        );
1745
1746        // If we use merge_with_schema the schema is respected
1747        let merged = left_batch
1748            .merge_with_schema(&right_batch, &output_schema)
1749            .unwrap();
1750        assert_eq!(merged.schema().as_ref(), &output_schema);
1751
1752        // If we use merge we get first-come first-serve based on the left batch
1753        let (naive_schema, _) = test_batch(
1754            &["a", "b", "c"],
1755            &[DataType::Int32, DataType::Int64, DataType::Int32],
1756        );
1757        let merged = left_batch.merge(&right_batch).unwrap();
1758        assert_eq!(merged.schema().as_ref(), &naive_schema);
1759    }
1760
1761    #[test]
1762    fn test_merge_list_struct() {
1763        let x_field = Arc::new(Field::new("x", DataType::Int32, true));
1764        let y_field = Arc::new(Field::new("y", DataType::Int32, true));
1765        let x_struct_field = Arc::new(Field::new(
1766            "item",
1767            DataType::Struct(Fields::from(vec![x_field.clone()])),
1768            true,
1769        ));
1770        let y_struct_field = Arc::new(Field::new(
1771            "item",
1772            DataType::Struct(Fields::from(vec![y_field.clone()])),
1773            true,
1774        ));
1775        let both_struct_field = Arc::new(Field::new(
1776            "item",
1777            DataType::Struct(Fields::from(vec![x_field.clone(), y_field.clone()])),
1778            true,
1779        ));
1780        let left_schema = Schema::new(vec![Field::new(
1781            "list_struct",
1782            DataType::List(x_struct_field.clone()),
1783            true,
1784        )]);
1785        let right_schema = Schema::new(vec![Field::new(
1786            "list_struct",
1787            DataType::List(y_struct_field.clone()),
1788            true,
1789        )]);
1790        let both_schema = Schema::new(vec![Field::new(
1791            "list_struct",
1792            DataType::List(both_struct_field.clone()),
1793            true,
1794        )]);
1795
1796        let x = Arc::new(Int32Array::from(vec![1]));
1797        let y = Arc::new(Int32Array::from(vec![2]));
1798        let x_struct = Arc::new(StructArray::new(
1799            Fields::from(vec![x_field.clone()]),
1800            vec![x.clone()],
1801            None,
1802        ));
1803        let y_struct = Arc::new(StructArray::new(
1804            Fields::from(vec![y_field.clone()]),
1805            vec![y.clone()],
1806            None,
1807        ));
1808        let both_struct = Arc::new(StructArray::new(
1809            Fields::from(vec![x_field.clone(), y_field.clone()]),
1810            vec![x.clone(), y],
1811            None,
1812        ));
1813        let both_null_struct = Arc::new(StructArray::new(
1814            Fields::from(vec![x_field, y_field]),
1815            vec![x, Arc::new(new_null_array(&DataType::Int32, 1))],
1816            None,
1817        ));
1818        let offsets = OffsetBuffer::from_lengths([1]);
1819        let x_s_list = ListArray::new(x_struct_field, offsets.clone(), x_struct, None);
1820        let y_s_list = ListArray::new(y_struct_field, offsets.clone(), y_struct, None);
1821        let both_list = ListArray::new(
1822            both_struct_field.clone(),
1823            offsets.clone(),
1824            both_struct,
1825            None,
1826        );
1827        let both_null_list = ListArray::new(both_struct_field, offsets, both_null_struct, None);
1828        let x_batch =
1829            RecordBatch::try_new(Arc::new(left_schema), vec![Arc::new(x_s_list)]).unwrap();
1830        let y_batch = RecordBatch::try_new(
1831            Arc::new(right_schema.clone()),
1832            vec![Arc::new(y_s_list.clone())],
1833        )
1834        .unwrap();
1835        let merged = x_batch.merge(&y_batch).unwrap();
1836        let expected =
1837            RecordBatch::try_new(Arc::new(both_schema.clone()), vec![Arc::new(both_list)]).unwrap();
1838        assert_eq!(merged, expected);
1839
1840        let y_null_list = new_null_array(y_s_list.data_type(), 1);
1841        let y_null_batch =
1842            RecordBatch::try_new(Arc::new(right_schema), vec![Arc::new(y_null_list.clone())])
1843                .unwrap();
1844        let expected =
1845            RecordBatch::try_new(Arc::new(both_schema), vec![Arc::new(both_null_list)]).unwrap();
1846        let merged = x_batch.merge(&y_null_batch).unwrap();
1847        assert_eq!(merged, expected);
1848    }
1849
1850    #[test]
1851    fn test_byte_width_opt() {
1852        assert_eq!(DataType::Int32.byte_width_opt(), Some(4));
1853        assert_eq!(DataType::Int64.byte_width_opt(), Some(8));
1854        assert_eq!(DataType::Float32.byte_width_opt(), Some(4));
1855        assert_eq!(DataType::Float64.byte_width_opt(), Some(8));
1856        assert_eq!(DataType::Utf8.byte_width_opt(), None);
1857        assert_eq!(DataType::Binary.byte_width_opt(), None);
1858        assert_eq!(
1859            DataType::List(Arc::new(Field::new("item", DataType::Int32, true))).byte_width_opt(),
1860            None
1861        );
1862        assert_eq!(
1863            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 3)
1864                .byte_width_opt(),
1865            Some(12)
1866        );
1867        assert_eq!(
1868            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 4)
1869                .byte_width_opt(),
1870            Some(16)
1871        );
1872        assert_eq!(
1873            DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Utf8, true)), 5)
1874                .byte_width_opt(),
1875            None
1876        );
1877    }
1878
1879    #[test]
1880    fn test_take_record_batch() {
1881        let schema = Arc::new(Schema::new(vec![
1882            Field::new("a", DataType::Int32, true),
1883            Field::new("b", DataType::Utf8, true),
1884        ]));
1885        let batch = RecordBatch::try_new(
1886            schema.clone(),
1887            vec![
1888                Arc::new(Int32Array::from_iter_values(0..20)),
1889                Arc::new(StringArray::from_iter_values(
1890                    (0..20).map(|i| format!("str-{}", i)),
1891                )),
1892            ],
1893        )
1894        .unwrap();
1895        let taken = batch.take(&(vec![1_u32, 5_u32, 10_u32].into())).unwrap();
1896        assert_eq!(
1897            taken,
1898            RecordBatch::try_new(
1899                schema,
1900                vec![
1901                    Arc::new(Int32Array::from(vec![1, 5, 10])),
1902                    Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])),
1903                ],
1904            )
1905            .unwrap()
1906        )
1907    }
1908
1909    #[test]
1910    fn test_schema_project_by_schema() {
1911        let metadata = [("key".to_string(), "value".to_string())];
1912        let schema = Arc::new(
1913            Schema::new(vec![
1914                Field::new("a", DataType::Int32, true),
1915                Field::new("b", DataType::Utf8, true),
1916            ])
1917            .with_metadata(metadata.clone().into()),
1918        );
1919        let batch = RecordBatch::try_new(
1920            schema,
1921            vec![
1922                Arc::new(Int32Array::from_iter_values(0..20)),
1923                Arc::new(StringArray::from_iter_values(
1924                    (0..20).map(|i| format!("str-{}", i)),
1925                )),
1926            ],
1927        )
1928        .unwrap();
1929
1930        // Empty schema
1931        let empty_schema = Schema::empty();
1932        let empty_projected = batch.project_by_schema(&empty_schema).unwrap();
1933        let expected_schema = empty_schema.with_metadata(metadata.clone().into());
1934        assert_eq!(
1935            empty_projected,
1936            RecordBatch::from(StructArray::new_empty_fields(batch.num_rows(), None))
1937                .with_schema(Arc::new(expected_schema))
1938                .unwrap()
1939        );
1940
1941        // Re-ordered schema
1942        let reordered_schema = Schema::new(vec![
1943            Field::new("b", DataType::Utf8, true),
1944            Field::new("a", DataType::Int32, true),
1945        ]);
1946        let reordered_projected = batch.project_by_schema(&reordered_schema).unwrap();
1947        let expected_schema = Arc::new(reordered_schema.with_metadata(metadata.clone().into()));
1948        assert_eq!(
1949            reordered_projected,
1950            RecordBatch::try_new(
1951                expected_schema,
1952                vec![
1953                    Arc::new(StringArray::from_iter_values(
1954                        (0..20).map(|i| format!("str-{}", i)),
1955                    )),
1956                    Arc::new(Int32Array::from_iter_values(0..20)),
1957                ],
1958            )
1959            .unwrap()
1960        );
1961
1962        // Sub schema
1963        let sub_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
1964        let sub_projected = batch.project_by_schema(&sub_schema).unwrap();
1965        let expected_schema = Arc::new(sub_schema.with_metadata(metadata.into()));
1966        assert_eq!(
1967            sub_projected,
1968            RecordBatch::try_new(
1969                expected_schema,
1970                vec![Arc::new(Int32Array::from_iter_values(0..20))],
1971            )
1972            .unwrap()
1973        );
1974    }
1975
1976    #[test]
1977    fn test_project_preserves_struct_validity() {
1978        // Test that projecting a struct array preserves its validity (fix for issue #4385)
1979        let fields = Fields::from(vec![
1980            Field::new("id", DataType::Int32, false),
1981            Field::new("value", DataType::Float32, true),
1982        ]);
1983
1984        // Create a struct array with validity
1985        let id_array = Int32Array::from(vec![1, 2, 3]);
1986        let value_array = Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]);
1987        let struct_array = StructArray::new(
1988            fields.clone(),
1989            vec![
1990                Arc::new(id_array) as ArrayRef,
1991                Arc::new(value_array) as ArrayRef,
1992            ],
1993            Some(vec![true, false, true].into()), // Second struct is null
1994        );
1995
1996        // Project the struct array
1997        let projected = project(&struct_array, &fields).unwrap();
1998
1999        // Verify the validity is preserved
2000        assert_eq!(projected.null_count(), 1);
2001        assert!(!projected.is_null(0));
2002        assert!(projected.is_null(1));
2003        assert!(!projected.is_null(2));
2004    }
2005
2006    #[test]
2007    fn test_merge_struct_with_different_validity() {
2008        // Test case from Weston's review comment
2009        // File 1 has height field with some nulls
2010        let height_array = Int32Array::from(vec![Some(500), None, Some(600), None]);
2011        let left_fields = Fields::from(vec![Field::new("height", DataType::Int32, true)]);
2012        let left_struct = StructArray::new(
2013            left_fields,
2014            vec![Arc::new(height_array) as ArrayRef],
2015            Some(vec![true, false, true, false].into()), // Rows 2 and 4 are null structs
2016        );
2017
2018        // File 2 has width field with some nulls
2019        let width_array = Int32Array::from(vec![Some(300), Some(200), None, None]);
2020        let right_fields = Fields::from(vec![Field::new("width", DataType::Int32, true)]);
2021        let right_struct = StructArray::new(
2022            right_fields,
2023            vec![Arc::new(width_array) as ArrayRef],
2024            Some(vec![true, true, false, false].into()), // Rows 3 and 4 are null structs
2025        );
2026
2027        // Merge the two structs
2028        let merged = merge(&left_struct, &right_struct);
2029
2030        // Expected:
2031        // Row 1: both non-null -> {width: 300, height: 500}
2032        // Row 2: left null, right non-null -> {width: 200, height: null}
2033        // Row 3: left non-null, right null -> {width: null, height: 600}
2034        // Row 4: both null -> null struct
2035
2036        assert_eq!(merged.null_count(), 1); // Only row 4 is null
2037        assert!(!merged.is_null(0));
2038        assert!(!merged.is_null(1));
2039        assert!(!merged.is_null(2));
2040        assert!(merged.is_null(3));
2041
2042        // Check field values
2043        let height_col = merged.column_by_name("height").unwrap();
2044        let height_values = height_col.as_any().downcast_ref::<Int32Array>().unwrap();
2045        assert_eq!(height_values.value(0), 500);
2046        assert!(height_values.is_null(1)); // height is null when left struct was null
2047        assert_eq!(height_values.value(2), 600);
2048
2049        let width_col = merged.column_by_name("width").unwrap();
2050        let width_values = width_col.as_any().downcast_ref::<Int32Array>().unwrap();
2051        assert_eq!(width_values.value(0), 300);
2052        assert_eq!(width_values.value(1), 200);
2053        assert!(width_values.is_null(2)); // width is null when right struct was null
2054    }
2055
2056    #[test]
2057    fn test_merge_null_typed_column_with_parent_validity() {
2058        // Reproduces ENT-990: panic in adjust_child_validity when a Null-typed column
2059        // exists on one side and the parent struct has null rows.
2060        // Arrow's Null type has no null bitmap, so passing one to ArrayData::try_new panics.
2061        let left_struct = StructArray::new(
2062            Fields::from(vec![Field::new("a", DataType::Int32, true)]),
2063            vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef],
2064            Some(vec![true, false].into()),
2065        );
2066        let right_struct = StructArray::new(
2067            Fields::from(vec![Field::new("b", DataType::Null, true)]),
2068            vec![Arc::new(NullArray::new(2)) as ArrayRef],
2069            Some(vec![true, false].into()),
2070        );
2071
2072        // Previously panicked: "Arrays of type Null cannot contain a null bitmask"
2073        let merged = merge(&left_struct, &right_struct);
2074        assert_eq!(merged.len(), 2);
2075        let b_col = merged.column_by_name("b").unwrap();
2076        // DataType::Null implies all-null by definition; no null bitmap is stored.
2077        assert_eq!(b_col.data_type(), &DataType::Null);
2078        assert_eq!(b_col.len(), 2);
2079    }
2080
2081    #[test]
2082    fn test_merge_with_schema_with_nullable_struct_list_schema_mismatch() {
2083        // left_list setup
2084        let left_company_id = Arc::new(Int32Array::from(vec![None, None]));
2085        let left_count = Arc::new(Int32Array::from(vec![None, None]));
2086        let left_struct = Arc::new(StructArray::new(
2087            Fields::from(vec![
2088                Field::new("company_id", DataType::Int32, true),
2089                Field::new("count", DataType::Int32, true),
2090            ]),
2091            vec![left_company_id, left_count],
2092            None,
2093        ));
2094        let left_list = Arc::new(ListArray::new(
2095            Arc::new(Field::new(
2096                "item",
2097                DataType::Struct(left_struct.fields().clone()),
2098                true,
2099            )),
2100            OffsetBuffer::from_lengths([2]),
2101            left_struct,
2102            None,
2103        ));
2104
2105        // Right List Setup
2106        let right_company_name = Arc::new(StringArray::from(vec!["Google", "Microsoft"]));
2107        let right_struct = Arc::new(StructArray::new(
2108            Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2109            vec![right_company_name],
2110            None,
2111        ));
2112        let right_list = Arc::new(ListArray::new(
2113            Arc::new(Field::new(
2114                "item",
2115                DataType::Struct(right_struct.fields().clone()),
2116                true,
2117            )),
2118            OffsetBuffer::from_lengths([2]),
2119            right_struct,
2120            None,
2121        ));
2122
2123        let target_fields = Fields::from(vec![Field::new(
2124            "companies",
2125            DataType::List(Arc::new(Field::new(
2126                "item",
2127                DataType::Struct(Fields::from(vec![
2128                    Field::new("company_id", DataType::Int32, true),
2129                    Field::new("company_name", DataType::Utf8, true),
2130                    Field::new("count", DataType::Int32, true),
2131                ])),
2132                true,
2133            ))),
2134            true,
2135        )]);
2136
2137        let left_batch = RecordBatch::try_new(
2138            Arc::new(Schema::new(vec![Field::new(
2139                "companies",
2140                left_list.data_type().clone(),
2141                true,
2142            )])),
2143            vec![left_list as ArrayRef],
2144        )
2145        .unwrap();
2146
2147        let right_batch = RecordBatch::try_new(
2148            Arc::new(Schema::new(vec![Field::new(
2149                "companies",
2150                right_list.data_type().clone(),
2151                true,
2152            )])),
2153            vec![right_list as ArrayRef],
2154        )
2155        .unwrap();
2156
2157        let merged = left_batch
2158            .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2159            .unwrap();
2160
2161        // Verify the merged structure
2162        let merged_list = merged
2163            .column_by_name("companies")
2164            .unwrap()
2165            .as_any()
2166            .downcast_ref::<ListArray>()
2167            .unwrap();
2168        let merged_struct = merged_list.values().as_struct();
2169
2170        // Should have all 3 fields
2171        assert_eq!(merged_struct.num_columns(), 3);
2172        assert!(merged_struct.column_by_name("company_id").is_some());
2173        assert!(merged_struct.column_by_name("company_name").is_some());
2174        assert!(merged_struct.column_by_name("count").is_some());
2175
2176        // Verify values
2177        let company_id = merged_struct
2178            .column_by_name("company_id")
2179            .unwrap()
2180            .as_any()
2181            .downcast_ref::<Int32Array>()
2182            .unwrap();
2183        assert!(company_id.is_null(0));
2184        assert!(company_id.is_null(1));
2185
2186        let company_name = merged_struct
2187            .column_by_name("company_name")
2188            .unwrap()
2189            .as_any()
2190            .downcast_ref::<StringArray>()
2191            .unwrap();
2192        assert_eq!(company_name.value(0), "Google");
2193        assert_eq!(company_name.value(1), "Microsoft");
2194
2195        let count = merged_struct
2196            .column_by_name("count")
2197            .unwrap()
2198            .as_any()
2199            .downcast_ref::<Int32Array>()
2200            .unwrap();
2201        assert!(count.is_null(0));
2202        assert!(count.is_null(1));
2203    }
2204
2205    #[test]
2206    fn test_merge_struct_lists() {
2207        test_merge_struct_lists_generic::<i32>();
2208    }
2209
2210    #[test]
2211    fn test_merge_struct_large_lists() {
2212        test_merge_struct_lists_generic::<i64>();
2213    }
2214
2215    fn test_merge_struct_lists_generic<O: OffsetSizeTrait>() {
2216        // left_list setup
2217        let left_company_id = Arc::new(Int32Array::from(vec![
2218            Some(1),
2219            Some(2),
2220            Some(3),
2221            Some(4),
2222            Some(5),
2223            Some(6),
2224            Some(7),
2225            Some(8),
2226            Some(9),
2227            Some(10),
2228            Some(11),
2229            Some(12),
2230            Some(13),
2231            Some(14),
2232            Some(15),
2233            Some(16),
2234            Some(17),
2235            Some(18),
2236            Some(19),
2237            Some(20),
2238        ]));
2239        let left_count = Arc::new(Int32Array::from(vec![
2240            Some(10),
2241            Some(20),
2242            Some(30),
2243            Some(40),
2244            Some(50),
2245            Some(60),
2246            Some(70),
2247            Some(80),
2248            Some(90),
2249            Some(100),
2250            Some(110),
2251            Some(120),
2252            Some(130),
2253            Some(140),
2254            Some(150),
2255            Some(160),
2256            Some(170),
2257            Some(180),
2258            Some(190),
2259            Some(200),
2260        ]));
2261        let left_struct = Arc::new(StructArray::new(
2262            Fields::from(vec![
2263                Field::new("company_id", DataType::Int32, true),
2264                Field::new("count", DataType::Int32, true),
2265            ]),
2266            vec![left_company_id, left_count],
2267            None,
2268        ));
2269
2270        let left_list = Arc::new(GenericListArray::<O>::new(
2271            Arc::new(Field::new(
2272                "item",
2273                DataType::Struct(left_struct.fields().clone()),
2274                true,
2275            )),
2276            OffsetBuffer::from_lengths([3, 1]),
2277            left_struct.clone(),
2278            None,
2279        ));
2280
2281        let left_list_struct = Arc::new(StructArray::new(
2282            Fields::from(vec![Field::new(
2283                "companies",
2284                if O::IS_LARGE {
2285                    DataType::LargeList(Arc::new(Field::new(
2286                        "item",
2287                        DataType::Struct(left_struct.fields().clone()),
2288                        true,
2289                    )))
2290                } else {
2291                    DataType::List(Arc::new(Field::new(
2292                        "item",
2293                        DataType::Struct(left_struct.fields().clone()),
2294                        true,
2295                    )))
2296                },
2297                true,
2298            )]),
2299            vec![left_list as ArrayRef],
2300            None,
2301        ));
2302
2303        // right_list setup
2304        let right_company_name = Arc::new(StringArray::from(vec![
2305            "Google",
2306            "Microsoft",
2307            "Apple",
2308            "Facebook",
2309        ]));
2310        let right_struct = Arc::new(StructArray::new(
2311            Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2312            vec![right_company_name],
2313            None,
2314        ));
2315        let right_list = Arc::new(GenericListArray::<O>::new(
2316            Arc::new(Field::new(
2317                "item",
2318                DataType::Struct(right_struct.fields().clone()),
2319                true,
2320            )),
2321            OffsetBuffer::from_lengths([3, 1]),
2322            right_struct.clone(),
2323            None,
2324        ));
2325
2326        let right_list_struct = Arc::new(StructArray::new(
2327            Fields::from(vec![Field::new(
2328                "companies",
2329                if O::IS_LARGE {
2330                    DataType::LargeList(Arc::new(Field::new(
2331                        "item",
2332                        DataType::Struct(right_struct.fields().clone()),
2333                        true,
2334                    )))
2335                } else {
2336                    DataType::List(Arc::new(Field::new(
2337                        "item",
2338                        DataType::Struct(right_struct.fields().clone()),
2339                        true,
2340                    )))
2341                },
2342                true,
2343            )]),
2344            vec![right_list as ArrayRef],
2345            None,
2346        ));
2347
2348        // prepare schema
2349        let target_fields = Fields::from(vec![Field::new(
2350            "companies",
2351            if O::IS_LARGE {
2352                DataType::LargeList(Arc::new(Field::new(
2353                    "item",
2354                    DataType::Struct(Fields::from(vec![
2355                        Field::new("company_id", DataType::Int32, true),
2356                        Field::new("company_name", DataType::Utf8, true),
2357                        Field::new("count", DataType::Int32, true),
2358                    ])),
2359                    true,
2360                )))
2361            } else {
2362                DataType::List(Arc::new(Field::new(
2363                    "item",
2364                    DataType::Struct(Fields::from(vec![
2365                        Field::new("company_id", DataType::Int32, true),
2366                        Field::new("company_name", DataType::Utf8, true),
2367                        Field::new("count", DataType::Int32, true),
2368                    ])),
2369                    true,
2370                )))
2371            },
2372            true,
2373        )]);
2374
2375        // merge left_list and right_list
2376        let merged_array = merge_with_schema(&left_list_struct, &right_list_struct, &target_fields);
2377        assert_eq!(merged_array.len(), 2);
2378    }
2379
2380    #[test]
2381    fn test_merge_with_schema_sliced_list_struct() {
2382        test_merge_with_schema_sliced_list_struct_generic::<i32>();
2383    }
2384
2385    #[test]
2386    fn test_merge_with_schema_sliced_large_list_struct() {
2387        test_merge_with_schema_sliced_list_struct_generic::<i64>();
2388    }
2389
2390    // Regression for #6580: merge_with_schema panicked when the left list was a
2391    // sliced view whose offsets did not start at zero (common after a filtered
2392    // scan). Cloning those offsets alongside `trimmed_values` produced offsets
2393    // larger than the trimmed child, panicking in `(Large)ListArray::new`.
2394    fn test_merge_with_schema_sliced_list_struct_generic<O: OffsetSizeTrait>() {
2395        let make_list_dtype = |item_field: Arc<Field>| {
2396            if O::IS_LARGE {
2397                DataType::LargeList(item_field)
2398            } else {
2399                DataType::List(item_field)
2400            }
2401        };
2402
2403        // Build a List<Struct> with two rows of 5 items each, then slice away
2404        // the first row so the remaining list's offsets start at 5, not 0.
2405        let struct_fields_a = Fields::from(vec![Field::new("a", DataType::Int32, true)]);
2406        let left_values = Arc::new(StructArray::new(
2407            struct_fields_a.clone(),
2408            vec![Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef],
2409            None,
2410        ));
2411        let full_list = GenericListArray::<O>::new(
2412            Arc::new(Field::new("item", DataType::Struct(struct_fields_a), true)),
2413            OffsetBuffer::<O>::from_lengths([5, 5]),
2414            left_values,
2415            None,
2416        );
2417        let sliced_left = full_list.slice(1, 1);
2418        assert_eq!(sliced_left.offsets()[0].as_usize(), 5);
2419        assert_eq!(sliced_left.offsets()[1].as_usize(), 10);
2420
2421        let struct_fields_b = Fields::from(vec![Field::new("b", DataType::Int32, true)]);
2422        let right_values = Arc::new(StructArray::new(
2423            struct_fields_b.clone(),
2424            vec![Arc::new(Int32Array::from_iter_values(100..105)) as ArrayRef],
2425            None,
2426        ));
2427        let right_list = GenericListArray::<O>::new(
2428            Arc::new(Field::new("item", DataType::Struct(struct_fields_b), true)),
2429            OffsetBuffer::<O>::from_lengths([5]),
2430            right_values,
2431            None,
2432        );
2433
2434        let target_item_field = Arc::new(Field::new(
2435            "item",
2436            DataType::Struct(Fields::from(vec![
2437                Field::new("a", DataType::Int32, true),
2438                Field::new("b", DataType::Int32, true),
2439            ])),
2440            true,
2441        ));
2442        let target_fields = Fields::from(vec![Field::new(
2443            "items",
2444            make_list_dtype(target_item_field),
2445            true,
2446        )]);
2447
2448        let left_batch = RecordBatch::try_new(
2449            Arc::new(Schema::new(vec![Field::new(
2450                "items",
2451                sliced_left.data_type().clone(),
2452                true,
2453            )])),
2454            vec![Arc::new(sliced_left) as ArrayRef],
2455        )
2456        .unwrap();
2457        let right_batch = RecordBatch::try_new(
2458            Arc::new(Schema::new(vec![Field::new(
2459                "items",
2460                right_list.data_type().clone(),
2461                true,
2462            )])),
2463            vec![Arc::new(right_list) as ArrayRef],
2464        )
2465        .unwrap();
2466
2467        let merged = left_batch
2468            .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2469            .unwrap();
2470
2471        let merged_list = merged
2472            .column_by_name("items")
2473            .unwrap()
2474            .as_any()
2475            .downcast_ref::<GenericListArray<O>>()
2476            .unwrap();
2477        assert_eq!(merged_list.len(), 1);
2478        assert_eq!(merged_list.value_length(0).as_usize(), 5);
2479        let merged_struct = merged_list.values().as_struct();
2480        assert_eq!(merged_struct.num_columns(), 2);
2481        let a = merged_struct
2482            .column_by_name("a")
2483            .unwrap()
2484            .as_any()
2485            .downcast_ref::<Int32Array>()
2486            .unwrap();
2487        // After shifting offsets to zero, values 5..10 should be first.
2488        let a_vals: Vec<i32> = a.iter().map(|v| v.unwrap()).collect();
2489        assert_eq!(a_vals, vec![5, 6, 7, 8, 9]);
2490    }
2491
2492    #[test]
2493    fn test_project_by_schema_list_struct_reorder() {
2494        // Test that project_by_schema correctly reorders fields inside List<Struct>
2495        // This is a regression test for issue #5702
2496
2497        // Source schema with inner struct fields in order: c, b, a
2498        let source_inner_struct = DataType::Struct(Fields::from(vec![
2499            Field::new("c", DataType::Utf8, true),
2500            Field::new("b", DataType::Utf8, true),
2501            Field::new("a", DataType::Utf8, true),
2502        ]));
2503        let source_schema = Arc::new(Schema::new(vec![
2504            Field::new("id", DataType::Int32, false),
2505            Field::new(
2506                "data",
2507                DataType::List(Arc::new(Field::new(
2508                    "item",
2509                    source_inner_struct.clone(),
2510                    true,
2511                ))),
2512                true,
2513            ),
2514        ]));
2515
2516        // Create source data with c, b, a order
2517        let c_array = StringArray::from(vec!["c1", "c2"]);
2518        let b_array = StringArray::from(vec!["b1", "b2"]);
2519        let a_array = StringArray::from(vec!["a1", "a2"]);
2520        let inner_struct = StructArray::from(vec![
2521            (
2522                Arc::new(Field::new("c", DataType::Utf8, true)),
2523                Arc::new(c_array) as ArrayRef,
2524            ),
2525            (
2526                Arc::new(Field::new("b", DataType::Utf8, true)),
2527                Arc::new(b_array) as ArrayRef,
2528            ),
2529            (
2530                Arc::new(Field::new("a", DataType::Utf8, true)),
2531                Arc::new(a_array) as ArrayRef,
2532            ),
2533        ]);
2534
2535        let list_array = ListArray::new(
2536            Arc::new(Field::new("item", source_inner_struct, true)),
2537            OffsetBuffer::from_lengths([1, 1]),
2538            Arc::new(inner_struct),
2539            None,
2540        );
2541
2542        let batch = RecordBatch::try_new(
2543            source_schema,
2544            vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(list_array)],
2545        )
2546        .unwrap();
2547
2548        // Target schema with inner struct fields in order: a, b, c
2549        let target_inner_struct = DataType::Struct(Fields::from(vec![
2550            Field::new("a", DataType::Utf8, true),
2551            Field::new("b", DataType::Utf8, true),
2552            Field::new("c", DataType::Utf8, true),
2553        ]));
2554        let target_schema = Schema::new(vec![
2555            Field::new("id", DataType::Int32, false),
2556            Field::new(
2557                "data",
2558                DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
2559                true,
2560            ),
2561        ]);
2562
2563        // Project should reorder the inner struct fields
2564        let projected = batch.project_by_schema(&target_schema).unwrap();
2565
2566        // Verify the schema is correct
2567        assert_eq!(projected.schema().as_ref(), &target_schema);
2568
2569        // Verify the data is correct by checking inner struct field order
2570        let projected_list = projected.column(1).as_list::<i32>();
2571        let projected_struct = projected_list.values().as_struct();
2572
2573        // Fields should now be in order: a, b, c
2574        assert_eq!(
2575            projected_struct.column_by_name("a").unwrap().as_ref(),
2576            &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2577        );
2578        assert_eq!(
2579            projected_struct.column_by_name("b").unwrap().as_ref(),
2580            &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2581        );
2582        assert_eq!(
2583            projected_struct.column_by_name("c").unwrap().as_ref(),
2584            &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2585        );
2586
2587        // Also verify positional access matches expected order (a=0, b=1, c=2)
2588        assert_eq!(
2589            projected_struct.column(0).as_ref(),
2590            &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2591        );
2592        assert_eq!(
2593            projected_struct.column(1).as_ref(),
2594            &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2595        );
2596        assert_eq!(
2597            projected_struct.column(2).as_ref(),
2598            &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2599        );
2600    }
2601
2602    #[test]
2603    fn test_project_by_schema_nested_list_struct() {
2604        // Test deeply nested List<Struct<List<Struct>>> projection
2605        let inner_struct = DataType::Struct(Fields::from(vec![
2606            Field::new("y", DataType::Int32, true),
2607            Field::new("x", DataType::Int32, true),
2608        ]));
2609        let source_schema = Arc::new(Schema::new(vec![Field::new(
2610            "outer",
2611            DataType::List(Arc::new(Field::new(
2612                "item",
2613                DataType::Struct(Fields::from(vec![
2614                    Field::new("b", DataType::Utf8, true),
2615                    Field::new(
2616                        "inner_list",
2617                        DataType::List(Arc::new(Field::new("item", inner_struct.clone(), true))),
2618                        true,
2619                    ),
2620                    Field::new("a", DataType::Utf8, true),
2621                ])),
2622                true,
2623            ))),
2624            true,
2625        )]));
2626
2627        // Create deeply nested data
2628        let y_array = Int32Array::from(vec![1, 2]);
2629        let x_array = Int32Array::from(vec![3, 4]);
2630        let innermost_struct = StructArray::from(vec![
2631            (
2632                Arc::new(Field::new("y", DataType::Int32, true)),
2633                Arc::new(y_array) as ArrayRef,
2634            ),
2635            (
2636                Arc::new(Field::new("x", DataType::Int32, true)),
2637                Arc::new(x_array) as ArrayRef,
2638            ),
2639        ]);
2640        let inner_list = ListArray::new(
2641            Arc::new(Field::new("item", inner_struct.clone(), true)),
2642            OffsetBuffer::from_lengths([2]),
2643            Arc::new(innermost_struct),
2644            None,
2645        );
2646
2647        let b_array = StringArray::from(vec!["b1"]);
2648        let a_array = StringArray::from(vec!["a1"]);
2649        let middle_struct = StructArray::from(vec![
2650            (
2651                Arc::new(Field::new("b", DataType::Utf8, true)),
2652                Arc::new(b_array) as ArrayRef,
2653            ),
2654            (
2655                Arc::new(Field::new(
2656                    "inner_list",
2657                    DataType::List(Arc::new(Field::new("item", inner_struct, true))),
2658                    true,
2659                )),
2660                Arc::new(inner_list) as ArrayRef,
2661            ),
2662            (
2663                Arc::new(Field::new("a", DataType::Utf8, true)),
2664                Arc::new(a_array) as ArrayRef,
2665            ),
2666        ]);
2667
2668        let outer_list = ListArray::new(
2669            Arc::new(Field::new("item", middle_struct.data_type().clone(), true)),
2670            OffsetBuffer::from_lengths([1]),
2671            Arc::new(middle_struct),
2672            None,
2673        );
2674
2675        let batch =
2676            RecordBatch::try_new(source_schema, vec![Arc::new(outer_list) as ArrayRef]).unwrap();
2677
2678        // Target schema with reordered fields at all levels
2679        let target_inner_struct = DataType::Struct(Fields::from(vec![
2680            Field::new("x", DataType::Int32, true), // x before y now
2681            Field::new("y", DataType::Int32, true),
2682        ]));
2683        let target_schema = Schema::new(vec![Field::new(
2684            "outer",
2685            DataType::List(Arc::new(Field::new(
2686                "item",
2687                DataType::Struct(Fields::from(vec![
2688                    Field::new("a", DataType::Utf8, true), // a before b now
2689                    Field::new(
2690                        "inner_list",
2691                        DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
2692                        true,
2693                    ),
2694                    Field::new("b", DataType::Utf8, true),
2695                ])),
2696                true,
2697            ))),
2698            true,
2699        )]);
2700
2701        let projected = batch.project_by_schema(&target_schema).unwrap();
2702
2703        // Verify schema
2704        assert_eq!(projected.schema().as_ref(), &target_schema);
2705
2706        // Verify deeply nested data is reordered correctly
2707        let outer_list = projected.column(0).as_list::<i32>();
2708        let middle_struct = outer_list.values().as_struct();
2709
2710        // Middle struct should have a first, then inner_list, then b
2711        assert_eq!(
2712            middle_struct.column(0).as_ref(),
2713            &StringArray::from(vec!["a1"]) as &dyn Array
2714        );
2715        assert_eq!(
2716            middle_struct.column(2).as_ref(),
2717            &StringArray::from(vec!["b1"]) as &dyn Array
2718        );
2719
2720        // Inner list's struct should have x first, then y
2721        let inner_list = middle_struct.column(1).as_list::<i32>();
2722        let innermost_struct = inner_list.values().as_struct();
2723        assert_eq!(
2724            innermost_struct.column(0).as_ref(),
2725            &Int32Array::from(vec![3, 4]) as &dyn Array
2726        );
2727        assert_eq!(
2728            innermost_struct.column(1).as_ref(),
2729            &Int32Array::from(vec![1, 2]) as &dyn Array
2730        );
2731    }
2732}