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

vortex_array/arrow/
datum.rs
