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

vortex_array/arrow/
datum.rs
