feat: expose make_deletions_null to python as include_deleted_rows (#3533)

Author: westonpace

python/python/lance/dataset.py

@@ -322,6 +322,7 @@ def scanner(
         io_buffer_size: Optional[int] = None,
         late_materialization: Optional[bool | List[str]] = None,
         use_scalar_index: Optional[bool] = None,
+        include_deleted_rows: Optional[bool] = None,
     ) -> LanceScanner:
         """Return a Scanner that can support various pushdowns.
 
@@ -414,6 +415,14 @@ def scanner(
         fast_search:  bool, default False
             If True, then the search will only be performed on the indexed data, which
             yields faster search time.
+        include_deleted_rows: bool, default False
+            If True, then rows that have been deleted, but are still present in the
+            fragment, will be returned.  These rows will have the _rowid column set
+            to null.  All other columns will reflect the value stored on disk and may
+            not be null.
+
+            Note: if this is a search operation, or a take operation (including scalar
+            indexed scans) then deleted rows cannot be returned.
 
         Notes
         -----
@@ -463,6 +472,7 @@ def setopt(opt, val):
         setopt(builder.use_stats, use_stats)
         setopt(builder.use_scalar_index, use_scalar_index)
         setopt(builder.fast_search, fast_search)
+        setopt(builder.include_deleted_rows, include_deleted_rows)
 
         # columns=None has a special meaning. we can't treat it as "user didn't specify"
         if self._default_scan_options is None:
@@ -543,6 +553,7 @@ def to_table(
         io_buffer_size: Optional[int] = None,
         late_materialization: Optional[bool | List[str]] = None,
         use_scalar_index: Optional[bool] = None,
+        include_deleted_rows: Optional[bool] = None,
     ) -> pa.Table:
         """Read the data into memory as a :py:class:`pyarrow.Table`
 
@@ -612,6 +623,14 @@ def to_table(
                 currently only supports a single column in the columns list.
             - query: str
                 The query string to search for.
+        include_deleted_rows: bool, optional, default False
+            If True, then rows that have been deleted, but are still present in the
+            fragment, will be returned.  These rows will have the _rowid column set
+            to null.  All other columns will reflect the value stored on disk and may
+            not be null.
+
+            Note: if this is a search operation, or a take operation (including scalar
+            indexed scans) then deleted rows cannot be returned.
 
         Notes
         -----
@@ -639,6 +658,7 @@ def to_table(
             use_stats=use_stats,
             fast_search=fast_search,
             full_text_query=full_text_query,
+            include_deleted_rows=include_deleted_rows,
         ).to_table()
 
     @property
@@ -2982,6 +3002,7 @@ def __init__(self, ds: LanceDataset):
         self._fast_search = False
         self._full_text_query = None
         self._use_scalar_index = None
+        self._include_deleted_rows = None
 
     def apply_defaults(self, default_opts: Dict[str, Any]) -> ScannerBuilder:
         for key, value in default_opts.items():
@@ -3259,6 +3280,15 @@ def fast_search(self, flag: bool) -> ScannerBuilder:
         self._fast_search = flag
         return self
 
+    def include_deleted_rows(self, flag: bool) -> ScannerBuilder:
+        """Include deleted rows
+
+        Rows which have been deleted, but are still present in the fragment, will be
+        returned.  These rows will have all columns (except _rowaddr) set to null
+        """
+        self._include_deleted_rows = flag
+        return self
+
     def full_text_search(
         self,
         query: str,
@@ -3296,6 +3326,7 @@ def to_scanner(self) -> LanceScanner:
             self._full_text_query,
             self._late_materialization,
             self._use_scalar_index,
+            self._include_deleted_rows,
         )
         return LanceScanner(scanner, self.ds)
 

python/python/lance/lance/__init__.pyi

@@ -206,6 +206,7 @@ class _Dataset:
         full_text_query: Optional[dict] = None,
         late_materialization: Optional[bool | List[str]] = None,
         use_scalar_index: Optional[bool] = None,
+        include_deleted_rows: Optional[bool] = None,
     ) -> _Scanner: ...
     def count_rows(self, filter: Optional[str] = None) -> int: ...
     def take(

python/python/tests/test_dataset.py

@@ -2261,6 +2261,44 @@ def test_scanner_schemas(tmp_path: Path):
     assert scanner.projected_schema == pa.schema([pa.field("a", pa.int64())])
 
 
+def test_scan_deleted_rows(tmp_path: Path):
+    base_dir = tmp_path / "dataset"
+    df = pd.DataFrame({"a": range(100), "b": range(100)})
+    ds = lance.write_dataset(df, base_dir, max_rows_per_file=25)
+    ds.create_scalar_index("b", "BTREE")
+    ds.delete("a < 30")
+
+    assert ds.count_rows() == 70
+
+    assert ds.scanner(with_row_id=True).to_table().num_rows == 70
+    with_deleted = ds.scanner(with_row_id=True, include_deleted_rows=True).to_table()
+
+    assert with_deleted.num_rows == 75
+
+    assert with_deleted.slice(0, 5) == pa.table(
+        {
+            "a": range(25, 30),
+            "b": range(25, 30),
+            "_rowid": pa.array([None] * 5, pa.uint64()),
+        }
+    )
+
+    assert (
+        ds.scanner(with_row_id=True, include_deleted_rows=True, filter="a < 32")
+        .to_table()
+        .num_rows
+        == 7
+    )
+
+    with pytest.raises(ValueError, match="Cannot include deleted rows"):
+        ds.scanner(
+            include_deleted_rows=True, with_row_id=True, filter="b < 30"
+        ).to_table()
+
+    with pytest.raises(ValueError, match="with_row_id is false"):
+        ds.scanner(include_deleted_rows=True, filter="a < 30").to_table()
+
+
 def test_custom_commit_lock(tmp_path: Path):
     called_lock = False
     called_release = False

python/python/tests/test_scalar_index.py

@@ -272,6 +272,11 @@ def test_full_text_search(dataset, with_position):
     for row in results:
         assert query in row.as_py()
 
+    with pytest.raises(ValueError, match="Cannot include deleted rows"):
+        dataset.to_table(
+            with_row_id=True, full_text_query=query, include_deleted_rows=True
+        )
+
 
 def test_filter_with_fts_index(dataset):
     dataset.create_scalar_index("doc", index_type="INVERTED", with_position=False)

python/python/tests/test_vector_index.py

@@ -1121,6 +1121,19 @@ def test_optimize_indices(indexed_dataset):
     assert len(indices) == 2
 
 
+def test_no_include_deleted_rows(indexed_dataset):
+    with pytest.raises(ValueError, match="Cannot include deleted rows"):
+        indexed_dataset.to_table(
+            nearest={
+                "column": "vector",
+                "q": np.random.randn(128),
+                "k": 10,
+            },
+            with_row_id=True,
+            include_deleted_rows=True,
+        )
+
+
 def test_drop_indices(indexed_dataset):
     idx_name = indexed_dataset.list_indices()[0]["name"]
 

python/src/dataset.rs

@@ -486,7 +486,7 @@ impl Dataset {
     }
 
     #[allow(clippy::too_many_arguments)]
-    #[pyo3(signature=(columns=None, columns_with_transform=None, filter=None, prefilter=None, limit=None, offset=None, nearest=None, batch_size=None, io_buffer_size=None, batch_readahead=None, fragment_readahead=None, scan_in_order=None, fragments=None, with_row_id=None, with_row_address=None, use_stats=None, substrait_filter=None, fast_search=None, full_text_query=None, late_materialization=None, use_scalar_index=None))]
+    #[pyo3(signature=(columns=None, columns_with_transform=None, filter=None, prefilter=None, limit=None, offset=None, nearest=None, batch_size=None, io_buffer_size=None, batch_readahead=None, fragment_readahead=None, scan_in_order=None, fragments=None, with_row_id=None, with_row_address=None, use_stats=None, substrait_filter=None, fast_search=None, full_text_query=None, late_materialization=None, use_scalar_index=None, include_deleted_rows=None))]
     fn scanner(
         self_: PyRef<'_, Self>,
         columns: Option<Vec<String>>,
@@ -510,6 +510,7 @@ impl Dataset {
         full_text_query: Option<&Bound<'_, PyDict>>,
         late_materialization: Option<PyObject>,
         use_scalar_index: Option<bool>,
+        include_deleted_rows: Option<bool>,
     ) -> PyResult<Scanner> {
         let mut scanner: LanceScanner = self_.ds.scan();
         match (columns, columns_with_transform) {
@@ -608,6 +609,10 @@ impl Dataset {
             scanner.fast_search();
         }
 
+        if let Some(true) = include_deleted_rows {
+            scanner.include_deleted_rows();
+        }
+
         if let Some(fragments) = fragments {
             let fragments = fragments
                 .into_iter()

rust/lance/src/dataset/scanner.rs

@@ -334,6 +334,9 @@ pub struct Scanner {
     /// This is essentially a weak consistency search. Users can run index or optimize index
     /// to make the index catch up with the latest data.
     fast_search: bool,
+
+    /// If true, the scanner will emit deleted rows
+    include_deleted_rows: bool,
 }
 
 fn escape_column_name(name: &str) -> String {
@@ -372,6 +375,7 @@ impl Scanner {
             fragments: None,
             fast_search: false,
             use_scalar_index: true,
+            include_deleted_rows: false,
         }
     }
 
@@ -551,6 +555,21 @@ impl Scanner {
         self
     }
 
+    /// Include deleted rows
+    ///
+    /// These are rows that have been deleted from the dataset but are still present in the
+    /// underlying storage.  These rows will have the `_rowid` column set to NULL.  The other columns
+    /// (include _rowaddr) will be set to their deleted values.
+    ///
+    /// This can be useful for generating aligned fragments or debugging
+    ///
+    /// Note: when entire fragments are deleted, the scanner will not emit any rows for that fragment
+    /// since the fragment is no longer present in the dataset.
+    pub fn include_deleted_rows(&mut self) -> &mut Self {
+        self.include_deleted_rows = true;
+        self
+    }
+
     /// Set the I/O buffer size
     ///
     /// This is the amount of RAM that will be reserved for holding I/O received from
@@ -1235,6 +1254,14 @@ impl Scanner {
                 location: location!(),
             });
         }
+
+        if self.include_deleted_rows && !self.with_row_id {
+            return Err(Error::InvalidInput {
+                source: "include_deleted_rows is set but with_row_id is false".into(),
+                location: location!(),
+            });
+        }
+
         if let Some(first_blob_col) = self
             .projection_plan
             .physical_schema
@@ -1318,6 +1345,13 @@ impl Scanner {
         // Stage 1: source (either an (K|A)NN search, full text search or or a (full|indexed) scan)
         let mut plan: Arc<dyn ExecutionPlan> = match (&self.nearest, &self.full_text_query) {
             (Some(_), None) => {
+                if self.include_deleted_rows {
+                    return Err(Error::InvalidInput {
+                        source: "Cannot include deleted rows in a nearest neighbor search".into(),
+                        location: location!(),
+                    });
+                }
+
                 // The source is an nearest neighbor search
                 if self.prefilter {
                     // If we are prefiltering then the knn node will take care of the filter
@@ -1332,6 +1366,13 @@ impl Scanner {
                 }
             }
             (None, Some(query)) => {
+                if self.include_deleted_rows {
+                    return Err(Error::InvalidInput {
+                        source: "Cannot include deleted rows in an FTS search".into(),
+                        location: location!(),
+                    });
+                }
+
                 // The source is an FTS search
                 if self.prefilter {
                     // If we are prefiltering then the fts node will take care of the filter
@@ -1357,6 +1398,14 @@ impl Scanner {
                 } else {
                     self.use_stats
                 };
+
+                if filter_plan.index_query.is_some() && self.include_deleted_rows {
+                    return Err(Error::InvalidInput {
+                        source: "Cannot include deleted rows in a scalar indexed scan".into(),
+                        location: location!(),
+                    });
+                }
+
                 match (
                     filter_plan.index_query.is_some(),
                     filter_plan.refine_expr.is_some(),
@@ -1406,7 +1455,7 @@ impl Scanner {
                         self.scan(
                             with_row_id,
                             self.with_row_address,
-                            false,
+                            self.include_deleted_rows,
                             scan_range,
                             eager_schema,
                         )