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::{BooleanBufferBuilder, 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>;
449}
450
451impl FixedSizeBinaryArrayExt for FixedSizeBinaryArray {
452 fn try_new_from_values(values: &UInt8Array, stride: i32) -> Result<Self> {
453 let data_type = DataType::FixedSizeBinary(stride);
454 let data = ArrayDataBuilder::new(data_type)
455 .len(values.len() / stride as usize)
456 .add_buffer(values.into_data().buffers()[0].clone())
457 .build()?;
458 Ok(Self::from(data))
459 }
460}
461
462pub fn as_fixed_size_binary_array(arr: &dyn Array) -> &FixedSizeBinaryArray {
463 arr.as_any().downcast_ref::<FixedSizeBinaryArray>().unwrap()
464}
465
466pub fn iter_str_array(arr: &dyn Array) -> Box<dyn Iterator<Item = Option<&str>> + Send + '_> {
467 match arr.data_type() {
468 DataType::Utf8 => Box::new(arr.as_string::<i32>().iter()),
469 DataType::LargeUtf8 => Box::new(arr.as_string::<i64>().iter()),
470 DataType::Utf8View => Box::new(arr.as_string_view().iter()),
471 _ => panic!(
472 "Expecting Utf8, LargeUtf8 or Utf8View, found {:?}",
473 arr.data_type()
474 ),
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 holds_struct(data_type: &DataType) -> bool {
916 match data_type {
917 DataType::Struct(_) => true,
918 DataType::List(item) | DataType::LargeList(item) => holds_struct(item.data_type()),
919 _ => false,
920 }
921}
922
923fn merge_list_values(left: &ArrayRef, right: &ArrayRef) -> ArrayRef {
926 match (left.data_type(), right.data_type()) {
927 (DataType::Struct(_), DataType::Struct(_)) => {
928 Arc::new(merge(left.as_struct(), right.as_struct()))
929 }
930 (DataType::List(_), DataType::List(_))
931 | (DataType::LargeList(_), DataType::LargeList(_))
932 if left.data_type() != right.data_type()
933 && holds_struct(left.data_type())
934 && holds_struct(right.data_type()) =>
935 {
936 merge_list_struct(left, right)
937 }
938 _ => left.clone(),
939 }
940}
941
942fn merge_list_struct_helper<T: OffsetSizeTrait>(
943 left: &dyn Array,
944 right: &dyn Array,
945) -> Arc<dyn Array> {
946 let left_list: &GenericListArray<T> = left.as_list();
947 let right_list: &GenericListArray<T> = right.as_list();
948 let item_name = match left.data_type() {
949 DataType::List(item) | DataType::LargeList(item) => item.name().clone(),
950 _ => unreachable!(),
951 };
952 let null_left = left_list.null_count() == left_list.len();
955 let null_right = !null_left && right_list.null_count() == right_list.len();
956 let (merged_items, list, items_nullable) = if null_left || null_right {
957 let not_null = if null_left { right_list } else { left_list };
958 let values_len = not_null.values().len();
959 let values = |list: &GenericListArray<T>, is_null: bool| {
960 if is_null {
961 new_null_array(list.values().data_type(), values_len)
962 } else {
963 list.values().clone()
964 }
965 };
966 let merged = merge_list_values(
967 &values(left_list, null_left),
968 &values(right_list, null_right),
969 );
970 (merged, not_null, true)
971 } else {
972 if left_list.offsets().inner() != right_list.offsets().inner() {
973 panic!("Attempt to merge list struct arrays which do not have same offsets");
974 }
975 let merged = merge_list_values(left_list.values(), right_list.values());
976 let nullable = match left.data_type() {
977 DataType::List(item) | DataType::LargeList(item) => item.is_nullable(),
978 _ => unreachable!(),
979 };
980 (merged, left_list, nullable)
981 };
982 let items_field = Arc::new(Field::new(
983 item_name,
984 merged_items.data_type().clone(),
985 items_nullable,
986 ));
987 Arc::new(GenericListArray::<T>::new(
988 items_field,
989 list.offsets().clone(),
990 merged_items,
991 list.nulls().cloned(),
992 ))
993}
994
995fn merge_list_struct(left: &dyn Array, right: &dyn Array) -> Arc<dyn Array> {
998 match (left.data_type(), right.data_type()) {
999 (DataType::List(_), DataType::List(_)) => merge_list_struct_helper::<i32>(left, right),
1000 (DataType::LargeList(_), DataType::LargeList(_)) => {
1001 merge_list_struct_helper::<i64>(left, right)
1002 }
1003 _ => unreachable!(),
1004 }
1005}
1006
1007fn merge_struct_validity(
1012 left_validity: Option<&arrow_buffer::NullBuffer>,
1013 right_validity: Option<&arrow_buffer::NullBuffer>,
1014) -> Option<arrow_buffer::NullBuffer> {
1015 match (left_validity, right_validity) {
1016 (None, _) | (_, None) => None,
1018 (Some(left), Some(right)) => {
1019 if left.null_count() == 0 || right.null_count() == 0 {
1020 return None;
1021 }
1022 if left.null_count() == left.len() {
1023 return Some(right.clone());
1024 }
1025 if right.null_count() == right.len() {
1026 return Some(left.clone());
1027 }
1028
1029 let left_buffer = left.inner();
1030 let right_buffer = right.inner();
1031
1032 let merged_buffer = left_buffer | right_buffer;
1035
1036 Some(arrow_buffer::NullBuffer::from(merged_buffer))
1037 }
1038 }
1039}
1040
1041fn merge_list_child_values(
1042 child_field: &Field,
1043 left_values: ArrayRef,
1044 right_values: ArrayRef,
1045) -> ArrayRef {
1046 match child_field.data_type() {
1047 DataType::Struct(child_fields) => Arc::new(merge_with_schema(
1048 left_values.as_struct(),
1049 right_values.as_struct(),
1050 child_fields,
1051 )) as ArrayRef,
1052 DataType::List(grandchild) => {
1053 let left_list = left_values
1054 .as_any()
1055 .downcast_ref::<ListArray>()
1056 .expect("left list values should be ListArray");
1057 let right_list = right_values
1058 .as_any()
1059 .downcast_ref::<ListArray>()
1060 .expect("right list values should be ListArray");
1061 let merged_values = merge_list_child_values(
1062 grandchild.as_ref(),
1063 left_list.values().clone(),
1064 right_list.values().clone(),
1065 );
1066 let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1067 Arc::new(ListArray::new(
1068 grandchild.clone(),
1069 left_list.offsets().clone(),
1070 merged_values,
1071 merged_validity,
1072 )) as ArrayRef
1073 }
1074 DataType::LargeList(grandchild) => {
1075 let left_list = left_values
1076 .as_any()
1077 .downcast_ref::<LargeListArray>()
1078 .expect("left list values should be LargeListArray");
1079 let right_list = right_values
1080 .as_any()
1081 .downcast_ref::<LargeListArray>()
1082 .expect("right list values should be LargeListArray");
1083 let merged_values = merge_list_child_values(
1084 grandchild.as_ref(),
1085 left_list.values().clone(),
1086 right_list.values().clone(),
1087 );
1088 let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1089 Arc::new(LargeListArray::new(
1090 grandchild.clone(),
1091 left_list.offsets().clone(),
1092 merged_values,
1093 merged_validity,
1094 )) as ArrayRef
1095 }
1096 DataType::FixedSizeList(grandchild, list_size) => {
1097 let left_list = left_values
1098 .as_any()
1099 .downcast_ref::<FixedSizeListArray>()
1100 .expect("left list values should be FixedSizeListArray");
1101 let right_list = right_values
1102 .as_any()
1103 .downcast_ref::<FixedSizeListArray>()
1104 .expect("right list values should be FixedSizeListArray");
1105 let merged_values = merge_list_child_values(
1106 grandchild.as_ref(),
1107 left_list.values().clone(),
1108 right_list.values().clone(),
1109 );
1110 let merged_validity = merge_struct_validity(left_list.nulls(), right_list.nulls());
1111 Arc::new(FixedSizeListArray::new(
1112 grandchild.clone(),
1113 *list_size,
1114 merged_values,
1115 merged_validity,
1116 )) as ArrayRef
1117 }
1118 _ => left_values.clone(),
1119 }
1120}
1121
1122fn adjust_child_validity(
1126 child: &ArrayRef,
1127 parent_validity: Option<&arrow_buffer::NullBuffer>,
1128) -> ArrayRef {
1129 let parent_validity = match parent_validity {
1131 None => return child.clone(),
1132 Some(p) if p.null_count() == 0 => return child.clone(), Some(p) => p,
1134 };
1135
1136 if child.data_type() == &DataType::Null {
1139 return child.clone();
1140 }
1141
1142 let child_validity = child.nulls();
1143
1144 let new_validity = match child_validity {
1146 None => {
1147 parent_validity.clone()
1149 }
1150 Some(child_nulls) => {
1151 let child_buffer = child_nulls.inner();
1152 let parent_buffer = parent_validity.inner();
1153
1154 let merged_buffer = child_buffer & parent_buffer;
1157
1158 arrow_buffer::NullBuffer::from(merged_buffer)
1159 }
1160 };
1161
1162 let null_bit_buffer = if new_validity.inner().offset() == child.offset() {
1165 new_validity.into_inner().into_inner()
1166 } else {
1167 let mut builder = BooleanBufferBuilder::new(child.offset() + child.len());
1168 builder.append_n(child.offset(), true);
1169 builder.append_buffer(new_validity.inner());
1170 builder.build().into_inner()
1171 };
1172
1173 arrow_array::make_array(
1175 arrow_data::ArrayData::try_new(
1176 child.data_type().clone(),
1177 child.len(),
1178 Some(null_bit_buffer),
1179 child.offset(),
1180 child.to_data().buffers().to_vec(),
1181 child.to_data().child_data().to_vec(),
1182 )
1183 .unwrap(),
1184 )
1185}
1186
1187fn merge(left_struct_array: &StructArray, right_struct_array: &StructArray) -> StructArray {
1188 let mut fields: Vec<Field> = vec![];
1189 let mut columns: Vec<ArrayRef> = vec![];
1190 let right_fields = right_struct_array.fields();
1191 let right_columns = right_struct_array.columns();
1192
1193 let left_validity = left_struct_array.nulls();
1195 let right_validity = right_struct_array.nulls();
1196
1197 let merged_validity = merge_struct_validity(left_validity, right_validity);
1199
1200 for (left_field, left_column) in left_struct_array
1202 .fields()
1203 .iter()
1204 .zip(left_struct_array.columns().iter())
1205 {
1206 match right_fields
1207 .iter()
1208 .position(|f| f.name() == left_field.name())
1209 {
1210 Some(right_index) => {
1212 let right_field = right_fields.get(right_index).unwrap();
1213 let right_column = right_columns.get(right_index).unwrap();
1214 match (left_field.data_type(), right_field.data_type()) {
1216 (DataType::Struct(_), DataType::Struct(_)) => {
1217 let left_sub_array = left_column.as_struct();
1218 let right_sub_array = right_column.as_struct();
1219 let merged_sub_array = merge(left_sub_array, right_sub_array);
1220 fields.push(Field::new(
1221 left_field.name(),
1222 merged_sub_array.data_type().clone(),
1223 left_field.is_nullable(),
1224 ));
1225 columns.push(Arc::new(merged_sub_array) as ArrayRef);
1226 }
1227 (DataType::List(left_list), DataType::List(right_list))
1228 | (DataType::LargeList(left_list), DataType::LargeList(right_list))
1229 if holds_struct(left_list.data_type())
1230 && holds_struct(right_list.data_type()) =>
1231 {
1232 if left_list.data_type() == right_list.data_type() {
1234 fields.push(left_field.as_ref().clone());
1235 columns.push(left_column.clone());
1236 } else {
1237 let merged_sub_array = merge_list_struct(&left_column, &right_column);
1241
1242 fields.push(Field::new(
1243 left_field.name(),
1244 merged_sub_array.data_type().clone(),
1245 left_field.is_nullable(),
1246 ));
1247 columns.push(merged_sub_array);
1248 }
1249 }
1250 _ => {
1252 fields.push(left_field.as_ref().clone());
1254 let adjusted_column = adjust_child_validity(left_column, left_validity);
1256 columns.push(adjusted_column);
1257 }
1258 }
1259 }
1260 None => {
1261 fields.push(left_field.as_ref().clone());
1262 let adjusted_column = adjust_child_validity(left_column, left_validity);
1264 columns.push(adjusted_column);
1265 }
1266 }
1267 }
1268
1269 right_fields
1271 .iter()
1272 .zip(right_columns.iter())
1273 .for_each(|(field, column)| {
1274 if !left_struct_array
1276 .fields()
1277 .iter()
1278 .any(|f| f.name() == field.name())
1279 {
1280 fields.push(field.as_ref().clone());
1281 let adjusted_column = adjust_child_validity(column, right_validity);
1284 columns.push(adjusted_column);
1285 }
1286 });
1287
1288 StructArray::try_new(Fields::from(fields), columns, merged_validity).unwrap()
1289}
1290
1291fn merge_with_schema(
1292 left_struct_array: &StructArray,
1293 right_struct_array: &StructArray,
1294 fields: &Fields,
1295) -> StructArray {
1296 fn same_type_kind(left: &DataType, right: &DataType) -> bool {
1298 match (left, right) {
1299 (DataType::Struct(_), DataType::Struct(_)) => true,
1300 (DataType::Struct(_), _) => false,
1301 (_, DataType::Struct(_)) => false,
1302 _ => true,
1303 }
1304 }
1305
1306 let mut output_fields: Vec<Field> = Vec::with_capacity(fields.len());
1307 let mut columns: Vec<ArrayRef> = Vec::with_capacity(fields.len());
1308
1309 let left_fields = left_struct_array.fields();
1310 let left_columns = left_struct_array.columns();
1311 let right_fields = right_struct_array.fields();
1312 let right_columns = right_struct_array.columns();
1313
1314 let left_validity = left_struct_array.nulls();
1316 let right_validity = right_struct_array.nulls();
1317
1318 let merged_validity = merge_struct_validity(left_validity, right_validity);
1320
1321 for field in fields {
1322 let left_match_idx = left_fields.iter().position(|f| {
1323 f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1324 });
1325 let right_match_idx = right_fields.iter().position(|f| {
1326 f.name() == field.name() && same_type_kind(f.data_type(), field.data_type())
1327 });
1328
1329 match (left_match_idx, right_match_idx) {
1330 (None, Some(right_idx)) => {
1331 output_fields.push(right_fields[right_idx].as_ref().clone());
1332 let adjusted_column =
1334 adjust_child_validity(&right_columns[right_idx], right_validity);
1335 columns.push(adjusted_column);
1336 }
1337 (Some(left_idx), None) => {
1338 output_fields.push(left_fields[left_idx].as_ref().clone());
1339 let adjusted_column = adjust_child_validity(&left_columns[left_idx], left_validity);
1341 columns.push(adjusted_column);
1342 }
1343 (Some(left_idx), Some(right_idx)) => {
1344 match field.data_type() {
1345 DataType::Struct(child_fields) => {
1346 let left_sub_array = left_columns[left_idx].as_struct();
1347 let right_sub_array = right_columns[right_idx].as_struct();
1348 let merged_sub_array =
1349 merge_with_schema(left_sub_array, right_sub_array, child_fields);
1350 output_fields.push(Field::new(
1351 field.name(),
1352 merged_sub_array.data_type().clone(),
1353 field.is_nullable(),
1354 ));
1355 columns.push(Arc::new(merged_sub_array) as ArrayRef);
1356 }
1357 DataType::List(child_field) => {
1358 let left_list = left_columns[left_idx]
1359 .as_any()
1360 .downcast_ref::<ListArray>()
1361 .unwrap();
1362 let right_list = right_columns[right_idx]
1363 .as_any()
1364 .downcast_ref::<ListArray>()
1365 .unwrap();
1366 let merged_values = merge_list_child_values(
1367 child_field.as_ref(),
1368 left_list.trimmed_values(),
1369 right_list.trimmed_values(),
1370 );
1371 let merged_validity =
1372 merge_struct_validity(left_list.nulls(), right_list.nulls());
1373 let merged_list = ListArray::new(
1377 child_field.clone(),
1378 left_list.trimmed_offsets(),
1379 merged_values,
1380 merged_validity,
1381 );
1382 output_fields.push(field.as_ref().clone());
1383 columns.push(Arc::new(merged_list) as ArrayRef);
1384 }
1385 DataType::LargeList(child_field) => {
1386 let left_list = left_columns[left_idx]
1387 .as_any()
1388 .downcast_ref::<LargeListArray>()
1389 .unwrap();
1390 let right_list = right_columns[right_idx]
1391 .as_any()
1392 .downcast_ref::<LargeListArray>()
1393 .unwrap();
1394 let merged_values = merge_list_child_values(
1395 child_field.as_ref(),
1396 left_list.trimmed_values(),
1397 right_list.trimmed_values(),
1398 );
1399 let merged_validity =
1400 merge_struct_validity(left_list.nulls(), right_list.nulls());
1401 let merged_list = LargeListArray::new(
1402 child_field.clone(),
1403 left_list.trimmed_offsets(),
1404 merged_values,
1405 merged_validity,
1406 );
1407 output_fields.push(field.as_ref().clone());
1408 columns.push(Arc::new(merged_list) as ArrayRef);
1409 }
1410 DataType::FixedSizeList(child_field, list_size) => {
1411 let left_list = left_columns[left_idx]
1412 .as_any()
1413 .downcast_ref::<FixedSizeListArray>()
1414 .unwrap();
1415 let right_list = right_columns[right_idx]
1416 .as_any()
1417 .downcast_ref::<FixedSizeListArray>()
1418 .unwrap();
1419 let merged_values = merge_list_child_values(
1420 child_field.as_ref(),
1421 left_list.values().clone(),
1422 right_list.values().clone(),
1423 );
1424 let merged_validity =
1425 merge_struct_validity(left_list.nulls(), right_list.nulls());
1426 let merged_list = FixedSizeListArray::new(
1427 child_field.clone(),
1428 *list_size,
1429 merged_values,
1430 merged_validity,
1431 );
1432 output_fields.push(field.as_ref().clone());
1433 columns.push(Arc::new(merged_list) as ArrayRef);
1434 }
1435 _ => {
1436 output_fields.push(left_fields[left_idx].as_ref().clone());
1437 let adjusted_column =
1439 adjust_child_validity(&left_columns[left_idx], left_validity);
1440 columns.push(adjusted_column);
1441 }
1442 }
1443 }
1444 (None, None) => {
1445 }
1447 }
1448 }
1449
1450 StructArray::try_new(Fields::from(output_fields), columns, merged_validity).unwrap()
1451}
1452
1453fn get_sub_array<'a>(array: &'a ArrayRef, components: &[&str]) -> Option<&'a ArrayRef> {
1454 if components.is_empty() {
1455 return Some(array);
1456 }
1457 if !matches!(array.data_type(), DataType::Struct(_)) {
1458 return None;
1459 }
1460 let struct_arr = array.as_struct();
1461 struct_arr
1462 .column_by_name(components[0])
1463 .and_then(|arr| get_sub_array(arr, &components[1..]))
1464}
1465
1466pub fn interleave_batches(
1470 batches: &[RecordBatch],
1471 indices: &[(usize, usize)],
1472) -> Result<RecordBatch> {
1473 let first_batch = batches.first().ok_or_else(|| {
1474 ArrowError::InvalidArgumentError("Cannot interleave zero RecordBatches".to_string())
1475 })?;
1476 let schema = first_batch.schema();
1477 let num_columns = first_batch.num_columns();
1478 let mut columns = Vec::with_capacity(num_columns);
1479 let mut chunks = Vec::with_capacity(batches.len());
1480
1481 for i in 0..num_columns {
1482 for batch in batches {
1483 chunks.push(batch.column(i).as_ref());
1484 }
1485 let new_column = interleave(&chunks, indices)?;
1486 columns.push(new_column);
1487 chunks.clear();
1488 }
1489
1490 RecordBatch::try_new(schema, columns)
1491}
1492
1493pub trait BufferExt {
1494 fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self;
1509
1510 fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self;
1519}
1520
1521fn is_pwr_two(n: u64) -> bool {
1522 n & (n - 1) == 0
1523}
1524
1525impl BufferExt for arrow_buffer::Buffer {
1526 fn from_bytes_bytes(bytes: bytes::Bytes, bytes_per_value: u64) -> Self {
1527 if is_pwr_two(bytes_per_value) && bytes.as_ptr().align_offset(bytes_per_value as usize) != 0
1528 {
1529 let size_bytes = bytes.len();
1531 Self::copy_bytes_bytes(bytes, size_bytes)
1532 } else {
1533 unsafe {
1536 Self::from_custom_allocation(
1537 NonNull::new(bytes.as_ptr() as _).expect("should be a valid pointer"),
1538 bytes.len(),
1539 Arc::new(bytes),
1540 )
1541 }
1542 }
1543 }
1544
1545 fn copy_bytes_bytes(bytes: bytes::Bytes, size_bytes: usize) -> Self {
1546 assert!(size_bytes >= bytes.len());
1547 let mut buf = MutableBuffer::with_capacity(size_bytes);
1548 let to_fill = size_bytes - bytes.len();
1549 buf.extend(bytes);
1550 buf.extend(std::iter::repeat_n(0_u8, to_fill));
1551
1552 buf.shrink_to_fit();
1555
1556 Self::from(buf)
1557 }
1558}
1559
1560#[cfg(test)]
1561mod tests {
1562 use super::*;
1563 use arrow_array::{
1564 BinaryArray, BinaryViewArray, FixedSizeBinaryArray, Float32Array, Int32Array,
1565 LargeBinaryArray, ListArray, NullArray, StringArray, StructArray, new_empty_array,
1566 new_null_array,
1567 };
1568 use arrow_buffer::OffsetBuffer;
1569
1570 #[rstest::rstest]
1571 fn test_adjust_child_validity_sliced_boolean(
1572 #[values(0, 1, 7, 8, 9, 1025)] child_offset: usize,
1573 #[values(0, 1, 7, 8, 9, 1025)] parent_offset: usize,
1574 #[values(false, true)] child_has_nulls: bool,
1575 ) {
1576 let mut values = vec![Some(false); child_offset];
1577 values.extend([
1578 Some(true),
1579 Some(false),
1580 if child_has_nulls { None } else { Some(true) },
1581 Some(false),
1582 ]);
1583 let child: ArrayRef =
1584 Arc::new(arrow_array::BooleanArray::from(values).slice(child_offset, 4));
1585 let mut validity = vec![false; parent_offset];
1586 validity.extend([true, false, true, true]);
1587 let parent = arrow_buffer::NullBuffer::from(validity).slice(parent_offset, 4);
1588
1589 let adjusted = adjust_child_validity(&child, Some(&parent));
1590 let adjusted = adjusted.as_boolean();
1591 assert_eq!(
1592 adjusted.iter().collect::<Vec<_>>(),
1593 vec![
1594 Some(true),
1595 None,
1596 if child_has_nulls { None } else { Some(true) },
1597 Some(false)
1598 ]
1599 );
1600 }
1601
1602 #[rstest::rstest]
1603 fn test_adjust_child_validity_sliced_parent(
1604 #[values(0, 1, 7, 8, 9, 1025)] parent_offset: usize,
1605 #[values(false, true)] child_has_nulls: bool,
1606 ) {
1607 let child: ArrayRef = Arc::new(Int32Array::from(vec![
1608 Some(1),
1609 Some(2),
1610 if child_has_nulls { None } else { Some(3) },
1611 Some(4),
1612 ]));
1613 let mut validity = vec![false; parent_offset];
1614 validity.extend([true, false, true, true]);
1615 let parent = arrow_buffer::NullBuffer::from(validity).slice(parent_offset, 4);
1616
1617 let adjusted = adjust_child_validity(&child, Some(&parent));
1618 let adjusted = adjusted.as_primitive::<arrow_array::types::Int32Type>();
1619 assert_eq!(
1620 adjusted.iter().collect::<Vec<_>>(),
1621 vec![
1622 Some(1),
1623 None,
1624 if child_has_nulls { None } else { Some(3) },
1625 Some(4)
1626 ]
1627 );
1628 }
1629
1630 #[test]
1631 fn test_convert_to_floating_point_preserves_inner_nulls() {
1632 let values = Int8Array::from(vec![Some(1), None, Some(3), Some(4)]);
1637 let fsl = FixedSizeListArray::new(
1638 Arc::new(Field::new("item", DataType::Int8, true)),
1639 2,
1640 Arc::new(values),
1641 None,
1642 );
1643
1644 let converted = fsl.convert_to_floating_point().unwrap();
1645
1646 assert_eq!(converted.len(), 2);
1647 let conv_values = converted
1648 .values()
1649 .as_any()
1650 .downcast_ref::<Float32Array>()
1651 .unwrap();
1652 assert_eq!(conv_values.len(), 4);
1653 assert_eq!(conv_values.value(0), 1.0);
1654 assert!(conv_values.is_null(1));
1655 assert_eq!(conv_values.value(2), 3.0);
1656 assert_eq!(conv_values.value(3), 4.0);
1657 }
1658
1659 #[test]
1660 fn test_convert_to_floating_point_preserves_inner_nulls_f64_arm() {
1661 let values = UInt8Array::from(vec![Some(10u8), None, Some(30), Some(40)]);
1665 let fsl = FixedSizeListArray::new(
1666 Arc::new(Field::new("item", DataType::UInt8, true)),
1667 2,
1668 Arc::new(values),
1669 None,
1670 );
1671
1672 let converted = fsl.convert_to_floating_point().unwrap();
1673
1674 assert_eq!(converted.len(), 2);
1675 let conv_values = converted
1676 .values()
1677 .as_any()
1678 .downcast_ref::<Float64Array>()
1679 .unwrap();
1680 assert_eq!(conv_values.len(), 4);
1681 assert_eq!(conv_values.value(0), 10.0);
1682 assert!(conv_values.is_null(1));
1683 assert_eq!(conv_values.value(2), 30.0);
1684 assert_eq!(conv_values.value(3), 40.0);
1685 }
1686
1687 #[test]
1688 fn test_merge_recursive() {
1689 let a_array = Int32Array::from(vec![Some(1), Some(2), Some(3)]);
1690 let e_array = Int32Array::from(vec![Some(4), Some(5), Some(6)]);
1691 let c_array = Int32Array::from(vec![Some(7), Some(8), Some(9)]);
1692 let d_array = StringArray::from(vec![Some("a"), Some("b"), Some("c")]);
1693
1694 let left_schema = Schema::new(vec![
1695 Field::new("a", DataType::Int32, true),
1696 Field::new(
1697 "b",
1698 DataType::Struct(vec![Field::new("c", DataType::Int32, true)].into()),
1699 true,
1700 ),
1701 ]);
1702 let left_batch = RecordBatch::try_new(
1703 Arc::new(left_schema),
1704 vec![
1705 Arc::new(a_array.clone()),
1706 Arc::new(StructArray::from(vec![(
1707 Arc::new(Field::new("c", DataType::Int32, true)),
1708 Arc::new(c_array.clone()) as ArrayRef,
1709 )])),
1710 ],
1711 )
1712 .unwrap();
1713
1714 let right_schema = Schema::new(vec![
1715 Field::new("e", DataType::Int32, true),
1716 Field::new(
1717 "b",
1718 DataType::Struct(vec![Field::new("d", DataType::Utf8, true)].into()),
1719 true,
1720 ),
1721 ]);
1722 let right_batch = RecordBatch::try_new(
1723 Arc::new(right_schema),
1724 vec![
1725 Arc::new(e_array.clone()),
1726 Arc::new(StructArray::from(vec![(
1727 Arc::new(Field::new("d", DataType::Utf8, true)),
1728 Arc::new(d_array.clone()) as ArrayRef,
1729 )])) as ArrayRef,
1730 ],
1731 )
1732 .unwrap();
1733
1734 let merged_schema = Schema::new(vec![
1735 Field::new("a", DataType::Int32, true),
1736 Field::new(
1737 "b",
1738 DataType::Struct(
1739 vec![
1740 Field::new("c", DataType::Int32, true),
1741 Field::new("d", DataType::Utf8, true),
1742 ]
1743 .into(),
1744 ),
1745 true,
1746 ),
1747 Field::new("e", DataType::Int32, true),
1748 ]);
1749 let merged_batch = RecordBatch::try_new(
1750 Arc::new(merged_schema),
1751 vec![
1752 Arc::new(a_array) as ArrayRef,
1753 Arc::new(StructArray::from(vec![
1754 (
1755 Arc::new(Field::new("c", DataType::Int32, true)),
1756 Arc::new(c_array) as ArrayRef,
1757 ),
1758 (
1759 Arc::new(Field::new("d", DataType::Utf8, true)),
1760 Arc::new(d_array) as ArrayRef,
1761 ),
1762 ])) as ArrayRef,
1763 Arc::new(e_array) as ArrayRef,
1764 ],
1765 )
1766 .unwrap();
1767
1768 let result = left_batch.merge(&right_batch).unwrap();
1769 assert_eq!(result, merged_batch);
1770 }
1771
1772 #[test]
1773 fn test_merge_with_schema() {
1774 fn test_batch(names: &[&str], types: &[DataType]) -> (Schema, RecordBatch) {
1775 let fields: Fields = names
1776 .iter()
1777 .zip(types)
1778 .map(|(name, ty)| Field::new(name.to_string(), ty.clone(), false))
1779 .collect();
1780 let schema = Schema::new(vec![Field::new(
1781 "struct",
1782 DataType::Struct(fields.clone()),
1783 false,
1784 )]);
1785 let children = types.iter().map(new_empty_array).collect::<Vec<_>>();
1786 let batch = RecordBatch::try_new(
1787 Arc::new(schema.clone()),
1788 vec![Arc::new(StructArray::new(fields, children, None)) as ArrayRef],
1789 );
1790 (schema, batch.unwrap())
1791 }
1792
1793 let (_, left_batch) = test_batch(&["a", "b"], &[DataType::Int32, DataType::Int64]);
1794 let (_, right_batch) = test_batch(&["c", "b"], &[DataType::Int32, DataType::Int64]);
1795 let (output_schema, _) = test_batch(
1796 &["b", "a", "c"],
1797 &[DataType::Int64, DataType::Int32, DataType::Int32],
1798 );
1799
1800 let merged = left_batch
1802 .merge_with_schema(&right_batch, &output_schema)
1803 .unwrap();
1804 assert_eq!(merged.schema().as_ref(), &output_schema);
1805
1806 let (naive_schema, _) = test_batch(
1808 &["a", "b", "c"],
1809 &[DataType::Int32, DataType::Int64, DataType::Int32],
1810 );
1811 let merged = left_batch.merge(&right_batch).unwrap();
1812 assert_eq!(merged.schema().as_ref(), &naive_schema);
1813 }
1814
1815 #[test]
1816 fn test_merge_list_struct() {
1817 let x_field = Arc::new(Field::new("x", DataType::Int32, true));
1818 let y_field = Arc::new(Field::new("y", DataType::Int32, true));
1819 let x_struct_field = Arc::new(Field::new(
1820 "item",
1821 DataType::Struct(Fields::from(vec![x_field.clone()])),
1822 true,
1823 ));
1824 let y_struct_field = Arc::new(Field::new(
1825 "item",
1826 DataType::Struct(Fields::from(vec![y_field.clone()])),
1827 true,
1828 ));
1829 let both_struct_field = Arc::new(Field::new(
1830 "item",
1831 DataType::Struct(Fields::from(vec![x_field.clone(), y_field.clone()])),
1832 true,
1833 ));
1834 let left_schema = Schema::new(vec![Field::new(
1835 "list_struct",
1836 DataType::List(x_struct_field.clone()),
1837 true,
1838 )]);
1839 let right_schema = Schema::new(vec![Field::new(
1840 "list_struct",
1841 DataType::List(y_struct_field.clone()),
1842 true,
1843 )]);
1844 let both_schema = Schema::new(vec![Field::new(
1845 "list_struct",
1846 DataType::List(both_struct_field.clone()),
1847 true,
1848 )]);
1849
1850 let x = Arc::new(Int32Array::from(vec![1]));
1851 let y = Arc::new(Int32Array::from(vec![2]));
1852 let x_struct = Arc::new(StructArray::new(
1853 Fields::from(vec![x_field.clone()]),
1854 vec![x.clone()],
1855 None,
1856 ));
1857 let y_struct = Arc::new(StructArray::new(
1858 Fields::from(vec![y_field.clone()]),
1859 vec![y.clone()],
1860 None,
1861 ));
1862 let both_struct = Arc::new(StructArray::new(
1863 Fields::from(vec![x_field.clone(), y_field.clone()]),
1864 vec![x.clone(), y],
1865 None,
1866 ));
1867 let both_null_struct = Arc::new(StructArray::new(
1868 Fields::from(vec![x_field, y_field]),
1869 vec![x, Arc::new(new_null_array(&DataType::Int32, 1))],
1870 None,
1871 ));
1872 let offsets = OffsetBuffer::from_lengths([1]);
1873 let x_s_list = ListArray::new(x_struct_field, offsets.clone(), x_struct, None);
1874 let y_s_list = ListArray::new(y_struct_field, offsets.clone(), y_struct, None);
1875 let both_list = ListArray::new(
1876 both_struct_field.clone(),
1877 offsets.clone(),
1878 both_struct,
1879 None,
1880 );
1881 let both_null_list = ListArray::new(both_struct_field, offsets, both_null_struct, None);
1882 let x_batch =
1883 RecordBatch::try_new(Arc::new(left_schema), vec![Arc::new(x_s_list)]).unwrap();
1884 let y_batch = RecordBatch::try_new(
1885 Arc::new(right_schema.clone()),
1886 vec![Arc::new(y_s_list.clone())],
1887 )
1888 .unwrap();
1889 let merged = x_batch.merge(&y_batch).unwrap();
1890 let expected =
1891 RecordBatch::try_new(Arc::new(both_schema.clone()), vec![Arc::new(both_list)]).unwrap();
1892 assert_eq!(merged, expected);
1893
1894 let y_null_list = new_null_array(y_s_list.data_type(), 1);
1895 let y_null_batch =
1896 RecordBatch::try_new(Arc::new(right_schema), vec![Arc::new(y_null_list.clone())])
1897 .unwrap();
1898 let expected =
1899 RecordBatch::try_new(Arc::new(both_schema), vec![Arc::new(both_null_list)]).unwrap();
1900 let merged = x_batch.merge(&y_null_batch).unwrap();
1901 assert_eq!(merged, expected);
1902 }
1903
1904 #[test]
1905 fn test_merge_all_null_list_adds_nested_child() {
1906 let strukt = |cols: Vec<(&str, ArrayRef)>| -> ArrayRef {
1909 let fields = cols
1910 .iter()
1911 .map(|(name, col)| Field::new(*name, col.data_type().clone(), true))
1912 .collect::<Vec<_>>();
1913 Arc::new(StructArray::new(
1914 fields.into(),
1915 cols.into_iter().map(|(_, col)| col).collect(),
1916 None,
1917 ))
1918 };
1919 let list = |values: ArrayRef, large: bool| -> ArrayRef {
1921 let field = Arc::new(Field::new("item", values.data_type().clone(), true));
1922 let lengths = [values.len(), 0];
1923 let nulls = Some(arrow_buffer::NullBuffer::from(vec![true, false]));
1924 if large {
1925 Arc::new(LargeListArray::new(
1926 field,
1927 OffsetBuffer::from_lengths(lengths),
1928 values,
1929 nulls,
1930 ))
1931 } else {
1932 Arc::new(ListArray::new(
1933 field,
1934 OffsetBuffer::from_lengths(lengths),
1935 values,
1936 nulls,
1937 ))
1938 }
1939 };
1940 let a = || Arc::new(Int32Array::from(vec![7])) as ArrayRef;
1941 let b = || new_null_array(&DataType::Utf8, 1);
1942 for large in [false, true] {
1943 let under_list = |inner: ArrayRef, with_t: bool| {
1945 let mut cols = vec![("in", inner)];
1946 if with_t {
1947 cols.insert(0, ("t", Arc::new(Int32Array::from(vec![1])) as ArrayRef));
1948 }
1949 list(strukt(cols), large)
1950 };
1951 let under_two_lists = |inner: ArrayRef, _: bool| {
1953 let inner_list: ArrayRef = if large {
1954 let field = Arc::new(Field::new("item", inner.data_type().clone(), true));
1955 Arc::new(LargeListArray::new(
1956 field,
1957 OffsetBuffer::from_lengths([1]),
1958 inner,
1959 None,
1960 ))
1961 } else {
1962 let field = Arc::new(Field::new("item", inner.data_type().clone(), true));
1963 Arc::new(ListArray::new(
1964 field,
1965 OffsetBuffer::from_lengths([1]),
1966 inner,
1967 None,
1968 ))
1969 };
1970 list(inner_list, large)
1971 };
1972 let shapes: [&dyn Fn(ArrayRef, bool) -> ArrayRef; 2] = [&under_list, &under_two_lists];
1973 for shape in shapes {
1974 let data = shape(strukt(vec![("a", a())]), true);
1975 let expected = shape(strukt(vec![("a", a()), ("b", b())]), true);
1976 let placeholder =
1977 new_null_array(shape(strukt(vec![("b", b())]), false).data_type(), 2);
1978 let batch = |col: ArrayRef| {
1979 RecordBatch::try_new(
1980 Arc::new(Schema::new(vec![Field::new(
1981 "c",
1982 col.data_type().clone(),
1983 true,
1984 )])),
1985 vec![col],
1986 )
1987 .unwrap()
1988 };
1989 let expected = batch(expected);
1990 for merged in [
1992 batch(data.clone())
1993 .merge(&batch(placeholder.clone()))
1994 .unwrap(),
1995 batch(placeholder.clone())
1996 .merge(&batch(data.clone()))
1997 .unwrap(),
1998 ] {
1999 let merged = merged.project_by_schema(expected.schema_ref()).unwrap();
2000 assert_eq!(merged, expected, "large={large}");
2001 }
2002 }
2003 }
2004 }
2005
2006 #[test]
2007 fn test_merge_list_struct_identical_schema() {
2008 let x_field = Arc::new(Field::new("x", DataType::Int32, true));
2012 let item_field = Arc::new(Field::new(
2013 "item",
2014 DataType::Struct(Fields::from(vec![x_field.clone()])),
2015 true,
2016 ));
2017 let schema = Arc::new(Schema::new(vec![Field::new(
2018 "list_struct",
2019 DataType::List(item_field.clone()),
2020 true,
2021 )]));
2022
2023 let build_list = |values: Vec<i32>| {
2024 let len = values.len();
2025 let item_struct = Arc::new(StructArray::new(
2026 Fields::from(vec![x_field.clone()]),
2027 vec![Arc::new(Int32Array::from(values))],
2028 None,
2029 ));
2030 ListArray::new(
2031 item_field.clone(),
2032 OffsetBuffer::from_lengths([len]),
2033 item_struct,
2034 None,
2035 )
2036 };
2037
2038 let left =
2041 RecordBatch::try_new(schema.clone(), vec![Arc::new(build_list(vec![1, 2]))]).unwrap();
2042 let right = RecordBatch::try_new(schema, vec![Arc::new(build_list(vec![3, 4]))]).unwrap();
2043
2044 let merged = left.merge(&right).unwrap();
2045
2046 assert_eq!(merged.num_columns(), 1);
2048 assert_eq!(merged, left);
2049 }
2050
2051 #[test]
2052 fn test_merge_nested_list_struct_identical_schema() {
2053 let x_field = Arc::new(Field::new("x", DataType::Int32, true));
2057 let item_field = Arc::new(Field::new(
2058 "item",
2059 DataType::Struct(Fields::from(vec![x_field.clone()])),
2060 true,
2061 ));
2062 let companies_field = Arc::new(Field::new(
2063 "companies",
2064 DataType::List(item_field.clone()),
2065 true,
2066 ));
2067 let schema = Arc::new(Schema::new(vec![Field::new(
2068 "outer",
2069 DataType::Struct(Fields::from(vec![companies_field.clone()])),
2070 true,
2071 )]));
2072
2073 let build_outer = |x: i32| {
2074 let item_struct = Arc::new(StructArray::new(
2076 Fields::from(vec![x_field.clone()]),
2077 vec![Arc::new(Int32Array::from(vec![x]))],
2078 None,
2079 ));
2080 let companies = Arc::new(ListArray::new(
2081 item_field.clone(),
2082 OffsetBuffer::from_lengths([1]),
2083 item_struct,
2084 None,
2085 ));
2086 StructArray::new(
2087 Fields::from(vec![companies_field.clone()]),
2088 vec![companies],
2089 None,
2090 )
2091 };
2092
2093 let left = RecordBatch::try_new(schema.clone(), vec![Arc::new(build_outer(10))]).unwrap();
2095 let right = RecordBatch::try_new(schema, vec![Arc::new(build_outer(20))]).unwrap();
2096
2097 let merged = left.merge(&right).unwrap();
2098
2099 assert_eq!(merged.column(0).as_struct().num_columns(), 1);
2102 assert_eq!(merged, left);
2103 }
2104
2105 #[test]
2106 fn test_byte_width_opt() {
2107 assert_eq!(DataType::Int32.byte_width_opt(), Some(4));
2108 assert_eq!(DataType::Int64.byte_width_opt(), Some(8));
2109 assert_eq!(DataType::Float32.byte_width_opt(), Some(4));
2110 assert_eq!(DataType::Float64.byte_width_opt(), Some(8));
2111 assert_eq!(DataType::Utf8.byte_width_opt(), None);
2112 assert_eq!(DataType::Binary.byte_width_opt(), None);
2113 assert_eq!(
2114 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))).byte_width_opt(),
2115 None
2116 );
2117 assert_eq!(
2118 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 3)
2119 .byte_width_opt(),
2120 Some(12)
2121 );
2122 assert_eq!(
2123 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Int32, true)), 4)
2124 .byte_width_opt(),
2125 Some(16)
2126 );
2127 assert_eq!(
2128 DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Utf8, true)), 5)
2129 .byte_width_opt(),
2130 None
2131 );
2132 }
2133
2134 #[test]
2135 fn test_take_record_batch() {
2136 let schema = Arc::new(Schema::new(vec![
2137 Field::new("a", DataType::Int32, true),
2138 Field::new("b", DataType::Utf8, true),
2139 ]));
2140 let batch = RecordBatch::try_new(
2141 schema.clone(),
2142 vec![
2143 Arc::new(Int32Array::from_iter_values(0..20)),
2144 Arc::new(StringArray::from_iter_values(
2145 (0..20).map(|i| format!("str-{}", i)),
2146 )),
2147 ],
2148 )
2149 .unwrap();
2150 let taken = batch.take(&(vec![1_u32, 5_u32, 10_u32].into())).unwrap();
2151 assert_eq!(
2152 taken,
2153 RecordBatch::try_new(
2154 schema,
2155 vec![
2156 Arc::new(Int32Array::from(vec![1, 5, 10])),
2157 Arc::new(StringArray::from(vec!["str-1", "str-5", "str-10"])),
2158 ],
2159 )
2160 .unwrap()
2161 )
2162 }
2163
2164 #[test]
2165 fn test_schema_project_by_schema() {
2166 let metadata = [("key".to_string(), "value".to_string())];
2167 let schema = Arc::new(
2168 Schema::new(vec![
2169 Field::new("a", DataType::Int32, true),
2170 Field::new("b", DataType::Utf8, true),
2171 ])
2172 .with_metadata(metadata.clone().into()),
2173 );
2174 let batch = RecordBatch::try_new(
2175 schema,
2176 vec![
2177 Arc::new(Int32Array::from_iter_values(0..20)),
2178 Arc::new(StringArray::from_iter_values(
2179 (0..20).map(|i| format!("str-{}", i)),
2180 )),
2181 ],
2182 )
2183 .unwrap();
2184
2185 let empty_schema = Schema::empty();
2187 let empty_projected = batch.project_by_schema(&empty_schema).unwrap();
2188 let expected_schema = empty_schema.with_metadata(metadata.clone().into());
2189 assert_eq!(
2190 empty_projected,
2191 RecordBatch::from(StructArray::new_empty_fields(batch.num_rows(), None))
2192 .with_schema(Arc::new(expected_schema))
2193 .unwrap()
2194 );
2195
2196 let reordered_schema = Schema::new(vec![
2198 Field::new("b", DataType::Utf8, true),
2199 Field::new("a", DataType::Int32, true),
2200 ]);
2201 let reordered_projected = batch.project_by_schema(&reordered_schema).unwrap();
2202 let expected_schema = Arc::new(reordered_schema.with_metadata(metadata.clone().into()));
2203 assert_eq!(
2204 reordered_projected,
2205 RecordBatch::try_new(
2206 expected_schema,
2207 vec![
2208 Arc::new(StringArray::from_iter_values(
2209 (0..20).map(|i| format!("str-{}", i)),
2210 )),
2211 Arc::new(Int32Array::from_iter_values(0..20)),
2212 ],
2213 )
2214 .unwrap()
2215 );
2216
2217 let sub_schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]);
2219 let sub_projected = batch.project_by_schema(&sub_schema).unwrap();
2220 let expected_schema = Arc::new(sub_schema.with_metadata(metadata.into()));
2221 assert_eq!(
2222 sub_projected,
2223 RecordBatch::try_new(
2224 expected_schema,
2225 vec![Arc::new(Int32Array::from_iter_values(0..20))],
2226 )
2227 .unwrap()
2228 );
2229 }
2230
2231 #[test]
2232 fn test_project_rebuilds_sliced_map() {
2233 let entry_fields = Fields::from(vec![
2237 Field::new("key", DataType::Utf8, false),
2238 Field::new("value", DataType::Int32, true),
2239 ]);
2240 let entries_field = Arc::new(Field::new(
2241 "entries",
2242 DataType::Struct(entry_fields.clone()),
2243 false,
2244 ));
2245 let entries = StructArray::new(
2246 entry_fields,
2247 vec![
2248 Arc::new(StringArray::from(vec!["k0", "k1", "k2"])) as ArrayRef,
2249 Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef,
2250 ],
2251 None,
2252 );
2253 let map = MapArray::new(
2254 entries_field.clone(),
2255 OffsetBuffer::new(vec![0, 1, 1, 3].into()),
2256 entries,
2257 Some(arrow_buffer::NullBuffer::from(vec![true, false, true])),
2258 false,
2259 );
2260 let schema = Arc::new(Schema::new(vec![Field::new(
2261 "m",
2262 DataType::Map(entries_field, false),
2263 true,
2264 )]));
2265 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(map) as ArrayRef]).unwrap();
2266
2267 let projected = batch
2269 .slice(1, 2)
2270 .project_by_schema(schema.as_ref())
2271 .unwrap();
2272 let map = projected.column(0).as_map();
2273 assert!(map.is_null(0));
2274 assert!(map.is_valid(1));
2275 assert_eq!(map.value_length(1), 2);
2276 assert_eq!(
2277 map.value(1)
2278 .column(1)
2279 .as_primitive::<arrow_array::types::Int32Type>()
2280 .values(),
2281 &[2, 3]
2282 );
2283 }
2284
2285 #[test]
2286 fn test_project_preserves_struct_validity() {
2287 let fields = Fields::from(vec![
2289 Field::new("id", DataType::Int32, false),
2290 Field::new("value", DataType::Float32, true),
2291 ]);
2292
2293 let id_array = Int32Array::from(vec![1, 2, 3]);
2295 let value_array = Float32Array::from(vec![Some(1.0), Some(2.0), Some(3.0)]);
2296 let struct_array = StructArray::new(
2297 fields.clone(),
2298 vec![
2299 Arc::new(id_array) as ArrayRef,
2300 Arc::new(value_array) as ArrayRef,
2301 ],
2302 Some(vec![true, false, true].into()), );
2304
2305 let projected = project(&struct_array, &fields).unwrap();
2307
2308 assert_eq!(projected.null_count(), 1);
2310 assert!(!projected.is_null(0));
2311 assert!(projected.is_null(1));
2312 assert!(!projected.is_null(2));
2313 }
2314
2315 #[test]
2316 fn test_merge_struct_with_different_validity() {
2317 let height_array = Int32Array::from(vec![Some(500), None, Some(600), None]);
2320 let left_fields = Fields::from(vec![Field::new("height", DataType::Int32, true)]);
2321 let left_struct = StructArray::new(
2322 left_fields,
2323 vec![Arc::new(height_array) as ArrayRef],
2324 Some(vec![true, false, true, false].into()), );
2326
2327 let width_array = Int32Array::from(vec![Some(300), Some(200), None, None]);
2329 let right_fields = Fields::from(vec![Field::new("width", DataType::Int32, true)]);
2330 let right_struct = StructArray::new(
2331 right_fields,
2332 vec![Arc::new(width_array) as ArrayRef],
2333 Some(vec![true, true, false, false].into()), );
2335
2336 let merged = merge(&left_struct, &right_struct);
2338
2339 assert_eq!(merged.null_count(), 1); assert!(!merged.is_null(0));
2347 assert!(!merged.is_null(1));
2348 assert!(!merged.is_null(2));
2349 assert!(merged.is_null(3));
2350
2351 let height_col = merged.column_by_name("height").unwrap();
2353 let height_values = height_col.as_any().downcast_ref::<Int32Array>().unwrap();
2354 assert_eq!(height_values.value(0), 500);
2355 assert!(height_values.is_null(1)); assert_eq!(height_values.value(2), 600);
2357
2358 let width_col = merged.column_by_name("width").unwrap();
2359 let width_values = width_col.as_any().downcast_ref::<Int32Array>().unwrap();
2360 assert_eq!(width_values.value(0), 300);
2361 assert_eq!(width_values.value(1), 200);
2362 assert!(width_values.is_null(2)); let all_null_left = StructArray::new(
2367 Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2368 vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef],
2369 Some(vec![false, false].into()),
2370 );
2371 let all_null_right = StructArray::new(
2372 Fields::from(vec![Field::new("width", DataType::Int32, true)]),
2373 vec![Arc::new(Int32Array::from(vec![None, None])) as ArrayRef],
2374 Some(vec![false, false].into()),
2375 );
2376
2377 let merged = merge(&all_null_left, &all_null_right);
2378 assert_eq!(merged.null_count(), 2);
2379
2380 let partial_left = StructArray::new(
2382 Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2383 vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef],
2384 Some(vec![true, false].into()),
2385 );
2386 let merged = merge(&partial_left, &all_null_right);
2387 assert!(!merged.is_null(0));
2388 assert!(merged.is_null(1));
2389
2390 let all_valid_left = StructArray::new(
2392 Fields::from(vec![Field::new("height", DataType::Int32, true)]),
2393 vec![Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef],
2394 None,
2395 );
2396 let merged = merge(&all_valid_left, &all_null_right);
2397 assert_eq!(merged.null_count(), 0);
2398
2399 let all_valid: arrow_buffer::NullBuffer = vec![true, true].into();
2401 let partial: arrow_buffer::NullBuffer = vec![true, false].into();
2402 assert!(merge_struct_validity(Some(&all_valid), Some(&partial)).is_none());
2403 assert!(merge_struct_validity(Some(&partial), Some(&all_valid)).is_none());
2404 }
2405
2406 #[test]
2407 fn test_merge_null_typed_column_with_parent_validity() {
2408 let left_struct = StructArray::new(
2412 Fields::from(vec![Field::new("a", DataType::Int32, true)]),
2413 vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef],
2414 Some(vec![true, false].into()),
2415 );
2416 let right_struct = StructArray::new(
2417 Fields::from(vec![Field::new("b", DataType::Null, true)]),
2418 vec![Arc::new(NullArray::new(2)) as ArrayRef],
2419 Some(vec![true, false].into()),
2420 );
2421
2422 let merged = merge(&left_struct, &right_struct);
2424 assert_eq!(merged.len(), 2);
2425 let b_col = merged.column_by_name("b").unwrap();
2426 assert_eq!(b_col.data_type(), &DataType::Null);
2428 assert_eq!(b_col.len(), 2);
2429 }
2430
2431 #[test]
2432 fn test_merge_with_schema_with_nullable_struct_list_schema_mismatch() {
2433 let left_company_id = Arc::new(Int32Array::from(vec![None, None]));
2435 let left_count = Arc::new(Int32Array::from(vec![None, None]));
2436 let left_struct = Arc::new(StructArray::new(
2437 Fields::from(vec![
2438 Field::new("company_id", DataType::Int32, true),
2439 Field::new("count", DataType::Int32, true),
2440 ]),
2441 vec![left_company_id, left_count],
2442 None,
2443 ));
2444 let left_list = Arc::new(ListArray::new(
2445 Arc::new(Field::new(
2446 "item",
2447 DataType::Struct(left_struct.fields().clone()),
2448 true,
2449 )),
2450 OffsetBuffer::from_lengths([2]),
2451 left_struct,
2452 None,
2453 ));
2454
2455 let right_company_name = Arc::new(StringArray::from(vec!["Google", "Microsoft"]));
2457 let right_struct = Arc::new(StructArray::new(
2458 Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2459 vec![right_company_name],
2460 None,
2461 ));
2462 let right_list = Arc::new(ListArray::new(
2463 Arc::new(Field::new(
2464 "item",
2465 DataType::Struct(right_struct.fields().clone()),
2466 true,
2467 )),
2468 OffsetBuffer::from_lengths([2]),
2469 right_struct,
2470 None,
2471 ));
2472
2473 let target_fields = Fields::from(vec![Field::new(
2474 "companies",
2475 DataType::List(Arc::new(Field::new(
2476 "item",
2477 DataType::Struct(Fields::from(vec![
2478 Field::new("company_id", DataType::Int32, true),
2479 Field::new("company_name", DataType::Utf8, true),
2480 Field::new("count", DataType::Int32, true),
2481 ])),
2482 true,
2483 ))),
2484 true,
2485 )]);
2486
2487 let left_batch = RecordBatch::try_new(
2488 Arc::new(Schema::new(vec![Field::new(
2489 "companies",
2490 left_list.data_type().clone(),
2491 true,
2492 )])),
2493 vec![left_list as ArrayRef],
2494 )
2495 .unwrap();
2496
2497 let right_batch = RecordBatch::try_new(
2498 Arc::new(Schema::new(vec![Field::new(
2499 "companies",
2500 right_list.data_type().clone(),
2501 true,
2502 )])),
2503 vec![right_list as ArrayRef],
2504 )
2505 .unwrap();
2506
2507 let merged = left_batch
2508 .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2509 .unwrap();
2510
2511 let merged_list = merged
2513 .column_by_name("companies")
2514 .unwrap()
2515 .as_any()
2516 .downcast_ref::<ListArray>()
2517 .unwrap();
2518 let merged_struct = merged_list.values().as_struct();
2519
2520 assert_eq!(merged_struct.num_columns(), 3);
2522 assert!(merged_struct.column_by_name("company_id").is_some());
2523 assert!(merged_struct.column_by_name("company_name").is_some());
2524 assert!(merged_struct.column_by_name("count").is_some());
2525
2526 let company_id = merged_struct
2528 .column_by_name("company_id")
2529 .unwrap()
2530 .as_any()
2531 .downcast_ref::<Int32Array>()
2532 .unwrap();
2533 assert!(company_id.is_null(0));
2534 assert!(company_id.is_null(1));
2535
2536 let company_name = merged_struct
2537 .column_by_name("company_name")
2538 .unwrap()
2539 .as_any()
2540 .downcast_ref::<StringArray>()
2541 .unwrap();
2542 assert_eq!(company_name.value(0), "Google");
2543 assert_eq!(company_name.value(1), "Microsoft");
2544
2545 let count = merged_struct
2546 .column_by_name("count")
2547 .unwrap()
2548 .as_any()
2549 .downcast_ref::<Int32Array>()
2550 .unwrap();
2551 assert!(count.is_null(0));
2552 assert!(count.is_null(1));
2553 }
2554
2555 #[test]
2556 fn test_merge_struct_lists() {
2557 test_merge_struct_lists_generic::<i32>();
2558 }
2559
2560 #[test]
2561 fn test_merge_struct_large_lists() {
2562 test_merge_struct_lists_generic::<i64>();
2563 }
2564
2565 fn test_merge_struct_lists_generic<O: OffsetSizeTrait>() {
2566 let left_company_id = Arc::new(Int32Array::from(vec![
2568 Some(1),
2569 Some(2),
2570 Some(3),
2571 Some(4),
2572 Some(5),
2573 Some(6),
2574 Some(7),
2575 Some(8),
2576 Some(9),
2577 Some(10),
2578 Some(11),
2579 Some(12),
2580 Some(13),
2581 Some(14),
2582 Some(15),
2583 Some(16),
2584 Some(17),
2585 Some(18),
2586 Some(19),
2587 Some(20),
2588 ]));
2589 let left_count = Arc::new(Int32Array::from(vec![
2590 Some(10),
2591 Some(20),
2592 Some(30),
2593 Some(40),
2594 Some(50),
2595 Some(60),
2596 Some(70),
2597 Some(80),
2598 Some(90),
2599 Some(100),
2600 Some(110),
2601 Some(120),
2602 Some(130),
2603 Some(140),
2604 Some(150),
2605 Some(160),
2606 Some(170),
2607 Some(180),
2608 Some(190),
2609 Some(200),
2610 ]));
2611 let left_struct = Arc::new(StructArray::new(
2612 Fields::from(vec![
2613 Field::new("company_id", DataType::Int32, true),
2614 Field::new("count", DataType::Int32, true),
2615 ]),
2616 vec![left_company_id, left_count],
2617 None,
2618 ));
2619
2620 let left_list = Arc::new(GenericListArray::<O>::new(
2621 Arc::new(Field::new(
2622 "item",
2623 DataType::Struct(left_struct.fields().clone()),
2624 true,
2625 )),
2626 OffsetBuffer::from_lengths([3, 1]),
2627 left_struct.clone(),
2628 None,
2629 ));
2630
2631 let left_list_struct = Arc::new(StructArray::new(
2632 Fields::from(vec![Field::new(
2633 "companies",
2634 if O::IS_LARGE {
2635 DataType::LargeList(Arc::new(Field::new(
2636 "item",
2637 DataType::Struct(left_struct.fields().clone()),
2638 true,
2639 )))
2640 } else {
2641 DataType::List(Arc::new(Field::new(
2642 "item",
2643 DataType::Struct(left_struct.fields().clone()),
2644 true,
2645 )))
2646 },
2647 true,
2648 )]),
2649 vec![left_list as ArrayRef],
2650 None,
2651 ));
2652
2653 let right_company_name = Arc::new(StringArray::from(vec![
2655 "Google",
2656 "Microsoft",
2657 "Apple",
2658 "Facebook",
2659 ]));
2660 let right_struct = Arc::new(StructArray::new(
2661 Fields::from(vec![Field::new("company_name", DataType::Utf8, true)]),
2662 vec![right_company_name],
2663 None,
2664 ));
2665 let right_list = Arc::new(GenericListArray::<O>::new(
2666 Arc::new(Field::new(
2667 "item",
2668 DataType::Struct(right_struct.fields().clone()),
2669 true,
2670 )),
2671 OffsetBuffer::from_lengths([3, 1]),
2672 right_struct.clone(),
2673 None,
2674 ));
2675
2676 let right_list_struct = Arc::new(StructArray::new(
2677 Fields::from(vec![Field::new(
2678 "companies",
2679 if O::IS_LARGE {
2680 DataType::LargeList(Arc::new(Field::new(
2681 "item",
2682 DataType::Struct(right_struct.fields().clone()),
2683 true,
2684 )))
2685 } else {
2686 DataType::List(Arc::new(Field::new(
2687 "item",
2688 DataType::Struct(right_struct.fields().clone()),
2689 true,
2690 )))
2691 },
2692 true,
2693 )]),
2694 vec![right_list as ArrayRef],
2695 None,
2696 ));
2697
2698 let target_fields = Fields::from(vec![Field::new(
2700 "companies",
2701 if O::IS_LARGE {
2702 DataType::LargeList(Arc::new(Field::new(
2703 "item",
2704 DataType::Struct(Fields::from(vec![
2705 Field::new("company_id", DataType::Int32, true),
2706 Field::new("company_name", DataType::Utf8, true),
2707 Field::new("count", DataType::Int32, true),
2708 ])),
2709 true,
2710 )))
2711 } else {
2712 DataType::List(Arc::new(Field::new(
2713 "item",
2714 DataType::Struct(Fields::from(vec![
2715 Field::new("company_id", DataType::Int32, true),
2716 Field::new("company_name", DataType::Utf8, true),
2717 Field::new("count", DataType::Int32, true),
2718 ])),
2719 true,
2720 )))
2721 },
2722 true,
2723 )]);
2724
2725 let merged_array = merge_with_schema(&left_list_struct, &right_list_struct, &target_fields);
2727 assert_eq!(merged_array.len(), 2);
2728 }
2729
2730 #[test]
2731 fn test_merge_with_schema_sliced_list_struct() {
2732 test_merge_with_schema_sliced_list_struct_generic::<i32>();
2733 }
2734
2735 #[test]
2736 fn test_merge_with_schema_sliced_large_list_struct() {
2737 test_merge_with_schema_sliced_list_struct_generic::<i64>();
2738 }
2739
2740 fn test_merge_with_schema_sliced_list_struct_generic<O: OffsetSizeTrait>() {
2745 let make_list_dtype = |item_field: Arc<Field>| {
2746 if O::IS_LARGE {
2747 DataType::LargeList(item_field)
2748 } else {
2749 DataType::List(item_field)
2750 }
2751 };
2752
2753 let struct_fields_a = Fields::from(vec![Field::new("a", DataType::Int32, true)]);
2756 let left_values = Arc::new(StructArray::new(
2757 struct_fields_a.clone(),
2758 vec![Arc::new(Int32Array::from_iter_values(0..10)) as ArrayRef],
2759 None,
2760 ));
2761 let full_list = GenericListArray::<O>::new(
2762 Arc::new(Field::new("item", DataType::Struct(struct_fields_a), true)),
2763 OffsetBuffer::<O>::from_lengths([5, 5]),
2764 left_values,
2765 None,
2766 );
2767 let sliced_left = full_list.slice(1, 1);
2768 assert_eq!(sliced_left.offsets()[0].as_usize(), 5);
2769 assert_eq!(sliced_left.offsets()[1].as_usize(), 10);
2770
2771 let struct_fields_b = Fields::from(vec![Field::new("b", DataType::Int32, true)]);
2772 let right_values = Arc::new(StructArray::new(
2773 struct_fields_b.clone(),
2774 vec![Arc::new(Int32Array::from_iter_values(100..105)) as ArrayRef],
2775 None,
2776 ));
2777 let right_list = GenericListArray::<O>::new(
2778 Arc::new(Field::new("item", DataType::Struct(struct_fields_b), true)),
2779 OffsetBuffer::<O>::from_lengths([5]),
2780 right_values,
2781 None,
2782 );
2783
2784 let target_item_field = Arc::new(Field::new(
2785 "item",
2786 DataType::Struct(Fields::from(vec![
2787 Field::new("a", DataType::Int32, true),
2788 Field::new("b", DataType::Int32, true),
2789 ])),
2790 true,
2791 ));
2792 let target_fields = Fields::from(vec![Field::new(
2793 "items",
2794 make_list_dtype(target_item_field),
2795 true,
2796 )]);
2797
2798 let left_batch = RecordBatch::try_new(
2799 Arc::new(Schema::new(vec![Field::new(
2800 "items",
2801 sliced_left.data_type().clone(),
2802 true,
2803 )])),
2804 vec![Arc::new(sliced_left) as ArrayRef],
2805 )
2806 .unwrap();
2807 let right_batch = RecordBatch::try_new(
2808 Arc::new(Schema::new(vec![Field::new(
2809 "items",
2810 right_list.data_type().clone(),
2811 true,
2812 )])),
2813 vec![Arc::new(right_list) as ArrayRef],
2814 )
2815 .unwrap();
2816
2817 let merged = left_batch
2818 .merge_with_schema(&right_batch, &Schema::new(target_fields.to_vec()))
2819 .unwrap();
2820
2821 let merged_list = merged
2822 .column_by_name("items")
2823 .unwrap()
2824 .as_any()
2825 .downcast_ref::<GenericListArray<O>>()
2826 .unwrap();
2827 assert_eq!(merged_list.len(), 1);
2828 assert_eq!(merged_list.value_length(0).as_usize(), 5);
2829 let merged_struct = merged_list.values().as_struct();
2830 assert_eq!(merged_struct.num_columns(), 2);
2831 let a = merged_struct
2832 .column_by_name("a")
2833 .unwrap()
2834 .as_any()
2835 .downcast_ref::<Int32Array>()
2836 .unwrap();
2837 let a_vals: Vec<i32> = a.iter().map(|v| v.unwrap()).collect();
2839 assert_eq!(a_vals, vec![5, 6, 7, 8, 9]);
2840 }
2841
2842 #[test]
2843 fn test_project_by_schema_list_struct_reorder() {
2844 let source_inner_struct = DataType::Struct(Fields::from(vec![
2849 Field::new("c", DataType::Utf8, true),
2850 Field::new("b", DataType::Utf8, true),
2851 Field::new("a", DataType::Utf8, true),
2852 ]));
2853 let source_schema = Arc::new(Schema::new(vec![
2854 Field::new("id", DataType::Int32, false),
2855 Field::new(
2856 "data",
2857 DataType::List(Arc::new(Field::new(
2858 "item",
2859 source_inner_struct.clone(),
2860 true,
2861 ))),
2862 true,
2863 ),
2864 ]));
2865
2866 let c_array = StringArray::from(vec!["c1", "c2"]);
2868 let b_array = StringArray::from(vec!["b1", "b2"]);
2869 let a_array = StringArray::from(vec!["a1", "a2"]);
2870 let inner_struct = StructArray::from(vec![
2871 (
2872 Arc::new(Field::new("c", DataType::Utf8, true)),
2873 Arc::new(c_array) as ArrayRef,
2874 ),
2875 (
2876 Arc::new(Field::new("b", DataType::Utf8, true)),
2877 Arc::new(b_array) as ArrayRef,
2878 ),
2879 (
2880 Arc::new(Field::new("a", DataType::Utf8, true)),
2881 Arc::new(a_array) as ArrayRef,
2882 ),
2883 ]);
2884
2885 let list_array = ListArray::new(
2886 Arc::new(Field::new("item", source_inner_struct, true)),
2887 OffsetBuffer::from_lengths([1, 1]),
2888 Arc::new(inner_struct),
2889 None,
2890 );
2891
2892 let batch = RecordBatch::try_new(
2893 source_schema,
2894 vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(list_array)],
2895 )
2896 .unwrap();
2897
2898 let target_inner_struct = DataType::Struct(Fields::from(vec![
2900 Field::new("a", DataType::Utf8, true),
2901 Field::new("b", DataType::Utf8, true),
2902 Field::new("c", DataType::Utf8, true),
2903 ]));
2904 let target_schema = Schema::new(vec![
2905 Field::new("id", DataType::Int32, false),
2906 Field::new(
2907 "data",
2908 DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
2909 true,
2910 ),
2911 ]);
2912
2913 let projected = batch.project_by_schema(&target_schema).unwrap();
2915
2916 assert_eq!(projected.schema().as_ref(), &target_schema);
2918
2919 let projected_list = projected.column(1).as_list::<i32>();
2921 let projected_struct = projected_list.values().as_struct();
2922
2923 assert_eq!(
2925 projected_struct.column_by_name("a").unwrap().as_ref(),
2926 &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2927 );
2928 assert_eq!(
2929 projected_struct.column_by_name("b").unwrap().as_ref(),
2930 &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2931 );
2932 assert_eq!(
2933 projected_struct.column_by_name("c").unwrap().as_ref(),
2934 &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2935 );
2936
2937 assert_eq!(
2939 projected_struct.column(0).as_ref(),
2940 &StringArray::from(vec!["a1", "a2"]) as &dyn Array
2941 );
2942 assert_eq!(
2943 projected_struct.column(1).as_ref(),
2944 &StringArray::from(vec!["b1", "b2"]) as &dyn Array
2945 );
2946 assert_eq!(
2947 projected_struct.column(2).as_ref(),
2948 &StringArray::from(vec!["c1", "c2"]) as &dyn Array
2949 );
2950 }
2951
2952 #[test]
2953 fn test_project_by_schema_nested_list_struct() {
2954 let inner_struct = DataType::Struct(Fields::from(vec![
2956 Field::new("y", DataType::Int32, true),
2957 Field::new("x", DataType::Int32, true),
2958 ]));
2959 let source_schema = Arc::new(Schema::new(vec![Field::new(
2960 "outer",
2961 DataType::List(Arc::new(Field::new(
2962 "item",
2963 DataType::Struct(Fields::from(vec![
2964 Field::new("b", DataType::Utf8, true),
2965 Field::new(
2966 "inner_list",
2967 DataType::List(Arc::new(Field::new("item", inner_struct.clone(), true))),
2968 true,
2969 ),
2970 Field::new("a", DataType::Utf8, true),
2971 ])),
2972 true,
2973 ))),
2974 true,
2975 )]));
2976
2977 let y_array = Int32Array::from(vec![1, 2]);
2979 let x_array = Int32Array::from(vec![3, 4]);
2980 let innermost_struct = StructArray::from(vec![
2981 (
2982 Arc::new(Field::new("y", DataType::Int32, true)),
2983 Arc::new(y_array) as ArrayRef,
2984 ),
2985 (
2986 Arc::new(Field::new("x", DataType::Int32, true)),
2987 Arc::new(x_array) as ArrayRef,
2988 ),
2989 ]);
2990 let inner_list = ListArray::new(
2991 Arc::new(Field::new("item", inner_struct.clone(), true)),
2992 OffsetBuffer::from_lengths([2]),
2993 Arc::new(innermost_struct),
2994 None,
2995 );
2996
2997 let b_array = StringArray::from(vec!["b1"]);
2998 let a_array = StringArray::from(vec!["a1"]);
2999 let middle_struct = StructArray::from(vec![
3000 (
3001 Arc::new(Field::new("b", DataType::Utf8, true)),
3002 Arc::new(b_array) as ArrayRef,
3003 ),
3004 (
3005 Arc::new(Field::new(
3006 "inner_list",
3007 DataType::List(Arc::new(Field::new("item", inner_struct, true))),
3008 true,
3009 )),
3010 Arc::new(inner_list) as ArrayRef,
3011 ),
3012 (
3013 Arc::new(Field::new("a", DataType::Utf8, true)),
3014 Arc::new(a_array) as ArrayRef,
3015 ),
3016 ]);
3017
3018 let outer_list = ListArray::new(
3019 Arc::new(Field::new("item", middle_struct.data_type().clone(), true)),
3020 OffsetBuffer::from_lengths([1]),
3021 Arc::new(middle_struct),
3022 None,
3023 );
3024
3025 let batch =
3026 RecordBatch::try_new(source_schema, vec![Arc::new(outer_list) as ArrayRef]).unwrap();
3027
3028 let target_inner_struct = DataType::Struct(Fields::from(vec![
3030 Field::new("x", DataType::Int32, true), Field::new("y", DataType::Int32, true),
3032 ]));
3033 let target_schema = Schema::new(vec![Field::new(
3034 "outer",
3035 DataType::List(Arc::new(Field::new(
3036 "item",
3037 DataType::Struct(Fields::from(vec![
3038 Field::new("a", DataType::Utf8, true), Field::new(
3040 "inner_list",
3041 DataType::List(Arc::new(Field::new("item", target_inner_struct, true))),
3042 true,
3043 ),
3044 Field::new("b", DataType::Utf8, true),
3045 ])),
3046 true,
3047 ))),
3048 true,
3049 )]);
3050
3051 let projected = batch.project_by_schema(&target_schema).unwrap();
3052
3053 assert_eq!(projected.schema().as_ref(), &target_schema);
3055
3056 let outer_list = projected.column(0).as_list::<i32>();
3058 let middle_struct = outer_list.values().as_struct();
3059
3060 assert_eq!(
3062 middle_struct.column(0).as_ref(),
3063 &StringArray::from(vec!["a1"]) as &dyn Array
3064 );
3065 assert_eq!(
3066 middle_struct.column(2).as_ref(),
3067 &StringArray::from(vec!["b1"]) as &dyn Array
3068 );
3069
3070 let inner_list = middle_struct.column(1).as_list::<i32>();
3072 let innermost_struct = inner_list.values().as_struct();
3073 assert_eq!(
3074 innermost_struct.column(0).as_ref(),
3075 &Int32Array::from(vec![3, 4]) as &dyn Array
3076 );
3077 assert_eq!(
3078 innermost_struct.column(1).as_ref(),
3079 &Int32Array::from(vec![1, 2]) as &dyn Array
3080 );
3081 }
3082
3083 #[test]
3084 fn test_iter_binary_array_accepts_binary_variants() {
3085 let binary = BinaryArray::from(vec![b"a".as_slice(), b"bc"]);
3086 assert_eq!(
3087 iter_binary_array(&binary).unwrap().collect::<Vec<_>>(),
3088 vec![Some(b"a".as_slice()), Some(b"bc".as_slice())]
3089 );
3090
3091 let large_binary = LargeBinaryArray::from(vec![b"x".as_slice(), b"yz"]);
3092 assert_eq!(
3093 iter_binary_array(&large_binary)
3094 .unwrap()
3095 .collect::<Vec<_>>(),
3096 vec![Some(b"x".as_slice()), Some(b"yz".as_slice())]
3097 );
3098
3099 let binary_view = BinaryViewArray::from(vec![b"1".as_slice(), b"23"]);
3100 assert_eq!(
3101 iter_binary_array(&binary_view).unwrap().collect::<Vec<_>>(),
3102 vec![Some(b"1".as_slice()), Some(b"23".as_slice())]
3103 );
3104
3105 let fixed_size =
3106 FixedSizeBinaryArray::try_from_iter([b"abcd", b"efgh"].into_iter()).unwrap();
3107 assert_eq!(
3108 iter_binary_array(&fixed_size).unwrap().collect::<Vec<_>>(),
3109 vec![Some(b"abcd".as_slice()), Some(b"efgh".as_slice())]
3110 );
3111 }
3112
3113 #[test]
3114 fn test_iter_binary_array_rejects_non_binary() {
3115 let int_array = Int32Array::from(vec![1, 2, 3]);
3116 let Err(error) = iter_binary_array(&int_array) else {
3117 panic!("expected an error for non-binary array");
3118 };
3119 assert!(error.to_string().contains("Expecting a binary type"));
3120 }
3121}