vortex_array/compute/
scalar_at.rs

1use vortex_error::{VortexExpect, VortexResult, vortex_bail, vortex_err};
2use vortex_scalar::Scalar;
3
4use crate::Array;
5use crate::encoding::Encoding;
6
7/// Implementation of scalar_at for an encoding.
8///
9/// SAFETY: the index is guaranteed to be within the bounds of the [`crate::ArrayRef`].
10pub trait ScalarAtFn<A> {
11    fn scalar_at(&self, array: A, index: usize) -> VortexResult<Scalar>;
12}
13
14impl<E: Encoding> ScalarAtFn<&dyn Array> for E
15where
16    E: for<'a> ScalarAtFn<&'a E::Array>,
17{
18    fn scalar_at(&self, array: &dyn Array, index: usize) -> VortexResult<Scalar> {
19        let array_ref = array
20            .as_any()
21            .downcast_ref::<E::Array>()
22            .vortex_expect("Failed to downcast array");
23        ScalarAtFn::scalar_at(self, array_ref, index)
24    }
25}
26
27pub fn scalar_at(array: &dyn Array, index: usize) -> VortexResult<Scalar> {
28    if index >= array.len() {
29        vortex_bail!(OutOfBounds: index, 0, array.len());
30    }
31
32    if !array.is_valid(index)? {
33        return Ok(Scalar::null(array.dtype().clone()));
34    }
35
36    let scalar = array
37        .vtable()
38        .scalar_at_fn()
39        .map(|f| f.scalar_at(array, index))
40        .unwrap_or_else(|| Err(vortex_err!(NotImplemented: "scalar_at", array.encoding())))?;
41
42    debug_assert_eq!(
43        scalar.dtype(),
44        array.dtype(),
45        "ScalarAt dtype mismatch {}",
46        array.encoding()
47    );
48
49    Ok(scalar)
50}