16498255725

Committed 24 Jul 2025 01:29PM UTC coverage: 81.085% (-0.006%) from 81.091%

Build # 16498255725

Build Type

Pull #4003

github

Committed by

web-flow

Commit Message

Merge bb86643ac into 3e5a1339f

Pull Request Pull Request #4003: fix: Decimal256 from arrow casts the array to Decimal256Type

Run Details

2 of 3 new or added lines in 1 file covered. (66.67%)

1 existing line in 1 file now uncovered.

42106 of 51928 relevant lines covered (81.09%)

173470.4 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

71.8

/vortex-array/src/arrow/convert.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use arrow_array::array::{
    Array as ArrowArray, ArrowPrimitiveType, BooleanArray as ArrowBooleanArray, GenericByteArray,
    NullArray as ArrowNullArray, OffsetSizeTrait, PrimitiveArray as ArrowPrimitiveArray,
    StructArray as ArrowStructArray,
};
use arrow_array::cast::{AsArray, as_null_array};
use arrow_array::types::{
    ByteArrayType, ByteViewType, Date32Type, Date64Type, Decimal128Type, Decimal256Type,
    Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type,
    Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
    TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
    TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
};
use arrow_array::{GenericByteViewArray, GenericListArray, RecordBatch, make_array};
use arrow_buffer::buffer::{NullBuffer, OffsetBuffer};
use arrow_buffer::{ArrowNativeType, BooleanBuffer, Buffer as ArrowBuffer, ScalarBuffer};
use arrow_schema::{DataType, TimeUnit as ArrowTimeUnit};
use itertools::Itertools;
use vortex_buffer::{Alignment, Buffer, ByteBuffer};
use vortex_dtype::datetime::TimeUnit;
use vortex_dtype::{DType, DecimalDType, NativePType, PType};
use vortex_error::{VortexExpect as _, vortex_panic};
use vortex_scalar::i256;

use crate::arrays::{
    BoolArray, DecimalArray, ListArray, NullArray, PrimitiveArray, StructArray, TemporalArray,
    VarBinArray, VarBinViewArray,
};
use crate::arrow::FromArrowArray;
use crate::validity::Validity;
use crate::{ArrayRef, IntoArray};

impl IntoArray for ArrowBuffer {
    fn into_array(self) -> ArrayRef {
        PrimitiveArray::from_byte_buffer(
            ByteBuffer::from_arrow_buffer(self, Alignment::of::<u8>()),
            PType::U8,
            Validity::NonNullable,
        )
        .into_array()
    }
}

impl IntoArray for BooleanBuffer {
    fn into_array(self) -> ArrayRef {
        BoolArray::new(self, Validity::NonNullable).into_array()
    }
}

impl<T> IntoArray for ScalarBuffer<T>
where
    T: ArrowNativeType + NativePType,
{
    fn into_array(self) -> ArrayRef {
        PrimitiveArray::new(
            Buffer::<T>::from_arrow_scalar_buffer(self),
            Validity::NonNullable,
        )
        .into_array()
    }
}

impl<O> IntoArray for OffsetBuffer<O>
where
    O: NativePType + OffsetSizeTrait,
{
    fn into_array(self) -> ArrayRef {
        let primitive = PrimitiveArray::new(
            Buffer::from_arrow_scalar_buffer(self.into_inner()),
            Validity::NonNullable,
        );

        primitive.into_array()
    }
}

macro_rules! impl_from_arrow_primitive {
    ($ty:path) => {
        impl FromArrowArray<&ArrowPrimitiveArray<$ty>> for ArrayRef {
            fn from_arrow(value: &ArrowPrimitiveArray<$ty>, nullable: bool) -> Self {
                let buffer = Buffer::from_arrow_scalar_buffer(value.values().clone());
                let validity = nulls(value.nulls(), nullable);
                PrimitiveArray::new(buffer, validity).into_array()
            }
        }
    };
}

impl_from_arrow_primitive!(Int8Type);
impl_from_arrow_primitive!(Int16Type);
impl_from_arrow_primitive!(Int32Type);
impl_from_arrow_primitive!(Int64Type);
impl_from_arrow_primitive!(UInt8Type);
impl_from_arrow_primitive!(UInt16Type);
impl_from_arrow_primitive!(UInt32Type);
impl_from_arrow_primitive!(UInt64Type);
impl_from_arrow_primitive!(Float16Type);
impl_from_arrow_primitive!(Float32Type);
impl_from_arrow_primitive!(Float64Type);

impl FromArrowArray<&ArrowPrimitiveArray<Decimal128Type>> for ArrayRef {
    fn from_arrow(array: &ArrowPrimitiveArray<Decimal128Type>, nullable: bool) -> Self {
        let decimal_type = DecimalDType::new(array.precision(), array.scale());
        let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
        let validity = nulls(array.nulls(), nullable);
        DecimalArray::new(buffer, decimal_type, validity).into_array()
    }
}

impl FromArrowArray<&ArrowPrimitiveArray<Decimal256Type>> for ArrayRef {
    fn from_arrow(array: &ArrowPrimitiveArray<Decimal256Type>, nullable: bool) -> Self {
        let decimal_type = DecimalDType::new(array.precision(), array.scale());
        let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());
        // SAFETY: Our i256 implementation has the same bit-pattern representation of the
        //  arrow_buffer::i256 type. It is safe to treat values held inside the buffer as values
        //  of either type.
        let buffer =
            unsafe { std::mem::transmute::<Buffer<arrow_buffer::i256>, Buffer<i256>>(buffer) };
        let validity = nulls(array.nulls(), nullable);
        DecimalArray::new(buffer, decimal_type, validity).into_array()
    }
}

macro_rules! impl_from_arrow_temporal {
    ($ty:path) => {
        impl FromArrowArray<&ArrowPrimitiveArray<$ty>> for ArrayRef {
            fn from_arrow(value: &ArrowPrimitiveArray<$ty>, nullable: bool) -> Self {
                temporal_array(value, nullable)
            }
        }
    };
}

// timestamp
impl_from_arrow_temporal!(TimestampSecondType);
impl_from_arrow_temporal!(TimestampMillisecondType);
impl_from_arrow_temporal!(TimestampMicrosecondType);
impl_from_arrow_temporal!(TimestampNanosecondType);

// time
impl_from_arrow_temporal!(Time32SecondType);
impl_from_arrow_temporal!(Time32MillisecondType);
impl_from_arrow_temporal!(Time64MicrosecondType);
impl_from_arrow_temporal!(Time64NanosecondType);

// date
impl_from_arrow_temporal!(Date32Type);
impl_from_arrow_temporal!(Date64Type);

fn temporal_array<T: ArrowPrimitiveType>(value: &ArrowPrimitiveArray<T>, nullable: bool) -> ArrayRef
where
    T::Native: NativePType,
{
    let arr = PrimitiveArray::new(
        Buffer::from_arrow_scalar_buffer(value.values().clone()),
        nulls(value.nulls(), nullable),
    )
    .into_array();

    match T::DATA_TYPE {
        DataType::Timestamp(time_unit, tz) => {
            let tz = tz.map(|s| s.to_string());
            TemporalArray::new_timestamp(arr, time_unit.into(), tz).into()
        }
        DataType::Time32(time_unit) => TemporalArray::new_time(arr, time_unit.into()).into(),
        DataType::Time64(time_unit) => TemporalArray::new_time(arr, time_unit.into()).into(),
        DataType::Date32 => TemporalArray::new_date(arr, TimeUnit::D).into(),
        DataType::Date64 => TemporalArray::new_date(arr, TimeUnit::Ms).into(),
        DataType::Duration(_) => unimplemented!(),
        DataType::Interval(_) => unimplemented!(),
        _ => vortex_panic!("Invalid temporal type: {}", T::DATA_TYPE),
    }
}

impl<T: ByteArrayType> FromArrowArray<&GenericByteArray<T>> for ArrayRef
where
    <T as ByteArrayType>::Offset: NativePType,
{
    fn from_arrow(value: &GenericByteArray<T>, nullable: bool) -> Self {
        let dtype = match T::DATA_TYPE {
            DataType::Binary | DataType::LargeBinary => DType::Binary(nullable.into()),
            DataType::Utf8 | DataType::LargeUtf8 => DType::Utf8(nullable.into()),
            _ => vortex_panic!("Invalid data type for ByteArray: {}", T::DATA_TYPE),
        };
        VarBinArray::try_new(
            value.offsets().clone().into_array(),
            ByteBuffer::from_arrow_buffer(value.values().clone(), Alignment::of::<u8>()),
            dtype,
            nulls(value.nulls(), nullable),
        )
        .vortex_expect("Failed to convert Arrow GenericByteArray to Vortex VarBinArray")
        .into_array()
    }
}

impl<T: ByteViewType> FromArrowArray<&GenericByteViewArray<T>> for ArrayRef {
    fn from_arrow(value: &GenericByteViewArray<T>, nullable: bool) -> Self {
        let dtype = match T::DATA_TYPE {
            DataType::BinaryView => DType::Binary(nullable.into()),
            DataType::Utf8View => DType::Utf8(nullable.into()),
            _ => vortex_panic!("Invalid data type for ByteViewArray: {}", T::DATA_TYPE),
        };

        let views_buffer = Buffer::from_byte_buffer(
            Buffer::from_arrow_scalar_buffer(value.views().clone()).into_byte_buffer(),
        );

        VarBinViewArray::try_new(
            views_buffer,
            value
                .data_buffers()
                .iter()
                .map(|b| ByteBuffer::from_arrow_buffer(b.clone(), Alignment::of::<u8>()))
                .collect::<Vec<_>>(),
            dtype,
            nulls(value.nulls(), nullable),
        )
        .vortex_expect("Failed to convert Arrow GenericByteViewArray to Vortex VarBinViewArray")
        .into_array()
    }
}

impl FromArrowArray<&ArrowBooleanArray> for ArrayRef {
    fn from_arrow(value: &ArrowBooleanArray, nullable: bool) -> Self {
        BoolArray::new(value.values().clone(), nulls(value.nulls(), nullable)).into_array()
    }
}

/// Strip out the nulls from this array and return a new array without nulls.
fn remove_nulls(data: arrow_data::ArrayData) -> arrow_data::ArrayData {
    if data.null_count() == 0 {
        // No nulls to remove, return the array as is
        return data;
    }

    let children = match data.data_type() {
        DataType::Struct(fields) => Some(
            fields
                .iter()
                .zip(data.child_data().iter())
                .map(|(field, child_data)| {
                    if field.is_nullable() {
                        child_data.clone()
                    } else {
                        remove_nulls(child_data.clone())
                    }
                })
                .collect_vec(),
        ),
        DataType::List(f)
        | DataType::LargeList(f)
        | DataType::ListView(f)
        | DataType::LargeListView(f)
        | DataType::FixedSizeList(f, _)
            if !f.is_nullable() =>
        {
            // All list types only have one child
            assert_eq!(
                data.child_data().len(),
                1,
                "List types should have one child"
            );
            Some(vec![remove_nulls(data.child_data()[0].clone())])
        }
        _ => None,
    };

    let mut builder = data.into_builder().nulls(None);
    if let Some(children) = children {
        builder = builder.child_data(children);
    }
    builder
        .build()
        .vortex_expect("reconstructing array without nulls")
}

impl FromArrowArray<&ArrowStructArray> for ArrayRef {
    fn from_arrow(value: &ArrowStructArray, nullable: bool) -> Self {
        StructArray::try_new(
            value.column_names().iter().copied().collect(),
            value
                .columns()
                .iter()
                .zip(value.fields())
                .map(|(c, field)| {
                    // Arrow pushes down nulls, even into non-nullable fields. So we strip them
                    // out here because Vortex is a little more strict.
                    if c.null_count() > 0 && !field.is_nullable() {
                        let stripped = make_array(remove_nulls(c.into_data()));
                        Self::from_arrow(stripped.as_ref(), false)
                    } else {
                        Self::from_arrow(c.as_ref(), field.is_nullable())
                    }
                })
                .collect(),
            value.len(),
            nulls(value.nulls(), nullable),
        )
        .vortex_expect("Failed to convert Arrow StructArray to Vortex StructArray")
        .into_array()
    }
}

impl<O: OffsetSizeTrait + NativePType> FromArrowArray<&GenericListArray<O>> for ArrayRef {
    fn from_arrow(value: &GenericListArray<O>, nullable: bool) -> Self {
        // Extract the validity of the underlying element array
        let elem_nullable = match value.data_type() {
            DataType::List(field) => field.is_nullable(),
            DataType::LargeList(field) => field.is_nullable(),
            dt => vortex_panic!("Invalid data type for ListArray: {dt}"),
        };
        ListArray::try_new(
            Self::from_arrow(value.values().as_ref(), elem_nullable),
            // offsets are always non-nullable
            value.offsets().clone().into_array(),
            nulls(value.nulls(), nullable),
        )
        .vortex_expect("Failed to convert Arrow StructArray to Vortex StructArray")
        .into_array()
    }
}

impl FromArrowArray<&ArrowNullArray> for ArrayRef {
    fn from_arrow(value: &ArrowNullArray, nullable: bool) -> Self {
        assert!(nullable);
        NullArray::new(value.len()).into_array()
    }
}

fn nulls(nulls: Option<&NullBuffer>, nullable: bool) -> Validity {
    if nullable {
        nulls
            .map(|nulls| {
                if nulls.null_count() == nulls.len() {
                    Validity::AllInvalid
                } else {
                    Validity::from(nulls.inner().clone())
                }
            })
            .unwrap_or_else(|| Validity::AllValid)
    } else {
        assert!(nulls.map(|x| x.null_count() == 0).unwrap_or(true));
        Validity::NonNullable
    }
}

impl FromArrowArray<&dyn ArrowArray> for ArrayRef {
    fn from_arrow(array: &dyn ArrowArray, nullable: bool) -> Self {
        match array.data_type() {
            DataType::Boolean => Self::from_arrow(array.as_boolean(), nullable),
            DataType::UInt8 => Self::from_arrow(array.as_primitive::<UInt8Type>(), nullable),
            DataType::UInt16 => Self::from_arrow(array.as_primitive::<UInt16Type>(), nullable),
            DataType::UInt32 => Self::from_arrow(array.as_primitive::<UInt32Type>(), nullable),
            DataType::UInt64 => Self::from_arrow(array.as_primitive::<UInt64Type>(), nullable),
            DataType::Int8 => Self::from_arrow(array.as_primitive::<Int8Type>(), nullable),
            DataType::Int16 => Self::from_arrow(array.as_primitive::<Int16Type>(), nullable),
            DataType::Int32 => Self::from_arrow(array.as_primitive::<Int32Type>(), nullable),
            DataType::Int64 => Self::from_arrow(array.as_primitive::<Int64Type>(), nullable),
            DataType::Float16 => Self::from_arrow(array.as_primitive::<Float16Type>(), nullable),
            DataType::Float32 => Self::from_arrow(array.as_primitive::<Float32Type>(), nullable),
            DataType::Float64 => Self::from_arrow(array.as_primitive::<Float64Type>(), nullable),
            DataType::Utf8 => Self::from_arrow(array.as_string::<i32>(), nullable),
            DataType::LargeUtf8 => Self::from_arrow(array.as_string::<i64>(), nullable),
            DataType::Binary => Self::from_arrow(array.as_binary::<i32>(), nullable),
            DataType::LargeBinary => Self::from_arrow(array.as_binary::<i64>(), nullable),
            DataType::BinaryView => Self::from_arrow(array.as_binary_view(), nullable),
            DataType::Utf8View => Self::from_arrow(array.as_string_view(), nullable),
            DataType::Struct(_) => Self::from_arrow(array.as_struct(), nullable),
            DataType::List(_) => Self::from_arrow(array.as_list::<i32>(), nullable),
            DataType::LargeList(_) => Self::from_arrow(array.as_list::<i64>(), nullable),
            DataType::Null => Self::from_arrow(as_null_array(array), nullable),
            DataType::Timestamp(u, _) => match u {
                ArrowTimeUnit::Second => {
                    Self::from_arrow(array.as_primitive::<TimestampSecondType>(), nullable)
                }
                ArrowTimeUnit::Millisecond => {
                    Self::from_arrow(array.as_primitive::<TimestampMillisecondType>(), nullable)
                }
                ArrowTimeUnit::Microsecond => {
                    Self::from_arrow(array.as_primitive::<TimestampMicrosecondType>(), nullable)
                }
                ArrowTimeUnit::Nanosecond => {
                    Self::from_arrow(array.as_primitive::<TimestampNanosecondType>(), nullable)
                }
            },
            DataType::Date32 => Self::from_arrow(array.as_primitive::<Date32Type>(), nullable),
            DataType::Date64 => Self::from_arrow(array.as_primitive::<Date64Type>(), nullable),
            DataType::Time32(u) => match u {
                ArrowTimeUnit::Second => {
                    Self::from_arrow(array.as_primitive::<Time32SecondType>(), nullable)
                }
                ArrowTimeUnit::Millisecond => {
                    Self::from_arrow(array.as_primitive::<Time32MillisecondType>(), nullable)
                }
                _ => unreachable!(),
            },
            DataType::Time64(u) => match u {
                ArrowTimeUnit::Microsecond => {
                    Self::from_arrow(array.as_primitive::<Time64MicrosecondType>(), nullable)
                }
                ArrowTimeUnit::Nanosecond => {
                    Self::from_arrow(array.as_primitive::<Time64NanosecondType>(), nullable)
                }
                _ => unreachable!(),
            },
            DataType::Decimal128(..) => {
                Self::from_arrow(array.as_primitive::<Decimal128Type>(), nullable)
            }
            DataType::Decimal256(..) => {
                Self::from_arrow(array.as_primitive::<Decimal256Type>(), nullable)
            }
            _ => vortex_panic!(
                "Array encoding not implemented for Arrow data type {}",
                array.data_type().clone()
            ),
        }
    }
}

impl FromArrowArray<RecordBatch> for ArrayRef {
    fn from_arrow(array: RecordBatch, nullable: bool) -> Self {
        ArrayRef::from_arrow(&arrow_array::StructArray::from(array), nullable)
    }
}

impl FromArrowArray<&RecordBatch> for ArrayRef {
    fn from_arrow(array: &RecordBatch, nullable: bool) -> Self {
        Self::from_arrow(array.clone(), nullable)
    }
}

#[cfg(test)]
mod tests {
    use arrow_array::new_null_array;
    use arrow_schema::{DataType, Field, Fields};

    use crate::ArrayRef;
    use crate::arrow::FromArrowArray as _;

    #[test]
    pub fn nullable_may_contain_non_nullable() {
        let null_struct_array_with_non_nullable_field = new_null_array(
            &DataType::Struct(Fields::from(vec![Field::new(
                "non_nullable_inner",
                DataType::Int32,
                false,
            )])),
            1,
        );
        ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true);
    }

    #[test]
    pub fn nullable_may_contain_deeply_nested_non_nullable() {
        let null_struct_array_with_non_nullable_field = new_null_array(
            &DataType::Struct(Fields::from(vec![Field::new(
                "non_nullable_inner",
                DataType::Struct(Fields::from(vec![Field::new(
                    "non_nullable_deeper_inner",
                    DataType::Int32,
                    false,
                )])),
                false,
            )])),
            1,
        );
        ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true);
    }

    #[test]
    #[should_panic]
    pub fn cannot_handle_nullable_struct_containing_non_nullable_dictionary() {
        let null_struct_array_with_non_nullable_field = new_null_array(
            &DataType::Struct(Fields::from(vec![Field::new(
                "non_nullable_deeper_inner",
                DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
                false,
            )])),
            1,
        );

        ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true);
    }
}

1	// SPDX-License-Identifier: Apache-2.0
2	// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4	use arrow_array::array::{
5	Array as ArrowArray, ArrowPrimitiveType, BooleanArray as ArrowBooleanArray, GenericByteArray,
6	NullArray as ArrowNullArray, OffsetSizeTrait, PrimitiveArray as ArrowPrimitiveArray,
7	StructArray as ArrowStructArray,
8	};
9	use arrow_array::cast::{AsArray, as_null_array};
10	use arrow_array::types::{
11	ByteArrayType, ByteViewType, Date32Type, Date64Type, Decimal128Type, Decimal256Type,
12	Float16Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type,
13	Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType,
14	TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType,
15	TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
16	};
17	use arrow_array::{GenericByteViewArray, GenericListArray, RecordBatch, make_array};
18	use arrow_buffer::buffer::{NullBuffer, OffsetBuffer};
19	use arrow_buffer::{ArrowNativeType, BooleanBuffer, Buffer as ArrowBuffer, ScalarBuffer};
20	use arrow_schema::{DataType, TimeUnit as ArrowTimeUnit};
21	use itertools::Itertools;
22	use vortex_buffer::{Alignment, Buffer, ByteBuffer};
23	use vortex_dtype::datetime::TimeUnit;
24	use vortex_dtype::{DType, DecimalDType, NativePType, PType};
25	use vortex_error::{VortexExpect as _, vortex_panic};
26	use vortex_scalar::i256;
27
28	use crate::arrays::{
29	BoolArray, DecimalArray, ListArray, NullArray, PrimitiveArray, StructArray, TemporalArray,
30	VarBinArray, VarBinViewArray,
31	};
32	use crate::arrow::FromArrowArray;
33	use crate::validity::Validity;
34	use crate::{ArrayRef, IntoArray};
35
36	impl IntoArray for ArrowBuffer {
37	fn into_array(self) -> ArrayRef {	×
38	PrimitiveArray::from_byte_buffer(	×
39	ByteBuffer::from_arrow_buffer(self, Alignment::of::<u8>()),	×
40	PType::U8,	×
41	Validity::NonNullable,	×
42	)	×
43	.into_array()	×
44	}	×
45	}
46
47	impl IntoArray for BooleanBuffer {
48	fn into_array(self) -> ArrayRef {	×
49	BoolArray::new(self, Validity::NonNullable).into_array()	×
50	}	×
51	}
52
53	impl<T> IntoArray for ScalarBuffer<T>
54	where
55	T: ArrowNativeType + NativePType,
56	{
57	fn into_array(self) -> ArrayRef {	×
58	PrimitiveArray::new(	×
59	Buffer::<T>::from_arrow_scalar_buffer(self),	×
60	Validity::NonNullable,	×
61	)	×
62	.into_array()	×
63	}	×
64	}
65
66	impl<O> IntoArray for OffsetBuffer<O>
67	where
68	O: NativePType + OffsetSizeTrait,
69	{
70	fn into_array(self) -> ArrayRef {	39✔
71	let primitive = PrimitiveArray::new(	39✔
72	Buffer::from_arrow_scalar_buffer(self.into_inner()),	39✔
73	Validity::NonNullable,	39✔
74	);
75
76	primitive.into_array()	39✔
77	}	39✔
78	}
79
80	macro_rules! impl_from_arrow_primitive {
81	($ty:path) => {
82	impl FromArrowArray<&ArrowPrimitiveArray<$ty>> for ArrayRef {
83	fn from_arrow(value: &ArrowPrimitiveArray<$ty>, nullable: bool) -> Self {	4,113✔
84	let buffer = Buffer::from_arrow_scalar_buffer(value.values().clone());	4,113✔
85	let validity = nulls(value.nulls(), nullable);	4,113✔
86	PrimitiveArray::new(buffer, validity).into_array()	4,113✔
87	}	4,113✔
88	}
89	};
90	}
91
92	impl_from_arrow_primitive!(Int8Type);
93	impl_from_arrow_primitive!(Int16Type);
94	impl_from_arrow_primitive!(Int32Type);
95	impl_from_arrow_primitive!(Int64Type);
96	impl_from_arrow_primitive!(UInt8Type);
97	impl_from_arrow_primitive!(UInt16Type);
98	impl_from_arrow_primitive!(UInt32Type);
99	impl_from_arrow_primitive!(UInt64Type);
100	impl_from_arrow_primitive!(Float16Type);
101	impl_from_arrow_primitive!(Float32Type);
102	impl_from_arrow_primitive!(Float64Type);
103
104	impl FromArrowArray<&ArrowPrimitiveArray<Decimal128Type>> for ArrayRef {
105	fn from_arrow(array: &ArrowPrimitiveArray<Decimal128Type>, nullable: bool) -> Self {	26✔
106	let decimal_type = DecimalDType::new(array.precision(), array.scale());	26✔
107	let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());	26✔
108	let validity = nulls(array.nulls(), nullable);	26✔
109	DecimalArray::new(buffer, decimal_type, validity).into_array()	26✔
110	}	26✔
111	}
112
113	impl FromArrowArray<&ArrowPrimitiveArray<Decimal256Type>> for ArrayRef {
114	fn from_arrow(array: &ArrowPrimitiveArray<Decimal256Type>, nullable: bool) -> Self {	×
115	let decimal_type = DecimalDType::new(array.precision(), array.scale());	×
116	let buffer = Buffer::from_arrow_scalar_buffer(array.values().clone());	×
117	// SAFETY: Our i256 implementation has the same bit-pattern representation of the
118	// arrow_buffer::i256 type. It is safe to treat values held inside the buffer as values
119	// of either type.
120	let buffer =	×
121	unsafe { std::mem::transmute::<Buffer<arrow_buffer::i256>, Buffer<i256>>(buffer) };	×
122	let validity = nulls(array.nulls(), nullable);	×
123	DecimalArray::new(buffer, decimal_type, validity).into_array()	×
124	}	×
125	}
126
127	macro_rules! impl_from_arrow_temporal {
128	($ty:path) => {
129	impl FromArrowArray<&ArrowPrimitiveArray<$ty>> for ArrayRef {
130	fn from_arrow(value: &ArrowPrimitiveArray<$ty>, nullable: bool) -> Self {	88✔
131	temporal_array(value, nullable)	88✔
132	}	88✔
133	}
134	};
135	}
136
137	// timestamp
138	impl_from_arrow_temporal!(TimestampSecondType);
139	impl_from_arrow_temporal!(TimestampMillisecondType);
140	impl_from_arrow_temporal!(TimestampMicrosecondType);
141	impl_from_arrow_temporal!(TimestampNanosecondType);
142
143	// time
144	impl_from_arrow_temporal!(Time32SecondType);
145	impl_from_arrow_temporal!(Time32MillisecondType);
146	impl_from_arrow_temporal!(Time64MicrosecondType);
147	impl_from_arrow_temporal!(Time64NanosecondType);
148
149	// date
150	impl_from_arrow_temporal!(Date32Type);
151	impl_from_arrow_temporal!(Date64Type);
152
153	fn temporal_array<T: ArrowPrimitiveType>(value: &ArrowPrimitiveArray<T>, nullable: bool) -> ArrayRef	88✔
154	where	88✔
155	T::Native: NativePType,	88✔
156	{
157	let arr = PrimitiveArray::new(	88✔
158	Buffer::from_arrow_scalar_buffer(value.values().clone()),	88✔
159	nulls(value.nulls(), nullable),	88✔
160	)
161	.into_array();	88✔
162
163	match T::DATA_TYPE {	88✔
164	DataType::Timestamp(time_unit, tz) => {	74✔
165	let tz = tz.map(\|s\| s.to_string());	74✔
166	TemporalArray::new_timestamp(arr, time_unit.into(), tz).into()	74✔
167	}
168	DataType::Time32(time_unit) => TemporalArray::new_time(arr, time_unit.into()).into(),	×
169	DataType::Time64(time_unit) => TemporalArray::new_time(arr, time_unit.into()).into(),	×
170	DataType::Date32 => TemporalArray::new_date(arr, TimeUnit::D).into(),	14✔
171	DataType::Date64 => TemporalArray::new_date(arr, TimeUnit::Ms).into(),	×
172	DataType::Duration(_) => unimplemented!(),	×
173	DataType::Interval(_) => unimplemented!(),	×
174	_ => vortex_panic!("Invalid temporal type: {}", T::DATA_TYPE),	×
175	}
176	}	88✔
177
178	impl<T: ByteArrayType> FromArrowArray<&GenericByteArray<T>> for ArrayRef
179	where
180	<T as ByteArrayType>::Offset: NativePType,
181	{
182	fn from_arrow(value: &GenericByteArray<T>, nullable: bool) -> Self {	38✔
183	let dtype = match T::DATA_TYPE {	38✔
184	DataType::Binary \| DataType::LargeBinary => DType::Binary(nullable.into()),	×
185	DataType::Utf8 \| DataType::LargeUtf8 => DType::Utf8(nullable.into()),	38✔
186	_ => vortex_panic!("Invalid data type for ByteArray: {}", T::DATA_TYPE),	×
187	};
188	VarBinArray::try_new(	38✔
189	value.offsets().clone().into_array(),	38✔
190	ByteBuffer::from_arrow_buffer(value.values().clone(), Alignment::of::<u8>()),	38✔
191	dtype,	38✔
192	nulls(value.nulls(), nullable),	38✔
193	)	38✔
194	.vortex_expect("Failed to convert Arrow GenericByteArray to Vortex VarBinArray")	38✔
195	.into_array()	38✔
196	}	38✔
197	}
198
199	impl<T: ByteViewType> FromArrowArray<&GenericByteViewArray<T>> for ArrayRef {
200	fn from_arrow(value: &GenericByteViewArray<T>, nullable: bool) -> Self {	1,975✔
201	let dtype = match T::DATA_TYPE {	1,975✔
202	DataType::BinaryView => DType::Binary(nullable.into()),	41✔
203	DataType::Utf8View => DType::Utf8(nullable.into()),	1,934✔
204	_ => vortex_panic!("Invalid data type for ByteViewArray: {}", T::DATA_TYPE),	×
205	};
206
207	let views_buffer = Buffer::from_byte_buffer(	1,975✔
208	Buffer::from_arrow_scalar_buffer(value.views().clone()).into_byte_buffer(),	1,975✔
209	);
210
211	VarBinViewArray::try_new(	1,975✔
212	views_buffer,	1,975✔
213	value	1,975✔
214	.data_buffers()	1,975✔
215	.iter()	1,975✔
216	.map(\|b\| ByteBuffer::from_arrow_buffer(b.clone(), Alignment::of::<u8>()))	2,149✔
217	.collect::<Vec<_>>(),	1,975✔
218	dtype,	1,975✔
219	nulls(value.nulls(), nullable),	1,975✔
220	)
221	.vortex_expect("Failed to convert Arrow GenericByteViewArray to Vortex VarBinViewArray")	1,975✔
222	.into_array()	1,975✔
223	}	1,975✔
224	}
225
226	impl FromArrowArray<&ArrowBooleanArray> for ArrayRef {
227	fn from_arrow(value: &ArrowBooleanArray, nullable: bool) -> Self {	8,214✔
228	BoolArray::new(value.values().clone(), nulls(value.nulls(), nullable)).into_array()	8,214✔
229	}	8,214✔
230	}
231
232	/// Strip out the nulls from this array and return a new array without nulls.
233	fn remove_nulls(data: arrow_data::ArrayData) -> arrow_data::ArrayData {	4✔
234	if data.null_count() == 0 {	4✔
235	// No nulls to remove, return the array as is
236	return data;	×
237	}	4✔
238
239	let children = match data.data_type() {	4✔
240	DataType::Struct(fields) => Some(	1✔
241	fields	1✔
242	.iter()	1✔
243	.zip(data.child_data().iter())	1✔
244	.map(\|(field, child_data)\| {	1✔
245	if field.is_nullable() {	1✔
246	child_data.clone()	×
247	} else {
248	remove_nulls(child_data.clone())	1✔
249	}
250	})	1✔
251	.collect_vec(),	1✔
252	),
253	DataType::List(f)	×
254	\| DataType::LargeList(f)	×
255	\| DataType::ListView(f)	×
256	\| DataType::LargeListView(f)	×
257	\| DataType::FixedSizeList(f, _)	×
258	if !f.is_nullable() =>	×
259	{
260	// All list types only have one child
261	assert_eq!(	×
262	data.child_data().len(),	×
263	1,
264	"List types should have one child"	×
265	);
266	Some(vec![remove_nulls(data.child_data()[0].clone())])	×
267	}
268	_ => None,	3✔
269	};
270
271	let mut builder = data.into_builder().nulls(None);	4✔
272	if let Some(children) = children {	4✔
273	builder = builder.child_data(children);	1✔
274	}	3✔
275	builder	4✔
276	.build()	4✔
277	.vortex_expect("reconstructing array without nulls")	4✔
278	}	4✔
279
280	impl FromArrowArray<&ArrowStructArray> for ArrayRef {
281	fn from_arrow(value: &ArrowStructArray, nullable: bool) -> Self {	134✔
282	StructArray::try_new(	134✔
283	value.column_names().iter().copied().collect(),	134✔
284	value	134✔
285	.columns()	134✔
286	.iter()	134✔
287	.zip(value.fields())	134✔
288	.map(\|(c, field)\| {	1,011✔
289	// Arrow pushes down nulls, even into non-nullable fields. So we strip them
290	// out here because Vortex is a little more strict.
291	if c.null_count() > 0 && !field.is_nullable() {	1,011✔
292	let stripped = make_array(remove_nulls(c.into_data()));	3✔
293	Self::from_arrow(stripped.as_ref(), false)	3✔
294	} else {
295	Self::from_arrow(c.as_ref(), field.is_nullable())	1,008✔
296	}
297	})	1,011✔
298	.collect(),	134✔
299	value.len(),	134✔
300	nulls(value.nulls(), nullable),	134✔
301	)
302	.vortex_expect("Failed to convert Arrow StructArray to Vortex StructArray")	134✔
303	.into_array()	134✔
304	}	134✔
305	}
306
307	impl<O: OffsetSizeTrait + NativePType> FromArrowArray<&GenericListArray<O>> for ArrayRef {
308	fn from_arrow(value: &GenericListArray<O>, nullable: bool) -> Self {	1✔
309	// Extract the validity of the underlying element array
310	let elem_nullable = match value.data_type() {	1✔
311	DataType::List(field) => field.is_nullable(),	1✔
312	DataType::LargeList(field) => field.is_nullable(),	×
313	dt => vortex_panic!("Invalid data type for ListArray: {dt}"),	×
314	};
315	ListArray::try_new(	1✔
316	Self::from_arrow(value.values().as_ref(), elem_nullable),	1✔
317	// offsets are always non-nullable	1✔
318	value.offsets().clone().into_array(),	1✔
319	nulls(value.nulls(), nullable),	1✔
320	)	1✔
321	.vortex_expect("Failed to convert Arrow StructArray to Vortex StructArray")	1✔
322	.into_array()	1✔
323	}	1✔
324	}
325
326	impl FromArrowArray<&ArrowNullArray> for ArrayRef {
327	fn from_arrow(value: &ArrowNullArray, nullable: bool) -> Self {	×
328	assert!(nullable);	×
329	NullArray::new(value.len()).into_array()	×
330	}	×
331	}
332
333	fn nulls(nulls: Option<&NullBuffer>, nullable: bool) -> Validity {	14,588✔
334	if nullable {	14,588✔
335	nulls	9,617✔
336	.map(\|nulls\| {	9,617✔
337	if nulls.null_count() == nulls.len() {	2,397✔
338	Validity::AllInvalid	5✔
339	} else {
340	Validity::from(nulls.inner().clone())	2,392✔
341	}
342	})	2,397✔
343	.unwrap_or_else(\|\| Validity::AllValid)	9,617✔
344	} else {
345	assert!(nulls.map(\|x\| x.null_count() == 0).unwrap_or(true));	4,971✔
346	Validity::NonNullable	4,971✔
347	}
348	}	14,588✔
349
350	impl FromArrowArray<&dyn ArrowArray> for ArrayRef {
351	fn from_arrow(array: &dyn ArrowArray, nullable: bool) -> Self {	6,244✔
352	match array.data_type() {	6,244✔
353	DataType::Boolean => Self::from_arrow(array.as_boolean(), nullable),	×
354	DataType::UInt8 => Self::from_arrow(array.as_primitive::<UInt8Type>(), nullable),	43✔
355	DataType::UInt16 => Self::from_arrow(array.as_primitive::<UInt16Type>(), nullable),	528✔
356	DataType::UInt32 => Self::from_arrow(array.as_primitive::<UInt32Type>(), nullable),	9✔
357	DataType::UInt64 => Self::from_arrow(array.as_primitive::<UInt64Type>(), nullable),	639✔
358	DataType::Int8 => Self::from_arrow(array.as_primitive::<Int8Type>(), nullable),	6✔
359	DataType::Int16 => Self::from_arrow(array.as_primitive::<Int16Type>(), nullable),	6✔
360	DataType::Int32 => Self::from_arrow(array.as_primitive::<Int32Type>(), nullable),	2,278✔
361	DataType::Int64 => Self::from_arrow(array.as_primitive::<Int64Type>(), nullable),	227✔
362	DataType::Float16 => Self::from_arrow(array.as_primitive::<Float16Type>(), nullable),	2✔
363	DataType::Float32 => Self::from_arrow(array.as_primitive::<Float32Type>(), nullable),	4✔
364	DataType::Float64 => Self::from_arrow(array.as_primitive::<Float64Type>(), nullable),	371✔
365	DataType::Utf8 => Self::from_arrow(array.as_string::<i32>(), nullable),	1✔
366	DataType::LargeUtf8 => Self::from_arrow(array.as_string::<i64>(), nullable),	37✔
367	DataType::Binary => Self::from_arrow(array.as_binary::<i32>(), nullable),	×
368	DataType::LargeBinary => Self::from_arrow(array.as_binary::<i64>(), nullable),	×
369	DataType::BinaryView => Self::from_arrow(array.as_binary_view(), nullable),	41✔
370	DataType::Utf8View => Self::from_arrow(array.as_string_view(), nullable),	1,934✔
371	DataType::Struct(_) => Self::from_arrow(array.as_struct(), nullable),	4✔
372	DataType::List(_) => Self::from_arrow(array.as_list::<i32>(), nullable),	×
373	DataType::LargeList(_) => Self::from_arrow(array.as_list::<i64>(), nullable),	×
374	DataType::Null => Self::from_arrow(as_null_array(array), nullable),	×
375	DataType::Timestamp(u, _) => match u {	74✔
376	ArrowTimeUnit::Second => {
377	Self::from_arrow(array.as_primitive::<TimestampSecondType>(), nullable)	×
378	}
379	ArrowTimeUnit::Millisecond => {
380	Self::from_arrow(array.as_primitive::<TimestampMillisecondType>(), nullable)	×
381	}
382	ArrowTimeUnit::Microsecond => {
383	Self::from_arrow(array.as_primitive::<TimestampMicrosecondType>(), nullable)	74✔
384	}
385	ArrowTimeUnit::Nanosecond => {
386	Self::from_arrow(array.as_primitive::<TimestampNanosecondType>(), nullable)	×
387	}
388	},
389	DataType::Date32 => Self::from_arrow(array.as_primitive::<Date32Type>(), nullable),	14✔
390	DataType::Date64 => Self::from_arrow(array.as_primitive::<Date64Type>(), nullable),	×
391	DataType::Time32(u) => match u {	×
392	ArrowTimeUnit::Second => {
393	Self::from_arrow(array.as_primitive::<Time32SecondType>(), nullable)	×
394	}
395	ArrowTimeUnit::Millisecond => {
396	Self::from_arrow(array.as_primitive::<Time32MillisecondType>(), nullable)	×
397	}
398	_ => unreachable!(),	×
399	},
400	DataType::Time64(u) => match u {	×
401	ArrowTimeUnit::Microsecond => {
402	Self::from_arrow(array.as_primitive::<Time64MicrosecondType>(), nullable)	×
403	}
404	ArrowTimeUnit::Nanosecond => {
405	Self::from_arrow(array.as_primitive::<Time64NanosecondType>(), nullable)	×
406	}
407	_ => unreachable!(),	×
408	},
409	DataType::Decimal128(..) => {
410	Self::from_arrow(array.as_primitive::<Decimal128Type>(), nullable)	26✔
411	}
412	DataType::Decimal256(..) => {
NEW 413	Self::from_arrow(array.as_primitive::<Decimal256Type>(), nullable)	×
414	}
415	_ => vortex_panic!(	×
416	"Array encoding not implemented for Arrow data type {}",	×
417	array.data_type().clone()	×
418	),
419	}
420	}	6,244✔
421	}
422
423	impl FromArrowArray<RecordBatch> for ArrayRef {
424	fn from_arrow(array: RecordBatch, nullable: bool) -> Self {	129✔
425	ArrayRef::from_arrow(&arrow_array::StructArray::from(array), nullable)	129✔
426	}	129✔
427	}
428
429	impl FromArrowArray<&RecordBatch> for ArrayRef {
430	fn from_arrow(array: &RecordBatch, nullable: bool) -> Self {	18✔
431	Self::from_arrow(array.clone(), nullable)	18✔
432	}	18✔
433	}
434
435	#[cfg(test)]
436	mod tests {
437	use arrow_array::new_null_array;
438	use arrow_schema::{DataType, Field, Fields};
439
440	use crate::ArrayRef;
441	use crate::arrow::FromArrowArray as _;
442
443	#[test]
444	pub fn nullable_may_contain_non_nullable() {	1✔
445	let null_struct_array_with_non_nullable_field = new_null_array(	1✔
446	&DataType::Struct(Fields::from(vec![Field::new(	1✔
447	"non_nullable_inner",	1✔
448	DataType::Int32,	1✔
449	false,	1✔
450	)])),	1✔
451	1,
452	);
453	ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true);	1✔
454	}	1✔
455
456	#[test]
457	pub fn nullable_may_contain_deeply_nested_non_nullable() {	1✔
458	let null_struct_array_with_non_nullable_field = new_null_array(	1✔
459	&DataType::Struct(Fields::from(vec![Field::new(	1✔
460	"non_nullable_inner",	1✔
461	DataType::Struct(Fields::from(vec![Field::new(	1✔
462	"non_nullable_deeper_inner",	1✔
463	DataType::Int32,	1✔
464	false,	1✔
465	)])),	1✔
466	false,	1✔
467	)])),	1✔
468	1,
469	);
470	ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true);	1✔
471	}	1✔
472
473	#[test]
474	#[should_panic]
475	pub fn cannot_handle_nullable_struct_containing_non_nullable_dictionary() {	1✔
476	let null_struct_array_with_non_nullable_field = new_null_array(	1✔
477	&DataType::Struct(Fields::from(vec![Field::new(	1✔
478	"non_nullable_deeper_inner",	1✔
479	DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),	1✔
480	false,	1✔
481	)])),	1✔
482	1,
483	);
484
485	ArrayRef::from_arrow(null_struct_array_with_non_nullable_field.as_ref(), true);	1✔
486	}	1✔
487	}

vortex-data / vortex / 16498255725

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous