Update DataSelector to return Bitmask instead of Vec<usize>

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())
         }
     }
 

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]

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())
         }
     }
 

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 {
@@ -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 gather_rows_mask(&self, mask: &Bitmask) -> Self {
+        assert_eq!(
+            mask.len(), self.n_rows,
+            "Matrix::gather_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>) {
@@ -1213,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);

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};
@@ -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: &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]

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)
         }
     }

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())
         }
     }
 

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]