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