1#![warn(clippy::undocumented_unsafe_blocks)]
9
10#[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
51pub const ARROW_EXT_NAME_KEY: &str = "ARROW:extension:name";
53
54pub const ARROW_EXT_META_KEY: &str = "ARROW:extension:metadata";
56
57pub const BLOB_META_KEY: &str = "lance-encoding:blob";
60pub const BLOB_V2_EXT_NAME: &str = "lance.blob.v2";
62pub const BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY: &str =
64 "lance-encoding:blob-dedicated-size-threshold";
65pub const BLOB_INLINE_SIZE_THRESHOLD_META_KEY: &str = "lance-encoding:blob-inline-size-threshold";
67pub 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 fn is_binary_like(&self) -> bool;
76
77 fn is_struct(&self) -> bool;
79
80 fn is_fixed_stride(&self) -> bool;
85
86 fn is_dictionary(&self) -> bool;
88
89 fn byte_width(&self) -> usize;
92
93 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
185pub 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 fn try_new_from_values<T: Array + 'static>(
255 values: T,
256 list_size: i32,
257 ) -> Result<FixedSizeListArray>;
258
259 fn sample(&self, n: usize) -> Result<FixedSizeListArray>;
273
274 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
427pub fn as_fixed_size_list_array(arr: &dyn Array) -> &FixedSizeListArray {
430 arr.as_any().downcast_ref::<FixedSizeListArray>().unwrap()
431}
432
433pub trait FixedSizeBinaryArrayExt {
434 fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result<FixedSizeBinaryArray>;
449}
450
451impl FixedSizeBinaryArrayExt for FixedSizeBinaryArray {
452 fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result<Self> {
453 let data_type = DataType::FixedSizeBinary(stride);
454 let data = ArrayDataBuilder::new(data_type)
455 .len(values.len() / stride as usize)
456 .add_buffer(values.into_data().buffers()[0].clone())
457 .build()?;
458 Ok(Self::from(data))
459 }
460}
461
462pub fn as_fixed_size_binary_array(arr: &dyn Array) -> &FixedSizeBinaryArray {
463 arr.as_any().downcast_ref::<FixedSizeBinaryArray>().unwrap()
464}
465
466pub fn iter_str_array(arr: &dyn Array) -> Box<dyn Iterator<Item = Option<&str>> + Send + '_> {
467 match arr.data_type() {
468 DataType::Utf8 => Box::new(arr.as_string::<i32>().iter()),
469 DataType::LargeUtf8 => Box::new(arr.as_string::<i64>().iter()),
470 DataType::Utf8View => Box::new(arr.as_string_view().iter()),
471 _ => panic!(
472 "Expecting Utf8, LargeUtf8 or Utf8View, found {:?}",
473 arr.data_type()
474 ),
475 }
476}
477
478pub fn iter_binary_array(
479 arr: &dyn Array,
480) -> Result<Box<dyn Iterator<Item = Option<&[u8]>> + Send + '_>> {
481 match arr.data_type() {
482 DataType::Binary => Ok(Box::new(arr.as_binary::<i32>().iter())),
483 DataType::LargeBinary => Ok(Box::new(arr.as_binary::<i64>().iter())),
484 DataType::BinaryView => Ok(Box::new(arr.as_binary_view().iter())),
485 DataType::FixedSizeBinary(_) => Ok(Box::new(arr.as_fixed_size_binary().iter())),
486 data_type => Err(ArrowError::InvalidArgumentError(format!(
487 "Expecting a binary type, found {data_type}"
488 ))),
489 }
490}
491
492pub trait RecordBatchExt {
494 fn try_with_column(&self, field: Field, arr: ArrayRef) -> Result<RecordBatch>;
524
525 fn try_with_column_at(&self, index: usize, field: Field, arr: ArrayRef) -> Result<RecordBatch>;
527
528 fn try_new_from_struct_array(&self, arr: StructArray) -> Result<RecordBatch>;
532
533 fn merge(&self, other: &RecordBatch) -> Result<RecordBatch>;
578
579 fn merge_with_schema(&self, other: &RecordBatch, schema: &Schema) -> Result<RecordBatch>;
589
590 fn drop_column(&self, name: &str) -> Result<RecordBatch>;
594
595 fn replace_column_by_name(&self, name: &str, column: Arc<dyn Array>) -> Result<RecordBatch>;
597
598 fn replace_column_schema_by_name(
600 &self,
601 name: &str,
602 new_data_type: DataType,
603 column: Arc<dyn Array>,
604 ) -> Result<RecordBatch>;
605
606 fn rename_column(&self, index: usize, new_name: &str) -> Result<RecordBatch>;
608
609 fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef>;
611
612 fn project_by_schema(&self, schema: &Schema) -> Result<RecordBatch>;
614
615 fn metadata(&self) -> &HashMap<String, String>;
617
618 fn add_metadata(&self, key: String, value: String) -> Result<RecordBatch> {
620 let mut metadata = self.metadata().clone();
621 metadata.insert(key, value);
622 self.with_metadata(metadata)
623 }
624
625 fn with_metadata(&self, metadata: HashMap<String, String>) -> Result<RecordBatch>;
627
628 fn take(&self, indices: &UInt32Array) -> Result<RecordBatch>;
630
631 fn shrink_to_fit(&self) -> Result<RecordBatch>;
633
634 fn sort_by_column(&self, column: usize, options: Option<SortOptions>) -> Result<RecordBatch>;
636}
637
638impl RecordBatchExt for RecordBatch {
639 fn try_with_column(&self, field: Field, arr: ArrayRef) -> Result<Self> {
640 let new_schema = Arc::new(self.schema().as_ref().try_with_column(field)?);
641 let mut new_columns = self.columns().to_vec();
642 new_columns.push(arr);
643 Self::try_new(new_schema, new_columns)
644 }
645
646 fn try_with_column_at(&self, index: usize, field: Field, arr: ArrayRef) -> Result<Self> {
647 let new_schema = Arc::new(self.schema().as_ref().try_with_column_at(index, field)?);
648 let mut new_columns = self.columns().to_vec();
649 new_columns.insert(index, arr);
650 Self::try_new(new_schema, new_columns)
651 }
652
653 fn try_new_from_struct_array(&self, arr: StructArray) -> Result<Self> {
654 let schema = Arc::new(Schema::new_with_metadata(
655 arr.fields().to_vec(),
656 self.schema().metadata.clone(),
657 ));
658 let batch = Self::from(arr);
659 batch.with_schema(schema)
660 }
661
662 fn merge(&self, other: &Self) -> Result<Self> {
663 if self.num_rows() != other.num_rows() {
664 return Err(ArrowError::InvalidArgumentError(format!(
665 "Attempt to merge two RecordBatch with different sizes: {} != {}",
666 self.num_rows(),
667 other.num_rows()
668 )));
669 }
670 let left_struct_array: StructArray = self.clone().into();
671 let right_struct_array: StructArray = other.clone().into();
672 self.try_new_from_struct_array(merge(&left_struct_array, &right_struct_array))
673 }
674
675 fn merge_with_schema(&self, other: &RecordBatch, schema: &Schema) -> Result<RecordBatch> {
676 if self.num_rows() != other.num_rows() {
677 return Err(ArrowError::InvalidArgumentError(format!(
678 "Attempt to merge two RecordBatch with different sizes: {} != {}",
679 self.num_rows(),
680 other.num_rows()
681 )));
682 }
683 let left_struct_array: StructArray = self.clone().into();
684 let right_struct_array: StructArray = other.clone().into();
685 self.try_new_from_struct_array(merge_with_schema(
686 &left_struct_array,
687 &right_struct_array,
688 schema.fields(),
689 ))
690 }
691
692 fn drop_column(&self, name: &str) -> Result<Self> {
693 let mut fields = vec![];
694 let mut columns = vec![];
695 for i in 0..self.schema().fields.len() {
696 if self.schema().field(i).name() != name {
697 fields.push(self.schema().field(i).clone());
698 columns.push(self.column(i).clone());
699 }
700 }
701 Self::try_new(
702 Arc::new(Schema::new_with_metadata(
703 fields,
704 self.schema().metadata().clone(),
705 )),
706 columns,
707 )
708 }
709
710 fn rename_column(&self, index: usize, new_name: &str) -> Result<RecordBatch> {
711 let mut fields = self.schema().fields().to_vec();
712 if index >= fields.len() {
713 return Err(ArrowError::InvalidArgumentError(format!(
714 "Index out of bounds: {}",
715 index
716 )));
717 }
718 fields[index] = Arc::new(Field::new(
719 new_name,
720 fields[index].data_type().clone(),
721 fields[index].is_nullable(),
722 ));
723 Self::try_new(
724 Arc::new(Schema::new_with_metadata(
725 fields,
726 self.schema().metadata().clone(),
727 )),
728 self.columns().to_vec(),
729 )
730 }
731
732 fn replace_column_by_name(&self, name: &str, column: Arc<dyn Array>) -> Result<RecordBatch> {
733 let mut columns = self.columns().to_vec();
734 let field_i = self
735 .schema()
736 .fields()
737 .iter()
738 .position(|f| f.name() == name)
739 .ok_or_else(|| ArrowError::SchemaError(format!("Field {} does not exist", name)))?;
740 columns[field_i] = column;
741 Self::try_new(self.schema(), columns)
742 }
743
744 fn replace_column_schema_by_name(
745 &self,
746 name: &str,
747 new_data_type: DataType,
748 column: Arc<dyn Array>,
749 ) -> Result<RecordBatch> {
750 let fields = self
751 .schema()
752 .fields()
753 .iter()
754 .map(|x| {
755 if x.name() != name {
756 x.clone()
757 } else {
758 let new_field = Field::new(name, new_data_type.clone(), x.is_nullable());
759 Arc::new(new_field)
760 }
761 })
762 .collect::<Vec<_>>();
763 let schema = Schema::new_with_metadata(fields, self.schema().metadata.clone());
764 let mut columns = self.columns().to_vec();
765 let field_i = self
766 .schema()
767 .fields()
768 .iter()
769 .position(|f| f.name() == name)
770 .ok_or_else(|| ArrowError::SchemaError(format!("Field {} does not exist", name)))?;
771 columns[field_i] = column;
772 Self::try_new(Arc::new(schema), columns)
773 }
774
775 fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef> {
776 let split = name.split('.').collect::<Vec<_>>();
777 if split.is_empty() {
778 return None;
779 }
780
781 self.column_by_name(split[0])
782 .and_then(|arr| get_sub_array(arr, &split[1..]))
783 }
784
785 fn project_by_schema(&self, schema: &Schema) -> Result<Self> {
786 let struct_array: StructArray = self.clone().into();
787 self.try_new_from_struct_array(project(&struct_array, schema.fields())?)
788 }
789
790 fn metadata(&self) -> &HashMap<String, String> {
791 self.schema_ref().metadata()
792 }
793
794 fn with_metadata(&self, metadata: HashMap<String, String>) -> Result<RecordBatch> {
795 let mut schema = self.schema_ref().as_ref().clone();
796 schema.metadata = metadata;
797 Self::try_new(schema.into(), self.columns().into())
798 }
799
800 fn take(&self, indices: &UInt32Array) -> Result<Self> {
801 let struct_array: StructArray = self.clone().into();
802 let taken = take(&struct_array, indices, None)?;
803 self.try_new_from_struct_array(taken.as_struct().clone())
804 }
805
806 fn shrink_to_fit(&self) -> Result<Self> {
807 crate::deepcopy::deep_copy_batch_sliced(self)
809 }
810
811 fn sort_by_column(&self, column: usize, options: Option<SortOptions>) -> Result<Self> {
812 if column >= self.num_columns() {
813 return Err(ArrowError::InvalidArgumentError(format!(
814 "Column index out of bounds: {}",
815 column
816 )));
817 }
818 let column = self.column(column);
819 let sorted = arrow_ord::sort::sort_to_indices(column, options, None)?;
820 self.take(&sorted)
821 }
822}
823
824fn project_array(array: &ArrayRef, target_field: &Field) -> Result<ArrayRef> {
827 match target_field.data_type() {
828 DataType::Struct(subfields) => {
829 let struct_arr = array.as_struct();
830 let projected = project(struct_arr, subfields)?;
831 Ok(Arc::new(projected))
832 }
833 DataType::List(inner_field) => {
834 let list_arr: &ListArray = array.as_list();
835 let projected_values = project_array(list_arr.values(), inner_field.as_ref())?;
836 Ok(Arc::new(ListArray::new(
837 inner_field.clone(),
838 list_arr.offsets().clone(),
839 projected_values,
840 list_arr.nulls().cloned(),
841 )))
842 }
843 DataType::LargeList(inner_field) => {
844 let list_arr: &LargeListArray = array.as_list();
845 let projected_values = project_array(list_arr.values(), inner_field.as_ref())?;
846 Ok(Arc::new(LargeListArray::new(
847 inner_field.clone(),
848 list_arr.offsets().clone(),
849 projected_values,
850 list_arr.nulls().cloned(),
851 )))
852 }
853 DataType::FixedSizeList(inner_field, size) => {
854 let list_arr = array.as_fixed_size_list();
855 let projected_values = project_array(list_arr.values(), inner_field.as_ref())?;
856 Ok(Arc::new(FixedSizeListArray::new(
857 inner_field.clone(),
858 *size,
859 projected_values,
860 list_arr.nulls().cloned(),
861 )))
862 }
863 DataType::Map(entries_field, sorted) if !entries_field.is_nullable() => {
867 let map_arr = array.as_map();
868 let DataType::Struct(entry_fields) = entries_field.data_type() else {
869 return Err(ArrowError::SchemaError(format!(
870 "Map entries field must be a struct, got {}",
871 entries_field.data_type()
872 )));
873 };
874 let projected_entries = project(map_arr.entries(), entry_fields)?;
875 Ok(Arc::new(MapArray::try_new(
880 entries_field.clone(),
881 map_arr.offsets().clone(),
882 projected_entries,
883 map_arr.nulls().cloned(),
884 *sorted,
885 )?))
886 }
887 _ => Ok(array.clone()),
888 }
889}
890
891fn project(struct_array: &StructArray, fields: &Fields) -> Result<StructArray> {
892 if fields.is_empty() {
893 return Ok(StructArray::new_empty_fields(
894 struct_array.len(),
895 struct_array.nulls().cloned(),
896 ));
897 }
898 let mut columns: Vec<ArrayRef> = vec![];
899 for field in fields.iter() {
900 if let Some(col) = struct_array.column_by_name(field.name()) {
901 let projected = project_array(col, field.as_ref())?;
902 columns.push(projected);
903 } else {
904 return Err(ArrowError::SchemaError(format!(
905 "field {} does not exist in the RecordBatch",
906 field.name()
907 )));
908 }
909 }
910 StructArray::try_new(fields.clone(), columns, struct_array.nulls().cloned())
912}
913
914fn holds_struct(data_type: &DataType) -> bool {
916 match data_type {
917 DataType::Struct(_) => true,
918 DataType::List(item) | DataType::LargeList(item) => holds_struct(item.data_type()),
919 _ => false,
920 }
921}
922
923fn merge_list_values(left: &ArrayRef, right: &ArrayRef) -> ArrayRef {
926 match (left.data_type(), right.data_type()) {
927 (DataType::Struct(_), DataType::Struct(_)) => {
928 Arc::new(merge(left.as_struct(), right.as_struct()))
929 }
930 (DataType::List(_), DataType::List(_))
931 | (DataType::LargeList(_), DataType::LargeList(_))
932 if left.data_type() != right.data_type()
933 && holds_struct(left.data_type())
934 && holds_struct(right.data_type()) =>
935 {
936 merge_list_struct(left, right)
937 }
938 _ => left.clone(),
939 }
940}
941
942fn merge_list_struct_helper<T: OffsetSizeTrait>(
943 left: &dyn Array,
944 right: &dyn Array,
945) -> Arc<dyn Array> {
946 let left_list: &GenericListArray<T> = left.as_list();
947 let right_list: &GenericListArray<T> = right.as_list();
948 let item_name = match left.data_type() {
949 DataType::List(item) | DataType::LargeList(item) => item.name().clone(),
950 _ => unreachable!(),
951 };
952 let null_left = left_list.null_count() == left_list.len();
955 let null_right = !null_left && right_list.null_count() == right_list.len();
956 let (merged_items, list, items_nullable) = if null_left || null_right {
957 let not_null = if null_left { right_list } else { left_list };
958 let values_len = not_null.values().len();
959 let values = |list: &GenericListArray<T>, is_null: bool| {
960 if is_null {
961 new_null_array(list.values().data_type(), values_len)
962 } else {
963 list.values().clone()
964 }
965 };
966 let merged = merge_list_values(
967 &values(left_list, null_left),
968 &values(right_list, null_right),
969 );
970 (merged, not_null, true)
971 } else {
972 if left_list.offsets().inner() != right_list.offsets().inner() {
973 panic!("Attempt to merge list struct arrays which do not have same offsets");
974 }
975 let merged = merge_list_values(left_list.values(), right_list.values());
976 let nullable = match left.data_type() {
977 DataType::List(item) | DataType::LargeList(item) => item.is_nullable(),
978 _ => unreachable!(),
979 };
980 (merged, left_list, nullable)
981 };
982 let items_field = Arc::new(Field::new(
983 item_name,
984 merged_items.data_type().clone(),
985 items_nullable,
986 ));
987 Arc::new(GenericListArray::<T>::new(
988 items_field,
989 list.offsets().clone(),
990 merged_items,
991 list.nulls().cloned(),
992 ))
993}
994
995fn merge_list_struct(left: &dyn Array, right: &dyn Array) -> Arc<dyn Array> {
998 match (left.data_type(), right.data_type()) {
999 (DataType::List(_), DataType::List(_)) => merge_list_struct_helper::<i32>(left, right),
1000 (DataType::LargeList(_), DataType::LargeList(_)) => {
1001 merge_list_struct_helper::<i64>(left, right)
1002 }
1003 _ => unreachable!(),
1004 }
1005}
1006
1007fn merge_struct_validity(
1012 left_validity: Option<&arrow_buffer::NullBuffer>,
1013 right_validity: Option<&arrow_buffer::NullBuffer>,
1014) -> Option<arrow_buffer::NullBuffer> {
1015 match (left_validity, right_validity) {
1016 (None, _) | (_, None) => None,
1018 (Some(left), Some(right)) => {
1019 if left.null_count() == 0 || right.null_count() == 0 {
1020 return None;
1021 }
1022 if left.null_count() == left.len() {
1023 return Some(right.clone());
1024 }
1025 if right.null_count() == right.len() {
1026 return Some(left.clone());
1027 }
1028
1029 let left_buffer = left.inner();
1030 let right_buffer = right.inner();
1031
1032 let merged_buffer = left_buffer | right_buffer;
1035
1036 Some(arrow_buffer::NullBuffer::from(merged_buffer))
1037 }
1038 }
1039}
1040
1041fn merge_list_child_values(
1042 child_field: &Field,
1043 left_values: ArrayRef,
1044 right_values: ArrayRef,
1045) -> ArrayRef {
1046 match child_field.data_type() {
1047 DataType::Struct(child_fields) => Arc::new(merge_with_schema(
1048 left_values.as_struct(),
1049 right_values.as_struct(),
1050 child_fields,
1051 )) as ArrayRef,
1052 DataType::List(grandchild) => {
1053 let left_list = left_values
1054 .as_any()
1055 .downcast_ref::<ListArray>()
1056 .expect("left list values should be ListArray");
1057 let right_list = right_values
1058 .as_any()
1059 .downcast_ref::<ListArray>()
1060 .expect("right list values should be ListArray");
1061 let merged_values = merge_list_child_values(
1062 grandchild.as_ref(),
1063 left_list.values().clone(),
1064 right_list.values().clone(),
1065 );
1066 let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1067 Arc::new(ListArray::new(
1068 grandchild.clone(),
1069 left_list.offsets().clone(),
1070 merged_values,
1071 merged_validity,
1072 )) as ArrayRef
1073 }
1074 DataType::LargeList(grandchild) => {
1075 let left_list = left_values
1076 .as_any()
1077 .downcast_ref::<LargeListArray>()
1078 .expect("left list values should be LargeListArray");
1079 let right_list = right_values
1080 .as_any()
1081 .downcast_ref::<LargeListArray>()
1082 .expect("right list values should be LargeListArray");
1083 let merged_values = merge_list_child_values(
1084 grandchild.as_ref(),
1085 left_list.values().clone(),
1086 right_list.values().clone(),
1087 );
1088 let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1089 Arc::new(LargeListArray::new(
1090 grandchild.clone(),
1091 left_list.offsets().clone(),
1092 merged_values,
1093 merged_validity,
1094 )) as ArrayRef
1095 }
1096 DataType::FixedSizeList(grandchild, list_size) => {
1097 let left_list = left_values
1098 .as_any()
1099 .downcast_ref::<FixedSizeListArray>()
1100 .expect("left list values should be FixedSizeListArray");
1101 let right_list = right_values
1102 .as_any()
1103 .downcast_ref::<FixedSizeListArray>()
1104 .expect("right list values should be FixedSizeListArray");
1105 let merged_values = merge_list_child_values(
1106 grandchild.as_ref(),
1107 left_list.values().clone(),
1108 right_list.values().clone(),
1109 );
1110 let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1111 Arc::new(FixedSizeListArray::new(
1112 grandchild.clone(),
1113 *list_size,
1114 merged_values,
1115 merged_validity,
1116 )) as ArrayRef
1117 }
1118 _ => left_values.clone(),
1119 }
1120}
1121
1122fn adjust_child_validity(
1126 child: &ArrayRef,
1127 parent_validity: Option<&arrow_buffer::NullBuffer>,
1128) -> ArrayRef {
1129 let parent_validity = match parent_validity {
1131 None => return child.clone(),
1132 Some(p) if p.null_count() == 0 => return child.clone(), Some(p) => p,
1134 };
1135
1136 if child.data_type() == &DataType::Null {
1139 return child.clone();
1140 }
1141
1142 let child_validity = child.nulls();
1143
1144 let new_validity = match child_validity {
1146 None => {
1147 parent_validity.clone()
1149 }
1150 Some(child_nulls) => {
1151 let child_buffer = child_nulls.inner();
1152 let parent_buffer = parent_validity.inner();
1153
1154 let merged_buffer = child_buffer & parent_buffer;
1157
1158 arrow_buffer::NullBuffer::from(merged_buffer)
1159 }
1160 };
1161
1162 arrow_array::make_array(
1164 arrow_data::ArrayData::try_new(
1165 child.data_type().clone(),
1166 child.len(),
1167 Some(new_validity.into_inner().into_inner()),
1168 child.offset(),
1169 child.to_data().buffers().to_vec(),
1170 child.to_data().child_data().to_vec(),
1171 )
1172 .unwrap(),
1173 )
1174}
1175
1176fn merge(left_struct_array: &StructArray, right_struct_array: &StructArray) -> StructArray {
1177 let mut fields: Vec<Field> = vec![];
1178 let mut columns: Vec<ArrayRef> = vec![];
1179 let right_fields = right_struct_array.fields();
1180 let right_columns = right_struct_array.columns();
1181
1182 let left_validity = left_struct_array.nulls();
1184 let right_validity = right_struct_array.nulls();
1185
1186 let merged_validity = merge_struct_validity(left_validity, right_validity);
1188
1189 for (left_field, left_column) in left_struct_array
1191 .fields()
1192 .iter()
1193 .zip(left_struct_array.columns().iter())
1194 {
1195 match right_fields
1196 .iter()
1197 .position(|f| f.name() == left_field.name())
1198 {
1199 Some(right_index) => {
1201 let right_field = right_fields.get(right_index).unwrap();
1202 let right_column = right_columns.get(right_index).unwrap();
1203 match (left_field.data_type(), right_field.data_type()) {
1205 (DataType::Struct(_), DataType::Struct(_)) => {
1206 let left_sub_array = left_column.as_struct();
1207 let right_sub_array = right_column.as_struct();
1208 let merged_sub_array = merge(left_sub_array, right_sub_array);
1209 fields.push(Field::new(
1210 left_field.name(),
1211 merged_sub_array.data_type().clone(),
1212 left_field.is_nullable(),
1213 ));
1214 columns.push(Arc::new(merged_sub_array) as ArrayRef);
1215 }
1216 (DataType::List(left_list), DataType::List(right_list))
1217 | (DataType::LargeList(left_list), DataType::LargeList(right_list))
1218 if holds_struct(left_list.data_type())
1219 && holds_struct(right_list.data_type()) =>
1220 {
1221 if left_list.data_type() == right_list.data_type() {
1223 fields.push(left_field.as_ref().clone());
1224 columns.push(left_column.clone());
1225 } else {
1226 let merged_sub_array = merge_list_struct(&left_column, &right_column);
1230
1231 fields.push(Field::new(
1232 left_field.name(),
1233 merged_sub_array.data_type().clone(),
1234 left_field.is_nullable(),
1235 ));
1236 columns.push(merged_sub_array);
1237 }
1238 }
1239 _ => {
1241 fields.push(left_field.as_ref().clone());
1243 let adjusted_column = adjust_child_validity(left_column, left_validity);
1245 columns.push(adjusted_column);
1246 }
1247 }
1248 }
1249 None => {
1250 fields.push(left_field.as_ref().clone());
1251 let adjusted_column = adjust_child_validity(left_column, left_validity);
1253 columns.push(adjusted_column);
1254 }
1255 }
1256 }
1257
1258 right_fields
1260 .iter()
1261 .zip(right_columns.iter())
1262 .for_each(|(field, column)| {
1263 if !left_struct_array
1265 .fields()
1266 .iter()
1267 .any(|f| f.name() == field.name())
1268 {
1269 fields.push(field.as_ref().clone());
1270 let adjusted_column = adjust_child_validity(column, right_validity);
1273 columns.push(adjusted_column);
1274 }
1275 });
1276
1277 StructArray::try_new(Fields::from(fields), columns, merged_validity).unwrap()
1278}
1279
1280fn merge_with_schema(
1281 left_struct_array: &StructArray,
1282 right_struct_array: &StructArray,
1283 fields: &Fields,
1284) -> StructArray {
1285 fn same_type_kind(left: &DataType, right: &DataType) -> bool {
1287 match (left, right) {
1288 (DataType::Struct(_), DataType::Struct(_)) => true,
1289 (DataType::Struct(_), _) => false,
1290 (_, DataType::Struct(_)) => false,
1291 _ => true,
1292 }
1293 }
1294
1295 let mut output_fields: Vec<Field> = Vec::with_capacity(fields.len());
1296 let mut columns: Vec<ArrayRef> = Vec::with_capacity(fields.len());
1297
1298 let left_fields = left_struct_array.fields();
1299 let left_columns = left_struct_array.columns();
1300 let right_fields = right_struct_array.fields();
1301 let right_columns = right_struct_array.columns();
1302
1303 let left_validity = left_struct_array.nulls();
1305 let right_validity = right_struct_array.nulls();
1306
1307 let merged_validity = merge_struct_validity(left_validity, right_validity);
1309
1310 for field in fields {
1311 let left_match_idx = left_fields.iter().position(|f| {
1312 f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1313 });
1314 let right_match_idx = right_fields.iter().position(|f| {
1315 f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1316 });
1317
1318 match (left_match_idx, right_match_idx) {
1319 (None, Some(right_idx)) => {
1320 output_fields.push(right_fields[right_idx].as_ref().clone());
1321 let adjusted_column =
1323 adjust_child_validity(&right_columns[right_idx], right_validity);
1324 columns.push(adjusted_column);
1325 }
1326 (Some(left_idx), None) => {
1327 output_fields.push(left_fields[left_idx].as_ref().clone());
1328 let adjusted_column = adjust_child_validity(&left_columns[left_idx], left_validity);
1330 columns.push(adjusted_column);
1331 }
1332 (Some(left_idx), Some(right_idx)) => {
1333 match field.data_type() {
1334 DataType::Struct(child_fields) => {
1335 let left_sub_array = left_columns[left_idx].as_struct();
1336 let right_sub_array = right_columns[right_idx].as_struct();
1337 let merged_sub_array =
1338 merge_with_schema(left_sub_array, right_sub_array, child_fields);
1339 output_fields.push(Field::new(
1340 field.name(),
1341 merged_sub_array.data_type().clone(),
1342 field.is_nullable(),
1343 ));
1344 columns.push(Arc::new(merged_sub_array) as ArrayRef);
1345 }
1346 DataType::List(child_field) => {
1347 let left_list = left_columns[left_idx]
1348 .as_any()
1349 .downcast_ref::<ListArray>()
1350 .unwrap();
1351 let right_list = right_columns[right_idx]
1352 .as_any()
1353 .downcast_ref::<ListArray>()
1354 .unwrap();
1355 let merged_values = merge_list_child_values(
1356 child_field.as_ref(),
1357 left_list.trimmed_values(),
1358 right_list.trimmed_values(),
1359 );
1360 let merged_validity =
1361 merge_struct_validity(left_list.nulls(), right_list.nulls());
1362 let merged_list = ListArray::new(
1366 child_field.clone(),
1367 left_list.trimmed_offsets(),
1368 merged_values,
1369 merged_validity,
1370 );
1371 output_fields.push(field.as_ref().clone());
1372 columns.push(Arc::new(merged_list) as ArrayRef);
1373 }
1374 DataType::LargeList(child_field) => {
1375 let left_list = left_columns[left_idx]
1376 .as_any()
1377 .downcast_ref::<LargeListArray>()
1378 .unwrap();
1379 let right_list = right_columns[right_idx]
1380 .as_any()
1381 .downcast_ref::<LargeListArray>()
1382 .unwrap();
1383 let merged_values = merge_list_child_values(
1384 child_field.as_ref(),
1385 left_list.trimmed_values(),
1386 right_list.trimmed_values(),
1387 );
1388 let merged_validity =
1389 merge_struct_validity(left_list.nulls(), right_list.nulls());
1390 let merged_list = LargeListArray::new(
1391 child_field.clone(),
1392 left_list.trimmed_offsets(),
1393 merged_values,
1394 merged_validity,
1395 );
1396 output_fields.push(field.as_ref().clone());
1397 columns.push(Arc::new(merged_list) as ArrayRef);
1398 }
1399 DataType::FixedSizeList(child_field, list_size) => {
1400 let left_list = left_columns[left_idx]
1401 .as_any()
1402 .downcast_ref::<FixedSizeListArray>()
1403 .unwrap();
1404 let right_list = right_columns[right_idx]
1405 .as_any()
1406 .downcast_ref::<FixedSizeListArray>()
1407 .unwrap();
1408 let merged_values = merge_list_child_values(
1409 child_field.as_ref(),
1410 left_list.values().clone(),
1411 right_list.values().clone(),
1412 );
1413 let merged_validity =
1414 merge_struct_validity(left_list.nulls(), right_list.nulls());
1415 let merged_list = FixedSizeListArray::new(
1416 child_field.clone(),
1417 *list_size,
1418 merged_values,
1419 merged_validity,
1420 );
1421 output_fields.push(field.as_ref().clone());
1422 columns.push(Arc::new(merged_list) as ArrayRef);
1423 }
1424 _ => {
1425 output_fields.push(left_fields[left_idx].as_ref().clone());
1426 let adjusted_column =
1428 adjust_child_validity(&left_columns[left_idx], left_validity);
1429 columns.push(adjusted_column);
1430 }
1431 }
1432 }
1433 (None, None) => {
1434 }
1436 }
1437 }
1438
1439 StructArray::try_new(Fields::from(output_fields), columns, merged_validity).unwrap()
1440}
1441
1442fn get_sub_array<'a>(array: &'a ArrayRef, components: &[&str]) -> Option<&'a ArrayRef> {
1443 if components.is_empty() {
1444 return Some(array);
1445 }
1446 if !matches!(array.data_type(), DataType::Struct(_)) {
1447 return None;
1448 }
1449 let struct_arr = array.as_struct();
1450 struct_arr
1451 .column_by_name(components[0])
1452 .and_then(|arr| get_sub_array(arr, &components[1..]))
1453}
1454
1455pub fn interleave_batches(
1459 batches: &[RecordBatch],
1460 indices: &[(usize, usize)],
1461) -> Result<RecordBatch> {
1462 let first_batch = batches.first().ok_or_else(|| {
1463 ArrowError::InvalidArgumentError("Cannot interleave zero RecordBatches".to_string())
1464 })?;
1465 let schema = first_batch.schema();
1466 let num_columns = first_batch.num_columns();
1467 let mut columns = Vec::with_capacity(num_columns);
1468 let mut chunks = Vec::with_capacity(batches.len());
1469
1470 for i in 0..num_columns {
1471 for batch in batches {
1472 chunks.push(batch.column(i).as_ref());
1473 }
1474 let new_column = interleave(&chunks, indices)?;
1475 columns.push(new_column);
1476 chunks.clear();
1477 }
1478
1479 RecordBatch::try_new(schema, columns)
1480}
1481
1482pub trait BufferExt {
1483 fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self;
1498
1499 fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self;
1508}
1509
1510fn is_pwr_two(n: u64) -> bool {
1511 n & (n - 1) == 0
1512}
1513
1514impl BufferExt for arrow_buffer::Buffer {
1515 fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self {
1516 if is_pwr_two(bytes_per_value) && bytes.as_ptr().align_offset(bytes_per_value as usize) != 0
1517 {
1518 let size_bytes = bytes.len();
1520 Self::copy_bytes_bytes(bytes, size_bytes)
1521 } else {
1522 unsafe {
1525 Self::from_custom_allocation(
1526 NonNull::new(bytes.as_ptr() as _).expect("should be a valid pointer"),
1527 bytes.len(),
1528 Arc::new(bytes),
1529 )
1530 }
1531 }
1532 }
1533
1534 fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self {
1535 assert!(size_bytes >= bytes.len());
1536 let mut buf = MutableBuffer::with_capacity(size_bytes);
1537 let to_fill = size_bytes - bytes.len();
1538 buf.extend(bytes);
1539 buf.extend(std::iter::repeat_n(0_u8, to_fill));
1540
1541 buf.shrink_to_fit();
1544
1545 Self::from(buf)
1546 }
1547}
1548
1549#[cfg(test)]
1550mod tests {
1551 use super::*;
1552 use arrow_array::{
1553 BinaryArray, BinaryViewArray, FixedSizeBinaryArray, Float32Array, Int32Array,
1554 LargeBinaryArray, ListArray, NullArray, StringArray, StructArray, new_empty_array,
1555 new_null_array,
1556 };
1557 use arrow_buffer::OffsetBuffer;
1558
1559 #[test]
1560 fn test_convert_to_floating_point_preserves_inner_nulls() {
1561 let values = Int8Array::from(vec![Some(1), None, Some(3), Some(4)]);
1566 let fsl = FixedSizeListArray::new(
1567 Arc::new(Field::new("item", DataType::Int8, true)),
1568 2,
1569 Arc::new(values),
1570 None,
1571 );
1572
1573 let converted = fsl.convert_to_floating_point().unwrap();
1574
1575 assert_eq!(converted.len(), 2);
1576 let conv_values = converted
1577 .values()
1578 .as_any()
1579 .downcast_ref::<Float32Array>()
1580 .unwrap();
1581 assert_eq!(conv_values.len(), 4);
1582 assert_eq!(conv_values.value(0), 1.0);
1583 assert!(conv_values.is_null(1));
1584 assert_eq!(conv_values.value(2), 3.0);
1585 assert_eq!(conv_values.value(3), 4.0);
1586 }
1587
1588 #[test]
1589 fn test_convert_to_floating_point_preserves_inner_nulls_f64_arm() {
1590 let values = UInt8Array::from(vec![Some(10u8), None, Some(30), Some(40)]);
1594 let fsl = FixedSizeListArray::new(
1595 Arc::new(Field::new("item", DataType::UInt8, true)),
1596 2,
1597 Arc::new(values),
1598 None,
1599 );
1600
1601 let converted = fsl.convert_to_floating_point().unwrap();
1602
1603 assert_eq!(converted.len(), 2);
1604 let conv_values = converted
1605 .values()
1606 .as_any()
1607 .downcast_ref::<Float64Array>()
1608 .unwrap();
1609 assert_eq!(conv_values.len(), 4);
1610 assert_eq!(conv_values.value(0), 10.0);
1611 assert!(conv_values.is_null(1));
1612 assert_eq!(conv_values.value(2), 30.0);
1613 assert_eq!(conv_values.value(3), 40.0);
1614 }
1615
1616 #[test]
1617 fn test_merge_recursive() {
1618 let a_array = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
1619 let e_array = Int32Array::from(vec![Some(4), Some(5), Some(6)]);
1620 let c_array = Int32Array::from(vec![Some(7), Some(8), Some(9)]);
1621 let d_array = StringArray::from(vec![Some("a"), Some("b"), Some("c")]);
1622
1623 let left_schema = Schema::new(vec![
1624 Field::new("a", DataType::Int32, true),
1625 Field::new(
1626 "b",
1627 DataType::Struct(vec![Field::new("c", DataType::Int32, true)].into()),
1628 true,
1629 ),
1630 ]);
1631 let left_batch = RecordBatch::try_new(
1632 Arc::new(left_schema),
1633 vec![
1634 Arc::new(a_array.clone()),
1635 Arc::new(StructArray::from(vec![(
1636 Arc::new(Field::new("c", DataType::Int32, true)),
1637 Arc::new(c_array.clone()) as ArrayRef,
1638 )])),
1639 ],
1640 )
1641 .unwrap();
1642
1643 let right_schema = Schema::new(vec![
1644 Field::new("e", DataType::Int32, true),
1645 Field::new(
1646 "b",
1647 DataType::Struct(vec![Field::new("d", DataType::Utf8, true)].into()),
1648 true,
1649 ),
1650 ]);
1651 let right_batch = RecordBatch::try_new(
1652 Arc::new(right_schema),
1653 vec![
1654 Arc::new(e_array.clone()),
1655 Arc::new(StructArray::from(vec![(
1656 Arc::new(Field::new("d", DataType::Utf8, true)),
1657 Arc::new(d_array.clone()) as ArrayRef,
1658 )])) as ArrayRef,
1659 ],
1660 )
1661 .unwrap();
1662
1663 let merged_schema = Schema::new(vec![
1664 Field::new("a", DataType::Int32, true),
1665 Field::new(
1666 "b",
1667 DataType::Struct(
1668 vec![
1669 Field::new("c", DataType::Int32, true),
1670 Field::new("d", DataType::Utf8, true),
1671 ]
1672 .into(),
1673 ),
1674 true,
1675 ),
1676 Field::new("e", DataType::Int32, true),
1677 ]);
1678 let merged_batch = RecordBatch::try_new(
1679 Arc::new(merged_schema),
1680 vec![
1681 Arc::new(a_array) as ArrayRef,
1682 Arc::new(StructArray::from(vec![
1683 (
1684 Arc::new(Field::new("c", DataType::Int32, true)),
1685 Arc::new(c_array) as ArrayRef,
1686 ),
1687 (
1688 Arc::new(Field::new("d", DataType::Utf8, true)),
1689 Arc::new(d_array) as ArrayRef,
1690 ),
1691 ])) as ArrayRef,
1692 Arc::new(e_array) as ArrayRef,
1693 ],
1694 )
1695 .unwrap();
1696
1697 let result = left_batch.merge(&right_batch).unwrap();
1698 assert_eq!(result, merged_batch);
1699 }
1700
1701 #[test]
1702 fn test_merge_with_schema() {
1703 fn test_batch(names: &[&str], types: &[DataType]) -> (Schema, RecordBatch) {
1704 let fields: Fields = names
1705 .iter()
1706 .zip(types)
1707 .map(|(name, ty)| Field::new(name.to_string(), ty.clone(), false))
1708 .collect();
1709 let schema = Schema::new(vec![Field::new(
1710 "struct",
1711 DataType::Struct(fields.clone()),
1712 false,
1713 )]);
1714 let children = types.iter().map(new_empty_array).collect::<Vec<_>>();
1715 let batch = RecordBatch::try_new(
1716 Arc::new(schema.clone()),
1717 vec![Arc::new(StructArray::new(fields, children, None)) as ArrayRef],
1718 );
1719 (schema, batch.unwrap())
1720 }
1721
1722 let (_, left_batch) = test_batch(&["a", "b"], &[DataType::Int32, DataType::Int64]);
1723 let (_, right_batch) = test_batch(&["c", "b"], &[DataType::Int32, DataType::Int64]);
1724 let (output_schema, _) = test_batch(
1725 &["b", "a", "c"],
1726 &[DataType::Int64, DataType::Int32, DataType::Int32],
1727 );
1728
1729 let merged = left_batch
1731 .merge_with_schema(&right_batch, &output_schema)
1732 .unwrap();
1733 assert_eq!(merged.schema().as_ref(), &output_schema);
1734
1735 let (naive_schema, _) = test_batch(
1737 &["a", "b", "c"],
1738 &[DataType::Int32, DataType::Int64, DataType::Int32],
1739 );
1740 let merged = left_batch.merge(&right_batch).unwrap();
1741 assert_eq!(merged.schema().as_ref(), &naive_schema);
1742 }
1743
1744 #[test]
1745 fn test_merge_list_struct() {
1746 let x_field = Arc::new(Field::new("x", DataType::Int32, true));
1747 let y_field = Arc::new(Field::new("y", DataType::Int32, true));
1748 let x_struct_field = Arc::new(Field::new(
1749 "item",
1750 DataType::Struct(Fields::from(vec![x_field.clone()])),
1751 true,
1752 ));
1753 let y_struct_field = Arc::new(Field::new(
1754 "item",
1755 DataType::Struct(Fields::from(vec![y_field.clone()])),
1756 true,
1757 ));
1758 let both_struct_field = Arc::new(Field::new(
1759 "item",
1760 DataType::Struct(Fields::from(vec![x_field.clone(), y_field.clone()])),
1761 true,
1762 ));
1763 let left_schema = Schema::new(vec![Field::new(
1764 "list_struct",
1765 DataType::List(x_struct_field.clone()),
1766 true,
1767 )]);
1768 let right_schema = Schema::new(vec![Field::new(
1769 "list_struct",
1770 DataType::List(y_struct_field.clone()),
1771 true,
1772 )]);
1773 let both_schema = Schema::new(vec![Field::new(
1774 "list_struct",
1775 DataType::List(both_struct_field.clone()),
1776 true,
1777 )]);
1778
1779 let x = Arc::new(Int32Array::from(vec![1]));
1780 let y = Arc::new(Int32Array::from(vec![2]));
1781 let x_struct = Arc::new(StructArray::new(
1782 Fields::from(vec![x_field.clone()]),
1783 vec![x.clone()],
1784 None,
1785 ));
1786 let y_struct = Arc::new(StructArray::new(
1787 Fields::from(vec![y_field.clone()]),
1788 vec![y.clone()],
1789 None,
1790 ));
1791 let both_struct = Arc::new(StructArray::new(
1792 Fields::from(vec![x_field.clone(), y_field.clone()]),
1793 vec![x.clone(), y],
1794 None,
1795 ));
1796 let both_null_struct = Arc::new(StructArray::new(
1797 Fields::from(vec![x_field, y_field]),
1798 vec![x, Arc::new(new_null_array(&DataType::Int32, 1))],
1799 None,
1800 ));
1801 let offsets = OffsetBuffer::from_lengths([1]);
1802 let x_s_list = ListArray::new(x_struct_field, offsets.clone(), x_struct, None);
1803 let y_s_list = ListArray::new(y_struct_field, offsets.clone(), y_struct, None);
1804 let both_list = ListArray::new(
1805 both_struct_field.clone(),
1806 offsets.clone(),
1807 both_struct,
1808 None,
1809 );
1810 let both_null_list = ListArray::new(both_struct_field, offsets, both_null_struct, None);
1811 let x_batch =
1812 RecordBatch::try_new(Arc::new(left_schema), vec![Arc::new(x_s_list)]).unwrap();
1813 let y_batch = RecordBatch::try_new(
1814 Arc::new(right_schema.clone()),
1815 vec![Arc::new(y_s_list.clone())],
1816 )
1817 .unwrap();
1818 let merged = x_batch.merge(&y_batch).unwrap();
1819 let expected =
1820 RecordBatch::try_new(Arc::new(both_schema.clone()), vec![Arc::new(both_list)]).unwrap();
1821 assert_eq!(merged, expected);
1822
1823 let y_null_list = new_null_array(y_s_list.data_type(), 1);
1824 let y_null_batch =
1825 RecordBatch::try_new(Arc::new(right_schema), vec![Arc::new(y_null_list.clone())])
1826 .unwrap();
1827 let expected =
1828 RecordBatch::try_new(Arc::new(both_schema), vec![Arc::new(both_null_list)]).unwrap();
1829 let merged = x_batch.merge(&y_null_batch).unwrap();
1830 assert_eq!(merged, expected);
1831 }
1832
1833 #[test]
1834 fn test_merge_all_null_list_adds_nested_child() {
1835 let strukt = |cols: Vec<(&str, ArrayRef)>| -> ArrayRef {
1838 let fields = cols
1839 .iter()
1840 .map(|(name, col)| Field::new(*name, col.data_type().clone(), true))
1841 .collect::<Vec<_>>();
1842 Arc::new(StructArray::new(
1843 fields.into(),
1844 cols.into_iter().map(|(_, col)| col).collect(),
1845 None,
1846 ))
1847 };
1848 let list = |values: ArrayRef, large: bool| -> ArrayRef {
1850 let field = Arc::new(Field::new("item", values.data_type().clone(), true));
1851 let lengths = [values.len(), 0];
1852 let nulls = Some(arrow_buffer::NullBuffer::from(vec![true, false]));
1853 if large {
1854 Arc::new(LargeListArray::new(
1855 field,
1856 OffsetBuffer::from_lengths(lengths),
1857 values,
1858 nulls,
1859 ))
1860 } else {
1861 Arc::new(ListArray::new(
1862 field,
1863 OffsetBuffer::from_lengths(lengths),
1864 values,
1865 nulls,
1866 ))
1867 }
1868 };
1869 let a = || Arc::new(Int32Array::from(vec![7])) as ArrayRef;
1870 let b = || new_null_array(&DataType::Utf8, 1);
1871 for large in [false, true] {
1872 let under_list = |inner: ArrayRef, with_t: bool| {
1874 let mut cols = vec![("in", inner)];
1875 if with_t {
1876 cols.insert(0, ("t", Arc::new(Int32Array::from(vec![1])) as ArrayRef));
1877 }
1878 list(strukt(cols), large)
1879 };
1880 let under_two_lists = |inner: ArrayRef, _: bool| {
1882 let inner_list: ArrayRef = if large {
1883 let field = Arc::new(Field::new("item", inner.data_type().clone(), true));
1884 Arc::new(LargeListArray::new(
1885 field,
1886 OffsetBuffer::from_lengths([1]),
1887 inner,
1888 None,
1889 ))
1890 } else {
1891 let field = Arc::new(Field::new("item", inner.data_type().clone(), true));
1892 Arc::new(ListArray::new(
1893 field,
1894 OffsetBuffer::from_lengths([1]),
1895 inner,
1896 None,
1897 ))
1898 };
1899 list(inner_list, large)
1900 };
1901 let shapes: [&dyn Fn(ArrayRef, bool) -> ArrayRef; 2] = [&under_list, &under_two_lists];
1902 for shape in shapes {
1903 let data = shape(strukt(vec![("a", a())]), true);
1904 let expected = shape(strukt(vec![("a", a()), ("b", b())]), true);
1905 let placeholder =
1906 new_null_array(shape(strukt(vec![("b", b())]), false).data_type(), 2);
1907 let batch = |col: ArrayRef| {
1908 RecordBatch::try_new(
1909 Arc::new(Schema::new(vec![Field::new(
1910 "c",
1911 col.data_type().clone(),
1912 true,
1913 )])),
1914 vec![col],
1915 )
1916 .unwrap()
1917 };
1918 let expected = batch(expected);
1919 for merged in [
1921 batch(data.clone())
1922 .merge(&batch(placeholder.clone()))
1923 .unwrap(),
1924 batch(placeholder.clone())
1925 .merge(&batch(data.clone()))
1926 .unwrap(),
1927 ] {
1928 let merged = merged.project_by_schema(expected.schema_ref()).unwrap();
1929 assert_eq!(merged, expected, "large={large}");
1930 }
1931 }
1932 }
1933 }
1934
1935 #[test]
1936 fn test_merge_list_struct_identical_schema() {
1937 let x_field = Arc::new(Field::new("x", DataType::Int32, true));
1941 let item_field = Arc::new(Field::new(
1942 "item",
1943 DataType::Struct(Fields::from(vec![x_field.clone()])),
1944 true,
1945 ));
1946 let schema = Arc::new(Schema::new(vec![Field::new(
1947 "list_struct",
1948 DataType::List(item_field.clone()),
1949 true,
1950 )]));
1951
1952 let build_list = |values: Vec<i32>| {
1953 let len = values.len();
1954 let item_struct = Arc::new(StructArray::new(
1955 Fields::from(vec![x_field.clone()]),
1956 vec![Arc::new(Int32Array::from(values))],
1957 None,
1958 ));
1959 ListArray::new(
1960 item_field.clone(),
1961 OffsetBuffer::from_lengths([len]),
1962 item_struct,
1963 None,
1964 )
1965 };
1966
1967 let left =
1970 RecordBatch::try_new(schema.clone(), vec![Arc::new(build_list(vec![1, 2]))]).unwrap();
1971 let right = RecordBatch::try_new(schema, vec![Arc::new(build_list(vec![3, 4]))]).unwrap();
1972
1973 let merged = left.merge(&right).unwrap();
1974
1975 assert_eq!(merged.num_columns(), 1);
1977 assert_eq!(merged, left);
1978 }
1979
1980 #[test]
1981 fn test_merge_nested_list_struct_identical_schema() {
1982 let x_field = Arc::new(Field::new("x", DataType::Int32, true));
1986 let item_field = Arc::new(Field::new(
1987 "item",
1988 DataType::Struct(Fields::from(vec![x_field.clone()])),
1989 true,
1990 ));
1991 let companies_field = Arc::new(Field::new(
1992 "companies",
1993 DataType::List(item_field.clone()),
1994 true,
1995 ));
1996 let schema = Arc::new(Schema::new(vec![Field::new(
1997 "outer",
1998 DataType::Struct(Fields::from(vec![companies_field.clone()])),
1999 true,
2000 )]));
2001
2002 let build_outer = |x: i32| {
2003 let item_struct = Arc::new(StructArray::new(
2005 Fields::from(vec![x_field.clone()]),
2006 vec![Arc::new(Int32Array::from(vec![x]))],
2007 None,
2008 ));
2009 let companies = Arc::new(ListArray::new(
2010 item_field.clone(),
2011 OffsetBuffer::from_lengths([1]),
2012 item_struct,
2013 None,
2014 ));
2015 StructArray::new(
2016 Fields::from(vec![companies_field.clone()]),
2017 vec![companies],
2018 None,
2019 )
2020 };
2021
2022 let left = RecordBatch::try_new(schema.clone(), vec![Arc::new(build_outer(10))]).unwrap();
2024 let right = RecordBatch::try_new(schema, vec![Arc::new(build_outer(20))]).unwrap();
2025
2026 let merged = left.merge(&right).unwrap();
2027
2028 assert_eq!(merged.column(0).as_struct().num_columns(), 1);
2031 assert_eq!(merged, left);
2032 }
2033
2034 #[test]
2035 fn test_byte_width_opt() {
2036 assert_eq!(DataType::Int32.byte_width_opt(), Some(4));
2037 assert_eq!(DataType::Int64.byte_width_opt(), Some(8));
2038 assert_eq!(DataType::Float32.byte_width_opt(), Some(4));
2039 assert_eq!(DataType::Float64.byte_width_opt(), Some(8));
2040 assert_eq!(DataType::Utf8.byte_width_opt(), None);
2041 assert_eq!(DataType::Binary.byte_width_opt(), None);
2042 assert_eq!(
2043 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))).byte_width_opt(),
2044 None
2045 );
2046 assert_eq!(
2047 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 3)
2048 .byte_width_opt(),
2049 Some(12)
2050 );
2051 assert_eq!(
2052 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 4)
2053 .byte_width_opt(),
2054 Some(16)
2055 );
2056 assert_eq!(
2057 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Utf8, true)), 5)
2058 .byte_width_opt(),
2059 None
2060 );
2061 }
2062
2063 #[test]
2064 fn test_take_record_batch() {
2065 let schema = Arc::new(Schema::new(vec![
2066 Field::new("a", DataType::Int32, true),
2067 Field::new("b", DataType::Utf8, true),
2068 ]));
2069 let batch = RecordBatch::try_new(
2070 schema.clone(),
2071 vec![
2072 Arc::new(Int32Array::from_iter_values(0..20)),
2073 Arc::new(StringArray::from_iter_values(
2074 (0..20).map(|i| format!("str-{}", i)),
2075 )),
2076 ],
2077 )
2078 .unwrap();
2079 let taken = batch.take(&(vec![1_u32, 5_u32, 10_u32].into())).unwrap();
2080 assert_eq!(
2081 taken,
2082 RecordBatch::try_new(
2083 schema,
2084 vec![
2085 Arc::new(Int32Array::from(vec![1, 5, 10])),
2086 Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])),
2087 ],
2088 )
2089 .unwrap()
2090 )
2091 }
2092
2093 #[test]
2094 fn test_schema_project_by_schema() {
2095 let metadata = [("key".to_string(), "value".to_string())];
2096 let schema = Arc::new(
2097 Schema::new(vec![
2098 Field::new("a", DataType::Int32, true),
2099 Field::new("b", DataType::Utf8, true),
2100 ])
2101 .with_metadata(metadata.clone().into()),
2102 );
2103 let batch = RecordBatch::try_new(
2104 schema,
2105 vec![
2106 Arc::new(Int32Array::from_iter_values(0..20)),
2107 Arc::new(StringArray::from_iter_values(
2108 (0..20).map(|i| format!("str-{}", i)),
2109 )),
2110 ],
2111 )
2112 .unwrap();
2113
2114 let empty_schema = Schema::empty();
2116 let empty_projected = batch.project_by_schema(&empty_schema).unwrap();
2117 let expected_schema = empty_schema.with_metadata(metadata.clone().into());
2118 assert_eq!(
2119 empty_projected,
2120 RecordBatch::from(StructArray::new_empty_fields(batch.num_rows(), None))
2121 .with_schema(Arc::new(expected_schema))
2122 .unwrap()
2123 );
2124
2125 let reordered_schema = Schema::new(vec![
2127 Field::new("b", DataType::Utf8, true),
2128 Field::new("a", DataType::Int32, true),
2129 ]);
2130 let reordered_projected = batch.project_by_schema(&reordered_schema).unwrap();
2131 let expected_schema = Arc::new(reordered_schema.with_metadata(metadata.clone().into()));
2132 assert_eq!(
2133 reordered_projected,
2134 RecordBatch::try_new(
2135 expected_schema,
2136 vec![
2137 Arc::new(StringArray::from_iter_values(
2138 (0..20).map(|i| format!("str-{}", i)),
2139 )),
2140 Arc::new(Int32Array::from_iter_values(0..20)),
2141 ],
2142 )
2143 .unwrap()
2144 );
2145
2146 let sub_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
2148 let sub_projected = batch.project_by_schema(&sub_schema).unwrap();
2149 let expected_schema = Arc::new(sub_schema.with_metadata(metadata.into()));
2150 assert_eq!(
2151 sub_projected,
2152 RecordBatch::try_new(
2153 expected_schema,
2154 vec![Arc::new(Int32Array::from_iter_values(0..20))],
2155 )
2156 .unwrap()
2157 );
2158 }
2159
2160 #[test]
2161 fn test_project_rebuilds_sliced_map() {
2162 let entry_fields = Fields::from(vec![
2166 Field::new("key", DataType::Utf8, false),
2167 Field::new("value", DataType::Int32, true),
2168 ]);
2169 let entries_field = Arc::new(Field::new(
2170 "entries",
2171 DataType::Struct(entry_fields.clone()),
2172 false,
2173 ));
2174 let entries = StructArray::new(
2175 entry_fields,
2176 vec![
2177 Arc::new(StringArray::from(vec!["k0", "k1", "k2"])) as ArrayRef,
2178 Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef,
2179 ],
2180 None,
2181 );
2182 let map = MapArray::new(
2183 entries_field.clone(),
2184 OffsetBuffer::new(vec![0, 1, 1, 3].into()),
2185 entries,
2186 Some(arrow_buffer::NullBuffer::from(vec![true, false, true])),
2187 false,
2188 );
2189 let schema = Arc::new(Schema::new(vec![Field::new(
2190 "m",
2191 DataType::Map(entries_field, false),
2192 true,
2193 )]));
2194 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(map) as ArrayRef]).unwrap();
2195
2196 let projected = batch
2198 .slice(1, 2)
2199 .project_by_schema(schema.as_ref())
2200 .unwrap();
2201 let map = projected.column(0).as_map();
2202 assert!(map.is_null(0));
2203 assert!(map.is_valid(1));
2204 assert_eq!(map.value_length(1), 2);
2205 assert_eq!(
2206 map.value(1)
2207 .column(1)
2208 .as_primitive::<arrow_array::types::Int32Type>()
2209 .values(),
2210 &[2, 3]
2211 );
2212 }
2213
2214 #[test]
2215 fn test_project_preserves_struct_validity() {
2216 let fields = Fields::from(vec![
2218 Field::new("id", DataType::Int32, false),
2219 Field::new("value", DataType::Float32, true),
2220 ]);
2221
2222 let id_array = Int32Array::from(vec![1, 2, 3]);
2224 let value_array = Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]);
2225 let struct_array = StructArray::new(
2226 fields.clone(),
2227 vec![
2228 Arc::new(id_array) as ArrayRef,
2229 Arc::new(value_array) as ArrayRef,
2230 ],
2231 Some(vec![true, false, true].into()), );
2233
2234 let projected = project(&struct_array, &fields).unwrap();
2236
2237 assert_eq!(projected.null_count(), 1);
2239 assert!(!projected.is_null(0));
2240 assert!(projected.is_null(1));
2241 assert!(!projected.is_null(2));
2242 }
2243
2244 #[test]
2245 fn test_merge_struct_with_different_validity() {
2246 let height_array = Int32Array::from(vec![Some(500), None, Some(600), None]);
2249 let left_fields = Fields::from(vec![Field::new("height", DataType::Int32, true)]);
2250 let left_struct = StructArray::new(
2251 left_fields,
2252 vec![Arc::new(height_array) as ArrayRef],
2253 Some(vec![true, false, true, false].into()), );
2255
2256 let width_array = Int32Array::from(vec![Some(300), Some(200), None, None]);
2258 let right_fields = Fields::from(vec![Field::new("width", DataType::Int32, true)]);
2259 let right_struct = StructArray::new(
2260 right_fields,
2261 vec![Arc::new(width_array) as ArrayRef],
2262 Some(vec![true, true, false, false].into()), );
2264
2265 let merged = merge(&left_struct, &right_struct);
2267
2268 assert_eq!(merged.null_count(), 1); assert!(!merged.is_null(0));
2276 assert!(!merged.is_null(1));
2277 assert!(!merged.is_null(2));
2278 assert!(merged.is_null(3));
2279
2280 let height_col = merged.column_by_name("height").unwrap();
2282 let height_values = height_col.as_any().downcast_ref::<Int32Array>().unwrap();
2283 assert_eq!(height_values.value(0), 500);
2284 assert!(height_values.is_null(1)); assert_eq!(height_values.value(2), 600);
2286
2287 let width_col = merged.column_by_name("width").unwrap();
2288 let width_values = width_col.as_any().downcast_ref::<Int32Array>().unwrap();
2289 assert_eq!(width_values.value(0), 300);
2290 assert_eq!(width_values.value(1), 200);
2291 assert!(width_values.is_null(2)); let all_null_left = StructArray::new(
2296 Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2297 vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef],
2298 Some(vec![false, false].into()),
2299 );
2300 let all_null_right = StructArray::new(
2301 Fields::from(vec![Field::new("width", DataType::Int32, true)]),
2302 vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef],
2303 Some(vec![false, false].into()),
2304 );
2305
2306 let merged = merge(&all_null_left, &all_null_right);
2307 assert_eq!(merged.null_count(), 2);
2308
2309 let partial_left = StructArray::new(
2311 Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2312 vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef],
2313 Some(vec![true, false].into()),
2314 );
2315 let merged = merge(&partial_left, &all_null_right);
2316 assert!(!merged.is_null(0));
2317 assert!(merged.is_null(1));
2318
2319 let all_valid_left = StructArray::new(
2321 Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2322 vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef],
2323 None,
2324 );
2325 let merged = merge(&all_valid_left, &all_null_right);
2326 assert_eq!(merged.null_count(), 0);
2327
2328 let all_valid: arrow_buffer::NullBuffer = vec![true, true].into();
2330 let partial: arrow_buffer::NullBuffer = vec![true, false].into();
2331 assert!(merge_struct_validity(Some(&all_valid), Some(&partial)).is_none());
2332 assert!(merge_struct_validity(Some(&partial), Some(&all_valid)).is_none());
2333 }
2334
2335 #[test]
2336 fn test_merge_null_typed_column_with_parent_validity() {
2337 let left_struct = StructArray::new(
2341 Fields::from(vec![Field::new("a", DataType::Int32, true)]),
2342 vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef],
2343 Some(vec![true, false].into()),
2344 );
2345 let right_struct = StructArray::new(
2346 Fields::from(vec![Field::new("b", DataType::Null, true)]),
2347 vec![Arc::new(NullArray::new(2)) as ArrayRef],
2348 Some(vec![true, false].into()),
2349 );
2350
2351 let merged = merge(&left_struct, &right_struct);
2353 assert_eq!(merged.len(), 2);
2354 let b_col = merged.column_by_name("b").unwrap();
2355 assert_eq!(b_col.data_type(), &DataType::Null);
2357 assert_eq!(b_col.len(), 2);
2358 }
2359
2360 #[test]
2361 fn test_merge_with_schema_with_nullable_struct_list_schema_mismatch() {
2362 let left_company_id = Arc::new(Int32Array::from(vec![None, None]));
2364 let left_count = Arc::new(Int32Array::from(vec![None, None]));
2365 let left_struct = Arc::new(StructArray::new(
2366 Fields::from(vec![
2367 Field::new("company_id", DataType::Int32, true),
2368 Field::new("count", DataType::Int32, true),
2369 ]),
2370 vec![left_company_id, left_count],
2371 None,
2372 ));
2373 let left_list = Arc::new(ListArray::new(
2374 Arc::new(Field::new(
2375 "item",
2376 DataType::Struct(left_struct.fields().clone()),
2377 true,
2378 )),
2379 OffsetBuffer::from_lengths([2]),
2380 left_struct,
2381 None,
2382 ));
2383
2384 let right_company_name = Arc::new(StringArray::from(vec!["Google", "Microsoft"]));
2386 let right_struct = Arc::new(StructArray::new(
2387 Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2388 vec![right_company_name],
2389 None,
2390 ));
2391 let right_list = Arc::new(ListArray::new(
2392 Arc::new(Field::new(
2393 "item",
2394 DataType::Struct(right_struct.fields().clone()),
2395 true,
2396 )),
2397 OffsetBuffer::from_lengths([2]),
2398 right_struct,
2399 None,
2400 ));
2401
2402 let target_fields = Fields::from(vec![Field::new(
2403 "companies",
2404 DataType::List(Arc::new(Field::new(
2405 "item",
2406 DataType::Struct(Fields::from(vec![
2407 Field::new("company_id", DataType::Int32, true),
2408 Field::new("company_name", DataType::Utf8, true),
2409 Field::new("count", DataType::Int32, true),
2410 ])),
2411 true,
2412 ))),
2413 true,
2414 )]);
2415
2416 let left_batch = RecordBatch::try_new(
2417 Arc::new(Schema::new(vec![Field::new(
2418 "companies",
2419 left_list.data_type().clone(),
2420 true,
2421 )])),
2422 vec![left_list as ArrayRef],
2423 )
2424 .unwrap();
2425
2426 let right_batch = RecordBatch::try_new(
2427 Arc::new(Schema::new(vec![Field::new(
2428 "companies",
2429 right_list.data_type().clone(),
2430 true,
2431 )])),
2432 vec![right_list as ArrayRef],
2433 )
2434 .unwrap();
2435
2436 let merged = left_batch
2437 .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2438 .unwrap();
2439
2440 let merged_list = merged
2442 .column_by_name("companies")
2443 .unwrap()
2444 .as_any()
2445 .downcast_ref::<ListArray>()
2446 .unwrap();
2447 let merged_struct = merged_list.values().as_struct();
2448
2449 assert_eq!(merged_struct.num_columns(), 3);
2451 assert!(merged_struct.column_by_name("company_id").is_some());
2452 assert!(merged_struct.column_by_name("company_name").is_some());
2453 assert!(merged_struct.column_by_name("count").is_some());
2454
2455 let company_id = merged_struct
2457 .column_by_name("company_id")
2458 .unwrap()
2459 .as_any()
2460 .downcast_ref::<Int32Array>()
2461 .unwrap();
2462 assert!(company_id.is_null(0));
2463 assert!(company_id.is_null(1));
2464
2465 let company_name = merged_struct
2466 .column_by_name("company_name")
2467 .unwrap()
2468 .as_any()
2469 .downcast_ref::<StringArray>()
2470 .unwrap();
2471 assert_eq!(company_name.value(0), "Google");
2472 assert_eq!(company_name.value(1), "Microsoft");
2473
2474 let count = merged_struct
2475 .column_by_name("count")
2476 .unwrap()
2477 .as_any()
2478 .downcast_ref::<Int32Array>()
2479 .unwrap();
2480 assert!(count.is_null(0));
2481 assert!(count.is_null(1));
2482 }
2483
2484 #[test]
2485 fn test_merge_struct_lists() {
2486 test_merge_struct_lists_generic::<i32>();
2487 }
2488
2489 #[test]
2490 fn test_merge_struct_large_lists() {
2491 test_merge_struct_lists_generic::<i64>();
2492 }
2493
2494 fn test_merge_struct_lists_generic<O: OffsetSizeTrait>() {
2495 let left_company_id = Arc::new(Int32Array::from(vec![
2497 Some(1),
2498 Some(2),
2499 Some(3),
2500 Some(4),
2501 Some(5),
2502 Some(6),
2503 Some(7),
2504 Some(8),
2505 Some(9),
2506 Some(10),
2507 Some(11),
2508 Some(12),
2509 Some(13),
2510 Some(14),
2511 Some(15),
2512 Some(16),
2513 Some(17),
2514 Some(18),
2515 Some(19),
2516 Some(20),
2517 ]));
2518 let left_count = Arc::new(Int32Array::from(vec![
2519 Some(10),
2520 Some(20),
2521 Some(30),
2522 Some(40),
2523 Some(50),
2524 Some(60),
2525 Some(70),
2526 Some(80),
2527 Some(90),
2528 Some(100),
2529 Some(110),
2530 Some(120),
2531 Some(130),
2532 Some(140),
2533 Some(150),
2534 Some(160),
2535 Some(170),
2536 Some(180),
2537 Some(190),
2538 Some(200),
2539 ]));
2540 let left_struct = Arc::new(StructArray::new(
2541 Fields::from(vec![
2542 Field::new("company_id", DataType::Int32, true),
2543 Field::new("count", DataType::Int32, true),
2544 ]),
2545 vec![left_company_id, left_count],
2546 None,
2547 ));
2548
2549 let left_list = Arc::new(GenericListArray::<O>::new(
2550 Arc::new(Field::new(
2551 "item",
2552 DataType::Struct(left_struct.fields().clone()),
2553 true,
2554 )),
2555 OffsetBuffer::from_lengths([3, 1]),
2556 left_struct.clone(),
2557 None,
2558 ));
2559
2560 let left_list_struct = Arc::new(StructArray::new(
2561 Fields::from(vec![Field::new(
2562 "companies",
2563 if O::IS_LARGE {
2564 DataType::LargeList(Arc::new(Field::new(
2565 "item",
2566 DataType::Struct(left_struct.fields().clone()),
2567 true,
2568 )))
2569 } else {
2570 DataType::List(Arc::new(Field::new(
2571 "item",
2572 DataType::Struct(left_struct.fields().clone()),
2573 true,
2574 )))
2575 },
2576 true,
2577 )]),
2578 vec![left_list as ArrayRef],
2579 None,
2580 ));
2581
2582 let right_company_name = Arc::new(StringArray::from(vec![
2584 "Google",
2585 "Microsoft",
2586 "Apple",
2587 "Facebook",
2588 ]));
2589 let right_struct = Arc::new(StructArray::new(
2590 Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2591 vec![right_company_name],
2592 None,
2593 ));
2594 let right_list = Arc::new(GenericListArray::<O>::new(
2595 Arc::new(Field::new(
2596 "item",
2597 DataType::Struct(right_struct.fields().clone()),
2598 true,
2599 )),
2600 OffsetBuffer::from_lengths([3, 1]),
2601 right_struct.clone(),
2602 None,
2603 ));
2604
2605 let right_list_struct = Arc::new(StructArray::new(
2606 Fields::from(vec![Field::new(
2607 "companies",
2608 if O::IS_LARGE {
2609 DataType::LargeList(Arc::new(Field::new(
2610 "item",
2611 DataType::Struct(right_struct.fields().clone()),
2612 true,
2613 )))
2614 } else {
2615 DataType::List(Arc::new(Field::new(
2616 "item",
2617 DataType::Struct(right_struct.fields().clone()),
2618 true,
2619 )))
2620 },
2621 true,
2622 )]),
2623 vec![right_list as ArrayRef],
2624 None,
2625 ));
2626
2627 let target_fields = Fields::from(vec![Field::new(
2629 "companies",
2630 if O::IS_LARGE {
2631 DataType::LargeList(Arc::new(Field::new(
2632 "item",
2633 DataType::Struct(Fields::from(vec![
2634 Field::new("company_id", DataType::Int32, true),
2635 Field::new("company_name", DataType::Utf8, true),
2636 Field::new("count", DataType::Int32, true),
2637 ])),
2638 true,
2639 )))
2640 } else {
2641 DataType::List(Arc::new(Field::new(
2642 "item",
2643 DataType::Struct(Fields::from(vec![
2644 Field::new("company_id", DataType::Int32, true),
2645 Field::new("company_name", DataType::Utf8, true),
2646 Field::new("count", DataType::Int32, true),
2647 ])),
2648 true,
2649 )))
2650 },
2651 true,
2652 )]);
2653
2654 let merged_array = merge_with_schema(&left_list_struct, &right_list_struct, &target_fields);
2656 assert_eq!(merged_array.len(), 2);
2657 }
2658
2659 #[test]
2660 fn test_merge_with_schema_sliced_list_struct() {
2661 test_merge_with_schema_sliced_list_struct_generic::<i32>();
2662 }
2663
2664 #[test]
2665 fn test_merge_with_schema_sliced_large_list_struct() {
2666 test_merge_with_schema_sliced_list_struct_generic::<i64>();
2667 }
2668
2669 fn test_merge_with_schema_sliced_list_struct_generic<O: OffsetSizeTrait>() {
2674 let make_list_dtype = |item_field: Arc<Field>| {
2675 if O::IS_LARGE {
2676 DataType::LargeList(item_field)
2677 } else {
2678 DataType::List(item_field)
2679 }
2680 };
2681
2682 let struct_fields_a = Fields::from(vec![Field::new("a", DataType::Int32, true)]);
2685 let left_values = Arc::new(StructArray::new(
2686 struct_fields_a.clone(),
2687 vec![Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef],
2688 None,
2689 ));
2690 let full_list = GenericListArray::<O>::new(
2691 Arc::new(Field::new("item", DataType::Struct(struct_fields_a), true)),
2692 OffsetBuffer::<O>::from_lengths([5, 5]),
2693 left_values,
2694 None,
2695 );
2696 let sliced_left = full_list.slice(1, 1);
2697 assert_eq!(sliced_left.offsets()[0].as_usize(), 5);
2698 assert_eq!(sliced_left.offsets()[1].as_usize(), 10);
2699
2700 let struct_fields_b = Fields::from(vec![Field::new("b", DataType::Int32, true)]);
2701 let right_values = Arc::new(StructArray::new(
2702 struct_fields_b.clone(),
2703 vec![Arc::new(Int32Array::from_iter_values(100..105)) as ArrayRef],
2704 None,
2705 ));
2706 let right_list = GenericListArray::<O>::new(
2707 Arc::new(Field::new("item", DataType::Struct(struct_fields_b), true)),
2708 OffsetBuffer::<O>::from_lengths([5]),
2709 right_values,
2710 None,
2711 );
2712
2713 let target_item_field = Arc::new(Field::new(
2714 "item",
2715 DataType::Struct(Fields::from(vec![
2716 Field::new("a", DataType::Int32, true),
2717 Field::new("b", DataType::Int32, true),
2718 ])),
2719 true,
2720 ));
2721 let target_fields = Fields::from(vec![Field::new(
2722 "items",
2723 make_list_dtype(target_item_field),
2724 true,
2725 )]);
2726
2727 let left_batch = RecordBatch::try_new(
2728 Arc::new(Schema::new(vec![Field::new(
2729 "items",
2730 sliced_left.data_type().clone(),
2731 true,
2732 )])),
2733 vec![Arc::new(sliced_left) as ArrayRef],
2734 )
2735 .unwrap();
2736 let right_batch = RecordBatch::try_new(
2737 Arc::new(Schema::new(vec![Field::new(
2738 "items",
2739 right_list.data_type().clone(),
2740 true,
2741 )])),
2742 vec![Arc::new(right_list) as ArrayRef],
2743 )
2744 .unwrap();
2745
2746 let merged = left_batch
2747 .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2748 .unwrap();
2749
2750 let merged_list = merged
2751 .column_by_name("items")
2752 .unwrap()
2753 .as_any()
2754 .downcast_ref::<GenericListArray<O>>()
2755 .unwrap();
2756 assert_eq!(merged_list.len(), 1);
2757 assert_eq!(merged_list.value_length(0).as_usize(), 5);
2758 let merged_struct = merged_list.values().as_struct();
2759 assert_eq!(merged_struct.num_columns(), 2);
2760 let a = merged_struct
2761 .column_by_name("a")
2762 .unwrap()
2763 .as_any()
2764 .downcast_ref::<Int32Array>()
2765 .unwrap();
2766 let a_vals: Vec<i32> = a.iter().map(|v| v.unwrap()).collect();
2768 assert_eq!(a_vals, vec![5, 6, 7, 8, 9]);
2769 }
2770
2771 #[test]
2772 fn test_project_by_schema_list_struct_reorder() {
2773 let source_inner_struct = DataType::Struct(Fields::from(vec![
2778 Field::new("c", DataType::Utf8, true),
2779 Field::new("b", DataType::Utf8, true),
2780 Field::new("a", DataType::Utf8, true),
2781 ]));
2782 let source_schema = Arc::new(Schema::new(vec![
2783 Field::new("id", DataType::Int32, false),
2784 Field::new(
2785 "data",
2786 DataType::List(Arc::new(Field::new(
2787 "item",
2788 source_inner_struct.clone(),
2789 true,
2790 ))),
2791 true,
2792 ),
2793 ]));
2794
2795 let c_array = StringArray::from(vec!["c1", "c2"]);
2797 let b_array = StringArray::from(vec!["b1", "b2"]);
2798 let a_array = StringArray::from(vec!["a1", "a2"]);
2799 let inner_struct = StructArray::from(vec![
2800 (
2801 Arc::new(Field::new("c", DataType::Utf8, true)),
2802 Arc::new(c_array) as ArrayRef,
2803 ),
2804 (
2805 Arc::new(Field::new("b", DataType::Utf8, true)),
2806 Arc::new(b_array) as ArrayRef,
2807 ),
2808 (
2809 Arc::new(Field::new("a", DataType::Utf8, true)),
2810 Arc::new(a_array) as ArrayRef,
2811 ),
2812 ]);
2813
2814 let list_array = ListArray::new(
2815 Arc::new(Field::new("item", source_inner_struct, true)),
2816 OffsetBuffer::from_lengths([1, 1]),
2817 Arc::new(inner_struct),
2818 None,
2819 );
2820
2821 let batch = RecordBatch::try_new(
2822 source_schema,
2823 vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(list_array)],
2824 )
2825 .unwrap();
2826
2827 let target_inner_struct = DataType::Struct(Fields::from(vec![
2829 Field::new("a", DataType::Utf8, true),
2830 Field::new("b", DataType::Utf8, true),
2831 Field::new("c", DataType::Utf8, true),
2832 ]));
2833 let target_schema = Schema::new(vec![
2834 Field::new("id", DataType::Int32, false),
2835 Field::new(
2836 "data",
2837 DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
2838 true,
2839 ),
2840 ]);
2841
2842 let projected = batch.project_by_schema(&target_schema).unwrap();
2844
2845 assert_eq!(projected.schema().as_ref(), &target_schema);
2847
2848 let projected_list = projected.column(1).as_list::<i32>();
2850 let projected_struct = projected_list.values().as_struct();
2851
2852 assert_eq!(
2854 projected_struct.column_by_name("a").unwrap().as_ref(),
2855 &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2856 );
2857 assert_eq!(
2858 projected_struct.column_by_name("b").unwrap().as_ref(),
2859 &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2860 );
2861 assert_eq!(
2862 projected_struct.column_by_name("c").unwrap().as_ref(),
2863 &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2864 );
2865
2866 assert_eq!(
2868 projected_struct.column(0).as_ref(),
2869 &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2870 );
2871 assert_eq!(
2872 projected_struct.column(1).as_ref(),
2873 &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2874 );
2875 assert_eq!(
2876 projected_struct.column(2).as_ref(),
2877 &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2878 );
2879 }
2880
2881 #[test]
2882 fn test_project_by_schema_nested_list_struct() {
2883 let inner_struct = DataType::Struct(Fields::from(vec![
2885 Field::new("y", DataType::Int32, true),
2886 Field::new("x", DataType::Int32, true),
2887 ]));
2888 let source_schema = Arc::new(Schema::new(vec![Field::new(
2889 "outer",
2890 DataType::List(Arc::new(Field::new(
2891 "item",
2892 DataType::Struct(Fields::from(vec![
2893 Field::new("b", DataType::Utf8, true),
2894 Field::new(
2895 "inner_list",
2896 DataType::List(Arc::new(Field::new("item", inner_struct.clone(), true))),
2897 true,
2898 ),
2899 Field::new("a", DataType::Utf8, true),
2900 ])),
2901 true,
2902 ))),
2903 true,
2904 )]));
2905
2906 let y_array = Int32Array::from(vec![1, 2]);
2908 let x_array = Int32Array::from(vec![3, 4]);
2909 let innermost_struct = StructArray::from(vec![
2910 (
2911 Arc::new(Field::new("y", DataType::Int32, true)),
2912 Arc::new(y_array) as ArrayRef,
2913 ),
2914 (
2915 Arc::new(Field::new("x", DataType::Int32, true)),
2916 Arc::new(x_array) as ArrayRef,
2917 ),
2918 ]);
2919 let inner_list = ListArray::new(
2920 Arc::new(Field::new("item", inner_struct.clone(), true)),
2921 OffsetBuffer::from_lengths([2]),
2922 Arc::new(innermost_struct),
2923 None,
2924 );
2925
2926 let b_array = StringArray::from(vec!["b1"]);
2927 let a_array = StringArray::from(vec!["a1"]);
2928 let middle_struct = StructArray::from(vec![
2929 (
2930 Arc::new(Field::new("b", DataType::Utf8, true)),
2931 Arc::new(b_array) as ArrayRef,
2932 ),
2933 (
2934 Arc::new(Field::new(
2935 "inner_list",
2936 DataType::List(Arc::new(Field::new("item", inner_struct, true))),
2937 true,
2938 )),
2939 Arc::new(inner_list) as ArrayRef,
2940 ),
2941 (
2942 Arc::new(Field::new("a", DataType::Utf8, true)),
2943 Arc::new(a_array) as ArrayRef,
2944 ),
2945 ]);
2946
2947 let outer_list = ListArray::new(
2948 Arc::new(Field::new("item", middle_struct.data_type().clone(), true)),
2949 OffsetBuffer::from_lengths([1]),
2950 Arc::new(middle_struct),
2951 None,
2952 );
2953
2954 let batch =
2955 RecordBatch::try_new(source_schema, vec![Arc::new(outer_list) as ArrayRef]).unwrap();
2956
2957 let target_inner_struct = DataType::Struct(Fields::from(vec![
2959 Field::new("x", DataType::Int32, true), Field::new("y", DataType::Int32, true),
2961 ]));
2962 let target_schema = Schema::new(vec![Field::new(
2963 "outer",
2964 DataType::List(Arc::new(Field::new(
2965 "item",
2966 DataType::Struct(Fields::from(vec![
2967 Field::new("a", DataType::Utf8, true), Field::new(
2969 "inner_list",
2970 DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
2971 true,
2972 ),
2973 Field::new("b", DataType::Utf8, true),
2974 ])),
2975 true,
2976 ))),
2977 true,
2978 )]);
2979
2980 let projected = batch.project_by_schema(&target_schema).unwrap();
2981
2982 assert_eq!(projected.schema().as_ref(), &target_schema);
2984
2985 let outer_list = projected.column(0).as_list::<i32>();
2987 let middle_struct = outer_list.values().as_struct();
2988
2989 assert_eq!(
2991 middle_struct.column(0).as_ref(),
2992 &StringArray::from(vec!["a1"]) as &dyn Array
2993 );
2994 assert_eq!(
2995 middle_struct.column(2).as_ref(),
2996 &StringArray::from(vec!["b1"]) as &dyn Array
2997 );
2998
2999 let inner_list = middle_struct.column(1).as_list::<i32>();
3001 let innermost_struct = inner_list.values().as_struct();
3002 assert_eq!(
3003 innermost_struct.column(0).as_ref(),
3004 &Int32Array::from(vec![3, 4]) as &dyn Array
3005 );
3006 assert_eq!(
3007 innermost_struct.column(1).as_ref(),
3008 &Int32Array::from(vec![1, 2]) as &dyn Array
3009 );
3010 }
3011
3012 #[test]
3013 fn test_iter_binary_array_accepts_binary_variants() {
3014 let binary = BinaryArray::from(vec![b"a".as_slice(), b"bc"]);
3015 assert_eq!(
3016 iter_binary_array(&binary).unwrap().collect::<Vec<_>>(),
3017 vec![Some(b"a".as_slice()), Some(b"bc".as_slice())]
3018 );
3019
3020 let large_binary = LargeBinaryArray::from(vec![b"x".as_slice(), b"yz"]);
3021 assert_eq!(
3022 iter_binary_array(&large_binary)
3023 .unwrap()
3024 .collect::<Vec<_>>(),
3025 vec![Some(b"x".as_slice()), Some(b"yz".as_slice())]
3026 );
3027
3028 let binary_view = BinaryViewArray::from(vec![b"1".as_slice(), b"23"]);
3029 assert_eq!(
3030 iter_binary_array(&binary_view).unwrap().collect::<Vec<_>>(),
3031 vec![Some(b"1".as_slice()), Some(b"23".as_slice())]
3032 );
3033
3034 let fixed_size =
3035 FixedSizeBinaryArray::try_from_iter([b"abcd", b"efgh"].into_iter()).unwrap();
3036 assert_eq!(
3037 iter_binary_array(&fixed_size).unwrap().collect::<Vec<_>>(),
3038 vec![Some(b"abcd".as_slice()), Some(b"efgh".as_slice())]
3039 );
3040 }
3041
3042 #[test]
3043 fn test_iter_binary_array_rejects_non_binary() {
3044 let int_array = Int32Array::from(vec![1, 2, 3]);
3045 let Err(error) = iter_binary_array(&int_array) else {
3046 panic!("expected an error for non-binary array");
3047 };
3048 assert!(error.to_string().contains("Expecting a binary type"));
3049 }
3050}