1use std::sync::Arc;
5
6use arrow_array::AnyDictionaryArray;
7use arrow_array::Array as ArrowArray;
8use arrow_array::ArrowPrimitiveType;
9use arrow_array::BooleanArray as ArrowBooleanArray;
10use arrow_array::DictionaryArray;
11use arrow_array::FixedSizeListArray as ArrowFixedSizeListArray;
12use arrow_array::GenericByteArray;
13use arrow_array::GenericByteViewArray;
14use arrow_array::GenericListArray;
15use arrow_array::GenericListViewArray;
16use arrow_array::MapArray as ArrowMapArray;
17use arrow_array::NullArray as ArrowNullArray;
18use arrow_array::OffsetSizeTrait;
19use arrow_array::PrimitiveArray as ArrowPrimitiveArray;
20use arrow_array::RecordBatch;
21use arrow_array::StructArray as ArrowStructArray;
22use arrow_array::cast::AsArray;
23use arrow_array::cast::as_null_array;
24use arrow_array::make_array;
25use arrow_array::types::ArrowDictionaryKeyType;
26use arrow_array::types::ByteArrayType;
27use arrow_array::types::ByteViewType;
28use arrow_array::types::Date32Type;
29use arrow_array::types::Date64Type;
30use arrow_array::types::Decimal32Type;
31use arrow_array::types::Decimal64Type;
32use arrow_array::types::Decimal128Type;
33use arrow_array::types::Decimal256Type;
34use arrow_array::types::Float16Type;
35use arrow_array::types::Float32Type;
36use arrow_array::types::Float64Type;
37use arrow_array::types::Int8Type;
38use arrow_array::types::Int16Type;
39use arrow_array::types::Int32Type;
40use arrow_array::types::Int64Type;
41use arrow_array::types::Time32MillisecondType;
42use arrow_array::types::Time32SecondType;
43use arrow_array::types::Time64MicrosecondType;
44use arrow_array::types::Time64NanosecondType;
45use arrow_array::types::TimestampMicrosecondType;
46use arrow_array::types::TimestampMillisecondType;
47use arrow_array::types::TimestampNanosecondType;
48use arrow_array::types::TimestampSecondType;
49use arrow_array::types::UInt8Type;
50use arrow_array::types::UInt16Type;
51use arrow_array::types::UInt32Type;
52use arrow_array::types::UInt64Type;
53use arrow_buffer::ArrowNativeType;
54use arrow_buffer::BooleanBuffer;
55use arrow_buffer::Buffer as ArrowBuffer;
56use arrow_buffer::ScalarBuffer;
57use arrow_buffer::buffer::NullBuffer;
58use arrow_buffer::buffer::OffsetBuffer;
59use arrow_schema::DataType;
60use arrow_schema::TimeUnit as ArrowTimeUnit;
61use vortex_array::ArrayRef;
62use vortex_array::IntoArray;
63use vortex_array::arrays::BoolArray;
64use vortex_array::arrays::DecimalArray;
65use vortex_array::arrays::DictArray;
66use vortex_array::arrays::FixedSizeListArray;
67use vortex_array::arrays::ListArray;
68use vortex_array::arrays::ListViewArray;
69use vortex_array::arrays::MapArray;
70use vortex_array::arrays::NullArray;
71use vortex_array::arrays::PrimitiveArray;
72use vortex_array::arrays::Struct;
73use vortex_array::arrays::StructArray;
74use vortex_array::arrays::TemporalArray;
75use vortex_array::arrays::VarBinArray;
76use vortex_array::arrays::VarBinViewArray;
77use vortex_array::arrays::struct_::StructArrayExt;
78use vortex_array::dtype::DType;
79use vortex_array::dtype::DecimalDType;
80use vortex_array::dtype::FieldNames;
81use vortex_array::dtype::IntegerPType;
82use vortex_array::dtype::MapDType;
83use vortex_array::dtype::NativePType;
84use vortex_array::dtype::Nullability;
85use vortex_array::dtype::PType;
86use vortex_array::dtype::i256;
87use vortex_array::extension::datetime::TimeUnit;
88use vortex_array::validity::Validity;
89use vortex_buffer::Alignment;
90use vortex_buffer::BitBuffer;
91use vortex_buffer::Buffer;
92use vortex_buffer::ByteBuffer;
93use vortex_error::VortexResult;
94use vortex_error::vortex_bail;
95use vortex_error::vortex_ensure;
96use vortex_error::vortex_ensure_eq;
97use vortex_error::vortex_err;
98use vortex_error::vortex_panic;
99
100use crate::FromArrowArray;
101use crate::dtype::FromArrowType;
102
103pub trait IntoVortexArray {
108 fn into_array(self) -> ArrayRef;
110}
111
112impl IntoVortexArray for ArrowBuffer {
113 fn into_array(self) -> ArrayRef {
114 PrimitiveArray::from_byte_buffer(
115 ByteBuffer::from_arrow_buffer(self, Alignment::of::<u8>()),
116 PType::U8,
117 Validity::NonNullable,
118 )
119 .into_array()
120 }
121}
122
123impl IntoVortexArray for BooleanBuffer {
124 fn into_array(self) -> ArrayRef {
125 BoolArray::new(self.into(), Validity::NonNullable).into_array()
126 }
127}
128
129impl<T> IntoVortexArray for ScalarBuffer<T>
130where
131 T: ArrowNativeType + NativePType,
132{
133 fn into_array(self) -> ArrayRef {
134 PrimitiveArray::new(
135 Buffer::<T>::from_arrow_scalar_buffer(self),
136 Validity::NonNullable,
137 )
138 .into_array()
139 }
140}
141
142impl<O> IntoVortexArray for OffsetBuffer<O>
143where
144 O: IntegerPType + OffsetSizeTrait,
145{
146 fn into_array(self) -> ArrayRef {
147 let primitive = PrimitiveArray::new(
148 Buffer::from_arrow_scalar_buffer(self.into_inner()),
149 Validity::NonNullable,
150 );
151
152 primitive.into_array()
153 }
154}
155
156macro_rules! impl_from_arrow_primitive {
157 ($T:path) => {
158 impl FromArrowArray<&ArrowPrimitiveArray<$T>> for ArrayRef {
159 fn from_arrow(value: &ArrowPrimitiveArray<$T>, nullable: bool) -> VortexResult<Self> {
160 let buffer = Buffer::from_arrow_scalar_buffer(value.values().clone());
161 let validity = nulls(value.nulls(), nullable)?;
162 Ok(PrimitiveArray::new(buffer, validity).into_array())
163 }
164 }
165 };
166}
167
168impl_from_arrow_primitive!(Int8Type);
169impl_from_arrow_primitive!(Int16Type);
170impl_from_arrow_primitive!(Int32Type);
171impl_from_arrow_primitive!(Int64Type);
172impl_from_arrow_primitive!(UInt8Type);
173impl_from_arrow_primitive!(UInt16Type);
174impl_from_arrow_primitive!(UInt32Type);
175impl_from_arrow_primitive!(UInt64Type);
176impl_from_arrow_primitive!(Float16Type);
177impl_from_arrow_primitive!(Float32Type);
178impl_from_arrow_primitive!(Float64Type);
179
180impl FromArrowArray<&ArrowPrimitiveArray<Decimal32Type>> for ArrayRef {
181 fn from_arrow(
182 array: &ArrowPrimitiveArray<Decimal32Type>,
183 nullable: bool,
184 ) -> VortexResult<Self> {
185 let decimal_type = DecimalDType::new(array.precision(), array.scale());
186 let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
187 let validity = nulls(array.nulls(), nullable)?;
188 Ok(DecimalArray::new(buffer, decimal_type, validity).into_array())
189 }
190}
191
192impl FromArrowArray<&ArrowPrimitiveArray<Decimal64Type>> for ArrayRef {
193 fn from_arrow(
194 array: &ArrowPrimitiveArray<Decimal64Type>,
195 nullable: bool,
196 ) -> VortexResult<Self> {
197 let decimal_type = DecimalDType::new(array.precision(), array.scale());
198 let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
199 let validity = nulls(array.nulls(), nullable)?;
200 Ok(DecimalArray::new(buffer, decimal_type, validity).into_array())
201 }
202}
203
204impl FromArrowArray<&ArrowPrimitiveArray<Decimal128Type>> for ArrayRef {
205 fn from_arrow(
206 array: &ArrowPrimitiveArray<Decimal128Type>,
207 nullable: bool,
208 ) -> VortexResult<Self> {
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}
215
216impl FromArrowArray<&ArrowPrimitiveArray<Decimal256Type>> for ArrayRef {
217 fn from_arrow(
218 array: &ArrowPrimitiveArray<Decimal256Type>,
219 nullable: bool,
220 ) -> VortexResult<Self> {
221 let decimal_type = DecimalDType::new(array.precision(), array.scale());
222 let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
223 let buffer =
227 unsafe { std::mem::transmute::<Buffer<arrow_buffer::i256>, Buffer<i256>>(buffer) };
228 let validity = nulls(array.nulls(), nullable)?;
229 Ok(DecimalArray::new(buffer, decimal_type, validity).into_array())
230 }
231}
232
233macro_rules! impl_from_arrow_temporal {
234 ($T:path) => {
235 impl FromArrowArray<&ArrowPrimitiveArray<$T>> for ArrayRef {
236 fn from_arrow(
237 value: &ArrowPrimitiveArray<$T>,
238 nullable: bool,
239 ) -> vortex_error::VortexResult<Self> {
240 temporal_array(value, nullable)
241 }
242 }
243 };
244}
245
246impl_from_arrow_temporal!(TimestampSecondType);
248impl_from_arrow_temporal!(TimestampMillisecondType);
249impl_from_arrow_temporal!(TimestampMicrosecondType);
250impl_from_arrow_temporal!(TimestampNanosecondType);
251
252impl_from_arrow_temporal!(Time32SecondType);
254impl_from_arrow_temporal!(Time32MillisecondType);
255impl_from_arrow_temporal!(Time64MicrosecondType);
256impl_from_arrow_temporal!(Time64NanosecondType);
257
258impl_from_arrow_temporal!(Date32Type);
260impl_from_arrow_temporal!(Date64Type);
261
262fn temporal_array<T: ArrowPrimitiveType>(
263 value: &ArrowPrimitiveArray<T>,
264 nullable: bool,
265) -> VortexResult<ArrayRef>
266where
267 T::Native: NativePType,
268{
269 let arr = PrimitiveArray::new(
270 Buffer::from_arrow_scalar_buffer(value.values().clone()),
271 nulls(value.nulls(), nullable)?,
272 )
273 .into_array();
274
275 Ok(match value.data_type() {
276 DataType::Timestamp(time_unit, tz) => {
277 TemporalArray::new_timestamp(arr, TimeUnit::from_arrow(time_unit), tz.clone()).into()
278 }
279 DataType::Time32(time_unit) => {
280 TemporalArray::new_time(arr, TimeUnit::from_arrow(time_unit)).into()
281 }
282 DataType::Time64(time_unit) => {
283 TemporalArray::new_time(arr, TimeUnit::from_arrow(time_unit)).into()
284 }
285 DataType::Date32 => TemporalArray::new_date(arr, TimeUnit::Days).into(),
286 DataType::Date64 => TemporalArray::new_date(arr, TimeUnit::Milliseconds).into(),
287 DataType::Duration(_) => unimplemented!(),
288 DataType::Interval(_) => unimplemented!(),
289 _ => vortex_panic!("Invalid temporal type: {}", value.data_type()),
290 })
291}
292
293impl<T: ByteArrayType> FromArrowArray<&GenericByteArray<T>> for ArrayRef
294where
295 <T as ByteArrayType>::Offset: IntegerPType,
296{
297 fn from_arrow(value: &GenericByteArray<T>, nullable: bool) -> VortexResult<Self> {
298 let dtype = match T::DATA_TYPE {
299 DataType::Binary | DataType::LargeBinary => DType::Binary(nullable.into()),
300 DataType::Utf8 | DataType::LargeUtf8 => DType::Utf8(nullable.into()),
301 dt => vortex_panic!("Invalid data type for ByteArray: {dt}"),
302 };
303 Ok(unsafe {
305 VarBinArray::new_unchecked(
306 value.offsets().clone().into_array(),
307 ByteBuffer::from_arrow_buffer(value.values().clone(), Alignment::of::<u8>()),
308 dtype,
309 nulls(value.nulls(), nullable)?,
310 )
311 }
312 .into_array())
313 }
314}
315
316impl<T: ByteViewType> FromArrowArray<&GenericByteViewArray<T>> for ArrayRef {
317 fn from_arrow(value: &GenericByteViewArray<T>, nullable: bool) -> VortexResult<Self> {
318 let dtype = match T::DATA_TYPE {
319 DataType::BinaryView => DType::Binary(nullable.into()),
320 DataType::Utf8View => DType::Utf8(nullable.into()),
321 dt => vortex_panic!("Invalid data type for ByteViewArray: {dt}"),
322 };
323
324 let views_buffer = Buffer::from_byte_buffer(
325 Buffer::from_arrow_scalar_buffer(value.views().clone()).into_byte_buffer(),
326 );
327
328 Ok(unsafe {
331 VarBinViewArray::new_unchecked(
332 views_buffer,
333 Arc::from(
334 value
335 .data_buffers()
336 .iter()
337 .map(|b| ByteBuffer::from_arrow_buffer(b.clone(), Alignment::of::<u8>()))
338 .collect::<Vec<_>>(),
339 ),
340 dtype,
341 nulls(value.nulls(), nullable)?,
342 )
343 .into_array()
344 })
345 }
346}
347
348impl FromArrowArray<&ArrowBooleanArray> for ArrayRef {
349 fn from_arrow(value: &ArrowBooleanArray, nullable: bool) -> VortexResult<Self> {
350 Ok(BoolArray::new(
351 value.values().clone().into(),
352 nulls(value.nulls(), nullable)?,
353 )
354 .into_array())
355 }
356}
357
358pub(crate) fn remove_nulls(data: arrow_data::ArrayData) -> VortexResult<arrow_data::ArrayData> {
360 if data.null_count() == 0 {
361 return Ok(data);
363 }
364
365 let children = match data.data_type() {
366 DataType::Struct(fields) => Some(
367 fields
368 .iter()
369 .zip(data.child_data().iter())
370 .map(|(field, child_data)| {
371 if field.is_nullable() {
372 Ok(child_data.clone())
373 } else {
374 remove_nulls(child_data.clone())
375 }
376 })
377 .collect::<VortexResult<Vec<_>>>()?,
378 ),
379 DataType::List(f)
380 | DataType::LargeList(f)
381 | DataType::ListView(f)
382 | DataType::LargeListView(f)
383 | DataType::FixedSizeList(f, _)
384 if !f.is_nullable() =>
385 {
386 vortex_ensure_eq!(
388 data.child_data().len(),
389 1,
390 "List types should have one child"
391 );
392 Some(vec![remove_nulls(data.child_data()[0].clone())?])
393 }
394 _ => None,
395 };
396
397 let mut builder = data.into_builder().nulls(None);
398 if let Some(children) = children {
399 builder = builder.child_data(children);
400 }
401 builder
402 .build()
403 .map_err(|e| vortex_err!("Failed to reconstruct Arrow array without nulls: {e}"))
404}
405
406impl FromArrowArray<&ArrowStructArray> for ArrayRef {
407 fn from_arrow(value: &ArrowStructArray, nullable: bool) -> VortexResult<Self> {
408 Ok(StructArray::try_new(
409 value.column_names().iter().copied().collect(),
410 value
411 .columns()
412 .iter()
413 .zip(value.fields())
414 .map(|(c, field)| {
415 if c.null_count() > 0 && !field.is_nullable() {
418 let stripped = make_array(remove_nulls(c.into_data())?);
419 Self::from_arrow(stripped.as_ref(), false)
420 } else {
421 Self::from_arrow(c.as_ref(), field.is_nullable())
422 }
423 })
424 .collect::<VortexResult<Vec<_>>>()?,
425 value.len(),
426 nulls(value.nulls(), nullable)?,
427 )?
428 .into_array())
429 }
430}
431
432impl<O: IntegerPType + OffsetSizeTrait> FromArrowArray<&GenericListArray<O>> for ArrayRef {
433 fn from_arrow(value: &GenericListArray<O>, nullable: bool) -> VortexResult<Self> {
434 let elements_are_nullable = match value.data_type() {
436 DataType::List(field) => field.is_nullable(),
437 DataType::LargeList(field) => field.is_nullable(),
438 dt => vortex_panic!("Invalid data type for ListArray: {dt}"),
439 };
440
441 let elements = Self::from_arrow(value.values().as_ref(), elements_are_nullable)?;
442
443 let offsets = value.offsets().clone().into_array();
445 let nulls = nulls(value.nulls(), nullable)?;
446
447 Ok(ListArray::try_new(elements, offsets, nulls)?.into_array())
448 }
449}
450
451impl<O: OffsetSizeTrait + NativePType> FromArrowArray<&GenericListViewArray<O>> for ArrayRef {
452 fn from_arrow(array: &GenericListViewArray<O>, nullable: bool) -> VortexResult<Self> {
453 let elements_are_nullable = match array.data_type() {
455 DataType::ListView(field) => field.is_nullable(),
456 DataType::LargeListView(field) => field.is_nullable(),
457 dt => vortex_panic!("Invalid data type for ListViewArray: {dt}"),
458 };
459
460 let elements = Self::from_arrow(array.values().as_ref(), elements_are_nullable)?;
461
462 let offsets = array.offsets().clone().into_array();
464 let sizes = array.sizes().clone().into_array();
465 let nulls = nulls(array.nulls(), nullable)?;
466
467 Ok(ListViewArray::try_new(elements, offsets, sizes, nulls)?.into_array())
468 }
469}
470
471impl FromArrowArray<&ArrowFixedSizeListArray> for ArrayRef {
472 fn from_arrow(array: &ArrowFixedSizeListArray, nullable: bool) -> VortexResult<Self> {
473 let DataType::FixedSizeList(field, list_size) = array.data_type() else {
474 vortex_panic!("Invalid data type for ListArray: {}", array.data_type());
475 };
476
477 Ok(FixedSizeListArray::try_new(
478 Self::from_arrow(array.values().as_ref(), field.is_nullable())?,
479 *list_size as u32,
480 nulls(array.nulls(), nullable)?,
481 array.len(),
482 )?
483 .into_array())
484 }
485}
486
487pub(crate) fn map_from_arrow_parts(
488 entries: ArrayRef,
489 offsets: &OffsetBuffer<i32>,
490 arrow_nulls: Option<&NullBuffer>,
491 keys_sorted: bool,
492 nullable: bool,
493) -> VortexResult<ArrayRef> {
494 let DType::Struct(struct_dtype, Nullability::NonNullable) = entries.dtype() else {
495 vortex_bail!(
496 "Arrow map entries must import as non-nullable struct, got {}",
497 entries.dtype()
498 );
499 };
500 vortex_ensure!(
501 struct_dtype.nfields() == 2,
502 "Arrow map entries struct must contain exactly two fields"
503 );
504
505 let key_dtype = struct_dtype
506 .field_by_index(0)
507 .ok_or_else(|| vortex_err!("Arrow map entries struct missing key field"))?;
508 vortex_ensure!(
509 !key_dtype.is_nullable(),
510 "Arrow map key field must be non-nullable"
511 );
512 let value_dtype = struct_dtype
513 .field_by_index(1)
514 .ok_or_else(|| vortex_err!("Arrow map entries struct missing value field"))?;
515 let map_dtype = MapDType::try_new(key_dtype, value_dtype, keys_sorted)?;
516 let entries_struct = entries.as_::<Struct>();
517 let entries = StructArray::try_new(
518 FieldNames::from(["key", "value"]),
519 vec![
520 entries_struct.unmasked_field(0).clone(),
521 entries_struct.unmasked_field(1).clone(),
522 ],
523 entries.len(),
524 Validity::NonNullable,
525 )?
526 .into_array();
527
528 let len = offsets.len() - 1;
529 let sizes = Buffer::<i32>::from_iter(offsets.windows(2).map(|window| window[1] - window[0]))
530 .into_array();
531 let offsets = offsets.inner().slice(0, len).into_array();
532 let validity = nulls(arrow_nulls, nullable)?;
533 let entries = ListViewArray::try_new(entries, offsets, sizes, validity)?;
534 let entries = unsafe { entries.with_zero_copy_to_list(true) };
538
539 Ok(MapArray::try_new(map_dtype, entries)?.into_array())
540}
541
542impl FromArrowArray<&ArrowMapArray> for ArrayRef {
543 fn from_arrow(array: &ArrowMapArray, nullable: bool) -> VortexResult<Self> {
544 let DataType::Map(_, keys_sorted) = array.data_type() else {
545 vortex_panic!("Invalid data type for MapArray: {}", array.data_type());
546 };
547 let entries = Self::from_arrow(array.entries(), false)?;
548 map_from_arrow_parts(
549 entries,
550 array.offsets(),
551 array.nulls(),
552 *keys_sorted,
553 nullable,
554 )
555 }
556}
557
558impl FromArrowArray<&ArrowNullArray> for ArrayRef {
559 fn from_arrow(value: &ArrowNullArray, nullable: bool) -> VortexResult<Self> {
560 vortex_ensure!(
561 nullable,
562 "Cannot convert an Arrow NullArray into a non-nullable Vortex array"
563 );
564 Ok(NullArray::new(value.len()).into_array())
565 }
566}
567
568impl<K: ArrowDictionaryKeyType> FromArrowArray<&DictionaryArray<K>> for DictArray {
569 fn from_arrow(array: &DictionaryArray<K>, nullable: bool) -> VortexResult<Self> {
570 let keys = AnyDictionaryArray::keys(array);
571 let keys = ArrayRef::from_arrow(keys, keys.is_nullable())?;
572 let values = ArrayRef::from_arrow(array.values().as_ref(), nullable)?;
573 Ok(unsafe { DictArray::new_unchecked(keys, values) })
575 }
576}
577
578pub(crate) fn nulls(nulls: Option<&NullBuffer>, nullable: bool) -> VortexResult<Validity> {
579 if nullable {
580 Ok(nulls
581 .map(|nulls| {
582 if nulls.null_count() == nulls.len() {
583 Validity::AllInvalid
584 } else {
585 Validity::from(BitBuffer::from(nulls.inner().clone()))
586 }
587 })
588 .unwrap_or(Validity::AllValid))
589 } else {
590 let null_count = nulls.map(NullBuffer::null_count).unwrap_or(0);
591 vortex_ensure_eq!(
592 null_count,
593 0,
594 "Cannot convert an Arrow array containing {null_count} nulls into a non-nullable Vortex array"
595 );
596 Ok(Validity::NonNullable)
597 }
598}
599
600impl FromArrowArray<&dyn ArrowArray> for ArrayRef {
601 fn from_arrow(array: &dyn ArrowArray, nullable: bool) -> VortexResult<Self> {
602 match array.data_type() {
603 DataType::Boolean => Self::from_arrow(array.as_boolean(), nullable),
604 DataType::UInt8 => Self::from_arrow(array.as_primitive::<UInt8Type>(), nullable),
605 DataType::UInt16 => Self::from_arrow(array.as_primitive::<UInt16Type>(), nullable),
606 DataType::UInt32 => Self::from_arrow(array.as_primitive::<UInt32Type>(), nullable),
607 DataType::UInt64 => Self::from_arrow(array.as_primitive::<UInt64Type>(), nullable),
608 DataType::Int8 => Self::from_arrow(array.as_primitive::<Int8Type>(), nullable),
609 DataType::Int16 => Self::from_arrow(array.as_primitive::<Int16Type>(), nullable),
610 DataType::Int32 => Self::from_arrow(array.as_primitive::<Int32Type>(), nullable),
611 DataType::Int64 => Self::from_arrow(array.as_primitive::<Int64Type>(), nullable),
612 DataType::Float16 => Self::from_arrow(array.as_primitive::<Float16Type>(), nullable),
613 DataType::Float32 => Self::from_arrow(array.as_primitive::<Float32Type>(), nullable),
614 DataType::Float64 => Self::from_arrow(array.as_primitive::<Float64Type>(), nullable),
615 DataType::Utf8 => Self::from_arrow(array.as_string::<i32>(), nullable),
616 DataType::LargeUtf8 => Self::from_arrow(array.as_string::<i64>(), nullable),
617 DataType::Binary => Self::from_arrow(array.as_binary::<i32>(), nullable),
618 DataType::LargeBinary => Self::from_arrow(array.as_binary::<i64>(), nullable),
619 DataType::BinaryView => Self::from_arrow(array.as_binary_view(), nullable),
620 DataType::Utf8View => Self::from_arrow(array.as_string_view(), nullable),
621 DataType::Struct(_) => Self::from_arrow(array.as_struct(), nullable),
622 DataType::List(_) => Self::from_arrow(array.as_list::<i32>(), nullable),
623 DataType::LargeList(_) => Self::from_arrow(array.as_list::<i64>(), nullable),
624 DataType::ListView(_) => Self::from_arrow(array.as_list_view::<i32>(), nullable),
625 DataType::LargeListView(_) => Self::from_arrow(array.as_list_view::<i64>(), nullable),
626 DataType::FixedSizeList(..) => Self::from_arrow(array.as_fixed_size_list(), nullable),
627 DataType::Map(..) => Self::from_arrow(array.as_map(), nullable),
628 DataType::Null => Self::from_arrow(as_null_array(array), nullable),
629 DataType::Timestamp(u, _) => match u {
630 ArrowTimeUnit::Second => {
631 Self::from_arrow(array.as_primitive::<TimestampSecondType>(), nullable)
632 }
633 ArrowTimeUnit::Millisecond => {
634 Self::from_arrow(array.as_primitive::<TimestampMillisecondType>(), nullable)
635 }
636 ArrowTimeUnit::Microsecond => {
637 Self::from_arrow(array.as_primitive::<TimestampMicrosecondType>(), nullable)
638 }
639 ArrowTimeUnit::Nanosecond => {
640 Self::from_arrow(array.as_primitive::<TimestampNanosecondType>(), nullable)
641 }
642 },
643 DataType::Date32 => Self::from_arrow(array.as_primitive::<Date32Type>(), nullable),
644 DataType::Date64 => Self::from_arrow(array.as_primitive::<Date64Type>(), nullable),
645 DataType::Time32(u) => match u {
646 ArrowTimeUnit::Second => {
647 Self::from_arrow(array.as_primitive::<Time32SecondType>(), nullable)
648 }
649 ArrowTimeUnit::Millisecond => {
650 Self::from_arrow(array.as_primitive::<Time32MillisecondType>(), nullable)
651 }
652 ArrowTimeUnit::Microsecond | ArrowTimeUnit::Nanosecond => unreachable!(),
653 },
654 DataType::Time64(u) => match u {
655 ArrowTimeUnit::Microsecond => {
656 Self::from_arrow(array.as_primitive::<Time64MicrosecondType>(), nullable)
657 }
658 ArrowTimeUnit::Nanosecond => {
659 Self::from_arrow(array.as_primitive::<Time64NanosecondType>(), nullable)
660 }
661 ArrowTimeUnit::Second | ArrowTimeUnit::Millisecond => unreachable!(),
662 },
663 DataType::Decimal32(..) => {
664 Self::from_arrow(array.as_primitive::<Decimal32Type>(), nullable)
665 }
666 DataType::Decimal64(..) => {
667 Self::from_arrow(array.as_primitive::<Decimal64Type>(), nullable)
668 }
669 DataType::Decimal128(..) => {
670 Self::from_arrow(array.as_primitive::<Decimal128Type>(), nullable)
671 }
672 DataType::Decimal256(..) => {
673 Self::from_arrow(array.as_primitive::<Decimal256Type>(), nullable)
674 }
675 DataType::Dictionary(key_type, _) => match key_type.as_ref() {
676 DataType::Int8 => Ok(DictArray::from_arrow(
677 array.as_dictionary::<Int8Type>(),
678 nullable,
679 )?
680 .into_array()),
681 DataType::Int16 => Ok(DictArray::from_arrow(
682 array.as_dictionary::<Int16Type>(),
683 nullable,
684 )?
685 .into_array()),
686 DataType::Int32 => Ok(DictArray::from_arrow(
687 array.as_dictionary::<Int32Type>(),
688 nullable,
689 )?
690 .into_array()),
691 DataType::Int64 => Ok(DictArray::from_arrow(
692 array.as_dictionary::<Int64Type>(),
693 nullable,
694 )?
695 .into_array()),
696 DataType::UInt8 => Ok(DictArray::from_arrow(
697 array.as_dictionary::<UInt8Type>(),
698 nullable,
699 )?
700 .into_array()),
701 DataType::UInt16 => Ok(DictArray::from_arrow(
702 array.as_dictionary::<UInt16Type>(),
703 nullable,
704 )?
705 .into_array()),
706 DataType::UInt32 => Ok(DictArray::from_arrow(
707 array.as_dictionary::<UInt32Type>(),
708 nullable,
709 )?
710 .into_array()),
711 DataType::UInt64 => Ok(DictArray::from_arrow(
712 array.as_dictionary::<UInt64Type>(),
713 nullable,
714 )?
715 .into_array()),
716 key_dt => vortex_bail!("Unsupported dictionary key type: {key_dt}"),
717 },
718 dt => vortex_bail!("Array encoding not implemented for Arrow data type {dt}"),
719 }
720 }
721}
722
723impl FromArrowArray<RecordBatch> for ArrayRef {
724 fn from_arrow(array: RecordBatch, nullable: bool) -> VortexResult<Self> {
725 ArrayRef::from_arrow(&arrow_array::StructArray::from(array), nullable)
726 }
727}
728
729impl FromArrowArray<&RecordBatch> for ArrayRef {
730 fn from_arrow(array: &RecordBatch, nullable: bool) -> VortexResult<Self> {
731 Self::from_arrow(array.clone(), nullable)
732 }
733}
734
735#[cfg(test)]
736mod tests {
737 use std::sync::Arc;
738
739 use arrow_array::Array as ArrowArray;
740 use arrow_array::BinaryArray;
741 use arrow_array::BooleanArray;
742 use arrow_array::Date32Array;
743 use arrow_array::Date64Array;
744 use arrow_array::FixedSizeListArray as ArrowFixedSizeListArray;
745 use arrow_array::Float32Array;
746 use arrow_array::Float64Array;
747 use arrow_array::GenericListViewArray;
748 use arrow_array::Int8Array;
749 use arrow_array::Int16Array;
750 use arrow_array::Int32Array;
751 use arrow_array::Int64Array;
752 use arrow_array::LargeBinaryArray;
753 use arrow_array::LargeStringArray;
754 use arrow_array::NullArray;
755 use arrow_array::RecordBatch;
756 use arrow_array::StringArray;
757 use arrow_array::StructArray;
758 use arrow_array::Time32MillisecondArray;
759 use arrow_array::Time32SecondArray;
760 use arrow_array::Time64MicrosecondArray;
761 use arrow_array::Time64NanosecondArray;
762 use arrow_array::TimestampMicrosecondArray;
763 use arrow_array::TimestampMillisecondArray;
764 use arrow_array::TimestampNanosecondArray;
765 use arrow_array::TimestampSecondArray;
766 use arrow_array::UInt8Array;
767 use arrow_array::UInt16Array;
768 use arrow_array::UInt32Array;
769 use arrow_array::UInt64Array;
770 use arrow_array::builder::BinaryViewBuilder;
771 use arrow_array::builder::Decimal128Builder;
772 use arrow_array::builder::Decimal256Builder;
773 use arrow_array::builder::Int32Builder;
774 use arrow_array::builder::LargeListBuilder;
775 use arrow_array::builder::ListBuilder;
776 use arrow_array::builder::StringViewBuilder;
777 use arrow_array::new_null_array;
778 use arrow_array::types::ArrowPrimitiveType;
779 use arrow_array::types::Float16Type;
780 use arrow_buffer::BooleanBuffer;
781 use arrow_buffer::Buffer as ArrowBuffer;
782 use arrow_buffer::OffsetBuffer;
783 use arrow_buffer::ScalarBuffer;
784 use arrow_schema::DataType;
785 use arrow_schema::Field;
786 use arrow_schema::Fields;
787 use arrow_schema::Schema;
788 use rstest::rstest;
789 use vortex_array::ArrayRef;
790 use vortex_array::arrays::Decimal;
791 use vortex_array::arrays::FixedSizeList;
792 use vortex_array::arrays::List;
793 use vortex_array::arrays::ListView;
794 use vortex_array::arrays::Primitive;
795 use vortex_array::arrays::Struct;
796 use vortex_array::arrays::VarBinView;
797 use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
798 use vortex_array::arrays::fixed_size_list::FixedSizeListArraySlotsExt;
799 use vortex_array::arrays::list::ListArraySlotsExt;
800 use vortex_array::arrays::listview::ListViewArraySlotsExt;
801 use vortex_array::arrays::struct_::StructArrayExt;
802 use vortex_array::dtype::DType;
803 use vortex_array::dtype::Nullability;
804 use vortex_array::dtype::PType;
805 use vortex_array::extension::datetime::TimeUnit;
806 use vortex_array::extension::datetime::Timestamp;
807
808 use crate::FromArrowArray as _;
809 use crate::IntoVortexArray as _;
810
811 #[rstest]
812 #[case::i8(
813 Arc::new(Int8Array::from(vec![Some(1), None, Some(3), Some(4)])),
814 Arc::new(Int8Array::from(vec![1, 2, 3, 4])),
815 PType::I8,
816 )]
817 #[case::i16(
818 Arc::new(Int16Array::from(vec![Some(100), None, Some(300), Some(400)])),
819 Arc::new(Int16Array::from(vec![100, 200, 300, 400])),
820 PType::I16,
821 )]
822 #[case::i32(
823 Arc::new(Int32Array::from(vec![Some(1000), None, Some(3000), Some(4000)])),
824 Arc::new(Int32Array::from(vec![1000, 2000, 3000, 4000])),
825 PType::I32,
826 )]
827 #[case::i64(
828 Arc::new(Int64Array::from(vec![Some(10000), None, Some(30000), Some(40000)])),
829 Arc::new(Int64Array::from(vec![10000_i64, 20000, 30000, 40000])),
830 PType::I64,
831 )]
832 #[case::u8(
833 Arc::new(UInt8Array::from(vec![Some(1), None, Some(3), Some(4)])),
834 Arc::new(UInt8Array::from(vec![1_u8, 2, 3, 4])),
835 PType::U8,
836 )]
837 #[case::u16(
838 Arc::new(UInt16Array::from(vec![Some(100), None, Some(300), Some(400)])),
839 Arc::new(UInt16Array::from(vec![100_u16, 200, 300, 400])),
840 PType::U16,
841 )]
842 #[case::u32(
843 Arc::new(UInt32Array::from(vec![Some(1000), None, Some(3000), Some(4000)])),
844 Arc::new(UInt32Array::from(vec![1000_u32, 2000, 3000, 4000])),
845 PType::U32,
846 )]
847 #[case::u64(
848 Arc::new(UInt64Array::from(vec![Some(10000), None, Some(30000), Some(40000)])),
849 Arc::new(UInt64Array::from(vec![10000_u64, 20000, 30000, 40000])),
850 PType::U64,
851 )]
852 #[case::f32(
853 Arc::new(Float32Array::from(vec![Some(1.5), None, Some(3.5), Some(4.5)])),
854 Arc::new(Float32Array::from(vec![1.5_f32, 2.5, 3.5, 4.5])),
855 PType::F32,
856 )]
857 #[case::f64(
858 Arc::new(Float64Array::from(vec![Some(1.5), None, Some(3.5), Some(4.5)])),
859 Arc::new(Float64Array::from(vec![1.5_f64, 2.5, 3.5, 4.5])),
860 PType::F64,
861 )]
862 fn test_primitive_array_conversion(
863 #[case] nullable: Arc<dyn ArrowArray>,
864 #[case] non_nullable: Arc<dyn ArrowArray>,
865 #[case] expected_ptype: PType,
866 ) {
867 let v_null = ArrayRef::from_arrow(nullable.as_ref(), true).unwrap();
868 let v_non_null = ArrayRef::from_arrow(non_nullable.as_ref(), false).unwrap();
869 assert_eq!(v_null.len(), 4);
870 assert_eq!(v_non_null.len(), 4);
871 assert_eq!(v_null.as_::<Primitive>().ptype(), expected_ptype);
872 assert_eq!(v_non_null.as_::<Primitive>().ptype(), expected_ptype);
873 }
874
875 #[test]
876 fn test_float16_array_conversion() {
877 let values = vec![
878 Some(<Float16Type as ArrowPrimitiveType>::Native::from_f32(1.5)),
879 None,
880 Some(<Float16Type as ArrowPrimitiveType>::Native::from_f32(3.5)),
881 ];
882 let arrow_array = arrow_array::PrimitiveArray::<Float16Type>::from(values);
883 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
884
885 let non_null_values = vec![
886 <Float16Type as ArrowPrimitiveType>::Native::from_f32(1.5),
887 <Float16Type as ArrowPrimitiveType>::Native::from_f32(2.5),
888 ];
889 let arrow_array_non_null =
890 arrow_array::PrimitiveArray::<Float16Type>::from(non_null_values);
891 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
892
893 assert_eq!(vortex_array.len(), 3);
894 assert_eq!(vortex_array_non_null.len(), 2);
895
896 let primitive_array = vortex_array.as_::<Primitive>();
898 assert_eq!(primitive_array.ptype(), PType::F16);
899
900 let primitive_array_non_null = vortex_array_non_null.as_::<Primitive>();
901 assert_eq!(primitive_array_non_null.ptype(), PType::F16);
902 }
903
904 #[test]
906 fn test_decimal128_array_conversion() {
907 let mut builder = Decimal128Builder::with_capacity(4);
908 builder.append_value(12345);
909 builder.append_null();
910 builder.append_value(67890);
911 builder.append_value(11111);
912 let decimal_array = builder.finish().with_precision_and_scale(10, 2).unwrap();
913
914 let vortex_array = ArrayRef::from_arrow(&decimal_array, true).unwrap();
915 assert_eq!(vortex_array.len(), 4);
916
917 let mut builder_non_null = Decimal128Builder::with_capacity(3);
918 builder_non_null.append_value(12345);
919 builder_non_null.append_value(67890);
920 builder_non_null.append_value(11111);
921 let decimal_array_non_null = builder_non_null
922 .finish()
923 .with_precision_and_scale(10, 2)
924 .unwrap();
925
926 let vortex_array_non_null = ArrayRef::from_arrow(&decimal_array_non_null, false).unwrap();
927 assert_eq!(vortex_array_non_null.len(), 3);
928
929 let decimal_vortex_array = vortex_array.as_::<Decimal>();
931 assert_eq!(decimal_vortex_array.decimal_dtype().precision(), 10);
932 assert_eq!(decimal_vortex_array.decimal_dtype().scale(), 2);
933
934 let decimal_vortex_array_non_null = vortex_array_non_null.as_::<Decimal>();
935 assert_eq!(
936 decimal_vortex_array_non_null.decimal_dtype().precision(),
937 10
938 );
939 assert_eq!(decimal_vortex_array_non_null.decimal_dtype().scale(), 2);
940 }
941
942 #[test]
943 fn test_decimal256_array_conversion() {
944 let mut builder = Decimal256Builder::with_capacity(4);
945 builder.append_value(arrow_buffer::i256::from_i128(12345));
946 builder.append_null();
947 builder.append_value(arrow_buffer::i256::from_i128(67890));
948 builder.append_value(arrow_buffer::i256::from_i128(11111));
949 let decimal_array = builder.finish().with_precision_and_scale(38, 10).unwrap();
950
951 let vortex_array = ArrayRef::from_arrow(&decimal_array, true).unwrap();
952 assert_eq!(vortex_array.len(), 4);
953
954 let mut builder_non_null = Decimal256Builder::with_capacity(3);
955 builder_non_null.append_value(arrow_buffer::i256::from_i128(12345));
956 builder_non_null.append_value(arrow_buffer::i256::from_i128(67890));
957 builder_non_null.append_value(arrow_buffer::i256::from_i128(11111));
958 let decimal_array_non_null = builder_non_null
959 .finish()
960 .with_precision_and_scale(38, 10)
961 .unwrap();
962
963 let vortex_array_non_null = ArrayRef::from_arrow(&decimal_array_non_null, false).unwrap();
964 assert_eq!(vortex_array_non_null.len(), 3);
965
966 let decimal_vortex_array = vortex_array.as_::<Decimal>();
968 assert_eq!(decimal_vortex_array.decimal_dtype().precision(), 38);
969 assert_eq!(decimal_vortex_array.decimal_dtype().scale(), 10);
970
971 let decimal_vortex_array_non_null = vortex_array_non_null.as_::<Decimal>();
972 assert_eq!(
973 decimal_vortex_array_non_null.decimal_dtype().precision(),
974 38
975 );
976 assert_eq!(decimal_vortex_array_non_null.decimal_dtype().scale(), 10);
977 }
978
979 #[rstest]
981 #[case::timestamp_second(
982 Arc::new(TimestampSecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
983 Arc::new(TimestampSecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
984 )]
985 #[case::timestamp_millisecond(
986 Arc::new(TimestampMillisecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
987 Arc::new(TimestampMillisecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
988 )]
989 #[case::timestamp_microsecond(
990 Arc::new(TimestampMicrosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
991 Arc::new(TimestampMicrosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
992 )]
993 #[case::timestamp_nanosecond(
994 Arc::new(TimestampNanosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
995 Arc::new(TimestampNanosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
996 )]
997 #[case::time32_second(
998 Arc::new(Time32SecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
999 Arc::new(Time32SecondArray::from(vec![1000_i32, 2000, 3000, 4000])),
1000 )]
1001 #[case::time32_millisecond(
1002 Arc::new(Time32MillisecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1003 Arc::new(Time32MillisecondArray::from(vec![1000_i32, 2000, 3000, 4000])),
1004 )]
1005 #[case::time64_microsecond(
1006 Arc::new(Time64MicrosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1007 Arc::new(Time64MicrosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
1008 )]
1009 #[case::time64_nanosecond(
1010 Arc::new(Time64NanosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
1011 Arc::new(Time64NanosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
1012 )]
1013 #[case::date32(
1014 Arc::new(Date32Array::from(vec![Some(18000), None, Some(18002), Some(18003)])),
1015 Arc::new(Date32Array::from(vec![18000_i32, 18001, 18002, 18003])),
1016 )]
1017 #[case::date64(
1018 Arc::new(Date64Array::from(vec![Some(1555200000000), None, Some(1555286400000), Some(1555372800000)]
1019 )),
1020 Arc::new(Date64Array::from(vec![1555200000000_i64, 1555213600000, 1555286400000, 1555372800000]
1021 )),
1022 )]
1023 fn test_temporal_array_conversion(
1024 #[case] nullable: Arc<dyn ArrowArray>,
1025 #[case] non_nullable: Arc<dyn ArrowArray>,
1026 ) {
1027 let v_null = ArrayRef::from_arrow(nullable.as_ref(), true).unwrap();
1028 let v_non_null = ArrayRef::from_arrow(non_nullable.as_ref(), false).unwrap();
1029 assert_eq!(v_null.len(), 4);
1030 assert_eq!(v_non_null.len(), 4);
1031 }
1032
1033 #[test]
1034 fn test_timestamp_timezone_microsecond_array_conversion() {
1035 let arrow_array =
1036 TimestampMicrosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])
1037 .with_timezone("UTC");
1038 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1039
1040 let arrow_array_non_null =
1041 TimestampMicrosecondArray::from(vec![1000_i64, 2000, 3000, 4000]).with_timezone("UTC");
1042 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1043
1044 assert_eq!(vortex_array.len(), 4);
1045 assert_eq!(
1046 vortex_array.dtype(),
1047 &DType::Extension(
1048 Timestamp::new_with_tz(
1049 TimeUnit::Microseconds,
1050 Some("UTC".into()),
1051 Nullability::Nullable
1052 )
1053 .erased()
1054 ),
1055 );
1056 assert_eq!(vortex_array_non_null.len(), 4);
1057 assert_eq!(
1058 vortex_array_non_null.dtype(),
1059 &DType::Extension(
1060 Timestamp::new_with_tz(
1061 TimeUnit::Microseconds,
1062 Some("UTC".into()),
1063 Nullability::NonNullable
1064 )
1065 .erased()
1066 )
1067 );
1068 }
1069
1070 #[rstest]
1072 #[case::utf8(
1073 Arc::new(StringArray::from(vec![Some("hello"), None, Some("world"), Some("test")])),
1074 Arc::new(StringArray::from(vec!["hello", "world", "test", "vortex"])),
1075 DType::Utf8(Nullability::NonNullable),
1076 )]
1077 #[case::large_utf8(
1078 Arc::new(LargeStringArray::from(vec![Some("hello"), None, Some("world"), Some("test")])),
1079 Arc::new(LargeStringArray::from(vec!["hello", "world", "test", "vortex"])),
1080 DType::Utf8(Nullability::NonNullable),
1081 )]
1082 #[case::binary(
1083 Arc::new(BinaryArray::from(vec![
1084 Some("hello".as_bytes()), None, Some("world".as_bytes()), Some("test".as_bytes()),
1085 ])),
1086 Arc::new(BinaryArray::from(vec![
1087 "hello".as_bytes(), "world".as_bytes(), "test".as_bytes(), "vortex".as_bytes(),
1088 ])),
1089 DType::Binary(Nullability::NonNullable),
1090 )]
1091 #[case::large_binary(
1092 Arc::new(LargeBinaryArray::from(vec![
1093 Some("hello".as_bytes()), None, Some("world".as_bytes()), Some("test".as_bytes()),
1094 ])),
1095 Arc::new(LargeBinaryArray::from(vec![
1096 "hello".as_bytes(), "world".as_bytes(), "test".as_bytes(), "vortex".as_bytes(),
1097 ])),
1098 DType::Binary(Nullability::NonNullable),
1099 )]
1100 fn test_string_binary_array_conversion(
1101 #[case] nullable: Arc<dyn ArrowArray>,
1102 #[case] non_nullable: Arc<dyn ArrowArray>,
1103 #[case] expected_non_nullable_dtype: DType,
1104 ) {
1105 let v_null = ArrayRef::from_arrow(nullable.as_ref(), true).unwrap();
1106 let v_non_null = ArrayRef::from_arrow(non_nullable.as_ref(), false).unwrap();
1107 assert_eq!(v_null.len(), 4);
1108 assert_eq!(v_non_null.len(), 4);
1109 assert_eq!(v_null.dtype(), &expected_non_nullable_dtype.as_nullable());
1110 assert_eq!(v_non_null.dtype(), &expected_non_nullable_dtype);
1111 }
1112
1113 #[test]
1114 fn test_utf8_view_array_conversion() {
1115 let mut builder = StringViewBuilder::new();
1116 builder.append_value("hello");
1117 builder.append_null();
1118 builder.append_value("world");
1119 builder.append_value("test");
1120 let arrow_array = builder.finish();
1121 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1122
1123 let mut builder_non_null = StringViewBuilder::new();
1124 builder_non_null.append_value("hello");
1125 builder_non_null.append_value("world");
1126 builder_non_null.append_value("test");
1127 builder_non_null.append_value("vortex");
1128 let arrow_array_non_null = builder_non_null.finish();
1129 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1130
1131 assert_eq!(vortex_array.len(), 4);
1132 assert_eq!(vortex_array_non_null.len(), 4);
1133
1134 let varbin_view_array = vortex_array.as_::<VarBinView>();
1136 assert_eq!(
1137 varbin_view_array.data_buffers().len(),
1138 arrow_array.data_buffers().len()
1139 );
1140 assert_eq!(varbin_view_array.dtype(), &DType::Utf8(true.into()));
1141
1142 let varbin_view_array_non_null = vortex_array_non_null.as_::<VarBinView>();
1143 assert_eq!(
1144 varbin_view_array_non_null.data_buffers().len(),
1145 arrow_array_non_null.data_buffers().len()
1146 );
1147 assert_eq!(
1148 varbin_view_array_non_null.dtype(),
1149 &DType::Utf8(false.into())
1150 );
1151 }
1152
1153 #[test]
1154 fn test_binary_view_array_conversion() {
1155 let mut builder = BinaryViewBuilder::new();
1156 builder.append_value(b"hello");
1157 builder.append_null();
1158 builder.append_value(b"world");
1159 builder.append_value(b"test");
1160 let arrow_array = builder.finish();
1161 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1162
1163 let mut builder_non_null = BinaryViewBuilder::new();
1164 builder_non_null.append_value(b"hello");
1165 builder_non_null.append_value(b"world");
1166 builder_non_null.append_value(b"test");
1167 builder_non_null.append_value(b"vortex");
1168 let arrow_array_non_null = builder_non_null.finish();
1169 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1170
1171 assert_eq!(vortex_array.len(), 4);
1172 assert_eq!(vortex_array_non_null.len(), 4);
1173
1174 let varbin_view_array = vortex_array.as_::<VarBinView>();
1176 assert_eq!(
1177 varbin_view_array.data_buffers().len(),
1178 arrow_array.data_buffers().len()
1179 );
1180 assert_eq!(varbin_view_array.dtype(), &DType::Binary(true.into()));
1181
1182 let varbin_view_array_non_null = vortex_array_non_null.as_::<VarBinView>();
1183 assert_eq!(
1184 varbin_view_array_non_null.data_buffers().len(),
1185 arrow_array_non_null.data_buffers().len()
1186 );
1187 assert_eq!(
1188 varbin_view_array_non_null.dtype(),
1189 &DType::Binary(false.into())
1190 );
1191 }
1192
1193 #[test]
1195 fn test_boolean_array_conversion() {
1196 let arrow_array = BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]);
1197 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1198
1199 let arrow_array_non_null = BooleanArray::from(vec![true, false, true, false]);
1200 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1201
1202 assert_eq!(vortex_array.len(), 4);
1203 assert_eq!(vortex_array_non_null.len(), 4);
1204 }
1205
1206 #[test]
1208 fn test_struct_array_conversion() {
1209 let fields = vec![
1210 Field::new("field1", DataType::Int32, true),
1211 Field::new("field2", DataType::Utf8, false),
1212 ];
1213 let schema = Fields::from(fields);
1214
1215 let field1_data = Int32Array::from(vec![Some(1), None, Some(3)]);
1216 let field2_data = StringArray::from(vec!["a", "b", "c"]);
1217
1218 let arrow_array = StructArray::new(
1219 schema.clone(),
1220 vec![Arc::new(field1_data), Arc::new(field2_data)],
1221 None,
1222 );
1223
1224 let vortex_array = ArrayRef::from_arrow(&arrow_array, false).unwrap();
1225 assert_eq!(vortex_array.len(), 3);
1226
1227 let struct_vortex_array = vortex_array.as_::<Struct>();
1229 assert_eq!(struct_vortex_array.names().len(), 2);
1230 assert_eq!(struct_vortex_array.names()[0], "field1");
1231 assert_eq!(struct_vortex_array.names()[1], "field2");
1232
1233 let nullable_array = StructArray::new(
1235 schema,
1236 vec![
1237 Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])),
1238 Arc::new(StringArray::from(vec!["a", "b", "c"])),
1239 ],
1240 Some(arrow_buffer::NullBuffer::new(BooleanBuffer::from(vec![
1241 true, false, true,
1242 ]))),
1243 );
1244
1245 let vortex_nullable_array = ArrayRef::from_arrow(&nullable_array, true).unwrap();
1246 assert_eq!(vortex_nullable_array.len(), 3);
1247
1248 let struct_vortex_nullable_array = vortex_nullable_array.as_::<Struct>();
1250 assert_eq!(struct_vortex_nullable_array.names().len(), 2);
1251 assert_eq!(struct_vortex_nullable_array.names()[0], "field1");
1252 assert_eq!(struct_vortex_nullable_array.names()[1], "field2");
1253 }
1254
1255 #[test]
1257 fn test_list_array_conversion() {
1258 let mut builder = ListBuilder::new(Int32Builder::new());
1259 builder.append_value([Some(1), None, Some(3)]);
1260 builder.append_null();
1261 builder.append_value([Some(4), Some(5)]);
1262 let arrow_array = builder.finish();
1263
1264 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1265 assert_eq!(vortex_array.len(), 3);
1266
1267 let list_vortex_array = vortex_array.as_::<List>();
1269 let offsets_array = list_vortex_array.offsets().as_::<Primitive>();
1270 assert_eq!(offsets_array.len(), 4); assert_eq!(offsets_array.ptype(), PType::I32);
1272
1273 let mut builder_non_null = ListBuilder::new(Int32Builder::new());
1275 builder_non_null.append_value([Some(1), None, Some(3)]);
1276 builder_non_null.append_value([Some(4), Some(5)]);
1277 let arrow_array_non_null = builder_non_null.finish();
1278
1279 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1280 assert_eq!(vortex_array_non_null.len(), 2);
1281
1282 let list_vortex_array_non_null = vortex_array_non_null.as_::<List>();
1284 let offsets_array_non_null = list_vortex_array_non_null.offsets().as_::<Primitive>();
1285 assert_eq!(offsets_array_non_null.len(), 3); assert_eq!(offsets_array_non_null.ptype(), PType::I32);
1287 }
1288
1289 #[test]
1290 fn test_large_list_array_conversion() {
1291 let mut builder = LargeListBuilder::new(Int32Builder::new());
1292 builder.append_value([Some(1), None, Some(3)]);
1293 builder.append_null();
1294 builder.append_value([Some(4), Some(5)]);
1295 let arrow_array = builder.finish();
1296
1297 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1298 assert_eq!(vortex_array.len(), 3);
1299
1300 let list_vortex_array = vortex_array.as_::<List>();
1302 let offsets_array = list_vortex_array.offsets().as_::<Primitive>();
1303 assert_eq!(offsets_array.len(), 4); assert_eq!(offsets_array.ptype(), PType::I64); let mut builder_non_null = LargeListBuilder::new(Int32Builder::new());
1308 builder_non_null.append_value([Some(1), None, Some(3)]);
1309 builder_non_null.append_value([Some(4), Some(5)]);
1310 let arrow_array_non_null = builder_non_null.finish();
1311
1312 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1313 assert_eq!(vortex_array_non_null.len(), 2);
1314
1315 let list_vortex_array_non_null = vortex_array_non_null.as_::<List>();
1317 let offsets_array_non_null = list_vortex_array_non_null.offsets().as_::<Primitive>();
1318 assert_eq!(offsets_array_non_null.len(), 3); assert_eq!(offsets_array_non_null.ptype(), PType::I64); }
1321
1322 #[test]
1323 fn test_fixed_size_list_array_conversion() {
1324 let values = Int32Array::from(vec![
1326 Some(1),
1327 Some(2),
1328 Some(3), Some(4),
1330 None,
1331 Some(6), Some(7),
1333 Some(8),
1334 Some(9), Some(10),
1336 Some(11),
1337 Some(12), ]);
1339
1340 let field = Arc::new(Field::new("item", DataType::Int32, true));
1342 let arrow_array =
1343 ArrowFixedSizeListArray::try_new(Arc::clone(&field), 3, Arc::new(values), None)
1344 .unwrap();
1345 let vortex_array = ArrayRef::from_arrow(&arrow_array, false).unwrap();
1346
1347 assert_eq!(vortex_array.len(), 4);
1348
1349 let fsl_vortex_array = vortex_array.as_::<FixedSizeList>();
1351 assert_eq!(fsl_vortex_array.list_size(), 3);
1352 assert_eq!(fsl_vortex_array.elements().len(), 12); let values_nullable = Int32Array::from(vec![
1356 Some(1),
1357 Some(2),
1358 Some(3), Some(4),
1360 None,
1361 Some(6), Some(7),
1363 Some(8),
1364 Some(9), ]);
1366
1367 let null_buffer =
1369 arrow_buffer::NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
1370
1371 let arrow_array_nullable = ArrowFixedSizeListArray::try_new(
1372 field,
1373 3,
1374 Arc::new(values_nullable),
1375 Some(null_buffer),
1376 )
1377 .unwrap();
1378 let vortex_array_nullable = ArrayRef::from_arrow(&arrow_array_nullable, true).unwrap();
1379
1380 assert_eq!(vortex_array_nullable.len(), 3);
1381
1382 let fsl_vortex_array_nullable = vortex_array_nullable.as_::<FixedSizeList>();
1384 assert_eq!(fsl_vortex_array_nullable.list_size(), 3);
1385 assert_eq!(fsl_vortex_array_nullable.elements().len(), 9); }
1387
1388 #[test]
1389 fn test_list_view_array_conversion() {
1390 let values = Int32Array::from(vec![
1392 Some(1),
1393 Some(2),
1394 Some(3), Some(4),
1396 Some(5), Some(6), Some(7),
1399 Some(8),
1400 Some(9),
1401 Some(10), ]);
1403
1404 let offsets = ScalarBuffer::from(vec![0i32, 3, 5, 6]);
1406 let sizes = ScalarBuffer::from(vec![3i32, 2, 1, 4]);
1407
1408 let field = Arc::new(Field::new("item", DataType::Int32, true));
1409 let arrow_array = GenericListViewArray::try_new(
1410 Arc::clone(&field),
1411 offsets.clone(),
1412 sizes.clone(),
1413 Arc::new(values.clone()),
1414 None,
1415 )
1416 .unwrap();
1417
1418 let vortex_array = ArrayRef::from_arrow(&arrow_array, false).unwrap();
1419 assert_eq!(vortex_array.len(), 4);
1420
1421 let list_view_vortex_array = vortex_array.as_::<ListView>();
1423 let offsets_array = list_view_vortex_array.offsets().as_::<Primitive>();
1424 let sizes_array = list_view_vortex_array.sizes().as_::<Primitive>();
1425
1426 assert_eq!(offsets_array.len(), 4);
1427 assert_eq!(offsets_array.ptype(), PType::I32);
1428 assert_eq!(sizes_array.len(), 4);
1429 assert_eq!(sizes_array.ptype(), PType::I32);
1430
1431 let null_buffer =
1433 arrow_buffer::NullBuffer::new(BooleanBuffer::from(vec![true, false, true, true]));
1434
1435 let arrow_array_nullable = GenericListViewArray::try_new(
1436 Arc::clone(&field),
1437 offsets,
1438 sizes,
1439 Arc::new(values.clone()),
1440 Some(null_buffer),
1441 )
1442 .unwrap();
1443
1444 let vortex_array_nullable = ArrayRef::from_arrow(&arrow_array_nullable, true).unwrap();
1445 assert_eq!(vortex_array_nullable.len(), 4);
1446
1447 let large_offsets = ScalarBuffer::from(vec![0i64, 3, 5, 6]);
1449 let large_sizes = ScalarBuffer::from(vec![3i64, 2, 1, 4]);
1450
1451 let large_arrow_array = GenericListViewArray::try_new(
1452 field,
1453 large_offsets,
1454 large_sizes,
1455 Arc::new(values),
1456 None,
1457 )
1458 .unwrap();
1459
1460 let large_vortex_array = ArrayRef::from_arrow(&large_arrow_array, false).unwrap();
1461 assert_eq!(large_vortex_array.len(), 4);
1462
1463 let large_list_view_vortex_array = large_vortex_array.as_::<ListView>();
1465 let large_offsets_array = large_list_view_vortex_array.offsets().as_::<Primitive>();
1466 let large_sizes_array = large_list_view_vortex_array.sizes().as_::<Primitive>();
1467
1468 assert_eq!(large_offsets_array.len(), 4);
1469 assert_eq!(large_offsets_array.ptype(), PType::I64); assert_eq!(large_sizes_array.len(), 4);
1471 assert_eq!(large_sizes_array.ptype(), PType::I64); }
1473
1474 #[test]
1476 fn test_null_array_conversion() {
1477 let arrow_array = NullArray::new(5);
1478 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1479 assert_eq!(vortex_array.len(), 5);
1480 }
1481
1482 #[test]
1484 fn test_arrow_buffer_conversion() {
1485 let data = vec![1u8, 2, 3, 4, 5];
1486 let arrow_buffer = ArrowBuffer::from_vec(data);
1487 let vortex_array = arrow_buffer.into_array();
1488 assert_eq!(vortex_array.len(), 5);
1489 }
1490
1491 #[test]
1492 fn test_boolean_buffer_conversion() {
1493 let data = vec![true, false, true, false, true];
1494 let boolean_buffer = BooleanBuffer::from(data);
1495 let vortex_array = boolean_buffer.into_array();
1496 assert_eq!(vortex_array.len(), 5);
1497 }
1498
1499 #[test]
1500 fn test_scalar_buffer_conversion() {
1501 let data = vec![1i32, 2, 3, 4, 5];
1502 let scalar_buffer = ScalarBuffer::from(data);
1503 let vortex_array = scalar_buffer.into_array();
1504 assert_eq!(vortex_array.len(), 5);
1505 }
1506
1507 #[test]
1508 fn test_offset_buffer_conversion() {
1509 let data = vec![0i32, 2, 5, 8, 10];
1510 let offset_buffer = OffsetBuffer::new(ScalarBuffer::from(data));
1511 let vortex_array = offset_buffer.into_array();
1512 assert_eq!(vortex_array.len(), 5);
1513 }
1514
1515 #[test]
1517 fn test_record_batch_conversion() {
1518 let schema = Arc::new(Schema::new(vec![
1519 Field::new("field1", DataType::Int32, false),
1520 Field::new("field2", DataType::Utf8, false),
1521 ]));
1522
1523 let field1_data = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
1524 let field2_data = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
1525
1526 let record_batch = RecordBatch::try_new(schema, vec![field1_data, field2_data]).unwrap();
1527
1528 let vortex_array = ArrayRef::from_arrow(record_batch, false).unwrap();
1529 assert_eq!(vortex_array.len(), 4);
1530
1531 let schema = Arc::new(Schema::new(vec![
1533 Field::new("field1", DataType::Int32, false),
1534 Field::new("field2", DataType::Utf8, false),
1535 ]));
1536
1537 let field1_data = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
1538 let field2_data = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
1539
1540 let record_batch = RecordBatch::try_new(schema, vec![field1_data, field2_data]).unwrap();
1541
1542 let vortex_array = ArrayRef::from_arrow(&record_batch, false).unwrap();
1543 assert_eq!(vortex_array.len(), 4);
1544 }
1545
1546 #[test]
1548 fn test_dyn_array_conversion() {
1549 let int_array = Int32Array::from(vec![1, 2, 3, 4]);
1550 let dyn_array: &dyn ArrowArray = &int_array;
1551 let vortex_array = ArrayRef::from_arrow(dyn_array, false).unwrap();
1552 assert_eq!(vortex_array.len(), 4);
1553
1554 let string_array = StringArray::from(vec!["a", "b", "c"]);
1555 let dyn_array: &dyn ArrowArray = &string_array;
1556 let vortex_array = ArrayRef::from_arrow(dyn_array, false).unwrap();
1557 assert_eq!(vortex_array.len(), 3);
1558
1559 let bool_array = BooleanArray::from(vec![true, false, true]);
1560 let dyn_array: &dyn ArrowArray = &bool_array;
1561 let vortex_array = ArrayRef::from_arrow(dyn_array, false).unwrap();
1562 assert_eq!(vortex_array.len(), 3);
1563 }
1564
1565 #[test]
1567 pub fn nullable_may_contain_non_nullable() {
1568 let null_struct_array_with_non_nullable_field = new_null_array(
1569 &DataType::Struct(Fields::from(vec![Field::new(
1570 "non_nullable_inner",
1571 DataType::Int32,
1572 false,
1573 )])),
1574 1,
1575 );
1576 ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap();
1577 }
1578
1579 #[test]
1580 pub fn nullable_may_contain_deeply_nested_non_nullable() {
1581 let null_struct_array_with_non_nullable_field = new_null_array(
1582 &DataType::Struct(Fields::from(vec![Field::new(
1583 "non_nullable_inner",
1584 DataType::Struct(Fields::from(vec![Field::new(
1585 "non_nullable_deeper_inner",
1586 DataType::Int32,
1587 false,
1588 )])),
1589 false,
1590 )])),
1591 1,
1592 );
1593 ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap();
1594 }
1595
1596 #[test]
1597 fn non_nullable_request_rejects_nulls() {
1598 let arrow_array = Int32Array::from(vec![Some(1), None, Some(3)]);
1601 assert!(ArrayRef::from_arrow(&arrow_array, false).is_err());
1602 }
1603
1604 #[test]
1605 fn non_nullable_request_rejects_null_array() {
1606 let arrow_array = NullArray::new(5);
1609 assert!(ArrayRef::from_arrow(&arrow_array, false).is_err());
1610 }
1611
1612 #[test]
1613 fn non_nullable_struct_with_nulls_errors() {
1614 let struct_array = new_null_array(
1617 &DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])),
1618 3,
1619 );
1620 assert!(ArrayRef::from_arrow(struct_array.as_ref(), false).is_err());
1621 }
1622
1623 #[test]
1624 fn non_nullable_list_with_nulls_errors() {
1625 let mut builder = ListBuilder::new(Int32Builder::new());
1628 builder.append_value([Some(1), Some(2)]);
1629 builder.append_null();
1630 let list = builder.finish();
1631 assert!(ArrayRef::from_arrow(&list, false).is_err());
1632 }
1633
1634 #[test]
1635 pub fn nullable_struct_containing_non_nullable_dictionary_with_nulls_errors() {
1636 let null_struct_array_with_non_nullable_field = new_null_array(
1640 &DataType::Struct(Fields::from(vec![Field::new(
1641 "non_nullable_deeper_inner",
1642 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
1643 false,
1644 )])),
1645 1,
1646 );
1647
1648 assert!(
1649 ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).is_err()
1650 );
1651 }
1652}