vortex_array/arrow/
datum.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_array::Array as ArrowArray;
5use arrow_array::ArrayRef as ArrowArrayRef;
6use arrow_array::Datum as ArrowDatum;
7use arrow_schema::DataType;
8use vortex_error::VortexExpect;
9use vortex_error::VortexResult;
10use vortex_error::vortex_panic;
11
12use crate::Array;
13use crate::ArrayRef;
14use crate::IntoArray;
15use crate::arrays::ConstantArray;
16use crate::arrays::ConstantVTable;
17use crate::arrow::FromArrowArray;
18use crate::arrow::IntoArrowArray;
19
20/// A wrapper around a generic Arrow array that can be used as a Datum in Arrow compute.
21#[derive(Debug)]
22pub struct Datum {
23    array: ArrowArrayRef,
24    is_scalar: bool,
25}
26
27impl Datum {
28    /// Create a new [`Datum`] from an [`ArrayRef`], which can then be passed to Arrow compute.
29    pub fn try_new(array: &dyn Array) -> VortexResult<Self> {
30        if array.is::<ConstantVTable>() {
31            Ok(Self {
32                array: array.slice(0..1)?.into_arrow_preferred()?,
33                is_scalar: true,
34            })
35        } else {
36            Ok(Self {
37                array: array.to_array().into_arrow_preferred()?,
38                is_scalar: false,
39            })
40        }
41    }
42
43    /// Create a new [`Datum`] from an [`Array`], which can then be passed to Arrow compute.
44    /// This not try and convert the array to a scalar if it is constant.
45    pub fn try_new_array(array: &dyn Array) -> VortexResult<Self> {
46        Ok(Self {
47            array: array.to_array().into_arrow_preferred()?,
48            is_scalar: false,
49        })
50    }
51
52    pub fn try_new_with_target_datatype(
53        array: &dyn Array,
54        target_datatype: &DataType,
55    ) -> VortexResult<Self> {
56        if array.is::<ConstantVTable>() {
57            Ok(Self {
58                array: array.slice(0..1)?.into_arrow(target_datatype)?,
59                is_scalar: true,
60            })
61        } else {
62            Ok(Self {
63                array: array.to_array().into_arrow(target_datatype)?,
64                is_scalar: false,
65            })
66        }
67    }
68
69    pub fn data_type(&self) -> &DataType {
70        self.array.data_type()
71    }
72}
73
74impl ArrowDatum for Datum {
75    fn get(&self) -> (&dyn ArrowArray, bool) {
76        (&self.array, self.is_scalar)
77    }
78}
79
80/// Convert an Arrow array to an Array with a specific length.
81/// This is useful for compute functions that delegate to Arrow using [Datum],
82/// which will return a scalar (length 1 Arrow array) if the input array is constant.
83///
84/// # Error
85///
86/// The provided array must have length
87pub fn from_arrow_array_with_len<A>(array: A, len: usize, nullable: bool) -> VortexResult<ArrayRef>
88where
89    ArrayRef: FromArrowArray<A>,
90{
91    let array = ArrayRef::from_arrow(array, nullable)?;
92    if array.len() == len {
93        return Ok(array);
94    }
95
96    if array.len() != 1 {
97        vortex_panic!(
98            "Array length mismatch, expected {} got {} for encoding {}",
99            len,
100            array.len(),
101            array.encoding_id()
102        );
103    }
104
105    Ok(ConstantArray::new(
106        array
107            .scalar_at(0)
108            .vortex_expect("array of length 1 must support scalar_at(0)"),
109        len,
110    )
111    .into_array())
112}
vortex_array/arrow/datum.rs

vortex_array/arrow/
datum.rs
