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