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::VortexSessionExecute;
13use vortex_array::arrays::Constant;
14use vortex_array::arrays::ConstantArray;
15use vortex_array::legacy_session;
16use vortex_error::VortexExpect;
17use vortex_error::VortexResult;
18use vortex_error::vortex_panic;
19
20use crate::ArrowSessionExt;
21#[allow(deprecated)]
22use crate::FromArrowArray;
23
24/// A wrapper around a generic Arrow array that can be used as a Datum in Arrow compute.
25#[derive(Debug)]
26pub struct Datum {
27    array: ArrowArrayRef,
28    is_scalar: bool,
29}
30
31impl Datum {
32    /// Create a new [`Datum`] from an [`ArrayRef`], which can then be passed to Arrow compute.
33    pub fn try_new(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
34        let session = ctx.session().clone();
35        if array.is::<Constant>() {
36            Ok(Self {
37                array: session
38                    .arrow()
39                    .execute_arrow(array.slice(0..1)?, None, ctx)?,
40                is_scalar: true,
41            })
42        } else {
43            Ok(Self {
44                array: session.arrow().execute_arrow(array.clone(), None, ctx)?,
45                is_scalar: false,
46            })
47        }
48    }
49
50    /// Create a new [`Datum`] from an `DynArrayData`, which can then be passed to Arrow compute.
51    /// This not try and convert the array to a scalar if it is constant.
52    pub fn try_new_array(array: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<Self> {
53        let session = ctx.session().clone();
54        Ok(Self {
55            array: session.arrow().execute_arrow(array.clone(), None, ctx)?,
56            is_scalar: false,
57        })
58    }
59
60    pub fn try_new_with_target_datatype(
61        array: &ArrayRef,
62        target_datatype: &DataType,
63        ctx: &mut ExecutionCtx,
64    ) -> VortexResult<Self> {
65        let session = ctx.session().clone();
66        let target_field = Field::new(
67            String::new(),
68            target_datatype.clone(),
69            array.dtype().is_nullable(),
70        );
71
72        if array.is::<Constant>() {
73            Ok(Self {
74                array: session.arrow().execute_arrow(
75                    array.slice(0..1)?,
76                    Some(&target_field),
77                    ctx,
78                )?,
79                is_scalar: true,
80            })
81        } else {
82            Ok(Self {
83                array: session
84                    .arrow()
85                    .execute_arrow(array.clone(), Some(&target_field), ctx)?,
86                is_scalar: false,
87            })
88        }
89    }
90
91    pub fn data_type(&self) -> &DataType {
92        self.array.data_type()
93    }
94}
95
96impl ArrowDatum for Datum {
97    fn get(&self) -> (&dyn ArrowArray, bool) {
98        (&self.array, self.is_scalar)
99    }
100}
101
102/// Convert an Arrow array to an Array with a specific length.
103/// This is useful for compute functions that delegate to Arrow using [Datum],
104/// which will return a scalar (length 1 Arrow array) if the input array is constant.
105///
106/// # Error
107///
108/// The provided array must have length
109#[deprecated(
110    note = "Relies on the hidden global `legacy_session()`; use `from_arrow_columnar` with an explicit `ExecutionCtx` instead"
111)]
112#[allow(clippy::disallowed_methods)]
113#[allow(deprecated)]
114pub fn from_arrow_array_with_len<A>(array: A, len: usize, nullable: bool) -> VortexResult<ArrayRef>
115where
116    ArrayRef: FromArrowArray<A>,
117{
118    let array = ArrayRef::from_arrow(array, nullable)?;
119    if array.len() == len {
120        return Ok(array);
121    }
122
123    if array.len() != 1 {
124        vortex_panic!(
125            "Array length mismatch, expected {} got {} for encoding {}",
126            len,
127            array.len(),
128            array.encoding_id()
129        );
130    }
131
132    Ok(ConstantArray::new(
133        array
134            .execute_scalar(0, &mut legacy_session().create_execution_ctx())
135            .vortex_expect("array of length 1 must support execute_scalar(0)"),
136        len,
137    )
138    .into_array())
139}
140
141/// Convert an Arrow array to an Array with a specific length, using the provided
142/// [`ExecutionCtx`].
143///
144/// This is useful for compute functions that delegate to Arrow using [Datum],
145/// which will return a scalar (length 1 Arrow array) if the input array is constant.
146///
147/// The array is imported through the [`ArrowSession`](crate::ArrowSession) of `ctx`, so nested
148/// extension fields reach their registered import plugins.
149///
150/// # Error
151///
152/// The provided array must have length `len` or `1`.
153pub fn from_arrow_columnar(
154    array: ArrowArrayRef,
155    len: usize,
156    nullable: bool,
157    ctx: &mut ExecutionCtx,
158) -> VortexResult<ArrayRef> {
159    let session = ctx.session().clone();
160    let array = session.arrow().from_arrow_array(array, nullable)?;
161    if array.len() == len {
162        return Ok(array);
163    }
164
165    if array.len() != 1 {
166        vortex_panic!(
167            "Array length mismatch, expected {} got {} for encoding {}",
168            len,
169            array.len(),
170            array.encoding_id()
171        );
172    }
173
174    Ok(ConstantArray::new(array.execute_scalar(0, ctx)?, len).into_array())
175}
vortex_arrow/datum.rs

vortex_arrow/
datum.rs
