diff --git a/python/python/ci_benchmarks/benchmarks/test_merge_insert.py b/python/python/ci_benchmarks/benchmarks/test_merge_insert.py index b1d0170db24..f2e6ad2fc03 100644 --- a/python/python/ci_benchmarks/benchmarks/test_merge_insert.py +++ b/python/python/ci_benchmarks/benchmarks/test_merge_insert.py @@ -95,12 +95,13 @@ def reset(self, cold: bool = False) -> lance.LanceDataset: def _open_target(name: str) -> Iterable[Target]: uri = get_dataset_uri(name) dataset = lance.dataset(uri) - base_version = dataset.tags.get_version(BASE_TAG) - if base_version is None: + tags = dataset.tags.list() + if BASE_TAG not in tags: pytest.skip( f"Dataset {name} has no {BASE_TAG} tag; " "run python/ci_benchmarks/datagen/gen_all.py" ) + base_version = tags[BASE_TAG]["version"] yield Target(uri=uri, dataset=dataset, base_version=base_version) diff --git a/python/python/ci_benchmarks/datagen/merge_insert.py b/python/python/ci_benchmarks/datagen/merge_insert.py index ebda5730aa2..33f165d8772 100644 --- a/python/python/ci_benchmarks/datagen/merge_insert.py +++ b/python/python/ci_benchmarks/datagen/merge_insert.py @@ -257,14 +257,18 @@ def _already_generated(uri: str, expected_rows: int) -> bool: A previous benchmark run may have left extra versions behind, so the row count is checked at the tagged version rather than at the latest one. + + Incomplete generations (dataset written, tag never created) and missing + tags both return False so the caller can overwrite and retag. """ try: ds = lance.dataset(uri) except ValueError: return False - base_version = ds.tags.get_version(BASE_TAG) - if base_version is None: + tags = ds.tags.list() + if BASE_TAG not in tags: return False + base_version = tags[BASE_TAG]["version"] return ds.checkout_version(base_version).count_rows() == expected_rows diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 3cdd44ea35e..1da97869668 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -7240,7 +7240,7 @@ def list(self) -> dict[str, Tag]: """ return self._ds.tags() - def get_version(self, tag: str) -> Optional[int]: + def get_version(self, tag: str) -> int: """ Get the version of a specific tag by name. @@ -7251,8 +7251,13 @@ def get_version(self, tag: str) -> Optional[int]: Returns ------- - int or None - The version number of the tag if it exists, otherwise None. + int + The version number of the tag. + + Raises + ------ + ValueError + If the tag does not exist. Use :meth:`list` to check for presence. """ return self._ds.get_version(tag)