fix: defer blank rows when a fully deleted batch has no row to copy

python/python/tests/test_dataset.py

@@ -2162,6 +2162,68 @@ def test_merge_with_commit(tmp_path: Path):
     assert tbl == expected
 
 
+@pytest.mark.parametrize(
+    ("delete_predicate", "expected_ids"),
+    [
+        pytest.param("id < 50", list(range(50, 150)), id="leading"),
+        pytest.param(
+            "id >= 50 AND id < 100",
+            list(range(50)) + list(range(100, 150)),
+            id="middle",
+        ),
+        pytest.param("id >= 100", list(range(100)), id="trailing"),
+    ],
+)
+def test_merge_columns_with_deleted_batch_commit(
+    tmp_path: Path, delete_predicate: str, expected_ids: list
+):
+    # A fully deleted read batch must still contribute its rows to the new data
+    # file, otherwise the fragment's data files disagree on the physical row
+    # count. The deleted run is placed at the start, middle, and end because the
+    # updater can only borrow a placeholder row from a batch that has live rows.
+    base_dir = tmp_path / "test"
+    table = pa.table({"id": range(150), "value": range(150)})
+    dataset = lance.write_dataset(table, base_dir, max_rows_per_file=200)
+
+    dataset.delete(delete_predicate)
+    assert dataset.count_rows() == 100
+
+    merged_frags = []
+    schema = None
+    for frag in dataset.get_fragments():
+        live_ids = frag.scanner(columns=["id"]).to_table()["id"].to_pylist()
+        right_table = pa.table(
+            {"merged": pa.array([row_id * 10 for row_id in live_ids], pa.int64())},
+            schema=pa.schema([pa.field("merged", pa.int64(), nullable=False)]),
+        )
+        merged, schema = frag.merge_columns(right_table, batch_size=50)
+        merged_frags.append(merged)
+
+    dataset = lance.LanceDataset.commit(
+        dataset.uri,
+        lance.LanceOperation.Merge(merged_frags, schema),
+        read_version=dataset.version,
+    )
+    dataset.validate()
+
+    assert dataset.to_table() == pa.table(
+        {
+            "id": expected_ids,
+            "value": expected_ids,
+            "merged": [row_id * 10 for row_id in expected_ids],
+        },
+        schema=pa.schema(
+            [
+                pa.field("id", pa.int64()),
+                pa.field("value", pa.int64()),
+                # The blanks written for the deleted rows are copies of a live row,
+                # so the merged column stays non-nullable end to end.
+                pa.field("merged", pa.int64(), nullable=False),
+            ]
+        ),
+    )
+
+
 def test_merge_with_schema_holes(tmp_path: Path):
     # Create table with 3 cols
     table = pa.table({"a": range(10)})

rust/lance/src/dataset/fragment.rs

@@ -6171,6 +6171,82 @@ mod tests {
         }
     }
 
+    /// A deletion vector naming a row the fragment does not have leaves the restorer
+    /// with rows it can never account for, so `Updater::next` has to refuse at the end
+    /// of the stream rather than let a data file short of those rows be written.
+    ///
+    /// `write_deletions` rejects an over-long vector, so the file is written directly
+    /// to get a fragment into this state.
+    #[tokio::test]
+    async fn test_updater_rejects_deletion_vector_past_end_of_fragment() {
+        let test_dir = TempStrDir::default();
+        let test_uri = &test_dir;
+        let mut dataset = create_dataset(test_uri, LanceFileVersion::Stable).await;
+
+        // Point a fragment's deletion file at a row it does not have. 200 rows are
+        // spread over several 40-row fragments, so 10_000 is past the end of any of
+        // them. Pick a fragment whose id is not zero, so the assertion below cannot
+        // pass on a message that dropped the id entirely.
+        let deletion_vector: DeletionVector = [10_000].into_iter().collect();
+        let fragment_index = 1;
+        let fragment_id = dataset.manifest.fragments[fragment_index].id;
+        assert_ne!(fragment_id, 0, "need a non-zero fragment id");
+        let deletion_file = write_deletion_file(
+            &dataset.base,
+            fragment_id,
+            dataset.version().version,
+            &deletion_vector,
+            dataset.object_store.as_ref(),
+        )
+        .await
+        .unwrap();
+        let mut fragments = dataset.manifest.fragments.as_ref().clone();
+        fragments[fragment_index].deletion_file = deletion_file;
+        let mut manifest = dataset.manifest.as_ref().clone();
+        manifest.fragments = Arc::new(fragments);
+        dataset.manifest = Arc::new(manifest);
+
+        let new_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+            "double_i",
+            DataType::Int32,
+            true,
+        )]));
+        let fragment = dataset.get_fragment(fragment_id as usize).unwrap();
+        let mut updater = fragment
+            .updater(Some(&["i"]), None, None, None)
+            .await
+            .unwrap();
+
+        // Every live row is handed back, so the loop only ends when next() gives up.
+        let err = loop {
+            match updater.next().await {
+                Ok(Some(batch)) => {
+                    let input_col = batch.column_by_name("i").unwrap();
+                    let result_col = mul(input_col, &Int32Array::new_scalar(2)).unwrap();
+                    let batch = RecordBatch::try_new(
+                        new_schema.clone(),
+                        vec![Arc::new(result_col) as ArrayRef],
+                    )
+                    .unwrap();
+                    updater.update(batch).await.unwrap();
+                }
+                Ok(None) => panic!("expected next() to refuse the unaccounted-for row"),
+                Err(err) => break err,
+            }
+        };
+
+        assert!(matches!(err, Error::NotSupported { .. }), "{err:?}");
+        let message = err.to_string();
+        assert!(
+            message.contains("unaccounted for"),
+            "expected the stream-ended wording, got: {message}"
+        );
+        assert!(
+            message.contains(&format!("fragment {fragment_id}")),
+            "message should name the fragment: {message}"
+        );
+    }
+
     #[rstest]
     #[tokio::test]
     async fn test_merge_fragment(

rust/lance/src/dataset/schema_evolution.rs

@@ -1298,24 +1298,34 @@ mod test {
         Ok(())
     }
 
+    /// Regression test: when an entire read batch has been deleted, the updater
+    /// yields a 0-row batch and the deleted rows must still be restored, because
+    /// every data file in a fragment has to keep the same physical row count.
+    ///
+    /// A single fragment holds 150 rows and 50 consecutive rows are deleted. Read
+    /// with batch_size=50 the deleted run lines up exactly with one read batch,
+    /// which therefore arrives empty. The run is placed at the start, in the
+    /// middle, and at the end because the restorer treats those positions
+    /// differently: a deleted run that trails a live batch is greedily appended to
+    /// it, while a run starting at row 0 has no preceding batch to absorb it.
+    #[rstest]
+    #[case::leading("i < 50", (50..150).collect::<Vec<i32>>())]
+    #[case::middle("i >= 50 AND i < 100", (0..50).chain(100..150).collect::<Vec<i32>>())]
+    #[case::trailing("i >= 100", (0..100).collect::<Vec<i32>>())]
     #[tokio::test]
-    async fn test_add_columns_with_fully_deleted_batch() -> Result<()> {
-        // Regression test: when an entire read batch has been deleted, the
-        // updater yields a 0-row batch. The inner loop then never runs and
-        // `batches` stays empty, so `concat_batches(&batches[0]..)` used to
-        // panic with "index out of bounds: the len is 0 but the index is 0".
-        //
-        // A single fragment holds 105 rows; deleting the trailing 5 rows means
-        // that, when read with batch_size=50, the third batch [100..105) is
-        // fully filtered out and produces an empty batch.
+    async fn test_add_columns_with_fully_deleted_batch(
+        #[case] delete_predicate: &str,
+        #[case] expected_live_ids: Vec<i32>,
+        #[values(true, false)] new_column_nullable: bool,
+    ) -> Result<()> {
         let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
             "i",
             DataType::Int32,
             false,
         )]));
         let batch = RecordBatch::try_new(
             schema.clone(),
-            vec![Arc::new(Int32Array::from_iter_values(0..105))],
+            vec![Arc::new(Int32Array::from_iter_values(0..150))],
         )?;
         let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
 
@@ -1331,28 +1341,32 @@ mod test {
         )
         .await?;
 
-        // Delete the entire trailing batch [100..105).
-        dataset.delete("i >= 100").await?;
+        dataset.delete(delete_predicate).await?;
         assert_eq!(dataset.count_rows(None).await?, 100);
 
         let new_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
             "j",
             DataType::Int32,
-            false,
+            new_column_nullable,
         )]));
         let new_batch = RecordBatch::try_new(
             new_schema.clone(),
             vec![Arc::new(Int32Array::from_iter_values(0..100))],
         )?;
         let reader = RecordBatchIterator::new(vec![Ok(new_batch)], new_schema.clone());
 
-        // Read with batch_size=50 so the deleted trailing rows form a full empty batch.
+        // Read with batch_size=50 so the deleted rows form a full empty batch.
         dataset
             .add_columns(NewColumnTransform::Reader(Box::new(reader)), None, Some(50))
             .await?;
+        dataset.validate().await?;
 
         let data = dataset.scan().try_into_batch().await?;
         assert_eq!(data.num_rows(), 100);
+        assert_eq!(
+            data.column_by_name("i").unwrap().as_ref(),
+            &Int32Array::from(expected_live_ids)
+        );
         assert_eq!(
             data.column_by_name("j").unwrap().as_ref(),
             &Int32Array::from_iter_values(0..100)
@@ -1361,6 +1375,73 @@ mod test {
         Ok(())
     }
 
+    /// A legacy fragment whose trailing row group is entirely deleted cannot defer its
+    /// blanks: that batch reaches `add_blanks` with no live row to copy, so the update
+    /// is refused rather than writing a data file short of the deleted rows. Deferring
+    /// is what a v2 fragment does instead, which
+    /// `test_add_columns_with_fully_deleted_batch`'s trailing case covers.
+    #[tokio::test]
+    async fn test_add_columns_legacy_trailing_deleted_batch_errors() -> Result<()> {
+        let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+            "i",
+            DataType::Int32,
+            false,
+        )]));
+        let batch = RecordBatch::try_new(
+            schema.clone(),
+            vec![Arc::new(Int32Array::from_iter_values(0..105))],
+        )?;
+        let reader = RecordBatchIterator::new(vec![Ok(batch)], schema.clone());
+
+        let test_dir = TempStrDir::default();
+        let test_uri = &test_dir;
+        let mut dataset = Dataset::write(
+            reader,
+            test_uri,
+            Some(WriteParams {
+                max_rows_per_file: 200,
+                max_rows_per_group: 50,
+                data_storage_version: Some(LanceFileVersion::Legacy),
+                ..Default::default()
+            }),
+        )
+        .await?;
+
+        // The last row group is [100, 105); deleting all of it leaves a trailing read
+        // batch with no live rows, which legacy files cannot defer past.
+        dataset.delete("i >= 100").await?;
+
+        let new_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new(
+            "j",
+            DataType::Int32,
+            true,
+        )]));
+        let new_batch = RecordBatch::try_new(
+            new_schema.clone(),
+            vec![Arc::new(Int32Array::from_iter_values(0..100))],
+        )?;
+        let reader = RecordBatchIterator::new(vec![Ok(new_batch)], new_schema.clone());
+
+        let err = dataset
+            .add_columns(NewColumnTransform::Reader(Box::new(reader)), None, None)
+            .await
+            .unwrap_err();
+
+        assert!(
+            matches!(err, Error::NotSupported { .. }),
+            "expected NotSupported, got {err:?}"
+        );
+        // Match add_blanks' own wording, not the shared "run compaction" tail: the
+        // stream-ended error in Updater::next carries that tail too, and this case
+        // fails before the stream ever runs out.
+        assert!(
+            err.to_string().contains("missing too many rows in merge"),
+            "expected the add_blanks rejection, got: {err}"
+        );
+
+        Ok(())
+    }
+
     #[rstest]
     #[tokio::test]
     async fn test_add_columns_cleans_up_blob_v2_data_on_stream_error(