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

vortex_array/arrow/
datum.rs
