diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index ae49b8cab3d..3bdec48ae9b 100644 --- a/python/python/tests/test_dataset.py +++ b/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)}) diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 87c63a1549d..070591a9ed0 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/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( diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 3643e8043ec..fd933fed3e3 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -1298,16 +1298,26 @@ 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::>())] + #[case::middle("i >= 50 AND i < 100", (0..50).chain(100..150).collect::>())] + #[case::trailing("i >= 100", (0..100).collect::>())] #[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, + #[values(true, false)] new_column_nullable: bool, + ) -> Result<()> { let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( "i", DataType::Int32, @@ -1315,7 +1325,7 @@ mod test { )])); 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,14 +1341,13 @@ 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(), @@ -1346,13 +1355,18 @@ mod test { )?; 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( diff --git a/rust/lance/src/dataset/updater.rs b/rust/lance/src/dataset/updater.rs index c5cd8db9600..a504a07b9e7 100644 --- a/rust/lance/src/dataset/updater.rs +++ b/rust/lance/src/dataset/updater.rs @@ -119,6 +119,10 @@ impl Updater { } /// Returns the next [`RecordBatch`] as input for updater. + /// + /// Every batch this hands out must be passed back to [`Self::update`] before the + /// next call: the deletion restorer advances there, so skipping it would leave + /// deleted rows unaccounted for and fail the stream at its end. pub async fn next(&mut self) -> Result> { if self.finished { return Ok(None); @@ -127,9 +131,26 @@ impl Updater { match batch { None => { if !self.deletion_restorer.is_exhausted() { - // This can happen only if there is a batch size (e.g. v1 file) and the - // last batch(es) are entirely deleted. - return Err(Error::not_supported_source("Missing too many rows in merge, run compaction to materialize deletions first".into())); + // The stream cannot supply rows the restorer still needs. In + // practice that means the deletion vector points at rows the + // stream never produced — an id past the fragment's physical row + // count, or fewer rows read than the fragment claims to have. + // + // Deferred blanks can also be outstanding here, but only if no + // batch after the deferral had a live row, i.e. the whole fragment + // is deleted; `write_deletions` drops such a fragment before it + // reaches an updater, so that path is defensive. A legacy file + // cannot defer at all — its fully deleted batch is refused + // earlier, by `add_blanks`. + // + // Don't name a count: the deletion-vector case owes no blanks yet, + // so a number here would read as zero rows owed. + return Err(Error::not_supported(format!( + "Fragment Updater: the input stream for fragment {} ended while \ + deleted rows were still unaccounted for, run compaction to \ + materialize deletions first", + self.fragment.id(), + ))); } self.finished = true; Ok(None) @@ -296,6 +317,10 @@ impl Updater { /// /// To do this we scan through the deletion vector in sorted order, merging deleted rows /// in as appropriate. +/// +/// Any method returning an error leaves the restorer mid-batch: the deletion vector +/// has been walked past rows that never made it into an output batch. Drop it and +/// start over rather than calling it again. struct DeletionRestorer { current_row_id: u32, @@ -305,6 +330,12 @@ struct DeletionRestorer { deletion_vector_iter: Option + Send>>, last_deleted_row_id: Option, + + /// Blank rows owed to batches that had no live row to copy a placeholder from + /// + /// See [`Self::restore`] for why they are deferred instead of materialized. + /// Only ever non-zero for non-legacy files, which are the only ones that defer. + pending_blank_rows: u32, } impl DeletionRestorer { @@ -314,11 +345,12 @@ impl DeletionRestorer { legacy_batch_size, deletion_vector_iter: Some(deletion_vector.into_sorted_iter()), last_deleted_row_id: None, + pending_blank_rows: 0, } } fn is_exhausted(&self) -> bool { - self.deletion_vector_iter.is_none() + self.deletion_vector_iter.is_none() && self.pending_blank_rows == 0 } fn is_full(batch_size: Option, num_rows: u32) -> bool { @@ -361,11 +393,14 @@ impl DeletionRestorer { let deletion_vector_iter = self.deletion_vector_iter.as_mut().unwrap(); // Now we need to walk through our deletion vector and figure out where to insert blanks - let mut next_deleted_id = if self.last_deleted_row_id.is_some() { - self.last_deleted_row_id - } else { - deletion_vector_iter.next() - }; + // Take the stashed id rather than peeking at it: leaving a consumed id in the + // field relies on the early return above to never read it again. `or_else` has + // to stay lazy — `or` would pull from the iterator even when a stash is waiting, + // silently dropping a deleted row. + let mut next_deleted_id = self + .last_deleted_row_id + .take() + .or_else(|| deletion_vector_iter.next()); loop { if let Some(next_deleted_id) = next_deleted_id { if next_deleted_id > last_row_id @@ -385,17 +420,65 @@ impl DeletionRestorer { } else { // Deleted row ids iterator is exhausted self.deletion_vector_iter = None; + // `is_exhausted` reads these two together, so a stash left behind here + // would make it report exhaustion while a deleted row is still owed. + debug_assert!(self.last_deleted_row_id.is_none()); return deleted; } next_deleted_id = deletion_vector_iter.next(); } } + /// Restore the deleted rows for one batch of live rows. + /// + /// Blanks are materialized by copying the batch's first live row (see + /// [`add_blanks`]), so a batch with no live rows has nothing to copy from. That + /// happens when a deleted run starts at physical row 0: there is no preceding + /// batch for [`Self::deleted_batch_offsets_in_range`] to append the run to, so + /// the run arrives as an empty batch carrying every one of its offsets. + /// + /// Rather than invent placeholder values for an arbitrary schema, we remember + /// how many blanks we owe and prepend them to the next batch that does have a + /// live row. Deleted rows sort before the live rows that follow them, so the + /// physical row order is preserved either way. fn restore(&mut self, batch: RecordBatch) -> Result { + // Holds by construction today — deferring is the only thing that sets + // pending_blank_rows and it is gated on non-legacy — so this documents the + // invariant the legacy row-count check below depends on rather than guarding + // against a state we can reach. + debug_assert!(self.pending_blank_rows == 0 || self.legacy_batch_size.is_none()); + // Because of deleted rows, the number of row ids in the batch might not // match the length. let deleted_batch_offsets = self.deleted_batch_offsets_in_range(batch.num_rows() as u32); - let batch = add_blanks(batch, &deleted_batch_offsets)?; + + // Legacy files must reproduce the original row group size, which deferring + // would break, so they keep reporting the pre-existing error instead. + if batch.num_rows() == 0 && self.legacy_batch_size.is_none() { + let deferred = deleted_batch_offsets.len() as u32; + self.pending_blank_rows += deferred; + self.current_row_id += deferred; + return Ok(batch); + } + + let pending_blank_rows = self.pending_blank_rows; + let batch_offsets = if pending_blank_rows == 0 { + deleted_batch_offsets + } else { + // The deferred blanks take the front of the batch, pushing the offsets + // computed for this batch back by that many rows. + let mut batch_offsets = + Vec::with_capacity(pending_blank_rows as usize + deleted_batch_offsets.len()); + batch_offsets.extend(0..pending_blank_rows); + batch_offsets.extend( + deleted_batch_offsets + .iter() + .map(|offset| offset + pending_blank_rows), + ); + batch_offsets + }; + + let batch = add_blanks(batch, &batch_offsets)?; if let Some(batch_size) = self.legacy_batch_size { // validation just in case, when the input has a fixed batch size then the @@ -410,12 +493,23 @@ impl DeletionRestorer { } } - self.current_row_id += batch.num_rows() as u32; + // The deferred blanks were counted when they were deferred. + self.current_row_id += batch.num_rows() as u32 - pending_blank_rows; + self.pending_blank_rows = 0; Ok(batch) } } /// Add blank rows where there are deleted rows +/// +/// `batch_offsets` must be strictly increasing, and no offset may require more +/// live rows before it than the batch has left: an offset is the position a blank +/// takes in the output, so either kind of violation asks for an impossible number +/// of live rows in between. +/// +/// Blanks copy the batch's first row, so the batch must have at least one row. +/// [`DeletionRestorer::restore`] defers blanks past an empty batch to keep that +/// true; only legacy files, which cannot defer, can still reach the error below. pub(crate) fn add_blanks(batch: RecordBatch, batch_offsets: &[u32]) -> Result { // Fast early return if batch_offsets.is_empty() { @@ -423,18 +517,38 @@ pub(crate) fn add_blanks(batch: RecordBatch, batch_offsets: &[u32]) -> Result::with_capacity(batch.num_rows() + batch_offsets.len()); let mut batch_pos = 0; let mut next_id = 0; - for batch_offset in batch_offsets { - let num_rows = *batch_offset - next_id; + for (idx, batch_offset) in batch_offsets.iter().enumerate() { + // A non-increasing offset panics in debug and wraps in release; reject it + // up front so the error names the real problem. + let num_rows = batch_offset.checked_sub(next_id).ok_or_else(|| { + Error::internal(format!( + "Fragment Updater: blank offsets must be strictly increasing, but offset \ + {batch_offset} (entry {idx} of {}) is below the expected minimum {next_id}", + batch_offsets.len() + )) + })?; + // An offset needing more live rows than remain would index past the batch: + // `take` runs unchecked below, so catch it here rather than letting it + // panic inside arrow or, worse, read the wrong rows. + if num_rows > num_live_rows - batch_pos { + return Err(Error::internal(format!( + "Fragment Updater: blank offset {batch_offset} (entry {idx} of \ + {}) needs {num_rows} more live rows before it, but {} of the batch's \ + {num_live_rows} are still unused", + batch_offsets.len(), + num_live_rows - batch_pos + ))); + } selection_vector.extend(batch_pos..batch_pos + num_rows); // For simplicity, we just use the first value for deleted rows // TODO: optimize this to use small value for each column. @@ -442,7 +556,7 @@ pub(crate) fn add_blanks(batch: RecordBatch, batch_offsets: &[u32]) -> Result Result()) + .into_batch_rows(RowCount::from(0)) + .unwrap(); + + // Assert the source, not just is_err: the batch-size check further down + // returns Internal, and the two are different failures. + let err = restorer.restore(empty).unwrap_err(); + assert!(matches!(err, Error::NotSupported { .. }), "{err:?}"); + } + + /// The v2 side of the same deletion vector: blanks owed by a batch with no live + /// row are deferred to a later batch that has one to copy. + #[test] + fn test_restore_deletes_leading_empty_batch() { + let mut restorer = super::DeletionRestorer::new((0..10).chain([15]).collect(), None); + + let empty = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(0)) + .unwrap(); + + // Nothing is written for the fully deleted batch itself. + assert_eq!(restorer.restore(empty.clone()).unwrap().num_rows(), 0); + assert!(!restorer.is_exhausted()); + + // A second empty batch must carry the debt through untouched: row 15 is + // out of its range, so it defers nothing of its own. + let restored = restorer.restore(empty).unwrap(); + assert_eq!(restored.num_rows(), 0); + assert!(!restorer.is_exhausted()); + + // The next batch covers row ids 10..15, so it owes the 10 deferred blanks + // in front of its own rows and one more for row 15 at the end. That last + // one is what pins the offset shift: without it the offsets would not be + // increasing. + let batch = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(5)) + .unwrap(); + let restored = restorer.restore(batch).unwrap(); + + assert_eq!(restored.num_rows(), 16); + let values = restored.column(0).as_primitive::(); + // Blanks copy the batch's first live row rather than inventing a value, + // which is what lets a non-nullable column through. + for i in 0..10 { + assert_eq!(values.value(i), 0); + } + for i in 0..5 { + assert_eq!(values.value(10 + i), i as i32); + } + assert_eq!(values.value(15), 0); + assert!(restorer.is_exhausted()); + } + + /// The debt itself has to keep the restorer from reporting exhaustion, not just + /// the deletion vector. With no row past the deleted run the iterator empties on + /// the first call, so only `pending_blank_rows` can hold `is_exhausted` back — + /// and it must, or `Updater::next` would accept a data file short by ten rows. + #[test] + fn test_restore_deletes_owes_blanks_after_vector_drains() { + let mut restorer = super::DeletionRestorer::new((0..10).collect(), None); + + let empty = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(0)) + .unwrap(); + + assert_eq!(restorer.restore(empty).unwrap().num_rows(), 0); + assert!(!restorer.is_exhausted()); + + let batch = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(5)) + .unwrap(); + assert_eq!(restorer.restore(batch).unwrap().num_rows(), 15); + assert!(restorer.is_exhausted()); + } + + /// A deletion vector naming a row the fragment does not have leaves the restorer + /// unexhausted with no blanks owed: the id stays stashed, so the iterator is never + /// drained. `Updater::next` relies on this to refuse rather than write a data file + /// missing that row, and the error must not claim a blank count for it. + #[test] + fn test_restore_deletes_not_exhausted_when_deletion_vector_overruns() { + let mut restorer = super::DeletionRestorer::new([100].into_iter().collect(), None); + + let batch = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(5)) + .unwrap(); + + // Row 100 is past this batch, so it is stashed rather than consumed and + // nothing is restored. No blanks are owed either — which is why the error in + // `Updater::next` cannot name a count. + assert_eq!(restorer.restore(batch).unwrap().num_rows(), 5); + assert!(!restorer.is_exhausted()); + } + + /// Deferred blanks are counted into `current_row_id` when they are deferred, so + /// consuming them must not count them again. A later deleted row is what makes + /// the double count observable: it lands at the wrong offset once the restorer + /// thinks the fragment is further along than it is. + #[test] + fn test_restore_deletes_does_not_double_count_deferred_blanks() { + let mut restorer = super::DeletionRestorer::new((0..10).chain([22]).collect(), None); + + let empty = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(0)) + .unwrap(); + assert_eq!(restorer.restore(empty).unwrap().num_rows(), 0); + + // Rows 10..20 are live, so this batch pays off the ten blanks and nothing + // else: row 22 is past its range and stays stashed. + let batch = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(10)) + .unwrap(); + assert_eq!(restorer.restore(batch).unwrap().num_rows(), 20); + assert!(!restorer.is_exhausted()); + + // Row 22 falls inside this batch's range, but only if current_row_id sits at + // 20. Counting the deferred blanks twice would have pushed it to 30, putting + // row 22 behind the batch and dropping its blank. + let batch = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(5)) + .unwrap(); + let restored = restorer.restore(batch).unwrap(); + + // Physical rows 20..25 arrive with row 22 deleted, so the blank lands third. + assert_eq!(restored.num_rows(), 6); + let values = restored.column(0).as_primitive::(); + assert_eq!(values.value(0), 0); + assert_eq!(values.value(1), 1); + assert_eq!(values.value(2), 0); + for i in 2..5 { + assert_eq!(values.value(1 + i), i as i32); + } + assert!(restorer.is_exhausted()); + } + #[test] fn test_add_blanks() { let batch = lance_datagen::gen_batch() @@ -538,4 +806,55 @@ mod tests { } assert_eq!(values.value(11), 0); } + + /// The ways a caller can hand `add_blanks` offsets it cannot satisfy. The + /// message keyword matters as much as the variant: most of these return + /// `Internal`, so matching only the variant would let one check stand in for + /// the other. + #[rstest] + #[case::empty_batch(0, &[0, 1, 2], "missing too many rows in merge")] + #[case::non_increasing(5, &[3, 1], "strictly increasing")] + #[case::equal_offsets(5, &[1, 1], "strictly increasing")] + #[case::past_end(5, &[100], "more live rows before it, but")] + // Rejected at the second offset with only three live rows left, so this is the + // only case that exercises the `- batch_pos` term: without it the remaining + // count reads as five and this offset slips through. + #[case::past_end_after_live_rows(5, &[2, 7], "more live rows before it, but")] + fn test_add_blanks_rejects_invalid_offsets( + #[case] num_rows: u64, + #[case] batch_offsets: &[u32], + #[case] expected_message: &str, + ) { + let batch = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(num_rows)) + .unwrap(); + + let err = add_blanks(batch, batch_offsets).unwrap_err(); + let message = err.to_string(); + assert!( + message.contains(expected_message), + "expected {expected_message:?} in {message:?}" + ); + } + + /// An offset equal to the batch length is the trailing-deletion shape: every + /// live row comes first, then the blanks. It has to be accepted, which is what + /// pins the bounds check to `>` rather than `>=`. + #[test] + fn test_add_blanks_at_end_of_batch() { + let batch = lance_datagen::gen_batch() + .col("x", lance_datagen::array::step::()) + .into_batch_rows(RowCount::from(5)) + .unwrap(); + + let with_blanks = add_blanks(batch, &[5]).unwrap(); + + assert_eq!(with_blanks.num_rows(), 6); + let values = with_blanks.column(0).as_primitive::(); + for i in 0..5 { + assert_eq!(values.value(i), i as i32); + } + assert_eq!(values.value(5), 0); + } }