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::fixed_size_list::FixedSizeListArraySlotsExt;
720 use vortex_array::arrays::list::ListArraySlotsExt;
721 use vortex_array::arrays::listview::ListViewArraySlotsExt;
722 use vortex_array::arrays::struct_::StructArrayExt;
723 use vortex_array::dtype::DType;
724 use vortex_array::dtype::Nullability;
725 use vortex_array::dtype::PType;
726 use vortex_array::extension::datetime::TimeUnit;
727 use vortex_array::extension::datetime::Timestamp;
728
729 use crate::FromArrowArray as _;
730 use crate::IntoVortexArray as _;
731
732 #[rstest]
733 #[case::i8(
734 Arc::new(Int8Array::from(vec![Some(1), None, Some(3), Some(4)])),
735 Arc::new(Int8Array::from(vec![1, 2, 3, 4])),
736 PType::I8,
737 )]
738 #[case::i16(
739 Arc::new(Int16Array::from(vec![Some(100), None, Some(300), Some(400)])),
740 Arc::new(Int16Array::from(vec![100, 200, 300, 400])),
741 PType::I16,
742 )]
743 #[case::i32(
744 Arc::new(Int32Array::from(vec![Some(1000), None, Some(3000), Some(4000)])),
745 Arc::new(Int32Array::from(vec![1000, 2000, 3000, 4000])),
746 PType::I32,
747 )]
748 #[case::i64(
749 Arc::new(Int64Array::from(vec![Some(10000), None, Some(30000), Some(40000)])),
750 Arc::new(Int64Array::from(vec![10000_i64, 20000, 30000, 40000])),
751 PType::I64,
752 )]
753 #[case::u8(
754 Arc::new(UInt8Array::from(vec![Some(1), None, Some(3), Some(4)])),
755 Arc::new(UInt8Array::from(vec![1_u8, 2, 3, 4])),
756 PType::U8,
757 )]
758 #[case::u16(
759 Arc::new(UInt16Array::from(vec![Some(100), None, Some(300), Some(400)])),
760 Arc::new(UInt16Array::from(vec![100_u16, 200, 300, 400])),
761 PType::U16,
762 )]
763 #[case::u32(
764 Arc::new(UInt32Array::from(vec![Some(1000), None, Some(3000), Some(4000)])),
765 Arc::new(UInt32Array::from(vec![1000_u32, 2000, 3000, 4000])),
766 PType::U32,
767 )]
768 #[case::u64(
769 Arc::new(UInt64Array::from(vec![Some(10000), None, Some(30000), Some(40000)])),
770 Arc::new(UInt64Array::from(vec![10000_u64, 20000, 30000, 40000])),
771 PType::U64,
772 )]
773 #[case::f32(
774 Arc::new(Float32Array::from(vec![Some(1.5), None, Some(3.5), Some(4.5)])),
775 Arc::new(Float32Array::from(vec![1.5_f32, 2.5, 3.5, 4.5])),
776 PType::F32,
777 )]
778 #[case::f64(
779 Arc::new(Float64Array::from(vec![Some(1.5), None, Some(3.5), Some(4.5)])),
780 Arc::new(Float64Array::from(vec![1.5_f64, 2.5, 3.5, 4.5])),
781 PType::F64,
782 )]
783 fn test_primitive_array_conversion(
784 #[case] nullable: Arc<dyn ArrowArray>,
785 #[case] non_nullable: Arc<dyn ArrowArray>,
786 #[case] expected_ptype: PType,
787 ) {
788 let v_null = ArrayRef::from_arrow(nullable.as_ref(), true).unwrap();
789 let v_non_null = ArrayRef::from_arrow(non_nullable.as_ref(), false).unwrap();
790 assert_eq!(v_null.len(), 4);
791 assert_eq!(v_non_null.len(), 4);
792 assert_eq!(v_null.as_::<Primitive>().ptype(), expected_ptype);
793 assert_eq!(v_non_null.as_::<Primitive>().ptype(), expected_ptype);
794 }
795
796 #[test]
797 fn test_float16_array_conversion() {
798 let values = vec![
799 Some(<Float16Type as ArrowPrimitiveType>::Native::from_f32(1.5)),
800 None,
801 Some(<Float16Type as ArrowPrimitiveType>::Native::from_f32(3.5)),
802 ];
803 let arrow_array = arrow_array::PrimitiveArray::<Float16Type>::from(values);
804 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
805
806 let non_null_values = vec![
807 <Float16Type as ArrowPrimitiveType>::Native::from_f32(1.5),
808 <Float16Type as ArrowPrimitiveType>::Native::from_f32(2.5),
809 ];
810 let arrow_array_non_null =
811 arrow_array::PrimitiveArray::<Float16Type>::from(non_null_values);
812 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
813
814 assert_eq!(vortex_array.len(), 3);
815 assert_eq!(vortex_array_non_null.len(), 2);
816
817 let primitive_array = vortex_array.as_::<Primitive>();
819 assert_eq!(primitive_array.ptype(), PType::F16);
820
821 let primitive_array_non_null = vortex_array_non_null.as_::<Primitive>();
822 assert_eq!(primitive_array_non_null.ptype(), PType::F16);
823 }
824
825 #[test]
827 fn test_decimal128_array_conversion() {
828 let mut builder = Decimal128Builder::with_capacity(4);
829 builder.append_value(12345);
830 builder.append_null();
831 builder.append_value(67890);
832 builder.append_value(11111);
833 let decimal_array = builder.finish().with_precision_and_scale(10, 2).unwrap();
834
835 let vortex_array = ArrayRef::from_arrow(&decimal_array, true).unwrap();
836 assert_eq!(vortex_array.len(), 4);
837
838 let mut builder_non_null = Decimal128Builder::with_capacity(3);
839 builder_non_null.append_value(12345);
840 builder_non_null.append_value(67890);
841 builder_non_null.append_value(11111);
842 let decimal_array_non_null = builder_non_null
843 .finish()
844 .with_precision_and_scale(10, 2)
845 .unwrap();
846
847 let vortex_array_non_null = ArrayRef::from_arrow(&decimal_array_non_null, false).unwrap();
848 assert_eq!(vortex_array_non_null.len(), 3);
849
850 let decimal_vortex_array = vortex_array.as_::<Decimal>();
852 assert_eq!(decimal_vortex_array.decimal_dtype().precision(), 10);
853 assert_eq!(decimal_vortex_array.decimal_dtype().scale(), 2);
854
855 let decimal_vortex_array_non_null = vortex_array_non_null.as_::<Decimal>();
856 assert_eq!(
857 decimal_vortex_array_non_null.decimal_dtype().precision(),
858 10
859 );
860 assert_eq!(decimal_vortex_array_non_null.decimal_dtype().scale(), 2);
861 }
862
863 #[test]
864 fn test_decimal256_array_conversion() {
865 let mut builder = Decimal256Builder::with_capacity(4);
866 builder.append_value(arrow_buffer::i256::from_i128(12345));
867 builder.append_null();
868 builder.append_value(arrow_buffer::i256::from_i128(67890));
869 builder.append_value(arrow_buffer::i256::from_i128(11111));
870 let decimal_array = builder.finish().with_precision_and_scale(38, 10).unwrap();
871
872 let vortex_array = ArrayRef::from_arrow(&decimal_array, true).unwrap();
873 assert_eq!(vortex_array.len(), 4);
874
875 let mut builder_non_null = Decimal256Builder::with_capacity(3);
876 builder_non_null.append_value(arrow_buffer::i256::from_i128(12345));
877 builder_non_null.append_value(arrow_buffer::i256::from_i128(67890));
878 builder_non_null.append_value(arrow_buffer::i256::from_i128(11111));
879 let decimal_array_non_null = builder_non_null
880 .finish()
881 .with_precision_and_scale(38, 10)
882 .unwrap();
883
884 let vortex_array_non_null = ArrayRef::from_arrow(&decimal_array_non_null, false).unwrap();
885 assert_eq!(vortex_array_non_null.len(), 3);
886
887 let decimal_vortex_array = vortex_array.as_::<Decimal>();
889 assert_eq!(decimal_vortex_array.decimal_dtype().precision(), 38);
890 assert_eq!(decimal_vortex_array.decimal_dtype().scale(), 10);
891
892 let decimal_vortex_array_non_null = vortex_array_non_null.as_::<Decimal>();
893 assert_eq!(
894 decimal_vortex_array_non_null.decimal_dtype().precision(),
895 38
896 );
897 assert_eq!(decimal_vortex_array_non_null.decimal_dtype().scale(), 10);
898 }
899
900 #[rstest]
902 #[case::timestamp_second(
903 Arc::new(TimestampSecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
904 Arc::new(TimestampSecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
905 )]
906 #[case::timestamp_millisecond(
907 Arc::new(TimestampMillisecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
908 Arc::new(TimestampMillisecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
909 )]
910 #[case::timestamp_microsecond(
911 Arc::new(TimestampMicrosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
912 Arc::new(TimestampMicrosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
913 )]
914 #[case::timestamp_nanosecond(
915 Arc::new(TimestampNanosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
916 Arc::new(TimestampNanosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
917 )]
918 #[case::time32_second(
919 Arc::new(Time32SecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
920 Arc::new(Time32SecondArray::from(vec![1000_i32, 2000, 3000, 4000])),
921 )]
922 #[case::time32_millisecond(
923 Arc::new(Time32MillisecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
924 Arc::new(Time32MillisecondArray::from(vec![1000_i32, 2000, 3000, 4000])),
925 )]
926 #[case::time64_microsecond(
927 Arc::new(Time64MicrosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
928 Arc::new(Time64MicrosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
929 )]
930 #[case::time64_nanosecond(
931 Arc::new(Time64NanosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])),
932 Arc::new(Time64NanosecondArray::from(vec![1000_i64, 2000, 3000, 4000])),
933 )]
934 #[case::date32(
935 Arc::new(Date32Array::from(vec![Some(18000), None, Some(18002), Some(18003)])),
936 Arc::new(Date32Array::from(vec![18000_i32, 18001, 18002, 18003])),
937 )]
938 #[case::date64(
939 Arc::new(Date64Array::from(vec![Some(1555200000000), None, Some(1555286400000), Some(1555372800000)]
940 )),
941 Arc::new(Date64Array::from(vec![1555200000000_i64, 1555213600000, 1555286400000, 1555372800000]
942 )),
943 )]
944 fn test_temporal_array_conversion(
945 #[case] nullable: Arc<dyn ArrowArray>,
946 #[case] non_nullable: Arc<dyn ArrowArray>,
947 ) {
948 let v_null = ArrayRef::from_arrow(nullable.as_ref(), true).unwrap();
949 let v_non_null = ArrayRef::from_arrow(non_nullable.as_ref(), false).unwrap();
950 assert_eq!(v_null.len(), 4);
951 assert_eq!(v_non_null.len(), 4);
952 }
953
954 #[test]
955 fn test_timestamp_timezone_microsecond_array_conversion() {
956 let arrow_array =
957 TimestampMicrosecondArray::from(vec![Some(1000), None, Some(3000), Some(4000)])
958 .with_timezone("UTC");
959 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
960
961 let arrow_array_non_null =
962 TimestampMicrosecondArray::from(vec![1000_i64, 2000, 3000, 4000]).with_timezone("UTC");
963 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
964
965 assert_eq!(vortex_array.len(), 4);
966 assert_eq!(
967 vortex_array.dtype(),
968 &DType::Extension(
969 Timestamp::new_with_tz(
970 TimeUnit::Microseconds,
971 Some("UTC".into()),
972 Nullability::Nullable
973 )
974 .erased()
975 ),
976 );
977 assert_eq!(vortex_array_non_null.len(), 4);
978 assert_eq!(
979 vortex_array_non_null.dtype(),
980 &DType::Extension(
981 Timestamp::new_with_tz(
982 TimeUnit::Microseconds,
983 Some("UTC".into()),
984 Nullability::NonNullable
985 )
986 .erased()
987 )
988 );
989 }
990
991 #[rstest]
993 #[case::utf8(
994 Arc::new(StringArray::from(vec![Some("hello"), None, Some("world"), Some("test")])),
995 Arc::new(StringArray::from(vec!["hello", "world", "test", "vortex"])),
996 DType::Utf8(Nullability::NonNullable),
997 )]
998 #[case::large_utf8(
999 Arc::new(LargeStringArray::from(vec![Some("hello"), None, Some("world"), Some("test")])),
1000 Arc::new(LargeStringArray::from(vec!["hello", "world", "test", "vortex"])),
1001 DType::Utf8(Nullability::NonNullable),
1002 )]
1003 #[case::binary(
1004 Arc::new(BinaryArray::from(vec![
1005 Some("hello".as_bytes()), None, Some("world".as_bytes()), Some("test".as_bytes()),
1006 ])),
1007 Arc::new(BinaryArray::from(vec![
1008 "hello".as_bytes(), "world".as_bytes(), "test".as_bytes(), "vortex".as_bytes(),
1009 ])),
1010 DType::Binary(Nullability::NonNullable),
1011 )]
1012 #[case::large_binary(
1013 Arc::new(LargeBinaryArray::from(vec![
1014 Some("hello".as_bytes()), None, Some("world".as_bytes()), Some("test".as_bytes()),
1015 ])),
1016 Arc::new(LargeBinaryArray::from(vec![
1017 "hello".as_bytes(), "world".as_bytes(), "test".as_bytes(), "vortex".as_bytes(),
1018 ])),
1019 DType::Binary(Nullability::NonNullable),
1020 )]
1021 fn test_string_binary_array_conversion(
1022 #[case] nullable: Arc<dyn ArrowArray>,
1023 #[case] non_nullable: Arc<dyn ArrowArray>,
1024 #[case] expected_non_nullable_dtype: DType,
1025 ) {
1026 let v_null = ArrayRef::from_arrow(nullable.as_ref(), true).unwrap();
1027 let v_non_null = ArrayRef::from_arrow(non_nullable.as_ref(), false).unwrap();
1028 assert_eq!(v_null.len(), 4);
1029 assert_eq!(v_non_null.len(), 4);
1030 assert_eq!(v_null.dtype(), &expected_non_nullable_dtype.as_nullable());
1031 assert_eq!(v_non_null.dtype(), &expected_non_nullable_dtype);
1032 }
1033
1034 #[test]
1035 fn test_utf8_view_array_conversion() {
1036 let mut builder = StringViewBuilder::new();
1037 builder.append_value("hello");
1038 builder.append_null();
1039 builder.append_value("world");
1040 builder.append_value("test");
1041 let arrow_array = builder.finish();
1042 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1043
1044 let mut builder_non_null = StringViewBuilder::new();
1045 builder_non_null.append_value("hello");
1046 builder_non_null.append_value("world");
1047 builder_non_null.append_value("test");
1048 builder_non_null.append_value("vortex");
1049 let arrow_array_non_null = builder_non_null.finish();
1050 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1051
1052 assert_eq!(vortex_array.len(), 4);
1053 assert_eq!(vortex_array_non_null.len(), 4);
1054
1055 let varbin_view_array = vortex_array.as_::<VarBinView>();
1057 assert_eq!(
1058 varbin_view_array.data_buffers().len(),
1059 arrow_array.data_buffers().len()
1060 );
1061 assert_eq!(varbin_view_array.dtype(), &DType::Utf8(true.into()));
1062
1063 let varbin_view_array_non_null = vortex_array_non_null.as_::<VarBinView>();
1064 assert_eq!(
1065 varbin_view_array_non_null.data_buffers().len(),
1066 arrow_array_non_null.data_buffers().len()
1067 );
1068 assert_eq!(
1069 varbin_view_array_non_null.dtype(),
1070 &DType::Utf8(false.into())
1071 );
1072 }
1073
1074 #[test]
1075 fn test_binary_view_array_conversion() {
1076 let mut builder = BinaryViewBuilder::new();
1077 builder.append_value(b"hello");
1078 builder.append_null();
1079 builder.append_value(b"world");
1080 builder.append_value(b"test");
1081 let arrow_array = builder.finish();
1082 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1083
1084 let mut builder_non_null = BinaryViewBuilder::new();
1085 builder_non_null.append_value(b"hello");
1086 builder_non_null.append_value(b"world");
1087 builder_non_null.append_value(b"test");
1088 builder_non_null.append_value(b"vortex");
1089 let arrow_array_non_null = builder_non_null.finish();
1090 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1091
1092 assert_eq!(vortex_array.len(), 4);
1093 assert_eq!(vortex_array_non_null.len(), 4);
1094
1095 let varbin_view_array = vortex_array.as_::<VarBinView>();
1097 assert_eq!(
1098 varbin_view_array.data_buffers().len(),
1099 arrow_array.data_buffers().len()
1100 );
1101 assert_eq!(varbin_view_array.dtype(), &DType::Binary(true.into()));
1102
1103 let varbin_view_array_non_null = vortex_array_non_null.as_::<VarBinView>();
1104 assert_eq!(
1105 varbin_view_array_non_null.data_buffers().len(),
1106 arrow_array_non_null.data_buffers().len()
1107 );
1108 assert_eq!(
1109 varbin_view_array_non_null.dtype(),
1110 &DType::Binary(false.into())
1111 );
1112 }
1113
1114 #[test]
1116 fn test_boolean_array_conversion() {
1117 let arrow_array = BooleanArray::from(vec![Some(true), None, Some(false), Some(true)]);
1118 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1119
1120 let arrow_array_non_null = BooleanArray::from(vec![true, false, true, false]);
1121 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1122
1123 assert_eq!(vortex_array.len(), 4);
1124 assert_eq!(vortex_array_non_null.len(), 4);
1125 }
1126
1127 #[test]
1129 fn test_struct_array_conversion() {
1130 let fields = vec![
1131 Field::new("field1", DataType::Int32, true),
1132 Field::new("field2", DataType::Utf8, false),
1133 ];
1134 let schema = Fields::from(fields);
1135
1136 let field1_data = Int32Array::from(vec![Some(1), None, Some(3)]);
1137 let field2_data = StringArray::from(vec!["a", "b", "c"]);
1138
1139 let arrow_array = StructArray::new(
1140 schema.clone(),
1141 vec![Arc::new(field1_data), Arc::new(field2_data)],
1142 None,
1143 );
1144
1145 let vortex_array = ArrayRef::from_arrow(&arrow_array, false).unwrap();
1146 assert_eq!(vortex_array.len(), 3);
1147
1148 let struct_vortex_array = vortex_array.as_::<Struct>();
1150 assert_eq!(struct_vortex_array.names().len(), 2);
1151 assert_eq!(struct_vortex_array.names()[0], "field1");
1152 assert_eq!(struct_vortex_array.names()[1], "field2");
1153
1154 let nullable_array = StructArray::new(
1156 schema,
1157 vec![
1158 Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])),
1159 Arc::new(StringArray::from(vec!["a", "b", "c"])),
1160 ],
1161 Some(arrow_buffer::NullBuffer::new(BooleanBuffer::from(vec![
1162 true, false, true,
1163 ]))),
1164 );
1165
1166 let vortex_nullable_array = ArrayRef::from_arrow(&nullable_array, true).unwrap();
1167 assert_eq!(vortex_nullable_array.len(), 3);
1168
1169 let struct_vortex_nullable_array = vortex_nullable_array.as_::<Struct>();
1171 assert_eq!(struct_vortex_nullable_array.names().len(), 2);
1172 assert_eq!(struct_vortex_nullable_array.names()[0], "field1");
1173 assert_eq!(struct_vortex_nullable_array.names()[1], "field2");
1174 }
1175
1176 #[test]
1178 fn test_list_array_conversion() {
1179 let mut builder = ListBuilder::new(Int32Builder::new());
1180 builder.append_value([Some(1), None, Some(3)]);
1181 builder.append_null();
1182 builder.append_value([Some(4), Some(5)]);
1183 let arrow_array = builder.finish();
1184
1185 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1186 assert_eq!(vortex_array.len(), 3);
1187
1188 let list_vortex_array = vortex_array.as_::<List>();
1190 let offsets_array = list_vortex_array.offsets().as_::<Primitive>();
1191 assert_eq!(offsets_array.len(), 4); assert_eq!(offsets_array.ptype(), PType::I32);
1193
1194 let mut builder_non_null = ListBuilder::new(Int32Builder::new());
1196 builder_non_null.append_value([Some(1), None, Some(3)]);
1197 builder_non_null.append_value([Some(4), Some(5)]);
1198 let arrow_array_non_null = builder_non_null.finish();
1199
1200 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1201 assert_eq!(vortex_array_non_null.len(), 2);
1202
1203 let list_vortex_array_non_null = vortex_array_non_null.as_::<List>();
1205 let offsets_array_non_null = list_vortex_array_non_null.offsets().as_::<Primitive>();
1206 assert_eq!(offsets_array_non_null.len(), 3); assert_eq!(offsets_array_non_null.ptype(), PType::I32);
1208 }
1209
1210 #[test]
1211 fn test_large_list_array_conversion() {
1212 let mut builder = LargeListBuilder::new(Int32Builder::new());
1213 builder.append_value([Some(1), None, Some(3)]);
1214 builder.append_null();
1215 builder.append_value([Some(4), Some(5)]);
1216 let arrow_array = builder.finish();
1217
1218 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1219 assert_eq!(vortex_array.len(), 3);
1220
1221 let list_vortex_array = vortex_array.as_::<List>();
1223 let offsets_array = list_vortex_array.offsets().as_::<Primitive>();
1224 assert_eq!(offsets_array.len(), 4); assert_eq!(offsets_array.ptype(), PType::I64); let mut builder_non_null = LargeListBuilder::new(Int32Builder::new());
1229 builder_non_null.append_value([Some(1), None, Some(3)]);
1230 builder_non_null.append_value([Some(4), Some(5)]);
1231 let arrow_array_non_null = builder_non_null.finish();
1232
1233 let vortex_array_non_null = ArrayRef::from_arrow(&arrow_array_non_null, false).unwrap();
1234 assert_eq!(vortex_array_non_null.len(), 2);
1235
1236 let list_vortex_array_non_null = vortex_array_non_null.as_::<List>();
1238 let offsets_array_non_null = list_vortex_array_non_null.offsets().as_::<Primitive>();
1239 assert_eq!(offsets_array_non_null.len(), 3); assert_eq!(offsets_array_non_null.ptype(), PType::I64); }
1242
1243 #[test]
1244 fn test_fixed_size_list_array_conversion() {
1245 let values = Int32Array::from(vec![
1247 Some(1),
1248 Some(2),
1249 Some(3), Some(4),
1251 None,
1252 Some(6), Some(7),
1254 Some(8),
1255 Some(9), Some(10),
1257 Some(11),
1258 Some(12), ]);
1260
1261 let field = Arc::new(Field::new("item", DataType::Int32, true));
1263 let arrow_array =
1264 ArrowFixedSizeListArray::try_new(Arc::clone(&field), 3, Arc::new(values), None)
1265 .unwrap();
1266 let vortex_array = ArrayRef::from_arrow(&arrow_array, false).unwrap();
1267
1268 assert_eq!(vortex_array.len(), 4);
1269
1270 let fsl_vortex_array = vortex_array.as_::<FixedSizeList>();
1272 assert_eq!(fsl_vortex_array.list_size(), 3);
1273 assert_eq!(fsl_vortex_array.elements().len(), 12); let values_nullable = Int32Array::from(vec![
1277 Some(1),
1278 Some(2),
1279 Some(3), Some(4),
1281 None,
1282 Some(6), Some(7),
1284 Some(8),
1285 Some(9), ]);
1287
1288 let null_buffer =
1290 arrow_buffer::NullBuffer::new(BooleanBuffer::from(vec![true, false, true]));
1291
1292 let arrow_array_nullable = ArrowFixedSizeListArray::try_new(
1293 field,
1294 3,
1295 Arc::new(values_nullable),
1296 Some(null_buffer),
1297 )
1298 .unwrap();
1299 let vortex_array_nullable = ArrayRef::from_arrow(&arrow_array_nullable, true).unwrap();
1300
1301 assert_eq!(vortex_array_nullable.len(), 3);
1302
1303 let fsl_vortex_array_nullable = vortex_array_nullable.as_::<FixedSizeList>();
1305 assert_eq!(fsl_vortex_array_nullable.list_size(), 3);
1306 assert_eq!(fsl_vortex_array_nullable.elements().len(), 9); }
1308
1309 #[test]
1310 fn test_list_view_array_conversion() {
1311 let values = Int32Array::from(vec![
1313 Some(1),
1314 Some(2),
1315 Some(3), Some(4),
1317 Some(5), Some(6), Some(7),
1320 Some(8),
1321 Some(9),
1322 Some(10), ]);
1324
1325 let offsets = ScalarBuffer::from(vec![0i32, 3, 5, 6]);
1327 let sizes = ScalarBuffer::from(vec![3i32, 2, 1, 4]);
1328
1329 let field = Arc::new(Field::new("item", DataType::Int32, true));
1330 let arrow_array = GenericListViewArray::try_new(
1331 Arc::clone(&field),
1332 offsets.clone(),
1333 sizes.clone(),
1334 Arc::new(values.clone()),
1335 None,
1336 )
1337 .unwrap();
1338
1339 let vortex_array = ArrayRef::from_arrow(&arrow_array, false).unwrap();
1340 assert_eq!(vortex_array.len(), 4);
1341
1342 let list_view_vortex_array = vortex_array.as_::<ListView>();
1344 let offsets_array = list_view_vortex_array.offsets().as_::<Primitive>();
1345 let sizes_array = list_view_vortex_array.sizes().as_::<Primitive>();
1346
1347 assert_eq!(offsets_array.len(), 4);
1348 assert_eq!(offsets_array.ptype(), PType::I32);
1349 assert_eq!(sizes_array.len(), 4);
1350 assert_eq!(sizes_array.ptype(), PType::I32);
1351
1352 let null_buffer =
1354 arrow_buffer::NullBuffer::new(BooleanBuffer::from(vec![true, false, true, true]));
1355
1356 let arrow_array_nullable = GenericListViewArray::try_new(
1357 Arc::clone(&field),
1358 offsets,
1359 sizes,
1360 Arc::new(values.clone()),
1361 Some(null_buffer),
1362 )
1363 .unwrap();
1364
1365 let vortex_array_nullable = ArrayRef::from_arrow(&arrow_array_nullable, true).unwrap();
1366 assert_eq!(vortex_array_nullable.len(), 4);
1367
1368 let large_offsets = ScalarBuffer::from(vec![0i64, 3, 5, 6]);
1370 let large_sizes = ScalarBuffer::from(vec![3i64, 2, 1, 4]);
1371
1372 let large_arrow_array = GenericListViewArray::try_new(
1373 field,
1374 large_offsets,
1375 large_sizes,
1376 Arc::new(values),
1377 None,
1378 )
1379 .unwrap();
1380
1381 let large_vortex_array = ArrayRef::from_arrow(&large_arrow_array, false).unwrap();
1382 assert_eq!(large_vortex_array.len(), 4);
1383
1384 let large_list_view_vortex_array = large_vortex_array.as_::<ListView>();
1386 let large_offsets_array = large_list_view_vortex_array.offsets().as_::<Primitive>();
1387 let large_sizes_array = large_list_view_vortex_array.sizes().as_::<Primitive>();
1388
1389 assert_eq!(large_offsets_array.len(), 4);
1390 assert_eq!(large_offsets_array.ptype(), PType::I64); assert_eq!(large_sizes_array.len(), 4);
1392 assert_eq!(large_sizes_array.ptype(), PType::I64); }
1394
1395 #[test]
1397 fn test_null_array_conversion() {
1398 let arrow_array = NullArray::new(5);
1399 let vortex_array = ArrayRef::from_arrow(&arrow_array, true).unwrap();
1400 assert_eq!(vortex_array.len(), 5);
1401 }
1402
1403 #[test]
1405 fn test_arrow_buffer_conversion() {
1406 let data = vec![1u8, 2, 3, 4, 5];
1407 let arrow_buffer = ArrowBuffer::from_vec(data);
1408 let vortex_array = arrow_buffer.into_array();
1409 assert_eq!(vortex_array.len(), 5);
1410 }
1411
1412 #[test]
1413 fn test_boolean_buffer_conversion() {
1414 let data = vec![true, false, true, false, true];
1415 let boolean_buffer = BooleanBuffer::from(data);
1416 let vortex_array = boolean_buffer.into_array();
1417 assert_eq!(vortex_array.len(), 5);
1418 }
1419
1420 #[test]
1421 fn test_scalar_buffer_conversion() {
1422 let data = vec![1i32, 2, 3, 4, 5];
1423 let scalar_buffer = ScalarBuffer::from(data);
1424 let vortex_array = scalar_buffer.into_array();
1425 assert_eq!(vortex_array.len(), 5);
1426 }
1427
1428 #[test]
1429 fn test_offset_buffer_conversion() {
1430 let data = vec![0i32, 2, 5, 8, 10];
1431 let offset_buffer = OffsetBuffer::new(ScalarBuffer::from(data));
1432 let vortex_array = offset_buffer.into_array();
1433 assert_eq!(vortex_array.len(), 5);
1434 }
1435
1436 #[test]
1438 fn test_record_batch_conversion() {
1439 let schema = Arc::new(Schema::new(vec![
1440 Field::new("field1", DataType::Int32, false),
1441 Field::new("field2", DataType::Utf8, false),
1442 ]));
1443
1444 let field1_data = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
1445 let field2_data = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
1446
1447 let record_batch = RecordBatch::try_new(schema, vec![field1_data, field2_data]).unwrap();
1448
1449 let vortex_array = ArrayRef::from_arrow(record_batch, false).unwrap();
1450 assert_eq!(vortex_array.len(), 4);
1451
1452 let schema = Arc::new(Schema::new(vec![
1454 Field::new("field1", DataType::Int32, false),
1455 Field::new("field2", DataType::Utf8, false),
1456 ]));
1457
1458 let field1_data = Arc::new(Int32Array::from(vec![1, 2, 3, 4]));
1459 let field2_data = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
1460
1461 let record_batch = RecordBatch::try_new(schema, vec![field1_data, field2_data]).unwrap();
1462
1463 let vortex_array = ArrayRef::from_arrow(&record_batch, false).unwrap();
1464 assert_eq!(vortex_array.len(), 4);
1465 }
1466
1467 #[test]
1469 fn test_dyn_array_conversion() {
1470 let int_array = Int32Array::from(vec![1, 2, 3, 4]);
1471 let dyn_array: &dyn ArrowArray = &int_array;
1472 let vortex_array = ArrayRef::from_arrow(dyn_array, false).unwrap();
1473 assert_eq!(vortex_array.len(), 4);
1474
1475 let string_array = StringArray::from(vec!["a", "b", "c"]);
1476 let dyn_array: &dyn ArrowArray = &string_array;
1477 let vortex_array = ArrayRef::from_arrow(dyn_array, false).unwrap();
1478 assert_eq!(vortex_array.len(), 3);
1479
1480 let bool_array = BooleanArray::from(vec![true, false, true]);
1481 let dyn_array: &dyn ArrowArray = &bool_array;
1482 let vortex_array = ArrayRef::from_arrow(dyn_array, false).unwrap();
1483 assert_eq!(vortex_array.len(), 3);
1484 }
1485
1486 #[test]
1488 pub fn nullable_may_contain_non_nullable() {
1489 let null_struct_array_with_non_nullable_field = new_null_array(
1490 &DataType::Struct(Fields::from(vec![Field::new(
1491 "non_nullable_inner",
1492 DataType::Int32,
1493 false,
1494 )])),
1495 1,
1496 );
1497 ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap();
1498 }
1499
1500 #[test]
1501 pub fn nullable_may_contain_deeply_nested_non_nullable() {
1502 let null_struct_array_with_non_nullable_field = new_null_array(
1503 &DataType::Struct(Fields::from(vec![Field::new(
1504 "non_nullable_inner",
1505 DataType::Struct(Fields::from(vec![Field::new(
1506 "non_nullable_deeper_inner",
1507 DataType::Int32,
1508 false,
1509 )])),
1510 false,
1511 )])),
1512 1,
1513 );
1514 ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).unwrap();
1515 }
1516
1517 #[test]
1518 fn non_nullable_request_rejects_nulls() {
1519 let arrow_array = Int32Array::from(vec![Some(1), None, Some(3)]);
1522 assert!(ArrayRef::from_arrow(&arrow_array, false).is_err());
1523 }
1524
1525 #[test]
1526 fn non_nullable_request_rejects_null_array() {
1527 let arrow_array = NullArray::new(5);
1530 assert!(ArrayRef::from_arrow(&arrow_array, false).is_err());
1531 }
1532
1533 #[test]
1534 fn non_nullable_struct_with_nulls_errors() {
1535 let struct_array = new_null_array(
1538 &DataType::Struct(Fields::from(vec![Field::new("a", DataType::Int32, true)])),
1539 3,
1540 );
1541 assert!(ArrayRef::from_arrow(struct_array.as_ref(), false).is_err());
1542 }
1543
1544 #[test]
1545 fn non_nullable_list_with_nulls_errors() {
1546 let mut builder = ListBuilder::new(Int32Builder::new());
1549 builder.append_value([Some(1), Some(2)]);
1550 builder.append_null();
1551 let list = builder.finish();
1552 assert!(ArrayRef::from_arrow(&list, false).is_err());
1553 }
1554
1555 #[test]
1556 pub fn nullable_struct_containing_non_nullable_dictionary_with_nulls_errors() {
1557 let null_struct_array_with_non_nullable_field = new_null_array(
1561 &DataType::Struct(Fields::from(vec![Field::new(
1562 "non_nullable_deeper_inner",
1563 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
1564 false,
1565 )])),
1566 1,
1567 );
1568
1569 assert!(
1570 ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true).is_err()
1571 );
1572 }
1573}