From 33a57f32c4472fc71d055f1e608a6d9a2f001086 Mon Sep 17 00:00:00 2001
From: Peter Bower <37089506+pbower@users.noreply.github.com>
Date: Tue, 13 Oct 2026 04:43:07 +0100
Subject: [PATCH 1/2] Update DataSelector to return Bitmask instead of
 Vec<usize>

---
 src/enums/array.rs                            |  15 +-
 src/structs/chunked/super_ndarray.rs          |  15 +-
 src/structs/field_array.rs                    |  14 +-
 src/structs/matrix.rs                         |  22 ++
 src/structs/ndarray.rs                        |  63 ++++-
 src/structs/table.rs                          |  14 +-
 src/structs/views/array_view.rs               |  17 +-
 .../views/chunked/super_ndarray_view.rs       |  18 +-
 src/structs/views/matrix_view.rs              |  44 +++-
 src/structs/views/ndarray_view.rs             |  18 +-
 src/structs/views/table_view.rs               |  12 +-
 src/traits/selection.rs                       | 228 ++++++++++++++----
 tests/test_selection.rs                       |  23 +-
 13 files changed, 361 insertions(+), 142 deletions(-)

diff --git a/src/enums/array.rs b/src/enums/array.rs
index 7d71363..4d18d90 100644
--- a/src/enums/array.rs
+++ b/src/enums/array.rs
@@ -7335,21 +7335,16 @@ impl crate::traits::selection::RowSelection for Array {
     /// For non-contiguous selections (index arrays), gathers into a new array.
     fn r<S: crate::traits::selection::DataSelector>(&self, selection: S) -> ArrayV {
         if selection.is_contiguous() {
-            // Contiguous selection (ranges): create a view
-            let indices = selection.resolve_indices(self.len());
-            if indices.is_empty() {
+            let (offset, len) = selection.resolve_range(self.len());
+            if len == 0 {
                 return ArrayV::new(self.clone(), 0, 0);
             }
-            let offset = indices[0];
-            let len = indices.len();
             ArrayV::new(self.clone(), offset, len)
         } else {
-            // Non-contiguous selection (index arrays): gather into new array
-            // For now, create a view over the whole array and use its gather
+            let mask = selection.resolve_mask(self.len());
             let full_view = ArrayV::new(self.clone(), 0, self.len());
-            let indices = selection.resolve_indices(self.len());
-            let gathered = full_view.gather_indices(&indices);
-            ArrayV::new(gathered, 0, indices.len())
+            let gathered = full_view.gather_mask(&mask);
+            ArrayV::new(gathered, 0, mask.count_ones())
         }
     }
 
diff --git a/src/structs/chunked/super_ndarray.rs b/src/structs/chunked/super_ndarray.rs
index 22be08b..122addf 100644
--- a/src/structs/chunked/super_ndarray.rs
+++ b/src/structs/chunked/super_ndarray.rs
@@ -44,7 +44,7 @@ use crate::structs::ndarray::NdArrayIter;
 #[cfg(feature = "views")]
 use crate::Vec64;
 #[cfg(all(feature = "views", feature = "select"))]
-use crate::structs::ndarray::gather_obs_impl;
+use crate::structs::ndarray::gather_obs_mask_impl;
 #[cfg(all(feature = "views", feature = "select"))]
 use crate::traits::selection::{AxisSelection, DataSelector, RowSelection};
 #[cfg(feature = "views")]
@@ -620,13 +620,13 @@ impl<T: Float> RowSelection for SuperNdArray<T> {
                 self.name.clone(),
             );
         }
-        let indices = selection.resolve_indices(self.n_obs());
         if selection.is_contiguous() {
-            let start = indices.first().copied().unwrap_or(0);
-            return self.slice(start, indices.len());
+            let (start, len) = selection.resolve_range(self.n_obs());
+            return self.slice(start, len);
         }
-        let gathered = gather_obs_impl(
-            &indices,
+        let mask = selection.resolve_mask(self.n_obs());
+        let gathered = gather_obs_mask_impl(
+            &mask,
             &self.shape(),
             Some(self.name.clone()),
             |idx| self.get(idx),
@@ -1001,10 +1001,11 @@ mod tests {
             ],
             "data",
         );
+        // Bitmask-based gathering yields observations in ascending index order.
         let v = snd.r(&[3, 0]);
         assert_eq!(v.n_slices(), 1);
         let vals: Vec<f64> = (&v).into_iter().collect();
-        assert_eq!(vals, vec![4.0, 1.0]);
+        assert_eq!(vals, vec![1.0, 4.0]);
     }
 
     #[test]
diff --git a/src/structs/field_array.rs b/src/structs/field_array.rs
index f711342..06c5ddd 100644
--- a/src/structs/field_array.rs
+++ b/src/structs/field_array.rs
@@ -702,18 +702,16 @@ impl RowSelection for FieldArray {
 
     fn r<S: DataSelector>(&self, selection: S) -> ArrayV {
         if selection.is_contiguous() {
-            // Contiguous selection (ranges): adjust offset and len
-            let indices = selection.resolve_indices(self.array.len());
-            if indices.is_empty() {
+            let (offset, len) = selection.resolve_range(self.array.len());
+            if len == 0 {
                 return ArrayV::new(self.array.clone(), 0, 0);
             }
-            ArrayV::new(self.array.clone(), indices[0], indices.len())
+            ArrayV::new(self.array.clone(), offset, len)
         } else {
-            // Non-contiguous selection (index arrays): gather into new array
+            let mask = selection.resolve_mask(self.array.len());
             let view = ArrayV::from(self.array.clone());
-            let indices = selection.resolve_indices(self.array.len());
-            let gathered_array = view.gather_indices(&indices);
-            ArrayV::new(gathered_array, 0, indices.len())
+            let gathered_array = view.gather_mask(&mask);
+            ArrayV::new(gathered_array, 0, mask.count_ones())
         }
     }
 
diff --git a/src/structs/matrix.rs b/src/structs/matrix.rs
index 0a7bc7a..84d6718 100644
--- a/src/structs/matrix.rs
+++ b/src/structs/matrix.rs
@@ -413,6 +413,28 @@ impl Matrix {
         dst
     }
 
+    /// Extract the rows at set mask bits into a new Matrix.
+    ///
+    /// The mask must have the same length as `n_rows`. Column stride alignment
+    /// is maintained in the result.
+    pub fn extract_rows_mask(&self, mask: &crate::Bitmask) -> Self {
+        assert_eq!(
+            mask.len(), self.n_rows,
+            "Matrix::extract_rows_mask: mask length {} does not match row count {}",
+            mask.len(), self.n_rows
+        );
+        let n_new = mask.count_ones();
+        let mut dst = Matrix::new(n_new, self.n_cols, self.name.clone());
+        for j in 0..self.n_cols {
+            let src_col = self.col(j);
+            let dst_col = dst.col_mut(j);
+            for (k, idx) in mask.iter_set().enumerate() {
+                dst_col[k] = src_col[idx];
+            }
+        }
+        dst
+    }
+
     /// Sets the matrix name.
     #[inline]
     pub fn set_name(&mut self, name: impl Into<String>) {
diff --git a/src/structs/ndarray.rs b/src/structs/ndarray.rs
index 0099154..54542a1 100644
--- a/src/structs/ndarray.rs
+++ b/src/structs/ndarray.rs
@@ -1474,17 +1474,17 @@ impl<T: Float> RowSelection for NdArray<T> {
     fn r<S: DataSelector>(&self, selection: S) -> NdArrayV<T> {
         assert!(self.ndim() > 0, "row selection requires an axis 0");
         let n_obs = self.shape()[0];
-        let indices = selection.resolve_indices(n_obs);
         if selection.is_contiguous() {
-            let start = indices.first().copied().unwrap_or(0);
-            let ranges: Vec<Range<usize>> = std::iter::once(start..start + indices.len())
+            let (start, len) = selection.resolve_range(n_obs);
+            let ranges: Vec<Range<usize>> = std::iter::once(start..start + len)
                 .chain(self.shape()[1..].iter().map(|&n| 0..n))
                 .collect();
             let refs: Vec<&dyn DataSelector> = ranges.iter().map(|r| r as _).collect();
             return self.slice(&refs);
         }
-        NdArrayV::from_ndarray(gather_obs_impl(
-            &indices,
+        let mask = selection.resolve_mask(n_obs);
+        NdArrayV::from_ndarray(gather_obs_mask_impl(
+            &mask,
             self.shape(),
             self.name.clone(),
             |idx| self.get(idx),
@@ -1501,6 +1501,7 @@ impl<T: Float> RowSelection for NdArray<T> {
 /// element through the provided accessor, so any stride layout gathers
 /// correctly.
 #[cfg(all(feature = "views", feature = "select"))]
+#[allow(dead_code)]
 pub(crate) fn gather_obs_impl<T: Float>(
     indices: &[usize],
     shape: &[usize],
@@ -1538,6 +1539,48 @@ pub(crate) fn gather_obs_impl<T: Float>(
     result
 }
 
+/// Mask-based variant of [`gather_obs_impl`]. Gathers axis-0 observations
+/// where the corresponding mask bit is set, without materialising an
+/// intermediate index list.
+#[cfg(all(feature = "views", feature = "select"))]
+pub(crate) fn gather_obs_mask_impl<T: Float>(
+    mask: &crate::Bitmask,
+    shape: &[usize],
+    name: Option<String>,
+    get: impl Fn(&[usize]) -> T,
+) -> NdArray<T> {
+    let n_selected = mask.count_ones();
+    let mut out_shape = shape.to_vec();
+    out_shape[0] = n_selected;
+    let total: usize = out_shape.iter().product();
+
+    let mut flat = Vec64::with_capacity(total);
+    let ndim = shape.len();
+    let mut idx = vec![0usize; ndim];
+    let inner_runs: usize = shape[1..].iter().product::<usize>().max(1);
+    for _ in 0..inner_runs {
+        for obs in mask.iter_set() {
+            idx[0] = obs;
+            flat.push(get(&idx));
+        }
+        // Advance the inner multi-index in column-major order.
+        let mut carry = true;
+        for d in 1..ndim {
+            if carry {
+                idx[d] += 1;
+                if idx[d] < shape[d] {
+                    carry = false;
+                } else {
+                    idx[d] = 0;
+                }
+            }
+        }
+    }
+    let mut result = NdArray::from_slice(&flat, &out_shape);
+    result.name = name;
+    result
+}
+
 // *** Bracket indexing: arr[col][row] ******************************
 
 /// `arr[i]` selects along the outermost stored axis.
@@ -2139,13 +2182,13 @@ mod tests {
     #[test]
     fn row_selection_gathers_indices() {
         let a = NdArray::from_slice(&[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], &[3, 2]);
+        // Bitmask-based gathering yields rows in ascending index order.
         let v = a.r(&[2, 0]);
         assert_eq!(v.shape(), &[2, 2]);
-        // Gathered rows follow selection order.
-        assert_eq!(v.get(&[0, 0]), 3.0);
-        assert_eq!(v.get(&[0, 1]), 6.0);
-        assert_eq!(v.get(&[1, 0]), 1.0);
-        assert_eq!(v.get(&[1, 1]), 4.0);
+        assert_eq!(v.get(&[0, 0]), 1.0);
+        assert_eq!(v.get(&[0, 1]), 4.0);
+        assert_eq!(v.get(&[1, 0]), 3.0);
+        assert_eq!(v.get(&[1, 1]), 6.0);
     }
 
     #[test]
diff --git a/src/structs/table.rs b/src/structs/table.rs
index e870415..6bd9148 100644
--- a/src/structs/table.rs
+++ b/src/structs/table.rs
@@ -1320,19 +1320,15 @@ impl RowSelection for Table {
 
     fn r<S: DataSelector>(&self, selection: S) -> TableV {
         if selection.is_contiguous() {
-            // Contiguous selection (ranges): create a properly windowed view
-            let indices = selection.resolve_indices(self.n_rows);
-            if indices.is_empty() {
+            let (offset, len) = selection.resolve_range(self.n_rows);
+            if len == 0 {
                 return TableV::from_table(self.clone(), 0, 0);
             }
-            let new_offset = indices[0];
-            let new_len = indices.len();
-            TableV::from_table(self.clone(), new_offset, new_len)
+            TableV::from_table(self.clone(), offset, len)
         } else {
-            // Non-contiguous selection (index arrays): materialise
-            let indices = selection.resolve_indices(self.n_rows);
+            let mask = selection.resolve_mask(self.n_rows);
             let table_v = TableV::from(self.clone());
-            let materialised_table = table_v.gather_rows(&indices);
+            let materialised_table = table_v.gather_rows_mask(&mask);
             TableV::from(materialised_table)
         }
     }
diff --git a/src/structs/views/array_view.rs b/src/structs/views/array_view.rs
index 87e596f..3651592 100644
--- a/src/structs/views/array_view.rs
+++ b/src/structs/views/array_view.rs
@@ -1961,19 +1961,16 @@ impl RowSelection for ArrayV {
 
     fn r<S: DataSelector>(&self, selection: S) -> ArrayV {
         if selection.is_contiguous() {
-            // Contiguous selection (ranges): adjust offset and len
-            let indices = selection.resolve_indices(self.len());
-            if indices.is_empty() {
+            let (start, len) = selection.resolve_range(self.len());
+            if len == 0 {
                 return ArrayV::new(self.array.clone(), self.offset, 0);
             }
-            let new_offset = self.offset + indices[0];
-            let new_len = indices.len();
-            ArrayV::new(self.array.clone(), new_offset, new_len)
+            let new_offset = self.offset + start;
+            ArrayV::new(self.array.clone(), new_offset, len)
         } else {
-            // Non-contiguous selection (index arrays): gather into new array
-            let indices = selection.resolve_indices(self.len());
-            let gathered_array = self.gather_indices(&indices);
-            ArrayV::new(gathered_array, 0, indices.len())
+            let mask = selection.resolve_mask(self.len());
+            let gathered_array = self.gather_mask(&mask);
+            ArrayV::new(gathered_array, 0, mask.count_ones())
         }
     }
 
diff --git a/src/structs/views/chunked/super_ndarray_view.rs b/src/structs/views/chunked/super_ndarray_view.rs
index 6919868..b6c650c 100644
--- a/src/structs/views/chunked/super_ndarray_view.rs
+++ b/src/structs/views/chunked/super_ndarray_view.rs
@@ -41,7 +41,7 @@ use crate::enums::shape_dim::ShapeDim;
 use crate::structs::chunked::super_ndarray::SuperNdArray;
 use crate::structs::ndarray::NdArray;
 #[cfg(feature = "select")]
-use crate::structs::ndarray::gather_obs_impl;
+use crate::structs::ndarray::gather_obs_mask_impl;
 use crate::structs::views::ndarray_view::NdArrayV;
 use crate::traits::concatenate::Concatenate;
 #[cfg(feature = "select")]
@@ -289,13 +289,13 @@ impl<T: Float> RowSelection for SuperNdArrayV<T> {
                 self.name.clone(),
             );
         }
-        let indices = selection.resolve_indices(self.n_obs());
         if selection.is_contiguous() {
-            let start = indices.first().copied().unwrap_or(0);
-            return self.slice(start, indices.len());
+            let (start, len) = selection.resolve_range(self.n_obs());
+            return self.slice(start, len);
         }
-        let gathered = gather_obs_impl(
-            &indices,
+        let mask = selection.resolve_mask(self.n_obs());
+        let gathered = gather_obs_mask_impl(
+            &mask,
             &self.shape(),
             Some(self.name.clone()),
             |idx| self.get(idx),
@@ -609,11 +609,11 @@ mod tests {
         let sub = v.r(1..4);
         assert_eq!(sub.n_obs(), 3);
         assert_eq!(sub.get(&[0, 0]), 2.0);
-        // Index selection gathers in order.
+        // Bitmask-based gathering yields observations in ascending index order.
         let picked = v.r(&[4, 0]);
         assert_eq!(picked.n_slices(), 1);
-        assert_eq!(picked.get(&[0, 1]), 50.0);
-        assert_eq!(picked.get(&[1, 0]), 1.0);
+        assert_eq!(picked.get(&[0, 0]), 1.0);
+        assert_eq!(picked.get(&[1, 1]), 50.0);
     }
 
     #[test]
diff --git a/src/structs/views/matrix_view.rs b/src/structs/views/matrix_view.rs
index e36f9eb..e50ba07 100644
--- a/src/structs/views/matrix_view.rs
+++ b/src/structs/views/matrix_view.rs
@@ -340,6 +340,27 @@ impl MatrixV {
             .collect();
         self.matrix.extract_rows(&absolute)
     }
+
+    /// Gathers the rows at set mask bits into a new owned `Matrix`.
+    ///
+    /// The mask is relative to this view's window and must match its row count.
+    pub fn gather_rows_mask(&self, mask: &crate::Bitmask) -> Matrix {
+        assert_eq!(
+            mask.len(), self.len,
+            "MatrixV::gather_rows_mask: mask length {} does not match row count {}",
+            mask.len(), self.len
+        );
+        let n_new = mask.count_ones();
+        let mut dst = Matrix::new(n_new, self.n_cols(), self.matrix.name.clone());
+        for j in 0..self.n_cols() {
+            let src_col = self.col(j);
+            let dst_col = dst.col_mut(j);
+            for (k, idx) in mask.iter_set().enumerate() {
+                dst_col[k] = src_col[idx];
+            }
+        }
+        dst
+    }
 }
 
 impl Shape for MatrixV {
@@ -456,12 +477,12 @@ impl RowSelection for Matrix {
     type View = MatrixV;
 
     fn r<S: DataSelector>(&self, selection: S) -> MatrixV {
-        let indices = selection.resolve_indices(self.n_rows);
         if selection.is_contiguous() {
-            let start = indices.first().copied().unwrap_or(0);
-            return MatrixV::from_matrix(self.clone(), start, indices.len());
+            let (start, len) = selection.resolve_range(self.n_rows);
+            return MatrixV::from_matrix(self.clone(), start, len);
         }
-        MatrixV::from(self.extract_rows(&indices))
+        let mask = selection.resolve_mask(self.n_rows);
+        MatrixV::from(self.extract_rows_mask(&mask))
     }
 
     fn get_row_count(&self) -> usize {
@@ -479,12 +500,12 @@ impl RowSelection for MatrixV {
     type View = MatrixV;
 
     fn r<S: DataSelector>(&self, selection: S) -> MatrixV {
-        let indices = selection.resolve_indices(self.len);
         if selection.is_contiguous() {
-            let start = indices.first().copied().unwrap_or(0);
-            return self.from_self(start, indices.len());
+            let (start, len) = selection.resolve_range(self.len);
+            return self.from_self(start, len);
         }
-        MatrixV::from(self.gather_rows(&indices))
+        let mask = selection.resolve_mask(self.len);
+        MatrixV::from(self.gather_rows_mask(&mask))
     }
 
     fn get_row_count(&self) -> usize {
@@ -803,7 +824,8 @@ mod tests {
         let selected = sample().r(&[3usize, 0][..]);
         assert_eq!(selected.n_rows(), 2);
         assert!(selected.spans_backing());
-        assert_eq!(selected.col(1), &[13.0, 10.0]);
+        // Bitmask-based gathering yields rows in ascending index order.
+        assert_eq!(selected.col(1), &[10.0, 13.0]);
     }
 
     #[cfg(feature = "select")]
@@ -816,9 +838,9 @@ mod tests {
         assert_eq!(chained.col(0), &[2.0, 3.0]);
         assert_eq!(chained.get_row_count(), 2);
 
-        // Non-contiguous selections gather rows into a new matrix.
+        // Non-contiguous selections gather rows in ascending index order.
         let picked = v.r(&[3usize, 1][..]);
-        assert_eq!(picked.col(0), &[4.0, 2.0]);
+        assert_eq!(picked.col(0), &[2.0, 4.0]);
     }
 
     #[cfg(feature = "size")]
diff --git a/src/structs/views/ndarray_view.rs b/src/structs/views/ndarray_view.rs
index b800a54..b4738a8 100644
--- a/src/structs/views/ndarray_view.rs
+++ b/src/structs/views/ndarray_view.rs
@@ -15,7 +15,7 @@ use crate::enums::shape_dim::ShapeDim;
 use crate::structs::ndarray::{NdArray, NdArrayIter, NdDims, offset_of_impl};
 use crate::traits::print::print_ndarray_body;
 #[cfg(feature = "select")]
-use crate::structs::ndarray::gather_obs_impl;
+use crate::structs::ndarray::gather_obs_mask_impl;
 #[cfg(feature = "select")]
 use crate::traits::selection::{AxisSelection, DataSelector, RowSelection};
 #[cfg(feature = "select")]
@@ -443,17 +443,17 @@ impl<T: Float> RowSelection for NdArrayV<T> {
     fn r<S: DataSelector>(&self, selection: S) -> NdArrayV<T> {
         assert!(self.ndim() > 0, "row selection requires an axis 0");
         let n_obs = self.dims.shape()[0];
-        let indices = selection.resolve_indices(n_obs);
         if selection.is_contiguous() {
-            let start = indices.first().copied().unwrap_or(0);
-            let ranges: Vec<Range<usize>> = std::iter::once(start..start + indices.len())
+            let (start, len) = selection.resolve_range(n_obs);
+            let ranges: Vec<Range<usize>> = std::iter::once(start..start + len)
                 .chain(self.dims.shape()[1..].iter().map(|&n| 0..n))
                 .collect();
             let refs: Vec<&dyn DataSelector> = ranges.iter().map(|r| r as _).collect();
             return self.slice(&refs);
         }
-        NdArrayV::from_ndarray(gather_obs_impl(
-            &indices,
+        let mask = selection.resolve_mask(n_obs);
+        NdArrayV::from_ndarray(gather_obs_mask_impl(
+            &mask,
             self.dims.shape(),
             self.source.name.clone(),
             |idx| self.get(idx),
@@ -731,10 +731,10 @@ mod tests {
         let sub = v.r(1..3);
         assert_eq!(sub.shape(), &[2, 2]);
         assert_eq!(sub.get(&[0, 1]), 5.0);
-        // Index arrays gather in order, and the source is unaffected.
+        // Bitmask-based gathering yields rows in ascending index order.
         let picked = v.r(&[2, 0]);
-        assert_eq!(picked.get(&[0, 0]), 3.0);
-        assert_eq!(picked.get(&[1, 1]), 4.0);
+        assert_eq!(picked.get(&[0, 0]), 1.0);
+        assert_eq!(picked.get(&[1, 1]), 6.0);
     }
 
     #[test]
diff --git a/src/structs/views/table_view.rs b/src/structs/views/table_view.rs
index 297c3c9..89dc103 100644
--- a/src/structs/views/table_view.rs
+++ b/src/structs/views/table_view.rs
@@ -818,16 +818,14 @@ impl RowSelection for TableV {
 
     fn r<S: DataSelector>(&self, selection: S) -> TableV {
         if selection.is_contiguous() {
-            // Contiguous selection (ranges): derive a properly windowed subview
-            let indices = selection.resolve_indices(self.len);
-            if indices.is_empty() {
+            let (start, len) = selection.resolve_range(self.len);
+            if len == 0 {
                 return self.from_self(0, 0);
             }
-            self.from_self(indices[0], indices.len())
+            self.from_self(start, len)
         } else {
-            // Non-contiguous selection (index arrays): materialise into Table
-            let indices = selection.resolve_indices(self.len);
-            let materialised_table = self.gather_rows(&indices);
+            let mask = selection.resolve_mask(self.len);
+            let materialised_table = self.gather_rows_mask(&mask);
             TableV::from(materialised_table)
         }
     }
diff --git a/src/traits/selection.rs b/src/traits/selection.rs
index 6ccba37..a7059c2 100644
--- a/src/traits/selection.rs
+++ b/src/traits/selection.rs
@@ -26,7 +26,7 @@
 //! - **Selection3D**: Future Extension for 3D selection
 //! - **Selection4D**: Future Extension for 4D selection
 
-use crate::Field;
+use crate::{Bitmask, Field};
 use std::ops::{Range, RangeFrom, RangeFull, RangeInclusive, RangeTo};
 use std::sync::Arc;
 
@@ -47,8 +47,17 @@ pub trait FieldSelector {
 
 /// Trait for types that can specify a data selection (index-based dimension)
 pub trait DataSelector {
-    /// Resolve this selection to indices (within the given count)
-    fn resolve_indices(&self, count: usize) -> Vec<usize>;
+    /// Resolve this selection to a bitmask with the selected positions set.
+    fn resolve_mask(&self, count: usize) -> Bitmask;
+
+    /// Resolve the start offset and length of a contiguous selection.
+    ///
+    /// Only valid when [`is_contiguous`](Self::is_contiguous) returns true.
+    /// The default panics because non-contiguous selectors have no single
+    /// range representation.
+    fn resolve_range(&self, _count: usize) -> (usize, usize) {
+        panic!("resolve_range called on a non-contiguous selector")
+    }
 
     /// Returns true if this selector represents a contiguous range.
     /// Range types (Range, RangeFrom, etc.) return true.
@@ -64,18 +73,14 @@ pub trait DataSelector {
     /// are rejected - gather through [`RowSelection::r`] instead.
     #[cfg(feature = "ndarray")]
     fn resolve_axis(&self, dim_size: usize) -> (usize, usize, bool) {
-        // usize::MAX so out-of-range indices survive to the bounds check
-        // below rather than being silently filtered.
-        let indices = self.resolve_indices(usize::MAX);
+        let mask = self.resolve_mask(dim_size);
+        let n = mask.count_ones();
         assert_eq!(
-            indices.len(), 1,
+            n, 1,
             "axis selection: index arrays take a single index or a contiguous range; gather with r() instead"
         );
-        assert!(
-            indices[0] < dim_size,
-            "axis selection: index {} out of bounds (size {})", indices[0], dim_size
-        );
-        (indices[0], indices[0] + 1, true)
+        let idx = mask.iter_set().next().unwrap();
+        (idx, idx + 1, true)
     }
 
     /// Produce an owned version of this selector.
@@ -430,12 +435,16 @@ impl FieldSelector for Box<dyn FieldSelector + Send + Sync> {
 
 /// Single data index
 impl DataSelector for usize {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        let mut mask = Bitmask::new_set_all(count, false);
         if *self < count {
-            vec![*self]
-        } else {
-            Vec::new()
+            mask.set_true(*self);
         }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        if *self < count { (*self, 1) } else { (0, 0) }
     }
 
     /// A single index is a contiguous window of length one.
@@ -460,11 +469,19 @@ impl DataSelector for usize {
 /// Single data index from a plain integer literal. Negative values
 /// resolve to nothing.
 impl DataSelector for i32 {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        let mut mask = Bitmask::new_set_all(count, false);
+        if *self >= 0 && (*self as usize) < count {
+            mask.set_true(*self as usize);
+        }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
         if *self >= 0 && (*self as usize) < count {
-            vec![*self as usize]
+            (*self as usize, 1)
         } else {
-            Vec::new()
+            (0, 0)
         }
     }
 
@@ -489,8 +506,14 @@ impl DataSelector for i32 {
 
 /// Multiple data indices
 impl DataSelector for &[usize] {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
-        self.iter().copied().filter(|&idx| idx < count).collect()
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        let mut mask = Bitmask::new_set_all(count, false);
+        for &idx in self.iter() {
+            if idx < count {
+                mask.set_true(idx);
+            }
+        }
+        mask
     }
 
     fn to_owned(&self) -> Box<dyn DataSelector + Send + Sync> {
@@ -500,8 +523,14 @@ impl DataSelector for &[usize] {
 
 /// Multiple data indices (array reference)
 impl<const N: usize> DataSelector for &[usize; N] {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
-        self.iter().copied().filter(|&idx| idx < count).collect()
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        let mut mask = Bitmask::new_set_all(count, false);
+        for &idx in self.iter() {
+            if idx < count {
+                mask.set_true(idx);
+            }
+        }
+        mask
     }
 
     fn to_owned(&self) -> Box<dyn DataSelector + Send + Sync> {
@@ -511,8 +540,14 @@ impl<const N: usize> DataSelector for &[usize; N] {
 
 /// Multiple data indices (Vec)
 impl DataSelector for Vec<usize> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
-        self.iter().copied().filter(|&idx| idx < count).collect()
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        let mut mask = Bitmask::new_set_all(count, false);
+        for &idx in self.iter() {
+            if idx < count {
+                mask.set_true(idx);
+            }
+        }
+        mask
     }
 
     fn to_owned(&self) -> Box<dyn DataSelector + Send + Sync> {
@@ -522,9 +557,18 @@ impl DataSelector for Vec<usize> {
 
 /// Data range selection
 impl DataSelector for Range<usize> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        let mut mask = Bitmask::new_set_all(count, false);
         let end = self.end.min(count);
-        (self.start..end).collect()
+        if self.start < end {
+            mask.set_range(self.start, end, true);
+        }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        let end = self.end.min(count);
+        if self.start < end { (self.start, end - self.start) } else { (0, 0) }
     }
 
     fn is_contiguous(&self) -> bool {
@@ -549,10 +593,20 @@ impl DataSelector for Range<usize> {
 /// Data range selection from plain integer literals. Negative bounds
 /// clamp to zero.
 impl DataSelector for Range<i32> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        let mut mask = Bitmask::new_set_all(count, false);
         let start = self.start.max(0) as usize;
         let end = (self.end.max(0) as usize).min(count);
-        (start..end).collect()
+        if start < end {
+            mask.set_range(start, end, true);
+        }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        let start = self.start.max(0) as usize;
+        let end = (self.end.max(0) as usize).min(count);
+        if start < end { (start, end - start) } else { (0, 0) }
     }
 
     fn is_contiguous(&self) -> bool {
@@ -576,8 +630,16 @@ impl DataSelector for Range<i32> {
 
 /// Data range from selection
 impl DataSelector for RangeFrom<usize> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
-        (self.start..count).collect()
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        let mut mask = Bitmask::new_set_all(count, false);
+        if self.start < count {
+            mask.set_range(self.start, count, true);
+        }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        if self.start < count { (self.start, count - self.start) } else { (0, 0) }
     }
 
     fn is_contiguous(&self) -> bool {
@@ -600,9 +662,18 @@ impl DataSelector for RangeFrom<usize> {
 
 /// Data range from selection from a plain integer literal.
 impl DataSelector for RangeFrom<i32> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        let start = self.start.max(0) as usize;
+        let mut mask = Bitmask::new_set_all(count, false);
+        if start < count {
+            mask.set_range(start, count, true);
+        }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
         let start = self.start.max(0) as usize;
-        (start..count).collect()
+        if start < count { (start, count - start) } else { (0, 0) }
     }
 
     fn is_contiguous(&self) -> bool {
@@ -625,9 +696,18 @@ impl DataSelector for RangeFrom<i32> {
 
 /// Data range to selection
 impl DataSelector for RangeTo<usize> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
         let end = self.end.min(count);
-        (0..end).collect()
+        let mut mask = Bitmask::new_set_all(count, false);
+        if end > 0 {
+            mask.set_range(0, end, true);
+        }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        let end = self.end.min(count);
+        (0, end)
     }
 
     fn is_contiguous(&self) -> bool {
@@ -651,9 +731,18 @@ impl DataSelector for RangeTo<usize> {
 /// Data range to selection from a plain integer literal. Negative bounds
 /// clamp to zero.
 impl DataSelector for RangeTo<i32> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
         let end = (self.end.max(0) as usize).min(count);
-        (0..end).collect()
+        let mut mask = Bitmask::new_set_all(count, false);
+        if end > 0 {
+            mask.set_range(0, end, true);
+        }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        let end = (self.end.max(0) as usize).min(count);
+        (0, end)
     }
 
     fn is_contiguous(&self) -> bool {
@@ -676,8 +765,12 @@ impl DataSelector for RangeTo<i32> {
 
 /// Data full range selection
 impl DataSelector for RangeFull {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
-        (0..count).collect()
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        Bitmask::new_set_all(count, true)
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        (0, count)
     }
 
     fn is_contiguous(&self) -> bool {
@@ -696,10 +789,20 @@ impl DataSelector for RangeFull {
 
 /// Data inclusive range selection
 impl DataSelector for RangeInclusive<usize> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
         let start = *self.start();
         let end = (*self.end() + 1).min(count);
-        (start..end).collect()
+        let mut mask = Bitmask::new_set_all(count, false);
+        if start < end {
+            mask.set_range(start, end, true);
+        }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        let start = *self.start();
+        let end = (*self.end() + 1).min(count);
+        if start < end { (start, end - start) } else { (0, 0) }
     }
 
     fn is_contiguous(&self) -> bool {
@@ -724,13 +827,26 @@ impl DataSelector for RangeInclusive<usize> {
 /// Data inclusive range selection from plain integer literals. Negative
 /// bounds clamp to zero.
 impl DataSelector for RangeInclusive<i32> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
         if *self.end() < 0 {
-            return Vec::new();
+            return Bitmask::new_set_all(count, false);
         }
         let start = (*self.start()).max(0) as usize;
         let end = (*self.end() as usize + 1).min(count);
-        (start..end).collect()
+        let mut mask = Bitmask::new_set_all(count, false);
+        if start < end {
+            mask.set_range(start, end, true);
+        }
+        mask
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        if *self.end() < 0 {
+            return (0, 0);
+        }
+        let start = (*self.start()).max(0) as usize;
+        let end = (*self.end() as usize + 1).min(count);
+        if start < end { (start, end - start) } else { (0, 0) }
     }
 
     fn is_contiguous(&self) -> bool {
@@ -752,10 +868,30 @@ impl DataSelector for RangeInclusive<i32> {
     }
 }
 
+/// Bitmask as a data selector, returned as-is.
+impl DataSelector for Bitmask {
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        assert_eq!(
+            self.len(), count,
+            "Bitmask::resolve_mask: mask length {} does not match count {}",
+            self.len(), count
+        );
+        self.to_owned_copy()
+    }
+
+    fn to_owned(&self) -> Box<dyn DataSelector + Send + Sync> {
+        Box::new(self.to_owned_copy())
+    }
+}
+
 /// Boxed data selector for owned selection
 impl DataSelector for Box<dyn DataSelector + Send + Sync> {
-    fn resolve_indices(&self, count: usize) -> Vec<usize> {
-        (**self).resolve_indices(count)
+    fn resolve_mask(&self, count: usize) -> Bitmask {
+        (**self).resolve_mask(count)
+    }
+
+    fn resolve_range(&self, count: usize) -> (usize, usize) {
+        (**self).resolve_range(count)
     }
 
     fn is_contiguous(&self) -> bool {
diff --git a/tests/test_selection.rs b/tests/test_selection.rs
index 48bc280..a85a600 100644
--- a/tests/test_selection.rs
+++ b/tests/test_selection.rs
@@ -294,18 +294,29 @@ fn create_test_table() -> Table {
 
 use minarrow::traits::selection::DataSelector;
 
-/// Asserts that the owned selector matches the original on indices and contiguity.
+/// Asserts that the owned selector produces the same mask and contiguity as
+/// the original.
 fn assert_owned_matches<S: DataSelector>(selector: S, count: usize) {
     let owned = selector.to_owned();
-    assert_eq!(owned.resolve_indices(count), selector.resolve_indices(count));
+    let orig_mask = selector.resolve_mask(count);
+    let owned_mask = owned.resolve_mask(count);
+    assert_eq!(orig_mask.len(), owned_mask.len());
+    for i in 0..count {
+        assert_eq!(orig_mask.get(i), owned_mask.get(i), "bit {} differs", i);
+    }
     assert_eq!(owned.is_contiguous(), selector.is_contiguous());
 }
 
+/// Collects the set-bit indices from a mask for assertion comparisons.
+fn mask_indices(mask: &minarrow::Bitmask) -> Vec<usize> {
+    mask.iter_set().collect()
+}
+
 #[test]
 fn test_data_selector_to_owned_single_index() {
     assert_owned_matches(2usize, 5);
     assert_owned_matches(2i32, 5);
-    assert_eq!(DataSelector::to_owned(&3usize).resolve_indices(5), vec![3]);
+    assert_eq!(mask_indices(&DataSelector::to_owned(&3usize).resolve_mask(5)), vec![3]);
 }
 
 #[test]
@@ -334,16 +345,16 @@ fn test_data_selector_to_owned_ranges() {
 #[test]
 fn test_data_selector_boxed_impl_forwards() {
     let boxed: Box<dyn DataSelector + Send + Sync> = Box::new(1usize..4);
-    assert_eq!(boxed.resolve_indices(5), vec![1, 2, 3]);
+    assert_eq!(mask_indices(&boxed.resolve_mask(5)), vec![1, 2, 3]);
     assert!(boxed.is_contiguous());
 
     let boxed_list: Box<dyn DataSelector + Send + Sync> = Box::new(vec![0usize, 2]);
-    assert_eq!(boxed_list.resolve_indices(5), vec![0, 2]);
+    assert_eq!(mask_indices(&boxed_list.resolve_mask(5)), vec![0, 2]);
     assert!(!boxed_list.is_contiguous());
 
     // A boxed selector clones into another boxed selector with the same behaviour.
     let again = DataSelector::to_owned(&boxed);
-    assert_eq!(again.resolve_indices(5), vec![1, 2, 3]);
+    assert_eq!(mask_indices(&again.resolve_mask(5)), vec![1, 2, 3]);
     assert!(again.is_contiguous());
 }
 

From 3f3c1a7323061d4ceeb93ac8d5c17d8a6730ac5f Mon Sep 17 00:00:00 2001
From: Peter Bower <37089506+pbower@users.noreply.github.com>
Date: Tue, 13 Oct 2026 05:22:13 +0100
Subject: [PATCH 2/2] Fix imports

---
 src/structs/matrix.rs            | 10 +++++-----
 src/structs/ndarray.rs           |  4 ++--
 src/structs/views/matrix_view.rs |  7 ++++---
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/src/structs/matrix.rs b/src/structs/matrix.rs
index 84d6718..9e4afa7 100644
--- a/src/structs/matrix.rs
+++ b/src/structs/matrix.rs
@@ -27,7 +27,7 @@ use crate::enums::shape_dim::ShapeDim;
 use crate::structs::buffer::Buffer;
 use crate::traits::print::print_float_grid;
 use crate::traits::{concatenate::Concatenate, shape::Shape};
-use crate::{Array, Field, FieldArray, FloatArray, NumericArray, Table, Vec64};
+use crate::{Array, Bitmask, Field, FieldArray, FloatArray, NumericArray, Table, Vec64};
 #[cfg(feature = "views")]
 use crate::{SharedBuffer, TableV};
 
@@ -400,7 +400,7 @@ impl Matrix {
 
     /// Extract rows by index into a new Matrix.
     /// Column stride alignment is maintained in the result.
-    pub fn extract_rows(&self, indices: &[usize]) -> Self {
+    pub fn gather_rows(&self, indices: &[usize]) -> Self {
         let n_new = indices.len();
         let mut dst = Matrix::new(n_new, self.n_cols, self.name.clone());
         for j in 0..self.n_cols {
@@ -417,10 +417,10 @@ impl Matrix {
     ///
     /// The mask must have the same length as `n_rows`. Column stride alignment
     /// is maintained in the result.
-    pub fn extract_rows_mask(&self, mask: &crate::Bitmask) -> Self {
+    pub fn gather_rows_mask(&self, mask: &Bitmask) -> Self {
         assert_eq!(
             mask.len(), self.n_rows,
-            "Matrix::extract_rows_mask: mask length {} does not match row count {}",
+            "Matrix::gather_rows_mask: mask length {} does not match row count {}",
             mask.len(), self.n_rows
         );
         let n_new = mask.count_ones();
@@ -1235,7 +1235,7 @@ mod try_as_matrix_zc_tests {
             &mut table.cols[0].array
         {
             let mut fa = (**arc_fa).clone();
-            let mut mask = crate::Bitmask::new_set_all(3, true);
+            let mut mask = Bitmask::new_set_all(3, true);
             mask.set(0, false);
             fa.null_mask = Some(mask);
             *arc_fa = std::sync::Arc::new(fa);
diff --git a/src/structs/ndarray.rs b/src/structs/ndarray.rs
index 54542a1..08abe63 100644
--- a/src/structs/ndarray.rs
+++ b/src/structs/ndarray.rs
@@ -77,7 +77,7 @@ use crate::traits::print::print_ndarray_body;
 use crate::traits::selection::{AxisSelection, DataSelector, RowSelection};
 use crate::traits::type_unions::Float;
 use crate::traits::{concatenate::Concatenate, shape::Shape};
-use crate::{Array, ArrowType, Field, FieldArray, FloatArray, NumericArray, Table, Vec64};
+use crate::{Array, ArrowType, Bitmask, Field, FieldArray, FloatArray, NumericArray, Table, Vec64};
 
 #[cfg(feature = "matrix")]
 use crate::structs::matrix::{Matrix, aligned_stride};
@@ -1544,7 +1544,7 @@ pub(crate) fn gather_obs_impl<T: Float>(
 /// intermediate index list.
 #[cfg(all(feature = "views", feature = "select"))]
 pub(crate) fn gather_obs_mask_impl<T: Float>(
-    mask: &crate::Bitmask,
+    mask: &Bitmask,
     shape: &[usize],
     name: Option<String>,
     get: impl Fn(&[usize]) -> T,
diff --git a/src/structs/views/matrix_view.rs b/src/structs/views/matrix_view.rs
index e50ba07..1c93cae 100644
--- a/src/structs/views/matrix_view.rs
+++ b/src/structs/views/matrix_view.rs
@@ -62,6 +62,7 @@ use std::sync::Arc;
 use crate::aliases::MatrixVT;
 use crate::enums::error::MinarrowError;
 use crate::enums::shape_dim::ShapeDim;
+use Bitmask;
 use crate::structs::matrix::Matrix;
 use crate::traits::concatenate::Concatenate;
 use crate::traits::print::print_float_grid;
@@ -338,13 +339,13 @@ impl MatrixV {
                 self.offset + i
             })
             .collect();
-        self.matrix.extract_rows(&absolute)
+        self.matrix.gather_rows(&absolute)
     }
 
     /// Gathers the rows at set mask bits into a new owned `Matrix`.
     ///
     /// The mask is relative to this view's window and must match its row count.
-    pub fn gather_rows_mask(&self, mask: &crate::Bitmask) -> Matrix {
+    pub fn gather_rows_mask(&self, mask: &Bitmask) -> Matrix {
         assert_eq!(
             mask.len(), self.len,
             "MatrixV::gather_rows_mask: mask length {} does not match row count {}",
@@ -482,7 +483,7 @@ impl RowSelection for Matrix {
             return MatrixV::from_matrix(self.clone(), start, len);
         }
         let mask = selection.resolve_mask(self.n_rows);
-        MatrixV::from(self.extract_rows_mask(&mask))
+        MatrixV::from(self.gather_rows_mask(&mask))
     }
 
     fn get_row_count(&self) -> usize {