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