fix

Signed-off-by: Mikhail Kot <mikhail@spiraldb.com>

Author: myrrc

benchmarks/lance-bench/src/random_access.rs

@@ -82,6 +82,7 @@ pub async fn nested_structs_lance() -> anyhow::Result<PathBuf> {
 pub struct LanceRandomAccessor {
     name: String,
     dataset: Dataset,
+    projection: ProjectionRequest,
 }
 
 impl LanceRandomAccessor {
@@ -92,9 +93,11 @@ impl LanceRandomAccessor {
                 .ok_or_else(|| anyhow!("Invalid dataset path"))?,
         )
         .await?;
+        let projection = ProjectionRequest::from_schema(dataset.schema().clone());
         Ok(Self {
             name: name.into(),
             dataset,
+            projection,
         })
     }
 }
@@ -110,8 +113,7 @@ impl RandomAccessor for LanceRandomAccessor {
     }
 
     async fn take(&self, indices: &[u64]) -> anyhow::Result<usize> {
-        let projection = ProjectionRequest::from_schema(self.dataset.schema().clone());
-        let result = self.dataset.take(indices, projection).await?;
+        let result = self.dataset.take(indices, self.projection.clone()).await?;
         Ok(result.num_rows())
     }
 }

vortex-layout/src/scan/repeated_scan.rs

@@ -17,10 +17,12 @@ use vortex_array::iter::ArrayIterator;
 use vortex_array::iter::ArrayIteratorAdapter;
 use vortex_array::stream::ArrayStream;
 use vortex_array::stream::ArrayStreamAdapter;
+use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
 use vortex_io::runtime::BlockingRuntime;
 use vortex_io::session::RuntimeSessionExt;
 use vortex_scan::selection::Selection;
+use vortex_scan::selection::selection_roaring_intersects;
 use vortex_session::VortexSession;
 use vortex_utils::parallelism::get_available_parallelism;
 
@@ -130,7 +132,17 @@ impl<A: 'static + Send> RepeatedScan<A> {
             mapper: Arc::clone(&self.map_fn),
         });
 
+        let selection_range: Option<Range<u64>> = match &self.selection {
+            Selection::IncludeByIndex(buf) if !buf.is_empty() => {
+                Some(buf[0]..buf[buf.len() - 1] + 1)
+            }
+            Selection::IncludeRoaring(roaring) if !roaring.is_empty() => {
+                Some(roaring.min().vortex_expect("empty")..roaring.max().vortex_expect("empty") + 1)
+            }
+            _ => None,
+        };
         let row_range = intersect_ranges(self.row_range.as_ref(), row_range);
+        let row_range = intersect_ranges(row_range.as_ref(), selection_range);
 
         let ranges = match &self.splits {
             Splits::Natural(btree_set) => {
@@ -167,17 +179,33 @@ impl<A: 'static + Send> RepeatedScan<A> {
 
         let mut limit = self.limit;
         let mut tasks = Vec::new();
+        let mut cursor = 0usize;
 
         for range in ranges {
             if range.start >= range.end {
                 continue;
             }
 
-            if limit.is_some_and(|l| l == 0) {
-                break;
+            let selection_intersects = match &self.selection {
+                Selection::IncludeByIndex(buf) => {
+                    // "ranges" and "buf" are sorted so we can use two pointers
+                    // to intersect.
+                    while cursor < buf.len() && buf[cursor] < range.start {
+                        cursor += 1;
+                    }
+                    cursor < buf.len() && buf[cursor] < range.end
+                }
+                Selection::IncludeRoaring(roaring) => selection_roaring_intersects(roaring, &range),
+                _ => true,
+            };
+            if !selection_intersects {
+                continue;
             }
 
             tasks.push(split_exec(Arc::clone(&ctx), range, limit.as_mut())?);
+            if limit.is_some_and(|l| l == 0) {
+                break;
+            }
         }
 
         Ok(tasks)

vortex-scan/src/selection.rs

@@ -152,6 +152,25 @@ impl Selection {
     }
 }
 
+#[cfg(not(target_arch = "wasm32"))]
+/// true if "roaring" contains an element that's in "range"
+pub fn selection_roaring_intersects(roaring: &roaring::RoaringTreemap, range: &Range<u64>) -> bool {
+    roaring.range_cardinality(range.clone()) > 0
+}
+
+// vortex-web's wasm has roaring 0.11.3 which doesn't have range_cardinality
+#[cfg(target_arch = "wasm32")]
+/// true if "roaring" contains an element that's in "range"
+pub fn selection_roaring_intersects(roaring: &roaring::RoaringTreemap, range: &Range<u64>) -> bool {
+    let before_end = roaring.rank(range.end - 1);
+    let before_start = if range.start > 0 {
+        roaring.rank(range.start - 1)
+    } else {
+        0
+    };
+    before_end > before_start
+}
+
 /// Find the positional range within row_indices that covers all rows in the given range.
 fn indices_range(range: &Range<u64>, row_indices: &[u64]) -> Option<Range<usize>> {
     if row_indices.first().is_some_and(|&first| first >= range.end)