perf: plan take operations through FilteredReadExec's range-read path

python/python/tests/test_mem_wal.py

@@ -111,7 +111,8 @@ def test_point_lookup_with_memtables(tmp_path):
     plan = planner.plan_lookup(pa.array([2], type=pa.int64()))
     assert plan.schema.names == ["id", "name"]
     assert plan.dataset_schema.names == ["id", "name"]
-    assert "Take" in plan.explain() or "Scan" in plan.explain()
+    explained = plan.explain()
+    assert "Take" in explained or "Scan" in explained or "LanceRead" in explained
     result = plan.to_table()
     assert len(result) == 1, "Expected exactly one row for id=2"
     assert result.column("name")[0].as_py() == "gen1_2", (

rust/lance-namespace-impls/src/dir.rs

@@ -12256,7 +12256,7 @@ mod tests {
             &plan_str,
             &[
                 "ProjectionExec: expr=[id@0 as id, name@2 as name",
-                "Take: columns=\"id, _rowid, (name)\"",
+                "projection=[name], source=stream(_rowid)",
                 "LanceRead: uri=",
                 "projection=[id]",
                 "row_id=true, row_addr=false",
@@ -12353,9 +12353,7 @@ mod tests {
                 "AnalyzeExec verbose=true",
                 "ProjectionExec: elapsed=",
                 "expr=[id@0 as id, name@2 as name",
-                "Take: elapsed=",
-                "columns=\"id, _rowid, (name)\"",
-                "CoalesceBatchesExec: elapsed=",
+                "projection=[name], source=stream(_rowid)",
                 "LanceRead: elapsed=",
                 "projection=[id]",
                 "row_id=true, row_addr=false",

rust/lance/src/dataset/mem_wal/memtable/flush.rs

@@ -2172,9 +2172,8 @@ mod tests {
         crate::utils::test::assert_plan_node_equals(
             plan,
             "ProjectionExec: expr=[id@2 as id, text@3 as text, _score@1 as _score]
-  Take: ...
-    CoalesceBatchesExec: ...
-      MatchQuery: column=text, query=hello",
+  LanceRead: ..., source=stream(_rowid)
+    MatchQuery: column=text, query=hello",
         )
         .await
         .unwrap();

rust/lance/src/dataset/write/merge_insert.rs

@@ -58,7 +58,9 @@ use crate::{
     },
     index::DatasetIndexInternalExt,
     io::exec::{
-        AddRowAddrExec, Planner, TakeExec, project,
+        AddRowAddrExec, Planner, TakeExec,
+        filtered_read::{FilteredReadExec, FilteredReadOptions},
+        project,
         scalar_index::{IndexLookup, MapIndexExec},
         utils::ReplayExec,
     },
@@ -734,14 +736,21 @@ impl MergeInsertJob {
             )?);
         }
 
-        // 4 - Take the mapped row ids
+        // 4 - Take the mapped row ids.  TakeExec stays for legacy storage
+        //     only: the v1 reader cannot serve a FilteredReadExec.
         let projection = self
             .dataset
             .empty_projection()
             .union_arrow_schema(schema.as_ref(), OnMissing::Error)?;
-        let mut target =
+        let mut target: Arc<dyn ExecutionPlan> = if self.dataset.is_legacy_storage() {
             Arc::new(TakeExec::try_new(self.dataset.clone(), index_mapper, projection)?.unwrap())
-                as Arc<dyn ExecutionPlan>;
+        } else {
+            Arc::new(FilteredReadExec::try_new(
+                self.dataset.clone(),
+                FilteredReadOptions::new(projection),
+                Some(index_mapper),
+            )?)
+        };
 
         // 5 - Take puts the row id and row addr at the beginning.  A full scan (used when there is
         //     no scalar index) puts the row id and addr at the end.  We need to match these up so

rust/lance/src/io/exec/count_pushdown.rs

@@ -147,6 +147,11 @@ fn try_rewrite(agg: &AggregateExec) -> DFResult<Option<Arc<dyn ExecutionPlan>>>
     };
 
     let options = filtered_read.options();
+    // We don't currently support count pushdown when the row selector
+    // is a row stream.
+    if filtered_read.row_stream_input().is_some() {
+        return Ok(None);
+    }
     // A refine filter is a residual the index couldn't fully evaluate — it
     // needs column data to apply, which we can't.
     if options.refine_filter.is_some() {