vortex_array/arrow/
datum.rs

1use arrow_array::{Array as ArrowArray, ArrayRef as ArrowArrayRef, Datum as ArrowDatum};
2use arrow_schema::DataType;
3use vortex_error::{VortexResult, vortex_panic};
4
5use crate::arrays::ConstantArray;
6use crate::arrow::{FromArrowArray, IntoArrowArray};
7use crate::compute::{scalar_at, slice};
8use crate::{Array, ArrayRef};
9
10/// A wrapper around a generic Arrow array that can be used as a Datum in Arrow compute.
11#[derive(Debug)]
12pub struct Datum {
13    array: ArrowArrayRef,
14    is_scalar: bool,
15}
16
17impl Datum {
18    /// Create a new [`Datum`] from an [`ArrayRef`], which can then be passed to Arrow compute.
19    pub fn try_new(array: &dyn Array) -> VortexResult<Self> {
20        if array.is_constant() {
21            Ok(Self {
22                array: slice(array, 0, 1)?.into_arrow_preferred()?,
23                is_scalar: true,
24            })
25        } else {
26            Ok(Self {
27                array: array.to_array().into_arrow_preferred()?,
28                is_scalar: false,
29            })
30        }
31    }
32
33    pub fn with_target_datatype(
34        array: &dyn Array,
35        target_datatype: &DataType,
36    ) -> VortexResult<Self> {
37        if array.is_constant() {
38            Ok(Self {
39                array: slice(array, 0, 1)?.into_arrow(target_datatype)?,
40                is_scalar: true,
41            })
42        } else {
43            Ok(Self {
44                array: array.to_array().into_arrow(target_datatype)?,
45                is_scalar: false,
46            })
47        }
48    }
49}
50
51impl ArrowDatum for Datum {
52    fn get(&self) -> (&dyn ArrowArray, bool) {
53        (&self.array, self.is_scalar)
54    }
55}
56
57/// Convert an Arrow array to an Array with a specific length.
58/// This is useful for compute functions that delegate to Arrow using [Datum],
59/// which will return a scalar (length 1 Arrow array) if the input array is constant.
60///
61/// Panics if the length of the array is not 1 and also not equal to the expected length.
62pub fn from_arrow_array_with_len<A>(array: A, len: usize, nullable: bool) -> VortexResult<ArrayRef>
63where
64    ArrayRef: FromArrowArray<A>,
65{
66    let array = ArrayRef::from_arrow(array, nullable);
67    if array.len() == len {
68        return Ok(array);
69    }
70
71    if array.len() != 1 {
72        vortex_panic!(
73            "Array length mismatch, expected {} got {} for encoding {}",
74            len,
75            array.len(),
76            array.encoding()
77        );
78    }
79
80    Ok(ConstantArray::new(scalar_at(&array, 0)?, len).into_array())
81}
vortex_array/arrow/datum.rs

vortex_array/arrow/
datum.rs
