From 807f992e0eb1503a83797431696f880233ba2173 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 3 Aug 2026 16:50:59 +0800 Subject: [PATCH 1/5] fix(ci): treat missing merge_insert base tag as not generated bench_regress has failed on every main push since #8052 because _already_generated assumed Tags.get_version returns None for a missing tag, while the API raises ValueError. Shared GCS datasets left without merge_insert_base then crash datagen before overwrite can recover. Use tags.list() for presence checks and correct the get_version docs. --- .../benchmarks/test_merge_insert.py | 5 +- .../ci_benchmarks/datagen/merge_insert.py | 8 ++- .../datagen/test_merge_insert_datagen.py | 55 +++++++++++++++++++ python/python/lance/dataset.py | 11 +++- 4 files changed, 72 insertions(+), 7 deletions(-) create mode 100644 python/python/ci_benchmarks/datagen/test_merge_insert_datagen.py 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/ci_benchmarks/datagen/test_merge_insert_datagen.py b/python/python/ci_benchmarks/datagen/test_merge_insert_datagen.py new file mode 100644 index 00000000000..06db0fff4b6 --- /dev/null +++ b/python/python/ci_benchmarks/datagen/test_merge_insert_datagen.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright The Lance Authors + +"""Unit tests for merge_insert datagen helpers.""" + +from __future__ import annotations + +import lance +import numpy as np + +from ci_benchmarks.datagen.merge_insert import ( + BASE_TAG, + NARROW_SCHEMA, + _already_generated, + _tag_base, + narrow_batch, +) + + +def test_already_generated_missing_dataset(tmp_path): + assert not _already_generated(str(tmp_path / "missing"), expected_rows=10) + + +def test_already_generated_missing_base_tag(tmp_path): + """Incomplete generation: dataset exists but the base tag was never written.""" + uri = str(tmp_path / "no_tag") + lance.write_dataset( + narrow_batch(np.arange(10, dtype=np.int64)), + uri, + schema=NARROW_SCHEMA, + ) + assert not _already_generated(uri, expected_rows=10) + + +def test_already_generated_wrong_row_count(tmp_path): + uri = str(tmp_path / "wrong_rows") + ds = lance.write_dataset( + narrow_batch(np.arange(10, dtype=np.int64)), + uri, + schema=NARROW_SCHEMA, + ) + _tag_base(ds) + assert not _already_generated(uri, expected_rows=20) + + +def test_already_generated_ready(tmp_path): + uri = str(tmp_path / "ready") + ds = lance.write_dataset( + narrow_batch(np.arange(10, dtype=np.int64)), + uri, + schema=NARROW_SCHEMA, + ) + _tag_base(ds) + assert _already_generated(uri, expected_rows=10) + assert BASE_TAG in lance.dataset(uri).tags.list() 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) From ed09675db49932c13d50f8d5344681d3592a9f9a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 3 Aug 2026 16:54:54 +0800 Subject: [PATCH 2/5] ci: run datagen unit tests on PR and before GCS gen Wire python/ci_benchmarks/datagen into make test so PR Python CI catches readiness-check regressions. Run the same suite before gen_all in bench_regress so pure-Python failures do not wait on shared GCS generation. --- .github/workflows/ci-benchmarks.yml | 6 ++++++ python/Makefile | 4 +++- python/python/ci_benchmarks/README.md | 13 +++++++++++++ 3 files changed, 22 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci-benchmarks.yml b/.github/workflows/ci-benchmarks.yml index aa6e73ae74f..b3bacab2bbe 100644 --- a/.github/workflows/ci-benchmarks.yml +++ b/.github/workflows/ci-benchmarks.yml @@ -53,6 +53,12 @@ jobs: run: | source venv/bin/activate make -C ../memtest build-release + # Unit tests first: gen_all touches shared GCS and is expensive to debug + # when a pure-Python readiness check is wrong. + - name: Run datagen unit tests + run: | + source venv/bin/activate + pytest python/ci_benchmarks/datagen -q - name: Generate datasets run: | source venv/bin/activate diff --git a/python/Makefile b/python/Makefile index d5077019f35..6e8097b629d 100644 --- a/python/Makefile +++ b/python/Makefile @@ -33,7 +33,9 @@ build: ## Build the local Rust extension with maturin $(UV_RUN) maturin develop --uv test: ## Run Python tests except recurring tests - pytest $(PYTEST_ARGS) python/tests + # Include ci_benchmarks/datagen unit tests so PR CI catches datagen + # regressions before main-only bench_regress hits shared GCS state. + pytest $(PYTEST_ARGS) python/tests python/ci_benchmarks/datagen integtest: ## Start LocalStack and run integration tests @if [ "$(KEEP_COMPOSE)" = "1" ]; then \ diff --git a/python/python/ci_benchmarks/README.md b/python/python/ci_benchmarks/README.md index 0245d29166f..643478e8f92 100644 --- a/python/python/ci_benchmarks/README.md +++ b/python/python/ci_benchmarks/README.md @@ -19,6 +19,19 @@ ci_benchmarks/ └── datasets.py # Dataset URI resolver (local vs GCS) ``` +## Datagen unit tests + +Logic under `datagen/` (for example dataset readiness checks) has unit tests in +`datagen/test_*.py`. These run as part of the normal Python test suite +(`make test` / PR CI) and again before dataset generation in +`ci-benchmarks.yml`, so pure-Python regressions fail before main-only GCS +generation. + +```bash +# from python/ +pytest python/ci_benchmarks/datagen +``` + ## Running Benchmarks Locally ### 1. Generate test datasets From 095ccc7b0a88b9fabbfdd61ced8e19c4bab5eb1f Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 3 Aug 2026 17:03:20 +0800 Subject: [PATCH 3/5] fix(ci): collect merge_insert datagen tests under python/tests Collecting python/ci_benchmarks/datagen in make test put the source tree on sys.path ahead of the installed wheel, so import lance resolved to pure-Python sources without the native extension and macOS CI failed with BlobDescriptor ImportError. Move the readiness-check unit tests into python/tests, load the installed lance package first, then import ci_benchmarks helpers. --- .github/workflows/ci-benchmarks.yml | 2 +- python/Makefile | 7 ++++--- python/python/ci_benchmarks/README.md | 16 ++++++++------ .../test_merge_insert_datagen.py | 21 ++++++++++++++++--- 4 files changed, 33 insertions(+), 13 deletions(-) rename python/python/{ci_benchmarks/datagen => tests}/test_merge_insert_datagen.py (63%) diff --git a/.github/workflows/ci-benchmarks.yml b/.github/workflows/ci-benchmarks.yml index b3bacab2bbe..75e1434eaec 100644 --- a/.github/workflows/ci-benchmarks.yml +++ b/.github/workflows/ci-benchmarks.yml @@ -58,7 +58,7 @@ jobs: - name: Run datagen unit tests run: | source venv/bin/activate - pytest python/ci_benchmarks/datagen -q + pytest python/tests/test_merge_insert_datagen.py -q - name: Generate datasets run: | source venv/bin/activate diff --git a/python/Makefile b/python/Makefile index 6e8097b629d..893f44906ac 100644 --- a/python/Makefile +++ b/python/Makefile @@ -33,9 +33,10 @@ build: ## Build the local Rust extension with maturin $(UV_RUN) maturin develop --uv test: ## Run Python tests except recurring tests - # Include ci_benchmarks/datagen unit tests so PR CI catches datagen - # regressions before main-only bench_regress hits shared GCS state. - pytest $(PYTEST_ARGS) python/tests python/ci_benchmarks/datagen + # Datagen readiness checks live under python/tests (e.g. + # test_merge_insert_datagen.py) so PR CI covers them without collecting + # ci_benchmarks/ as a path (which would shadow the installed wheel). + pytest $(PYTEST_ARGS) python/tests integtest: ## Start LocalStack and run integration tests @if [ "$(KEEP_COMPOSE)" = "1" ]; then \ diff --git a/python/python/ci_benchmarks/README.md b/python/python/ci_benchmarks/README.md index 643478e8f92..8a036ce9b88 100644 --- a/python/python/ci_benchmarks/README.md +++ b/python/python/ci_benchmarks/README.md @@ -21,15 +21,19 @@ ci_benchmarks/ ## Datagen unit tests -Logic under `datagen/` (for example dataset readiness checks) has unit tests in -`datagen/test_*.py`. These run as part of the normal Python test suite -(`make test` / PR CI) and again before dataset generation in -`ci-benchmarks.yml`, so pure-Python regressions fail before main-only GCS -generation. +Logic under `datagen/` (for example dataset readiness checks) is covered by +unit tests under `python/tests/` (e.g. `test_merge_insert_datagen.py`). Those +run as part of the normal Python suite (`make test` / PR CI) and again before +dataset generation in `ci-benchmarks.yml`, so pure-Python regressions fail +before main-only GCS generation. + +They are intentionally **not** collected via `pytest python/ci_benchmarks/...` +in PR CI: that path puts the source tree on `sys.path` and shadows the +installed `lance` wheel extension. ```bash # from python/ -pytest python/ci_benchmarks/datagen +pytest python/tests/test_merge_insert_datagen.py ``` ## Running Benchmarks Locally diff --git a/python/python/ci_benchmarks/datagen/test_merge_insert_datagen.py b/python/python/tests/test_merge_insert_datagen.py similarity index 63% rename from python/python/ci_benchmarks/datagen/test_merge_insert_datagen.py rename to python/python/tests/test_merge_insert_datagen.py index 06db0fff4b6..07c3449bde7 100644 --- a/python/python/ci_benchmarks/datagen/test_merge_insert_datagen.py +++ b/python/python/tests/test_merge_insert_datagen.py @@ -1,14 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors -"""Unit tests for merge_insert datagen helpers.""" +"""Unit tests for merge_insert CI datagen readiness checks. + +These live under ``python/tests`` (not ``ci_benchmarks/``) so PR CI can collect +them without putting the source tree on ``sys.path`` ahead of the installed +wheel. Import ``lance`` first, then expose ``ci_benchmarks`` for the helpers. +""" from __future__ import annotations -import lance +import sys +from pathlib import Path + +# Load the installed native extension before adding the source tree for +# ``ci_benchmarks``. Otherwise ``import lance`` resolves to pure-Python sources +# without ``lance.lance`` and collection fails with BlobDescriptor ImportError. +import lance # noqa: F401 import numpy as np -from ci_benchmarks.datagen.merge_insert import ( +_PYTHON_SRC = Path(__file__).resolve().parents[1] +if str(_PYTHON_SRC) not in sys.path: + sys.path.insert(0, str(_PYTHON_SRC)) + +from ci_benchmarks.datagen.merge_insert import ( # noqa: E402 BASE_TAG, NARROW_SCHEMA, _already_generated, From fa50defc4f8fe4d830277dc24ed06849e0c63926 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 3 Aug 2026 17:28:14 +0800 Subject: [PATCH 4/5] fix(ci): restore sys.path after importing ci_benchmarks helpers Leaving the source tree on sys.path was inherited by multiprocessing spawn children (fragment progress and torch DataLoader workers), which then imported pure-Python lance without the native extension. --- .../python/tests/test_merge_insert_datagen.py | 43 ++++++++++++------- 1 file changed, 27 insertions(+), 16 deletions(-) diff --git a/python/python/tests/test_merge_insert_datagen.py b/python/python/tests/test_merge_insert_datagen.py index 07c3449bde7..e369ac11cfb 100644 --- a/python/python/tests/test_merge_insert_datagen.py +++ b/python/python/tests/test_merge_insert_datagen.py @@ -4,8 +4,13 @@ """Unit tests for merge_insert CI datagen readiness checks. These live under ``python/tests`` (not ``ci_benchmarks/``) so PR CI can collect -them without putting the source tree on ``sys.path`` ahead of the installed -wheel. Import ``lance`` first, then expose ``ci_benchmarks`` for the helpers. +them without putting the source tree on ``sys.path`` for the whole suite. + +``ci_benchmarks`` is not installed with the wheel, so we temporarily expose the +source tree only while importing the helpers, then restore ``sys.path``. Leaving +the path in place breaks later tests that use ``multiprocessing`` spawn: the +child inherits ``sys.path`` and resolves pure-Python ``lance`` without the +native extension. """ from __future__ import annotations @@ -13,23 +18,29 @@ import sys from pathlib import Path -# Load the installed native extension before adding the source tree for -# ``ci_benchmarks``. Otherwise ``import lance`` resolves to pure-Python sources -# without ``lance.lance`` and collection fails with BlobDescriptor ImportError. +# Load the installed native extension before temporarily adding the source tree +# for ``ci_benchmarks``. import lance # noqa: F401 import numpy as np -_PYTHON_SRC = Path(__file__).resolve().parents[1] -if str(_PYTHON_SRC) not in sys.path: - sys.path.insert(0, str(_PYTHON_SRC)) - -from ci_benchmarks.datagen.merge_insert import ( # noqa: E402 - BASE_TAG, - NARROW_SCHEMA, - _already_generated, - _tag_base, - narrow_batch, -) +_PYTHON_SRC = str(Path(__file__).resolve().parents[1]) +_path_added = _PYTHON_SRC not in sys.path +if _path_added: + sys.path.insert(0, _PYTHON_SRC) +try: + from ci_benchmarks.datagen.merge_insert import ( # noqa: E402 + BASE_TAG, + NARROW_SCHEMA, + _already_generated, + _tag_base, + narrow_batch, + ) +finally: + if _path_added: + try: + sys.path.remove(_PYTHON_SRC) + except ValueError: + pass def test_already_generated_missing_dataset(tmp_path): From 14a13f93d69a6912c67b0b2b743589da2ee5b3c3 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 3 Aug 2026 17:32:10 +0800 Subject: [PATCH 5/5] revert: drop optional merge_insert datagen unit tests ci_benchmarks is not part of the installed wheel, so PR-suite coverage needed sys.path hacks that broke multiprocessing spawn. The readiness check fix itself is sufficient; skip the non-essential test wiring. --- .github/workflows/ci-benchmarks.yml | 6 -- python/Makefile | 3 - python/python/ci_benchmarks/README.md | 17 ---- .../python/tests/test_merge_insert_datagen.py | 81 ------------------- 4 files changed, 107 deletions(-) delete mode 100644 python/python/tests/test_merge_insert_datagen.py diff --git a/.github/workflows/ci-benchmarks.yml b/.github/workflows/ci-benchmarks.yml index 75e1434eaec..aa6e73ae74f 100644 --- a/.github/workflows/ci-benchmarks.yml +++ b/.github/workflows/ci-benchmarks.yml @@ -53,12 +53,6 @@ jobs: run: | source venv/bin/activate make -C ../memtest build-release - # Unit tests first: gen_all touches shared GCS and is expensive to debug - # when a pure-Python readiness check is wrong. - - name: Run datagen unit tests - run: | - source venv/bin/activate - pytest python/tests/test_merge_insert_datagen.py -q - name: Generate datasets run: | source venv/bin/activate diff --git a/python/Makefile b/python/Makefile index 893f44906ac..d5077019f35 100644 --- a/python/Makefile +++ b/python/Makefile @@ -33,9 +33,6 @@ build: ## Build the local Rust extension with maturin $(UV_RUN) maturin develop --uv test: ## Run Python tests except recurring tests - # Datagen readiness checks live under python/tests (e.g. - # test_merge_insert_datagen.py) so PR CI covers them without collecting - # ci_benchmarks/ as a path (which would shadow the installed wheel). pytest $(PYTEST_ARGS) python/tests integtest: ## Start LocalStack and run integration tests diff --git a/python/python/ci_benchmarks/README.md b/python/python/ci_benchmarks/README.md index 8a036ce9b88..0245d29166f 100644 --- a/python/python/ci_benchmarks/README.md +++ b/python/python/ci_benchmarks/README.md @@ -19,23 +19,6 @@ ci_benchmarks/ └── datasets.py # Dataset URI resolver (local vs GCS) ``` -## Datagen unit tests - -Logic under `datagen/` (for example dataset readiness checks) is covered by -unit tests under `python/tests/` (e.g. `test_merge_insert_datagen.py`). Those -run as part of the normal Python suite (`make test` / PR CI) and again before -dataset generation in `ci-benchmarks.yml`, so pure-Python regressions fail -before main-only GCS generation. - -They are intentionally **not** collected via `pytest python/ci_benchmarks/...` -in PR CI: that path puts the source tree on `sys.path` and shadows the -installed `lance` wheel extension. - -```bash -# from python/ -pytest python/tests/test_merge_insert_datagen.py -``` - ## Running Benchmarks Locally ### 1. Generate test datasets diff --git a/python/python/tests/test_merge_insert_datagen.py b/python/python/tests/test_merge_insert_datagen.py deleted file mode 100644 index e369ac11cfb..00000000000 --- a/python/python/tests/test_merge_insert_datagen.py +++ /dev/null @@ -1,81 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright The Lance Authors - -"""Unit tests for merge_insert CI datagen readiness checks. - -These live under ``python/tests`` (not ``ci_benchmarks/``) so PR CI can collect -them without putting the source tree on ``sys.path`` for the whole suite. - -``ci_benchmarks`` is not installed with the wheel, so we temporarily expose the -source tree only while importing the helpers, then restore ``sys.path``. Leaving -the path in place breaks later tests that use ``multiprocessing`` spawn: the -child inherits ``sys.path`` and resolves pure-Python ``lance`` without the -native extension. -""" - -from __future__ import annotations - -import sys -from pathlib import Path - -# Load the installed native extension before temporarily adding the source tree -# for ``ci_benchmarks``. -import lance # noqa: F401 -import numpy as np - -_PYTHON_SRC = str(Path(__file__).resolve().parents[1]) -_path_added = _PYTHON_SRC not in sys.path -if _path_added: - sys.path.insert(0, _PYTHON_SRC) -try: - from ci_benchmarks.datagen.merge_insert import ( # noqa: E402 - BASE_TAG, - NARROW_SCHEMA, - _already_generated, - _tag_base, - narrow_batch, - ) -finally: - if _path_added: - try: - sys.path.remove(_PYTHON_SRC) - except ValueError: - pass - - -def test_already_generated_missing_dataset(tmp_path): - assert not _already_generated(str(tmp_path / "missing"), expected_rows=10) - - -def test_already_generated_missing_base_tag(tmp_path): - """Incomplete generation: dataset exists but the base tag was never written.""" - uri = str(tmp_path / "no_tag") - lance.write_dataset( - narrow_batch(np.arange(10, dtype=np.int64)), - uri, - schema=NARROW_SCHEMA, - ) - assert not _already_generated(uri, expected_rows=10) - - -def test_already_generated_wrong_row_count(tmp_path): - uri = str(tmp_path / "wrong_rows") - ds = lance.write_dataset( - narrow_batch(np.arange(10, dtype=np.int64)), - uri, - schema=NARROW_SCHEMA, - ) - _tag_base(ds) - assert not _already_generated(uri, expected_rows=20) - - -def test_already_generated_ready(tmp_path): - uri = str(tmp_path / "ready") - ds = lance.write_dataset( - narrow_batch(np.arange(10, dtype=np.int64)), - uri, - schema=NARROW_SCHEMA, - ) - _tag_base(ds) - assert _already_generated(uri, expected_rows=10) - assert BASE_TAG in lance.dataset(uri).tags.list()