parquet/arrow/arrow_reader/selection/

algebra.rs

// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership
// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements

//! Set algebra backing [`RowSelection::and_then`], [`RowSelection::intersection`]
//! and [`RowSelection::union`]
//!
//! Each operation has two implementations, picked by the backing of its
//! operands: a merge of the [`RowSelector`] runs, and a bitwise variant over
//! [`BooleanBuffer`] masks.

use super::boolean::{MaskSelection, boolean_mask_from_selectors};
use super::{MaskRunIter, RowSelection, RowSelectionInner, RowSelector};
use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder, MutableBuffer, bit_util};
use std::cmp::Ordering;
use std::iter::Peekable;

// Use word-at-a-time expansion for masks with at least 8192 rows, roughly
// 75% outer selectivity and 5% inner selectivity. Smaller or sparser inputs
// use set indices to avoid scanning every output word.
// In `outer.and_then(inner)`, outer selects from the original rows, and inner
// selects from those selected rows: inner.len() == outer.count_set_bits().
// The selectivity constants are divisors: at most 1/4 dropped, at least 1/20 kept.
const AND_THEN_DENSE_MASK_MIN_LEN: usize = 8192;
const AND_THEN_DENSE_MASK_MAX_DROPPED_FRACTION: usize = 4;
const AND_THEN_DENSE_MASK_MIN_INNER_FRACTION: usize = 20;

/// Applies `second` to the rows selected by `first`, both selector-backed.
pub(super) fn and_then_row_selections(
    first: &[RowSelector],
    second: &[RowSelector],
) -> RowSelection {
    let mut selectors = vec![];
    let mut first = first.iter().copied().peekable();
    let mut second = second.iter().copied().peekable();
    and_then_iter(&mut selectors, &mut first, &mut second);
    RowSelection::from_selectors(selectors)
}

/// Applies the mask `second` to the rows selected by the selector-backed `first`.
///
/// The mask is streamed as [`RowSelector`] runs, so it is never materialized.
pub(super) fn and_then_selectors_with_mask(
    first: &[RowSelector],
    second: &BooleanBuffer,
) -> RowSelection {
    let mut selectors = vec![];
    let mut first = first.iter().copied().peekable();
    let mut second = MaskRunIter::new(second).peekable();
    and_then_iter(&mut selectors, &mut first, &mut second);
    RowSelection::from_selectors(selectors)
}

fn and_then_iter<I, J>(
    selectors: &mut Vec<RowSelector>,
    first: &mut Peekable<I>,
    second: &mut Peekable<J>,
) where
    I: Iterator<Item = RowSelector>,
    J: Iterator<Item = RowSelector>,
{
    let mut to_skip = 0;
    while let Some(b) = second.peek_mut() {
        let a = first
            .peek_mut()
            .expect("selection exceeds the number of selected rows");

        if b.row_count == 0 {
            second.next().unwrap();
            continue;
        }

        if a.row_count == 0 {
            first.next().unwrap();
            continue;
        }

        if a.skip {
            // Records were skipped when producing second
            to_skip += a.row_count;
            first.next().unwrap();
            continue;
        }

        let skip = b.skip;
        let to_process = a.row_count.min(b.row_count);

        a.row_count -= to_process;
        b.row_count -= to_process;

        match skip {
            true => to_skip += to_process,
            false => {
                if to_skip != 0 {
                    selectors.push(RowSelector::skip(to_skip));
                    to_skip = 0;
                }
                selectors.push(RowSelector::select(to_process))
            }
        }
    }

    for v in first {
        if v.row_count != 0 {
            assert!(
                v.skip,
                "selection contains less than the number of selected rows"
            );
            to_skip += v.row_count
        }
    }

    if to_skip != 0 {
        selectors.push(RowSelector::skip(to_skip));
    }
}

/// Combine two lists of `RowSelection` return the intersection of them
/// For example:
/// self:      NNYYYYNNYYNYN
/// other:     NYNNNNNNY
///
/// returned:  NNNNNNNNYYNYN
pub(super) fn intersect_row_selections(
    left: &[RowSelector],
    right: &[RowSelector],
) -> RowSelection {
    combine_row_selections(left, right, |l, r| l || r)
}

/// Combine two lists of `RowSelector` return the union of them
/// For example:
/// self:      NNYYYYNNYYNYN
/// other:     NYNNNNNNY
///
/// returned:  NYYYYYNNYYNYN
///
/// This can be removed from here once RowSelection::union is in parquet::arrow
pub(super) fn union_row_selections(left: &[RowSelector], right: &[RowSelector]) -> RowSelection {
    combine_row_selections(left, right, |l, r| l && r)
}

/// Combines overlapping selector runs, passing the longer side's tail through unchanged.
fn combine_row_selections(
    left: &[RowSelector],
    right: &[RowSelector],
    combine_skip: impl Fn(bool, bool) -> bool,
) -> RowSelection {
    let mut l_iter = left.iter().copied().peekable();
    let mut r_iter = right.iter().copied().peekable();

    let iter = std::iter::from_fn(move || {
        loop {
            let l = l_iter.peek_mut();
            let r = r_iter.peek_mut();

            match (l, r) {
                (Some(a), _) if a.row_count == 0 => {
                    l_iter.next().unwrap();
                }
                (_, Some(b)) if b.row_count == 0 => {
                    r_iter.next().unwrap();
                }
                (Some(l), Some(r)) => {
                    let skip = combine_skip(l.skip, r.skip);
                    let row_count = l.row_count.min(r.row_count);

                    if l.row_count < r.row_count {
                        r.row_count -= row_count;
                        l_iter.next();
                    } else {
                        l.row_count -= row_count;
                        r_iter.next();
                    }

                    return Some(RowSelector { row_count, skip });
                }
                (Some(_), None) => return l_iter.next(),
                (None, Some(_)) => return r_iter.next(),
                (None, None) => return None,
            }
        }
    });

    iter.collect()
}

/// Combines a mask and selectors without expanding the mask into runs.
pub(super) fn combine_mixed_selection(
    selectors: &[RowSelector],
    mask: &MaskSelection,
    op: impl FnOnce(&BooleanBuffer, &BooleanBuffer) -> BooleanBuffer,
) -> RowSelection {
    // An empty operand leaves the other unchanged.
    if selectors.is_empty() {
        return RowSelection::from_mask_selection(mask.clone());
    }
    if mask.mask().is_empty() {
        return RowSelection::from_selectors(selectors.to_vec());
    }
    let selectors = boolean_mask_from_selectors(selectors);
    RowSelection::from_boolean_buffer(op(mask.mask(), &selectors))
}

/// Bitwise AND of two mask-backed selections. Longer side's tail passes through.
pub(super) fn intersect_masks(l: &BooleanBuffer, r: &BooleanBuffer) -> BooleanBuffer {
    if l.len() == r.len() {
        return combine_equal_length_masks(l, r, |a, b| a & b);
    }
    combine_unequal_length_masks(l, r, |a, b| a & b)
}

/// Bitwise OR of two mask-backed selections. Longer side's tail passes through.
pub(super) fn union_masks(l: &BooleanBuffer, r: &BooleanBuffer) -> BooleanBuffer {
    if l.len() == r.len() {
        return combine_equal_length_masks(l, r, |a, b| a | b);
    }
    combine_unequal_length_masks(l, r, |a, b| a | b)
}

/// Combines two masks of equal length with the bitwise operation `op`.
///
/// `BitAnd`/`BitOr` on `&BooleanBuffer` normalise the result to a zero bit offset,
/// which costs a second allocation and a shifting copy when both operands
/// share the same non-zero sub-64-bit alignment, causing
/// `from_bitwise_binary_op` to return a non-zero-offset result. Building the
/// buffer directly keeps the offset, as the unequal-length path does.
fn combine_equal_length_masks<F>(l: &BooleanBuffer, r: &BooleanBuffer, op: F) -> BooleanBuffer
where
    F: FnMut(u64, u64) -> u64,
{
    BooleanBuffer::from_bitwise_binary_op(
        l.values(),
        l.offset(),
        r.values(),
        r.offset(),
        l.len(),
        op,
    )
}

/// Combines two masks of unequal length with the bitwise operation `op`,
/// passing the longer side's tail through unchanged.
///
/// The longer mask is copied once into a [`MutableBuffer`] and `op` is then
/// applied in place over the common prefix. This avoids materialising the
/// prefix into its own buffer and copying both prefix and tail again through a
/// [`BooleanBufferBuilder`].
///
/// Neither the mask offsets nor the prefix length are assumed to be byte
/// aligned: the copy keeps the longer mask's offset within its first byte so it
/// stays a plain byte copy, and that offset is carried over to the result. Only
/// the longer mask's own byte range is copied, so the result does not retain the
/// backing allocation it was sliced from.
fn combine_unequal_length_masks<F>(l: &BooleanBuffer, r: &BooleanBuffer, op: F) -> BooleanBuffer
where
    F: FnMut(u64, u64) -> u64,
{
    let (longer, shorter) = if l.len() > r.len() { (l, r) } else { (r, l) };

    let sub_byte_offset = longer.offset() % 8;
    let start_byte = longer.offset() / 8;
    let end_byte = bit_util::ceil(longer.offset() + longer.len(), 8);
    let bytes = &longer.values()[start_byte..end_byte];
    let mut buffer = MutableBuffer::new(bytes.len());
    buffer.extend_from_slice(bytes);

    bit_util::apply_bitwise_binary_op(
        buffer.as_slice_mut(),
        sub_byte_offset,
        shorter.values(),
        shorter.offset(),
        shorter.len(),
        op,
    );

    BooleanBuffer::new(buffer.into(), sub_byte_offset, longer.len())
}

/// Applies `other` to the selected rows of `mask`, preserving the original row domain.
pub(super) fn and_then_mask(mask: &BooleanBuffer, other: &RowSelection) -> BooleanBuffer {
    match &other.inner {
        RowSelectionInner::Mask(other_mask) => and_then_masks(mask, other_mask.mask()),
        RowSelectionInner::Selectors(selectors) => {
            and_then_mask_from_selectors(mask, selectors.iter().copied())
        }
    }
}

fn and_then_mask_from_selectors<I>(mask: &BooleanBuffer, other: I) -> BooleanBuffer
where
    I: IntoIterator<Item = RowSelector>,
{
    let mut builder = BooleanBufferBuilder::new(mask.len());
    let mut other_iter = other.into_iter();
    let mut current = other_iter.next();
    let mut cursor = 0usize;

    // Iterate only over the set positions in `mask`; the gaps of unset bits
    // are filled in bulk with `append_n` instead of bit-by-bit.
    for set_idx in mask.set_indices() {
        if set_idx > cursor {
            builder.append_n(set_idx - cursor, false);
        }
        cursor = set_idx + 1;

        while current.as_ref().is_some_and(|s| s.row_count == 0) {
            current = other_iter.next();
        }
        let selector = current
            .as_mut()
            .expect("selection contains less than the number of selected rows");
        let selected = !selector.skip;
        selector.row_count -= 1;
        builder.append(selected);
    }
    if cursor < mask.len() {
        builder.append_n(mask.len() - cursor, false);
    }

    if current.is_some_and(|s| s.row_count != 0) || other_iter.any(|s| s.row_count != 0) {
        panic!("selection exceeds the number of selected rows");
    }

    builder.finish()
}

fn and_then_masks(mask: &BooleanBuffer, other: &BooleanBuffer) -> BooleanBuffer {
    let selected_count = mask.count_set_bits();
    match other.len().cmp(&selected_count) {
        Ordering::Less => panic!("selection contains less than the number of selected rows"),
        Ordering::Greater => panic!("selection exceeds the number of selected rows"),
        Ordering::Equal => {}
    }

    let other_true_count = other.count_set_bits();
    if other_true_count == 0 {
        return BooleanBuffer::new_unset(mask.len());
    }
    if other_true_count == selected_count {
        return mask.clone();
    }

    if should_use_dense_mask(mask.len(), selected_count, other_true_count) {
        and_then_dense_masks(mask, other)
    } else {
        and_then_sparse_masks(mask, other)
    }
}

/// Checks the length and selectivity thresholds for dense expansion.
#[inline]
fn should_use_dense_mask(mask_len: usize, selected_count: usize, other_true_count: usize) -> bool {
    mask_len >= AND_THEN_DENSE_MASK_MIN_LEN
        && mask_len - selected_count <= mask_len.div_ceil(AND_THEN_DENSE_MASK_MAX_DROPPED_FRACTION)
        && other_true_count >= selected_count / AND_THEN_DENSE_MASK_MIN_INNER_FRACTION
}

/// Maps each set bit in `other` to the corresponding set bit in `mask`.
/// Requires `other.len() == mask.count_set_bits()`.
fn and_then_sparse_masks(mask: &BooleanBuffer, other: &BooleanBuffer) -> BooleanBuffer {
    debug_assert_eq!(other.len(), mask.count_set_bits());

    let mut builder = BooleanBufferBuilder::new(mask.len());
    let mut outer_set_indices = mask.set_indices();
    let mut next_selected_ordinal = 0usize;
    let mut cursor = 0usize;

    for selected_ordinal in other.set_indices() {
        let skip = selected_ordinal - next_selected_ordinal;
        let set_idx = outer_set_indices
            .nth(skip)
            .expect("validated other length matches selected row count");
        if set_idx > cursor {
            builder.append_n(set_idx - cursor, false);
        }
        builder.append(true);
        cursor = set_idx + 1;
        next_selected_ordinal = selected_ordinal + 1;
    }

    if cursor < mask.len() {
        builder.append_n(mask.len() - cursor, false);
    }

    builder.finish()
}

/// Scatters the next `mask_word.count_ones()` bits of `other` into each mask word.
/// Requires `other.len() == mask.count_set_bits()`.
#[inline(never)]
fn and_then_dense_masks(mask: &BooleanBuffer, other: &BooleanBuffer) -> BooleanBuffer {
    debug_assert_eq!(other.len(), mask.count_set_bits());

    let mut other_chunks = other.bit_chunks().iter_padded();
    let mut other_remaining = other.len();
    let mut pending = 0_u128;
    let mut pending_len = 0;
    let mut output = MutableBuffer::with_capacity(mask.len().div_ceil(8));

    for mask_word in mask
        .bit_chunks()
        .iter_padded()
        .take(mask.len().div_ceil(64))
    {
        let selected = mask_word.count_ones() as usize;
        while pending_len < selected {
            let chunk = other_chunks
                .next()
                .expect("validated other length matches selected row count");
            let chunk_len = other_remaining.min(64);
            pending |= (chunk as u128) << pending_len;
            pending_len += chunk_len;
            other_remaining -= chunk_len;
        }

        let values = pending as u64;
        output.extend_from_slice(&bit_util::expand(values, mask_word).to_le_bytes());
        pending >>= selected;
        pending_len -= selected;
    }

    debug_assert_eq!(other_remaining, 0);
    debug_assert_eq!(pending_len, 0);
    output.truncate(mask.len().div_ceil(8));
    BooleanBuffer::new(output.into(), 0, mask.len())
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow_array::BooleanArray;
    use rand::{RngExt, SeedableRng, rng, rngs::StdRng};

    #[test]
    fn test_and() {
        let mut a = RowSelection::from(vec![
            RowSelector::skip(12),
            RowSelector::select(23),
            RowSelector::skip(3),
            RowSelector::select(5),
        ]);

        let b = RowSelection::from(vec![
            RowSelector::select(5),
            RowSelector::skip(4),
            RowSelector::select(15),
            RowSelector::skip(4),
        ]);

        let mut expected = RowSelection::from(vec![
            RowSelector::skip(12),
            RowSelector::select(5),
            RowSelector::skip(4),
            RowSelector::select(14),
            RowSelector::skip(3),
            RowSelector::select(1),
            RowSelector::skip(4),
        ]);

        assert_eq!(a.and_then(&b), expected);

        a.split_off(7);
        expected.split_off(7);
        assert_eq!(a.and_then(&b), expected);

        let a = RowSelection::from(vec![RowSelector::select(5), RowSelector::skip(3)]);

        let b = RowSelection::from(vec![
            RowSelector::select(2),
            RowSelector::skip(1),
            RowSelector::select(1),
            RowSelector::skip(1),
        ]);

        assert_eq!(
            a.and_then(&b).selectors(),
            vec![
                RowSelector::select(2),
                RowSelector::skip(1),
                RowSelector::select(1),
                RowSelector::skip(4)
            ]
        );
    }

    #[test]
    #[should_panic(expected = "selection exceeds the number of selected rows")]
    fn test_and_longer() {
        let a = RowSelection::from(vec![
            RowSelector::select(3),
            RowSelector::skip(33),
            RowSelector::select(3),
            RowSelector::skip(33),
        ]);
        let b = RowSelection::from(vec![RowSelector::select(36)]);
        a.and_then(&b);
    }

    #[test]
    #[should_panic(expected = "selection contains less than the number of selected rows")]
    fn test_and_shorter() {
        let a = RowSelection::from(vec![
            RowSelector::select(3),
            RowSelector::skip(33),
            RowSelector::select(3),
            RowSelector::skip(33),
        ]);
        let b = RowSelection::from(vec![RowSelector::select(3)]);
        a.and_then(&b);
    }

    #[test]
    fn test_intersect_row_selection_and_combine() {
        // a size equal b size
        let a = vec![
            RowSelector::select(5),
            RowSelector::skip(4),
            RowSelector::select(1),
        ];
        let b = vec![
            RowSelector::select(8),
            RowSelector::skip(1),
            RowSelector::select(1),
        ];

        let res = intersect_row_selections(&a, &b);
        assert_eq!(
            res.selectors(),
            vec![
                RowSelector::select(5),
                RowSelector::skip(4),
                RowSelector::select(1),
            ],
        );

        // a size larger than b size
        let a = vec![
            RowSelector::select(3),
            RowSelector::skip(33),
            RowSelector::select(3),
            RowSelector::skip(33),
        ];
        let b = vec![RowSelector::select(36), RowSelector::skip(36)];
        let res = intersect_row_selections(&a, &b);
        assert_eq!(
            res.selectors(),
            vec![RowSelector::select(3), RowSelector::skip(69)]
        );

        // a size less than b size
        let a = vec![RowSelector::select(3), RowSelector::skip(7)];
        let b = vec![
            RowSelector::select(2),
            RowSelector::skip(2),
            RowSelector::select(2),
            RowSelector::skip(2),
            RowSelector::select(2),
        ];
        let res = intersect_row_selections(&a, &b);
        assert_eq!(
            res.selectors(),
            vec![RowSelector::select(2), RowSelector::skip(8)]
        );

        let a = vec![RowSelector::select(3), RowSelector::skip(7)];
        let b = vec![
            RowSelector::select(2),
            RowSelector::skip(2),
            RowSelector::select(2),
            RowSelector::skip(2),
            RowSelector::select(2),
        ];
        let res = intersect_row_selections(&a, &b);
        assert_eq!(
            res.selectors(),
            vec![RowSelector::select(2), RowSelector::skip(8)]
        );
    }

    #[test]
    #[cfg_attr(miri, ignore)] // Takes too long
    fn test_and_fuzz() {
        let mut rand = rng();
        for _ in 0..100 {
            let a_len = rand.random_range(10..100);
            let a_bools: Vec<_> = (0..a_len).map(|_| rand.random_bool(0.2)).collect();
            let a = RowSelection::from_filters(&[BooleanArray::from(a_bools.clone())]);

            let b_len: usize = a_bools.iter().map(|x| *x as usize).sum();
            let b_bools: Vec<_> = (0..b_len).map(|_| rand.random_bool(0.8)).collect();
            let b = RowSelection::from_filters(&[BooleanArray::from(b_bools.clone())]);

            let mut expected_bools = vec![false; a_len];

            let mut iter_b = b_bools.iter();
            for (idx, b) in a_bools.iter().enumerate() {
                if *b && *iter_b.next().unwrap() {
                    expected_bools[idx] = true;
                }
            }

            let expected = RowSelection::from_filters(&[BooleanArray::from(expected_bools)]);

            assert_eq!(a_len, expected.total_row_count());

            assert_eq!(a.and_then(&b), expected);
        }
    }

    #[test]
    fn test_intersection() {
        let selection = RowSelection::from(vec![RowSelector::select(1048576)]);
        let result = selection.intersection(&selection);
        assert_eq!(result, selection);

        let a = RowSelection::from(vec![
            RowSelector::skip(10),
            RowSelector::select(10),
            RowSelector::skip(10),
            RowSelector::select(20),
        ]);

        let b = RowSelection::from(vec![
            RowSelector::skip(20),
            RowSelector::select(20),
            RowSelector::skip(10),
        ]);

        let result = a.intersection(&b);
        assert_eq!(
            result.selectors(),
            vec![
                RowSelector::skip(30),
                RowSelector::select(10),
                RowSelector::skip(10)
            ]
        );
    }

    #[test]
    fn test_union() {
        let selection = RowSelection::from(vec![RowSelector::select(1048576)]);
        let result = selection.union(&selection);
        assert_eq!(result, selection);

        // NYNYY
        let a = RowSelection::from(vec![
            RowSelector::skip(10),
            RowSelector::select(10),
            RowSelector::skip(10),
            RowSelector::select(20),
        ]);

        // NNYYNYN
        let b = RowSelection::from(vec![
            RowSelector::skip(20),
            RowSelector::select(20),
            RowSelector::skip(10),
            RowSelector::select(10),
            RowSelector::skip(10),
        ]);

        let result = a.union(&b);

        // NYYYYYN
        assert_eq!(
            result.iter().copied().collect::<Vec<_>>(),
            vec![
                RowSelector::skip(10),
                RowSelector::select(50),
                RowSelector::skip(10),
            ]
        );
    }

    #[test]
    fn test_mask_and_then_preserves_backing() {
        let outer_bits = vec![false, true, true, false, true, false, true];
        let inner_bits = vec![true, false, true, false];
        let outer_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(outer_bits.clone()));
        let inner = RowSelection::from_filters(&[BooleanArray::from(inner_bits.clone())]);

        let result = outer_mask.and_then(&inner);
        assert!(result.as_mask().is_some());

        let outer_selectors = RowSelection::from_filters(&[BooleanArray::from(outer_bits)]);
        let expected = outer_selectors.and_then(&inner);
        assert_eq!(result, expected);

        let result_mask = result.as_mask().unwrap();
        let actual_bits: Vec<_> = (0..result_mask.len())
            .map(|i| result_mask.value(i))
            .collect();
        assert_eq!(
            actual_bits,
            vec![false, true, false, false, true, false, false]
        );
    }

    #[test]
    fn test_mask_and_then_mask_preserves_backing() {
        let outer_bits = vec![false, true, true, false, true, false, true, true];
        let inner_bits = vec![false, true, false, true, false];
        let outer_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(outer_bits.clone()));
        let inner_mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(inner_bits));

        let result = outer_mask.and_then(&inner_mask);
        assert!(result.as_mask().is_some());

        let outer_selectors = RowSelection::from_filters(&[BooleanArray::from(outer_bits)]);
        let inner_selectors = RowSelection::from_filters(&[BooleanArray::from(vec![
            false, true, false, true, false,
        ])]);
        assert_eq!(result, outer_selectors.and_then(&inner_selectors));

        let result_mask = result.as_mask().unwrap();
        let actual_bits: Vec<_> = (0..result_mask.len())
            .map(|i| result_mask.value(i))
            .collect();
        assert_eq!(
            actual_bits,
            vec![false, false, true, false, false, false, true, false]
        );
    }

    struct AndThenTest {
        outer: BooleanBuffer,
        inner: BooleanBuffer,
        dense_mask_expected: bool,
    }

    impl AndThenTest {
        fn run(self) {
            let Self {
                outer,
                inner,
                dense_mask_expected,
            } = self;
            assert_eq!(inner.len(), outer.count_set_bits());
            assert_eq!(
                should_use_dense_mask(outer.len(), inner.len(), inner.count_set_bits()),
                dense_mask_expected
            );

            let mut inner_idx = 0;
            let expected = BooleanBuffer::from_iter((0..outer.len()).map(|i| {
                if !outer.value(i) {
                    return false;
                }
                let value = inner.value(inner_idx);
                inner_idx += 1;
                value
            }));

            if dense_mask_expected {
                assert_eq!(and_then_dense_masks(&outer, &inner), expected);
            }

            let outer = RowSelection::from_boolean_buffer(outer);
            let inner = RowSelection::from_boolean_buffer(inner);
            let actual = outer.and_then(&inner);
            assert_eq!(actual.as_mask().unwrap(), &expected);
        }
    }

    #[test]
    fn test_dense_mask_and_then_mask_with_offsets() {
        for len in [8192, 8193, 10_000] {
            for outer_stride in [100, 20, 4] {
                let outer_offset = 3;
                let outer = BooleanBuffer::from_iter(
                    (0..len + outer_offset).map(|i| (i + 1) % outer_stride != 0),
                )
                .slice(outer_offset, len);

                let inner_offset = 5;
                let inner_len = outer.count_set_bits();
                let inner = BooleanBuffer::from_iter(
                    (0..inner_len + inner_offset).map(|i| (i + 2) % 97 != 0),
                )
                .slice(inner_offset, inner_len);

                AndThenTest {
                    outer,
                    inner,
                    dense_mask_expected: true,
                }
                .run();
            }
        }
    }

    #[test]
    fn test_dense_mask_and_then_mask_randomized_word_boundaries() {
        let mut rng = StdRng::seed_from_u64(0x67e4_a91b_239d_c805);

        for (len, outer_offset, inner_offset) in [(8192, 0, 0), (8192, 3, 5), (8193, 7, 1)] {
            let outer =
                BooleanBuffer::from_iter((0..len + outer_offset).map(|_| rng.random_bool(0.8)))
                    .slice(outer_offset, len);
            let inner_len = outer.count_set_bits();
            let inner = BooleanBuffer::from_iter(
                (0..inner_len + inner_offset).map(|_| rng.random_bool(0.5)),
            )
            .slice(inner_offset, inner_len);

            AndThenTest {
                outer,
                inner,
                dense_mask_expected: true,
            }
            .run();
        }
    }

    #[test]
    fn test_dense_mask_and_then_policy() {
        let len = 8192;
        let selected = len * 3 / 4;
        let min_inner = selected / AND_THEN_DENSE_MASK_MIN_INNER_FRACTION;

        assert!(!should_use_dense_mask(len - 1, selected, min_inner));
        assert!(!should_use_dense_mask(len, selected - 1, min_inner));
        assert!(!should_use_dense_mask(len, selected, min_inner - 1));
        assert!(should_use_dense_mask(len, selected, min_inner));
    }

    #[test]
    fn test_selector_and_then_mask() {
        let outer =
            RowSelection::from_filters(&[BooleanArray::from(vec![false, true, true, false, true])]);
        let inner = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true]));

        let result = outer.and_then(&inner);
        assert!(result.as_mask().is_none());
        assert_eq!(
            result,
            RowSelection::from_filters(&[BooleanArray::from(vec![
                false, true, false, false, true,
            ])])
        );
    }

    #[test]
    fn test_mask_and_then_none_selected_returns_all_unset() {
        let outer = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
            false, true, true, false, true,
        ]));
        let inner =
            RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![false, false, false]));

        let result = outer.and_then(&inner);
        let mask = result.as_mask().unwrap();
        assert_eq!(mask.len(), 5);
        assert_eq!(mask.count_set_bits(), 0);
    }

    #[test]
    fn test_mask_intersection_uses_bitwise() {
        let a_bits = vec![true, true, false, true, false, true];
        let b_bits = vec![true, false, true, true, true, false];
        let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone()));
        let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone()));

        let r = a.intersection(&b);
        assert!(r.as_mask().is_some());

        let expected: Vec<bool> = a_bits.iter().zip(&b_bits).map(|(x, y)| *x && *y).collect();
        let expected_sel = RowSelection::from_filters(&[BooleanArray::from(expected)]);
        assert_eq!(r, expected_sel);
    }

    #[test]
    fn test_mask_union_uses_bitwise() {
        let a_bits = vec![true, false, false, true, false, false];
        let b_bits = vec![false, true, false, false, true, false];
        let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits.clone()));
        let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits.clone()));

        let r = a.union(&b);
        assert!(r.as_mask().is_some());

        let expected: Vec<bool> = a_bits.iter().zip(&b_bits).map(|(x, y)| *x || *y).collect();
        let expected_sel = RowSelection::from_filters(&[BooleanArray::from(expected)]);
        assert_eq!(r, expected_sel);
    }

    #[test]
    fn test_mixed_mask_selector_intersection_and_union() {
        let mask_bits = vec![true, false, true, false, true, false];
        let selector_bits = vec![false, true, true, false, false, true];
        let mask = RowSelection::from_boolean_buffer(BooleanBuffer::from(mask_bits.clone()));
        let selectors = RowSelection::from_filters(&[BooleanArray::from(selector_bits.clone())]);

        let intersection_bits: Vec<_> = mask_bits
            .iter()
            .zip(&selector_bits)
            .map(|(x, y)| *x && *y)
            .collect();
        let expected_intersection =
            RowSelection::from_filters(&[BooleanArray::from(intersection_bits)]);
        for result in [mask.intersection(&selectors), selectors.intersection(&mask)] {
            assert!(result.as_mask().is_some());
            assert_eq!(result, expected_intersection);
        }

        let union_bits: Vec<_> = mask_bits
            .iter()
            .zip(&selector_bits)
            .map(|(x, y)| *x || *y)
            .collect();
        let expected_union = RowSelection::from_filters(&[BooleanArray::from(union_bits)]);
        for result in [mask.union(&selectors), selectors.union(&mask)] {
            assert!(result.as_mask().is_some());
            assert_eq!(result, expected_union);
        }
    }

    #[test]
    fn test_mask_intersection_uneven_passes_tail_through() {
        let a_bits = vec![true, true, true, true, true];
        let b_bits = vec![true, false, true];
        let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits));
        let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits));

        let r = a.intersection(&b);
        let r_mask = r.as_mask().unwrap();
        assert_eq!(r_mask.len(), 5);
        let bits: Vec<bool> = (0..5).map(|i| r_mask.value(i)).collect();
        assert_eq!(bits, vec![true, false, true, true, true]);

        // Swapped operands: the right side is longer and its tail passes through.
        let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, true]));
        let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
            true, true, true, false, true,
        ]));
        let r = a.intersection(&b);
        let r_mask = r.as_mask().unwrap();
        assert_eq!(r_mask.len(), 5);
        let bits: Vec<bool> = (0..5).map(|i| r_mask.value(i)).collect();
        assert_eq!(bits, vec![true, false, true, false, true]);
    }

    #[test]
    fn test_mask_union_uneven_passes_tail_through() {
        let a_bits = vec![true, false, true];
        let b_bits = vec![false, true, false, true, false];
        let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(a_bits));
        let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(b_bits));

        let r = a.union(&b);
        let r_mask = r.as_mask().unwrap();
        assert_eq!(r_mask.len(), 5);
        let bits: Vec<bool> = (0..5).map(|i| r_mask.value(i)).collect();
        assert_eq!(bits, vec![true, true, true, true, false]);

        let a = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![
            false, true, false, false, true,
        ]));
        let b = RowSelection::from_boolean_buffer(BooleanBuffer::from(vec![true, false, false]));
        let r = a.union(&b);
        let r_mask = r.as_mask().unwrap();
        let bits: Vec<bool> = (0..5).map(|i| r_mask.value(i)).collect();
        assert_eq!(bits, vec![true, true, false, false, true]);
    }

    /// Expected result of combining two masks of possibly differing lengths:
    /// `op` over the common prefix, then the longer side's tail unchanged.
    fn expected_combined(l: &[bool], r: &[bool], op: fn(bool, bool) -> bool) -> Vec<bool> {
        let common = l.len().min(r.len());
        let longer = if l.len() > r.len() { l } else { r };
        (0..common)
            .map(|i| op(l[i], r[i]))
            .chain(longer[common..].iter().copied())
            .collect()
    }

    fn assert_mask_eq(actual: &BooleanBuffer, expected: &[bool], context: &str) {
        assert_eq!(actual.len(), expected.len(), "{context}: length");
        let actual: Vec<bool> = actual.iter().collect();
        assert_eq!(actual, expected, "{context}");
    }

    #[test]
    fn test_mixed_mask_selector_algebra_with_offsets_and_uneven_lengths() {
        for rows in [1, 7, 8, 63, 64, 65, 137] {
            for (mask_len, selector_len) in [(rows, rows), (rows, 137), (137, rows)] {
                for offset in [0, 1, 7, 63] {
                    for mask_bits in [
                        vec![false; mask_len],
                        vec![true; mask_len],
                        (0..mask_len).map(|i| i < mask_len / 2).collect(),
                        (0..mask_len).map(|i| i % 2 == 0).collect(),
                    ] {
                        let buffer = BooleanBuffer::from_iter(
                            std::iter::repeat_n(false, offset).chain(mask_bits.iter().copied()),
                        )
                        .slice(offset, mask_len);
                        let mask = RowSelection::from_boolean_buffer(buffer);
                        for selector_bits in [
                            vec![false; selector_len],
                            vec![true; selector_len],
                            (0..selector_len).map(|i| i % 2 == 0).collect(),
                            (0..selector_len).map(|i| (i / 11) % 2 == 0).collect(),
                        ] {
                            let selectors = RowSelection::from_filters(&[BooleanArray::from(
                                selector_bits.clone(),
                            )]);
                            let context = format!(
                                "mask_len={mask_len} selector_len={selector_len} offset={offset}"
                            );
                            let intersection =
                                expected_combined(&mask_bits, &selector_bits, |a, b| a && b);
                            let union =
                                expected_combined(&mask_bits, &selector_bits, |a, b| a || b);
                            for (left, right) in [(&mask, &selectors), (&selectors, &mask)] {
                                assert_mask_eq(
                                    left.intersection(right).as_mask().unwrap(),
                                    &intersection,
                                    &context,
                                );
                                assert_mask_eq(
                                    left.union(right).as_mask().unwrap(),
                                    &union,
                                    &context,
                                );
                            }
                        }
                    }
                }
            }
        }
    }

    #[test]
    fn test_mixed_algebra_empty_operand_preserves_tail() {
        let empty = RowSelection::default();
        let buffer = BooleanBuffer::from_iter((0..64).map(|i| i % 3 == 0));
        let mask = RowSelection::from_boolean_buffer(buffer.slice(5, 40));
        let selectors = RowSelection::from(vec![
            RowSelector::skip(3),
            RowSelector::select(7),
            RowSelector::skip(4),
        ]);
        let empty_mask = RowSelection::from_boolean_buffer(BooleanBuffer::new_unset(0));
        for (left, right, expected) in [
            (&empty, &mask, &mask),
            (&mask, &empty, &mask),
            (&empty_mask, &selectors, &selectors),
            (&selectors, &empty_mask, &selectors),
            (&empty, &empty_mask, &empty_mask),
            (&empty_mask, &empty, &empty_mask),
        ] {
            for result in [left.intersection(right), left.union(right)] {
                assert_eq!(&result, expected);
                match (result.as_mask(), expected.as_mask()) {
                    (Some(actual), Some(expected)) => {
                        assert_eq!(actual.inner().as_ptr(), expected.inner().as_ptr());
                        assert_eq!(actual.offset(), expected.offset());
                    }
                    (None, None) => {}
                    _ => panic!("empty operand changed the other operand's representation"),
                }
            }
        }
    }

    #[test]
    #[cfg_attr(miri, ignore)] // Takes too long
    fn test_mask_algebra_with_offsets() {
        // Offsets and lengths that are not byte (or word) aligned on either side,
        // so the common prefix can start and end mid byte. Covers both the equal
        // and uneven length paths.
        let base: Vec<bool> = (0..600).map(|i| i % 7 == 0 || i % 3 == 1).collect();
        let other: Vec<bool> = (0..600).map(|i| i % 5 == 2 || i % 11 == 4).collect();
        let base = BooleanBuffer::from(base);
        let other = BooleanBuffer::from(other);

        for l_offset in [0, 1, 5, 8, 13, 64, 67] {
            for r_offset in [0, 1, 3, 8, 60, 64, 70] {
                for (l_len, r_len) in [
                    (0, 9),
                    (9, 0),
                    (1, 200),
                    (200, 1),
                    (63, 130),
                    (321, 65),
                    (0, 0),
                    (1, 1),
                    (63, 63),
                    (64, 64),
                    (200, 200),
                    (321, 321),
                ] {
                    let l = base.slice(l_offset, l_len);
                    let r = other.slice(r_offset, r_len);
                    let l_bits: Vec<bool> = l.iter().collect();
                    let r_bits: Vec<bool> = r.iter().collect();
                    let context =
                        format!("l_offset={l_offset} r_offset={r_offset} lens=({l_len},{r_len})");

                    assert_mask_eq(
                        &intersect_masks(&l, &r),
                        &expected_combined(&l_bits, &r_bits, |a, b| a && b),
                        &format!("intersect {context}"),
                    );
                    assert_mask_eq(
                        &union_masks(&l, &r),
                        &expected_combined(&l_bits, &r_bits, |a, b| a || b),
                        &format!("union {context}"),
                    );
                }
            }
        }
    }

    #[test]
    #[cfg_attr(miri, ignore)] // Takes too long
    fn test_mask_algebra_does_not_retain_backing_buffer() {
        // A short slice of a long mask must not keep the long allocation alive,
        // including when the other operand is empty and contributes nothing.
        // Comparing capacities rather than lengths, since a shallow slice reports a
        // short length while still holding the original allocation through its `Arc`.
        let long = BooleanBuffer::from((0..80_000).map(|i| i % 3 == 0).collect::<Vec<bool>>());
        assert!(long.inner().capacity() >= 10_000);

        for (l, r) in [
            (long.slice(5, 40), BooleanBuffer::new_unset(0)),
            (long.slice(5, 40), BooleanBuffer::new_set(7)),
            (BooleanBuffer::new_set(7), long.slice(5, 40)),
        ] {
            for combined in [intersect_masks(&l, &r), union_masks(&l, &r)] {
                assert!(
                    combined.inner().capacity() < long.inner().capacity(),
                    "result retained the original backing allocation"
                );
            }
        }
    }

    #[test]
    #[cfg_attr(miri, ignore)] // Takes too long
    fn test_mask_algebra_fuzz() {
        let mut rng = rng();
        for _ in 0..200 {
            let l_offset = rng.random_range(0..70);
            let r_offset = rng.random_range(0..70);
            let l_len = rng.random_range(0..300);
            // Bias towards equal lengths so that path is hit often
            let r_len = match rng.random_bool(0.25) {
                true => l_len,
                false => rng.random_range(0..300),
            };

            let l_bits: Vec<bool> = (0..l_offset + l_len)
                .map(|_| rng.random_bool(0.5))
                .collect();
            let r_bits: Vec<bool> = (0..r_offset + r_len)
                .map(|_| rng.random_bool(0.5))
                .collect();
            let l = BooleanBuffer::from(l_bits).slice(l_offset, l_len);
            let r = BooleanBuffer::from(r_bits).slice(r_offset, r_len);
            let l_bits: Vec<bool> = l.iter().collect();
            let r_bits: Vec<bool> = r.iter().collect();

            assert_mask_eq(
                &intersect_masks(&l, &r),
                &expected_combined(&l_bits, &r_bits, |a, b| a && b),
                "intersect",
            );
            assert_mask_eq(
                &union_masks(&l, &r),
                &expected_combined(&l_bits, &r_bits, |a, b| a || b),
                "union",
            );
        }
    }

    #[test]
    fn dense_composition_with_zero_words_and_offsets() {
        for mask_len in [8192, 8193, 65537] {
            for (mask_offset, inner_offset) in [(0, 0), (3, 5), (63, 65), (65, 129)] {
                for stride in [4, 100] {
                    let mask =
                        BooleanBuffer::from_iter((0..mask_offset + mask_len + 73).map(|i| {
                            !(mask_offset..mask_offset + mask_len).contains(&i)
                                || (i - mask_offset) % stride != 0
                        }))
                        .slice(mask_offset, mask_len);
                    let len = mask.count_set_bits();
                    for start in [0, 63, len / 2] {
                        for select_last in [false, true] {
                            let count = len / 10;
                            let inner =
                                BooleanBuffer::from_iter((0..inner_offset + len + 73).map(|i| {
                                    if !(inner_offset..inner_offset + len).contains(&i) {
                                        return true;
                                    }
                                    let j = i - inner_offset;
                                    (start..start + count).contains(&j)
                                        || select_last && j == len - 1
                                }))
                                .slice(inner_offset, len);
                            AndThenTest {
                                outer: mask.clone(),
                                inner,
                                dense_mask_expected: true,
                            }
                            .run();
                        }
                    }
                }
            }
        }
    }
}