vortex_array/search_sorted/

primitive.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright the Vortex contributors

use std::cell::RefCell;
use std::cmp::Ordering;
use std::marker::PhantomData;

use vortex_error::VortexResult;

use crate::ArrayRef;
use crate::ExecutionCtx;
use crate::RepeatedArrayProbe;
use crate::arrays::Primitive;
use crate::arrays::primitive::PrimitiveArrayExt;
use crate::dtype::NativePType;
use crate::search_sorted::IndexOrd;

/// A [`SearchSorted`](crate::search_sorted::SearchSorted) adapter over a sorted primitive-typed
/// array, comparing elements as the native type `T`.
///
/// Reads go through a [`RepeatedArrayProbe`], so the encoding state and validity resolved on the
/// first comparison are reused by the rest of the search rather than rebuilt per probe.
///
/// Values can be searched as `T`, `Option<T>`, or `usize`. Searching as `T` or `usize` treats
/// null elements as `T::zero()`; use `Option<T>` when the array may contain nulls, in which case
/// nulls sort before all non-null values.
pub struct SearchSortedPrimitiveArray<'a, T> {
    reader: Reader<'a, T>,
    len: usize,
    ctx: RefCell<&'a mut ExecutionCtx>,
    _ptype: PhantomData<T>,
}

/// Where a comparison reads its element from.
enum Reader<'a, T> {
    /// The array's own buffer. Only for a canonical, non-nullable, host-backed array, where
    /// every element is a `T` at a known offset and none of them is null.
    Values(&'a [T]),
    /// Anything else, through one probe reused by the whole search.
    Probe(RefCell<RepeatedArrayProbe>),
}

impl<'a, T: NativePType> SearchSortedPrimitiveArray<'a, T> {
    /// Wraps `array` for searching, panicking if the array's [`PType`](crate::dtype::PType) is
    /// not `T::PTYPE`.
    pub fn new(array: &'a ArrayRef, ctx: &'a mut ExecutionCtx) -> Self {
        assert_eq!(
            array.dtype().as_ptype(),
            T::PTYPE,
            "Array PType must match primitive type"
        );
        Self {
            reader: Self::reader(array),
            len: array.len(),
            ctx: RefCell::new(ctx),
            _ptype: PhantomData,
        }
    }

    /// Reads the buffer directly when the array can offer one, and probes otherwise.
    ///
    /// A nullable array is excluded because the buffer holds an arbitrary value where an element
    /// is null, and a device-backed array because its buffer is not addressable here.
    fn reader(array: &'a ArrayRef) -> Reader<'a, T> {
        if !array.dtype().is_nullable()
            && let Some(primitive) = array.as_opt::<Primitive>()
            && primitive.buffer_handle().as_host_opt().is_some()
            && let Some(values) = primitive.data().as_slice::<T>().get(..array.len())
        {
            return Reader::Values(values);
        }
        Reader::Probe(RefCell::new(array.repeated_probe()))
    }

    /// Returns the value at `idx`, or `None` if the element is null.
    ///
    /// The probe and the context are separate cells, so the two borrows never overlap.
    fn typed_value(&self, idx: usize) -> VortexResult<Option<T>> {
        let probe = match &self.reader {
            Reader::Values(values) => return Ok(Some(values[idx])),
            Reader::Probe(probe) => probe,
        };
        Ok(probe
            .borrow_mut()
            .execute_scalar(idx, &mut self.ctx.borrow_mut())?
            .as_primitive()
            .typed_value::<T>())
    }

    /// Returns the value at `idx`, with nulls mapped to `T::zero()`.
    fn value(&self, idx: usize) -> VortexResult<T> {
        Ok(self.typed_value(idx)?.unwrap_or_else(|| T::zero()))
    }
}

impl<T: NativePType> IndexOrd<T> for SearchSortedPrimitiveArray<'_, T> {
    fn index_cmp(&self, idx: usize, elem: &T) -> VortexResult<Option<Ordering>> {
        let value = self.value(idx)?;
        Ok(Some(value.total_compare(*elem)))
    }

    fn index_len(&self) -> usize {
        self.len
    }
}

impl<T: NativePType> IndexOrd<Option<T>> for SearchSortedPrimitiveArray<'_, T> {
    fn index_cmp(&self, idx: usize, elem: &Option<T>) -> VortexResult<Option<Ordering>> {
        // A null element reads back as a null scalar, so one read answers both whether the
        // element is valid and what its value is.
        let value = self.typed_value(idx)?;

        Ok(match (value, elem.as_ref()) {
            (Some(l), Some(r)) => Some(l.total_compare(*r)),
            (Some(_), None) => Some(Ordering::Greater),
            (None, Some(_)) => Some(Ordering::Less),
            (None, None) => Some(Ordering::Equal),
        })
    }

    fn index_len(&self) -> usize {
        self.len
    }
}

impl<T: NativePType> IndexOrd<usize> for SearchSortedPrimitiveArray<'_, T> {
    fn index_cmp(&self, idx: usize, elem: &usize) -> VortexResult<Option<Ordering>> {
        let value = self.value(idx)?;

        let Some(elem_t) = T::from_usize(*elem) else {
            return Ok(Some(Ordering::Less));
        };

        Ok(Some(value.total_compare(elem_t)))
    }

    fn index_len(&self) -> usize {
        self.len
    }
}

#[cfg(test)]
mod tests {
    use vortex_buffer::buffer;
    use vortex_error::VortexResult;

    use super::Reader;
    use crate::IntoArray;
    use crate::array_session;
    use crate::arrays::PrimitiveArray;
    use crate::executor::VortexSessionExecute;
    use crate::search_sorted::SearchResult;
    use crate::search_sorted::SearchSorted;
    use crate::search_sorted::SearchSortedPrimitiveArray;
    use crate::search_sorted::SearchSortedSide;
    use crate::validity::Validity;

    #[test]
    fn search_sorted_reads_canonical_values_directly() -> VortexResult<()> {
        let mut ctx = array_session().create_execution_ctx();
        let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::NonNullable).into_array();
        let searcher = SearchSortedPrimitiveArray::<i32>::new(&array, &mut ctx);
        assert!(matches!(searcher.reader, Reader::Values(_)));
        assert_eq!(
            searcher.search_sorted(&3i32, SearchSortedSide::Left)?,
            SearchResult::Found(2)
        );
        Ok(())
    }

    #[test]
    fn search_sorted_nullable_values_respects_validity() -> VortexResult<()> {
        let mut ctx = array_session().create_execution_ctx();
        let array = PrimitiveArray::new(
            buffer![99i32, 2, 3, 4],
            Validity::from_iter([false, true, true, true]),
        )
        .into_array();
        let searcher = SearchSortedPrimitiveArray::<i32>::new(&array, &mut ctx);
        assert_eq!(
            searcher.search_sorted(&None::<i32>, SearchSortedSide::Left)?,
            SearchResult::Found(0)
        );
        assert_eq!(
            searcher.search_sorted(&Some(3i32), SearchSortedSide::Left)?,
            SearchResult::Found(2)
        );
        Ok(())
    }

    #[test]
    fn search_sorted_optional_value() -> VortexResult<()> {
        let array = PrimitiveArray::new(buffer![1i32, 2, 3, 4], Validity::AllValid).into_array();
        let mut ctx = array_session().create_execution_ctx();
        let res = SearchSortedPrimitiveArray::<i32>::new(&array, &mut ctx)
            .search_sorted(&Some(3i32), SearchSortedSide::Left)?;
        assert_eq!(res, SearchResult::Found(2));
        Ok(())
    }

    #[test]
    fn search_sorted_optional_value_with_nulls() -> VortexResult<()> {
        let array =
            PrimitiveArray::from_option_iter([None, None, Some(2i32), Some(3)]).into_array();
        let mut ctx = array_session().create_execution_ctx();
        let searcher = SearchSortedPrimitiveArray::<i32>::new(&array, &mut ctx);
        assert_eq!(
            searcher.search_sorted(&Some(2i32), SearchSortedSide::Left)?,
            SearchResult::Found(2)
        );
        assert_eq!(
            searcher.search_sorted(&None, SearchSortedSide::Right)?,
            SearchResult::Found(2)
        );
        Ok(())
    }
}