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, 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>;
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
478pub trait RecordBatchExt {
480 fn try_with_column(&self, field: Field, arr: ArrayRef) -> Result<RecordBatch>;
510
511 fn try_with_column_at(&self, index: usize, field: Field, arr: ArrayRef) -> Result<RecordBatch>;
513
514 fn try_new_from_struct_array(&self, arr: StructArray) -> Result<RecordBatch>;
518
519 fn merge(&self, other: &RecordBatch) -> Result<RecordBatch>;
564
565 fn merge_with_schema(&self, other: &RecordBatch, schema: &Schema) -> Result<RecordBatch>;
575
576 fn drop_column(&self, name: &str) -> Result<RecordBatch>;
580
581 fn replace_column_by_name(&self, name: &str, column: Arc<dyn Array>) -> Result<RecordBatch>;
583
584 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 fn rename_column(&self, index: usize, new_name: &str) -> Result<RecordBatch>;
594
595 fn column_by_qualified_name(&self, name: &str) -> Option<&ArrayRef>;
597
598 fn project_by_schema(&self, schema: &Schema) -> Result<RecordBatch>;
600
601 fn metadata(&self) -> &HashMap<String, String>;
603
604 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 fn with_metadata(&self, metadata: HashMap<String, String>) -> Result<RecordBatch>;
613
614 fn take(&self, indices: &UInt32Array) -> Result<RecordBatch>;
616
617 fn shrink_to_fit(&self) -> Result<RecordBatch>;
619
620 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 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
810fn project_array(array: &ArrayRef, target_field: &Field) -> Result<ArrayRef> {
813 match target_field.data_type() {
814 DataType::Struct(subfields) => {
815 let struct_arr = array.as_struct();
816 let projected = project(struct_arr, subfields)?;
817 Ok(Arc::new(projected))
818 }
819 DataType::List(inner_field) => {
820 let list_arr: &ListArray = array.as_list();
821 let projected_values = project_array(list_arr.values(), inner_field.as_ref())?;
822 Ok(Arc::new(ListArray::new(
823 inner_field.clone(),
824 list_arr.offsets().clone(),
825 projected_values,
826 list_arr.nulls().cloned(),
827 )))
828 }
829 DataType::LargeList(inner_field) => {
830 let list_arr: &LargeListArray = array.as_list();
831 let projected_values = project_array(list_arr.values(), inner_field.as_ref())?;
832 Ok(Arc::new(LargeListArray::new(
833 inner_field.clone(),
834 list_arr.offsets().clone(),
835 projected_values,
836 list_arr.nulls().cloned(),
837 )))
838 }
839 DataType::FixedSizeList(inner_field, size) => {
840 let list_arr = array.as_fixed_size_list();
841 let projected_values = project_array(list_arr.values(), inner_field.as_ref())?;
842 Ok(Arc::new(FixedSizeListArray::new(
843 inner_field.clone(),
844 *size,
845 projected_values,
846 list_arr.nulls().cloned(),
847 )))
848 }
849 _ => Ok(array.clone()),
850 }
851}
852
853fn project(struct_array: &StructArray, fields: &Fields) -> Result<StructArray> {
854 if fields.is_empty() {
855 return Ok(StructArray::new_empty_fields(
856 struct_array.len(),
857 struct_array.nulls().cloned(),
858 ));
859 }
860 let mut columns: Vec<ArrayRef> = vec![];
861 for field in fields.iter() {
862 if let Some(col) = struct_array.column_by_name(field.name()) {
863 let projected = project_array(col, field.as_ref())?;
864 columns.push(projected);
865 } else {
866 return Err(ArrowError::SchemaError(format!(
867 "field {} does not exist in the RecordBatch",
868 field.name()
869 )));
870 }
871 }
872 StructArray::try_new(fields.clone(), columns, struct_array.nulls().cloned())
874}
875
876fn lists_have_same_offsets_helper<T: OffsetSizeTrait>(left: &dyn Array, right: &dyn Array) -> bool {
877 let left_list: &GenericListArray<T> = left.as_list();
878 let right_list: &GenericListArray<T> = right.as_list();
879 left_list.offsets().inner() == right_list.offsets().inner()
880}
881
882fn merge_list_structs_helper<T: OffsetSizeTrait>(
883 left: &dyn Array,
884 right: &dyn Array,
885 items_field_name: impl Into<String>,
886 items_nullable: bool,
887) -> Arc<dyn Array> {
888 let left_list: &GenericListArray<T> = left.as_list();
889 let right_list: &GenericListArray<T> = right.as_list();
890 let left_struct = left_list.values();
891 let right_struct = right_list.values();
892 let left_struct_arr = left_struct.as_struct();
893 let right_struct_arr = right_struct.as_struct();
894 let merged_items = Arc::new(merge(left_struct_arr, right_struct_arr));
895 let items_field = Arc::new(Field::new(
896 items_field_name,
897 merged_items.data_type().clone(),
898 items_nullable,
899 ));
900 Arc::new(GenericListArray::<T>::new(
901 items_field,
902 left_list.offsets().clone(),
903 merged_items,
904 left_list.nulls().cloned(),
905 ))
906}
907
908fn merge_list_struct_null_helper<T: OffsetSizeTrait>(
909 left: &dyn Array,
910 right: &dyn Array,
911 not_null: &dyn Array,
912 items_field_name: impl Into<String>,
913) -> Arc<dyn Array> {
914 let left_list: &GenericListArray<T> = left.as_list::<T>();
915 let not_null_list = not_null.as_list::<T>();
916 let right_list = right.as_list::<T>();
917
918 let left_struct = left_list.values().as_struct();
919 let not_null_struct: &StructArray = not_null_list.values().as_struct();
920 let right_struct = right_list.values().as_struct();
921
922 let values_len = not_null_list.values().len();
923 let mut merged_fields =
924 Vec::with_capacity(not_null_struct.num_columns() + right_struct.num_columns());
925 let mut merged_columns =
926 Vec::with_capacity(not_null_struct.num_columns() + right_struct.num_columns());
927
928 for (_, field) in left_struct.columns().iter().zip(left_struct.fields()) {
929 merged_fields.push(field.clone());
930 if let Some(val) = not_null_struct.column_by_name(field.name()) {
931 merged_columns.push(val.clone());
932 } else {
933 merged_columns.push(new_null_array(field.data_type(), values_len))
934 }
935 }
936 for (_, field) in right_struct
937 .columns()
938 .iter()
939 .zip(right_struct.fields())
940 .filter(|(_, field)| left_struct.column_by_name(field.name()).is_none())
941 {
942 merged_fields.push(field.clone());
943 if let Some(val) = not_null_struct.column_by_name(field.name()) {
944 merged_columns.push(val.clone());
945 } else {
946 merged_columns.push(new_null_array(field.data_type(), values_len));
947 }
948 }
949
950 let merged_struct = Arc::new(StructArray::new(
951 Fields::from(merged_fields),
952 merged_columns,
953 not_null_struct.nulls().cloned(),
954 ));
955 let items_field = Arc::new(Field::new(
956 items_field_name,
957 merged_struct.data_type().clone(),
958 true,
959 ));
960 Arc::new(GenericListArray::<T>::new(
961 items_field,
962 not_null_list.offsets().clone(),
963 merged_struct,
964 not_null_list.nulls().cloned(),
965 ))
966}
967
968fn merge_list_struct_null(
969 left: &dyn Array,
970 right: &dyn Array,
971 not_null: &dyn Array,
972) -> Arc<dyn Array> {
973 match left.data_type() {
974 DataType::List(left_field) => {
975 merge_list_struct_null_helper::<i32>(left, right, not_null, left_field.name())
976 }
977 DataType::LargeList(left_field) => {
978 merge_list_struct_null_helper::<i64>(left, right, not_null, left_field.name())
979 }
980 _ => unreachable!(),
981 }
982}
983
984fn merge_list_struct(left: &dyn Array, right: &dyn Array) -> Arc<dyn Array> {
985 if left.null_count() == left.len() {
989 return merge_list_struct_null(left, right, right);
990 } else if right.null_count() == right.len() {
991 return merge_list_struct_null(left, right, left);
992 }
993 match (left.data_type(), right.data_type()) {
994 (DataType::List(left_field), DataType::List(_)) => {
995 if !lists_have_same_offsets_helper::<i32>(left, right) {
996 panic!("Attempt to merge list struct arrays which do not have same offsets");
997 }
998 merge_list_structs_helper::<i32>(
999 left,
1000 right,
1001 left_field.name(),
1002 left_field.is_nullable(),
1003 )
1004 }
1005 (DataType::LargeList(left_field), DataType::LargeList(_)) => {
1006 if !lists_have_same_offsets_helper::<i64>(left, right) {
1007 panic!("Attempt to merge list struct arrays which do not have same offsets");
1008 }
1009 merge_list_structs_helper::<i64>(
1010 left,
1011 right,
1012 left_field.name(),
1013 left_field.is_nullable(),
1014 )
1015 }
1016 _ => unreachable!(),
1017 }
1018}
1019
1020fn merge_struct_validity(
1025 left_validity: Option<&arrow_buffer::NullBuffer>,
1026 right_validity: Option<&arrow_buffer::NullBuffer>,
1027) -> Option<arrow_buffer::NullBuffer> {
1028 match (left_validity, right_validity) {
1029 (None, _) | (_, None) => None,
1031 (Some(left), Some(right)) => {
1032 if left.null_count() == 0 || right.null_count() == 0 {
1033 return None;
1034 }
1035 if left.null_count() == left.len() {
1036 return Some(right.clone());
1037 }
1038 if right.null_count() == right.len() {
1039 return Some(left.clone());
1040 }
1041
1042 let left_buffer = left.inner();
1043 let right_buffer = right.inner();
1044
1045 let merged_buffer = left_buffer | right_buffer;
1048
1049 Some(arrow_buffer::NullBuffer::from(merged_buffer))
1050 }
1051 }
1052}
1053
1054fn merge_list_child_values(
1055 child_field: &Field,
1056 left_values: ArrayRef,
1057 right_values: ArrayRef,
1058) -> ArrayRef {
1059 match child_field.data_type() {
1060 DataType::Struct(child_fields) => Arc::new(merge_with_schema(
1061 left_values.as_struct(),
1062 right_values.as_struct(),
1063 child_fields,
1064 )) as ArrayRef,
1065 DataType::List(grandchild) => {
1066 let left_list = left_values
1067 .as_any()
1068 .downcast_ref::<ListArray>()
1069 .expect("left list values should be ListArray");
1070 let right_list = right_values
1071 .as_any()
1072 .downcast_ref::<ListArray>()
1073 .expect("right list values should be ListArray");
1074 let merged_values = merge_list_child_values(
1075 grandchild.as_ref(),
1076 left_list.values().clone(),
1077 right_list.values().clone(),
1078 );
1079 let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1080 Arc::new(ListArray::new(
1081 grandchild.clone(),
1082 left_list.offsets().clone(),
1083 merged_values,
1084 merged_validity,
1085 )) as ArrayRef
1086 }
1087 DataType::LargeList(grandchild) => {
1088 let left_list = left_values
1089 .as_any()
1090 .downcast_ref::<LargeListArray>()
1091 .expect("left list values should be LargeListArray");
1092 let right_list = right_values
1093 .as_any()
1094 .downcast_ref::<LargeListArray>()
1095 .expect("right list values should be LargeListArray");
1096 let merged_values = merge_list_child_values(
1097 grandchild.as_ref(),
1098 left_list.values().clone(),
1099 right_list.values().clone(),
1100 );
1101 let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1102 Arc::new(LargeListArray::new(
1103 grandchild.clone(),
1104 left_list.offsets().clone(),
1105 merged_values,
1106 merged_validity,
1107 )) as ArrayRef
1108 }
1109 DataType::FixedSizeList(grandchild, list_size) => {
1110 let left_list = left_values
1111 .as_any()
1112 .downcast_ref::<FixedSizeListArray>()
1113 .expect("left list values should be FixedSizeListArray");
1114 let right_list = right_values
1115 .as_any()
1116 .downcast_ref::<FixedSizeListArray>()
1117 .expect("right list values should be FixedSizeListArray");
1118 let merged_values = merge_list_child_values(
1119 grandchild.as_ref(),
1120 left_list.values().clone(),
1121 right_list.values().clone(),
1122 );
1123 let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1124 Arc::new(FixedSizeListArray::new(
1125 grandchild.clone(),
1126 *list_size,
1127 merged_values,
1128 merged_validity,
1129 )) as ArrayRef
1130 }
1131 _ => left_values.clone(),
1132 }
1133}
1134
1135fn adjust_child_validity(
1139 child: &ArrayRef,
1140 parent_validity: Option<&arrow_buffer::NullBuffer>,
1141) -> ArrayRef {
1142 let parent_validity = match parent_validity {
1144 None => return child.clone(),
1145 Some(p) if p.null_count() == 0 => return child.clone(), Some(p) => p,
1147 };
1148
1149 if child.data_type() == &DataType::Null {
1152 return child.clone();
1153 }
1154
1155 let child_validity = child.nulls();
1156
1157 let new_validity = match child_validity {
1159 None => {
1160 parent_validity.clone()
1162 }
1163 Some(child_nulls) => {
1164 let child_buffer = child_nulls.inner();
1165 let parent_buffer = parent_validity.inner();
1166
1167 let merged_buffer = child_buffer & parent_buffer;
1170
1171 arrow_buffer::NullBuffer::from(merged_buffer)
1172 }
1173 };
1174
1175 arrow_array::make_array(
1177 arrow_data::ArrayData::try_new(
1178 child.data_type().clone(),
1179 child.len(),
1180 Some(new_validity.into_inner().into_inner()),
1181 child.offset(),
1182 child.to_data().buffers().to_vec(),
1183 child.to_data().child_data().to_vec(),
1184 )
1185 .unwrap(),
1186 )
1187}
1188
1189fn merge(left_struct_array: &StructArray, right_struct_array: &StructArray) -> StructArray {
1190 let mut fields: Vec<Field> = vec![];
1191 let mut columns: Vec<ArrayRef> = vec![];
1192 let right_fields = right_struct_array.fields();
1193 let right_columns = right_struct_array.columns();
1194
1195 let left_validity = left_struct_array.nulls();
1197 let right_validity = right_struct_array.nulls();
1198
1199 let merged_validity = merge_struct_validity(left_validity, right_validity);
1201
1202 for (left_field, left_column) in left_struct_array
1204 .fields()
1205 .iter()
1206 .zip(left_struct_array.columns().iter())
1207 {
1208 match right_fields
1209 .iter()
1210 .position(|f| f.name() == left_field.name())
1211 {
1212 Some(right_index) => {
1214 let right_field = right_fields.get(right_index).unwrap();
1215 let right_column = right_columns.get(right_index).unwrap();
1216 match (left_field.data_type(), right_field.data_type()) {
1218 (DataType::Struct(_), DataType::Struct(_)) => {
1219 let left_sub_array = left_column.as_struct();
1220 let right_sub_array = right_column.as_struct();
1221 let merged_sub_array = merge(left_sub_array, right_sub_array);
1222 fields.push(Field::new(
1223 left_field.name(),
1224 merged_sub_array.data_type().clone(),
1225 left_field.is_nullable(),
1226 ));
1227 columns.push(Arc::new(merged_sub_array) as ArrayRef);
1228 }
1229 (DataType::List(left_list), DataType::List(right_list))
1230 if left_list.data_type().is_struct()
1231 && right_list.data_type().is_struct() =>
1232 {
1233 if left_list.data_type() == right_list.data_type() {
1235 fields.push(left_field.as_ref().clone());
1236 columns.push(left_column.clone());
1237 }
1238 let merged_sub_array = merge_list_struct(&left_column, &right_column);
1242
1243 fields.push(Field::new(
1244 left_field.name(),
1245 merged_sub_array.data_type().clone(),
1246 left_field.is_nullable(),
1247 ));
1248 columns.push(merged_sub_array);
1249 }
1250 _ => {
1252 fields.push(left_field.as_ref().clone());
1254 let adjusted_column = adjust_child_validity(left_column, left_validity);
1256 columns.push(adjusted_column);
1257 }
1258 }
1259 }
1260 None => {
1261 fields.push(left_field.as_ref().clone());
1262 let adjusted_column = adjust_child_validity(left_column, left_validity);
1264 columns.push(adjusted_column);
1265 }
1266 }
1267 }
1268
1269 right_fields
1271 .iter()
1272 .zip(right_columns.iter())
1273 .for_each(|(field, column)| {
1274 if !left_struct_array
1276 .fields()
1277 .iter()
1278 .any(|f| f.name() == field.name())
1279 {
1280 fields.push(field.as_ref().clone());
1281 let adjusted_column = adjust_child_validity(column, right_validity);
1284 columns.push(adjusted_column);
1285 }
1286 });
1287
1288 StructArray::try_new(Fields::from(fields), columns, merged_validity).unwrap()
1289}
1290
1291fn merge_with_schema(
1292 left_struct_array: &StructArray,
1293 right_struct_array: &StructArray,
1294 fields: &Fields,
1295) -> StructArray {
1296 fn same_type_kind(left: &DataType, right: &DataType) -> bool {
1298 match (left, right) {
1299 (DataType::Struct(_), DataType::Struct(_)) => true,
1300 (DataType::Struct(_), _) => false,
1301 (_, DataType::Struct(_)) => false,
1302 _ => true,
1303 }
1304 }
1305
1306 let mut output_fields: Vec<Field> = Vec::with_capacity(fields.len());
1307 let mut columns: Vec<ArrayRef> = Vec::with_capacity(fields.len());
1308
1309 let left_fields = left_struct_array.fields();
1310 let left_columns = left_struct_array.columns();
1311 let right_fields = right_struct_array.fields();
1312 let right_columns = right_struct_array.columns();
1313
1314 let left_validity = left_struct_array.nulls();
1316 let right_validity = right_struct_array.nulls();
1317
1318 let merged_validity = merge_struct_validity(left_validity, right_validity);
1320
1321 for field in fields {
1322 let left_match_idx = left_fields.iter().position(|f| {
1323 f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1324 });
1325 let right_match_idx = right_fields.iter().position(|f| {
1326 f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1327 });
1328
1329 match (left_match_idx, right_match_idx) {
1330 (None, Some(right_idx)) => {
1331 output_fields.push(right_fields[right_idx].as_ref().clone());
1332 let adjusted_column =
1334 adjust_child_validity(&right_columns[right_idx], right_validity);
1335 columns.push(adjusted_column);
1336 }
1337 (Some(left_idx), None) => {
1338 output_fields.push(left_fields[left_idx].as_ref().clone());
1339 let adjusted_column = adjust_child_validity(&left_columns[left_idx], left_validity);
1341 columns.push(adjusted_column);
1342 }
1343 (Some(left_idx), Some(right_idx)) => {
1344 match field.data_type() {
1345 DataType::Struct(child_fields) => {
1346 let left_sub_array = left_columns[left_idx].as_struct();
1347 let right_sub_array = right_columns[right_idx].as_struct();
1348 let merged_sub_array =
1349 merge_with_schema(left_sub_array, right_sub_array, child_fields);
1350 output_fields.push(Field::new(
1351 field.name(),
1352 merged_sub_array.data_type().clone(),
1353 field.is_nullable(),
1354 ));
1355 columns.push(Arc::new(merged_sub_array) as ArrayRef);
1356 }
1357 DataType::List(child_field) => {
1358 let left_list = left_columns[left_idx]
1359 .as_any()
1360 .downcast_ref::<ListArray>()
1361 .unwrap();
1362 let right_list = right_columns[right_idx]
1363 .as_any()
1364 .downcast_ref::<ListArray>()
1365 .unwrap();
1366 let merged_values = merge_list_child_values(
1367 child_field.as_ref(),
1368 left_list.trimmed_values(),
1369 right_list.trimmed_values(),
1370 );
1371 let merged_validity =
1372 merge_struct_validity(left_list.nulls(), right_list.nulls());
1373 let merged_list = ListArray::new(
1377 child_field.clone(),
1378 left_list.trimmed_offsets(),
1379 merged_values,
1380 merged_validity,
1381 );
1382 output_fields.push(field.as_ref().clone());
1383 columns.push(Arc::new(merged_list) as ArrayRef);
1384 }
1385 DataType::LargeList(child_field) => {
1386 let left_list = left_columns[left_idx]
1387 .as_any()
1388 .downcast_ref::<LargeListArray>()
1389 .unwrap();
1390 let right_list = right_columns[right_idx]
1391 .as_any()
1392 .downcast_ref::<LargeListArray>()
1393 .unwrap();
1394 let merged_values = merge_list_child_values(
1395 child_field.as_ref(),
1396 left_list.trimmed_values(),
1397 right_list.trimmed_values(),
1398 );
1399 let merged_validity =
1400 merge_struct_validity(left_list.nulls(), right_list.nulls());
1401 let merged_list = LargeListArray::new(
1402 child_field.clone(),
1403 left_list.trimmed_offsets(),
1404 merged_values,
1405 merged_validity,
1406 );
1407 output_fields.push(field.as_ref().clone());
1408 columns.push(Arc::new(merged_list) as ArrayRef);
1409 }
1410 DataType::FixedSizeList(child_field, list_size) => {
1411 let left_list = left_columns[left_idx]
1412 .as_any()
1413 .downcast_ref::<FixedSizeListArray>()
1414 .unwrap();
1415 let right_list = right_columns[right_idx]
1416 .as_any()
1417 .downcast_ref::<FixedSizeListArray>()
1418 .unwrap();
1419 let merged_values = merge_list_child_values(
1420 child_field.as_ref(),
1421 left_list.values().clone(),
1422 right_list.values().clone(),
1423 );
1424 let merged_validity =
1425 merge_struct_validity(left_list.nulls(), right_list.nulls());
1426 let merged_list = FixedSizeListArray::new(
1427 child_field.clone(),
1428 *list_size,
1429 merged_values,
1430 merged_validity,
1431 );
1432 output_fields.push(field.as_ref().clone());
1433 columns.push(Arc::new(merged_list) as ArrayRef);
1434 }
1435 _ => {
1436 output_fields.push(left_fields[left_idx].as_ref().clone());
1437 let adjusted_column =
1439 adjust_child_validity(&left_columns[left_idx], left_validity);
1440 columns.push(adjusted_column);
1441 }
1442 }
1443 }
1444 (None, None) => {
1445 }
1447 }
1448 }
1449
1450 StructArray::try_new(Fields::from(output_fields), columns, merged_validity).unwrap()
1451}
1452
1453fn get_sub_array<'a>(array: &'a ArrayRef, components: &[&str]) -> Option<&'a ArrayRef> {
1454 if components.is_empty() {
1455 return Some(array);
1456 }
1457 if !matches!(array.data_type(), DataType::Struct(_)) {
1458 return None;
1459 }
1460 let struct_arr = array.as_struct();
1461 struct_arr
1462 .column_by_name(components[0])
1463 .and_then(|arr| get_sub_array(arr, &components[1..]))
1464}
1465
1466pub fn interleave_batches(
1470 batches: &[RecordBatch],
1471 indices: &[(usize, usize)],
1472) -> Result<RecordBatch> {
1473 let first_batch = batches.first().ok_or_else(|| {
1474 ArrowError::InvalidArgumentError("Cannot interleave zero RecordBatches".to_string())
1475 })?;
1476 let schema = first_batch.schema();
1477 let num_columns = first_batch.num_columns();
1478 let mut columns = Vec::with_capacity(num_columns);
1479 let mut chunks = Vec::with_capacity(batches.len());
1480
1481 for i in 0..num_columns {
1482 for batch in batches {
1483 chunks.push(batch.column(i).as_ref());
1484 }
1485 let new_column = interleave(&chunks, indices)?;
1486 columns.push(new_column);
1487 chunks.clear();
1488 }
1489
1490 RecordBatch::try_new(schema, columns)
1491}
1492
1493pub trait BufferExt {
1494 fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self;
1509
1510 fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self;
1519}
1520
1521fn is_pwr_two(n: u64) -> bool {
1522 n & (n - 1) == 0
1523}
1524
1525impl BufferExt for arrow_buffer::Buffer {
1526 fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self {
1527 if is_pwr_two(bytes_per_value) && bytes.as_ptr().align_offset(bytes_per_value as usize) != 0
1528 {
1529 let size_bytes = bytes.len();
1531 Self::copy_bytes_bytes(bytes, size_bytes)
1532 } else {
1533 unsafe {
1536 Self::from_custom_allocation(
1537 NonNull::new(bytes.as_ptr() as _).expect("should be a valid pointer"),
1538 bytes.len(),
1539 Arc::new(bytes),
1540 )
1541 }
1542 }
1543 }
1544
1545 fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self {
1546 assert!(size_bytes >= bytes.len());
1547 let mut buf = MutableBuffer::with_capacity(size_bytes);
1548 let to_fill = size_bytes - bytes.len();
1549 buf.extend(bytes);
1550 buf.extend(std::iter::repeat_n(0_u8, to_fill));
1551
1552 buf.shrink_to_fit();
1555
1556 Self::from(buf)
1557 }
1558}
1559
1560#[cfg(test)]
1561mod tests {
1562 use super::*;
1563 use arrow_array::{Float32Array, Int32Array, NullArray, StructArray};
1564 use arrow_array::{ListArray, StringArray, new_empty_array, new_null_array};
1565 use arrow_buffer::OffsetBuffer;
1566
1567 #[test]
1568 fn test_convert_to_floating_point_preserves_inner_nulls() {
1569 let values = Int8Array::from(vec![Some(1), None, Some(3), Some(4)]);
1574 let fsl = FixedSizeListArray::new(
1575 Arc::new(Field::new("item", DataType::Int8, true)),
1576 2,
1577 Arc::new(values),
1578 None,
1579 );
1580
1581 let converted = fsl.convert_to_floating_point().unwrap();
1582
1583 assert_eq!(converted.len(), 2);
1584 let conv_values = converted
1585 .values()
1586 .as_any()
1587 .downcast_ref::<Float32Array>()
1588 .unwrap();
1589 assert_eq!(conv_values.len(), 4);
1590 assert_eq!(conv_values.value(0), 1.0);
1591 assert!(conv_values.is_null(1));
1592 assert_eq!(conv_values.value(2), 3.0);
1593 assert_eq!(conv_values.value(3), 4.0);
1594 }
1595
1596 #[test]
1597 fn test_convert_to_floating_point_preserves_inner_nulls_f64_arm() {
1598 let values = UInt8Array::from(vec![Some(10u8), None, Some(30), Some(40)]);
1602 let fsl = FixedSizeListArray::new(
1603 Arc::new(Field::new("item", DataType::UInt8, true)),
1604 2,
1605 Arc::new(values),
1606 None,
1607 );
1608
1609 let converted = fsl.convert_to_floating_point().unwrap();
1610
1611 assert_eq!(converted.len(), 2);
1612 let conv_values = converted
1613 .values()
1614 .as_any()
1615 .downcast_ref::<Float64Array>()
1616 .unwrap();
1617 assert_eq!(conv_values.len(), 4);
1618 assert_eq!(conv_values.value(0), 10.0);
1619 assert!(conv_values.is_null(1));
1620 assert_eq!(conv_values.value(2), 30.0);
1621 assert_eq!(conv_values.value(3), 40.0);
1622 }
1623
1624 #[test]
1625 fn test_merge_recursive() {
1626 let a_array = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
1627 let e_array = Int32Array::from(vec![Some(4), Some(5), Some(6)]);
1628 let c_array = Int32Array::from(vec![Some(7), Some(8), Some(9)]);
1629 let d_array = StringArray::from(vec![Some("a"), Some("b"), Some("c")]);
1630
1631 let left_schema = Schema::new(vec![
1632 Field::new("a", DataType::Int32, true),
1633 Field::new(
1634 "b",
1635 DataType::Struct(vec![Field::new("c", DataType::Int32, true)].into()),
1636 true,
1637 ),
1638 ]);
1639 let left_batch = RecordBatch::try_new(
1640 Arc::new(left_schema),
1641 vec![
1642 Arc::new(a_array.clone()),
1643 Arc::new(StructArray::from(vec![(
1644 Arc::new(Field::new("c", DataType::Int32, true)),
1645 Arc::new(c_array.clone()) as ArrayRef,
1646 )])),
1647 ],
1648 )
1649 .unwrap();
1650
1651 let right_schema = Schema::new(vec![
1652 Field::new("e", DataType::Int32, true),
1653 Field::new(
1654 "b",
1655 DataType::Struct(vec![Field::new("d", DataType::Utf8, true)].into()),
1656 true,
1657 ),
1658 ]);
1659 let right_batch = RecordBatch::try_new(
1660 Arc::new(right_schema),
1661 vec![
1662 Arc::new(e_array.clone()),
1663 Arc::new(StructArray::from(vec![(
1664 Arc::new(Field::new("d", DataType::Utf8, true)),
1665 Arc::new(d_array.clone()) as ArrayRef,
1666 )])) as ArrayRef,
1667 ],
1668 )
1669 .unwrap();
1670
1671 let merged_schema = Schema::new(vec![
1672 Field::new("a", DataType::Int32, true),
1673 Field::new(
1674 "b",
1675 DataType::Struct(
1676 vec![
1677 Field::new("c", DataType::Int32, true),
1678 Field::new("d", DataType::Utf8, true),
1679 ]
1680 .into(),
1681 ),
1682 true,
1683 ),
1684 Field::new("e", DataType::Int32, true),
1685 ]);
1686 let merged_batch = RecordBatch::try_new(
1687 Arc::new(merged_schema),
1688 vec![
1689 Arc::new(a_array) as ArrayRef,
1690 Arc::new(StructArray::from(vec![
1691 (
1692 Arc::new(Field::new("c", DataType::Int32, true)),
1693 Arc::new(c_array) as ArrayRef,
1694 ),
1695 (
1696 Arc::new(Field::new("d", DataType::Utf8, true)),
1697 Arc::new(d_array) as ArrayRef,
1698 ),
1699 ])) as ArrayRef,
1700 Arc::new(e_array) as ArrayRef,
1701 ],
1702 )
1703 .unwrap();
1704
1705 let result = left_batch.merge(&right_batch).unwrap();
1706 assert_eq!(result, merged_batch);
1707 }
1708
1709 #[test]
1710 fn test_merge_with_schema() {
1711 fn test_batch(names: &[&str], types: &[DataType]) -> (Schema, RecordBatch) {
1712 let fields: Fields = names
1713 .iter()
1714 .zip(types)
1715 .map(|(name, ty)| Field::new(name.to_string(), ty.clone(), false))
1716 .collect();
1717 let schema = Schema::new(vec![Field::new(
1718 "struct",
1719 DataType::Struct(fields.clone()),
1720 false,
1721 )]);
1722 let children = types.iter().map(new_empty_array).collect::<Vec<_>>();
1723 let batch = RecordBatch::try_new(
1724 Arc::new(schema.clone()),
1725 vec![Arc::new(StructArray::new(fields, children, None)) as ArrayRef],
1726 );
1727 (schema, batch.unwrap())
1728 }
1729
1730 let (_, left_batch) = test_batch(&["a", "b"], &[DataType::Int32, DataType::Int64]);
1731 let (_, right_batch) = test_batch(&["c", "b"], &[DataType::Int32, DataType::Int64]);
1732 let (output_schema, _) = test_batch(
1733 &["b", "a", "c"],
1734 &[DataType::Int64, DataType::Int32, DataType::Int32],
1735 );
1736
1737 let merged = left_batch
1739 .merge_with_schema(&right_batch, &output_schema)
1740 .unwrap();
1741 assert_eq!(merged.schema().as_ref(), &output_schema);
1742
1743 let (naive_schema, _) = test_batch(
1745 &["a", "b", "c"],
1746 &[DataType::Int32, DataType::Int64, DataType::Int32],
1747 );
1748 let merged = left_batch.merge(&right_batch).unwrap();
1749 assert_eq!(merged.schema().as_ref(), &naive_schema);
1750 }
1751
1752 #[test]
1753 fn test_merge_list_struct() {
1754 let x_field = Arc::new(Field::new("x", DataType::Int32, true));
1755 let y_field = Arc::new(Field::new("y", DataType::Int32, true));
1756 let x_struct_field = Arc::new(Field::new(
1757 "item",
1758 DataType::Struct(Fields::from(vec![x_field.clone()])),
1759 true,
1760 ));
1761 let y_struct_field = Arc::new(Field::new(
1762 "item",
1763 DataType::Struct(Fields::from(vec![y_field.clone()])),
1764 true,
1765 ));
1766 let both_struct_field = Arc::new(Field::new(
1767 "item",
1768 DataType::Struct(Fields::from(vec![x_field.clone(), y_field.clone()])),
1769 true,
1770 ));
1771 let left_schema = Schema::new(vec![Field::new(
1772 "list_struct",
1773 DataType::List(x_struct_field.clone()),
1774 true,
1775 )]);
1776 let right_schema = Schema::new(vec![Field::new(
1777 "list_struct",
1778 DataType::List(y_struct_field.clone()),
1779 true,
1780 )]);
1781 let both_schema = Schema::new(vec![Field::new(
1782 "list_struct",
1783 DataType::List(both_struct_field.clone()),
1784 true,
1785 )]);
1786
1787 let x = Arc::new(Int32Array::from(vec![1]));
1788 let y = Arc::new(Int32Array::from(vec![2]));
1789 let x_struct = Arc::new(StructArray::new(
1790 Fields::from(vec![x_field.clone()]),
1791 vec![x.clone()],
1792 None,
1793 ));
1794 let y_struct = Arc::new(StructArray::new(
1795 Fields::from(vec![y_field.clone()]),
1796 vec![y.clone()],
1797 None,
1798 ));
1799 let both_struct = Arc::new(StructArray::new(
1800 Fields::from(vec![x_field.clone(), y_field.clone()]),
1801 vec![x.clone(), y],
1802 None,
1803 ));
1804 let both_null_struct = Arc::new(StructArray::new(
1805 Fields::from(vec![x_field, y_field]),
1806 vec![x, Arc::new(new_null_array(&DataType::Int32, 1))],
1807 None,
1808 ));
1809 let offsets = OffsetBuffer::from_lengths([1]);
1810 let x_s_list = ListArray::new(x_struct_field, offsets.clone(), x_struct, None);
1811 let y_s_list = ListArray::new(y_struct_field, offsets.clone(), y_struct, None);
1812 let both_list = ListArray::new(
1813 both_struct_field.clone(),
1814 offsets.clone(),
1815 both_struct,
1816 None,
1817 );
1818 let both_null_list = ListArray::new(both_struct_field, offsets, both_null_struct, None);
1819 let x_batch =
1820 RecordBatch::try_new(Arc::new(left_schema), vec![Arc::new(x_s_list)]).unwrap();
1821 let y_batch = RecordBatch::try_new(
1822 Arc::new(right_schema.clone()),
1823 vec![Arc::new(y_s_list.clone())],
1824 )
1825 .unwrap();
1826 let merged = x_batch.merge(&y_batch).unwrap();
1827 let expected =
1828 RecordBatch::try_new(Arc::new(both_schema.clone()), vec![Arc::new(both_list)]).unwrap();
1829 assert_eq!(merged, expected);
1830
1831 let y_null_list = new_null_array(y_s_list.data_type(), 1);
1832 let y_null_batch =
1833 RecordBatch::try_new(Arc::new(right_schema), vec![Arc::new(y_null_list.clone())])
1834 .unwrap();
1835 let expected =
1836 RecordBatch::try_new(Arc::new(both_schema), vec![Arc::new(both_null_list)]).unwrap();
1837 let merged = x_batch.merge(&y_null_batch).unwrap();
1838 assert_eq!(merged, expected);
1839 }
1840
1841 #[test]
1842 fn test_byte_width_opt() {
1843 assert_eq!(DataType::Int32.byte_width_opt(), Some(4));
1844 assert_eq!(DataType::Int64.byte_width_opt(), Some(8));
1845 assert_eq!(DataType::Float32.byte_width_opt(), Some(4));
1846 assert_eq!(DataType::Float64.byte_width_opt(), Some(8));
1847 assert_eq!(DataType::Utf8.byte_width_opt(), None);
1848 assert_eq!(DataType::Binary.byte_width_opt(), None);
1849 assert_eq!(
1850 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))).byte_width_opt(),
1851 None
1852 );
1853 assert_eq!(
1854 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 3)
1855 .byte_width_opt(),
1856 Some(12)
1857 );
1858 assert_eq!(
1859 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 4)
1860 .byte_width_opt(),
1861 Some(16)
1862 );
1863 assert_eq!(
1864 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Utf8, true)), 5)
1865 .byte_width_opt(),
1866 None
1867 );
1868 }
1869
1870 #[test]
1871 fn test_take_record_batch() {
1872 let schema = Arc::new(Schema::new(vec![
1873 Field::new("a", DataType::Int32, true),
1874 Field::new("b", DataType::Utf8, true),
1875 ]));
1876 let batch = RecordBatch::try_new(
1877 schema.clone(),
1878 vec![
1879 Arc::new(Int32Array::from_iter_values(0..20)),
1880 Arc::new(StringArray::from_iter_values(
1881 (0..20).map(|i| format!("str-{}", i)),
1882 )),
1883 ],
1884 )
1885 .unwrap();
1886 let taken = batch.take(&(vec![1_u32, 5_u32, 10_u32].into())).unwrap();
1887 assert_eq!(
1888 taken,
1889 RecordBatch::try_new(
1890 schema,
1891 vec![
1892 Arc::new(Int32Array::from(vec![1, 5, 10])),
1893 Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])),
1894 ],
1895 )
1896 .unwrap()
1897 )
1898 }
1899
1900 #[test]
1901 fn test_schema_project_by_schema() {
1902 let metadata = [("key".to_string(), "value".to_string())];
1903 let schema = Arc::new(
1904 Schema::new(vec![
1905 Field::new("a", DataType::Int32, true),
1906 Field::new("b", DataType::Utf8, true),
1907 ])
1908 .with_metadata(metadata.clone().into()),
1909 );
1910 let batch = RecordBatch::try_new(
1911 schema,
1912 vec![
1913 Arc::new(Int32Array::from_iter_values(0..20)),
1914 Arc::new(StringArray::from_iter_values(
1915 (0..20).map(|i| format!("str-{}", i)),
1916 )),
1917 ],
1918 )
1919 .unwrap();
1920
1921 let empty_schema = Schema::empty();
1923 let empty_projected = batch.project_by_schema(&empty_schema).unwrap();
1924 let expected_schema = empty_schema.with_metadata(metadata.clone().into());
1925 assert_eq!(
1926 empty_projected,
1927 RecordBatch::from(StructArray::new_empty_fields(batch.num_rows(), None))
1928 .with_schema(Arc::new(expected_schema))
1929 .unwrap()
1930 );
1931
1932 let reordered_schema = Schema::new(vec![
1934 Field::new("b", DataType::Utf8, true),
1935 Field::new("a", DataType::Int32, true),
1936 ]);
1937 let reordered_projected = batch.project_by_schema(&reordered_schema).unwrap();
1938 let expected_schema = Arc::new(reordered_schema.with_metadata(metadata.clone().into()));
1939 assert_eq!(
1940 reordered_projected,
1941 RecordBatch::try_new(
1942 expected_schema,
1943 vec![
1944 Arc::new(StringArray::from_iter_values(
1945 (0..20).map(|i| format!("str-{}", i)),
1946 )),
1947 Arc::new(Int32Array::from_iter_values(0..20)),
1948 ],
1949 )
1950 .unwrap()
1951 );
1952
1953 let sub_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
1955 let sub_projected = batch.project_by_schema(&sub_schema).unwrap();
1956 let expected_schema = Arc::new(sub_schema.with_metadata(metadata.into()));
1957 assert_eq!(
1958 sub_projected,
1959 RecordBatch::try_new(
1960 expected_schema,
1961 vec![Arc::new(Int32Array::from_iter_values(0..20))],
1962 )
1963 .unwrap()
1964 );
1965 }
1966
1967 #[test]
1968 fn test_project_preserves_struct_validity() {
1969 let fields = Fields::from(vec![
1971 Field::new("id", DataType::Int32, false),
1972 Field::new("value", DataType::Float32, true),
1973 ]);
1974
1975 let id_array = Int32Array::from(vec![1, 2, 3]);
1977 let value_array = Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]);
1978 let struct_array = StructArray::new(
1979 fields.clone(),
1980 vec![
1981 Arc::new(id_array) as ArrayRef,
1982 Arc::new(value_array) as ArrayRef,
1983 ],
1984 Some(vec![true, false, true].into()), );
1986
1987 let projected = project(&struct_array, &fields).unwrap();
1989
1990 assert_eq!(projected.null_count(), 1);
1992 assert!(!projected.is_null(0));
1993 assert!(projected.is_null(1));
1994 assert!(!projected.is_null(2));
1995 }
1996
1997 #[test]
1998 fn test_merge_struct_with_different_validity() {
1999 let height_array = Int32Array::from(vec![Some(500), None, Some(600), None]);
2002 let left_fields = Fields::from(vec![Field::new("height", DataType::Int32, true)]);
2003 let left_struct = StructArray::new(
2004 left_fields,
2005 vec![Arc::new(height_array) as ArrayRef],
2006 Some(vec![true, false, true, false].into()), );
2008
2009 let width_array = Int32Array::from(vec![Some(300), Some(200), None, None]);
2011 let right_fields = Fields::from(vec![Field::new("width", DataType::Int32, true)]);
2012 let right_struct = StructArray::new(
2013 right_fields,
2014 vec![Arc::new(width_array) as ArrayRef],
2015 Some(vec![true, true, false, false].into()), );
2017
2018 let merged = merge(&left_struct, &right_struct);
2020
2021 assert_eq!(merged.null_count(), 1); assert!(!merged.is_null(0));
2029 assert!(!merged.is_null(1));
2030 assert!(!merged.is_null(2));
2031 assert!(merged.is_null(3));
2032
2033 let height_col = merged.column_by_name("height").unwrap();
2035 let height_values = height_col.as_any().downcast_ref::<Int32Array>().unwrap();
2036 assert_eq!(height_values.value(0), 500);
2037 assert!(height_values.is_null(1)); assert_eq!(height_values.value(2), 600);
2039
2040 let width_col = merged.column_by_name("width").unwrap();
2041 let width_values = width_col.as_any().downcast_ref::<Int32Array>().unwrap();
2042 assert_eq!(width_values.value(0), 300);
2043 assert_eq!(width_values.value(1), 200);
2044 assert!(width_values.is_null(2)); let all_null_left = StructArray::new(
2049 Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2050 vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef],
2051 Some(vec![false, false].into()),
2052 );
2053 let all_null_right = StructArray::new(
2054 Fields::from(vec![Field::new("width", DataType::Int32, true)]),
2055 vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef],
2056 Some(vec![false, false].into()),
2057 );
2058
2059 let merged = merge(&all_null_left, &all_null_right);
2060 assert_eq!(merged.null_count(), 2);
2061
2062 let partial_left = StructArray::new(
2064 Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2065 vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef],
2066 Some(vec![true, false].into()),
2067 );
2068 let merged = merge(&partial_left, &all_null_right);
2069 assert!(!merged.is_null(0));
2070 assert!(merged.is_null(1));
2071
2072 let all_valid_left = StructArray::new(
2074 Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2075 vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef],
2076 None,
2077 );
2078 let merged = merge(&all_valid_left, &all_null_right);
2079 assert_eq!(merged.null_count(), 0);
2080
2081 let all_valid: arrow_buffer::NullBuffer = vec![true, true].into();
2083 let partial: arrow_buffer::NullBuffer = vec![true, false].into();
2084 assert!(merge_struct_validity(Some(&all_valid), Some(&partial)).is_none());
2085 assert!(merge_struct_validity(Some(&partial), Some(&all_valid)).is_none());
2086 }
2087
2088 #[test]
2089 fn test_merge_null_typed_column_with_parent_validity() {
2090 let left_struct = StructArray::new(
2094 Fields::from(vec![Field::new("a", DataType::Int32, true)]),
2095 vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef],
2096 Some(vec![true, false].into()),
2097 );
2098 let right_struct = StructArray::new(
2099 Fields::from(vec![Field::new("b", DataType::Null, true)]),
2100 vec![Arc::new(NullArray::new(2)) as ArrayRef],
2101 Some(vec![true, false].into()),
2102 );
2103
2104 let merged = merge(&left_struct, &right_struct);
2106 assert_eq!(merged.len(), 2);
2107 let b_col = merged.column_by_name("b").unwrap();
2108 assert_eq!(b_col.data_type(), &DataType::Null);
2110 assert_eq!(b_col.len(), 2);
2111 }
2112
2113 #[test]
2114 fn test_merge_with_schema_with_nullable_struct_list_schema_mismatch() {
2115 let left_company_id = Arc::new(Int32Array::from(vec![None, None]));
2117 let left_count = Arc::new(Int32Array::from(vec![None, None]));
2118 let left_struct = Arc::new(StructArray::new(
2119 Fields::from(vec![
2120 Field::new("company_id", DataType::Int32, true),
2121 Field::new("count", DataType::Int32, true),
2122 ]),
2123 vec![left_company_id, left_count],
2124 None,
2125 ));
2126 let left_list = Arc::new(ListArray::new(
2127 Arc::new(Field::new(
2128 "item",
2129 DataType::Struct(left_struct.fields().clone()),
2130 true,
2131 )),
2132 OffsetBuffer::from_lengths([2]),
2133 left_struct,
2134 None,
2135 ));
2136
2137 let right_company_name = Arc::new(StringArray::from(vec!["Google", "Microsoft"]));
2139 let right_struct = Arc::new(StructArray::new(
2140 Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2141 vec![right_company_name],
2142 None,
2143 ));
2144 let right_list = Arc::new(ListArray::new(
2145 Arc::new(Field::new(
2146 "item",
2147 DataType::Struct(right_struct.fields().clone()),
2148 true,
2149 )),
2150 OffsetBuffer::from_lengths([2]),
2151 right_struct,
2152 None,
2153 ));
2154
2155 let target_fields = Fields::from(vec![Field::new(
2156 "companies",
2157 DataType::List(Arc::new(Field::new(
2158 "item",
2159 DataType::Struct(Fields::from(vec![
2160 Field::new("company_id", DataType::Int32, true),
2161 Field::new("company_name", DataType::Utf8, true),
2162 Field::new("count", DataType::Int32, true),
2163 ])),
2164 true,
2165 ))),
2166 true,
2167 )]);
2168
2169 let left_batch = RecordBatch::try_new(
2170 Arc::new(Schema::new(vec![Field::new(
2171 "companies",
2172 left_list.data_type().clone(),
2173 true,
2174 )])),
2175 vec![left_list as ArrayRef],
2176 )
2177 .unwrap();
2178
2179 let right_batch = RecordBatch::try_new(
2180 Arc::new(Schema::new(vec![Field::new(
2181 "companies",
2182 right_list.data_type().clone(),
2183 true,
2184 )])),
2185 vec![right_list as ArrayRef],
2186 )
2187 .unwrap();
2188
2189 let merged = left_batch
2190 .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2191 .unwrap();
2192
2193 let merged_list = merged
2195 .column_by_name("companies")
2196 .unwrap()
2197 .as_any()
2198 .downcast_ref::<ListArray>()
2199 .unwrap();
2200 let merged_struct = merged_list.values().as_struct();
2201
2202 assert_eq!(merged_struct.num_columns(), 3);
2204 assert!(merged_struct.column_by_name("company_id").is_some());
2205 assert!(merged_struct.column_by_name("company_name").is_some());
2206 assert!(merged_struct.column_by_name("count").is_some());
2207
2208 let company_id = merged_struct
2210 .column_by_name("company_id")
2211 .unwrap()
2212 .as_any()
2213 .downcast_ref::<Int32Array>()
2214 .unwrap();
2215 assert!(company_id.is_null(0));
2216 assert!(company_id.is_null(1));
2217
2218 let company_name = merged_struct
2219 .column_by_name("company_name")
2220 .unwrap()
2221 .as_any()
2222 .downcast_ref::<StringArray>()
2223 .unwrap();
2224 assert_eq!(company_name.value(0), "Google");
2225 assert_eq!(company_name.value(1), "Microsoft");
2226
2227 let count = merged_struct
2228 .column_by_name("count")
2229 .unwrap()
2230 .as_any()
2231 .downcast_ref::<Int32Array>()
2232 .unwrap();
2233 assert!(count.is_null(0));
2234 assert!(count.is_null(1));
2235 }
2236
2237 #[test]
2238 fn test_merge_struct_lists() {
2239 test_merge_struct_lists_generic::<i32>();
2240 }
2241
2242 #[test]
2243 fn test_merge_struct_large_lists() {
2244 test_merge_struct_lists_generic::<i64>();
2245 }
2246
2247 fn test_merge_struct_lists_generic<O: OffsetSizeTrait>() {
2248 let left_company_id = Arc::new(Int32Array::from(vec![
2250 Some(1),
2251 Some(2),
2252 Some(3),
2253 Some(4),
2254 Some(5),
2255 Some(6),
2256 Some(7),
2257 Some(8),
2258 Some(9),
2259 Some(10),
2260 Some(11),
2261 Some(12),
2262 Some(13),
2263 Some(14),
2264 Some(15),
2265 Some(16),
2266 Some(17),
2267 Some(18),
2268 Some(19),
2269 Some(20),
2270 ]));
2271 let left_count = Arc::new(Int32Array::from(vec![
2272 Some(10),
2273 Some(20),
2274 Some(30),
2275 Some(40),
2276 Some(50),
2277 Some(60),
2278 Some(70),
2279 Some(80),
2280 Some(90),
2281 Some(100),
2282 Some(110),
2283 Some(120),
2284 Some(130),
2285 Some(140),
2286 Some(150),
2287 Some(160),
2288 Some(170),
2289 Some(180),
2290 Some(190),
2291 Some(200),
2292 ]));
2293 let left_struct = Arc::new(StructArray::new(
2294 Fields::from(vec![
2295 Field::new("company_id", DataType::Int32, true),
2296 Field::new("count", DataType::Int32, true),
2297 ]),
2298 vec![left_company_id, left_count],
2299 None,
2300 ));
2301
2302 let left_list = Arc::new(GenericListArray::<O>::new(
2303 Arc::new(Field::new(
2304 "item",
2305 DataType::Struct(left_struct.fields().clone()),
2306 true,
2307 )),
2308 OffsetBuffer::from_lengths([3, 1]),
2309 left_struct.clone(),
2310 None,
2311 ));
2312
2313 let left_list_struct = Arc::new(StructArray::new(
2314 Fields::from(vec![Field::new(
2315 "companies",
2316 if O::IS_LARGE {
2317 DataType::LargeList(Arc::new(Field::new(
2318 "item",
2319 DataType::Struct(left_struct.fields().clone()),
2320 true,
2321 )))
2322 } else {
2323 DataType::List(Arc::new(Field::new(
2324 "item",
2325 DataType::Struct(left_struct.fields().clone()),
2326 true,
2327 )))
2328 },
2329 true,
2330 )]),
2331 vec![left_list as ArrayRef],
2332 None,
2333 ));
2334
2335 let right_company_name = Arc::new(StringArray::from(vec![
2337 "Google",
2338 "Microsoft",
2339 "Apple",
2340 "Facebook",
2341 ]));
2342 let right_struct = Arc::new(StructArray::new(
2343 Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2344 vec![right_company_name],
2345 None,
2346 ));
2347 let right_list = Arc::new(GenericListArray::<O>::new(
2348 Arc::new(Field::new(
2349 "item",
2350 DataType::Struct(right_struct.fields().clone()),
2351 true,
2352 )),
2353 OffsetBuffer::from_lengths([3, 1]),
2354 right_struct.clone(),
2355 None,
2356 ));
2357
2358 let right_list_struct = Arc::new(StructArray::new(
2359 Fields::from(vec![Field::new(
2360 "companies",
2361 if O::IS_LARGE {
2362 DataType::LargeList(Arc::new(Field::new(
2363 "item",
2364 DataType::Struct(right_struct.fields().clone()),
2365 true,
2366 )))
2367 } else {
2368 DataType::List(Arc::new(Field::new(
2369 "item",
2370 DataType::Struct(right_struct.fields().clone()),
2371 true,
2372 )))
2373 },
2374 true,
2375 )]),
2376 vec![right_list as ArrayRef],
2377 None,
2378 ));
2379
2380 let target_fields = Fields::from(vec![Field::new(
2382 "companies",
2383 if O::IS_LARGE {
2384 DataType::LargeList(Arc::new(Field::new(
2385 "item",
2386 DataType::Struct(Fields::from(vec![
2387 Field::new("company_id", DataType::Int32, true),
2388 Field::new("company_name", DataType::Utf8, true),
2389 Field::new("count", DataType::Int32, true),
2390 ])),
2391 true,
2392 )))
2393 } else {
2394 DataType::List(Arc::new(Field::new(
2395 "item",
2396 DataType::Struct(Fields::from(vec![
2397 Field::new("company_id", DataType::Int32, true),
2398 Field::new("company_name", DataType::Utf8, true),
2399 Field::new("count", DataType::Int32, true),
2400 ])),
2401 true,
2402 )))
2403 },
2404 true,
2405 )]);
2406
2407 let merged_array = merge_with_schema(&left_list_struct, &right_list_struct, &target_fields);
2409 assert_eq!(merged_array.len(), 2);
2410 }
2411
2412 #[test]
2413 fn test_merge_with_schema_sliced_list_struct() {
2414 test_merge_with_schema_sliced_list_struct_generic::<i32>();
2415 }
2416
2417 #[test]
2418 fn test_merge_with_schema_sliced_large_list_struct() {
2419 test_merge_with_schema_sliced_list_struct_generic::<i64>();
2420 }
2421
2422 fn test_merge_with_schema_sliced_list_struct_generic<O: OffsetSizeTrait>() {
2427 let make_list_dtype = |item_field: Arc<Field>| {
2428 if O::IS_LARGE {
2429 DataType::LargeList(item_field)
2430 } else {
2431 DataType::List(item_field)
2432 }
2433 };
2434
2435 let struct_fields_a = Fields::from(vec![Field::new("a", DataType::Int32, true)]);
2438 let left_values = Arc::new(StructArray::new(
2439 struct_fields_a.clone(),
2440 vec![Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef],
2441 None,
2442 ));
2443 let full_list = GenericListArray::<O>::new(
2444 Arc::new(Field::new("item", DataType::Struct(struct_fields_a), true)),
2445 OffsetBuffer::<O>::from_lengths([5, 5]),
2446 left_values,
2447 None,
2448 );
2449 let sliced_left = full_list.slice(1, 1);
2450 assert_eq!(sliced_left.offsets()[0].as_usize(), 5);
2451 assert_eq!(sliced_left.offsets()[1].as_usize(), 10);
2452
2453 let struct_fields_b = Fields::from(vec![Field::new("b", DataType::Int32, true)]);
2454 let right_values = Arc::new(StructArray::new(
2455 struct_fields_b.clone(),
2456 vec![Arc::new(Int32Array::from_iter_values(100..105)) as ArrayRef],
2457 None,
2458 ));
2459 let right_list = GenericListArray::<O>::new(
2460 Arc::new(Field::new("item", DataType::Struct(struct_fields_b), true)),
2461 OffsetBuffer::<O>::from_lengths([5]),
2462 right_values,
2463 None,
2464 );
2465
2466 let target_item_field = Arc::new(Field::new(
2467 "item",
2468 DataType::Struct(Fields::from(vec![
2469 Field::new("a", DataType::Int32, true),
2470 Field::new("b", DataType::Int32, true),
2471 ])),
2472 true,
2473 ));
2474 let target_fields = Fields::from(vec![Field::new(
2475 "items",
2476 make_list_dtype(target_item_field),
2477 true,
2478 )]);
2479
2480 let left_batch = RecordBatch::try_new(
2481 Arc::new(Schema::new(vec![Field::new(
2482 "items",
2483 sliced_left.data_type().clone(),
2484 true,
2485 )])),
2486 vec![Arc::new(sliced_left) as ArrayRef],
2487 )
2488 .unwrap();
2489 let right_batch = RecordBatch::try_new(
2490 Arc::new(Schema::new(vec![Field::new(
2491 "items",
2492 right_list.data_type().clone(),
2493 true,
2494 )])),
2495 vec![Arc::new(right_list) as ArrayRef],
2496 )
2497 .unwrap();
2498
2499 let merged = left_batch
2500 .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2501 .unwrap();
2502
2503 let merged_list = merged
2504 .column_by_name("items")
2505 .unwrap()
2506 .as_any()
2507 .downcast_ref::<GenericListArray<O>>()
2508 .unwrap();
2509 assert_eq!(merged_list.len(), 1);
2510 assert_eq!(merged_list.value_length(0).as_usize(), 5);
2511 let merged_struct = merged_list.values().as_struct();
2512 assert_eq!(merged_struct.num_columns(), 2);
2513 let a = merged_struct
2514 .column_by_name("a")
2515 .unwrap()
2516 .as_any()
2517 .downcast_ref::<Int32Array>()
2518 .unwrap();
2519 let a_vals: Vec<i32> = a.iter().map(|v| v.unwrap()).collect();
2521 assert_eq!(a_vals, vec![5, 6, 7, 8, 9]);
2522 }
2523
2524 #[test]
2525 fn test_project_by_schema_list_struct_reorder() {
2526 let source_inner_struct = DataType::Struct(Fields::from(vec![
2531 Field::new("c", DataType::Utf8, true),
2532 Field::new("b", DataType::Utf8, true),
2533 Field::new("a", DataType::Utf8, true),
2534 ]));
2535 let source_schema = Arc::new(Schema::new(vec![
2536 Field::new("id", DataType::Int32, false),
2537 Field::new(
2538 "data",
2539 DataType::List(Arc::new(Field::new(
2540 "item",
2541 source_inner_struct.clone(),
2542 true,
2543 ))),
2544 true,
2545 ),
2546 ]));
2547
2548 let c_array = StringArray::from(vec!["c1", "c2"]);
2550 let b_array = StringArray::from(vec!["b1", "b2"]);
2551 let a_array = StringArray::from(vec!["a1", "a2"]);
2552 let inner_struct = StructArray::from(vec![
2553 (
2554 Arc::new(Field::new("c", DataType::Utf8, true)),
2555 Arc::new(c_array) as ArrayRef,
2556 ),
2557 (
2558 Arc::new(Field::new("b", DataType::Utf8, true)),
2559 Arc::new(b_array) as ArrayRef,
2560 ),
2561 (
2562 Arc::new(Field::new("a", DataType::Utf8, true)),
2563 Arc::new(a_array) as ArrayRef,
2564 ),
2565 ]);
2566
2567 let list_array = ListArray::new(
2568 Arc::new(Field::new("item", source_inner_struct, true)),
2569 OffsetBuffer::from_lengths([1, 1]),
2570 Arc::new(inner_struct),
2571 None,
2572 );
2573
2574 let batch = RecordBatch::try_new(
2575 source_schema,
2576 vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(list_array)],
2577 )
2578 .unwrap();
2579
2580 let target_inner_struct = DataType::Struct(Fields::from(vec![
2582 Field::new("a", DataType::Utf8, true),
2583 Field::new("b", DataType::Utf8, true),
2584 Field::new("c", DataType::Utf8, true),
2585 ]));
2586 let target_schema = Schema::new(vec![
2587 Field::new("id", DataType::Int32, false),
2588 Field::new(
2589 "data",
2590 DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
2591 true,
2592 ),
2593 ]);
2594
2595 let projected = batch.project_by_schema(&target_schema).unwrap();
2597
2598 assert_eq!(projected.schema().as_ref(), &target_schema);
2600
2601 let projected_list = projected.column(1).as_list::<i32>();
2603 let projected_struct = projected_list.values().as_struct();
2604
2605 assert_eq!(
2607 projected_struct.column_by_name("a").unwrap().as_ref(),
2608 &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2609 );
2610 assert_eq!(
2611 projected_struct.column_by_name("b").unwrap().as_ref(),
2612 &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2613 );
2614 assert_eq!(
2615 projected_struct.column_by_name("c").unwrap().as_ref(),
2616 &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2617 );
2618
2619 assert_eq!(
2621 projected_struct.column(0).as_ref(),
2622 &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2623 );
2624 assert_eq!(
2625 projected_struct.column(1).as_ref(),
2626 &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2627 );
2628 assert_eq!(
2629 projected_struct.column(2).as_ref(),
2630 &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2631 );
2632 }
2633
2634 #[test]
2635 fn test_project_by_schema_nested_list_struct() {
2636 let inner_struct = DataType::Struct(Fields::from(vec![
2638 Field::new("y", DataType::Int32, true),
2639 Field::new("x", DataType::Int32, true),
2640 ]));
2641 let source_schema = Arc::new(Schema::new(vec![Field::new(
2642 "outer",
2643 DataType::List(Arc::new(Field::new(
2644 "item",
2645 DataType::Struct(Fields::from(vec![
2646 Field::new("b", DataType::Utf8, true),
2647 Field::new(
2648 "inner_list",
2649 DataType::List(Arc::new(Field::new("item", inner_struct.clone(), true))),
2650 true,
2651 ),
2652 Field::new("a", DataType::Utf8, true),
2653 ])),
2654 true,
2655 ))),
2656 true,
2657 )]));
2658
2659 let y_array = Int32Array::from(vec![1, 2]);
2661 let x_array = Int32Array::from(vec![3, 4]);
2662 let innermost_struct = StructArray::from(vec![
2663 (
2664 Arc::new(Field::new("y", DataType::Int32, true)),
2665 Arc::new(y_array) as ArrayRef,
2666 ),
2667 (
2668 Arc::new(Field::new("x", DataType::Int32, true)),
2669 Arc::new(x_array) as ArrayRef,
2670 ),
2671 ]);
2672 let inner_list = ListArray::new(
2673 Arc::new(Field::new("item", inner_struct.clone(), true)),
2674 OffsetBuffer::from_lengths([2]),
2675 Arc::new(innermost_struct),
2676 None,
2677 );
2678
2679 let b_array = StringArray::from(vec!["b1"]);
2680 let a_array = StringArray::from(vec!["a1"]);
2681 let middle_struct = StructArray::from(vec![
2682 (
2683 Arc::new(Field::new("b", DataType::Utf8, true)),
2684 Arc::new(b_array) as ArrayRef,
2685 ),
2686 (
2687 Arc::new(Field::new(
2688 "inner_list",
2689 DataType::List(Arc::new(Field::new("item", inner_struct, true))),
2690 true,
2691 )),
2692 Arc::new(inner_list) as ArrayRef,
2693 ),
2694 (
2695 Arc::new(Field::new("a", DataType::Utf8, true)),
2696 Arc::new(a_array) as ArrayRef,
2697 ),
2698 ]);
2699
2700 let outer_list = ListArray::new(
2701 Arc::new(Field::new("item", middle_struct.data_type().clone(), true)),
2702 OffsetBuffer::from_lengths([1]),
2703 Arc::new(middle_struct),
2704 None,
2705 );
2706
2707 let batch =
2708 RecordBatch::try_new(source_schema, vec![Arc::new(outer_list) as ArrayRef]).unwrap();
2709
2710 let target_inner_struct = DataType::Struct(Fields::from(vec![
2712 Field::new("x", DataType::Int32, true), Field::new("y", DataType::Int32, true),
2714 ]));
2715 let target_schema = Schema::new(vec![Field::new(
2716 "outer",
2717 DataType::List(Arc::new(Field::new(
2718 "item",
2719 DataType::Struct(Fields::from(vec![
2720 Field::new("a", DataType::Utf8, true), Field::new(
2722 "inner_list",
2723 DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
2724 true,
2725 ),
2726 Field::new("b", DataType::Utf8, true),
2727 ])),
2728 true,
2729 ))),
2730 true,
2731 )]);
2732
2733 let projected = batch.project_by_schema(&target_schema).unwrap();
2734
2735 assert_eq!(projected.schema().as_ref(), &target_schema);
2737
2738 let outer_list = projected.column(0).as_list::<i32>();
2740 let middle_struct = outer_list.values().as_struct();
2741
2742 assert_eq!(
2744 middle_struct.column(0).as_ref(),
2745 &StringArray::from(vec!["a1"]) as &dyn Array
2746 );
2747 assert_eq!(
2748 middle_struct.column(2).as_ref(),
2749 &StringArray::from(vec!["b1"]) as &dyn Array
2750 );
2751
2752 let inner_list = middle_struct.column(1).as_list::<i32>();
2754 let innermost_struct = inner_list.values().as_struct();
2755 assert_eq!(
2756 innermost_struct.column(0).as_ref(),
2757 &Int32Array::from(vec![3, 4]) as &dyn Array
2758 );
2759 assert_eq!(
2760 innermost_struct.column(1).as_ref(),
2761 &Int32Array::from(vec![1, 2]) as &dyn Array
2762 );
2763 }
2764}