Skip to main content

lance_arrow/
lib.rs

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