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    pub fn with_target_datatype(
36        array: &dyn Array,
37        target_datatype: &DataType,
38    ) -> VortexResult<Self> {
39        if array.is_constant() {
40            Ok(Self {
41                array: array.slice(0, 1).into_arrow(target_datatype)?,
42                is_scalar: true,
43            })
44        } else {
45            Ok(Self {
46                array: array.to_array().into_arrow(target_datatype)?,
47                is_scalar: false,
48            })
49        }
50    }
51}
52
53impl ArrowDatum for Datum {
54    fn get(&self) -> (&dyn ArrowArray, bool) {
55        (&self.array, self.is_scalar)
56    }
57}
58
59/// Convert an Arrow array to an Array with a specific length.
60/// This is useful for compute functions that delegate to Arrow using [Datum],
61/// which will return a scalar (length 1 Arrow array) if the input array is constant.
62///
63/// # Error
64///
65/// The provided array must have length
66pub fn from_arrow_array_with_len<A>(array: A, len: usize, nullable: bool) -> ArrayRef
67where
68    ArrayRef: FromArrowArray<A>,
69{
70    let array = ArrayRef::from_arrow(array, nullable);
71    if array.len() == len {
72        return array;
73    }
74
75    if array.len() != 1 {
76        vortex_panic!(
77            "Array length mismatch, expected {} got {} for encoding {}",
78            len,
79            array.len(),
80            array.encoding_id()
81        );
82    }
83
84    ConstantArray::new(array.scalar_at(0), len).into_array()
85}
vortex_array/arrow/datum.rs

vortex_array/arrow/
datum.rs
