1#![allow(deprecated)]
12
13use std::ops::Range;
14use std::sync::Arc;
15
16use arrow_array::AnyDictionaryArray;
17use arrow_array::Array as ArrowArray;
18use arrow_array::ArrowPrimitiveType;
19use arrow_array::BooleanArray as ArrowBooleanArray;
20use arrow_array::DictionaryArray;
21use arrow_array::FixedSizeListArray as ArrowFixedSizeListArray;
22use arrow_array::GenericByteArray;
23use arrow_array::GenericByteViewArray;
24use arrow_array::GenericListArray;
25use arrow_array::GenericListViewArray;
26use arrow_array::MapArray as ArrowMapArray;
27use arrow_array::NullArray as ArrowNullArray;
28use arrow_array::OffsetSizeTrait;
29use arrow_array::PrimitiveArray as ArrowPrimitiveArray;
30use arrow_array::RecordBatch;
31use arrow_array::StructArray as ArrowStructArray;
32use arrow_array::cast::AsArray;
33use arrow_array::cast::as_null_array;
34use arrow_array::make_array;
35use arrow_array::types::ArrowDictionaryKeyType;
36use arrow_array::types::ByteArrayType;
37use arrow_array::types::ByteViewType;
38use arrow_array::types::Date32Type;
39use arrow_array::types::Date64Type;
40use arrow_array::types::Decimal32Type;
41use arrow_array::types::Decimal64Type;
42use arrow_array::types::Decimal128Type;
43use arrow_array::types::Decimal256Type;
44use arrow_array::types::Float16Type;
45use arrow_array::types::Float32Type;
46use arrow_array::types::Float64Type;
47use arrow_array::types::Int8Type;
48use arrow_array::types::Int16Type;
49use arrow_array::types::Int32Type;
50use arrow_array::types::Int64Type;
51use arrow_array::types::Time32MillisecondType;
52use arrow_array::types::Time32SecondType;
53use arrow_array::types::Time64MicrosecondType;
54use arrow_array::types::Time64NanosecondType;
55use arrow_array::types::TimestampMicrosecondType;
56use arrow_array::types::TimestampMillisecondType;
57use arrow_array::types::TimestampNanosecondType;
58use arrow_array::types::TimestampSecondType;
59use arrow_array::types::UInt8Type;
60use arrow_array::types::UInt16Type;
61use arrow_array::types::UInt32Type;
62use arrow_array::types::UInt64Type;
63use arrow_buffer::ArrowNativeType;
64use arrow_buffer::BooleanBuffer;
65use arrow_buffer::Buffer as ArrowBuffer;
66use arrow_buffer::ScalarBuffer;
67use arrow_buffer::buffer::NullBuffer;
68use arrow_buffer::buffer::OffsetBuffer;
69use arrow_schema::DataType;
70use arrow_schema::TimeUnit as ArrowTimeUnit;
71use vortex_array::ArrayRef;
72use vortex_array::IntoArray;
73use vortex_array::arrays::BoolArray;
74use vortex_array::arrays::DecimalArray;
75use vortex_array::arrays::DictArray;
76use vortex_array::arrays::FixedSizeListArray;
77use vortex_array::arrays::ListArray;
78use vortex_array::arrays::ListViewArray;
79use vortex_array::arrays::MapArray;
80use vortex_array::arrays::NullArray;
81use vortex_array::arrays::PrimitiveArray;
82use vortex_array::arrays::Struct;
83use vortex_array::arrays::StructArray;
84use vortex_array::arrays::TemporalArray;
85use vortex_array::arrays::VarBinArray;
86use vortex_array::arrays::VarBinViewArray;
87use vortex_array::arrays::struct_::StructArrayExt;
88use vortex_array::dtype::DType;
89use vortex_array::dtype::DecimalDType;
90use vortex_array::dtype::FieldNames;
91use vortex_array::dtype::IntegerPType;
92use vortex_array::dtype::MapDType;
93use vortex_array::dtype::NativePType;
94use vortex_array::dtype::Nullability;
95use vortex_array::dtype::PType;
96use vortex_array::dtype::i256;
97use vortex_array::extension::datetime::TimeUnit;
98use vortex_array::validity::Validity;
99use vortex_buffer::Alignment;
100use vortex_buffer::BitBuffer;
101use vortex_buffer::Buffer;
102use vortex_buffer::ByteBuffer;
103use vortex_error::VortexResult;
104use vortex_error::vortex_bail;
105use vortex_error::vortex_ensure;
106use vortex_error::vortex_ensure_eq;
107use vortex_error::vortex_err;
108use vortex_error::vortex_panic;
109
110use crate::FromArrowArray;
111use crate::dtype::from_arrow_time_unit;
112
113pub trait IntoVortexArray {
118 fn into_array(self) -> ArrayRef;
120}
121
122impl IntoVortexArray for ArrowBuffer {
123 fn into_array(self) -> ArrayRef {
124 PrimitiveArray::from_byte_buffer(
125 ByteBuffer::from_arrow_buffer(self, Alignment::of::<u8>()),
126 PType::U8,
127 Validity::NonNullable,
128 )
129 .into_array()
130 }
131}
132
133impl IntoVortexArray for BooleanBuffer {
134 fn into_array(self) -> ArrayRef {
135 BoolArray::new(self.into(), Validity::NonNullable).into_array()
136 }
137}
138
139impl<T> IntoVortexArray for ScalarBuffer<T>
140where
141 T: ArrowNativeType + NativePType,
142{
143 fn into_array(self) -> ArrayRef {
144 PrimitiveArray::new(
145 Buffer::<T>::from_arrow_scalar_buffer(self),
146 Validity::NonNullable,
147 )
148 .into_array()
149 }
150}
151
152impl<O> IntoVortexArray for OffsetBuffer<O>
153where
154 O: IntegerPType + OffsetSizeTrait,
155{
156 fn into_array(self) -> ArrayRef {
157 let primitive = PrimitiveArray::new(
158 Buffer::from_arrow_scalar_buffer(self.into_inner()),
159 Validity::NonNullable,
160 );
161
162 primitive.into_array()
163 }
164}
165
166pub fn from_arrow_primitive<T: ArrowPrimitiveType>(
171 value: &ArrowPrimitiveArray<T>,
172 nullable: bool,
173) -> VortexResult<ArrayRef>
174where
175 T::Native: NativePType,
176{
177 let buffer = Buffer::from_arrow_scalar_buffer(value.values().clone());
178 let validity = nulls(value.nulls(), nullable)?;
179 Ok(PrimitiveArray::new(buffer, validity).into_array())
180}
181
182macro_rules! impl_from_arrow_primitive {
183 ($T:path) => {
184 impl FromArrowArray<&ArrowPrimitiveArray<$T>> for ArrayRef {
185 fn from_arrow(value: &ArrowPrimitiveArray<$T>, nullable: bool) -> VortexResult<Self> {
186 from_arrow_primitive(value, nullable)
187 }
188 }
189 };
190}
191
192impl_from_arrow_primitive!(Int8Type);
193impl_from_arrow_primitive!(Int16Type);
194impl_from_arrow_primitive!(Int32Type);
195impl_from_arrow_primitive!(Int64Type);
196impl_from_arrow_primitive!(UInt8Type);
197impl_from_arrow_primitive!(UInt16Type);
198impl_from_arrow_primitive!(UInt32Type);
199impl_from_arrow_primitive!(UInt64Type);
200impl_from_arrow_primitive!(Float16Type);
201impl_from_arrow_primitive!(Float32Type);
202impl_from_arrow_primitive!(Float64Type);
203
204pub fn from_arrow_decimal32(
206 array: &ArrowPrimitiveArray<Decimal32Type>,
207 nullable: bool,
208) -> VortexResult<ArrayRef> {
209 let decimal_type = DecimalDType::new(array.precision(), array.scale());
210 let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
211 let validity = nulls(array.nulls(), nullable)?;
212 Ok(DecimalArray::new(buffer, decimal_type, validity).into_array())
213}
214
215impl FromArrowArray<&ArrowPrimitiveArray<Decimal32Type>> for ArrayRef {
216 fn from_arrow(
217 array: &ArrowPrimitiveArray<Decimal32Type>,
218 nullable: bool,
219 ) -> VortexResult<Self> {
220 from_arrow_decimal32(array, nullable)
221 }
222}
223
224pub fn from_arrow_decimal64(
226 array: &ArrowPrimitiveArray<Decimal64Type>,
227 nullable: bool,
228) -> VortexResult<ArrayRef> {
229 let decimal_type = DecimalDType::new(array.precision(), array.scale());
230 let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
231 let validity = nulls(array.nulls(), nullable)?;
232 Ok(DecimalArray::new(buffer, decimal_type, validity).into_array())
233}
234
235impl FromArrowArray<&ArrowPrimitiveArray<Decimal64Type>> for ArrayRef {
236 fn from_arrow(
237 array: &ArrowPrimitiveArray<Decimal64Type>,
238 nullable: bool,
239 ) -> VortexResult<Self> {
240 from_arrow_decimal64(array, nullable)
241 }
242}
243
244pub fn from_arrow_decimal128(
246 array: &ArrowPrimitiveArray<Decimal128Type>,
247 nullable: bool,
248) -> VortexResult<ArrayRef> {
249 let decimal_type = DecimalDType::new(array.precision(), array.scale());
250 let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
251 let validity = nulls(array.nulls(), nullable)?;
252 Ok(DecimalArray::new(buffer, decimal_type, validity).into_array())
253}
254
255impl FromArrowArray<&ArrowPrimitiveArray<Decimal128Type>> for ArrayRef {
256 fn from_arrow(
257 array: &ArrowPrimitiveArray<Decimal128Type>,
258 nullable: bool,
259 ) -> VortexResult<Self> {
260 from_arrow_decimal128(array, nullable)
261 }
262}
263
264pub fn from_arrow_decimal256(
266 array: &ArrowPrimitiveArray<Decimal256Type>,
267 nullable: bool,
268) -> VortexResult<ArrayRef> {
269 let decimal_type = DecimalDType::new(array.precision(), array.scale());
270 let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
271 let buffer = unsafe { std::mem::transmute::<Buffer<arrow_buffer::i256>, Buffer<i256>>(buffer) };
275 let validity = nulls(array.nulls(), nullable)?;
276 Ok(DecimalArray::new(buffer, decimal_type, validity).into_array())
277}
278
279impl FromArrowArray<&ArrowPrimitiveArray<Decimal256Type>> for ArrayRef {
280 fn from_arrow(
281 array: &ArrowPrimitiveArray<Decimal256Type>,
282 nullable: bool,
283 ) -> VortexResult<Self> {
284 from_arrow_decimal256(array, nullable)
285 }
286}
287
288macro_rules! impl_from_arrow_temporal {
289 ($T:path) => {
290 impl FromArrowArray<&ArrowPrimitiveArray<$T>> for ArrayRef {
291 fn from_arrow(
292 value: &ArrowPrimitiveArray<$T>,
293 nullable: bool,
294 ) -> vortex_error::VortexResult<Self> {
295 from_arrow_temporal(value, nullable)
296 }
297 }
298 };
299}
300
301impl_from_arrow_temporal!(TimestampSecondType);
303impl_from_arrow_temporal!(TimestampMillisecondType);
304impl_from_arrow_temporal!(TimestampMicrosecondType);
305impl_from_arrow_temporal!(TimestampNanosecondType);
306
307impl_from_arrow_temporal!(Time32SecondType);
309impl_from_arrow_temporal!(Time32MillisecondType);
310impl_from_arrow_temporal!(Time64MicrosecondType);
311impl_from_arrow_temporal!(Time64NanosecondType);
312
313impl_from_arrow_temporal!(Date32Type);
315impl_from_arrow_temporal!(Date64Type);
316
317pub fn from_arrow_temporal<T: ArrowPrimitiveType>(
320 value: &ArrowPrimitiveArray<T>,
321 nullable: bool,
322) -> VortexResult<ArrayRef>
323where
324 T::Native: NativePType,
325{
326 let arr = PrimitiveArray::new(
327 Buffer::from_arrow_scalar_buffer(value.values().clone()),
328 nulls(value.nulls(), nullable)?,
329 )
330 .into_array();
331
332 Ok(match value.data_type() {
333 DataType::Timestamp(time_unit, tz) => {
334 TemporalArray::new_timestamp(arr, from_arrow_time_unit(*time_unit), tz.clone()).into()
335 }
336 DataType::Time32(time_unit) => {
337 TemporalArray::new_time(arr, from_arrow_time_unit(*time_unit)).into()
338 }
339 DataType::Time64(time_unit) => {
340 TemporalArray::new_time(arr, from_arrow_time_unit(*time_unit)).into()
341 }
342 DataType::Date32 => TemporalArray::new_date(arr, TimeUnit::Days).into(),
343 DataType::Date64 => TemporalArray::new_date(arr, TimeUnit::Milliseconds).into(),
344 DataType::Duration(_) | DataType::Interval(_) => {
345 vortex_bail!("Arrow {} arrays are not supported", value.data_type())
346 }
347 _ => vortex_panic!("Invalid temporal type: {}", value.data_type()),
348 })
349}
350
351pub(crate) fn trim_offsets<O: OffsetSizeTrait>(
356 offsets: &OffsetBuffer<O>,
357) -> (OffsetBuffer<O>, Range<usize>) {
358 let first = offsets[0];
359 let referenced = first.as_usize()..offsets[offsets.len() - 1].as_usize();
360 if first.as_usize() == 0 {
361 return (offsets.clone(), referenced);
362 }
363 let rebased = OffsetBuffer::new(offsets.iter().map(|&offset| offset - first).collect());
364 (rebased, referenced)
365}
366
367pub fn from_arrow_bytes<T: ByteArrayType>(
370 value: &GenericByteArray<T>,
371 nullable: bool,
372) -> VortexResult<ArrayRef>
373where
374 <T as ByteArrayType>::Offset: IntegerPType,
375{
376 let dtype = match T::DATA_TYPE {
377 DataType::Binary | DataType::LargeBinary => DType::Binary(nullable.into()),
378 DataType::Utf8 | DataType::LargeUtf8 => DType::Utf8(nullable.into()),
379 dt => vortex_panic!("Invalid data type for ByteArray: {dt}"),
380 };
381 let (offsets, referenced) = trim_offsets(value.offsets());
382 let bytes = value
383 .values()
384 .slice_with_length(referenced.start, referenced.len());
385 Ok(unsafe {
387 VarBinArray::new_unchecked(
388 offsets.into_array(),
389 ByteBuffer::from_arrow_buffer(bytes, Alignment::of::<u8>()),
390 dtype,
391 nulls(value.nulls(), nullable)?,
392 )
393 }
394 .into_array())
395}
396
397impl<T: ByteArrayType> FromArrowArray<&GenericByteArray<T>> for ArrayRef
398where
399 <T as ByteArrayType>::Offset: IntegerPType,
400{
401 fn from_arrow(value: &GenericByteArray<T>, nullable: bool) -> VortexResult<Self> {
402 from_arrow_bytes(value, nullable)
403 }
404}
405
406pub fn from_arrow_byte_view<T: ByteViewType>(
408 value: &GenericByteViewArray<T>,
409 nullable: bool,
410) -> VortexResult<ArrayRef> {
411 let dtype = match T::DATA_TYPE {
412 DataType::BinaryView => DType::Binary(nullable.into()),
413 DataType::Utf8View => DType::Utf8(nullable.into()),
414 dt => vortex_panic!("Invalid data type for ByteViewArray: {dt}"),
415 };
416
417 let views_buffer = Buffer::from_byte_buffer(
418 Buffer::from_arrow_scalar_buffer(value.views().clone()).into_byte_buffer(),
419 );
420
421 Ok(unsafe {
424 VarBinViewArray::new_unchecked(
425 views_buffer,
426 Arc::from(
427 value
428 .data_buffers()
429 .iter()
430 .map(|b| ByteBuffer::from_arrow_buffer(b.clone(), Alignment::of::<u8>()))
431 .collect::<Vec<_>>(),
432 ),
433 dtype,
434 nulls(value.nulls(), nullable)?,
435 )
436 .into_array()
437 })
438}
439
440impl<T: ByteViewType> FromArrowArray<&GenericByteViewArray<T>> for ArrayRef {
441 fn from_arrow(value: &GenericByteViewArray<T>, nullable: bool) -> VortexResult<Self> {
442 from_arrow_byte_view(value, nullable)
443 }
444}
445
446pub fn from_arrow_boolean(value: &ArrowBooleanArray, nullable: bool) -> VortexResult<ArrayRef> {
448 Ok(BoolArray::new(
449 value.values().clone().into(),
450 nulls(value.nulls(), nullable)?,
451 )
452 .into_array())
453}
454
455impl FromArrowArray<&ArrowBooleanArray> for ArrayRef {
456 fn from_arrow(value: &ArrowBooleanArray, nullable: bool) -> VortexResult<Self> {
457 from_arrow_boolean(value, nullable)
458 }
459}
460
461pub(crate) fn remove_nulls(data: arrow_data::ArrayData) -> VortexResult<arrow_data::ArrayData> {
463 if data.null_count() == 0 {
464 return Ok(data);
466 }
467
468 let children = match data.data_type() {
469 DataType::Struct(fields) => Some(
470 fields
471 .iter()
472 .zip(data.child_data().iter())
473 .map(|(field, child_data)| {
474 if field.is_nullable() {
475 Ok(child_data.clone())
476 } else {
477 remove_nulls(child_data.clone())
478 }
479 })
480 .collect::<VortexResult<Vec<_>>>()?,
481 ),
482 DataType::List(f)
483 | DataType::LargeList(f)
484 | DataType::ListView(f)
485 | DataType::LargeListView(f)
486 | DataType::FixedSizeList(f, _)
487 if !f.is_nullable() =>
488 {
489 vortex_ensure_eq!(
491 data.child_data().len(),
492 1,
493 "List types should have one child"
494 );
495 Some(vec![remove_nulls(data.child_data()[0].clone())?])
496 }
497 _ => None,
498 };
499
500 let mut builder = data.into_builder().nulls(None);
501 if let Some(children) = children {
502 builder = builder.child_data(children);
503 }
504 builder
505 .build()
506 .map_err(|e| vortex_err!("Failed to reconstruct Arrow array without nulls: {e}"))
507}
508
509pub fn from_arrow_struct(value: &ArrowStructArray, nullable: bool) -> VortexResult<ArrayRef> {
511 Ok(StructArray::try_new(
512 value.column_names().iter().copied().collect(),
513 value
514 .columns()
515 .iter()
516 .zip(value.fields())
517 .map(|(c, field)| {
518 if c.null_count() > 0 && !field.is_nullable() {
521 let stripped = make_array(remove_nulls(c.into_data())?);
522 from_arrow_dyn(stripped.as_ref(), false)
523 } else {
524 from_arrow_dyn(c.as_ref(), field.is_nullable())
525 }
526 })
527 .collect::<VortexResult<Vec<_>>>()?,
528 value.len(),
529 nulls(value.nulls(), nullable)?,
530 )?
531 .into_array())
532}
533
534impl FromArrowArray<&ArrowStructArray> for ArrayRef {
535 fn from_arrow(value: &ArrowStructArray, nullable: bool) -> VortexResult<Self> {
536 from_arrow_struct(value, nullable)
537 }
538}
539
540pub fn from_arrow_list<O: IntegerPType + OffsetSizeTrait>(
542 value: &GenericListArray<O>,
543 nullable: bool,
544) -> VortexResult<ArrayRef> {
545 let elements_are_nullable = match value.data_type() {
547 DataType::List(field) => field.is_nullable(),
548 DataType::LargeList(field) => field.is_nullable(),
549 dt => vortex_panic!("Invalid data type for ListArray: {dt}"),
550 };
551
552 let (offsets, referenced) = trim_offsets(value.offsets());
553 let elements = from_arrow_dyn(
554 value
555 .values()
556 .slice(referenced.start, referenced.len())
557 .as_ref(),
558 elements_are_nullable,
559 )?;
560 let nulls = nulls(value.nulls(), nullable)?;
561
562 Ok(ListArray::try_new(elements, offsets.into_array(), nulls)?.into_array())
564}
565
566impl<O: IntegerPType + OffsetSizeTrait> FromArrowArray<&GenericListArray<O>> for ArrayRef {
567 fn from_arrow(value: &GenericListArray<O>, nullable: bool) -> VortexResult<Self> {
568 from_arrow_list(value, nullable)
569 }
570}
571
572pub fn from_arrow_list_view<O: OffsetSizeTrait + NativePType>(
574 array: &GenericListViewArray<O>,
575 nullable: bool,
576) -> VortexResult<ArrayRef> {
577 let elements_are_nullable = match array.data_type() {
579 DataType::ListView(field) => field.is_nullable(),
580 DataType::LargeListView(field) => field.is_nullable(),
581 dt => vortex_panic!("Invalid data type for ListViewArray: {dt}"),
582 };
583
584 let elements = from_arrow_dyn(array.values().as_ref(), elements_are_nullable)?;
585
586 let offsets = array.offsets().clone().into_array();
588 let sizes = array.sizes().clone().into_array();
589 let nulls = nulls(array.nulls(), nullable)?;
590
591 Ok(ListViewArray::try_new(elements, offsets, sizes, nulls)?.into_array())
592}
593
594impl<O: OffsetSizeTrait + NativePType> FromArrowArray<&GenericListViewArray<O>> for ArrayRef {
595 fn from_arrow(array: &GenericListViewArray<O>, nullable: bool) -> VortexResult<Self> {
596 from_arrow_list_view(array, nullable)
597 }
598}
599
600pub fn from_arrow_fixed_size_list(
602 array: &ArrowFixedSizeListArray,
603 nullable: bool,
604) -> VortexResult<ArrayRef> {
605 let DataType::FixedSizeList(field, list_size) = array.data_type() else {
606 vortex_panic!("Invalid data type for ListArray: {}", array.data_type());
607 };
608
609 Ok(FixedSizeListArray::try_new(
610 from_arrow_dyn(array.values().as_ref(), field.is_nullable())?,
611 *list_size as u32,
612 nulls(array.nulls(), nullable)?,
613 array.len(),
614 )?
615 .into_array())
616}
617
618impl FromArrowArray<&ArrowFixedSizeListArray> for ArrayRef {
619 fn from_arrow(array: &ArrowFixedSizeListArray, nullable: bool) -> VortexResult<Self> {
620 from_arrow_fixed_size_list(array, nullable)
621 }
622}
623
624pub(crate) fn map_from_arrow_parts(
625 entries: ArrayRef,
626 offsets: &OffsetBuffer<i32>,
627 arrow_nulls: Option<&NullBuffer>,
628 keys_sorted: bool,
629 nullable: bool,
630) -> VortexResult<ArrayRef> {
631 let DType::Struct(struct_dtype, Nullability::NonNullable) = entries.dtype() else {
632 vortex_bail!(
633 "Arrow map entries must import as non-nullable struct, got {}",
634 entries.dtype()
635 );
636 };
637 vortex_ensure_eq!(
638 struct_dtype.nfields(),
639 2,
640 "Arrow map entries struct must contain exactly two fields"
641 );
642
643 let key_dtype = struct_dtype
644 .field_by_index(0)
645 .ok_or_else(|| vortex_err!("Arrow map entries struct missing key field"))?;
646 vortex_ensure!(
647 !key_dtype.is_nullable(),
648 "Arrow map key field must be non-nullable"
649 );
650 let value_dtype = struct_dtype
651 .field_by_index(1)
652 .ok_or_else(|| vortex_err!("Arrow map entries struct missing value field"))?;
653 let map_dtype = MapDType::try_new(key_dtype, value_dtype, keys_sorted)?;
654 let entries_struct = entries.as_::<Struct>();
655 let entries = StructArray::try_new(
656 FieldNames::from(["key", "value"]),
657 vec![
658 entries_struct.unmasked_field(0).clone(),
659 entries_struct.unmasked_field(1).clone(),
660 ],
661 entries.len(),
662 Validity::NonNullable,
663 )?
664 .into_array();
665
666 let len = offsets.len() - 1;
667 let sizes = Buffer::<i32>::from_iter(offsets.windows(2).map(|window| window[1] - window[0]))
668 .into_array();
669 let offsets = offsets.inner().slice(0, len).into_array();
670 let validity = nulls(arrow_nulls, nullable)?;
671 let entries = ListViewArray::try_new(entries, offsets, sizes, validity)?;
672 let entries = unsafe { entries.with_zero_copy_to_list(true) };
676
677 Ok(MapArray::try_new(map_dtype, entries)?.into_array())
678}
679
680pub fn from_arrow_map(array: &ArrowMapArray, nullable: bool) -> VortexResult<ArrayRef> {
682 let DataType::Map(_, keys_sorted) = array.data_type() else {
683 vortex_panic!("Invalid data type for MapArray: {}", array.data_type());
684 };
685 let entries = from_arrow_struct(array.entries(), false)?;
686 map_from_arrow_parts(
687 entries,
688 array.offsets(),
689 array.nulls(),
690 *keys_sorted,
691 nullable,
692 )
693}
694
695impl FromArrowArray<&ArrowMapArray> for ArrayRef {
696 fn from_arrow(array: &ArrowMapArray, nullable: bool) -> VortexResult<Self> {
697 from_arrow_map(array, nullable)
698 }
699}
700
701pub fn from_arrow_null(value: &ArrowNullArray, nullable: bool) -> VortexResult<ArrayRef> {
703 vortex_ensure!(
704 nullable,
705 "Cannot convert an Arrow NullArray into a non-nullable Vortex array"
706 );
707 Ok(NullArray::new(value.len()).into_array())
708}
709
710impl FromArrowArray<&ArrowNullArray> for ArrayRef {
711 fn from_arrow(value: &ArrowNullArray, nullable: bool) -> VortexResult<Self> {
712 from_arrow_null(value, nullable)
713 }
714}
715
716pub fn from_arrow_dictionary<K: ArrowDictionaryKeyType>(
718 array: &DictionaryArray<K>,
719 nullable: bool,
720) -> VortexResult<DictArray> {
721 let keys = AnyDictionaryArray::keys(array);
722 let keys = from_arrow_dyn(keys, keys.is_nullable())?;
723 let values = from_arrow_dyn(array.values().as_ref(), nullable)?;
724 Ok(unsafe { DictArray::new_unchecked(keys, values) })
726}
727
728impl<K: ArrowDictionaryKeyType> FromArrowArray<&DictionaryArray<K>> for DictArray {
729 fn from_arrow(array: &DictionaryArray<K>, nullable: bool) -> VortexResult<Self> {
730 from_arrow_dictionary(array, nullable)
731 }
732}
733
734pub fn nulls(nulls: Option<&NullBuffer>, nullable: bool) -> VortexResult<Validity> {
741 if nullable {
742 Ok(nulls
743 .map(|nulls| {
744 if nulls.null_count() == nulls.len() {
745 Validity::AllInvalid
746 } else {
747 Validity::from(BitBuffer::from(nulls.inner().clone()))
748 }
749 })
750 .unwrap_or(Validity::AllValid))
751 } else {
752 let null_count = nulls.map(NullBuffer::null_count).unwrap_or(0);
753 vortex_ensure_eq!(
754 null_count,
755 0,
756 "Cannot convert an Arrow array containing nulls into a non-nullable Vortex array"
757 );
758 Ok(Validity::NonNullable)
759 }
760}
761
762pub fn from_arrow_dyn(array: &dyn ArrowArray, nullable: bool) -> VortexResult<ArrayRef> {
765 match array.data_type() {
766 DataType::Boolean => from_arrow_boolean(array.as_boolean(), nullable),
767 DataType::UInt8 => from_arrow_primitive(array.as_primitive::<UInt8Type>(), nullable),
768 DataType::UInt16 => from_arrow_primitive(array.as_primitive::<UInt16Type>(), nullable),
769 DataType::UInt32 => from_arrow_primitive(array.as_primitive::<UInt32Type>(), nullable),
770 DataType::UInt64 => from_arrow_primitive(array.as_primitive::<UInt64Type>(), nullable),
771 DataType::Int8 => from_arrow_primitive(array.as_primitive::<Int8Type>(), nullable),
772 DataType::Int16 => from_arrow_primitive(array.as_primitive::<Int16Type>(), nullable),
773 DataType::Int32 => from_arrow_primitive(array.as_primitive::<Int32Type>(), nullable),
774 DataType::Int64 => from_arrow_primitive(array.as_primitive::<Int64Type>(), nullable),
775 DataType::Float16 => from_arrow_primitive(array.as_primitive::<Float16Type>(), nullable),
776 DataType::Float32 => from_arrow_primitive(array.as_primitive::<Float32Type>(), nullable),
777 DataType::Float64 => from_arrow_primitive(array.as_primitive::<Float64Type>(), nullable),
778 DataType::Utf8 => from_arrow_bytes(array.as_string::<i32>(), nullable),
779 DataType::LargeUtf8 => from_arrow_bytes(array.as_string::<i64>(), nullable),
780 DataType::Binary => from_arrow_bytes(array.as_binary::<i32>(), nullable),
781 DataType::LargeBinary => from_arrow_bytes(array.as_binary::<i64>(), nullable),
782 DataType::BinaryView => from_arrow_byte_view(array.as_binary_view(), nullable),
783 DataType::Utf8View => from_arrow_byte_view(array.as_string_view(), nullable),
784 DataType::Struct(_) => from_arrow_struct(array.as_struct(), nullable),
785 DataType::List(_) => from_arrow_list(array.as_list::<i32>(), nullable),
786 DataType::LargeList(_) => from_arrow_list(array.as_list::<i64>(), nullable),
787 DataType::ListView(_) => from_arrow_list_view(array.as_list_view::<i32>(), nullable),
788 DataType::LargeListView(_) => from_arrow_list_view(array.as_list_view::<i64>(), nullable),
789 DataType::FixedSizeList(..) => {
790 from_arrow_fixed_size_list(array.as_fixed_size_list(), nullable)
791 }
792 DataType::Map(..) => from_arrow_map(array.as_map(), nullable),
793 DataType::Null => from_arrow_null(as_null_array(array), nullable),
794 DataType::Timestamp(u, _) => match u {
795 ArrowTimeUnit::Second => {
796 from_arrow_temporal(array.as_primitive::<TimestampSecondType>(), nullable)
797 }
798 ArrowTimeUnit::Millisecond => {
799 from_arrow_temporal(array.as_primitive::<TimestampMillisecondType>(), nullable)
800 }
801 ArrowTimeUnit::Microsecond => {
802 from_arrow_temporal(array.as_primitive::<TimestampMicrosecondType>(), nullable)
803 }
804 ArrowTimeUnit::Nanosecond => {
805 from_arrow_temporal(array.as_primitive::<TimestampNanosecondType>(), nullable)
806 }
807 },
808 DataType::Date32 => from_arrow_temporal(array.as_primitive::<Date32Type>(), nullable),
809 DataType::Date64 => from_arrow_temporal(array.as_primitive::<Date64Type>(), nullable),
810 DataType::Time32(u) => match u {
811 ArrowTimeUnit::Second => {
812 from_arrow_temporal(array.as_primitive::<Time32SecondType>(), nullable)
813 }
814 ArrowTimeUnit::Millisecond => {
815 from_arrow_temporal(array.as_primitive::<Time32MillisecondType>(), nullable)
816 }
817 ArrowTimeUnit::Microsecond | ArrowTimeUnit::Nanosecond => unreachable!(),
818 },
819 DataType::Time64(u) => match u {
820 ArrowTimeUnit::Microsecond => {
821 from_arrow_temporal(array.as_primitive::<Time64MicrosecondType>(), nullable)
822 }
823 ArrowTimeUnit::Nanosecond => {
824 from_arrow_temporal(array.as_primitive::<Time64NanosecondType>(), nullable)
825 }
826 ArrowTimeUnit::Second | ArrowTimeUnit::Millisecond => unreachable!(),
827 },
828 DataType::Decimal32(..) => {
829 from_arrow_decimal32(array.as_primitive::<Decimal32Type>(), nullable)
830 }
831 DataType::Decimal64(..) => {
832 from_arrow_decimal64(array.as_primitive::<Decimal64Type>(), nullable)
833 }
834 DataType::Decimal128(..) => {
835 from_arrow_decimal128(array.as_primitive::<Decimal128Type>(), nullable)
836 }
837 DataType::Decimal256(..) => {
838 from_arrow_decimal256(array.as_primitive::<Decimal256Type>(), nullable)
839 }
840 DataType::Dictionary(key_type, _) => match key_type.as_ref() {
841 DataType::Int8 => Ok(from_arrow_dictionary(
842 array.as_dictionary::<Int8Type>(),
843 nullable,
844 )?
845 .into_array()),
846 DataType::Int16 => Ok(from_arrow_dictionary(
847 array.as_dictionary::<Int16Type>(),
848 nullable,
849 )?
850 .into_array()),
851 DataType::Int32 => Ok(from_arrow_dictionary(
852 array.as_dictionary::<Int32Type>(),
853 nullable,
854 )?
855 .into_array()),
856 DataType::Int64 => Ok(from_arrow_dictionary(
857 array.as_dictionary::<Int64Type>(),
858 nullable,
859 )?
860 .into_array()),
861 DataType::UInt8 => Ok(from_arrow_dictionary(
862 array.as_dictionary::<UInt8Type>(),
863 nullable,
864 )?
865 .into_array()),
866 DataType::UInt16 => Ok(from_arrow_dictionary(
867 array.as_dictionary::<UInt16Type>(),
868 nullable,
869 )?
870 .into_array()),
871 DataType::UInt32 => Ok(from_arrow_dictionary(
872 array.as_dictionary::<UInt32Type>(),
873 nullable,
874 )?
875 .into_array()),
876 DataType::UInt64 => Ok(from_arrow_dictionary(
877 array.as_dictionary::<UInt64Type>(),
878 nullable,
879 )?
880 .into_array()),
881 key_dt => vortex_bail!("Unsupported dictionary key type: {key_dt}"),
882 },
883 dt => vortex_bail!("Array encoding not implemented for Arrow data type {dt}"),
884 }
885}
886
887impl FromArrowArray<&dyn ArrowArray> for ArrayRef {
888 fn from_arrow(array: &dyn ArrowArray, nullable: bool) -> VortexResult<Self> {
889 from_arrow_dyn(array, nullable)
890 }
891}
892
893pub fn from_arrow_batch(batch: &RecordBatch, nullable: bool) -> VortexResult<ArrayRef> {
895 from_arrow_struct(&arrow_array::StructArray::from(batch.clone()), nullable)
896}
897
898impl FromArrowArray<RecordBatch> for ArrayRef {
899 fn from_arrow(array: RecordBatch, nullable: bool) -> VortexResult<Self> {
900 from_arrow_batch(&array, nullable)
901 }
902}
903
904impl FromArrowArray<&RecordBatch> for ArrayRef {
905 fn from_arrow(array: &RecordBatch, nullable: bool) -> VortexResult<Self> {
906 from_arrow_batch(array, nullable)
907 }
908}
909
910#[cfg(test)]
911mod tests {
912 use std::ops::Range;
913 use std::sync::Arc;
914
915 use arrow_array::Array as ArrowArray;
916 use arrow_array::BinaryArray;
917 use arrow_array::BooleanArray;
918 use arrow_array::Date32Array;
919 use arrow_array::Date64Array;
920 use arrow_array::FixedSizeListArray as ArrowFixedSizeListArray;
921 use arrow_array::Float32Array;
922 use arrow_array::Float64Array;
923 use arrow_array::GenericListViewArray;
924 use arrow_array::Int8Array;
925 use arrow_array::Int16Array;
926 use arrow_array::Int32Array;
927 use arrow_array::Int64Array;
928 use arrow_array::LargeBinaryArray;
929 use arrow_array::LargeStringArray;
930 use arrow_array::NullArray;
931 use arrow_array::RecordBatch;
932 use arrow_array::StringArray;
933 use arrow_array::StructArray;
934 use arrow_array::Time32MillisecondArray;
935 use arrow_array::Time32SecondArray;
936 use arrow_array::Time64MicrosecondArray;
937 use arrow_array::Time64NanosecondArray;
938 use arrow_array::TimestampMicrosecondArray;
939 use arrow_array::TimestampMillisecondArray;
940 use arrow_array::TimestampNanosecondArray;
941 use arrow_array::TimestampSecondArray;
942 use arrow_array::UInt8Array;
943 use arrow_array::UInt16Array;
944 use arrow_array::UInt32Array;
945 use arrow_array::UInt64Array;
946 use arrow_array::builder::BinaryViewBuilder;
947 use arrow_array::builder::Decimal128Builder;
948 use arrow_array::builder::Decimal256Builder;
949 use arrow_array::builder::Int32Builder;
950 use arrow_array::builder::Int64Builder;
951 use arrow_array::builder::LargeListBuilder;
952 use arrow_array::builder::ListBuilder;
953 use arrow_array::builder::MapBuilder as ArrowMapBuilder;
954 use arrow_array::builder::StringBuilder;
955 use arrow_array::builder::StringViewBuilder;
956 use arrow_array::new_null_array;
957 use arrow_array::types::ArrowPrimitiveType;
958 use arrow_array::types::Float16Type;
959 use arrow_buffer::BooleanBuffer;
960 use arrow_buffer::Buffer as ArrowBuffer;
961 use arrow_buffer::OffsetBuffer;
962 use arrow_buffer::ScalarBuffer;
963 use arrow_schema::DataType;
964 use arrow_schema::Field;
965 use arrow_schema::Fields;
966 use arrow_schema::Schema;
967 use rstest::rstest;
968 use vortex_array::ArrayRef;
969 use vortex_array::IntoArray;
970 use vortex_array::VortexSessionExecute as _;
971 use vortex_array::array_session;
972 use vortex_array::arrays::BoolArray;
973 use vortex_array::arrays::Decimal;
974 use vortex_array::arrays::DecimalArray;
975 use vortex_array::arrays::FixedSizeList;
976 use vortex_array::arrays::FixedSizeListArray as VortexFixedSizeListArray;
977 use vortex_array::arrays::List;
978 use vortex_array::arrays::ListArray;
979 use vortex_array::arrays::ListView;
980 use vortex_array::arrays::ListViewArray;
981 use vortex_array::arrays::NullArray as VortexNullArray;
982 use vortex_array::arrays::Primitive;
983 use vortex_array::arrays::PrimitiveArray;
984 use vortex_array::arrays::Struct;
985 use vortex_array::arrays::StructArray as VortexStructArray;
986 use vortex_array::arrays::TemporalArray;
987 use vortex_array::arrays::VarBinArray;
988 use vortex_array::arrays::VarBinView;
989 use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
990 use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
991 use vortex_array::arrays::list::ListArraySlotsExt;
992 use vortex_array::arrays::listview::ListViewArraySlotsExt;
993 use vortex_array::arrays::struct_::StructArrayExt;
994 use vortex_array::assert_arrays_eq;
995 use vortex_array::dtype::DType;
996 use vortex_array::dtype::DecimalDType;
997 use vortex_array::dtype::Nullability;
998 use vortex_array::dtype::PType;
999 use vortex_array::dtype::i256;
1000 use vortex_array::extension::datetime::TimeUnit;
1001 use vortex_array::extension::datetime::Timestamp;
1002 use vortex_array::scalar::Scalar;
1003 use vortex_array::validity::Validity;
1004 use vortex_buffer::buffer;
1005 use vortex_error::VortexResult;
1006
1007 use crate::FromArrowArray as _;
1008 use crate::IntoVortexArray as _;
1009 use crate::convert::from_arrow_batch;
1010 use crate::convert::from_arrow_dyn;
1011
1012 #[rstest]
1013 #[case::i8(
1014 Arc::new(Int8Array::from(vec![Some(1), None, Some(3), Some(4)])),
1015 Arc::new(Int8Array::from(vec![1, 2, 3, 4])),
1016 PType::I8,
1017 )]
1018 #[case::i16(
1019 Arc::new(Int16Array::from(vec![Some(100), None, Some(300), Some(400)])),
1020 Arc::new(Int16Array::from(vec![100, 200, 300, 400])),
1021 PType::I16,
1022 )]
1023 #[case::i32(
1024 Arc::new(Int32Array::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1025 Arc::new(Int32Array::from(vec![1000, 2000, 3000, 4000])),
1026 PType::I32,
1027 )]
1028 #[case::i64(
1029 Arc::new(Int64Array::from(vec![Some(10000), None, Some(30000), Some(40000)])),
1030 Arc::new(Int64Array::from(vec![10000_i64, 20000, 30000, 40000])),
1031 PType::I64,
1032 )]
1033 #[case::u8(
1034 Arc::new(UInt8Array::from(vec![Some(1), None, Some(3), Some(4)])),
1035 Arc::new(UInt8Array::from(vec![1_u8, 2, 3, 4])),
1036 PType::U8,
1037 )]
1038 #[case::u16(
1039 Arc::new(UInt16Array::from(vec![Some(100), None, Some(300), Some(400)])),
1040 Arc::new(UInt16Array::from(vec![100_u16, 200, 300, 400])),
1041 PType::U16,
1042 )]
1043 #[case::u32(
1044 Arc::new(UInt32Array::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1045 Arc::new(UInt32Array::from(vec![1000_u32, 2000, 3000, 4000])),
1046 PType::U32,
1047 )]
1048 #[case::u64(
1049 Arc::new(UInt64Array::from(vec![Some(10000), None, Some(30000), Some(40000)])),
1050 Arc::new(UInt64Array::from(vec![10000_u64, 20000, 30000, 40000])),
1051 PType::U64,
1052 )]
1053 #[case::f32(
1054 Arc::new(Float32Array::from(vec![Some(1.5), None, Some(3.5), Some(4.5)])),
1055 Arc::new(Float32Array::from(vec![1.5_f32, 2.5, 3.5, 4.5])),
1056 PType::F32,
1057 )]
1058 #[case::f64(
1059 Arc::new(Float64Array::from(vec![Some(1.5), None, Some(3.5), Some(4.5)])),
1060 Arc::new(Float64Array::from(vec![1.5_f64, 2.5, 3.5, 4.5])),
1061 PType::F64,
1062 )]
1063 fn test_primitive_array_conversion(
1064 #[case] nullable: Arc<dyn ArrowArray>,
1065 #[case] non_nullable: Arc<dyn ArrowArray>,
1066 #[case] expected_ptype: PType,
1067 ) {
1068 let v_null = ArrayRef::from_arrow(nullable.as_ref(), true).unwrap();
1069 let v_non_null = ArrayRef::from_arrow(non_nullable.as_ref(), false).unwrap();
1070 assert_eq!(v_null.len(), 4);
1071 assert_eq!(v_non_null.len(), 4);
1072 assert_eq!(v_null.as_::<Primitive>().ptype(), expected_ptype);
1073 assert_eq!(v_non_null.as_::<Primitive>().ptype(), expected_ptype);
1074 let (expected_nullable, expected_non_nullable) = match expected_ptype {
1075 PType::I8 => (
1076 PrimitiveArray::from_option_iter([Some(1i8), None, Some(3), Some(4)]),
1077 PrimitiveArray::from_iter([1i8, 2, 3, 4]),
1078 ),
1079 PType::I16 => (
1080 PrimitiveArray::from_option_iter([Some(100i16), None, Some(300), Some(400)]),
1081 PrimitiveArray::from_iter([100i16, 200, 300, 400]),
1082 ),
1083 PType::I32 => (
1084 PrimitiveArray::from_option_iter([Some(1000i32), None, Some(3000), Some(4000)]),
1085 PrimitiveArray::from_iter([1000i32, 2000, 3000, 4000]),
1086 ),
1087 PType::I64 => (
1088 PrimitiveArray::from_option_iter([Some(10000i64), None, Some(30000), Some(40000)]),
1089 PrimitiveArray::from_iter([10000i64, 20000, 30000, 40000]),
1090 ),
1091 PType::U8 => (
1092 PrimitiveArray::from_option_iter([Some(1u8), None, Some(3), Some(4)]),
1093 PrimitiveArray::from_iter([1u8, 2, 3, 4]),
1094 ),
1095 PType::U16 => (
1096 PrimitiveArray::from_option_iter([Some(100u16), None, Some(300), Some(400)]),
1097 PrimitiveArray::from_iter([100u16, 200, 300, 400]),
1098 ),
1099 PType::U32 => (
1100 PrimitiveArray::from_option_iter([Some(1000u32), None, Some(3000), Some(4000)]),
1101 PrimitiveArray::from_iter([1000u32, 2000, 3000, 4000]),
1102 ),
1103 PType::U64 => (
1104 PrimitiveArray::from_option_iter([Some(10000u64), None, Some(30000), Some(40000)]),
1105 PrimitiveArray::from_iter([10000u64, 20000, 30000, 40000]),
1106 ),
1107 PType::F32 => (
1108 PrimitiveArray::from_option_iter([Some(1.5f32), None, Some(3.5), Some(4.5)]),
1109 PrimitiveArray::from_iter([1.5f32, 2.5, 3.5, 4.5]),
1110 ),
1111 PType::F64 => (
1112 PrimitiveArray::from_option_iter([Some(1.5f64), None, Some(3.5), Some(4.5)]),
1113 PrimitiveArray::from_iter([1.5f64, 2.5, 3.5, 4.5]),
1114 ),
1115 _ => unreachable!(),
1116 };
1117 assert_arrays_eq!(
1118 v_null,
1119 expected_nullable,
1120 &mut array_session().create_execution_ctx()
1121 );
1122 assert_arrays_eq!(
1123 v_non_null,
1124 expected_non_nullable,
1125 &mut array_session().create_execution_ctx()
1126 );
1127 }
1128
1129 #[test]
1130 fn test_float16_array_conversion() {
1131 let values = vec![
1132 Some(<Float16Type as ArrowPrimitiveType>::Native::from_f32(1.5)),
1133 None,
1134 Some(<Float16Type as ArrowPrimitiveType>::Native::from_f32(3.5)),
1135 ];
1136 let arrow_array = arrow_array::PrimitiveArray::<Float16Type>::from(values);
1137 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1138
1139 let non_null_values = vec![
1140 <Float16Type as ArrowPrimitiveType>::Native::from_f32(1.5),
1141 <Float16Type as ArrowPrimitiveType>::Native::from_f32(2.5),
1142 ];
1143 let arrow_array_non_null =
1144 arrow_array::PrimitiveArray::<Float16Type>::from(non_null_values);
1145 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1146
1147 assert_eq!(vortex_array.len(), 3);
1148 assert_eq!(vortex_array_non_null.len(), 2);
1149
1150 let primitive_array = vortex_array.as_::<Primitive>();
1152 assert_eq!(primitive_array.ptype(), PType::F16);
1153
1154 let primitive_array_non_null = vortex_array_non_null.as_::<Primitive>();
1155 assert_eq!(primitive_array_non_null.ptype(), PType::F16);
1156 assert_arrays_eq!(
1157 vortex_array,
1158 PrimitiveArray::from_option_iter([
1159 Some(<Float16Type as ArrowPrimitiveType>::Native::from_f32(1.5)),
1160 None,
1161 Some(<Float16Type as ArrowPrimitiveType>::Native::from_f32(3.5))
1162 ]),
1163 &mut array_session().create_execution_ctx()
1164 );
1165 assert_arrays_eq!(
1166 vortex_array_non_null,
1167 PrimitiveArray::from_iter([
1168 <Float16Type as ArrowPrimitiveType>::Native::from_f32(1.5),
1169 <Float16Type as ArrowPrimitiveType>::Native::from_f32(2.5)
1170 ]),
1171 &mut array_session().create_execution_ctx()
1172 );
1173 }
1174
1175 #[test]
1177 fn test_decimal128_array_conversion() {
1178 let mut builder = Decimal128Builder::with_capacity(4);
1179 builder.append_value(12345);
1180 builder.append_null();
1181 builder.append_value(67890);
1182 builder.append_value(11111);
1183 let decimal_array = builder.finish().with_precision_and_scale(10, 2).unwrap();
1184
1185 let vortex_array = ArrayRef::from_arrow(&decimal_array, true).unwrap();
1186 assert_eq!(vortex_array.len(), 4);
1187
1188 let mut builder_non_null = Decimal128Builder::with_capacity(3);
1189 builder_non_null.append_value(12345);
1190 builder_non_null.append_value(67890);
1191 builder_non_null.append_value(11111);
1192 let decimal_array_non_null = builder_non_null
1193 .finish()
1194 .with_precision_and_scale(10, 2)
1195 .unwrap();
1196
1197 let vortex_array_non_null = ArrayRef::from_arrow(&decimal_array_non_null, false).unwrap();
1198 assert_eq!(vortex_array_non_null.len(), 3);
1199
1200 let decimal_vortex_array = vortex_array.as_::<Decimal>();
1202 assert_eq!(decimal_vortex_array.decimal_dtype().precision(), 10);
1203 assert_eq!(decimal_vortex_array.decimal_dtype().scale(), 2);
1204
1205 let decimal_vortex_array_non_null = vortex_array_non_null.as_::<Decimal>();
1206 assert_eq!(
1207 decimal_vortex_array_non_null.decimal_dtype().precision(),
1208 10
1209 );
1210 assert_eq!(decimal_vortex_array_non_null.decimal_dtype().scale(), 2);
1211 assert_arrays_eq!(
1212 vortex_array,
1213 DecimalArray::from_option_iter(
1214 [Some(12345i128), None, Some(67890i128), Some(11111i128)],
1215 DecimalDType::new(10, 2)
1216 ),
1217 &mut array_session().create_execution_ctx()
1218 );
1219 assert_arrays_eq!(
1220 vortex_array_non_null,
1221 DecimalArray::from_iter([12345i128, 67890i128, 11111i128], DecimalDType::new(10, 2)),
1222 &mut array_session().create_execution_ctx()
1223 );
1224 }
1225
1226 #[test]
1227 fn test_decimal256_array_conversion() {
1228 let mut builder = Decimal256Builder::with_capacity(4);
1229 builder.append_value(arrow_buffer::i256::from_i128(12345));
1230 builder.append_null();
1231 builder.append_value(arrow_buffer::i256::from_i128(67890));
1232 builder.append_value(arrow_buffer::i256::from_i128(11111));
1233 let decimal_array = builder.finish().with_precision_and_scale(38, 10).unwrap();
1234
1235 let vortex_array = ArrayRef::from_arrow(&decimal_array, true).unwrap();
1236 assert_eq!(vortex_array.len(), 4);
1237
1238 let mut builder_non_null = Decimal256Builder::with_capacity(3);
1239 builder_non_null.append_value(arrow_buffer::i256::from_i128(12345));
1240 builder_non_null.append_value(arrow_buffer::i256::from_i128(67890));
1241 builder_non_null.append_value(arrow_buffer::i256::from_i128(11111));
1242 let decimal_array_non_null = builder_non_null
1243 .finish()
1244 .with_precision_and_scale(38, 10)
1245 .unwrap();
1246
1247 let vortex_array_non_null = ArrayRef::from_arrow(&decimal_array_non_null, false).unwrap();
1248 assert_eq!(vortex_array_non_null.len(), 3);
1249
1250 let decimal_vortex_array = vortex_array.as_::<Decimal>();
1252 assert_eq!(decimal_vortex_array.decimal_dtype().precision(), 38);
1253 assert_eq!(decimal_vortex_array.decimal_dtype().scale(), 10);
1254
1255 let decimal_vortex_array_non_null = vortex_array_non_null.as_::<Decimal>();
1256 assert_eq!(
1257 decimal_vortex_array_non_null.decimal_dtype().precision(),
1258 38
1259 );
1260 assert_eq!(decimal_vortex_array_non_null.decimal_dtype().scale(), 10);
1261 assert_arrays_eq!(
1262 vortex_array,
1263 DecimalArray::from_option_iter(
1264 [
1265 Some(i256::from_i128(12345)),
1266 None,
1267 Some(i256::from_i128(67890)),
1268 Some(i256::from_i128(11111))
1269 ],
1270 DecimalDType::new(38, 10)
1271 ),
1272 &mut array_session().create_execution_ctx()
1273 );
1274 assert_arrays_eq!(
1275 vortex_array_non_null,
1276 DecimalArray::from_iter(
1277 [
1278 i256::from_i128(12345),
1279 i256::from_i128(67890),
1280 i256::from_i128(11111)
1281 ],
1282 DecimalDType::new(38, 10)
1283 ),
1284 &mut array_session().create_execution_ctx()
1285 );
1286 }
1287
1288 #[rstest]
1290 #[case::timestamp_second(
1291 Arc::new(TimestampSecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1292 Arc::new(TimestampSecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
1293 )]
1294 #[case::timestamp_millisecond(
1295 Arc::new(TimestampMillisecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1296 Arc::new(TimestampMillisecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
1297 )]
1298 #[case::timestamp_microsecond(
1299 Arc::new(TimestampMicrosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1300 Arc::new(TimestampMicrosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
1301 )]
1302 #[case::timestamp_nanosecond(
1303 Arc::new(TimestampNanosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1304 Arc::new(TimestampNanosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
1305 )]
1306 #[case::time32_second(
1307 Arc::new(Time32SecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1308 Arc::new(Time32SecondArray::from(vec![1000_i32, 2000, 3000, 4000])),
1309 )]
1310 #[case::time32_millisecond(
1311 Arc::new(Time32MillisecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1312 Arc::new(Time32MillisecondArray::from(vec![1000_i32, 2000, 3000, 4000])),
1313 )]
1314 #[case::time64_microsecond(
1315 Arc::new(Time64MicrosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1316 Arc::new(Time64MicrosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
1317 )]
1318 #[case::time64_nanosecond(
1319 Arc::new(Time64NanosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1320 Arc::new(Time64NanosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
1321 )]
1322 #[case::date32(
1323 Arc::new(Date32Array::from(vec![Some(18000), None, Some(18002), Some(18003)])),
1324 Arc::new(Date32Array::from(vec![18000_i32, 18001, 18002, 18003])),
1325 )]
1326 #[case::date64(
1327 Arc::new(Date64Array::from(vec![Some(1555200000000), None, Some(1555286400000), Some(1555372800000)]
1328 )),
1329 Arc::new(Date64Array::from(vec![1555200000000_i64, 1555213600000, 1555286400000, 1555372800000]
1330 )),
1331 )]
1332 fn test_temporal_array_conversion(
1333 #[case] nullable: Arc<dyn ArrowArray>,
1334 #[case] non_nullable: Arc<dyn ArrowArray>,
1335 ) {
1336 let v_null = ArrayRef::from_arrow(nullable.as_ref(), true).unwrap();
1337 let v_non_null = ArrayRef::from_arrow(non_nullable.as_ref(), false).unwrap();
1338 assert_eq!(v_null.len(), 4);
1339 assert_eq!(v_non_null.len(), 4);
1340 let expected = |nullable: bool| {
1341 let i32_values = if nullable {
1342 PrimitiveArray::from_option_iter([Some(1000i32), None, Some(3000), Some(4000)])
1343 } else {
1344 PrimitiveArray::from_iter([1000i32, 2000, 3000, 4000])
1345 }
1346 .into_array();
1347 let i64_values = if nullable {
1348 PrimitiveArray::from_option_iter([Some(1000i64), None, Some(3000), Some(4000)])
1349 } else {
1350 PrimitiveArray::from_iter([1000i64, 2000, 3000, 4000])
1351 }
1352 .into_array();
1353 match non_nullable.data_type() {
1354 DataType::Timestamp(unit, _) => {
1355 let unit = match unit {
1356 arrow_schema::TimeUnit::Second => TimeUnit::Seconds,
1357 arrow_schema::TimeUnit::Millisecond => TimeUnit::Milliseconds,
1358 arrow_schema::TimeUnit::Microsecond => TimeUnit::Microseconds,
1359 arrow_schema::TimeUnit::Nanosecond => TimeUnit::Nanoseconds,
1360 };
1361 TemporalArray::new_timestamp(i64_values, unit, None).into_array()
1362 }
1363 DataType::Time32(arrow_schema::TimeUnit::Second) => {
1364 TemporalArray::new_time(i32_values, TimeUnit::Seconds).into_array()
1365 }
1366 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
1367 TemporalArray::new_time(i32_values, TimeUnit::Milliseconds).into_array()
1368 }
1369 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
1370 TemporalArray::new_time(i64_values, TimeUnit::Microseconds).into_array()
1371 }
1372 DataType::Time64(arrow_schema::TimeUnit::Nanosecond) => {
1373 TemporalArray::new_time(i64_values, TimeUnit::Nanoseconds).into_array()
1374 }
1375 DataType::Date32 => TemporalArray::new_date(
1376 if nullable {
1377 PrimitiveArray::from_option_iter([
1378 Some(18000i32),
1379 None,
1380 Some(18002),
1381 Some(18003),
1382 ])
1383 } else {
1384 PrimitiveArray::from_iter([18000i32, 18001, 18002, 18003])
1385 }
1386 .into_array(),
1387 TimeUnit::Days,
1388 )
1389 .into_array(),
1390 DataType::Date64 => TemporalArray::new_date(
1391 if nullable {
1392 PrimitiveArray::from_option_iter([
1393 Some(1555200000000i64),
1394 None,
1395 Some(1555286400000),
1396 Some(1555372800000),
1397 ])
1398 } else {
1399 PrimitiveArray::from_iter([
1400 1555200000000i64,
1401 1555213600000,
1402 1555286400000,
1403 1555372800000,
1404 ])
1405 }
1406 .into_array(),
1407 TimeUnit::Milliseconds,
1408 )
1409 .into_array(),
1410 _ => unreachable!(),
1411 }
1412 };
1413 assert_arrays_eq!(
1414 v_null,
1415 expected(true),
1416 &mut array_session().create_execution_ctx()
1417 );
1418 assert_arrays_eq!(
1419 v_non_null,
1420 expected(false),
1421 &mut array_session().create_execution_ctx()
1422 );
1423 }
1424
1425 #[test]
1426 fn test_timestamp_timezone_microsecond_array_conversion() {
1427 let arrow_array =
1428 TimestampMicrosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])
1429 .with_timezone("UTC");
1430 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1431
1432 let arrow_array_non_null =
1433 TimestampMicrosecondArray::from(vec![1000_i64, 2000, 3000, 4000]).with_timezone("UTC");
1434 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1435
1436 assert_eq!(vortex_array.len(), 4);
1437 assert_eq!(
1438 vortex_array.dtype(),
1439 &DType::Extension(
1440 Timestamp::new_with_tz(
1441 TimeUnit::Microseconds,
1442 Some("UTC".into()),
1443 Nullability::Nullable
1444 )
1445 .erased()
1446 ),
1447 );
1448 assert_eq!(vortex_array_non_null.len(), 4);
1449 assert_eq!(
1450 vortex_array_non_null.dtype(),
1451 &DType::Extension(
1452 Timestamp::new_with_tz(
1453 TimeUnit::Microseconds,
1454 Some("UTC".into()),
1455 Nullability::NonNullable
1456 )
1457 .erased()
1458 )
1459 );
1460 assert_arrays_eq!(
1461 vortex_array,
1462 TemporalArray::new_timestamp(
1463 PrimitiveArray::from_option_iter([Some(1000i64), None, Some(3000), Some(4000)])
1464 .into_array(),
1465 TimeUnit::Microseconds,
1466 Some("UTC".into())
1467 ),
1468 &mut array_session().create_execution_ctx()
1469 );
1470 assert_arrays_eq!(
1471 vortex_array_non_null,
1472 TemporalArray::new_timestamp(
1473 buffer![1000i64, 2000, 3000, 4000].into_array(),
1474 TimeUnit::Microseconds,
1475 Some("UTC".into())
1476 ),
1477 &mut array_session().create_execution_ctx()
1478 );
1479 }
1480
1481 #[rstest]
1483 #[case::utf8(
1484 Arc::new(StringArray::from(vec![Some("hello"), None, Some("world"), Some("test")])),
1485 Arc::new(StringArray::from(vec!["hello", "world", "test", "vortex"])),
1486 DType::Utf8(Nullability::NonNullable),
1487 )]
1488 #[case::large_utf8(
1489 Arc::new(LargeStringArray::from(vec![Some("hello"), None, Some("world"), Some("test")])),
1490 Arc::new(LargeStringArray::from(vec!["hello", "world", "test", "vortex"])),
1491 DType::Utf8(Nullability::NonNullable),
1492 )]
1493 #[case::binary(
1494 Arc::new(BinaryArray::from(vec![
1495 Some("hello".as_bytes()), None, Some("world".as_bytes()), Some("test".as_bytes()),
1496 ])),
1497 Arc::new(BinaryArray::from(vec![
1498 "hello".as_bytes(), "world".as_bytes(), "test".as_bytes(), "vortex".as_bytes(),
1499 ])),
1500 DType::Binary(Nullability::NonNullable),
1501 )]
1502 #[case::large_binary(
1503 Arc::new(LargeBinaryArray::from(vec![
1504 Some("hello".as_bytes()), None, Some("world".as_bytes()), Some("test".as_bytes()),
1505 ])),
1506 Arc::new(LargeBinaryArray::from(vec![
1507 "hello".as_bytes(), "world".as_bytes(), "test".as_bytes(), "vortex".as_bytes(),
1508 ])),
1509 DType::Binary(Nullability::NonNullable),
1510 )]
1511 fn test_string_binary_array_conversion(
1512 #[case] nullable: Arc<dyn ArrowArray>,
1513 #[case] non_nullable: Arc<dyn ArrowArray>,
1514 #[case] expected_non_nullable_dtype: DType,
1515 ) {
1516 let v_null = ArrayRef::from_arrow(nullable.as_ref(), true).unwrap();
1517 let v_non_null = ArrayRef::from_arrow(non_nullable.as_ref(), false).unwrap();
1518 assert_eq!(v_null.len(), 4);
1519 assert_eq!(v_non_null.len(), 4);
1520 assert_eq!(v_null.dtype(), &expected_non_nullable_dtype.as_nullable());
1521 assert_eq!(v_non_null.dtype(), &expected_non_nullable_dtype);
1522 assert_arrays_eq!(
1523 v_null,
1524 VarBinArray::from_iter(
1525 [Some("hello"), None, Some("world"), Some("test")],
1526 expected_non_nullable_dtype.as_nullable()
1527 ),
1528 &mut array_session().create_execution_ctx()
1529 );
1530 assert_arrays_eq!(
1531 v_non_null,
1532 VarBinArray::from_iter_nonnull(
1533 ["hello", "world", "test", "vortex"],
1534 expected_non_nullable_dtype
1535 ),
1536 &mut array_session().create_execution_ctx()
1537 );
1538 }
1539
1540 #[test]
1541 fn test_utf8_view_array_conversion() {
1542 let mut builder = StringViewBuilder::new();
1543 builder.append_value("hello");
1544 builder.append_null();
1545 builder.append_value("world");
1546 builder.append_value("test");
1547 let arrow_array = builder.finish();
1548 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1549
1550 let mut builder_non_null = StringViewBuilder::new();
1551 builder_non_null.append_value("hello");
1552 builder_non_null.append_value("world");
1553 builder_non_null.append_value("test");
1554 builder_non_null.append_value("vortex");
1555 let arrow_array_non_null = builder_non_null.finish();
1556 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1557
1558 assert_eq!(vortex_array.len(), 4);
1559 assert_eq!(vortex_array_non_null.len(), 4);
1560
1561 let varbin_view_array = vortex_array.as_::<VarBinView>();
1563 assert_eq!(
1564 varbin_view_array.data_buffers().len(),
1565 arrow_array.data_buffers().len()
1566 );
1567 assert_eq!(varbin_view_array.dtype(), &DType::Utf8(true.into()));
1568
1569 let varbin_view_array_non_null = vortex_array_non_null.as_::<VarBinView>();
1570 assert_eq!(
1571 varbin_view_array_non_null.data_buffers().len(),
1572 arrow_array_non_null.data_buffers().len()
1573 );
1574 assert_eq!(
1575 varbin_view_array_non_null.dtype(),
1576 &DType::Utf8(false.into())
1577 );
1578 assert_arrays_eq!(
1579 vortex_array,
1580 VarBinArray::from_iter(
1581 [Some("hello"), None, Some("world"), Some("test")],
1582 DType::Utf8(Nullability::NonNullable).as_nullable()
1583 ),
1584 &mut array_session().create_execution_ctx()
1585 );
1586 assert_arrays_eq!(
1587 vortex_array_non_null,
1588 VarBinArray::from_iter_nonnull(
1589 ["hello", "world", "test", "vortex"],
1590 DType::Utf8(Nullability::NonNullable)
1591 ),
1592 &mut array_session().create_execution_ctx()
1593 );
1594 }
1595
1596 #[test]
1597 fn test_binary_view_array_conversion() {
1598 let mut builder = BinaryViewBuilder::new();
1599 builder.append_value(b"hello");
1600 builder.append_null();
1601 builder.append_value(b"world");
1602 builder.append_value(b"test");
1603 let arrow_array = builder.finish();
1604 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1605
1606 let mut builder_non_null = BinaryViewBuilder::new();
1607 builder_non_null.append_value(b"hello");
1608 builder_non_null.append_value(b"world");
1609 builder_non_null.append_value(b"test");
1610 builder_non_null.append_value(b"vortex");
1611 let arrow_array_non_null = builder_non_null.finish();
1612 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1613
1614 assert_eq!(vortex_array.len(), 4);
1615 assert_eq!(vortex_array_non_null.len(), 4);
1616
1617 let varbin_view_array = vortex_array.as_::<VarBinView>();
1619 assert_eq!(
1620 varbin_view_array.data_buffers().len(),
1621 arrow_array.data_buffers().len()
1622 );
1623 assert_eq!(varbin_view_array.dtype(), &DType::Binary(true.into()));
1624
1625 let varbin_view_array_non_null = vortex_array_non_null.as_::<VarBinView>();
1626 assert_eq!(
1627 varbin_view_array_non_null.data_buffers().len(),
1628 arrow_array_non_null.data_buffers().len()
1629 );
1630 assert_eq!(
1631 varbin_view_array_non_null.dtype(),
1632 &DType::Binary(false.into())
1633 );
1634 assert_arrays_eq!(
1635 vortex_array,
1636 VarBinArray::from_iter(
1637 [Some("hello"), None, Some("world"), Some("test")],
1638 DType::Binary(Nullability::NonNullable).as_nullable()
1639 ),
1640 &mut array_session().create_execution_ctx()
1641 );
1642 assert_arrays_eq!(
1643 vortex_array_non_null,
1644 VarBinArray::from_iter_nonnull(
1645 ["hello", "world", "test", "vortex"],
1646 DType::Binary(Nullability::NonNullable)
1647 ),
1648 &mut array_session().create_execution_ctx()
1649 );
1650 }
1651
1652 #[test]
1654 fn test_boolean_array_conversion() {
1655 let arrow_array = BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]);
1656 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1657
1658 let arrow_array_non_null = BooleanArray::from(vec![true, false, true, false]);
1659 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1660
1661 assert_eq!(vortex_array.len(), 4);
1662 assert_eq!(vortex_array_non_null.len(), 4);
1663 assert_arrays_eq!(
1664 vortex_array,
1665 BoolArray::from_iter([Some(true), None, Some(false), Some(true)]),
1666 &mut array_session().create_execution_ctx()
1667 );
1668 assert_arrays_eq!(
1669 vortex_array_non_null,
1670 BoolArray::from_iter([true, false, true, false]),
1671 &mut array_session().create_execution_ctx()
1672 );
1673 }
1674
1675 #[test]
1677 fn test_struct_array_conversion() {
1678 let fields = vec![
1679 Field::new("field1", DataType::Int32, true),
1680 Field::new("field2", DataType::Utf8, false),
1681 ];
1682 let schema = Fields::from(fields);
1683
1684 let field1_data = Int32Array::from(vec![Some(1), None, Some(3)]);
1685 let field2_data = StringArray::from(vec!["a", "b", "c"]);
1686
1687 let arrow_array = StructArray::new(
1688 schema.clone(),
1689 vec![Arc::new(field1_data), Arc::new(field2_data)],
1690 None,
1691 );
1692
1693 let vortex_array = ArrayRef::from_arrow(&arrow_array, false).unwrap();
1694 assert_eq!(vortex_array.len(), 3);
1695
1696 let struct_vortex_array = vortex_array.as_::<Struct>();
1698 assert_eq!(struct_vortex_array.names().len(), 2);
1699 assert_eq!(struct_vortex_array.names()[0], "field1");
1700 assert_eq!(struct_vortex_array.names()[1], "field2");
1701
1702 let nullable_array = StructArray::new(
1704 schema,
1705 vec![
1706 Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])),
1707 Arc::new(StringArray::from(vec!["a", "b", "c"])),
1708 ],
1709 Some(arrow_buffer::NullBuffer::new(BooleanBuffer::from(vec![
1710 true, false, true,
1711 ]))),
1712 );
1713
1714 let vortex_nullable_array = ArrayRef::from_arrow(&nullable_array, true).unwrap();
1715 assert_eq!(vortex_nullable_array.len(), 3);
1716
1717 let struct_vortex_nullable_array = vortex_nullable_array.as_::<Struct>();
1719 assert_eq!(struct_vortex_nullable_array.names().len(), 2);
1720 assert_eq!(struct_vortex_nullable_array.names()[0], "field1");
1721 assert_eq!(struct_vortex_nullable_array.names()[1], "field2");
1722 let expected = |validity| {
1723 VortexStructArray::new(
1724 ["field1", "field2"].into(),
1725 vec![
1726 PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(),
1727 VarBinArray::from_iter_nonnull(
1728 ["a", "b", "c"],
1729 DType::Utf8(Nullability::NonNullable),
1730 )
1731 .into_array(),
1732 ],
1733 3,
1734 validity,
1735 )
1736 };
1737 assert_arrays_eq!(
1738 vortex_array,
1739 expected(Validity::NonNullable),
1740 &mut array_session().create_execution_ctx()
1741 );
1742 assert_arrays_eq!(
1743 vortex_nullable_array,
1744 expected(Validity::from_iter([true, false, true])),
1745 &mut array_session().create_execution_ctx()
1746 );
1747 }
1748
1749 #[test]
1751 fn test_list_array_conversion() {
1752 let mut builder = ListBuilder::new(Int32Builder::new());
1753 builder.append_value([Some(1), None, Some(3)]);
1754 builder.append_null();
1755 builder.append_value([Some(4), Some(5)]);
1756 let arrow_array = builder.finish();
1757
1758 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1759 assert_eq!(vortex_array.len(), 3);
1760
1761 let list_vortex_array = vortex_array.as_::<List>();
1763 let offsets_array = list_vortex_array.offsets().as_::<Primitive>();
1764 assert_eq!(offsets_array.len(), 4); assert_eq!(offsets_array.ptype(), PType::I32);
1766
1767 let mut builder_non_null = ListBuilder::new(Int32Builder::new());
1769 builder_non_null.append_value([Some(1), None, Some(3)]);
1770 builder_non_null.append_value([Some(4), Some(5)]);
1771 let arrow_array_non_null = builder_non_null.finish();
1772
1773 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1774 assert_eq!(vortex_array_non_null.len(), 2);
1775
1776 let list_vortex_array_non_null = vortex_array_non_null.as_::<List>();
1778 let offsets_array_non_null = list_vortex_array_non_null.offsets().as_::<Primitive>();
1779 assert_eq!(offsets_array_non_null.len(), 3); assert_eq!(offsets_array_non_null.ptype(), PType::I32);
1781
1782 let expected_elements =
1783 PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), Some(5)])
1784 .into_array();
1785 let expected = ListArray::new(
1786 expected_elements.clone(),
1787 buffer![0u32, 3, 3, 5].into_array(),
1788 Validity::from_iter([true, false, true]),
1789 );
1790 assert_arrays_eq!(
1791 vortex_array,
1792 expected,
1793 &mut array_session().create_execution_ctx()
1794 );
1795 let expected = ListArray::new(
1796 expected_elements,
1797 buffer![0u32, 3, 5].into_array(),
1798 Validity::NonNullable,
1799 );
1800 assert_arrays_eq!(
1801 vortex_array_non_null,
1802 expected,
1803 &mut array_session().create_execution_ctx()
1804 );
1805 }
1806
1807 #[test]
1808 fn test_sliced_offset_arrays_import_only_their_rows() -> VortexResult<()> {
1809 let strings = |rows: Range<i64>| {
1810 StringArray::from_iter_values(rows.map(|row| format!("value-{row}")))
1811 };
1812 let lists = |rows: Range<i64>| {
1813 let mut builder = ListBuilder::new(Int64Builder::new());
1814 for row in rows {
1815 builder.append_value([Some(row), Some(row + 1)]);
1816 }
1817 builder.finish()
1818 };
1819 let mut ctx = array_session().create_execution_ctx();
1820
1821 let sliced = ArrayRef::from_arrow(&strings(0..1000).slice(500, 2), false)?;
1823 let fresh = ArrayRef::from_arrow(&strings(500..502), false)?;
1824 assert_eq!(sliced.nbytes(), fresh.nbytes());
1825 let expected = VarBinArray::from_iter_nonnull(
1826 ["value-500", "value-501"],
1827 DType::Utf8(Nullability::NonNullable),
1828 );
1829 assert_arrays_eq!(sliced, expected, &mut ctx);
1830 assert_arrays_eq!(fresh, expected, &mut ctx);
1831
1832 let sliced = ArrayRef::from_arrow(&lists(0..1000).slice(500, 2), false)?;
1833 let fresh = ArrayRef::from_arrow(&lists(500..502), false)?;
1834 assert_eq!(sliced.nbytes(), fresh.nbytes());
1835 let expected = ListArray::new(
1836 PrimitiveArray::from_option_iter([Some(500i64), Some(501), Some(501), Some(502)])
1837 .into_array(),
1838 buffer![0u32, 2, 4].into_array(),
1839 Validity::NonNullable,
1840 );
1841 assert_arrays_eq!(sliced, expected, &mut ctx);
1842 assert_arrays_eq!(fresh, expected, &mut ctx);
1843 Ok(())
1844 }
1845
1846 #[test]
1847 fn test_large_list_array_conversion() {
1848 let mut builder = LargeListBuilder::new(Int32Builder::new());
1849 builder.append_value([Some(1), None, Some(3)]);
1850 builder.append_null();
1851 builder.append_value([Some(4), Some(5)]);
1852 let arrow_array = builder.finish();
1853
1854 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1855 assert_eq!(vortex_array.len(), 3);
1856
1857 let list_vortex_array = vortex_array.as_::<List>();
1859 let offsets_array = list_vortex_array.offsets().as_::<Primitive>();
1860 assert_eq!(offsets_array.len(), 4); assert_eq!(offsets_array.ptype(), PType::I64); let mut builder_non_null = LargeListBuilder::new(Int32Builder::new());
1865 builder_non_null.append_value([Some(1), None, Some(3)]);
1866 builder_non_null.append_value([Some(4), Some(5)]);
1867 let arrow_array_non_null = builder_non_null.finish();
1868
1869 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1870 assert_eq!(vortex_array_non_null.len(), 2);
1871
1872 let list_vortex_array_non_null = vortex_array_non_null.as_::<List>();
1874 let offsets_array_non_null = list_vortex_array_non_null.offsets().as_::<Primitive>();
1875 assert_eq!(offsets_array_non_null.len(), 3); assert_eq!(offsets_array_non_null.ptype(), PType::I64); let expected_elements =
1879 PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), Some(5)])
1880 .into_array();
1881 let expected = ListArray::new(
1882 expected_elements.clone(),
1883 buffer![0u32, 3, 3, 5].into_array(),
1884 Validity::from_iter([true, false, true]),
1885 );
1886 assert_arrays_eq!(
1887 vortex_array,
1888 expected,
1889 &mut array_session().create_execution_ctx()
1890 );
1891 let expected = ListArray::new(
1892 expected_elements,
1893 buffer![0u32, 3, 5].into_array(),
1894 Validity::NonNullable,
1895 );
1896 assert_arrays_eq!(
1897 vortex_array_non_null,
1898 expected,
1899 &mut array_session().create_execution_ctx()
1900 );
1901 }
1902
1903 #[test]
1904 fn test_fixed_size_list_array_conversion() {
1905 let values = Int32Array::from(vec![
1907 Some(1),
1908 Some(2),
1909 Some(3), Some(4),
1911 None,
1912 Some(6), Some(7),
1914 Some(8),
1915 Some(9), Some(10),
1917 Some(11),
1918 Some(12), ]);
1920
1921 let field = Arc::new(Field::new("item", DataType::Int32, true));
1923 let arrow_array =
1924 ArrowFixedSizeListArray::try_new(Arc::clone(&field), 3, Arc::new(values), None)
1925 .unwrap();
1926 let vortex_array = ArrayRef::from_arrow(&arrow_array, false).unwrap();
1927
1928 assert_eq!(vortex_array.len(), 4);
1929
1930 let fsl_vortex_array = vortex_array.as_::<FixedSizeList>();
1932 assert_eq!(fsl_vortex_array.list_size(), 3);
1933 assert_eq!(fsl_vortex_array.elements().len(), 12); let values_nullable = Int32Array::from(vec![
1937 Some(1),
1938 Some(2),
1939 Some(3), Some(4),
1941 None,
1942 Some(6), Some(7),
1944 Some(8),
1945 Some(9), ]);
1947
1948 let null_buffer =
1950 arrow_buffer::NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
1951
1952 let arrow_array_nullable = ArrowFixedSizeListArray::try_new(
1953 field,
1954 3,
1955 Arc::new(values_nullable),
1956 Some(null_buffer),
1957 )
1958 .unwrap();
1959 let vortex_array_nullable = ArrayRef::from_arrow(&arrow_array_nullable, true).unwrap();
1960
1961 assert_eq!(vortex_array_nullable.len(), 3);
1962
1963 let fsl_vortex_array_nullable = vortex_array_nullable.as_::<FixedSizeList>();
1965 assert_eq!(fsl_vortex_array_nullable.list_size(), 3);
1966 assert_eq!(fsl_vortex_array_nullable.elements().len(), 9); assert_arrays_eq!(
1968 vortex_array,
1969 VortexFixedSizeListArray::new(
1970 PrimitiveArray::from_option_iter([
1971 Some(1i32),
1972 Some(2),
1973 Some(3),
1974 Some(4),
1975 None,
1976 Some(6),
1977 Some(7),
1978 Some(8),
1979 Some(9),
1980 Some(10),
1981 Some(11),
1982 Some(12)
1983 ])
1984 .into_array(),
1985 3,
1986 Validity::NonNullable,
1987 4
1988 ),
1989 &mut array_session().create_execution_ctx()
1990 );
1991 assert_arrays_eq!(
1992 vortex_array_nullable,
1993 VortexFixedSizeListArray::new(
1994 PrimitiveArray::from_option_iter([
1995 Some(1i32),
1996 Some(2),
1997 Some(3),
1998 Some(4),
1999 None,
2000 Some(6),
2001 Some(7),
2002 Some(8),
2003 Some(9)
2004 ])
2005 .into_array(),
2006 3,
2007 Validity::from_iter([true, false, true]),
2008 3
2009 ),
2010 &mut array_session().create_execution_ctx()
2011 );
2012 }
2013
2014 #[test]
2015 fn test_list_view_array_conversion() {
2016 let values = Int32Array::from(vec![
2018 Some(1),
2019 Some(2),
2020 Some(3), Some(4),
2022 Some(5), Some(6), Some(7),
2025 Some(8),
2026 Some(9),
2027 Some(10), ]);
2029
2030 let offsets = ScalarBuffer::from(vec![0i32, 3, 5, 6]);
2032 let sizes = ScalarBuffer::from(vec![3i32, 2, 1, 4]);
2033
2034 let field = Arc::new(Field::new("item", DataType::Int32, true));
2035 let arrow_array = GenericListViewArray::try_new(
2036 Arc::clone(&field),
2037 offsets.clone(),
2038 sizes.clone(),
2039 Arc::new(values.clone()),
2040 None,
2041 )
2042 .unwrap();
2043
2044 let vortex_array = ArrayRef::from_arrow(&arrow_array, false).unwrap();
2045 assert_eq!(vortex_array.len(), 4);
2046
2047 let list_view_vortex_array = vortex_array.as_::<ListView>();
2049 let offsets_array = list_view_vortex_array.offsets().as_::<Primitive>();
2050 let sizes_array = list_view_vortex_array.sizes().as_::<Primitive>();
2051
2052 assert_eq!(offsets_array.len(), 4);
2053 assert_eq!(offsets_array.ptype(), PType::I32);
2054 assert_eq!(sizes_array.len(), 4);
2055 assert_eq!(sizes_array.ptype(), PType::I32);
2056
2057 let null_buffer =
2059 arrow_buffer::NullBuffer::new(BooleanBuffer::from(vec![true, false, true, true]));
2060
2061 let arrow_array_nullable = GenericListViewArray::try_new(
2062 Arc::clone(&field),
2063 offsets,
2064 sizes,
2065 Arc::new(values.clone()),
2066 Some(null_buffer),
2067 )
2068 .unwrap();
2069
2070 let vortex_array_nullable = ArrayRef::from_arrow(&arrow_array_nullable, true).unwrap();
2071 assert_eq!(vortex_array_nullable.len(), 4);
2072
2073 let large_offsets = ScalarBuffer::from(vec![0i64, 3, 5, 6]);
2075 let large_sizes = ScalarBuffer::from(vec![3i64, 2, 1, 4]);
2076
2077 let large_arrow_array = GenericListViewArray::try_new(
2078 field,
2079 large_offsets,
2080 large_sizes,
2081 Arc::new(values),
2082 None,
2083 )
2084 .unwrap();
2085
2086 let large_vortex_array = ArrayRef::from_arrow(&large_arrow_array, false).unwrap();
2087 assert_eq!(large_vortex_array.len(), 4);
2088
2089 let large_list_view_vortex_array = large_vortex_array.as_::<ListView>();
2091 let large_offsets_array = large_list_view_vortex_array.offsets().as_::<Primitive>();
2092 let large_sizes_array = large_list_view_vortex_array.sizes().as_::<Primitive>();
2093
2094 assert_eq!(large_offsets_array.len(), 4);
2095 assert_eq!(large_offsets_array.ptype(), PType::I64); assert_eq!(large_sizes_array.len(), 4);
2097 assert_eq!(large_sizes_array.ptype(), PType::I64); let expected = |validity| {
2099 ListViewArray::new(
2100 PrimitiveArray::from_option_iter((1i32..=10).map(Some)).into_array(),
2101 buffer![0u32, 3, 5, 6].into_array(),
2102 buffer![3u32, 2, 1, 4].into_array(),
2103 validity,
2104 )
2105 };
2106 assert_arrays_eq!(
2107 vortex_array,
2108 expected(Validity::NonNullable),
2109 &mut array_session().create_execution_ctx()
2110 );
2111 assert_arrays_eq!(
2112 vortex_array_nullable,
2113 expected(Validity::from_iter([true, false, true, true])),
2114 &mut array_session().create_execution_ctx()
2115 );
2116 assert_arrays_eq!(
2117 large_vortex_array,
2118 expected(Validity::NonNullable),
2119 &mut array_session().create_execution_ctx()
2120 );
2121 }
2122
2123 #[test]
2125 fn test_null_array_conversion() {
2126 let arrow_array = NullArray::new(5);
2127 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
2128 assert_eq!(vortex_array.len(), 5);
2129 assert_arrays_eq!(
2130 vortex_array,
2131 VortexNullArray::new(5),
2132 &mut array_session().create_execution_ctx()
2133 );
2134 }
2135
2136 #[test]
2138 fn test_arrow_buffer_conversion() {
2139 let data = vec![1u8, 2, 3, 4, 5];
2140 let arrow_buffer = ArrowBuffer::from_vec(data);
2141 let vortex_array = arrow_buffer.into_array();
2142 assert_eq!(vortex_array.len(), 5);
2143 assert_arrays_eq!(
2144 vortex_array,
2145 buffer![1u8, 2, 3, 4, 5].into_array(),
2146 &mut array_session().create_execution_ctx()
2147 );
2148 }
2149
2150 #[test]
2151 fn test_boolean_buffer_conversion() {
2152 let data = vec![true, false, true, false, true];
2153 let boolean_buffer = BooleanBuffer::from(data);
2154 let vortex_array = boolean_buffer.into_array();
2155 assert_eq!(vortex_array.len(), 5);
2156 assert_arrays_eq!(
2157 vortex_array,
2158 BoolArray::from_iter([true, false, true, false, true]),
2159 &mut array_session().create_execution_ctx()
2160 );
2161 }
2162
2163 #[test]
2164 fn test_scalar_buffer_conversion() {
2165 let data = vec![1i32, 2, 3, 4, 5];
2166 let scalar_buffer = ScalarBuffer::from(data);
2167 let vortex_array = scalar_buffer.into_array();
2168 assert_eq!(vortex_array.len(), 5);
2169 assert_arrays_eq!(
2170 vortex_array,
2171 buffer![1i32, 2, 3, 4, 5].into_array(),
2172 &mut array_session().create_execution_ctx()
2173 );
2174 }
2175
2176 #[test]
2177 fn test_offset_buffer_conversion() {
2178 let data = vec![0i32, 2, 5, 8, 10];
2179 let offset_buffer = OffsetBuffer::new(ScalarBuffer::from(data));
2180 let vortex_array = offset_buffer.into_array();
2181 assert_eq!(vortex_array.len(), 5);
2182 assert_arrays_eq!(
2183 vortex_array,
2184 buffer![0i32, 2, 5, 8, 10].into_array(),
2185 &mut array_session().create_execution_ctx()
2186 );
2187 }
2188
2189 #[test]
2191 fn test_record_batch_conversion() {
2192 let schema = Arc::new(Schema::new(vec![
2193 Field::new("field1", DataType::Int32, false),
2194 Field::new("field2", DataType::Utf8, false),
2195 ]));
2196
2197 let field1_data = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
2198 let field2_data = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
2199
2200 let record_batch = RecordBatch::try_new(schema, vec![field1_data, field2_data]).unwrap();
2201
2202 let vortex_array = ArrayRef::from_arrow(record_batch, false).unwrap();
2203 let expected = VortexStructArray::new(
2204 ["field1", "field2"].into(),
2205 vec![
2206 buffer![1i32, 2, 3, 4].into_array(),
2207 VarBinArray::from_iter_nonnull(
2208 ["a", "b", "c", "d"],
2209 DType::Utf8(Nullability::NonNullable),
2210 )
2211 .into_array(),
2212 ],
2213 4,
2214 Validity::NonNullable,
2215 );
2216 assert_arrays_eq!(
2217 vortex_array,
2218 expected,
2219 &mut array_session().create_execution_ctx()
2220 );
2221
2222 let schema = Arc::new(Schema::new(vec![
2224 Field::new("field1", DataType::Int32, false),
2225 Field::new("field2", DataType::Utf8, false),
2226 ]));
2227
2228 let field1_data = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
2229 let field2_data = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
2230
2231 let record_batch = RecordBatch::try_new(schema, vec![field1_data, field2_data]).unwrap();
2232
2233 let vortex_array = ArrayRef::from_arrow(&record_batch, false).unwrap();
2234 let expected = VortexStructArray::new(
2235 ["field1", "field2"].into(),
2236 vec![
2237 buffer![1i32, 2, 3, 4].into_array(),
2238 VarBinArray::from_iter_nonnull(
2239 ["a", "b", "c", "d"],
2240 DType::Utf8(Nullability::NonNullable),
2241 )
2242 .into_array(),
2243 ],
2244 4,
2245 Validity::NonNullable,
2246 );
2247 assert_arrays_eq!(
2248 vortex_array,
2249 expected,
2250 &mut array_session().create_execution_ctx()
2251 );
2252 }
2253
2254 #[test]
2256 fn test_dyn_array_conversion() {
2257 let int_array = Int32Array::from(vec![1, 2, 3, 4]);
2258 let dyn_array: &dyn ArrowArray = &int_array;
2259 let vortex_array = ArrayRef::from_arrow(dyn_array, false).unwrap();
2260 assert_arrays_eq!(
2261 vortex_array,
2262 buffer![1i32, 2, 3, 4].into_array(),
2263 &mut array_session().create_execution_ctx()
2264 );
2265
2266 let string_array = StringArray::from(vec!["a", "b", "c"]);
2267 let dyn_array: &dyn ArrowArray = &string_array;
2268 let vortex_array = ArrayRef::from_arrow(dyn_array, false).unwrap();
2269 assert_arrays_eq!(
2270 vortex_array,
2271 VarBinArray::from_iter_nonnull(["a", "b", "c"], DType::Utf8(Nullability::NonNullable)),
2272 &mut array_session().create_execution_ctx()
2273 );
2274
2275 let bool_array = BooleanArray::from(vec![true, false, true]);
2276 let dyn_array: &dyn ArrowArray = &bool_array;
2277 let vortex_array = ArrayRef::from_arrow(dyn_array, false).unwrap();
2278 assert_arrays_eq!(
2279 vortex_array,
2280 BoolArray::from_iter([true, false, true]),
2281 &mut array_session().create_execution_ctx()
2282 );
2283 }
2284
2285 #[test]
2287 pub fn nullable_may_contain_non_nullable() {
2288 let null_struct_array_with_non_nullable_field = new_null_array(
2289 &DataType::Struct(Fields::from(vec![Field::new(
2290 "non_nullable_inner",
2291 DataType::Int32,
2292 false,
2293 )])),
2294 1,
2295 );
2296 let actual =
2297 ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap();
2298 assert_arrays_eq!(
2299 actual,
2300 VortexStructArray::new(
2301 ["non_nullable_inner"].into(),
2302 vec![buffer![0i32].into_array()],
2303 1,
2304 Validity::AllInvalid
2305 ),
2306 &mut array_session().create_execution_ctx()
2307 );
2308 }
2309
2310 #[test]
2311 pub fn nullable_may_contain_deeply_nested_non_nullable() {
2312 let null_struct_array_with_non_nullable_field = new_null_array(
2313 &DataType::Struct(Fields::from(vec![Field::new(
2314 "non_nullable_inner",
2315 DataType::Struct(Fields::from(vec![Field::new(
2316 "non_nullable_deeper_inner",
2317 DataType::Int32,
2318 false,
2319 )])),
2320 false,
2321 )])),
2322 1,
2323 );
2324 let actual =
2325 ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap();
2326 assert_arrays_eq!(
2327 actual,
2328 VortexStructArray::new(
2329 ["non_nullable_inner"].into(),
2330 vec![
2331 VortexStructArray::new(
2332 ["non_nullable_deeper_inner"].into(),
2333 vec![buffer![0i32].into_array()],
2334 1,
2335 Validity::NonNullable
2336 )
2337 .into_array()
2338 ],
2339 1,
2340 Validity::AllInvalid
2341 ),
2342 &mut array_session().create_execution_ctx()
2343 );
2344 }
2345
2346 #[test]
2347 fn dyn_dispatch_imports_map() -> VortexResult<()> {
2348 let mut builder = ArrowMapBuilder::new(None, StringBuilder::new(), Int32Builder::new());
2349 builder.keys().append_value("joe");
2350 builder.values().append_value(1);
2351 builder.append(true)?;
2352 builder.append(false)?;
2353 let arrow = builder.finish();
2354
2355 let expected = DType::map(
2356 DType::Utf8(Nullability::NonNullable),
2357 DType::Primitive(PType::I32, Nullability::Nullable),
2358 false,
2359 Nullability::Nullable,
2360 )?;
2361
2362 let direct = from_arrow_dyn(&arrow, true)?;
2364 assert_eq!(direct.dtype(), &expected);
2365
2366 let batch = RecordBatch::try_from_iter_with_nullable([(
2368 "maps",
2369 Arc::new(arrow) as Arc<dyn ArrowArray>,
2370 true,
2371 )])?;
2372 let imported = from_arrow_batch(&batch, false)?;
2373 assert_eq!(
2374 imported.as_::<Struct>().unmasked_field(0).dtype(),
2375 &expected
2376 );
2377
2378 let row = Scalar::try_map(
2379 expected.clone(),
2380 [(
2381 Scalar::utf8("joe", Nullability::NonNullable),
2382 Scalar::primitive(1i32, Nullability::Nullable),
2383 )],
2384 )?;
2385 let mut ctx = array_session().create_execution_ctx();
2386 for actual in [&direct, imported.as_::<Struct>().unmasked_field(0)] {
2387 assert_eq!(actual.len(), 2);
2388 assert_eq!(actual.execute_scalar(0, &mut ctx)?, row);
2389 assert_eq!(
2390 actual.execute_scalar(1, &mut ctx)?,
2391 Scalar::null(expected.clone())
2392 );
2393 }
2394 Ok(())
2395 }
2396
2397 #[test]
2398 fn non_nullable_request_rejects_nulls() {
2399 let arrow_array = Int32Array::from(vec![Some(1), None, Some(3)]);
2402 assert!(ArrayRef::from_arrow(&arrow_array, false).is_err());
2403 }
2404
2405 #[test]
2406 fn non_nullable_request_rejects_null_array() {
2407 let arrow_array = NullArray::new(5);
2410 assert!(ArrayRef::from_arrow(&arrow_array, false).is_err());
2411 }
2412
2413 #[test]
2414 fn non_nullable_struct_with_nulls_errors() {
2415 let struct_array = new_null_array(
2418 &DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])),
2419 3,
2420 );
2421 assert!(ArrayRef::from_arrow(struct_array.as_ref(), false).is_err());
2422 }
2423
2424 #[test]
2425 fn non_nullable_list_with_nulls_errors() {
2426 let mut builder = ListBuilder::new(Int32Builder::new());
2429 builder.append_value([Some(1), Some(2)]);
2430 builder.append_null();
2431 let list = builder.finish();
2432 assert!(ArrayRef::from_arrow(&list, false).is_err());
2433 }
2434
2435 #[test]
2436 pub fn nullable_struct_containing_non_nullable_dictionary_with_nulls_errors() {
2437 let null_struct_array_with_non_nullable_field = new_null_array(
2441 &DataType::Struct(Fields::from(vec![Field::new(
2442 "non_nullable_deeper_inner",
2443 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2444 false,
2445 )])),
2446 1,
2447 );
2448
2449 assert!(
2450 ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).is_err()
2451 );
2452 }
2453}