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