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")