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

vortex_array/arrow/
datum.rs
