From ca09b40afeab3074411ca33c9ca36e6cca4b46b2 Mon Sep 17 00:00:00 2001 From: hello-peter-tang <132056702+hello-peter-tang@users.noreply.github.com> Date: Tue, 3 Nov 2026 10:42:14 +0800 Subject: [PATCH] feat: add validate_fragments for distributed fragment validation Add `daft_lance.validate_fragments`, a read-only operation that validates selected Lance fragments without scanning rows. The driver resolves one dataset version and enumerates fragment metadata; Daft workers reopen that pinned snapshot through the existing DatasetOpenContext and call Lance's native LanceFragment.validate(), so validation runs with bounded concurrency on local and distributed runners. The returned DataFrame has one row per selected fragment with version, fragment_id, valid and error. Worker open/authentication failures carry a distinct prefix so they are not reported as corruption. Unknown fragment IDs and empty selections raise ValueError; the dataset version never changes. Reuses Lance's validator rather than reimplementing one. Closes #63. --- daft_lance/__init__.py | 2 + daft_lance/_lance.py | 87 ++++++++++ daft_lance/lance_fragment_validation.py | 119 +++++++++++++ tests/io/lance/test_fragment_validation.py | 187 +++++++++++++++++++++ 4 files changed, 395 insertions(+) create mode 100644 daft_lance/lance_fragment_validation.py create mode 100644 tests/io/lance/test_fragment_validation.py diff --git a/daft_lance/__init__.py b/daft_lance/__init__.py index 5819b62..ddb02c1 100644 --- a/daft_lance/__init__.py +++ b/daft_lance/__init__.py @@ -14,6 +14,7 @@ read_lance, update_columns, update_columns_df, + validate_fragments, write_lance, ) from .lance_update_column import UpdateColumnsResult @@ -31,5 +32,6 @@ "take_blobs", "update_columns", "update_columns_df", + "validate_fragments", "write_lance", ] diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 728defd..5dd8e20 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -15,6 +15,7 @@ from .lance_compaction import compact_files_internal from .lance_data_sink import LanceDataSink +from .lance_fragment_validation import validate_fragments_internal from .lance_merge_column import merge_columns_from_df, merge_columns_internal from .lance_scalar_index import create_scalar_index_internal, optimize_indices_internal from .lance_scan import LanceScanOperator @@ -1098,6 +1099,92 @@ def optimize_indices( ) +@PublicAPI +def validate_fragments( + uri: str | pathlib.Path | None = None, + io_config: IOConfig | None = None, + *, + table_id: list[str] | None = None, + namespace_impl: str | None = None, + namespace_properties: dict[str, str] | None = None, + storage_options: dict[str, Any] | None = None, + version: int | str | None = None, + asof: str | None = None, + block_size: int | None = None, + index_cache_size: int | None = None, + metadata_cache_size_bytes: int | None = None, + fragment_ids: Sequence[int] | None = None, + concurrency: int | None = None, +) -> DataFrame: + """Validate selected Lance fragments without scanning table rows. + + The driver resolves one dataset version and enumerates fragment metadata; + Daft workers reopen that pinned snapshot and call Lance's native + ``LanceFragment.validate()``, which checks fragment file and metadata + consistency, including deletion vectors. The operation is read-only: it + never commits, repairs, or changes the dataset version. + + Args: + uri: The URI of the Lance table (supports remote URLs to object stores such as `s3://` or `gs://`) + io_config: A custom IOConfig to use when accessing Lance data. Defaults to None. + table_id: Table identifier within the namespace, e.g. ["catalog", "schema", "table"]. + Mutually exclusive with ``uri``. + namespace_impl: Lance Namespace implementation, e.g. "dir" or "rest". + namespace_properties: Properties for connecting to the namespace, e.g. + {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest". + storage_options: Extra options for storage connection. + version: If specified, validate a specific version of the Lance dataset. + asof: If specified, validate the latest version created on or earlier than the given value. + block_size: Block size in bytes. Provide a hint for the size of the minimal I/O request. + index_cache_size: Index cache size. + metadata_cache_size_bytes: Size of the metadata cache in bytes. + fragment_ids: Optional subset of fragment IDs to validate. Every worker + reopens the same pinned version, so validation only covers the + listed fragments. Duplicates are ignored; unknown IDs and an empty + list raise ``ValueError``. When None, every fragment in the pinned + manifest is validated. + concurrency: Maximum number of concurrent validation tasks. Defaults + to Daft's runner default. Must be a positive integer. + + Returns: + DataFrame: One row per selected fragment with the columns + ``version`` (Int64), ``fragment_id`` (Int64), ``valid`` (Boolean), + and ``error`` (String). ``error`` is null for valid fragments; + otherwise it holds the validation exception. A table-open or + authentication failure on a worker is reported with a + ``failed to open the dataset at version ...`` prefix so it can be + distinguished from detected corruption. + + Raises: + ValueError: If ``fragment_ids`` contains IDs that are not in the + pinned version, or if ``concurrency`` is not positive. + + Note: + ``LanceFragment.validate`` requires PyLance >= 11; daft-lance already + requires a newer PyLance (>= 12). + """ + dataset_handle = construct_lance_dataset_handle( + uri, + storage_options=storage_options, + io_config=io_config, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + version=version, + asof=asof, + block_size=block_size, + index_cache_size=index_cache_size, + metadata_cache_size_bytes=metadata_cache_size_bytes, + ) + + return validate_fragments_internal( + dataset_handle.dataset, + dataset_handle.worker_open_context(), + fragment_ids=fragment_ids, + concurrency=concurrency, + ) + + @PublicAPI def compact_files( uri: str | pathlib.Path | None = None, diff --git a/daft_lance/lance_fragment_validation.py b/daft_lance/lance_fragment_validation.py new file mode 100644 index 0000000..c66187b --- /dev/null +++ b/daft_lance/lance_fragment_validation.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +import logging +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import lance +import pyarrow as pa + +import daft + +if TYPE_CHECKING: + from daft_lance.namespace import DatasetOpenContext + +logger = logging.getLogger(__name__) + +#: One row per selected fragment; ``error`` is null when the fragment is valid. +EMPTY_REPORT = pa.table( + { + "version": pa.array([], type=pa.int64()), + "fragment_id": pa.array([], type=pa.int64()), + "valid": pa.array([], type=pa.bool_()), + "error": pa.array([], type=pa.string()), + } +) + + +class FragmentValidationUDF: + """Validate Lance fragments on a Daft worker against a pinned snapshot. + + The dataset is opened once per worker instance, not per fragment, so the + pinned-manifest read is amortised across the fragments that worker owns. + """ + + def __init__(self, open_context: DatasetOpenContext, version: int) -> None: + self.open_context = open_context + self.version = version + self._lance_ds: lance.LanceDataset | None = None + + def _dataset(self) -> lance.LanceDataset: + if self._lance_ds is None: + self._lance_ds = self.open_context.open_pinned() + return self._lance_ds + + def __call__(self, fragment_id: int) -> str | None: + try: + dataset = self._dataset() + except Exception as exc: + # A table-open/authentication failure is not evidence of corruption, + # so report it with a distinct prefix instead of a validation error. + return f"failed to open the dataset at version {self.version}: {type(exc).__name__}: {exc}" + + try: + fragment = dataset.get_fragment(int(fragment_id)) + if fragment is None: + return f"fragment {fragment_id} is not present in pinned version {self.version}" + fragment.validate() + except Exception as exc: + return f"{type(exc).__name__}: {exc}" + return None + + +def _selected_fragment_ids( + available: list[int], + fragment_ids: Sequence[int] | None, + version: int, +) -> list[int]: + if fragment_ids is None: + return available + + requested = list(dict.fromkeys(fragment_ids)) + if not requested: + raise ValueError("fragment_ids must not be empty; pass None to validate every fragment") + available_set = set(available) + unknown = [fragment_id for fragment_id in requested if fragment_id not in available_set] + if unknown: + raise ValueError( + f"Fragment IDs {unknown} do not exist in dataset version {version}. " + f"Available fragment IDs: {sorted(available_set)}." + ) + return requested + + +def validate_fragments_internal( + lance_ds: lance.LanceDataset, + open_context: DatasetOpenContext, + *, + fragment_ids: Sequence[int] | None = None, + concurrency: int | None = None, +) -> daft.DataFrame: + """Validate selected fragments of the pinned snapshot ``open_context`` describes. + + ``lance_ds`` is the driver's live dataset and is only used to enumerate + fragment metadata and resolve the version; ``open_context`` is what workers + reopen from. + """ + if concurrency is not None and concurrency <= 0: + raise ValueError(f"concurrency must be positive, got {concurrency}") + + version = lance_ds.version + # ``get_fragments`` reads fragment metadata from the pinned manifest; it + # does not scan rows, so fragments with no live rows are still listed. + available = [fragment.fragment_id for fragment in lance_ds.get_fragments()] + selected = _selected_fragment_ids(available, fragment_ids, version) + logger.info("Validating %d of %d fragments at dataset version %d", len(selected), len(available), version) + + if not selected: + return daft.from_arrow(EMPTY_REPORT) + + df = daft.from_pylist([{"fragment_id": fragment_id} for fragment_id in selected]) + validate_fragment = daft.cls(FragmentValidationUDF, max_concurrency=concurrency) + validator = validate_fragment(open_context, version) + df = df.select(df["fragment_id"], validator(df["fragment_id"]).alias("error")) + return df.select( + daft.lit(version).alias("version"), + df["fragment_id"], + df["error"].is_null().alias("valid"), + df["error"], + ) diff --git a/tests/io/lance/test_fragment_validation.py b/tests/io/lance/test_fragment_validation.py new file mode 100644 index 0000000..655db8a --- /dev/null +++ b/tests/io/lance/test_fragment_validation.py @@ -0,0 +1,187 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import lance +import pyarrow as pa +import pytest + +import daft +import daft_lance +from daft_lance.lance_fragment_validation import FragmentValidationUDF +from daft_lance.namespace import DatasetOpenContext + + +def _write_fragments(path: Path, num_fragments: int = 3) -> lance.LanceDataset: + for fragment in range(num_fragments): + table = pa.table({"id": list(range(fragment * 10, (fragment + 1) * 10))}) + lance.write_dataset( + table, + str(path), + mode="create" if fragment == 0 else "append", + max_rows_per_file=10, + max_rows_per_group=10, + ) + return lance.dataset(str(path)) + + +def _rows_by_fragment(report: daft.DataFrame) -> dict[int, dict[str, Any]]: + return {row["fragment_id"]: row for row in report.to_pylist()} + + +def test_validates_every_fragment_of_a_healthy_dataset(tmp_path: Path) -> None: + path = tmp_path / "healthy.lance" + dataset = _write_fragments(path) + + report = daft_lance.validate_fragments(str(path)) + + assert report.schema().column_names() == ["version", "fragment_id", "valid", "error"] + rows = report.to_pylist() + assert [row["fragment_id"] for row in rows] == [0, 1, 2] + assert {row["version"] for row in rows} == {dataset.version} + assert all(row["valid"] for row in rows) + assert all(row["error"] is None for row in rows) + # Validation is read-only. + assert lance.dataset(str(path)).version == dataset.version + + +def test_validates_only_selected_fragments(tmp_path: Path) -> None: + path = tmp_path / "subset.lance" + _write_fragments(path) + + rows = daft_lance.validate_fragments(str(path), fragment_ids=[2, 0, 2]).to_pylist() + + assert sorted(row["fragment_id"] for row in rows) == [0, 2] + assert all(row["valid"] for row in rows) + + +def test_unknown_fragment_ids_raise(tmp_path: Path) -> None: + path = tmp_path / "unknown.lance" + _write_fragments(path) + + with pytest.raises(ValueError, match="do not exist"): + daft_lance.validate_fragments(str(path), fragment_ids=[999]) + + +def test_empty_fragment_selection_is_rejected(tmp_path: Path) -> None: + path = tmp_path / "empty.lance" + _write_fragments(path) + + with pytest.raises(ValueError, match="must not be empty"): + daft_lance.validate_fragments(str(path), fragment_ids=[]) + + +def test_non_positive_concurrency_is_rejected(tmp_path: Path) -> None: + path = tmp_path / "concurrency.lance" + _write_fragments(path, num_fragments=1) + + with pytest.raises(ValueError, match="concurrency must be positive"): + daft_lance.validate_fragments(str(path), concurrency=0) + + +def test_fragment_with_multiple_column_files_is_valid(tmp_path: Path) -> None: + path = tmp_path / "multi_file.lance" + dataset = _write_fragments(path, num_fragments=1) + dataset.add_columns({"doubled": "id * 2"}) + + fragment = lance.dataset(str(path)).get_fragment(0) + assert fragment is not None + assert len(fragment.metadata.files) == 2 + + rows = daft_lance.validate_fragments(str(path)).to_pylist() + assert [row["valid"] for row in rows] == [True] + + +def test_fragment_with_deletion_vector_is_valid(tmp_path: Path) -> None: + path = tmp_path / "deletions.lance" + dataset = _write_fragments(path) + dataset.delete("id < 5") + + assert [fragment.fragment_id for fragment in lance.dataset(str(path)).get_fragments()] == [0, 1, 2] + rows = daft_lance.validate_fragments(str(path)).to_pylist() + + assert len(rows) == 3 + assert all(row["valid"] for row in rows) + + +def test_fully_deleted_fragment_is_not_reported_as_corrupt(tmp_path: Path) -> None: + path = tmp_path / "fully_deleted.lance" + dataset = _write_fragments(path) + # Deleting every row of a fragment removes it from the manifest, so the + # report follows the manifest rather than failing on the absent fragment. + dataset.delete("id < 10") + remaining = [fragment.fragment_id for fragment in lance.dataset(str(path)).get_fragments()] + assert 0 not in remaining + + rows = daft_lance.validate_fragments(str(path)).to_pylist() + + assert sorted(row["fragment_id"] for row in rows) == remaining + assert all(row["valid"] for row in rows) + + +def test_missing_data_file_is_attributed_to_its_fragment(tmp_path: Path) -> None: + path = tmp_path / "missing_file.lance" + dataset = _write_fragments(path) + victim = dataset.get_fragment(1) + assert victim is not None + (path / "data" / victim.metadata.files[0].path).unlink() + + rows = _rows_by_fragment(daft_lance.validate_fragments(str(path))) + + assert rows[0]["valid"] is True + assert rows[1]["valid"] is False + assert rows[1]["error"] + assert rows[2]["valid"] is True + + +def test_validation_snapshot_is_pinned_against_concurrent_appends(tmp_path: Path) -> None: + path = tmp_path / "pinned.lance" + dataset = _write_fragments(path, num_fragments=2) + + planned = daft_lance.validate_fragments(str(path)) + # Append between planning and execution: the worker must still validate + # the snapshot the driver planned against. + lance.write_dataset(pa.table({"id": [100]}), str(path), mode="append") + assert lance.dataset(str(path)).version > dataset.version + + rows = planned.to_pylist() + assert {row["version"] for row in rows} == {dataset.version} + assert sorted(row["fragment_id"] for row in rows) == [0, 1] + + +def test_explicit_version_validates_that_snapshot(tmp_path: Path) -> None: + path = tmp_path / "versioned.lance" + first = _write_fragments(path, num_fragments=1) + first_version = first.version + lance.write_dataset(pa.table({"id": [99]}), str(path), mode="append") + assert lance.dataset(str(path)).version > first_version + + rows = daft_lance.validate_fragments(str(path), version=first_version).to_pylist() + + assert {row["version"] for row in rows} == {first_version} + assert len(rows) == 1 + + +def test_namespace_addressed_table(tmp_path: Path) -> None: + namespace: dict[str, Any] = {"namespace_impl": "dir", "namespace_properties": {"root": str(tmp_path)}} + daft_lance.write_lance( + daft.from_pydict({"id": [1, 2, 3]}), + table_id=["ns_table"], + mode="create", + **namespace, + ).collect() + + rows = daft_lance.validate_fragments(table_id=["ns_table"], **namespace).to_pylist() + + assert len(rows) == 1 + assert rows[0]["valid"] is True + + +def test_worker_open_failure_is_distinguished_from_corruption(tmp_path: Path) -> None: + context = DatasetOpenContext(uri=str(tmp_path / "missing.lance"), version=7) + + message = FragmentValidationUDF(context, context.version)(0) + + assert message is not None + assert message.startswith("failed to open the dataset at version 7")