vortex_scan/

selection.rs

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

//! Defines a selection mask over a scan.

use std::ops::Not;
use std::ops::Range;

use vortex_buffer::BitBufferMut;
use vortex_error::vortex_panic;
use vortex_mask::Mask;

use crate::row_mask::RowMask;
use crate::strict_sorted_buffer::StrictSortedBuffer;

/// A selection identifies a set of rows to include in the scan (in addition to applying any
/// filter predicates).
#[derive(Default, Clone, Debug)]
pub enum Selection {
    /// No selection, all rows are included.
    #[default]
    All,
    /// A selection of sorted, unique rows to include by index.
    IncludeByIndex(StrictSortedBuffer<u64>),
    /// A selection of sorted, unique rows to exclude by index.
    ExcludeByIndex(StrictSortedBuffer<u64>),
    /// A selection of rows to include using a [`roaring::RoaringTreemap`].
    IncludeRoaring(roaring::RoaringTreemap),
    /// A selection of rows to exclude using a [`roaring::RoaringTreemap`].
    ExcludeRoaring(roaring::RoaringTreemap),
}

impl Selection {
    /// Return the row count for this selection.
    pub fn row_count(&self, total_rows: u64) -> u64 {
        match self {
            Selection::All => total_rows,
            Selection::IncludeByIndex(include) => include.len() as u64,
            Selection::ExcludeByIndex(exclude) => total_rows.saturating_sub(exclude.len() as u64),
            Selection::IncludeRoaring(roaring) => roaring.len(),
            Selection::ExcludeRoaring(roaring) => total_rows.saturating_sub(roaring.len()),
        }
    }

    /// Count the selected rows that fall within `range`.
    pub fn row_count_in_range(&self, range: &Range<u64>) -> u64 {
        let range_len = range.end.saturating_sub(range.start);
        match self {
            Selection::All => range_len,
            Selection::IncludeByIndex(include) => count_in_range(range, include),
            Selection::ExcludeByIndex(exclude) => {
                range_len.saturating_sub(count_in_range(range, exclude))
            }
            Selection::IncludeRoaring(roaring) => roaring.range_cardinality(range.clone()),
            Selection::ExcludeRoaring(roaring) => {
                range_len.saturating_sub(roaring.range_cardinality(range.clone()))
            }
        }
    }

    /// Extract the [`RowMask`] for the given range from this selection.
    pub fn row_mask(&self, range: &Range<u64>) -> RowMask {
        if range.start >= range.end {
            return RowMask::new(0, Mask::AllFalse(0));
        }

        // Saturating subtraction to prevent underflow, though range should be valid
        let range_diff = range.end.saturating_sub(range.start);
        let range_len = usize::try_from(range_diff).unwrap_or_else(|_| {
            // If the range is too large for usize, cap it at usize::MAX
            // This is a defensive measure; in practice, ranges should be reasonable
            tracing::warn!(
                "Range length {} exceeds usize::MAX, capping at usize::MAX",
                range_diff
            );
            usize::MAX
        });

        match self {
            Selection::All => RowMask::new(range.start, Mask::new_true(range_len)),
            Selection::IncludeByIndex(include) => {
                RowMask::new(range.start, index_mask(range, range_len, include))
            }
            Selection::ExcludeByIndex(exclude) => {
                RowMask::new(range.start, index_mask(range, range_len, exclude).not())
            }
            Selection::IncludeRoaring(roaring) => {
                RowMask::new(range.start, roaring_mask(roaring, range, range_len))
            }
            Selection::ExcludeRoaring(roaring) => {
                RowMask::new(range.start, roaring_mask(roaring, range, range_len).not())
            }
        }
    }
}

/// Build the mask of positions within `range` that are set in `roaring`.
///
/// Seeks straight to `range` rather than intersecting the whole treemap, and short-circuits splits
/// that are entirely outside or inside the selection.
fn roaring_mask(roaring: &roaring::RoaringTreemap, range: &Range<u64>, range_len: usize) -> Mask {
    let selected = roaring.range_cardinality(range.clone());
    if selected == 0 {
        return Mask::new_false(range_len);
    }

    if selected == range_len as u64 {
        return Mask::new_true(range_len);
    }

    let mut bits = BitBufferMut::new_unset(range_len);
    let mut iter = roaring.iter();
    iter.advance_to(range.start);
    for idx in iter.take_while(|&idx| idx < range.end) {
        bits.set(relativize(range, idx));
    }

    Mask::from_buffer(bits.freeze())
}

/// Count the entries of the sorted `row_indices` that fall within `range`.
fn count_in_range(range: &Range<u64>, row_indices: &[u64]) -> u64 {
    indices_range(range, row_indices).map_or(0, |idx_range| idx_range.len() as u64)
}

/// Build the mask of positions within `range` that are named by the given sorted row indices.
fn index_mask(range: &Range<u64>, range_len: usize, row_indices: &[u64]) -> Mask {
    indices_range(range, row_indices)
        .map(|idx_range| {
            Mask::from_indices(
                range_len,
                row_indices[idx_range]
                    .iter()
                    .map(|&idx| relativize(range, idx)),
            )
        })
        .unwrap_or_else(|| Mask::new_false(range_len))
}

/// Shift an absolute row index to be relative to the start of `range`.
///
/// Panics if the index is not a `usize`-sized offset into `range`. Callers have already narrowed
/// the indices to `range`, whose own length had to fit in a `usize` to size the mask, so a failure
/// here means that invariant was broken rather than that the index was merely out of bounds.
fn relativize(range: &Range<u64>, idx: u64) -> usize {
    idx.checked_sub(range.start)
        .and_then(|relative| usize::try_from(relative).ok())
        .unwrap_or_else(|| {
            vortex_panic!(
                "index {:?} is not a usize offset into range {:?}",
                idx,
                range
            )
        })
}

/// Find the positional range within row_indices that covers all rows in the given range.
fn indices_range(range: &Range<u64>, row_indices: &[u64]) -> Option<Range<usize>> {
    if row_indices.first().is_some_and(|&first| first >= range.end)
        || row_indices.last().is_some_and(|&last| range.start > last)
    {
        return None;
    }

    // For the given row range, find the indices that are within the row_indices.
    let start_idx = row_indices
        .binary_search(&range.start)
        .unwrap_or_else(|x| x);
    let end_idx = row_indices.binary_search(&range.end).unwrap_or_else(|x| x);

    (start_idx != end_idx).then_some(start_idx..end_idx)
}

#[cfg(test)]
mod tests {
    use vortex_buffer::Buffer;

    use super::Selection;
    use crate::strict_sorted_buffer::StrictSortedBuffer;

    fn strict_sorted(indices: impl IntoIterator<Item = u64>) -> StrictSortedBuffer<u64> {
        StrictSortedBuffer::try_new(Buffer::from_iter(indices))
            .expect("test indices should be strictly increasing")
    }

    fn include(indices: impl IntoIterator<Item = u64>) -> Selection {
        Selection::IncludeByIndex(strict_sorted(indices))
    }

    fn exclude(indices: impl IntoIterator<Item = u64>) -> Selection {
        Selection::ExcludeByIndex(strict_sorted(indices))
    }

    #[test]
    fn test_row_count_in_range() {
        let range = 2..8;
        assert_eq!(Selection::All.row_count_in_range(&range), 6);
        // Indices 3, 5, 7 fall inside the range; 1 and 9 do not.
        assert_eq!(include([1, 3, 5, 7, 9]).row_count_in_range(&range), 3);
        assert_eq!(exclude([1, 3, 5, 7, 9]).row_count_in_range(&range), 3);
        assert_eq!(include([10, 11]).row_count_in_range(&range), 0);
        assert_eq!(include([1, 3]).row_count_in_range(&(8..8)), 0);
    }

    #[test]
    fn test_row_count_in_range_roaring() {
        let mut roaring = roaring::RoaringTreemap::new();
        roaring.insert_range(1..10u64);

        let range = 2..8;
        assert_eq!(
            Selection::IncludeRoaring(roaring.clone()).row_count_in_range(&range),
            6
        );
        assert_eq!(
            Selection::ExcludeRoaring(roaring).row_count_in_range(&range),
            0
        );
    }

    #[test]
    fn test_row_count_in_range_roaring_boundaries() {
        let roaring = [0, 3, u64::from(u32::MAX) + 1, u64::MAX - 1, u64::MAX]
            .into_iter()
            .collect::<roaring::RoaringTreemap>();
        let include = Selection::IncludeRoaring(roaring.clone());
        let exclude = Selection::ExcludeRoaring(roaring);

        for (range, expected) in [
            (0..0, 0),
            (0..1, 1),
            (1..3, 0),
            (3..4, 1),
            (u64::from(u32::MAX)..u64::from(u32::MAX) + 2, 1),
            (u64::MAX - 1..u64::MAX, 1),
            (0..u64::MAX, 4),
        ] {
            assert_eq!(include.row_count_in_range(&range), expected);
            assert_eq!(
                exclude.row_count_in_range(&range),
                range.end - range.start - expected
            );
        }
    }

    #[test]
    fn test_row_mask_all() {
        let selection = include([1, 3, 5, 7]);
        let range = 1..8;
        let row_mask = selection.row_mask(&range);

        assert_eq!(row_mask.mask().values().unwrap().indices(), &[0, 2, 4, 6]);
    }

    #[test]
    fn test_row_mask_slice() {
        let selection = include([1, 3, 5, 7]);
        let range = 3..6;
        let row_mask = selection.row_mask(&range);

        assert_eq!(row_mask.mask().values().unwrap().indices(), &[0, 2]);
    }

    #[test]
    fn test_row_mask_exclusive() {
        let selection = include([1, 3, 5, 7]);
        let range = 3..5;
        let row_mask = selection.row_mask(&range);

        assert_eq!(row_mask.mask().values().unwrap().indices(), &[0]);
    }

    #[test]
    fn test_row_mask_all_false() {
        let selection = include([1, 3, 5, 7]);
        let range = 8..10;
        let row_mask = selection.row_mask(&range);

        assert!(row_mask.mask().all_false());
    }

    #[test]
    fn test_row_mask_all_true() {
        let selection = include([1, 3, 4, 5, 6]);
        let range = 3..7;
        let row_mask = selection.row_mask(&range);

        assert!(row_mask.mask().all_true());
    }

    #[test]
    fn test_row_mask_zero() {
        let selection = include([0]);
        let range = 0..5;
        let row_mask = selection.row_mask(&range);

        assert_eq!(row_mask.mask().values().unwrap().indices(), &[0]);
    }

    mod roaring_tests {
        use roaring::RoaringTreemap;

        use super::*;

        #[test]
        fn test_roaring_include_basic() {
            let mut roaring = RoaringTreemap::new();
            roaring.insert(1);
            roaring.insert(3);
            roaring.insert(5);
            roaring.insert(7);

            let selection = Selection::IncludeRoaring(roaring);
            let range = 1..8;
            let row_mask = selection.row_mask(&range);

            assert_eq!(row_mask.mask().values().unwrap().indices(), &[0, 2, 4, 6]);
        }

        #[test]
        fn test_roaring_include_slice() {
            let mut roaring = RoaringTreemap::new();
            roaring.insert(1);
            roaring.insert(3);
            roaring.insert(5);
            roaring.insert(7);

            let selection = Selection::IncludeRoaring(roaring);
            let range = 3..6;
            let row_mask = selection.row_mask(&range);

            assert_eq!(row_mask.mask().values().unwrap().indices(), &[0, 2]);
        }

        #[test]
        fn test_roaring_include_disjoint() {
            let mut roaring = RoaringTreemap::new();
            roaring.insert(1);
            roaring.insert(3);
            roaring.insert(5);
            roaring.insert(7);

            let selection = Selection::IncludeRoaring(roaring);
            let range = 8..10;
            let row_mask = selection.row_mask(&range);

            assert!(row_mask.mask().all_false());
        }

        #[test]
        fn test_roaring_include_large_range() {
            let mut roaring = RoaringTreemap::new();
            // Insert a large number of indices
            for i in (0..1000000).step_by(2) {
                roaring.insert(i);
            }

            let selection = Selection::IncludeRoaring(roaring);
            let range = 1000..2000;
            let row_mask = selection.row_mask(&range);

            assert_eq!(row_mask.row_range(), range);
            assert_eq!(
                row_mask.mask().values().unwrap().indices(),
                &(0..1000).step_by(2).collect::<Vec<_>>()
            );
        }

        #[test]
        fn test_roaring_exclude_basic() {
            let mut roaring = RoaringTreemap::new();
            roaring.insert(1);
            roaring.insert(3);
            roaring.insert(5);

            let selection = Selection::ExcludeRoaring(roaring);
            let range = 0..7;
            let row_mask = selection.row_mask(&range);

            // Should exclude indices 1, 3, 5, so we get 0, 2, 4, 6
            assert_eq!(row_mask.mask().values().unwrap().indices(), &[0, 2, 4, 6]);
        }

        #[test]
        fn test_roaring_exclude_all() {
            let mut roaring = RoaringTreemap::new();
            // Exclude all indices in range
            for i in 10..20 {
                roaring.insert(i);
            }

            let selection = Selection::ExcludeRoaring(roaring);
            let range = 10..20;
            let row_mask = selection.row_mask(&range);

            assert!(row_mask.mask().all_false());
        }

        #[test]
        fn test_roaring_exclude_none() {
            let mut roaring = RoaringTreemap::new();
            roaring.insert(100);
            roaring.insert(101);

            let selection = Selection::ExcludeRoaring(roaring);
            let range = 0..10;
            let row_mask = selection.row_mask(&range);

            // Nothing to exclude in this range
            assert!(row_mask.mask().all_true());
        }

        #[test]
        fn test_roaring_exclude_partial() {
            let mut roaring = RoaringTreemap::new();
            roaring.insert(5);
            roaring.insert(6);
            roaring.insert(7);
            roaring.insert(15); // Outside range

            let selection = Selection::ExcludeRoaring(roaring);
            let range = 5..10;
            let row_mask = selection.row_mask(&range);

            // Should exclude 5, 6, 7 (mapped to 0, 1, 2), keep 8, 9 (mapped to 3, 4)
            assert_eq!(row_mask.mask().values().unwrap().indices(), &[3, 4]);
        }

        #[test]
        fn test_roaring_include_empty() {
            let roaring = RoaringTreemap::new();
            let selection = Selection::IncludeRoaring(roaring);
            let range = 0..100;
            let row_mask = selection.row_mask(&range);

            assert!(row_mask.mask().all_false());
        }

        #[test]
        fn test_roaring_exclude_empty() {
            let roaring = RoaringTreemap::new();
            let selection = Selection::ExcludeRoaring(roaring);
            let range = 0..100;
            let row_mask = selection.row_mask(&range);

            assert!(row_mask.mask().all_true());
        }

        #[test]
        fn test_roaring_include_boundary() {
            let mut roaring = RoaringTreemap::new();
            roaring.insert(0);
            roaring.insert(99);

            let selection = Selection::IncludeRoaring(roaring);
            let range = 0..100;
            let row_mask = selection.row_mask(&range);

            assert_eq!(row_mask.mask().values().unwrap().indices(), &[0, 99]);
        }

        #[test]
        fn test_roaring_include_range_insertion() {
            let mut roaring = RoaringTreemap::new();
            // Use insert_range for efficiency
            roaring.insert_range(10..20);
            roaring.insert_range(30..40);

            let selection = Selection::IncludeRoaring(roaring);
            let range = 15..35;
            let row_mask = selection.row_mask(&range);

            // Should include 15-19 (mapped to 0-4) and 30-34 (mapped to 15-19)
            let expected: Vec<usize> = (0..5).chain(15..20).collect();
            assert_eq!(row_mask.mask().values().unwrap().indices(), &expected);
        }

        #[test]
        fn test_roaring_overflow_protection() {
            let mut roaring = RoaringTreemap::new();
            // Insert very large indices
            roaring.insert(u64::MAX - 1);
            roaring.insert(u64::MAX);

            let selection = Selection::IncludeRoaring(roaring);
            let range = u64::MAX - 10..u64::MAX;
            let row_mask = selection.row_mask(&range);

            assert_eq!(row_mask.row_range(), range);
            assert_eq!(row_mask.mask().values().unwrap().indices(), &[9]);
        }

        #[test]
        fn test_roaring_exclude_overflow_protection() {
            let mut roaring = RoaringTreemap::new();
            roaring.insert(u64::MAX - 1);

            let selection = Selection::ExcludeRoaring(roaring);
            let range = u64::MAX - 10..u64::MAX;
            let row_mask = selection.row_mask(&range);

            assert_eq!(row_mask.row_range(), range);
            assert_eq!(
                row_mask.mask().values().unwrap().indices(),
                &(0..9).collect::<Vec<_>>()
            );
        }

        #[test]
        fn test_roaring_include_vs_buffer_equivalence() {
            // Test that RoaringTreemap and Buffer produce same results
            let indices = vec![1, 3, 5, 7, 9];

            let buffer_selection = include(indices.clone());

            let mut roaring = RoaringTreemap::new();
            for idx in &indices {
                roaring.insert(*idx);
            }
            let roaring_selection = Selection::IncludeRoaring(roaring);

            let range = 0..12;
            let buffer_mask = buffer_selection.row_mask(&range);
            let roaring_mask = roaring_selection.row_mask(&range);

            for mask in [buffer_mask, roaring_mask] {
                assert_eq!(mask.row_range(), range);
                assert_eq!(mask.mask().values().unwrap().indices(), &[1, 3, 5, 7, 9]);
            }
        }

        #[test]
        fn test_roaring_exclude_vs_buffer_equivalence() {
            // Test that ExcludeRoaring and ExcludeByIndex produce same results
            let indices = vec![2, 4, 6, 8];

            let buffer_selection = exclude(indices.clone());

            let mut roaring = RoaringTreemap::new();
            for idx in &indices {
                roaring.insert(*idx);
            }
            let roaring_selection = Selection::ExcludeRoaring(roaring);

            let range = 0..10;
            let buffer_mask = buffer_selection.row_mask(&range);
            let roaring_mask = roaring_selection.row_mask(&range);

            for mask in [buffer_mask, roaring_mask] {
                assert_eq!(mask.row_range(), range);
                assert_eq!(mask.mask().values().unwrap().indices(), &[0, 1, 3, 5, 7, 9]);
            }
        }
    }
}