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