vortex_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 arrow_schema::Field;
9use vortex_array::ArrayRef;
10use vortex_array::ExecutionCtx;
11use vortex_array::IntoArray;
12use vortex_array::arrays::Constant;
13use vortex_array::arrays::ConstantArray;
14use vortex_error::VortexResult;
15use vortex_error::vortex_panic;
16
17use crate::ArrowSessionExt;
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, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
29        let session = ctx.session().clone();
30        if array.is::<Constant>() {
31            Ok(Self {
32                array: session
33                    .arrow()
34                    .execute_arrow(array.slice(0..1)?, None, ctx)?,
35                is_scalar: true,
36            })
37        } else {
38            Ok(Self {
39                array: session.arrow().execute_arrow(array.clone(), None, ctx)?,
40                is_scalar: false,
41            })
42        }
43    }
44
45    /// Create a new [`Datum`] from an `DynArrayData`, 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        let session = ctx.session().clone();
49        Ok(Self {
50            array: session.arrow().execute_arrow(array.clone(), None, ctx)?,
51            is_scalar: false,
52        })
53    }
54
55    pub fn try_new_with_target_datatype(
56        array: &ArrayRef,
57        target_datatype: &DataType,
58        ctx: &mut ExecutionCtx,
59    ) -> VortexResult<Self> {
60        let session = ctx.session().clone();
61        let target_field = Field::new(
62            String::new(),
63            target_datatype.clone(),
64            array.dtype().is_nullable(),
65        );
66
67        if array.is::<Constant>() {
68            Ok(Self {
69                array: session.arrow().execute_arrow(
70                    array.slice(0..1)?,
71                    Some(&target_field),
72                    ctx,
73                )?,
74                is_scalar: true,
75            })
76        } else {
77            Ok(Self {
78                array: session
79                    .arrow()
80                    .execute_arrow(array.clone(), Some(&target_field), ctx)?,
81                is_scalar: false,
82            })
83        }
84    }
85
86    pub fn data_type(&self) -> &DataType {
87        self.array.data_type()
88    }
89}
90
91impl ArrowDatum for Datum {
92    fn get(&self) -> (&dyn ArrowArray, bool) {
93        (&self.array, self.is_scalar)
94    }
95}
96
97/// Convert an Arrow array to an Array with a specific length, using the provided
98/// [`ExecutionCtx`].
99///
100/// This is useful for compute functions that delegate to Arrow using [Datum],
101/// which will return a scalar (length 1 Arrow array) if the input array is constant.
102///
103/// The array is imported through the [`ArrowSession`](crate::ArrowSession) of `ctx`, so nested
104/// extension fields reach their registered import plugins.
105///
106/// # Error
107///
108/// The provided array must have length `len` or `1`.
109pub fn from_arrow_columnar(
110    array: ArrowArrayRef,
111    len: usize,
112    nullable: bool,
113    ctx: &mut ExecutionCtx,
114) -> VortexResult<ArrayRef> {
115    let session = ctx.session().clone();
116    let array = session.arrow().from_arrow_array(array, nullable)?;
117    if array.len() == len {
118        return Ok(array);
119    }
120
121    if array.len() != 1 {
122        vortex_panic!(
123            "Array length mismatch, expected {} got {} for encoding {}",
124            len,
125            array.len(),
126            array.encoding_id()
127        );
128    }
129
130    Ok(ConstantArray::new(array.execute_scalar(0, ctx)?, len).into_array())
131}
vortex_arrow/datum.rs

vortex_arrow/
datum.rs
