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

vortex_array/arrow/
datum.rs
