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