diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index a72a7613241..f2fe185a49e 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -4504,16 +4504,20 @@ def test_late_materialization_param(tmp_path: Path): ) filt = "filter % 2 == 0" - assert "(values)" in dataset.scanner( - filter=filt, late_materialization=None - ).explain_plan(True) + # A late-materialized column is fetched by a row-stream read + # (`projection=[values], source=stream`); an eager column appears in the + # scan projection + late = "projection=[values], source=stream" + assert late in dataset.scanner(filter=filt, late_materialization=None).explain_plan( + True + ) assert ", values" in dataset.scanner( filter=filt, late_materialization=False ).explain_plan(True) - assert "(values)" in dataset.scanner( - filter=filt, late_materialization=True - ).explain_plan(True) - assert "(values)" in dataset.scanner( + assert late in dataset.scanner(filter=filt, late_materialization=True).explain_plan( + True + ) + assert late in dataset.scanner( filter=filt, late_materialization=["values"] ).explain_plan(True) assert ", values" in dataset.scanner( diff --git a/python/python/tests/test_mem_wal.py b/python/python/tests/test_mem_wal.py index 95596a7123e..bcfeac47e32 100644 --- a/python/python/tests/test_mem_wal.py +++ b/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", ( diff --git a/python/python/tests/test_scalar_index.py b/python/python/tests/test_scalar_index.py index c6c0952578e..9d42926ff11 100644 --- a/python/python/tests/test_scalar_index.py +++ b/python/python/tests/test_scalar_index.py @@ -624,13 +624,13 @@ def make_fts_search(ds): plan = make_vec_search(ds).explain_plan() assert "ScalarIndexQuery" in plan assert "KNNVectorDistance" not in plan - assert "LanceRead" not in plan + assert "num_fragments" not in plan # no scan; the take prints as LanceRead assert make_vec_search(ds).to_table().num_rows == 6 plan = make_fts_search(ds).explain_plan() assert "ScalarIndexQuery" in plan assert "KNNVectorDistance" not in plan - assert "LanceRead" not in plan + assert "num_fragments" not in plan # no scan; the take prints as LanceRead assert make_fts_search(ds).to_table().num_rows == 6 # Add new data (including 6 more results) diff --git a/rust/lance-namespace-impls/src/dir.rs b/rust/lance-namespace-impls/src/dir.rs index dc2d83cf278..8710139bb78 100644 --- a/rust/lance-namespace-impls/src/dir.rs +++ b/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", diff --git a/rust/lance/src/dataset/mem_wal/memtable/flush.rs b/rust/lance/src/dataset/mem_wal/memtable/flush.rs index 410823c31db..9c9d94975ce 100644 --- a/rust/lance/src/dataset/mem_wal/memtable/flush.rs +++ b/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(); diff --git a/rust/lance/src/dataset/scanner.rs b/rust/lance/src/dataset/scanner.rs index 7c7a082ad7e..aab12a67931 100644 --- a/rust/lance/src/dataset/scanner.rs +++ b/rust/lance/src/dataset/scanner.rs @@ -4827,12 +4827,52 @@ impl Scanner { } /// Take row indices produced by input plan from the dataset (with projection) + /// + /// Planned as a [`FilteredReadExec`] row-stream read; legacy (v1) storage + /// keeps using [`TakeExec`]. #[allow(deprecated)] fn take( &self, input: Arc, output_projection: Projection, ) -> Result> { + let fields_to_take = output_projection + .clone() + .subtract_arrow_schema(input.schema().as_ref(), OnMissing::Ignore)?; + if !fields_to_take.has_data_fields() + && !fields_to_take.with_row_id + && !fields_to_take.with_row_addr + { + // No new columns needed + return Ok(input); + } + + let input_schema = input.schema(); + let has_row_id = input_schema.column_with_name(ROW_ID).is_some(); + let has_row_addr = input_schema.column_with_name(ROW_ADDR).is_some(); + // The v1 reader cannot serve a FilteredReadExec + if !self.dataset.is_legacy_storage() && (has_row_id || has_row_addr) { + // Pass the full (un-subtracted) target so a rebuild against a + // different child re-derives what to fetch, and preserve carried + // identity columns (downstream nodes may key off them; the final + // ProjectionExec trims for free) + let mut projection = output_projection; + projection.with_row_id |= has_row_id; + projection.with_row_addr |= has_row_addr; + let mut read_options = FilteredReadOptions::new(projection); + if let Some(batch_size) = self.batch_size { + read_options = read_options.with_batch_size(batch_size as u32); + } + if let Some(fragments) = &self.fragments { + read_options = read_options.with_fragments(Arc::new(fragments.clone())); + } + return Ok(Arc::new(FilteredReadExec::try_new( + self.dataset.clone(), + read_options, + Some(input), + )?)); + } + let coalesced = Arc::new(CoalesceBatchesExec::new( input.clone(), self.get_batch_size(), @@ -10188,9 +10228,8 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri..., projection=[i], row_id=true, row_addr=false, ordered=true, range=None" } else { "ProjectionExec: expr=[s@2 as s] - Take: columns=\"i, _rowid, (s)\" - CoalesceBatchesExec: target_batch_size=8192 - LanceRead: ..., projection=[i], num_fragments=2, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10) AND i < Int32(20), refine_filter=i > Int32(10) AND i < Int32(20)" + LanceRead: uri=..., projection=[s], source=stream(_rowid) + LanceRead: ..., projection=[i], num_fragments=2, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10) AND i < Int32(20), refine_filter=i > Int32(10) AND i < Int32(20)" }; assert_plan_equals( &dataset.dataset, @@ -10214,10 +10253,9 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri..., projection=[i, s], row_id=true, row_addr=false, ordered=true, range=None" } else { "ProjectionExec: expr=[i@0 as i, s@1 as s, vec@3 as vec] - Take: columns=\"i, s, _rowid, (vec)\" - CoalesceBatchesExec: target_batch_size=8192 - LanceRead: uri=..., projection=[i, s], num_fragments=2, range_before=None, range_after=None, \ - row_id=true, row_addr=false, full_filter=s IS NOT NULL, refine_filter=s IS NOT NULL" + LanceRead: uri=..., projection=[vec], source=stream(_rowid) + LanceRead: uri=..., projection=[i, s], num_fragments=2, range_before=None, range_after=None, \ + row_id=true, row_addr=false, full_filter=s IS NOT NULL, refine_filter=s IS NOT NULL" }; assert_plan_equals( &dataset.dataset, @@ -10257,10 +10295,9 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri..., projection=[s], row_id=true, row_addr=false, ordered=true, range=None" } else { "ProjectionExec: expr=[i@2 as i, s@0 as s, vec@3 as vec] - Take: columns=\"s, _rowid, (i), (vec)\" - CoalesceBatchesExec: target_batch_size=8192 - LanceRead: uri=..., projection=[s], num_fragments=2, range_before=None, \ - range_after=None, row_id=true, row_addr=false, full_filter=s IS NOT NULL, refine_filter=s IS NOT NULL" + LanceRead: uri=..., projection=[i, vec], source=stream(_rowid) + LanceRead: uri=..., projection=[s], num_fragments=2, range_before=None, \ + range_after=None, row_id=true, row_addr=false, full_filter=s IS NOT NULL, refine_filter=s IS NOT NULL" }; assert_plan_equals( &dataset.dataset, @@ -10282,10 +10319,9 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri..., projection=[s, vec], row_id=true, row_addr=false, ordered=true, range=None" } else { "ProjectionExec: expr=[i@3 as i, s@0 as s, vec@1 as vec] - Take: columns=\"s, vec, _rowid, (i)\" - CoalesceBatchesExec: target_batch_size=8192 - LanceRead: uri=..., projection=[s, vec], num_fragments=2, range_before=None, range_after=None, \ - row_id=true, row_addr=false, full_filter=s IS NOT NULL, refine_filter=s IS NOT NULL" + LanceRead: uri=..., projection=[i], source=stream(_rowid) + LanceRead: uri=..., projection=[s, vec], num_fragments=2, range_before=None, range_after=None, \ + row_id=true, row_addr=false, full_filter=s IS NOT NULL, refine_filter=s IS NOT NULL" }; assert_plan_equals( &dataset.dataset, @@ -10328,13 +10364,12 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri=..., projection=[vec], row_id=true, row_addr=false, ordered=false, range=None" } else { "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@0 as vec, _distance@2 as _distance] - Take: columns=\"vec, _rowid, _distance, (i), (s)\" - CoalesceBatchesExec: target_batch_size=8192 - FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=5), expr=... - KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vec], num_fragments=2, range_before=None, range_after=None, \ - row_id=true, row_addr=false, full_filter=--, refine_filter=--" + LanceRead: uri=..., projection=[i, s], source=stream(_rowid) + FilterExec: _distance@2 IS NOT NULL + SortExec: TopK(fetch=5), expr=... + KNNVectorDistance: metric=l2 + LanceRead: uri=..., projection=[vec], num_fragments=2, range_before=None, range_after=None, \ + row_id=true, row_addr=false, full_filter=--, refine_filter=--" }; assert_plan_equals( &dataset.dataset, @@ -10358,14 +10393,13 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri=..., projection=[vec], row_id=true, row_addr=false, ordered=false, range=None" } else { "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@0 as vec, _distance@2 as _distance] - Take: columns=\"vec, _rowid, _distance, (i), (s)\" - CoalesceBatchesExec: target_batch_size=8192 - GlobalLimitExec: skip=0, fetch=1 - FilterExec: _distance@2 IS NOT NULL - SortExec: TopK(fetch=5), expr=... - KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vec], num_fragments=2, range_before=None, range_after=None, \ - row_id=true, row_addr=false, full_filter=--, refine_filter=--" + LanceRead: uri=..., projection=[i, s], source=stream(_rowid) + GlobalLimitExec: skip=0, fetch=1 + FilterExec: _distance@2 IS NOT NULL + SortExec: TopK(fetch=5), expr=... + KNNVectorDistance: metric=l2 + LanceRead: uri=..., projection=[vec], num_fragments=2, range_before=None, range_after=None, \ + row_id=true, row_addr=false, full_filter=--, refine_filter=--" }; assert_plan_equals( &dataset.dataset, @@ -10378,13 +10412,20 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") // --------------------------------------------------------------------- dataset.make_vector_index().await?; log::info!("Test case: Basic ANN"); - let expected = + let expected = if data_storage_version == LanceFileVersion::Legacy { "ProjectionExec: expr=[i@2 as i, s@3 as s, vec@4 as vec, _distance@0 as _distance] Take: columns=\"_distance, _rowid, (i), (s), (vec)\" CoalesceBatchesExec: target_batch_size=8192 SortExec: TopK(fetch=42), expr=... ANNSubIndex: name=..., k=42, deltas=1, metric=L2 - ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1"; + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + } else { + "ProjectionExec: expr=[i@2 as i, s@3 as s, vec@4 as vec, _distance@0 as _distance] + LanceRead: uri=..., projection=[i, s, vec], source=stream(_rowid) + SortExec: TopK(fetch=42), expr=... + ANNSubIndex: name=..., k=42, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + }; assert_plan_equals( &dataset.dataset, |scan| scan.nearest("vec", &q, 42), @@ -10393,7 +10434,7 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") .await?; log::info!("Test case: ANN with refine"); - let expected = + let expected = if data_storage_version == LanceFileVersion::Legacy { "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] Take: columns=\"_rowid, vec, _distance, (i), (s)\" CoalesceBatchesExec: target_batch_size=8192 @@ -10404,7 +10445,18 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") CoalesceBatchesExec: target_batch_size=8192 SortExec: TopK(fetch=40), expr=... ANNSubIndex: name=..., k=40, deltas=1, metric=L2 - ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1"; + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + } else { + "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] + LanceRead: uri=..., projection=[i, s], source=stream(_rowid) + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=10), expr=... + KNNVectorDistance: metric=l2 + LanceRead: uri=..., projection=[vec], source=stream(_rowid) + SortExec: TopK(fetch=40), expr=... + ANNSubIndex: name=..., k=40, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + }; assert_plan_equals( &dataset.dataset, |scan| Ok(scan.nearest("vec", &q, 10)?.refine(4)), @@ -10424,13 +10476,12 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri=..., projection=[vec], row_id=true, row_addr=false, ordered=false, range=None" } else { "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@0 as vec, _distance@2 as _distance] - Take: columns=\"vec, _rowid, _distance, (i), (s)\" - CoalesceBatchesExec: target_batch_size=8192 - FilterExec: _distance@... IS NOT NULL - SortExec: TopK(fetch=13), expr=... - KNNVectorDistance: metric=l2 - LanceRead: uri=..., projection=[vec], num_fragments=2, range_before=None, range_after=None, \ - row_id=true, row_addr=false, full_filter=--, refine_filter=--" + LanceRead: uri=..., projection=[i, s], source=stream(_rowid) + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=13), expr=... + KNNVectorDistance: metric=l2 + LanceRead: uri=..., projection=[vec], num_fragments=2, range_before=None, range_after=None, \ + row_id=true, row_addr=false, full_filter=--, refine_filter=--" }; assert_plan_equals( &dataset.dataset, @@ -10440,7 +10491,8 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") .await?; log::info!("Test case: ANN with postfilter"); - let expected = "ProjectionExec: expr=[s@3 as s, vec@4 as vec, _distance@0 as _distance, _rowid@1 as _rowid] + let expected = if data_storage_version == LanceFileVersion::Legacy { + "ProjectionExec: expr=[s@3 as s, vec@4 as vec, _distance@0 as _distance, _rowid@1 as _rowid] Take: columns=\"_distance, _rowid, i, (s), (vec)\" CoalesceBatchesExec: target_batch_size=8192 FilterExec: i@2 > 10 @@ -10448,7 +10500,16 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") CoalesceBatchesExec: target_batch_size=8192 SortExec: TopK(fetch=17), expr=... ANNSubIndex: name=..., k=17, deltas=1, metric=L2 - ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1"; + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + } else { + "ProjectionExec: expr=[s@3 as s, vec@4 as vec, _distance@0 as _distance, _rowid@1 as _rowid] + LanceRead: uri=..., projection=[s, vec], source=stream(_rowid) + FilterExec: i@2 > 10 + LanceRead: uri=..., projection=[i], source=stream(_rowid) + SortExec: TopK(fetch=17), expr=... + ANNSubIndex: name=..., k=17, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + }; assert_plan_equals( &dataset.dataset, |scan| { @@ -10474,13 +10535,12 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri=..., projection=[i], row_id=true, row_addr=false, ordered=false, range=None" } else { "ProjectionExec: expr=[i@2 as i, s@3 as s, vec@4 as vec, _distance@0 as _distance] - Take: columns=\"_distance, _rowid, (i), (s), (vec)\" - CoalesceBatchesExec: target_batch_size=8192 - SortExec: TopK(fetch=17), expr=... - ANNSubIndex: name=..., k=17, deltas=1, metric=L2 - ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 - LanceRead: uri=..., projection=[], num_fragments=2, range_before=None, range_after=None, \ - row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10) + LanceRead: uri=..., projection=[i, s, vec], source=stream(_rowid) + SortExec: TopK(fetch=17), expr=... + ANNSubIndex: name=..., k=17, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 + LanceRead: uri=..., projection=[], num_fragments=2, range_before=None, range_after=None, \ + row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10) " }; assert_plan_equals( @@ -10497,7 +10557,8 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") dataset.append_new_data().await?; log::info!("Test case: Combined KNN/ANN"); - let expected = "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] + let expected = if data_storage_version == LanceFileVersion::Legacy { + "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] Take: columns=\"_rowid, vec, _distance, (i), (s)\" CoalesceBatchesExec: target_batch_size=8192 FilterExec: _distance@... IS NOT NULL @@ -10514,7 +10575,25 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") CoalesceBatchesExec: target_batch_size=8192 SortExec: TopK(fetch=6), expr=... ANNSubIndex: name=..., k=6, deltas=1, metric=L2 - ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1"; + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + } else { + "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] + LanceRead: uri=..., projection=[i, s], source=stream(_rowid) + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=6), expr=... + KNNVectorDistance: metric=l2 + CoalescePartitionsExec + UnionExec + ProjectionExec: expr=[_distance@2 as _distance, _rowid@1 as _rowid, vec@0 as vec] + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=6), expr=... + KNNVectorDistance: metric=l2 + LanceScan: uri=..., projection=[vec], row_id=true, row_addr=false, ordered=false, range=None + LanceRead: uri=..., projection=[vec], source=stream(_rowid) + SortExec: TopK(fetch=6), expr=... + ANNSubIndex: name=..., k=6, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + }; assert_plan_equals( &dataset.dataset, |scan| scan.nearest("vec", &q, 6), @@ -10526,7 +10605,8 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") // new data and with filter log::info!("Test case: Combined KNN/ANN with postfilter"); - let expected = "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] + let expected = if data_storage_version == LanceFileVersion::Legacy { + "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] Take: columns=\"_rowid, vec, _distance, i, (s)\" CoalesceBatchesExec: target_batch_size=8192 FilterExec: i@3 > 10 @@ -10546,7 +10626,27 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") CoalesceBatchesExec: target_batch_size=8192 SortExec: TopK(fetch=15), expr=... ANNSubIndex: name=..., k=15, deltas=1, metric=L2 - ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1"; + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + } else { + "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] + LanceRead: uri=..., projection=[s], source=stream(_rowid) + FilterExec: i@3 > 10 + LanceRead: uri=..., projection=[i], source=stream(_rowid) + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=15), expr=... + KNNVectorDistance: metric=l2 + CoalescePartitionsExec + UnionExec + ProjectionExec: expr=[_distance@2 as _distance, _rowid@1 as _rowid, vec@0 as vec] + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=15), expr=... + KNNVectorDistance: metric=l2 + LanceScan: uri=..., projection=[vec], row_id=true, row_addr=false, ordered=false, range=None + LanceRead: uri=..., projection=[vec], source=stream(_rowid) + SortExec: TopK(fetch=15), expr=... + ANNSubIndex: name=..., k=15, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1" + }; assert_plan_equals( &dataset.dataset, |scan| scan.nearest("vec", &q, 15)?.filter("i > 10"), @@ -10580,26 +10680,24 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri=..., projection=[i], row_id=true, row_addr=false, ordered=false, range=None" } else { "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] - Take: columns=\"_rowid, vec, _distance, (i), (s)\" - CoalesceBatchesExec: target_batch_size=8192 - FilterExec: _distance@... IS NOT NULL - SortExec: TopK(fetch=5), expr=... - KNNVectorDistance: metric=l2 - CoalescePartitionsExec - UnionExec - ProjectionExec: expr=[_distance@3 as _distance, _rowid@2 as _rowid, vec@0 as vec] - FilterExec: _distance@... IS NOT NULL - SortExec: TopK(fetch=5), expr=... - KNNVectorDistance: metric=l2 - FilterExec: i@1 > 10 - LanceScan: uri=..., projection=[vec, i], row_id=true, row_addr=false, ordered=false, range=None - Take: columns=\"_distance, _rowid, (vec)\" - CoalesceBatchesExec: target_batch_size=8192 - SortExec: TopK(fetch=5), expr=... - ANNSubIndex: name=..., k=5, deltas=1, metric=L2 - ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 - LanceRead: uri=..., projection=[], num_fragments=2, range_before=None, range_after=None, \ - row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10)" + LanceRead: uri=..., projection=[i, s], source=stream(_rowid) + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=5), expr=... + KNNVectorDistance: metric=l2 + CoalescePartitionsExec + UnionExec + ProjectionExec: expr=[_distance@3 as _distance, _rowid@2 as _rowid, vec@0 as vec] + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=5), expr=... + KNNVectorDistance: metric=l2 + FilterExec: i@1 > 10 + LanceScan: uri=..., projection=[vec, i], row_id=true, row_addr=false, ordered=false, range=None + LanceRead: uri=..., projection=[vec], source=stream(_rowid) + SortExec: TopK(fetch=5), expr=... + ANNSubIndex: name=..., k=5, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 + LanceRead: uri=..., projection=[], num_fragments=2, range_before=None, range_after=None, \ + row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10)" }; assert_plan_equals( &dataset.dataset, @@ -10622,14 +10720,22 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") dataset.make_scalar_index().await?; log::info!("Test case: ANN with scalar index"); - let expected = + let expected = if data_storage_version == LanceFileVersion::Legacy { "ProjectionExec: expr=[i@2 as i, s@3 as s, vec@4 as vec, _distance@0 as _distance] Take: columns=\"_distance, _rowid, (i), (s), (vec)\" CoalesceBatchesExec: target_batch_size=8192 SortExec: TopK(fetch=5), expr=... ANNSubIndex: name=..., k=5, deltas=1, metric=L2 ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 - ScalarIndexQuery: query=[i > 10]@i_idx(BTree)"; + ScalarIndexQuery: query=[i > 10]@i_idx(BTree)" + } else { + "ProjectionExec: expr=[i@2 as i, s@3 as s, vec@4 as vec, _distance@0 as _distance] + LanceRead: uri=..., projection=[i, s, vec], source=stream(_rowid) + SortExec: TopK(fetch=5), expr=... + ANNSubIndex: name=..., k=5, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 + ScalarIndexQuery: query=[i > 10]@i_idx(BTree)" + }; assert_plan_equals( &dataset.dataset, |scan| { @@ -10654,13 +10760,12 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri=..., projection=[i], row_id=true, row_addr=false, ordered=false, range=None" } else { "ProjectionExec: expr=[i@2 as i, s@3 as s, vec@4 as vec, _distance@0 as _distance] - Take: columns=\"_distance, _rowid, (i), (s), (vec)\" - CoalesceBatchesExec: target_batch_size=8192 - SortExec: TopK(fetch=5), expr=... - ANNSubIndex: name=..., k=5, deltas=1, metric=L2 - ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 - LanceRead: uri=..., projection=[], num_fragments=3, range_before=None, \ - range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10)" + LanceRead: uri=..., projection=[i, s, vec], source=stream(_rowid) + SortExec: TopK(fetch=5), expr=... + ANNSubIndex: name=..., k=5, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 + LanceRead: uri=..., projection=[], num_fragments=3, range_before=None, \ + range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10)" }; assert_plan_equals( &dataset.dataset, @@ -10678,7 +10783,8 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") dataset.append_new_data().await?; log::info!("Test case: Combined KNN/ANN with scalar index"); - let expected = "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] + let expected = if data_storage_version == LanceFileVersion::Legacy { + "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] Take: columns=\"_rowid, vec, _distance, (i), (s)\" CoalesceBatchesExec: target_batch_size=8192 FilterExec: _distance@... IS NOT NULL @@ -10697,7 +10803,27 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") SortExec: TopK(fetch=8), expr=... ANNSubIndex: name=..., k=8, deltas=1, metric=L2 ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 - ScalarIndexQuery: query=[i > 10]@i_idx(BTree)"; + ScalarIndexQuery: query=[i > 10]@i_idx(BTree)" + } else { + "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] + LanceRead: uri=..., projection=[i, s], source=stream(_rowid) + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=8), expr=... + KNNVectorDistance: metric=l2 + CoalescePartitionsExec + UnionExec + ProjectionExec: expr=[_distance@3 as _distance, _rowid@2 as _rowid, vec@0 as vec] + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=8), expr=... + KNNVectorDistance: metric=l2 + FilterExec: i@1 > 10 + LanceScan: uri=..., projection=[vec, i], row_id=true, row_addr=false, ordered=false, range=None + LanceRead: uri=..., projection=[vec], source=stream(_rowid) + SortExec: TopK(fetch=8), expr=... + ANNSubIndex: name=..., k=8, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 + ScalarIndexQuery: query=[i > 10]@i_idx(BTree)" + }; assert_plan_equals( &dataset.dataset, |scan| { @@ -10714,7 +10840,8 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") log::info!( "Test case: Combined KNN/ANN with updated scalar index and outdated vector index" ); - let expected = "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] + let expected = if data_storage_version == LanceFileVersion::Legacy { + "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] Take: columns=\"_rowid, vec, _distance, (i), (s)\" CoalesceBatchesExec: target_batch_size=8192 FilterExec: _distance@... IS NOT NULL @@ -10733,7 +10860,27 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") SortExec: TopK(fetch=11), expr=... ANNSubIndex: name=..., k=11, deltas=1, metric=L2 ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 - ScalarIndexQuery: query=[i > 10]@i_idx(BTree)"; + ScalarIndexQuery: query=[i > 10]@i_idx(BTree)" + } else { + "ProjectionExec: expr=[i@3 as i, s@4 as s, vec@1 as vec, _distance@2 as _distance] + LanceRead: uri=..., projection=[i, s], source=stream(_rowid) + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=11), expr=... + KNNVectorDistance: metric=l2 + CoalescePartitionsExec + UnionExec + ProjectionExec: expr=[_distance@3 as _distance, _rowid@2 as _rowid, vec@0 as vec] + FilterExec: _distance@... IS NOT NULL + SortExec: TopK(fetch=11), expr=... + KNNVectorDistance: metric=l2 + FilterExec: i@1 > 10 + LanceScan: uri=..., projection=[vec, i], row_id=true, row_addr=false, ordered=false, range=None + LanceRead: uri=..., projection=[vec], source=stream(_rowid) + SortExec: TopK(fetch=11), expr=... + ANNSubIndex: name=..., k=11, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1 + ScalarIndexQuery: query=[i > 10]@i_idx(BTree)" + }; dataset.make_scalar_index().await?; assert_plan_equals( &dataset.dataset, @@ -10779,10 +10926,9 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") .filter("i > 10") }, "ProjectionExec: expr=[s@2 as s] - Take: columns=\"i, _rowid, (s)\" - CoalesceBatchesExec: target_batch_size=8192 - LanceRead: uri=..., projection=[i], num_fragments=4, range_before=None, \ - range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10)", + LanceRead: uri=..., projection=[s], source=stream(_rowid) + LanceRead: uri=..., projection=[i], num_fragments=4, range_before=None, \ + range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=i > Int32(10)", ) .await?; } @@ -10891,10 +11037,16 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") // All rows are indexed dataset.make_fts_index().await?; log::info!("Test case: Full text search (match query)"); - let expected = r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] + let expected = if data_storage_version == LanceFileVersion::Legacy { + r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] Take: columns="_rowid, _score, (s)" CoalesceBatchesExec: target_batch_size=8192 - MatchQuery: column=s, query=hello"#; + MatchQuery: column=s, query=hello"# + } else { + r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] + LanceRead: uri=..., projection=[s], source=stream(_rowid) + MatchQuery: column=s, query=hello"# + }; assert_plan_equals( &dataset.dataset, |scan| { @@ -10907,10 +11059,16 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") .await?; log::info!("Test case: Full text search (phrase query)"); - let expected = r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] + let expected = if data_storage_version == LanceFileVersion::Legacy { + r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] Take: columns="_rowid, _score, (s)" CoalesceBatchesExec: target_batch_size=8192 - PhraseQuery: column=s, query=hello world"#; + PhraseQuery: column=s, query=hello world"# + } else { + r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] + LanceRead: uri=..., projection=[s], source=stream(_rowid) + PhraseQuery: column=s, query=hello world"# + }; assert_plan_equals( &dataset.dataset, |scan| { @@ -10924,12 +11082,20 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") .await?; log::info!("Test case: Full text search (boost query)"); - let expected = r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] + let expected = if data_storage_version == LanceFileVersion::Legacy { + r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] Take: columns="_rowid, _score, (s)" CoalesceBatchesExec: target_batch_size=8192 BoostQuery: negative_boost=1 MatchQuery: column=s, query=hello - MatchQuery: column=s, query=world"#; + MatchQuery: column=s, query=world"# + } else { + r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] + LanceRead: uri=..., projection=[s], source=stream(_rowid) + BoostQuery: negative_boost=1 + MatchQuery: column=s, query=hello + MatchQuery: column=s, query=world"# + }; assert_plan_equals( &dataset.dataset, |scan| { @@ -10960,11 +11126,10 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri=..., projection=[i], row_id=true, row_addr=false, ordered=false, range=None"# } else { r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] - Take: columns="_rowid, _score, (s)" - CoalesceBatchesExec: target_batch_size=8192 - MatchQuery: column=s, query=hello - LanceRead: uri=..., projection=[], num_fragments=5, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=-- - ScalarIndexQuery: query=[i > 10]@i_idx(BTree)"# + LanceRead: uri=..., projection=[s], source=stream(_rowid) + MatchQuery: column=s, query=hello + LanceRead: uri=..., projection=[], num_fragments=5, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=-- + ScalarIndexQuery: query=[i > 10]@i_idx(BTree)"# }; assert_plan_equals( &dataset.dataset, @@ -10996,14 +11161,13 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri=..., projection=[s], row_id=true, row_addr=false, ordered=true, range=None"# } else { r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] - Take: columns="_rowid, _score, (s)" - CoalesceBatchesExec: target_batch_size=8192 - SortExec: expr=[_score@1 DESC NULLS LAST], preserve_partitioning=[false] - CoalescePartitionsExec - UnionExec - MatchQuery: column=s, query=hello - FlatMatchQuery: column=s, query=hello - LanceRead: uri=..., projection=[s], num_fragments=1, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=--, refine_filter=--"# + LanceRead: uri=..., projection=[s], source=stream(_rowid) + SortExec: expr=[_score@1 DESC NULLS LAST], preserve_partitioning=[false] + CoalescePartitionsExec + UnionExec + MatchQuery: column=s, query=hello + FlatMatchQuery: column=s, query=hello + LanceRead: uri=..., projection=[s], num_fragments=1, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=--, refine_filter=--"# }; dataset.append_new_data().await?; assert_plan_equals( @@ -11018,10 +11182,16 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") .await?; log::info!("Test case: Full text search with unindexed rows and fast_search"); - let expected = r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] + let expected = if data_storage_version == LanceFileVersion::Legacy { + r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] Take: columns="_rowid, _score, (s)" CoalesceBatchesExec: target_batch_size=8192 - MatchQuery: column=s, query=hello"#; + MatchQuery: column=s, query=hello"# + } else { + r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] + LanceRead: uri=..., projection=[s], source=stream(_rowid) + MatchQuery: column=s, query=hello"# + }; assert_plan_equals( &dataset.dataset, |scan| { @@ -11066,17 +11236,16 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") LanceScan: uri=..., projection=[i, s], row_id=true, row_addr=false, ordered=false, range=None"# } else { r#"ProjectionExec: expr=[s@2 as s, _score@1 as _score, _rowid@0 as _rowid] - Take: columns="_rowid, _score, (s)" - CoalesceBatchesExec: target_batch_size=8192 - SortExec: expr=[_score@1 DESC NULLS LAST], preserve_partitioning=[false] - CoalescePartitionsExec - UnionExec - MatchQuery: column=s, query=hello - LanceRead: uri=..., projection=[], num_fragments=5, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=-- - ScalarIndexQuery: query=[i > 10]@i_idx(BTree) - FlatMatchQuery: column=s, query=hello - LanceRead: uri=..., projection=[s], num_fragments=1, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=-- - ScalarIndexQuery: query=[i > 10]@i_idx(BTree)"# + LanceRead: uri=..., projection=[s], source=stream(_rowid) + SortExec: expr=[_score@1 DESC NULLS LAST], preserve_partitioning=[false] + CoalescePartitionsExec + UnionExec + MatchQuery: column=s, query=hello + LanceRead: uri=..., projection=[], num_fragments=5, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=-- + ScalarIndexQuery: query=[i > 10]@i_idx(BTree) + FlatMatchQuery: column=s, query=hello + LanceRead: uri=..., projection=[s], num_fragments=1, range_before=None, range_after=None, row_id=true, row_addr=false, full_filter=i > Int32(10), refine_filter=-- + ScalarIndexQuery: query=[i > 10]@i_idx(BTree)"# }; assert_plan_equals( &dataset.dataset, @@ -11153,11 +11322,10 @@ full_filter=name LIKE Utf8(\"test%2\"), refine_filter=name LIKE Utf8(\"test%2\") SortExec: TopK(fetch=34), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST]... KNNVectorDistance: metric=l2 LanceScan: uri=..., projection=[vec], row_id=true, row_addr=false, ordered=false, range=None - Take: columns=\"_distance, _rowid, (vec)\" - CoalesceBatchesExec: target_batch_size=8192 - SortExec: TopK(fetch=34), expr=[_distance@0 ASC NULLS LAST, _rowid@1 ASC NULLS LAST]... - ANNSubIndex: name=idx, k=34, deltas=1, metric=L2 - ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1", + LanceRead: uri=..., projection=[vec], source=stream(_rowid) + SortExec: TopK(fetch=34), expr=[_distance@0 ASC NULLS LAST, _rowid@1 ASC NULLS LAST]... + ANNSubIndex: name=idx, k=34, deltas=1, metric=L2 + ANNIvfPartition: uuid=..., minimum_nprobes=1, maximum_nprobes=None, deltas=1", ) .await .unwrap(); diff --git a/rust/lance/src/dataset/write/merge_insert.rs b/rust/lance/src/dataset/write/merge_insert.rs index 19b41284e0b..b937de6a174 100644 --- a/rust/lance/src/dataset/write/merge_insert.rs +++ b/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, }, @@ -724,24 +726,35 @@ impl MergeInsertJob { index_mapper_input, )); - // If requested, add row addresses to the output - if add_row_addr { - let pos = index_mapper.schema().fields().len(); // Add to end - index_mapper = Arc::new(AddRowAddrExec::try_new( - index_mapper, - self.dataset.clone(), - pos, - )?); - } - - // 4 - Take the mapped row ids + // 4 - Take the mapped row ids (TakeExec stays for legacy storage: + // 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 = if self.dataset.is_legacy_storage() { + if add_row_addr { + let pos = index_mapper.schema().fields().len(); // Add to end + index_mapper = Arc::new(AddRowAddrExec::try_new( + index_mapper, + self.dataset.clone(), + pos, + )?); + } Arc::new(TakeExec::try_new(self.dataset.clone(), index_mapper, projection)?.unwrap()) - as Arc; + } else { + // Keep the mapped row ids; the read synthesizes the row addresses + // if requested (no AddRowAddrExec needed) + let mut projection = projection.with_row_id(); + if add_row_addr { + projection = projection.with_row_addr(); + } + 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 diff --git a/rust/lance/src/io/exec/count_pushdown.rs b/rust/lance/src/io/exec/count_pushdown.rs index d5d90b5881a..c3f2d9fc096 100644 --- a/rust/lance/src/io/exec/count_pushdown.rs +++ b/rust/lance/src/io/exec/count_pushdown.rs @@ -147,6 +147,11 @@ fn try_rewrite(agg: &AggregateExec) -> DFResult>> }; 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() { diff --git a/rust/lance/src/io/exec/filtered_read.rs b/rust/lance/src/io/exec/filtered_read.rs index 059fa4e2b36..48356392c9b 100644 --- a/rust/lance/src/io/exec/filtered_read.rs +++ b/rust/lance/src/io/exec/filtered_read.rs @@ -5,9 +5,16 @@ use std::collections::{BTreeMap, HashMap}; use std::pin::Pin; use std::sync::Mutex; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::{ops::Range, sync::Arc}; +use std::task::Poll; +use std::{ + ops::{Range, RangeInclusive}, + sync::Arc, +}; -use arrow_array::RecordBatch; +use arrow::compute::BatchCoalescer; +use arrow_array::cast::AsArray; +use arrow_array::types::UInt64Type; +use arrow_array::{Array, BooleanArray, RecordBatch, UInt32Array}; use arrow_schema::SchemaRef; use datafusion::common::runtime::SpawnedTask; use datafusion::common::stats::Precision; @@ -31,7 +38,9 @@ use lance_core::datatypes::OnMissing; use lance_core::utils::deletion::DeletionVector; use lance_core::utils::futures::FinallyStreamExt; use lance_core::utils::tokio::get_num_compute_intensive_cpus; -use lance_core::{Error, Result, datatypes::Projection}; +use lance_core::{ + Error, ROW_ADDR, ROW_ADDR_FIELD, ROW_ID, ROW_ID_FIELD, Result, datatypes::Projection, +}; use lance_datafusion::planner::Planner; use lance_datafusion::utils::{ ExecutionPlanMetricsSetExt, FRAGMENTS_SCANNED_METRIC, RANGES_SCANNED_METRIC, @@ -41,7 +50,8 @@ use lance_file::reader::FileReaderOptions; use lance_index::scalar::expression::FilterPlan; use lance_io::scheduler::{ScanScheduler, SchedulerConfig}; use lance_select::{ - IndexExprResult, RowAddrSelection, RowAddrTreeMap, bitmap_to_ranges, ranges_to_bitmap, + IndexExprResult, RowAddrMask, RowAddrSelection, RowAddrTreeMap, bitmap_to_ranges, + ranges_to_bitmap, result::IndexExprResultWireFormat, }; use lance_table::format::Fragment; use lance_table::rowids::RowIdSequence; @@ -346,7 +356,31 @@ impl FilteredReadStream { plan: FilteredReadInternalPlan, ) -> DataFusionResult { let global_metrics = Arc::new(FilteredReadGlobalMetrics::new(metrics)); + let loaded_fragments = Self::load_all_fragments(&dataset, &options).await?; + let scan_scheduler = Self::make_scan_scheduler(&dataset, &options); + Ok(Self::new_shared( + dataset, + options, + global_metrics, + plan, + scan_scheduler, + &loaded_fragments, + 0, + )) + } + /// Shared with the row-stream path, which injects its per-query + /// scheduler, cached fragments, and a per-batch priority offset + #[allow(clippy::too_many_arguments)] + fn new_shared( + dataset: Arc, + options: FilteredReadOptions, + global_metrics: Arc, + plan: FilteredReadInternalPlan, + scan_scheduler: Arc, + loaded_fragments: &[LoadedFragment], + priority_offset: u32, + ) -> Self { let threading_mode = options.threading_mode; let io_parallelism = dataset.object_store.io_parallelism(); @@ -355,62 +389,31 @@ impl FilteredReadStream { .unwrap_or_else(|| (*DEFAULT_FRAGMENT_READAHEAD).unwrap_or(io_parallelism * 2)) .max(1); - let fragments = options - .fragments - .clone() - .unwrap_or_else(|| dataset.fragments().clone()); - log::debug!( "Filtered read on {} fragments with frag_readahead={} and io_parallelism={}", - fragments.len(), + loaded_fragments.len(), fragment_readahead, io_parallelism ); - // Ideally we don't need to collect here but if we don't we get "implementation of FnOnce is - // not general enough" false positives from rustc - let frag_futs = fragments - .iter() - .map(|frag| { - Result::Ok(Self::load_fragment( - dataset.clone(), - frag.clone(), - options.with_deleted_rows, - )) - }) - .collect::>(); - let loaded_fragments = futures::stream::iter(frag_futs) - // Cannot use unordered because we need to populate logical_offset based on user-provided order - .try_buffered(io_parallelism) - .try_collect::>() - .await?; - let output_schema = Arc::new(options.projection.to_arrow_schema()); - let obj_store = dataset.object_store.clone(); - // Explicit options take precedence; otherwise fall back to the - // LANCE_DEFAULT_IO_BUFFER_SIZE env var if set; otherwise max_bandwidth. - let scheduler_config = if let Some(io_buffer_size_bytes) = options - .io_buffer_size_bytes - .or_else(get_default_io_buffer_size_override) - { - SchedulerConfig::new(io_buffer_size_bytes) - } else { - SchedulerConfig::max_bandwidth(obj_store.as_ref()) - }; - let scan_scheduler = ScanScheduler::new(obj_store, scheduler_config); - // Get scan_range_after_filter from the plan let scan_range_after_filter = plan.scan_range_after_filter.clone(); // Convert plan to scoped fragments for I/O - let scoped_fragments = Self::plan_to_scoped_fragments( + let mut scoped_fragments = Self::plan_to_scoped_fragments( &plan, - &loaded_fragments, + loaded_fragments, &dataset, &options, scan_scheduler.clone(), ); + if priority_offset != 0 { + for scoped in &mut scoped_fragments { + scoped.priority = scoped.priority.saturating_add(priority_offset); + } + } let global_metrics_clone = global_metrics.clone(); @@ -429,7 +432,7 @@ impl FilteredReadStream { .buffered(fragment_readahead); let task_stream = fragment_streams.try_flatten().boxed(); - Ok(Self { + Self { output_schema, task_stream: Arc::new(AsyncMutex::new(task_stream)), scan_scheduler, @@ -437,7 +440,60 @@ impl FilteredReadStream { active_partitions_counter: Arc::new(AtomicUsize::new(0)), threading_mode, scan_range_after_filter, - }) + } + } + + /// Drain the entire read into batches (used by the row-stream path, + /// which is the stream's only consumer and records metrics per batch) + async fn collect_all(&self, decode_parallelism: usize) -> Result> { + let mut task_stream = self.task_stream.lock().await; + (&mut *task_stream) + .try_buffered(decode_parallelism) + .try_collect() + .await + } + + async fn load_all_fragments( + dataset: &Arc, + options: &FilteredReadOptions, + ) -> Result> { + let io_parallelism = dataset.object_store.io_parallelism(); + let fragments = options + .fragments + .clone() + .unwrap_or_else(|| dataset.fragments().clone()); + // Ideally we don't need to collect here but if we don't we get "implementation of FnOnce is + // not general enough" false positives from rustc + let frag_futs = fragments + .iter() + .map(|frag| { + Result::Ok(Self::load_fragment( + dataset.clone(), + frag.clone(), + options.with_deleted_rows, + )) + }) + .collect::>(); + futures::stream::iter(frag_futs) + // Cannot use unordered because we need to populate logical_offset based on user-provided order + .try_buffered(io_parallelism) + .try_collect::>() + .await + } + + /// Create the I/O scheduler for a read (explicit option → env override → + /// max bandwidth) + fn make_scan_scheduler(dataset: &Dataset, options: &FilteredReadOptions) -> Arc { + let obj_store = dataset.object_store.clone(); + let scheduler_config = if let Some(io_buffer_size_bytes) = options + .io_buffer_size_bytes + .or_else(get_default_io_buffer_size_override) + { + SchedulerConfig::new(io_buffer_size_bytes) + } else { + SchedulerConfig::max_bandwidth(obj_store.as_ref()) + }; + ScanScheduler::new(obj_store, scheduler_config) } async fn load_fragment( @@ -1270,7 +1326,7 @@ pub struct FilteredReadOptions { pub scan_range_before_filter: Option>, /// The range of rows to read after applying the filter. pub scan_range_after_filter: Option>, - /// Include deleted rows in the scan + /// Include deleted rows in the scan; they are returned with a null row id pub with_deleted_rows: bool, /// The maximum number of rows per batch pub batch_size: Option, @@ -1504,7 +1560,7 @@ pub struct FilteredReadExec { options: FilteredReadOptions, properties: Arc, metrics: ExecutionPlanMetricsSet, - index_input: Option>, + input: RowSelector, // Precomputed internal plan plan: Arc>, // When execute is first called we will initialize the FilteredReadStream. In order to support @@ -1512,6 +1568,49 @@ pub struct FilteredReadExec { running_stream: Arc>>, } +/// Describes which rows a [`FilteredReadExec`] should read +#[derive(Debug)] +enum RowSelector { + /// Every live row of the dataset (no input plan) + AllRows, + /// A set of rows: one serialized [`IndexExprResult`] batch. Output is in + /// storage order and deduplicated. + RowSet(Arc), + /// A stream of rows: record batches with a `_rowid`/`_rowaddr` column + /// and other payload columns (just carried) + RowStream(Arc), +} + +impl RowSelector { + fn row_set_plan(&self) -> Option<&Arc> { + match self { + Self::RowSet(plan) => Some(plan), + _ => None, + } + } + + fn child(&self) -> Option<&Arc> { + match self { + Self::AllRows => None, + Self::RowSet(plan) => Some(plan), + Self::RowStream(source) => Some(&source.plan), + } + } +} + +/// State derived at construction for a row-stream source +#[derive(Debug)] +struct RowStreamSource { + plan: Arc, + /// The stream column identifying rows: [`ROW_ID`] or [`ROW_ADDR`] + key_column: &'static str, + /// Options for the internal fragment read; carries the projection that + /// reflects the actual columns to read (plus the alignment key column) + read_options: FilteredReadOptions, + /// The schema for newly read columns + new_fields_schema: SchemaRef, +} + /// Public plan for distributed execution - uses bitmap for flexibility #[derive(Clone)] pub struct FilteredReadPlan { @@ -1558,11 +1657,170 @@ impl FilteredReadInternalPlan { } impl FilteredReadExec { + /// Create a new filtered read pub fn try_new( + dataset: Arc, + options: FilteredReadOptions, + input: Option>, + ) -> Result { + match input { + Some(input) if Self::is_index_query_schema(input.schema().as_ref()) => { + Self::try_new_scan(dataset, options, Some(input)) + } + Some(input) => Self::try_new_row_stream(dataset, options, input), + None => Self::try_new_scan(dataset, options, None), + } + } + + /// Whether `schema` is one of the serialized [`IndexExprResult`] wire + /// layouts (see [`IndexExprResultWireFormat`]) + fn is_index_query_schema(schema: &arrow_schema::Schema) -> bool { + [ + IndexExprResultWireFormat::TwoMask, + IndexExprResultWireFormat::ThreeVariant, + ] + .iter() + .any(|format| schema.fields() == format.schema().fields()) + } + + /// The input columns that carry through to the output: identity columns + /// appear iff their flag is requested, ordinary columns always carry + fn carried_schema(input_schema: &arrow_schema::Schema, projection: &Projection) -> SchemaRef { + Arc::new(arrow_schema::Schema::new( + input_schema + .fields() + .iter() + .filter(|f| { + (f.name() != ROW_ID || projection.with_row_id) + && (f.name() != ROW_ADDR || projection.with_row_addr) + }) + .cloned() + .collect::>(), + )) + } + + /// Construct a read over a row-stream source + fn try_new_row_stream( + dataset: Arc, + options: FilteredReadOptions, + input: Arc, + ) -> Result { + if dataset.is_legacy_storage() { + return Err(Error::not_supported_source( + "taking rows through FilteredReadExec requires the v2 storage format" + .to_string() + .into(), + )); + } + if options.refine_filter.is_some() || options.full_filter.is_some() { + return Err(Error::invalid_input_source( + "filters are not supported when taking rows from an input plan".into(), + )); + } + // A limit is safer to apply upstream, on the cheap keyed rows + if options.scan_range_before_filter.is_some() || options.scan_range_after_filter.is_some() { + return Err(Error::invalid_input_source( + "scan ranges are not supported when taking rows from an input plan".into(), + )); + } + // Row-stream reads do not support deleted rows yet; deleted rows are + // excluded from the output by default + if options.with_deleted_rows || options.only_indexed_fragments { + return Err(Error::invalid_input_source( + "with_deleted_rows / only_indexed_fragments are not supported when taking rows from an input plan".into(), + )); + } + let input_schema = input.schema(); + let key_column = if input_schema.column_with_name(ROW_ID).is_some() { + ROW_ID + } else if input_schema.column_with_name(ROW_ADDR).is_some() { + ROW_ADDR + } else { + return Err(Error::invalid_input_source( + format!( + "a row-stream input plan must have a column named '{}' or '{}'", + ROW_ADDR, ROW_ID + ) + .into(), + )); + }; + + let fields_to_read = options + .projection + .clone() + .subtract_arrow_schema(input_schema.as_ref(), OnMissing::Ignore)?; + let synthesize_row_id = fields_to_read.with_row_id; + let synthesize_row_addr = fields_to_read.with_row_addr; + if !fields_to_read.has_data_fields() && !synthesize_row_id && !synthesize_row_addr { + return Err(Error::invalid_input_source( + "the input plan already contains every projected field; there is nothing to read" + .into(), + )); + } + + let carried_schema = Self::carried_schema(input_schema.as_ref(), &options.projection); + + // Output = carried columns ⊕ fetched fields ⊕ synthesized identity + let output_schema = Arc::new(arrow_schema::Schema::from( + &super::TakeExec::calculate_output_schema( + dataset.schema(), + carried_schema.as_ref(), + &fields_to_read, + ), + )); + + // Partitioning and emission behavior follow the input + let properties = Arc::new( + input + .properties() + .as_ref() + .clone() + .with_eq_properties(EquivalenceProperties::new(output_schema)), + ); + + let bare_schema = arrow_schema::Schema::from(&fields_to_read.to_bare_schema()); + let mut new_fields = bare_schema.fields().iter().cloned().collect::>(); + if synthesize_row_id { + new_fields.push(Arc::new(ROW_ID_FIELD.clone())); + } + if synthesize_row_addr { + new_fields.push(Arc::new(ROW_ADDR_FIELD.clone())); + } + let new_fields_schema = Arc::new(arrow_schema::Schema::new(new_fields)); + + // fields_to_read keeps the synthesis flags; add the key column on top + let mut read_options = options.clone(); + read_options.projection = if key_column == ROW_ID { + fields_to_read.with_row_id() + } else { + fields_to_read.with_row_addr() + }; + + Ok(Self { + dataset, + options, + properties, + metrics: ExecutionPlanMetricsSet::new(), + input: RowSelector::RowStream(Arc::new(RowStreamSource { + plan: input, + key_column, + read_options, + new_fields_schema, + })), + plan: Arc::new(OnceCell::new()), + running_stream: Arc::new(AsyncMutex::new(None)), + }) + } + + fn try_new_scan( dataset: Arc, mut options: FilteredReadOptions, index_input: Option>, ) -> Result { + let input = match index_input { + Some(plan) => RowSelector::RowSet(plan), + None => RowSelector::AllRows, + }; if options.with_deleted_rows { // Ensure we have the row id column if with_deleted_rows is set options.projection = options.projection.with_row_id(); @@ -1577,7 +1835,7 @@ impl FilteredReadExec { // Validate that there's a filter when using scan_range_after_filter if options.full_filter.is_none() && options.refine_filter.is_none() - && index_input.is_none() + && input.row_set_plan().is_none() { return Err(Error::invalid_input_source("scan_range_after_filter requires a filter to be applied. Use scan_range_before_filter for unfiltered scans." .into())); @@ -1616,7 +1874,7 @@ impl FilteredReadExec { properties, running_stream: Arc::new(AsyncMutex::new(None)), metrics, - index_input, + input, plan: Arc::new(OnceCell::new()), }) } @@ -1715,11 +1973,18 @@ impl FilteredReadExec { /// Get the existing plan or create it if it doesn't exist pub async fn get_or_create_plan(&self, ctx: Arc) -> Result { + if self.row_stream_input().is_some() { + return Err(Error::not_supported_source( + "a FilteredReadExec with a row-stream source does not have a precomputable plan" + .to_string() + .into(), + )); + } let internal_plan = Self::get_or_create_plan_impl( &self.plan, self.dataset.clone(), &self.options, - self.index_input.as_ref(), + self.input.row_set_plan(), 0, ctx, ) @@ -1751,7 +2016,7 @@ impl FilteredReadExec { .as_ref() .and_then(|o| o.batch_size_bytes); let metrics = self.metrics.clone(); - let index_input = self.index_input.clone(); + let index_input = self.input.row_set_plan().cloned(); let plan_cell = self.plan.clone(); let stream = futures::stream::once(async move { @@ -1796,29 +2061,496 @@ impl FilteredReadExec { }) .try_flatten(); - Box::pin(RecordBatchStreamAdapter::new(self.schema(), stream)) - } + Box::pin(RecordBatchStreamAdapter::new(self.schema(), stream)) + } + + pub fn dataset(&self) -> &Arc { + &self.dataset + } + + pub fn options(&self) -> &FilteredReadOptions { + &self.options + } + + pub fn index_input(&self) -> Option<&Arc> { + self.input.row_set_plan() + } + + pub fn row_stream_input(&self) -> Option<&Arc> { + match &self.input { + RowSelector::RowStream(source) => Some(&source.plan), + _ => None, + } + } + + /// Return the pre-computed plan if one exists, without triggering initialization. + pub fn plan(&self) -> Option { + self.plan.get().map(|p| p.to_external_plan()) + } + + fn execute_row_stream( + &self, + source: &Arc, + partition: usize, + context: Arc, + ) -> DataFusionResult { + let input_stream = source.plan.execute(partition, context)?; + let dataset = self.dataset.clone(); + let source = source.clone(); + let carried_schema = + Self::carried_schema(source.plan.schema().as_ref(), &self.options.projection); + let output_schema = self.schema(); + let metrics = self.metrics.clone(); + + let lazy_stream = futures::stream::once(async move { + let row_stream_read = Arc::new(RowStreamRead::new( + dataset, + source, + carried_schema, + output_schema, + &metrics, + partition, + )); + row_stream_read.apply(input_stream) + }) + .flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + lazy_stream, + ))) + } +} + +/// How many batches run ahead of the consumer (a prefetch depth, not a +/// parallelism setting) +const ROW_STREAM_PREFETCH_BATCHES: usize = 64; + +/// Fragment metadata, loaded on the first batch and reused afterwards +struct StreamFragments { + /// All dataset (or scoped) fragments, in dataset order + fragments: Vec, + /// Fragment id → position in `fragments` + positions: HashMap, + /// Each fragment's row-id span, for skipping fragments a batch cannot + /// touch (None = empty fragment) + id_spans: Vec>>, +} + +impl StreamFragments { + fn get(&self, fragment_id: u32) -> Option<&LoadedFragment> { + self.positions + .get(&fragment_id) + .map(|position| &self.fragments[*position]) + } +} + +/// Executes a [`FilteredReadExec`] over a row-stream source +struct RowStreamRead { + dataset: Arc, + source: Arc, + /// The input columns that carry through to the output + carried_schema: SchemaRef, + output_schema: SchemaRef, + scan_scheduler: Arc, + loaded_fragments: OnceCell, + global_metrics: Arc, + baseline_metrics: BaselineMetrics, +} + +impl RowStreamRead { + fn new( + dataset: Arc, + source: Arc, + carried_schema: SchemaRef, + output_schema: SchemaRef, + metrics: &ExecutionPlanMetricsSet, + partition: usize, + ) -> Self { + let scan_scheduler = + FilteredReadStream::make_scan_scheduler(&dataset, &source.read_options); + Self { + dataset, + source, + carried_schema, + output_schema, + scan_scheduler, + loaded_fragments: OnceCell::new(), + global_metrics: Arc::new(FilteredReadGlobalMetrics::new(metrics)), + baseline_metrics: BaselineMetrics::new(metrics, partition), + } + } + + async fn load_fragments(&self) -> Result<&StreamFragments> { + self.loaded_fragments + .get_or_try_init(|| async { + let fragments = FilteredReadStream::load_all_fragments( + &self.dataset, + &self.source.read_options, + ) + .await?; + let positions = fragments + .iter() + .enumerate() + .map(|(position, fragment)| (fragment.fragment.id() as u32, position)) + .collect(); + let id_spans = fragments + .iter() + .map(|fragment| fragment.row_id_sequence.row_id_range()) + .collect(); + Ok(StreamFragments { + fragments, + positions, + id_spans, + }) + }) + .await + } + + /// Build a batch's read ranges directly from physical row addresses + fn plan_batch_from_addresses( + addrs: &RowAddrTreeMap, + fragments: &StreamFragments, + ) -> FilteredReadInternalPlan { + let mut rows: BTreeMap>> = BTreeMap::new(); + for (fragment_id, requested) in addrs.iter() { + // Unknown fragments (e.g. fully deleted) drop like stale keys + let Some(fragment) = fragments.get(*fragment_id) else { + continue; + }; + let requested = match requested { + RowAddrSelection::Full => vec![0..fragment.num_physical_rows], + RowAddrSelection::Partial(bitmap) => bitmap_to_ranges(bitmap), + }; + let valid = FilteredReadStream::full_frag_range( + fragment.num_physical_rows, + &fragment.deletion_vector, + ); + let matched = FilteredReadStream::intersect_ranges(&valid, &requested); + if !matched.is_empty() { + rows.insert(*fragment_id, matched); + } + } + FilteredReadInternalPlan { + rows, + filters: HashMap::new(), + scan_range_after_filter: None, + } + } + + /// Build a batch's read ranges by resolving stable row ids through the + /// fragments' row-id sequences + fn plan_batch_from_row_ids( + ids: RowAddrTreeMap, + keys: &arrow_array::PrimitiveArray, + fragments: &StreamFragments, + ) -> FilteredReadInternalPlan { + let mut rows: BTreeMap>> = BTreeMap::new(); + let (Some(min_key), Some(max_key)) = (arrow::compute::min(keys), arrow::compute::max(keys)) + else { + // Every key is null + return FilteredReadInternalPlan { + rows, + filters: HashMap::new(), + scan_range_after_filter: None, + }; + }; + let requested = RowAddrMask::from_allowed(ids); + for (fragment, id_span) in fragments.fragments.iter().zip(&fragments.id_spans) { + // Only fragments whose id span overlaps the batch's key range + // can hold requested rows + let Some(id_span) = id_span else { continue }; + if *id_span.end() < min_key || *id_span.start() > max_key { + continue; + } + let offsets = fragment.row_id_sequence.mask_to_offset_ranges(&requested); + if offsets.is_empty() { + continue; + } + let valid = FilteredReadStream::full_frag_range( + fragment.num_physical_rows, + &fragment.deletion_vector, + ); + let matched = FilteredReadStream::intersect_ranges(&valid, &offsets); + if !matched.is_empty() { + rows.insert(fragment.fragment.id() as u32, matched); + } + } + FilteredReadInternalPlan { + rows, + filters: HashMap::new(), + scan_range_after_filter: None, + } + } + + fn key_array<'a>( + &self, + batch: &'a RecordBatch, + producer: &str, + ) -> DataFusionResult<&'a arrow_array::PrimitiveArray> { + let keys = batch + .column_by_name(self.source.key_column) + .ok_or_else(|| { + DataFusionError::Internal(format!( + "the row-stream {} is missing the '{}' column", + producer, self.source.key_column + )) + })?; + keys.as_primitive_opt::().ok_or_else(|| { + DataFusionError::Internal(format!( + "expected the row-stream column '{}' to be UInt64 but it was {}", + self.source.key_column, + keys.data_type() + )) + }) + } + + async fn plan_batch( + &self, + keys: &arrow_array::PrimitiveArray, + ) -> DataFusionResult { + let compute_timer = self.baseline_metrics.elapsed_compute().timer(); + // Null keys are excluded; attach_columns drops their rows + let batch_keys = if keys.null_count() == 0 { + RowAddrTreeMap::from_iter(keys.values().iter().copied()) + } else { + RowAddrTreeMap::from_iter(keys.iter().flatten()) + }; + drop(compute_timer); + + let fragments = self.load_fragments().await?; + // Row ids equal row addresses when the dataset does not use stable + // row ids, so either key resolves directly by position + if self.source.key_column == ROW_ADDR || !self.dataset.manifest.uses_stable_row_ids() { + Ok(Self::plan_batch_from_addresses(&batch_keys, fragments)) + } else { + Ok(Self::plan_batch_from_row_ids(batch_keys, keys, fragments)) + } + } + + /// Read the batch's planned ranges through the same executor as a scan, + /// returning the rows in storage order, deduplicated, with the key + /// column included + async fn read_batch( + &self, + internal_plan: FilteredReadInternalPlan, + batch_index: u32, + ) -> DataFusionResult { + let fragments = &self.load_fragments().await?.fragments; + // I/O priority: earlier batches strictly first (output emits in batch + // order), fragments keep dataset order within a batch + let priority_offset = batch_index.saturating_mul(fragments.len() as u32); + let read = FilteredReadStream::new_shared( + self.dataset.clone(), + self.source.read_options.clone(), + self.global_metrics.clone(), + internal_plan, + self.scan_scheduler.clone(), + fragments, + priority_offset, + ); + let decode_parallelism = match self.source.read_options.threading_mode { + FilteredReadThreadingMode::OnePartitionMultipleThreads(n) => n, + FilteredReadThreadingMode::MultiplePartitions(n) => n, + }; + let read_batches = read.collect_all(decode_parallelism.max(1)).await?; + Ok(arrow::compute::concat_batches( + &read.output_schema, + read_batches.iter(), + )?) + } + + /// Align the read rows back to the batch's row order and merge the + /// fetched columns on + fn attach_columns( + &self, + batch: RecordBatch, + read_data: RecordBatch, + ) -> DataFusionResult { + let _compute_timer = self.baseline_metrics.elapsed_compute().timer(); + let keys = self.key_array(&batch, "input")?; + let read_keys = self.key_array(&read_data, "read")?; + + // Fast path: one read row per input row with an identical key + // sequence — already aligned, skip the hash map and the permutation + if keys.null_count() == 0 + && read_data.num_rows() == batch.num_rows() + && read_keys.values() == keys.values() + { + let new_data = read_data.project_by_schema(self.source.new_fields_schema.as_ref())?; + let carried = batch.project_by_schema(self.carried_schema.as_ref())?; + return Ok(carried.merge_with_schema(&new_data, self.output_schema.as_ref())?); + } + + let key_to_index: HashMap = read_keys + .values() + .iter() + .enumerate() + .map(|(index, key)| (*key, index as u32)) + .collect(); + + // Sizes differ only when some input keys have no live row (null or + // stale keys): drop those input rows first + let batch = if read_data.num_rows() != batch.num_rows() { + let matched: BooleanArray = keys + .iter() + .map(|key| key.map(|key| key_to_index.contains_key(&key))) + .collect(); + arrow::compute::filter_record_batch(&batch, &matched)? + } else { + batch + }; + if batch.num_rows() == 0 { + return Ok(RecordBatch::new_empty(self.output_schema.clone())); + } - pub fn dataset(&self) -> &Arc { - &self.dataset + // Gather the read rows into input order — every remaining key hits + let keys = self.key_array(&batch, "input")?; + let indices = + UInt32Array::from_iter_values(keys.values().iter().map(|key| key_to_index[key])); + let new_data = arrow_select::take::take_record_batch(&read_data, &indices)?; + let new_data = new_data.project_by_schema(self.source.new_fields_schema.as_ref())?; + let carried = batch.project_by_schema(self.carried_schema.as_ref())?; + Ok(carried.merge_with_schema(&new_data, self.output_schema.as_ref())?) } - pub fn options(&self) -> &FilteredReadOptions { - &self.options + async fn execute_batch( + self: Arc, + batch: RecordBatch, + batch_index: u32, + ) -> DataFusionResult { + if batch.num_rows() == 0 { + return Ok(RecordBatch::new_empty(self.output_schema.clone())); + } + let internal_plan = self.plan_batch(self.key_array(&batch, "input")?).await?; + let read_data = self.read_batch(internal_plan, batch_index).await?; + self.attach_columns(batch, read_data) } - pub fn index_input(&self) -> Option<&Arc> { - self.index_input.as_ref() + /// Cut the input into batches of exactly `target` rows (final batch holds + /// the remainder): merging small batches amortizes per-batch planning, + /// slicing oversized ones bounds a batch's memory. Order is preserved. + fn coalesce_batches( + input: SendableRecordBatchStream, + target: usize, + ) -> impl Stream> { + struct CoalesceState { + input: SendableRecordBatchStream, + coalescer: BatchCoalescer, + exhausted: bool, + } + let coalescer = BatchCoalescer::new(input.schema(), target); + futures::stream::try_unfold( + CoalesceState { + input, + coalescer, + exhausted: false, + }, + |mut state| async move { + loop { + if let Some(batch) = state.coalescer.next_completed_batch() { + return Ok(Some((batch, state))); + } + if state.exhausted { + return Ok(None); + } + match state.input.try_next().await? { + Some(batch) => state.coalescer.push_batch(batch)?, + None => { + state.exhausted = true; + state.coalescer.finish_buffered_batch()?; + } + } + } + }, + ) } - /// Return the pre-computed plan if one exists, without triggering initialization. - pub fn plan(&self) -> Option { - self.plan.get().map(|p| p.to_external_plan()) + fn apply( + self: Arc, + input: SendableRecordBatchStream, + ) -> impl Stream> { + let batch_target_rows = self + .source + .read_options + .batch_size + .map(|batch_size| batch_size as usize) + .unwrap_or_else(|| get_default_batch_size().unwrap_or(BATCH_SIZE_FALLBACK)); + let on_result = self.clone(); + let on_done = self.clone(); + Self::coalesce_batches(input, batch_target_rows) + .enumerate() + .map(move |(batch_index, batch)| { + let batch = batch?; + let this = self.clone(); + DataFusionResult::Ok( + // SpawnedTask aborts on drop: cancelling the query + // cancels in-flight batches + SpawnedTask::spawn( + this.execute_batch(batch, batch_index as u32) + .in_current_span(), + ) + .map(|res| match res { + Ok(result) => result, + Err(join_error) => Err(DataFusionError::External(Box::new(join_error))), + }), + ) + }) + .boxed() + .try_buffered(ROW_STREAM_PREFETCH_BATCHES) + .map(move |result| { + on_result + .global_metrics + .io_metrics + .record(&on_result.scan_scheduler); + match on_result + .baseline_metrics + .record_poll(Poll::Ready(Some(result))) + { + Poll::Ready(Some(result)) => result, + _ => unreachable!("record_poll returned a different poll state"), + } + }) + .finally(move || { + on_done.baseline_metrics.done(); + on_done + .global_metrics + .io_metrics + .record(&on_done.scan_scheduler); + }) } } impl DisplayAs for FilteredReadExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result { + if let RowSelector::RowStream(source) = &self.input { + let columns = source + .new_fields_schema + .fields + .iter() + .map(|f| f.name().as_str()) + .collect::>() + .join(", "); + return match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "LanceRead: uri={}, projection=[{}], source=stream({})", + self.dataset.data_dir(), + columns, + source.key_column, + ) + } + DisplayFormatType::TreeRender => { + write!( + f, + "LanceRead\nuri={}\nprojection=[{}]\nsource=stream({})", + self.dataset.data_dir(), + columns, + source.key_column, + ) + } + }; + } let columns = self .options .projection @@ -1901,13 +2633,19 @@ impl ExecutionPlan for FilteredReadExec { } fn children(&self) -> Vec<&Arc> { - if let Some(index_input) = &self.index_input { - vec![index_input] + if let Some(child) = self.input.child() { + vec![child] } else { vec![] } } + fn benefits_from_input_partitioning(&self) -> Vec { + // Partitioning a row-stream read would create multiple I/O schedulers + // (RAM heavy); the other selectors have no row input + vec![false; self.children().len()] + } + fn metrics(&self) -> Option { Some(self.metrics.clone_inner()) } @@ -1916,6 +2654,13 @@ impl ExecutionPlan for FilteredReadExec { &self, partition: Option, ) -> datafusion::error::Result { + if let RowSelector::RowStream(source) = &self.input { + // At most one output row per input row + return Ok(Statistics { + num_rows: source.plan.partition_statistics(partition)?.num_rows, + ..Statistics::new_unknown(self.schema().as_ref()) + }); + } let fragments = self .options .fragments @@ -2038,18 +2783,12 @@ impl ExecutionPlan for FilteredReadExec { Error::internal("A FilteredReadExec cannot have two children".to_string()).into(), )) } else { - let index_input = children.into_iter().next(); - Ok(Arc::new(Self { - dataset: self.dataset.clone(), - options: self.options.clone(), - properties: self.properties.clone(), - metrics: self.metrics.clone(), - // Seems unlikely this would already be initialized but clear it - // out just in case - running_stream: Arc::new(AsyncMutex::new(None)), - index_input, - plan: Arc::new(OnceCell::new()), - })) + // Rebuild via try_new so the selector and derived state are + // re-derived from the new child's schema + let child = children.into_iter().next(); + let rebuilt = Self::try_new(self.dataset.clone(), self.options.clone(), child) + .map_err(|e| DataFusionError::External(e.into()))?; + Ok(Arc::new(rebuilt)) } } @@ -2058,10 +2797,16 @@ impl ExecutionPlan for FilteredReadExec { partition: usize, context: Arc, ) -> DataFusionResult { - Ok(self.obtain_stream(partition, context)) + match &self.input { + RowSelector::RowStream(source) => self.execute_row_stream(source, partition, context), + _ => Ok(self.obtain_stream(partition, context)), + } } fn fetch(&self) -> Option { + if self.row_stream_input().is_some() { + return None; + } if self.options.full_filter.is_none() { self.options .scan_range_before_filter @@ -2076,13 +2821,15 @@ impl ExecutionPlan for FilteredReadExec { } fn supports_limit_pushdown(&self) -> bool { - // This is to push the limit through the node and into an upstream node. - // The only upstream node is the index search and we can't push the limit - // to that node. - false + // A limit pushes through to a row-stream input (one output row per + // input row); the other selectors have no node to push it to + self.row_stream_input().is_some() } fn with_fetch(&self, limit: Option) -> Option> { + if self.row_stream_input().is_some() { + return None; + } // TODO: Support multiple partitions in the future by coordinating limits across partitions if matches!( self.options.threading_mode, @@ -2113,7 +2860,7 @@ impl ExecutionPlan for FilteredReadExec { match Self::try_new( self.dataset.clone(), updated_options, - self.index_input.clone(), + self.input.row_set_plan().cloned(), ) { Ok(exec) => Some(Arc::new(exec)), Err(e) => { @@ -2376,7 +3123,7 @@ mod tests { )) } - /// Round-trip every interval shape through the arrow wire format and + /// Batch-trip every interval shape through the arrow wire format and /// confirm the endpoints survive. Exercises both /// `IndexExprResult::serialize` and `EvaluatedIndex::try_from_arrow` /// so the schema names stay in sync. @@ -2410,15 +3157,15 @@ mod tests { // robust, but the canonical builders preserve representation. assert_eq!( decoded.index_result.lower, original.lower, - "{name}: lower endpoint changed across round-trip", + "{name}: lower endpoint changed across batch-trip", ); assert_eq!( decoded.index_result.upper, original.upper, - "{name}: upper endpoint changed across round-trip", + "{name}: upper endpoint changed across batch-trip", ); assert_eq!( decoded.applicable_fragments, frags, - "{name}: applicable fragments changed across round-trip", + "{name}: applicable fragments changed across batch-trip", ); } } @@ -2765,6 +3512,48 @@ mod tests { assert_eq!(num_rows, 300); } + /// A stale (not rebuilt after a delete) index hit drops on the live view + /// and returns as a null-_rowid tombstone with with_deleted_rows + #[test_log::test(tokio::test)] + async fn test_with_deleted_rows_stale_index() { + let fixture = Arc::new(TestFixture::new().await); + let base_options = FilteredReadOptions::basic_full_read(&fixture.dataset); + + // Row 220 is deletion-vector-deleted but still in the index + let filter_plan = fixture.filter_plan("fully_indexed == 220", true).await; + + // Live view: the stale index hit drops + fixture + .test_plan( + base_options.clone().with_filter_plan(filter_plan), + &u32s(vec![]), + ) + .await; + + // Physical view: the tombstone returns + let filter_plan = fixture.filter_plan("fully_indexed == 220", true).await; + let options = base_options + .with_deleted_rows() + .unwrap() + .with_filter_plan(filter_plan); + let plan = fixture.make_plan(options).await; + let stream = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); + let batches = stream.try_collect::>().await.unwrap(); + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, 1); + let batch = batches.iter().find(|b| b.num_rows() > 0).unwrap(); + let values = batch + .column_by_name("fully_indexed") + .unwrap() + .as_primitive::(); + assert_eq!(values.value(0), 220); + let row_ids = batch + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::(); + assert!(row_ids.is_null(0)); + } + #[test] fn test_dv_to_ranges() { let dv = Arc::new(DeletionVector::from_iter(vec![1])); @@ -3051,7 +3840,7 @@ mod tests { let options = base_options.with_filter_plan(filter_plan); let plan = fixture.make_plan(options).await; - assert!(plan.index_input.is_some()); + assert!(plan.index_input().is_some()); assert!(plan.options().refine_filter.is_some()); let limited_plan = plan.with_fetch(Some(10)).unwrap(); @@ -3780,7 +4569,7 @@ mod tests { /// Test that direct execution gives the same result as get_plan + execute_with_plan #[test_log::test(tokio::test)] - async fn test_plan_round_trip() { + async fn test_plan_batch_trip() { let fixture = TestFixture::new().await; let ctx = Arc::new(TaskContext::default()); @@ -3893,4 +4682,793 @@ mod tests { assert_eq!(default_result.num_rows(), capped_result.num_rows()); } + + // Row-stream selector tests + + mod row_stream { + use super::*; + use arrow_array::{Float32Array, StringArray, UInt64Array}; + use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema}; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use lance_datafusion::exec::OneShotExec; + use rstest::rstest; + + use crate::dataset::{Dataset, WriteParams}; + use crate::utils::test::NoContextTestFixture; + + struct TakeFixture { + dataset: Arc, + _tmp_dir: TempStrDir, + } + + /// 30 rows across 3 fragments with columns i, s, and struct{x, y} + async fn take_fixture(stable_row_ids: bool) -> TakeFixture { + let struct_fields = Fields::from(vec![ + Arc::new(ArrowField::new("x", DataType::Int32, false)), + Arc::new(ArrowField::new("y", DataType::Int32, false)), + ]); + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("i", DataType::Int32, false), + ArrowField::new("s", DataType::Utf8, false), + ArrowField::new("struct", DataType::Struct(struct_fields.clone()), false), + ])); + let batches: Vec = (0..3) + .map(|batch_id| { + let value_range = batch_id * 10..batch_id * 10 + 10; + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(value_range.clone())), + Arc::new(StringArray::from_iter_values( + value_range.clone().map(|v| format!("s-{v}")), + )), + Arc::new(arrow_array::StructArray::new( + struct_fields.clone(), + vec![ + Arc::new(Int32Array::from_iter(value_range.clone())), + Arc::new(Int32Array::from_iter(value_range)), + ], + None, + )), + ], + ) + .unwrap() + }) + .collect(); + + let tmp_dir = TempStrDir::default(); + let uri = tmp_dir.as_str(); + let params = WriteParams { + max_rows_per_file: 10, + enable_stable_row_ids: stable_row_ids, + ..Default::default() + }; + let reader = RecordBatchIterator::new(batches.into_iter().map(Ok), schema); + Dataset::write(reader, uri, Some(params)).await.unwrap(); + TakeFixture { + dataset: Arc::new(Dataset::open(uri).await.unwrap()), + _tmp_dir: tmp_dir, + } + } + + /// Wrap batches of (payload, key) rows into an input plan + fn rows_input(batches: Vec) -> Arc { + let schema = batches[0].schema(); + let stream = futures::stream::iter(batches.into_iter().map(Ok)); + let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); + Arc::new(OneShotExec::new(stream)) + } + + fn take_plan( + dataset: &Arc, + input: Arc, + columns: &[&str], + ) -> Result { + let projection = dataset + .empty_projection() + .union_columns(columns, OnMissing::Error) + .unwrap(); + FilteredReadExec::try_new( + dataset.clone(), + FilteredReadOptions::new(projection), + Some(input), + ) + } + + fn take_plan_sized( + dataset: &Arc, + input: Arc, + columns: &[&str], + batch_size: u32, + ) -> Result { + let projection = dataset + .empty_projection() + .union_columns(columns, OnMissing::Error) + .unwrap(); + FilteredReadExec::try_new( + dataset.clone(), + FilteredReadOptions::new(projection).with_batch_size(batch_size), + Some(input), + ) + } + + async fn run(plan: &FilteredReadExec) -> Vec { + plan.execute(0, Arc::new(TaskContext::default())) + .unwrap() + .try_collect::>() + .await + .unwrap() + } + + /// Output preserves the input's row order, duplicates, and payload + #[rstest] + #[case::by_row_addr(false, ROW_ADDR)] + #[case::by_row_id(false, ROW_ID)] + #[case::stable_by_row_addr(true, ROW_ADDR)] + #[case::stable_by_row_id(true, ROW_ID)] + #[tokio::test] + async fn take_preserves_order_dups_and_payload( + #[case] stable_row_ids: bool, + #[case] key: &str, + ) { + let fixture = take_fixture(stable_row_ids).await; + + // Stable row ids are assigned sequentially on write, so the id of + // row `i` is `i`; without them id == address + let addr = |frag: u64, off: u64| (frag << 32) | off; + let keys: Vec = if key == ROW_ID && stable_row_ids { + vec![21, 3, 15, 21, 0] + } else { + vec![ + addr(2, 1), // i = 21 + addr(0, 3), // i = 3 + addr(1, 5), // i = 15 + addr(2, 1), // i = 21 (duplicate) + addr(0, 0), // i = 0 + ] + }; + let expected_i: Vec = vec![21, 3, 15, 21, 0]; + + let input_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("payload", DataType::Float32, false), + ArrowField::new(key, DataType::UInt64, true), + ])); + let payload: Vec = (0..keys.len()).map(|v| v as f32 * 0.5).collect(); + let batch = RecordBatch::try_new( + input_schema.clone(), + vec![ + Arc::new(Float32Array::from(payload.clone())), + Arc::new(UInt64Array::from(keys.clone())), + ], + ) + .unwrap(); + let batches = vec![batch.slice(0, 3), batch.slice(3, 2)]; + + let plan = take_plan(&fixture.dataset, rows_input(batches), &["s", "i"]).unwrap(); + assert!(plan.row_stream_input().is_some()); + // Input columns, then new fields; the unrequested key is stripped + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>(), + vec!["payload", "i", "s"] + ); + + let result = run(&plan).await; + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_rows(), 5); + let result = concat_batches(&plan.schema(), &result).unwrap(); + + let i_col = result.column_by_name("i").unwrap(); + assert_eq!( + i_col.as_primitive::().values(), + &expected_i[..] + ); + let s_col = result.column_by_name("s").unwrap().as_string::(); + for (row, i) in expected_i.iter().enumerate() { + assert_eq!(s_col.value(row), format!("s-{i}")); + } + let payload_col = result + .column_by_name("payload") + .unwrap() + .as_primitive::(); + assert_eq!(payload_col.values(), &payload[..]); + } + + /// Tiny input batches merge into batches and oversized ones slice, + /// preserving order across the boundaries + #[tokio::test] + async fn take_coalesces_input_to_batch_size() { + let fixture = take_fixture(false).await; + + let addr = |frag: u64, off: u64| (frag << 32) | off; + let keys: Vec = vec![ + addr(2, 3), // i = 23 + addr(0, 1), // i = 1 + addr(1, 4), // i = 14 + addr(0, 7), // i = 7 + addr(2, 0), // i = 20 + addr(1, 1), // i = 11 + addr(0, 2), // i = 2 + ]; + let expected_i: Vec = vec![23, 1, 14, 7, 20, 11, 2]; + let input_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + ROW_ADDR, + DataType::UInt64, + true, + )])); + let batch = RecordBatch::try_new( + input_schema.clone(), + vec![Arc::new(UInt64Array::from(keys.clone()))], + ) + .unwrap(); + + let assert_batches = |result: Vec, schema: SchemaRef| { + assert_eq!( + result.iter().map(|b| b.num_rows()).collect::>(), + vec![3, 3, 1] + ); + let merged = concat_batches(&schema, &result).unwrap(); + let i_col = merged.column_by_name("i").unwrap(); + assert_eq!( + i_col.as_primitive::().values(), + &expected_i[..] + ); + }; + + // Seven one-row batches merge into batches of exactly 3, 3, 1 + let tiny = (0..7).map(|i| batch.slice(i, 1)).collect::>(); + let plan = take_plan_sized(&fixture.dataset, rows_input(tiny), &["i"], 3).unwrap(); + assert_batches(run(&plan).await, plan.schema()); + + // One oversized batch is sliced into the same batches + let plan = + take_plan_sized(&fixture.dataset, rows_input(vec![batch]), &["i"], 3).unwrap(); + assert_batches(run(&plan).await, plan.schema()); + } + + /// Storage-ordered input exercises the aligned fast path + #[rstest] + #[case::by_row_addr(ROW_ADDR)] + #[case::by_row_id(ROW_ID)] + #[tokio::test] + async fn take_aligned_input_fast_path(#[case] key: &str) { + let fixture = take_fixture(false).await; + + let addr = |frag: u64, off: u64| (frag << 32) | off; + let keys: Vec = vec![ + addr(0, 0), // i = 0 + addr(0, 3), // i = 3 + addr(1, 5), // i = 15 + addr(2, 1), // i = 21 + addr(2, 9), // i = 29 + ]; + let expected_i: Vec = vec![0, 3, 15, 21, 29]; + + let input_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("payload", DataType::Float32, false), + ArrowField::new(key, DataType::UInt64, true), + ])); + let payload: Vec = (0..keys.len()).map(|v| v as f32 * 0.5).collect(); + let batch = RecordBatch::try_new( + input_schema, + vec![ + Arc::new(Float32Array::from(payload.clone())), + Arc::new(UInt64Array::from(keys)), + ], + ) + .unwrap(); + + let plan = take_plan(&fixture.dataset, rows_input(vec![batch]), &["i"]).unwrap(); + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + assert_eq!(result.num_rows(), 5); + let i_col = result.column_by_name("i").unwrap(); + assert_eq!( + i_col.as_primitive::().values(), + &expected_i[..] + ); + let payload_col = result + .column_by_name("payload") + .unwrap() + .as_primitive::(); + assert_eq!(payload_col.values(), &payload[..]); + } + + /// A fragment-scoped take reads from the scoped fragments only; keys + /// pointing outside the scope drop like stale rows + #[tokio::test] + async fn take_scoped_to_fragments() { + let fixture = take_fixture(false).await; + let subset = Arc::new(vec![fixture.dataset.fragments()[1].clone()]); + + let addr = |frag: u64, off: u64| (frag << 32) | off; + let keys: Vec = vec![ + addr(1, 2), // i = 12, inside the scoped fragment + addr(0, 3), // i = 3, outside the scope + ]; + let input_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + ROW_ADDR, + DataType::UInt64, + true, + )])); + let batch = RecordBatch::try_new(input_schema, vec![Arc::new(UInt64Array::from(keys))]) + .unwrap(); + + let projection = fixture + .dataset + .empty_projection() + .union_columns(["i"], OnMissing::Error) + .unwrap(); + let plan = FilteredReadExec::try_new( + fixture.dataset.clone(), + FilteredReadOptions::new(projection).with_fragments(subset), + Some(rows_input(vec![batch])), + ) + .unwrap(); + + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + assert_eq!(result.num_rows(), 1); + let i_col = result + .column_by_name("i") + .unwrap() + .as_primitive::(); + assert_eq!(i_col.value(0), 12); + } + + /// Identity flags: requested-but-missing columns are synthesized, + /// carried ones kept, unrequested carried ones stripped + #[rstest] + #[case::unstable(false)] + #[case::stable(true)] + #[tokio::test] + async fn take_identity_flags(#[case] stable_row_ids: bool) { + let fixture = take_fixture(stable_row_ids).await; + let addr = |frag: u64, off: u64| (frag << 32) | off; + + let ids: Vec = if stable_row_ids { + vec![21, 3] + } else { + vec![addr(2, 1), addr(0, 3)] + }; + let expected_addrs: Vec = vec![addr(2, 1), addr(0, 3)]; + let expected_i: Vec = vec![21, 3]; + let id_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + ROW_ID, + DataType::UInt64, + true, + )])); + let id_batch = + RecordBatch::try_new(id_schema, vec![Arc::new(UInt64Array::from(ids.clone()))]) + .unwrap(); + + // Keep the carried _rowid and synthesize _rowaddr + let projection = fixture + .dataset + .empty_projection() + .union_columns(["i"], OnMissing::Error) + .unwrap() + .with_row_id() + .with_row_addr(); + let plan = FilteredReadExec::try_new( + fixture.dataset.clone(), + FilteredReadOptions::new(projection), + Some(rows_input(vec![id_batch.clone()])), + ) + .unwrap(); + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>(), + vec![ROW_ID, "i", ROW_ADDR] + ); + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + let id_col = result + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::(); + assert_eq!(id_col.values(), &ids[..]); + let addr_col = result + .column_by_name(ROW_ADDR) + .unwrap() + .as_primitive::(); + assert_eq!(addr_col.values(), &expected_addrs[..]); + let i_col = result + .column_by_name("i") + .unwrap() + .as_primitive::(); + assert_eq!(i_col.values(), &expected_i[..]); + + // Synthesize _rowaddr but strip the unrequested carried _rowid + let projection = fixture + .dataset + .empty_projection() + .union_columns(["i"], OnMissing::Error) + .unwrap() + .with_row_addr(); + let plan = FilteredReadExec::try_new( + fixture.dataset.clone(), + FilteredReadOptions::new(projection), + Some(rows_input(vec![id_batch.clone()])), + ) + .unwrap(); + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>(), + vec!["i", ROW_ADDR] + ); + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + let addr_col = result + .column_by_name(ROW_ADDR) + .unwrap() + .as_primitive::(); + assert_eq!(addr_col.values(), &expected_addrs[..]); + + // Address-keyed input, synthesize _rowid (the reverse direction) + let addr_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + ROW_ADDR, + DataType::UInt64, + true, + )])); + let addr_batch = RecordBatch::try_new( + addr_schema, + vec![Arc::new(UInt64Array::from(expected_addrs.clone()))], + ) + .unwrap(); + let projection = fixture + .dataset + .empty_projection() + .union_columns(["i"], OnMissing::Error) + .unwrap() + .with_row_id(); + let plan = FilteredReadExec::try_new( + fixture.dataset.clone(), + FilteredReadOptions::new(projection), + Some(rows_input(vec![addr_batch])), + ) + .unwrap(); + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>(), + vec!["i", ROW_ID] + ); + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + let id_col = result + .column_by_name(ROW_ID) + .unwrap() + .as_primitive::(); + assert_eq!(id_col.values(), &ids[..]); + + // Fetch nothing, synthesize only (the AddRowAddrExec shape) + let projection = fixture + .dataset + .empty_projection() + .with_row_id() + .with_row_addr(); + let plan = FilteredReadExec::try_new( + fixture.dataset.clone(), + FilteredReadOptions::new(projection), + Some(rows_input(vec![id_batch])), + ) + .unwrap(); + assert_eq!( + plan.schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect::>(), + vec![ROW_ID, ROW_ADDR] + ); + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + let addr_col = result + .column_by_name(ROW_ADDR) + .unwrap() + .as_primitive::(); + assert_eq!(addr_col.values(), &expected_addrs[..]); + } + + /// New sub-fields merge into an existing struct column + #[tokio::test] + async fn take_merges_nested_struct() { + let fixture = take_fixture(false).await; + + let data = fixture + .dataset + .scan() + .project(&["struct"]) + .unwrap() + .with_row_id() + .try_into_batch() + .await + .unwrap(); + // Rebuild the input with only struct.y so struct.x must be taken + let full_struct = data.column_by_name("struct").unwrap().as_struct(); + let y_only = arrow_array::StructArray::new( + Fields::from(vec![Arc::new(ArrowField::new("y", DataType::Int32, false))]), + vec![full_struct.column_by_name("y").unwrap().clone()], + None, + ); + let input_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("struct", y_only.data_type().clone(), false), + ArrowField::new(ROW_ID, DataType::UInt64, true), + ])); + let data = RecordBatch::try_new( + input_schema, + vec![ + Arc::new(y_only), + data.column_by_name(ROW_ID).unwrap().clone(), + ], + ) + .unwrap(); + + let projection = fixture + .dataset + .empty_projection() + .union_column("struct.x", OnMissing::Error) + .unwrap(); + let plan = FilteredReadExec::try_new( + fixture.dataset.clone(), + FilteredReadOptions::new(projection), + Some(rows_input(vec![data])), + ) + .unwrap(); + + let expected_struct_type = DataType::Struct(Fields::from(vec![ + Arc::new(ArrowField::new("x", DataType::Int32, false)), + Arc::new(ArrowField::new("y", DataType::Int32, false)), + ])); + assert_eq!( + plan.schema().field_with_name("struct").unwrap().data_type(), + &expected_struct_type + ); + + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + assert_eq!(result.num_rows(), 30); + let struct_col = result.column_by_name("struct").unwrap().as_struct(); + assert_eq!( + struct_col.column_by_name("x").unwrap(), + struct_col.column_by_name("y").unwrap() + ); + } + + /// Input rows whose key no longer exists (deleted rows) are dropped + #[rstest] + #[case::by_row_addr(false, ROW_ADDR)] + #[case::by_row_id(false, ROW_ID)] + #[case::stable_by_row_addr(true, ROW_ADDR)] + #[case::stable_by_row_id(true, ROW_ID)] + #[tokio::test] + async fn take_drops_stale_keys(#[case] stable_row_ids: bool, #[case] key: &str) { + let fixture = take_fixture(stable_row_ids).await; + let mut dataset = fixture.dataset.as_ref().clone(); + dataset.delete("i = 15").await.unwrap(); + let dataset = Arc::new(dataset); + + let addr = |frag: u64, off: u64| (frag << 32) | off; + // The pre-delete identifiers of rows 15 (now deleted) and 16 + let keys: Vec = if key == ROW_ID && stable_row_ids { + vec![15, 16] + } else { + vec![addr(1, 5), addr(1, 6)] + }; + let input_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + key, + DataType::UInt64, + true, + )])); + let batch = RecordBatch::try_new(input_schema, vec![Arc::new(UInt64Array::from(keys))]) + .unwrap(); + + let plan = take_plan(&dataset, rows_input(vec![batch]), &["i"]).unwrap(); + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + assert_eq!(result.num_rows(), 1); + let i_col = result + .column_by_name("i") + .unwrap() + .as_primitive::(); + assert_eq!(i_col.value(0), 16); + } + + /// Keys of a fully deleted fragment (gone from the manifest) are + /// dropped like stale rows + #[rstest] + #[case::by_row_addr(false, ROW_ADDR)] + #[case::by_row_id(false, ROW_ID)] + #[case::stable_by_row_addr(true, ROW_ADDR)] + #[case::stable_by_row_id(true, ROW_ID)] + #[tokio::test] + async fn take_drops_keys_of_deleted_fragment( + #[case] stable_row_ids: bool, + #[case] key: &str, + ) { + let fixture = take_fixture(stable_row_ids).await; + let mut dataset = fixture.dataset.as_ref().clone(); + dataset.delete("i >= 10 and i < 20").await.unwrap(); + let dataset = Arc::new(dataset); + + let addr = |frag: u64, off: u64| (frag << 32) | off; + // The pre-delete identifiers of row 15 (fragment 1, now gone + // from the manifest) and row 20 (fragment 2, still live) + let keys: Vec = if key == ROW_ID && stable_row_ids { + vec![15, 20] + } else { + vec![addr(1, 5), addr(2, 0)] + }; + let input_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + key, + DataType::UInt64, + true, + )])); + let batch = RecordBatch::try_new(input_schema, vec![Arc::new(UInt64Array::from(keys))]) + .unwrap(); + + let plan = take_plan(&dataset, rows_input(vec![batch]), &["i"]).unwrap(); + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + assert_eq!(result.num_rows(), 1); + let i_col = result + .column_by_name("i") + .unwrap() + .as_primitive::(); + assert_eq!(i_col.value(0), 20); + } + + /// After delete + compaction the stable row-id sequences are no + /// longer simple contiguous ranges; ids must still resolve to the + /// moved rows and deleted ids must still drop + #[tokio::test] + async fn take_stable_ids_after_compaction() { + use crate::dataset::optimize::{CompactionOptions, compact_files}; + + let fixture = take_fixture(true).await; + let mut dataset = fixture.dataset.as_ref().clone(); + // Punch holes, then rewrite all fragments into one + dataset.delete("i % 3 = 0").await.unwrap(); + compact_files(&mut dataset, CompactionOptions::default(), None) + .await + .unwrap(); + let dataset = Arc::new(dataset); + assert_eq!(dataset.get_fragments().len(), 1); + + // Survivors in scattered order, plus a compacted-away id (15) + let keys: Vec = vec![25, 1, 15, 14]; + let input_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + ROW_ID, + DataType::UInt64, + true, + )])); + let batch = RecordBatch::try_new(input_schema, vec![Arc::new(UInt64Array::from(keys))]) + .unwrap(); + + let plan = take_plan(&dataset, rows_input(vec![batch]), &["i"]).unwrap(); + let result = concat_batches(&plan.schema(), &run(&plan).await).unwrap(); + let i_col = result + .column_by_name("i") + .unwrap() + .as_primitive::(); + assert_eq!(i_col.values(), &[25, 1, 14]); + } + + /// with_deleted_rows is rejected for a row-stream read + #[rstest] + #[case::unstable(false)] + #[case::stable(true)] + #[tokio::test] + async fn take_rejects_with_deleted_rows(#[case] stable_row_ids: bool) { + let fixture = take_fixture(stable_row_ids).await; + let input_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + ROW_ADDR, + DataType::UInt64, + true, + )])); + let batch = + RecordBatch::try_new(input_schema, vec![Arc::new(UInt64Array::from(vec![0_u64]))]) + .unwrap(); + let projection = fixture + .dataset + .empty_projection() + .union_columns(["i"], OnMissing::Error) + .unwrap(); + let err = FilteredReadExec::try_new( + fixture.dataset, + FilteredReadOptions::new(projection) + .with_deleted_rows() + .unwrap(), + Some(rows_input(vec![batch])), + ) + .unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + assert!(err.to_string().contains("with_deleted_rows")); + } + + /// Construction errors: no key column, nothing to read + #[tokio::test] + async fn take_construction_errors() { + let fixture = take_fixture(false).await; + let no_key_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "payload", + DataType::UInt64, + true, + )])); + let batch = RecordBatch::try_new( + no_key_schema, + vec![Arc::new(UInt64Array::from(vec![0_u64]))], + ) + .unwrap(); + let err = take_plan(&fixture.dataset, rows_input(vec![batch]), &["s"]).unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + assert!(err.to_string().contains("must have a column")); + + // Taking fields the input already has: nothing to read + let with_s_schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new(ROW_ADDR, DataType::UInt64, true), + ArrowField::new("s", DataType::Utf8, false), + ])); + let with_s_batch = RecordBatch::try_new( + with_s_schema, + vec![ + Arc::new(UInt64Array::from(vec![0_u64])), + Arc::new(StringArray::from(vec!["x"])), + ], + ) + .unwrap(); + let err = + take_plan(&fixture.dataset, rows_input(vec![with_s_batch]), &["s"]).unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + assert!(err.to_string().contains("nothing to read")); + } + + /// with_new_children re-derives the row-stream source and preserves the schema + #[tokio::test] + async fn take_with_new_children() { + let fixture = take_fixture(false).await; + let input_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + ROW_ID, + DataType::UInt64, + true, + )])); + let batch = RecordBatch::try_new( + input_schema, + vec![Arc::new(UInt64Array::from(vec![0_u64, 1]))], + ) + .unwrap(); + let input = rows_input(vec![batch]); + let plan: Arc = + Arc::new(take_plan(&fixture.dataset, input.clone(), &["s"]).unwrap()); + let rebuilt = plan.clone().with_new_children(vec![input]).unwrap(); + assert_eq!(plan.schema(), rebuilt.schema()); + assert!( + rebuilt + .as_any() + .downcast_ref::() + .unwrap() + .row_stream_input() + .is_some() + ); + } + + /// Take-mode nodes can be created and executed without an active + /// tokio runtime (required for DataFusion foreign table providers) + #[test] + fn no_context_take_rows() { + use lance_datafusion::datagen::DatafusionDatagenExt; + use lance_datagen::{BatchCount, RowCount}; + + let fixture = NoContextTestFixture::new(); + let dataset = Arc::new(fixture.dataset); + let input = lance_datagen::gen_batch() + .col(ROW_ID, lance_datagen::array::step::()) + .into_df_exec(RowCount::from(50), BatchCount::from(2)); + let plan = take_plan(&dataset, input, &["text"]).unwrap(); + plan.execute(0, Arc::new(TaskContext::default())).unwrap(); + } + } } diff --git a/rust/lance/src/io/exec/filtered_read_proto.rs b/rust/lance/src/io/exec/filtered_read_proto.rs index 4eb329506aa..9f08bd102e7 100644 --- a/rust/lance/src/io/exec/filtered_read_proto.rs +++ b/rust/lance/src/io/exec/filtered_read_proto.rs @@ -46,6 +46,15 @@ pub async fn filtered_read_exec_to_proto( exec: &FilteredReadExec, state: &SessionState, ) -> Result { + if exec.row_stream_input().is_some() { + // TODO: define a proto representation for row-stream sources so these + // reads can participate in distributed plan pushdown + return Err(Error::not_supported_source( + "a FilteredReadExec with a row-stream source cannot be serialized" + .to_string() + .into(), + )); + } let table = table_identifier_from_dataset(exec.dataset()).await?; // Use the pruned dataset schema for filter encoding — filters can reference columns // outside the projection (e.g. SELECT name WHERE age > 10), and some dataset columns diff --git a/rust/lance/src/io/exec/optimizer.rs b/rust/lance/src/io/exec/optimizer.rs index 72488f3a14e..c5b7a7fab9c 100644 --- a/rust/lance/src/io/exec/optimizer.rs +++ b/rust/lance/src/io/exec/optimizer.rs @@ -6,6 +6,7 @@ use std::sync::Arc; use super::TakeExec; +use super::filtered_read::FilteredReadExec; use arrow_schema::Schema as ArrowSchema; #[allow(deprecated)] use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; @@ -18,17 +19,38 @@ use datafusion::{ }; use datafusion_physical_expr::{PhysicalExpr, expressions::Column}; -/// Rule that eliminates [TakeExec] nodes that are immediately followed by another [TakeExec]. +/// Rule that eliminates take nodes that are immediately followed by another +/// take node, fetching the union of the columns in a single node instead. +/// +/// A "take" is either a [TakeExec] (legacy storage) or a [FilteredReadExec] +/// with a row-stream source (see `FilteredReadExec::row_stream_input`); the +/// scanner emits stacked takes in some plan shapes (e.g. filter columns then +/// projection columns). #[derive(Debug)] pub struct CoalesceTake; impl CoalesceTake { + /// Whether `plan` is a take node this rule knows how to collapse + fn as_take(plan: &Arc) -> Option<&dyn ExecutionPlan> { + if plan.as_any().is::() { + Some(plan.as_ref()) + } else if let Some(filtered_read) = plan.as_any().downcast_ref::() { + filtered_read + .row_stream_input() + .is_some() + .then_some(plan.as_ref()) + } else { + None + } + } + fn field_order_differs(old_schema: &ArrowSchema, new_schema: &ArrowSchema) -> bool { - old_schema - .fields - .iter() - .zip(&new_schema.fields) - .any(|(old, new)| old.name() != new.name()) + old_schema.fields.len() != new_schema.fields.len() + || old_schema + .fields + .iter() + .zip(&new_schema.fields) + .any(|(old, new)| old.name() != new.name()) } fn remap_collapsed_output( @@ -47,24 +69,39 @@ impl CoalesceTake { Arc::new(ProjectionExec::try_new(project_exprs, plan).unwrap()) } + /// Collapse two stacked takes into one, or return None when the rebuilt + /// node would not produce every column of the original output (the + /// rebuild re-derives what to fetch from the outer take's projection, so + /// a column only the inner take fetched can go missing if the outer + /// projection doesn't cover it) fn collapse_takes( - inner_take: &TakeExec, - outer_take: &TakeExec, + inner_take: &dyn ExecutionPlan, + outer_take: &dyn ExecutionPlan, outer_exec: Arc, - ) -> Arc { + ) -> Option> { let inner_take_input = inner_take.children()[0].clone(); let old_output_schema = outer_take.schema(); - let collapsed = outer_exec - .with_new_children(vec![inner_take_input]) - .unwrap(); + let collapsed = outer_exec.with_new_children(vec![inner_take_input]).ok()?; let new_output_schema = collapsed.schema(); + if old_output_schema + .fields() + .iter() + .any(|field| new_output_schema.field_with_name(field.name()).is_err()) + { + return None; + } + // It's possible that collapsing the take can change the field order. This disturbs DF's planner and // so we must restore it. if Self::field_order_differs(&old_output_schema, &new_output_schema) { - Self::remap_collapsed_output(&old_output_schema, &new_output_schema, collapsed) + Some(Self::remap_collapsed_output( + &old_output_schema, + &new_output_schema, + collapsed, + )) } else { - collapsed + Some(collapsed) } } } @@ -78,26 +115,25 @@ impl PhysicalOptimizerRule for CoalesceTake { ) -> DFResult> { Ok(plan .transform_down(|plan| { - if let Some(outer_take) = plan.as_any().downcast_ref::() { + if let Some(outer_take) = Self::as_take(&plan) { let child = outer_take.children()[0]; - // Case 1: TakeExec -> TakeExec - if let Some(inner_take) = child.as_any().downcast_ref::() { - return Ok(Transformed::yes(Self::collapse_takes( - inner_take, - outer_take, - plan.clone(), - ))); - // Case 2: TakeExec -> CoalesceBatchesExec -> TakeExec + // Case 1: take -> take + if let Some(inner_take) = Self::as_take(child) { + if let Some(collapsed) = + Self::collapse_takes(inner_take, outer_take, plan.clone()) + { + return Ok(Transformed::yes(collapsed)); + } + // Case 2: take -> CoalesceBatchesExec -> take } else if let Some(exec_child) = child.as_any().downcast_ref::() { let inner_child = exec_child.children()[0].clone(); - if let Some(inner_take) = inner_child.as_any().downcast_ref::() { - return Ok(Transformed::yes(Self::collapse_takes( - inner_take, - outer_take, - plan.clone(), - ))); + if let Some(inner_take) = Self::as_take(&inner_child) + && let Some(collapsed) = + Self::collapse_takes(inner_take, outer_take, plan.clone()) + { + return Ok(Transformed::yes(collapsed)); } } } @@ -184,3 +220,196 @@ pub fn get_physical_optimizer() -> PhysicalOptimizer { Arc::new(datafusion::physical_optimizer::enforce_distribution::EnforceDistribution::new()), ]) } + +#[cfg(test)] +mod tests { + use super::*; + + use arrow_array::cast::AsArray; + use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator, StringArray, UInt64Array}; + use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; + use arrow_select::concat::concat_batches; + use datafusion::execution::TaskContext; + use datafusion::physical_plan::stream::RecordBatchStreamAdapter; + use futures::TryStreamExt; + use lance_core::ROW_ID; + use lance_core::datatypes::OnMissing; + use lance_core::utils::tempfile::TempStrDir; + use lance_datafusion::exec::OneShotExec; + use lance_file::version::LanceFileVersion; + + use crate::dataset::{Dataset, WriteParams}; + use crate::io::exec::filtered_read::{FilteredReadExec, FilteredReadOptions}; + + /// 20 rows, one fragment, columns i (Int32) and s (Utf8) + async fn fixture(storage_version: LanceFileVersion) -> (Arc, TempStrDir) { + let schema = Arc::new(ArrowSchema::new(vec![ + ArrowField::new("i", DataType::Int32, false), + ArrowField::new("s", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..20)), + Arc::new(StringArray::from_iter_values( + (0..20).map(|v| format!("s-{v}")), + )), + ], + ) + .unwrap(); + let tmp_dir = TempStrDir::default(); + let uri = tmp_dir.as_str(); + let params = WriteParams { + data_storage_version: Some(storage_version), + ..Default::default() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + Dataset::write(reader, uri, Some(params)).await.unwrap(); + (Arc::new(Dataset::open(uri).await.unwrap()), tmp_dir) + } + + /// An input plan producing one batch of `_rowid` keys + fn keys_input(keys: Vec) -> Arc { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + ROW_ID, + DataType::UInt64, + true, + )])); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(UInt64Array::from(keys))]).unwrap(); + let stream = futures::stream::iter(vec![Ok(batch)]); + let stream = Box::pin(RecordBatchStreamAdapter::new(schema, stream)); + Arc::new(OneShotExec::new(stream)) + } + + fn row_stream_take( + dataset: &Arc, + input: Arc, + columns: &[&str], + ) -> Arc { + // Mirror Scanner::take: full target projection, carried identity kept + let mut projection = dataset + .empty_projection() + .union_columns(columns, OnMissing::Error) + .unwrap(); + projection.with_row_id = true; + Arc::new( + FilteredReadExec::try_new( + dataset.clone(), + FilteredReadOptions::new(projection), + Some(input), + ) + .unwrap(), + ) + } + + async fn run(plan: &Arc) -> RecordBatch { + let stream = plan.execute(0, Arc::new(TaskContext::default())).unwrap(); + let schema = stream.schema(); + let batches: Vec<_> = stream.try_collect().await.unwrap(); + concat_batches(&schema, batches.iter()).unwrap() + } + + fn count_takes(plan: &Arc) -> usize { + let self_count = CoalesceTake::as_take(plan).map(|_| 1).unwrap_or(0); + self_count + + plan + .children() + .iter() + .map(|child| count_takes(child)) + .sum::() + } + + /// Stacked row-stream takes collapse into one node fetching both takes' + /// columns, preserving the output schema and values + #[tokio::test] + async fn collapse_row_stream_takes() { + let (dataset, _tmp) = fixture(LanceFileVersion::Stable).await; + + // OneShotExec inputs are single-use: build the plan fresh per run + let build = |dataset: &Arc| { + let inner = row_stream_take(dataset, keys_input(vec![3, 1, 4]), &["s"]); + row_stream_take(dataset, inner, &["i", "s"]) + }; + let outer = build(&dataset); + let expected_schema = outer.schema(); + assert_eq!(count_takes(&outer), 2); + let expected = run(&outer).await; + + let optimized = CoalesceTake + .optimize(build(&dataset), &ConfigOptions::default()) + .unwrap(); + assert_eq!(count_takes(&optimized), 1); + assert_eq!(optimized.schema(), expected_schema); + + let result = run(&optimized).await; + assert_eq!(result, expected); + let i_col = result + .column_by_name("i") + .unwrap() + .as_primitive::(); + assert_eq!(i_col.values(), &[3, 1, 4]); + } + + /// When the outer take's projection does not cover a column the inner + /// take fetched, the collapse is skipped instead of dropping the column + #[tokio::test] + async fn collapse_skipped_when_column_would_drop() { + let (dataset, _tmp) = fixture(LanceFileVersion::Stable).await; + + // Outer target deliberately omits the inner take's "s" + let build = |dataset: &Arc| { + let inner = row_stream_take(dataset, keys_input(vec![3, 1, 4]), &["s"]); + row_stream_take(dataset, inner, &["i"]) + }; + let outer = build(&dataset); + assert_eq!(count_takes(&outer), 2); + let expected = run(&outer).await; + + let optimized = CoalesceTake + .optimize(build(&dataset), &ConfigOptions::default()) + .unwrap(); + assert_eq!(count_takes(&optimized), 2); + assert_eq!(run(&optimized).await, expected); + assert!(expected.column_by_name("s").is_some()); + } + + /// Legacy TakeExec pairs still collapse (through the CoalesceBatchesExec + /// the scanner inserts on that path) + #[tokio::test] + async fn collapse_legacy_takes() { + let (dataset, _tmp) = fixture(LanceFileVersion::Legacy).await; + + let build = |dataset: &Arc| -> Arc { + let inner_proj = dataset + .empty_projection() + .union_columns(["s"], OnMissing::Error) + .unwrap(); + let inner: Arc = Arc::new( + TakeExec::try_new(dataset.clone(), keys_input(vec![3, 1, 4]), inner_proj) + .unwrap() + .unwrap(), + ); + let outer_proj = dataset + .empty_projection() + .union_columns(["i", "s"], OnMissing::Error) + .unwrap(); + Arc::new( + TakeExec::try_new(dataset.clone(), inner, outer_proj) + .unwrap() + .unwrap(), + ) + }; + let outer = build(&dataset); + let expected_schema = outer.schema(); + assert_eq!(count_takes(&outer), 2); + let expected = run(&outer).await; + + let optimized = CoalesceTake + .optimize(build(&dataset), &ConfigOptions::default()) + .unwrap(); + assert_eq!(count_takes(&optimized), 1); + assert_eq!(optimized.schema(), expected_schema); + assert_eq!(run(&optimized).await, expected); + } +} diff --git a/rust/lance/src/io/exec/take.rs b/rust/lance/src/io/exec/take.rs index c3642cdb043..5f1dabe4abe 100644 --- a/rust/lance/src/io/exec/take.rs +++ b/rust/lance/src/io/exec/take.rs @@ -521,7 +521,7 @@ impl TakeExec { /// /// If this happens the order of the new nested fields will match the order defined in /// the dataset schema. - fn calculate_output_schema( + pub(crate) fn calculate_output_schema( dataset_schema: &Schema, input_schema: &ArrowSchema, projection: &Projection,