diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index def635d58a6..348689554ee 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -6283,10 +6283,13 @@ def batch_readahead(self, nbatches: Optional[int] = None) -> ScannerBuilder: """ Set the maximum number of batches to decode concurrently. - This parameter must be greater than zero. + This parameter must be a positive integer. """ - if nbatches is not None and int(nbatches) <= 0: - raise ValueError("batch_readahead must be greater than 0") + if nbatches is not None: + if not isinstance(nbatches, int) or isinstance(nbatches, bool): + raise ValueError("batch_readahead must be an integer") + if nbatches <= 0: + raise ValueError("batch_readahead must be greater than 0") self._batch_readahead = nbatches return self @@ -7062,6 +7065,28 @@ def optimize_indices(self, **kwargs): """ self._dataset._ds.optimize_indices(**kwargs) + def cleanup_frag_reuse_index(self) -> None: + """Prune obsolete generations from the ``__lance_frag_reuse`` system index. + + The fragment-reuse index is created by :meth:`compact_files` when + ``defer_index_remap=True``. It retains one generation per deferred + compaction so indices can be remapped lazily. Once every index has caught + up to a generation (for example after rebuilding it or calling + :meth:`optimize_indices`), that generation is obsolete and can be pruned. + + This is safe to call at any time: a generation that an index has not yet + caught up to is always retained. + + Examples + -------- + >>> import lance + >>> import pyarrow as pa + >>> data = pa.table({"id": [1, 2, 3]}) + >>> dataset = lance.write_dataset(data, "memory://frag_reuse_example") + >>> dataset.optimize.cleanup_frag_reuse_index() + """ + self._dataset._ds.cleanup_frag_reuse_index() + def enable_auto_cleanup(self, auto_cleanup_config: AutoCleanupConfig, **kwargs): """Enable autocleaning for an existing dataset. diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index 03ec9f46db7..130e6342590 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -555,6 +555,7 @@ class _Dataset: ) -> Tuple[_Dataset, Transaction]: ... def validate(self): ... def migrate_manifest_paths_v2(self): ... + def cleanup_frag_reuse_index(self) -> None: ... def drop_columns(self, columns: List[str]): ... def add_columns_from_reader( self, reader: pa.RecordBatchReader, batch_size: Optional[int] = None diff --git a/python/python/tests/test_dataset.py b/python/python/tests/test_dataset.py index 898bdf93069..ddd0783a49a 100644 --- a/python/python/tests/test_dataset.py +++ b/python/python/tests/test_dataset.py @@ -1166,6 +1166,16 @@ def test_to_batches(tmp_path: Path): assert sorted_batches == table +@pytest.mark.parametrize("batch_readahead", [1.5, "1", True]) +def test_batch_readahead_requires_integer(tmp_path: Path, batch_readahead): + table = pa.Table.from_pydict({"a": range(10)}) + base_dir = tmp_path / "test" + lance.write_dataset(table, base_dir) + + with pytest.raises(ValueError, match="batch_readahead must be an integer"): + lance.dataset(base_dir).scanner(batch_readahead=batch_readahead) + + def test_list_from_parquet(tmp_path: Path): # This is a regression for GH-1482, the parquet reader creates # list fields with the name 'element' instead of 'item'. We should diff --git a/python/python/tests/test_optimize.py b/python/python/tests/test_optimize.py index 5d88af2566c..7006f7c1539 100644 --- a/python/python/tests/test_optimize.py +++ b/python/python/tests/test_optimize.py @@ -324,6 +324,62 @@ def test_defer_index_remap(tmp_path: Path): assert any(idx.name == "__lance_frag_reuse" for idx in indices) +def test_cleanup_frag_reuse_index(tmp_path: Path): + """cleanup_frag_reuse_index prunes only the reuse generations that every user + index has caught up to, and never the ones still in use. + + Setup: 6 fragments, one BTREE scalar index. Compact with + defer_index_remap=True so the frag-reuse index is populated. While the + scalar index has not caught up, cleanup must retain the generation. After + rebuilding the scalar index (create_scalar_index with replace=True) it is + newer than the reuse generation, so cleanup may prune it and num_versions + drops to 0. + """ + base_dir = tmp_path / "dataset" + data = pa.table({"i": range(6_000), "val": range(6_000)}) + dataset = lance.write_dataset(data, base_dir, max_rows_per_file=1_000) + dataset.create_scalar_index("i", "BTREE") + + dataset.delete("i < 500") + dataset.optimize.compact_files( + target_rows_per_fragment=2_000, defer_index_remap=True, num_threads=1 + ) + + dataset = lance.dataset(base_dir) + assert any( + idx.name == "__lance_frag_reuse" for idx in dataset.describe_indices() + ), "precondition: defer_index_remap must have created the frag-reuse index" + + before_stats = dataset.stats.index_stats("__lance_frag_reuse") + versions_before = before_stats["num_versions"] + assert versions_before >= 1, ( + "precondition: frag-reuse index must have at least one version before cleanup" + ) + + # Negative case: index not yet remapped, so cleanup must retain the generation. + dataset.optimize.cleanup_frag_reuse_index() + dataset = lance.dataset(base_dir) + retained_stats = dataset.stats.index_stats("__lance_frag_reuse") + assert retained_stats["num_versions"] == versions_before, ( + f"cleanup_frag_reuse_index must retain generations that an index has not " + f"caught up to, but num_versions went from {versions_before} to " + f"{retained_stats['num_versions']}" + ) + + # Positive case: rebuilding catches the index up, so cleanup can now prune all. + dataset.create_scalar_index("i", "BTREE", replace=True) + dataset = lance.dataset(base_dir) + + dataset.optimize.cleanup_frag_reuse_index() + + dataset = lance.dataset(base_dir) + after_stats = dataset.stats.index_stats("__lance_frag_reuse") + assert after_stats["num_versions"] == 0, ( + f"cleanup_frag_reuse_index should have pruned all reuse generations " + f"but num_versions={after_stats['num_versions']}" + ) + + @pytest.mark.parametrize("use_commit_options", [True, False]) def test_defer_index_remap_via_commit_options(tmp_path: Path, use_commit_options: bool): """Compaction.commit respects defer_index_remap passed in options. diff --git a/python/src/dataset.rs b/python/src/dataset.rs index a361e0b63a8..c6db0deba4c 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2995,6 +2995,17 @@ impl Dataset { Ok(()) } + fn cleanup_frag_reuse_index(&mut self) -> PyResult<()> { + let mut new_self = self.ds.as_ref().clone(); + rt().block_on( + None, + lance::dataset::index::frag_reuse::cleanup_frag_reuse_index(&mut new_self), + )? + .infer_error()?; + self.ds = Arc::new(new_self); + Ok(()) + } + fn drop_columns(&mut self, columns: Vec) -> PyResult<()> { let mut new_self = self.ds.as_ref().clone(); let columns: Vec<_> = columns.iter().map(|s| s.as_str()).collect(); diff --git a/rust/lance/src/dataset/index/frag_reuse.rs b/rust/lance/src/dataset/index/frag_reuse.rs index ceebe456bbf..975d11f054c 100644 --- a/rust/lance/src/dataset/index/frag_reuse.rs +++ b/rust/lance/src/dataset/index/frag_reuse.rs @@ -38,27 +38,13 @@ pub async fn cleanup_frag_reuse_index(dataset: &mut Dataset) -> lance_core::Resu return Ok(()); }; - let frag_reuse_details = load_frag_reuse_index_details(dataset, frag_reuse_index_meta) - .await - .unwrap(); + // Surface corrupt/stale metadata as a normal error instead of panicking on unwrap. + let frag_reuse_details = load_frag_reuse_index_details(dataset, frag_reuse_index_meta).await?; - let mut retained_versions = Vec::new(); + let mut retained_versions = Vec::with_capacity(frag_reuse_details.versions.len()); let mut fragment_bitmaps = RoaringBitmap::new(); for version in frag_reuse_details.versions.iter() { - let check_results = indices - .iter() - .map(|idx| is_index_remap_caught_up(version, idx)) - .collect::>(); - - if check_results - .iter() - .any(|r| matches!(r, Err(Error::InvalidInput { .. }))) - { - // If the check fails, the reuse version is likely corrupted, do not retain it. - continue; - } - - if !check_results.into_iter().all(|r| r.unwrap()) { + if !is_reuse_version_caught_up(version, &indices)? { fragment_bitmaps.extend(version.new_frag_bitmap()); retained_versions.push(version.clone()); } @@ -97,6 +83,32 @@ pub async fn cleanup_frag_reuse_index(dataset: &mut Dataset) -> lance_core::Resu Ok(()) } +/// Decide whether a reuse version can be safely cleaned up given the current indices. +/// +/// A reuse version is only droppable once every index has caught up to it. This +/// mirrors the load/remap path. `recalculate_fragment_bitmap` treats a rewrite group +/// that straddles indexed and non-indexed fragments as a recoverable state that aborts +/// the remap and preserves the reuse generation, rather than as corruption to discard. +/// We treat the same `InvalidInput` straddling state as "not caught up" so the reuse +/// generation is kept. Dropping it would leave the affected indexes unable to remap +/// their stale fragments. Any other error is surfaced to the caller. +fn is_reuse_version_caught_up( + frag_reuse_version: &FragReuseVersion, + indices: &[IndexMetadata], +) -> lance_core::Result { + for idx in indices.iter() { + match is_index_remap_caught_up(frag_reuse_version, idx) { + Ok(true) => {} + // The index is behind or straddles indexed and non-indexed data, so the + // reuse version is still needed and must be kept, as the load path does. + Ok(false) | Err(Error::InvalidInput { .. }) => return Ok(false), + // Any other error is unexpected, so surface it instead of panicking. + Err(e) => return Err(e), + } + } + Ok(true) +} + fn is_index_remap_caught_up( frag_reuse_version: &FragReuseVersion, index_meta: &IndexMetadata, @@ -155,6 +167,7 @@ mod tests { use arrow_array::types::{Float32Type, Int32Type}; use lance_datagen::Dimension; use lance_index::IndexType; + use lance_index::frag_reuse::{FragDigest, FragReuseGroup}; use lance_index::scalar::ScalarIndexParams; #[tokio::test] @@ -212,6 +225,22 @@ mod tests { is_index_remap_caught_up(&frag_reuse_details.versions[0], scalar_index).unwrap() ); + // Cleanup must not prune a reuse version while an index has not caught up to it. + cleanup_frag_reuse_index(&mut dataset).await.unwrap(); + let frag_reuse_index_meta = dataset + .load_index_by_name(FRAG_REUSE_INDEX_NAME) + .await + .unwrap() + .expect("Fragment reuse index must be available"); + let frag_reuse_details = load_frag_reuse_index_details(&dataset, &frag_reuse_index_meta) + .await + .unwrap(); + assert_eq!( + frag_reuse_details.versions.len(), + 1, + "reuse version must be retained while an index is not caught up" + ); + // Remap and check index is caught up remapping::remap_column_index(&mut dataset, &["i"], index_name.clone()) .await @@ -244,6 +273,165 @@ mod tests { )); } + /// Corrupt or stale fragment-reuse metadata must surface as a normal error, + /// not a panic. Regression for the `load_frag_reuse_index_details().unwrap()` + /// that used to crash the (Python) caller on unloadable metadata. + #[tokio::test] + async fn test_cleanup_frag_reuse_index_corrupt_metadata_errors() { + let mut dataset = lance_datagen::gen_batch() + .col("i", lance_datagen::array::step::()) + .into_ram_dataset(FragmentCount::from(6), FragmentRowCount::from(1000)) + .await + .unwrap(); + + dataset + .create_index( + &["i"], + IndexType::Scalar, + Some("scalar".into()), + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + + compact_files( + &mut dataset, + CompactionOptions { + target_rows_per_fragment: 2_000, + defer_index_remap: true, + ..Default::default() + }, + None, + ) + .await + .unwrap(); + + let frag_reuse_index_meta = dataset + .load_index_by_name(FRAG_REUSE_INDEX_NAME) + .await + .unwrap() + .expect("Fragment reuse index must be available"); + + // Replace the reuse index with one whose details cannot be loaded. + let corrupt_meta = IndexMetadata { + uuid: uuid::Uuid::new_v4(), + index_details: None, + ..frag_reuse_index_meta.clone() + }; + let transaction = Transaction::new( + dataset.manifest.version, + Operation::CreateIndex { + new_indices: vec![corrupt_meta], + removed_indices: vec![frag_reuse_index_meta], + }, + None, + ); + dataset + .apply_commit(transaction, &Default::default(), &Default::default()) + .await + .unwrap(); + + let err = cleanup_frag_reuse_index(&mut dataset) + .await + .expect_err("corrupt frag-reuse metadata must surface as an error, not a panic"); + assert!( + matches!(err, Error::Index { .. }), + "expected an Error::Index for corrupt frag-reuse metadata, got: {err}" + ); + assert!( + err.to_string().contains("fragment reuse index"), + "error message should identify the fragment reuse index, got: {err}" + ); + } + + fn frag_digest(id: u64) -> FragDigest { + FragDigest { + id, + physical_rows: 1000, + num_deleted_rows: 0, + } + } + + /// Build a reuse version with a single rewrite group rewriting `old_ids` into + /// `new_ids` at the given dataset version. + fn reuse_version(dataset_version: u64, old_ids: &[u64], new_ids: &[u64]) -> FragReuseVersion { + FragReuseVersion { + dataset_version, + groups: vec![FragReuseGroup { + changed_row_addrs: Vec::new(), + old_frags: old_ids.iter().copied().map(frag_digest).collect(), + new_frags: new_ids.iter().copied().map(frag_digest).collect(), + }], + } + } + + fn index_meta_with_bitmap(dataset_version: u64, frag_ids: &[u32]) -> IndexMetadata { + IndexMetadata { + uuid: uuid::Uuid::new_v4(), + fields: vec![0], + name: "straddle_idx".into(), + dataset_version, + fragment_bitmap: Some(RoaringBitmap::from_iter(frag_ids.iter().copied())), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + } + } + + /// A rewrite group that straddles indexed and non-indexed fragments must not be + /// treated as corrupt-and-droppable by cleanup. The load/remap path + /// (`recalculate_fragment_bitmap`) aborts on this state and preserves the reuse + /// generation, so cleanup must report the version as "not caught up" and retain + /// it. A buggy "drop on straddle" implementation would report it as caught up and + /// silently prune the still-needed generation. + #[test] + fn test_straddling_group_is_not_caught_up() { + // Reuse version at dataset version 5 rewrites old fragments [1, 2] -> [10]. + let version = reuse_version(5, &[1, 2], &[10]); + + // Index is at a newer dataset version so the version short-circuit does not + // fire, but its bitmap only covers fragment 1 of the group. It straddles + // indexed (1) and non-indexed (2) data, which makes + // `is_index_remap_caught_up` return `Err(InvalidInput)`. + let index = index_meta_with_bitmap(7, &[1]); + assert!( + matches!( + is_index_remap_caught_up(&version, &index), + Err(Error::InvalidInput { .. }) + ), + "a straddling group must produce InvalidInput at the per-index check" + ); + + // The retention decision must treat that straddle as "not caught up" so the + // reuse generation is retained, mirroring the load path. + assert!( + !is_reuse_version_caught_up(&version, std::slice::from_ref(&index)).unwrap(), + "cleanup must retain (not drop) a reuse version whose group straddles \ + indexed and non-indexed fragments" + ); + } + + /// Sanity check the opposite direction: an index is "caught up" once it has + /// already remapped the group's old fragments away, i.e. its bitmap no longer + /// contains any of them (and it is at/after the reuse version). Such a reuse + /// version is droppable. + #[test] + fn test_remapped_group_is_caught_up() { + let version = reuse_version(5, &[1, 2], &[10]); + // Bitmap points at the new fragment only. The old fragments [1, 2] have + // already been remapped away, so this index is caught up. + let index = index_meta_with_bitmap(7, &[10]); + assert!(is_reuse_version_caught_up(&version, std::slice::from_ref(&index)).unwrap()); + + // An index still covering all old fragments has not remapped yet, so the + // version is still needed and not caught up. + let stale_index = index_meta_with_bitmap(7, &[1, 2]); + assert!(!is_reuse_version_caught_up(&version, std::slice::from_ref(&stale_index)).unwrap()); + } + /// With more than one index on the table, remapping every index must catch /// all of them up so the reuse index can be trimmed. /// diff --git a/rust/lance/src/index/frag_reuse.rs b/rust/lance/src/index/frag_reuse.rs index 23a8fec5145..000f4e2751d 100644 --- a/rust/lance/src/index/frag_reuse.rs +++ b/rust/lance/src/index/frag_reuse.rs @@ -50,16 +50,41 @@ pub async fn load_frag_reuse_index_details( .join(index.uuid.to_string()) .join(external_file.path.clone()); + // Use checked arithmetic and bounds validation so that malformed + // metadata surfaces as a proper `Error` (and therefore a Python + // exception) instead of panicking or building an invalid range. + let offset = usize::try_from(external_file.offset).map_err(|_| { + Error::index(format!( + "Fragment reuse external file offset {} does not fit in usize", + external_file.offset + )) + })?; + let size = usize::try_from(external_file.size).map_err(|_| { + Error::index(format!( + "Fragment reuse external file size {} does not fit in usize", + external_file.size + )) + })?; + let end = offset.checked_add(size).ok_or_else(|| { + Error::index(format!( + "Fragment reuse external file range overflows: offset {offset} + size {size}" + )) + })?; + // the file content will be cached in the index cache later // so we do not put it to the file cache - let range = external_file.offset as usize - ..(external_file.offset as usize + external_file.size as usize); - let data = dataset - .object_store - .open(&file_path) - .await? - .get_range(range) - .await?; + let reader = dataset.object_store.open(&file_path).await?; + let file_size = reader.size().await.map_err(|e| { + Error::index(format!( + "Failed to read size of fragment reuse external file {file_path}: {e}" + )) + })?; + if end > file_size { + return Err(Error::index(format!( + "Fragment reuse external file range {offset}..{end} is out of bounds for file {file_path} of size {file_size}" + ))); + } + let data = reader.get_range(offset..end).await?; let pb_sequence = InlineContent::decode(data)?; Ok(Arc::new(FragReuseIndexDetails::try_from(pb_sequence)?)) @@ -176,3 +201,84 @@ pub(crate) async fn build_frag_reuse_index_metadata( files: None, }) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::utils::test::{DatagenExt, FragmentCount, FragmentRowCount}; + use arrow_array::types::Int32Type; + use rstest::rstest; + + /// Size of the valid external details file written by the malformed-input test. + const PAYLOAD_LEN: usize = 16; + + /// Build index metadata pointing at an external details file with the given + /// `offset`/`size`, so we can exercise the offset/size arithmetic in + /// `load_frag_reuse_index_details` with crafted (potentially malformed) values. + fn external_details_meta(uuid: Uuid, offset: u64, size: u64) -> IndexMetadata { + let proto = FragmentReuseIndexDetails { + content: Some(Content::External(ExternalFile { + path: FRAG_REUSE_DETAILS_FILE_NAME.to_owned(), + offset, + size, + })), + }; + IndexMetadata { + uuid, + name: FRAG_REUSE_INDEX_NAME.to_string(), + fields: vec![], + dataset_version: 0, + fragment_bitmap: Some(RoaringBitmap::new()), + index_details: Some(Arc::new(prost_types::Any::from_msg(&proto).unwrap())), + index_version: 0, + created_at: None, + base_id: None, + files: None, + } + } + + /// Malformed external metadata must surface as a normal Error instead of + /// panicking or building an invalid object-store range. Each case crafts a + /// different bad offset/size against a valid PAYLOAD_LEN-byte details file + /// and asserts both the error variant and an identifying message fragment. + #[rstest] + #[case::out_of_bounds(0, PAYLOAD_LEN as u64 + 1, "out of bounds")] + #[case::offset_past_eof(PAYLOAD_LEN as u64 + 100, 1, "out of bounds")] + #[case::overflow(u64::MAX, u64::MAX, "overflow")] + #[tokio::test] + async fn test_load_external_details_malformed_offset_size_errors( + #[case] offset: u64, + #[case] size: u64, + #[case] expected_msg: &str, + ) { + let dataset = lance_datagen::gen_batch() + .col("i", lance_datagen::array::step::()) + .into_ram_dataset(FragmentCount::from(1), FragmentRowCount::from(10)) + .await + .unwrap(); + + // Write a small, valid-looking external details file. + let uuid = Uuid::new_v4(); + let file_path = dataset + .indices_dir() + .join(uuid.to_string()) + .join(FRAG_REUSE_DETAILS_FILE_NAME); + let payload = vec![0u8; PAYLOAD_LEN]; + let mut writer = dataset.object_store.create(&file_path).await.unwrap(); + writer.write_all(&payload).await.unwrap(); + writer.shutdown().await.unwrap(); + + let meta = external_details_meta(uuid, offset, size); + let err = load_frag_reuse_index_details(&dataset, &meta) + .await + .expect_err("malformed external metadata must surface as an error, not a panic"); + assert!( + matches!(err, Error::Index { .. }), + "expected an Error::Index, got: {err:?}" + ); + assert!( + err.to_string().to_lowercase().contains(expected_msg), + "expected error message to contain {expected_msg:?}, got: {err}" + ); + } +}