diff --git a/README.md b/README.md index 7517ef7..661de0d 100644 --- a/README.md +++ b/README.md @@ -36,12 +36,42 @@ create_scalar_index("s3://bucket/my_dataset", column="name", index_type="INVERTE create_scalar_index("s3://bucket/my_dataset", column="ts", index_type="ZONEMAP") ``` +Types without a distributed path (e.g. `RTREE`) raise `ValueError` — call +pylance directly (`lance.dataset(uri).create_scalar_index(...)`) for +single-node indexing. + `replace=True` (the default) atomically swaps an existing index of the same name: the segment commit retires the old overlapped segments in the same -transaction as the new ones. `replace=False` rejects an existing name. Types -without a distributed path (e.g. `RTREE`) raise `ValueError` — call pylance -directly (`lance.dataset(uri).create_scalar_index(...)`) for single-node -indexing. +transaction as the new ones (the index type may change on a full rebuild; +the column may not — drop the index or use a different name for that). +`replace=False` rejects an existing name. + +#### Partial Builds and Incremental Backfill + +Pass `fragment_ids` to index only a subset of fragments today and backfill +the rest later. Fragments already covered by the existing same-name index +are skipped, and committed segments for other fragments are preserved +untouched: + +```python +# Index the first two fragments now. +create_scalar_index("s3://bucket/my_dataset", column="name", index_type="INVERTED", fragment_ids=[0, 1]) + +# After appending data, backfill only the new fragments. +create_scalar_index("s3://bucket/my_dataset", column="name", index_type="INVERTED", fragment_ids=[2, 3]) +``` + +Unknown fragment IDs and an empty list raise `ValueError`; duplicate IDs are +ignored. Indexing an empty dataset raises `ValueError` (no silent no-op), and +generated index names follow pylance's convention (`_idx`); before +committing, the driver validates that the built segments cover every +scheduled fragment exactly once and reference no fragments that a concurrent +compaction removed. A same-name index keeps its column and, on backfill, its +index type — mismatches are rejected by Lance's build/commit APIs. Backfill appends and never replaces existing segments, so +`replace=False` does not apply to it; a fully covered request is a no-op +that leaves the dataset version unchanged (rebuild with `replace=True` and +no `fragment_ids` instead). + > **Breaking change (from 0.5.0):** the `segmented` parameter was removed — > the distributed segment-index workflow is now the only code path, and the diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index f6dbe87..32aa78c 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -412,6 +412,7 @@ def create_scalar_index( fragment_group_size: int | None = None, num_partitions: int | None = None, max_concurrency: int | None = None, + fragment_ids: list[int] | None = None, **kwargs: Any, ) -> None: """Build a distributed scalar index using Daft's distributed execution. @@ -438,8 +439,12 @@ def create_scalar_index( replace: Whether to replace an existing index with the same name. Defaults to True, matching pylance. Replacement is atomic: the coordinator's segment commit retires the old index's overlapped - segments in the same transaction as the new ones. With - ``replace=False`` an existing index of the same name is rejected. + segments in the same transaction as the new ones (the index type + may change on a full rebuild). With ``replace=False`` + an existing index of the same name is rejected — except when + ``fragment_ids`` is given, which appends coverage for the + requested fragments (backfill) and never replaces existing + segments, so the ``replace`` flag does not apply. storage_options: Storage options for the dataset. version: Version of the dataset to use. asof: Timestamp to use for time travel queries. @@ -454,6 +459,17 @@ def create_scalar_index( greater than 1 enable additional parallelism on distributed runners; values <= 1 or None will use the default partitioning. max_concurrency: Maximum number of concurrent tasks to use for processing fragment batches. If None, Daft will use its default concurrency setting. Must be a positive integer. + fragment_ids: Optional subset of fragment IDs to index. Only the listed + fragments are scheduled; fragments already covered by an existing + same-name, same-column, same-type index are skipped and the + remaining ones are appended as new segments, preserving committed + segments untouched (``replace`` does not apply to this append + path, and a fully covered request is a no-op that leaves the + dataset version unchanged — pass ``replace=True`` without + ``fragment_ids`` to rebuild instead). Same-name indexes must keep + their column and, on backfill, their index type — Lance's build + and commit APIs reject mismatches. Duplicates are ignored; unknown + IDs and an empty list raise ``ValueError``. **kwargs: Additional keyword arguments forwarded to Lance's index segment creation API. Returns: @@ -496,6 +512,10 @@ def create_scalar_index( Refuse to overwrite an existing index: >>> daft_lance.create_scalar_index("s3://my-bucket/dataset/", column="title", replace=False) + + Index only half the fragments now and backfill the rest later: + >>> daft_lance.create_scalar_index("s3://my-bucket/dataset/", column="title", fragment_ids=[0, 1]) + >>> daft_lance.create_scalar_index("s3://my-bucket/dataset/", column="title", fragment_ids=[2, 3]) """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config @@ -525,6 +545,7 @@ def create_scalar_index( fragment_group_size=fragment_group_size, num_partitions=num_partitions, max_concurrency=max_concurrency, + fragment_ids=fragment_ids, **kwargs, ) diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index 47eec95..85d1509 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -102,6 +102,39 @@ def _existing_index_names(lance_ds: lance.LanceDataset) -> set[str]: return set() +def _existing_index_coverage(lance_ds: lance.LanceDataset, name: str) -> set[int] | None: + """Return the fragment IDs covered by an existing index, or None if absent. + + The coverage is the union of the fragment IDs covered by the index's + committed segments. Returns ``None`` when no index with that name exists. + When the manifest cannot be described but the name is visible through the + deprecated ``list_indices``, returns an empty set so callers still treat + the index as existing. Column/type compatibility is not checked here: + Lance's build and commit APIs reject incompatible combinations, and + duplicating those rules in string space has caused false rejections + before (e.g. 'LabelList' vs 'LABEL_LIST'). + """ + try: + descriptions = lance_ds.describe_indices() + except Exception: + logger.warning("describe_indices() failed; checking '%s' via list_indices", name, exc_info=True) + try: + if any(cast(dict[str, Any], idx).get("name") == name for idx in lance_ds.list_indices()): + return set() + except Exception: + pass + return None + + for desc in descriptions: + if desc.name != name: + continue + covered: set[int] = set() + for segment in desc.segments or []: + covered.update(segment.fragment_ids or ()) + return covered + return None + + def create_scalar_index_internal( lance_ds: lance.LanceDataset, open_context: DatasetOpenContext, @@ -113,6 +146,7 @@ def create_scalar_index_internal( fragment_group_size: int | None = None, num_partitions: int | None = None, max_concurrency: int | None = None, + fragment_ids: list[int] | None = None, **kwargs: Any, ) -> None: """Internal implementation of distributed scalar index creation. @@ -137,6 +171,13 @@ def create_scalar_index_internal( path raise ``ValueError`` instead of silently falling back to single-node Lance indexing — callers wanting single-node execution should call pylance directly. + + ``fragment_ids`` restricts the build to a subset of the dataset's + fragments. When the named index already exists, already-covered fragments + are skipped and only the remaining ones are built and appended; untouched + committed segments are preserved. This is the incremental backfill path + for newly appended fragments. Column/type compatibility of a same-name + index is validated by Lance's build/commit APIs, not duplicated here. """ if not column: raise ValueError("Column name cannot be empty") @@ -191,9 +232,40 @@ def create_scalar_index_internal( case _: pass - # Generate index name if not provided + # Generate index name if not provided (matches pylance's convention) if name is None: - name = f"{column}_{index_type.lower()}_idx" + name = f"{column}_idx" + + fragments = lance_ds.get_fragments() + available_fragment_ids = {fragment.fragment_id for fragment in fragments} + + # Validate and normalize the requested fragment subset, if any. + requested_fragment_ids: set[int] | None = None + if fragment_ids is not None: + if len(fragment_ids) == 0: + raise ValueError("fragment_ids must be a non-empty list of fragment IDs; pass None to index all fragments.") + unique_ids = list(dict.fromkeys(fragment_ids)) + duplicates = sorted({fid for fid in unique_ids if fragment_ids.count(fid) > 1}) + if duplicates: + logger.warning("Duplicate fragment_ids %s were given; each fragment is scheduled once.", duplicates) + unknown_ids = sorted(fid for fid in unique_ids if fid not in available_fragment_ids) + if unknown_ids: + raise ValueError( + f"fragment_ids {unknown_ids} do not exist in the dataset. " + f"Available fragment IDs: {sorted(available_fragment_ids)}" + ) + requested_fragment_ids = set(unique_ids) + + existing_coverage = _existing_index_coverage(lance_ds, name) + if existing_coverage is not None: + # Column/type compatibility of a same-name index is validated by + # Lance itself: the build API rejects a different column ("already + # exists with different fields") and the commit API rejects appending + # segments of a different type. Duplicating those rules here would be + # a string-space copy that can drift from the real type system (it + # already caused false rejections once), so no pre-check is kept. + if not replace and requested_fragment_ids is None: + raise ValueError(f"Index with name '{name}' already exists. Set replace=True to replace it.") # Replacement rides on Lance core's atomic overlap replacement: # commit_existing_index_segments retires committed segments whose fragments @@ -202,17 +274,33 @@ def create_scalar_index_internal( # longer overlap any live fragment cannot be retired this way, but normal # operations never produce them (compaction rewrites coverage; delete # retires fully-dead segments) and any that appear are healed by - # ``optimize_indices``. - handler_replace = False - if name in _existing_index_names(lance_ds): - if not replace: - raise ValueError(f"Index with name '{name}' already exists. Set replace=True to replace it.") - # Workers open the pinned snapshot where the same-name index still - # exists; building against that name requires replace=True. - handler_replace = True - - fragments = lance_ds.get_fragments() - fragment_ids_to_use = [fragment.fragment_id for fragment in fragments] + # ``optimize_indices``. Workers open the pinned snapshot where a + # same-name index still exists; building against that name always + # requires replace=True. + handler_replace = existing_coverage is not None + if existing_coverage is not None and requested_fragment_ids is not None: + # Incremental backfill: skip fragments already covered by committed + # segments; only the remainder is built and appended (non-overlapping + # segments are appended, not swapped). + covered = existing_coverage & available_fragment_ids + already_covered = requested_fragment_ids & covered + to_build = requested_fragment_ids - covered + if already_covered: + logger.info( + "Fragments %s are already covered by index '%s'; skipping them.", + sorted(already_covered), + name, + ) + if not to_build: + logger.info("All requested fragments are already covered by index '%s'; nothing to build.", name) + return + requested_fragment_ids = to_build + + if requested_fragment_ids is not None: + fragments = [fragment for fragment in fragments if fragment.fragment_id in requested_fragment_ids] + fragment_ids_to_use = sorted( + requested_fragment_ids if requested_fragment_ids is not None else (f.fragment_id for f in fragments) + ) # Adjust fragment grouping size if fragment_group_size is None: @@ -232,8 +320,7 @@ def create_scalar_index_internal( # Configure maximum concurrency for fragment batches if not fragment_data: - logger.info("No fragments found for dataset at %s; skipping scalar index creation.", open_context.uri) - return + raise ValueError(f"Dataset at {open_context.uri} contains no fragments") logger.info( "Starting distributed scalar index creation: column=%s, type=%s, name=%s, fragment_group_size=%s, max_concurrency=%s", @@ -250,7 +337,7 @@ def create_scalar_index_internal( index_type=index_type, name=name, fragment_data=fragment_data, - fragment_ids_to_use=fragment_ids_to_use, + expected_fragment_ids=fragment_ids_to_use, num_partitions=num_partitions, max_concurrency=max_concurrency, handler_replace=handler_replace, @@ -265,7 +352,7 @@ def _create_segmented_index( index_type: str, name: str, fragment_data: list[dict[str, list[int]]], - fragment_ids_to_use: list[int], + expected_fragment_ids: list[int] | None = None, num_partitions: int | None, max_concurrency: int | None, handler_replace: bool = False, @@ -316,6 +403,7 @@ def _create_segmented_index( # Reload dataset to pick up the latest version (segment files were written # by workers against the version that was current at their invocation time). lance_ds = open_context.open_latest() + _validate_segments_against_manifest(lance_ds, index_metas, expected_fragment_ids) logger.info( "Collected %d index segments; committing as segmented index %s", @@ -325,3 +413,54 @@ def _create_segmented_index( lance_ds.commit_existing_index_segments(name, column, index_metas) logger.info("Segmented index %s committed successfully", name) + + +def _validate_segments_against_manifest( + lance_ds: lance.LanceDataset, + index_metas: list[lance.Index | lance.indices.IndexSegment], + expected_fragment_ids: list[int] | None = None, +) -> None: + """Validate worker-built segments before the commit. + + Three checks, each failing loudly instead of committing a broken index: + + - Dead fragments: a segment references a fragment ID that no longer exists + in the manifest (a concurrent compaction rewrote it while the index was + being built). Lance would still accept the commit and the index would + permanently reference dead fragment IDs. + - Overlapping coverage: two segments cover the same fragment. Every + fragment must be covered exactly once. + - Incomplete coverage (when ``expected_fragment_ids`` is given): the union + of the segments' coverage must equal the scheduled fragment set, so a + silently lost worker result cannot produce a partial index. + """ + live_fragment_ids = {fragment.fragment_id for fragment in lance_ds.get_fragments()} + covered: set[int] = set() + duplicate_ids: set[int] = set() + dead_ids: set[int] = set() + for meta in index_metas: + for fragment_id in getattr(meta, "fragment_ids", None) or (): + if fragment_id in covered: + duplicate_ids.add(fragment_id) + covered.add(fragment_id) + if fragment_id not in live_fragment_ids: + dead_ids.add(fragment_id) + if dead_ids: + raise RuntimeError( + f"Cannot commit index segments: fragments {sorted(dead_ids)} no longer exist in the " + "dataset (they were most likely rewritten by a concurrent compaction while the index " + "was being built). Re-run the index build against the current dataset version." + ) + if duplicate_ids: + raise RuntimeError( + f"Cannot commit index segments: fragments {sorted(duplicate_ids)} are covered by more " + "than one segment; every fragment must be covered exactly once." + ) + if expected_fragment_ids is not None: + missing = set(expected_fragment_ids) - covered + if missing: + raise RuntimeError( + f"Cannot commit index segments: fragments {sorted(missing)} were scheduled but are " + "not covered by any built segment (a worker result was likely lost). Re-run the " + "index build." + ) diff --git a/tests/io/lance/test_lance_distributed_index_types.py b/tests/io/lance/test_lance_distributed_index_types.py index 2855801..92eb207 100644 --- a/tests/io/lance/test_lance_distributed_index_types.py +++ b/tests/io/lance/test_lance_distributed_index_types.py @@ -50,7 +50,7 @@ def test_new_type_builds_distributed_with_full_metadata( described = lance.dataset(uri).describe_indices() assert len(described) == 1 desc = described[0] - assert desc.name == f"{column}_{index_type.lower()}_idx" + assert desc.name == f"{column}_idx" # The whole point of the segment workflow: metadata must be complete, not "Unknown". assert desc.index_type == expected_type assert desc.type_url != "" diff --git a/tests/io/lance/test_lance_partial_fragment_index.py b/tests/io/lance/test_lance_partial_fragment_index.py new file mode 100644 index 0000000..602e802 --- /dev/null +++ b/tests/io/lance/test_lance_partial_fragment_index.py @@ -0,0 +1,400 @@ +from __future__ import annotations + +import warnings +from pathlib import Path +from typing import Any, cast + +import lance +import pytest + +from daft.dependencies import pa +from daft_lance import create_scalar_index +from daft_lance.lance_scalar_index import SegmentedFragmentIndexHandler + +warnings.filterwarnings("ignore", category=DeprecationWarning, module="lance") + + +def _make_dataset(path: Path, n_rows: int = 80, rows_per_file: int = 20) -> str: + table = pa.table( + { + "id": list(range(n_rows)), + "name": [f"name-{i % 8}" for i in range(n_rows)], + } + ) + lance.write_dataset(table, str(path), mode="create", max_rows_per_file=rows_per_file) + return str(path) + + +def _fragment_ids(uri: str) -> list[int]: + return sorted(f.fragment_id for f in lance.dataset(uri).get_fragments()) + + +def _covered_fragments(uri: str, index_name: str) -> set[int]: + for desc in lance.dataset(uri).describe_indices(): + if desc.name == index_name: + covered: set[int] = set() + for segment in desc.segments: + covered.update(segment.fragment_ids) + return covered + raise AssertionError(f"index {index_name} not found") + + +def _sorted_rows(uri: str, predicate: str) -> dict[str, list[Any]]: + table = lance.dataset(uri).scanner(filter=predicate).to_table().sort_by("id") + return table.to_pydict() + + +def test_partial_build_then_backfill_matches_full_build(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "incremental.lance") + full_uri = _make_dataset(tmp_path / "full.lance") + + # Index the first half of the fragments. + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 1]) + assert _covered_fragments(uri, "name_idx") == {0, 1} + # The partial index is already query-correct: uncovered fragments fall + # back to scans, including predicates that hit only uncovered fragments + # and zero-hit predicates. + for predicate in ["name = 'name-3'", "id >= 40", "name = 'nope'"]: + assert _sorted_rows(uri, predicate) == _sorted_rows(full_uri, predicate), predicate + + # Backfill the remaining fragments under the same index name. + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 1, 2, 3]) + assert _covered_fragments(uri, "name_idx") == {0, 1, 2, 3} + # Backfill appends: the original {0,1} segment is preserved, not rebuilt. + segments = [sorted(s.fragment_ids) for d in lance.dataset(uri).describe_indices() for s in d.segments] + assert segments == [[0, 1], [2, 3]] + + # Reference: one-shot full build on an identical dataset. + create_scalar_index(full_uri, column="name", index_type="INVERTED") + + for predicate in ["name = 'name-3'", "name in ('name-0', 'name-7')", "id >= 40"]: + assert _sorted_rows(uri, predicate) == _sorted_rows(full_uri, predicate), predicate + + +def test_partial_build_btree(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "btree.lance") + + create_scalar_index(uri, column="id", index_type="BTREE", fragment_ids=[0, 2]) + assert _covered_fragments(uri, "id_idx") == {0, 2} + + create_scalar_index(uri, column="id", index_type="BTREE", fragment_ids=_fragment_ids(uri)) + assert _covered_fragments(uri, "id_idx") == {0, 1, 2, 3} + + assert lance.dataset(uri).scanner(filter="id = 42").to_table().num_rows == 1 + + +def test_backfill_after_appending_new_fragments(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "append.lance", n_rows=40) + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 1]) + covered_before = _covered_fragments(uri, "name_idx") + + # Append data; Lance mints new fragment IDs for it. + extra = pa.table({"id": list(range(100, 140)), "name": [f"name-{100 + i}" for i in range(40)]}) + lance.write_dataset(extra, uri, mode="append", max_rows_per_file=20) + + new_fragments = set(_fragment_ids(uri)) - covered_before + assert new_fragments, "append must add fragments" + + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=sorted(new_fragments)) + assert _covered_fragments(uri, "name_idx") == set(_fragment_ids(uri)) + + found = lance.dataset(uri).scanner(filter="name = 'name-103'").to_table().num_rows + assert found == 1 + + +def test_duplicate_fragment_ids_are_deduplicated(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "dupes.lance") + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 0, 1, 1, 1]) + assert _covered_fragments(uri, "name_idx") == {0, 1} + + +def test_unknown_fragment_ids_raise(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "oob.lance") + with pytest.raises(ValueError, match=r"\[99\].*do not exist"): + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 99]) + + +def test_empty_fragment_ids_raise(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "empty.lance") + with pytest.raises(ValueError, match="non-empty"): + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[]) + + +def test_backfill_with_fully_covered_fragments_is_noop(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "noop.lance") + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 1]) + version_before = lance.dataset(uri).version + + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 1]) + + assert lance.dataset(uri).version == version_before + assert _covered_fragments(uri, "name_idx") == {0, 1} + + +def test_name_reuse_without_fragment_ids_replaces_atomically(tmp_path: Path) -> None: + """Default replace=True rebuilds the whole index atomically (no error).""" + uri = _make_dataset(tmp_path / "reuse.lance") + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0]) + version_before = lance.dataset(uri).version + + create_scalar_index(uri, column="name", index_type="INVERTED") + + latest = lance.dataset(uri) + assert latest.version == version_before + 1 + segments = [sorted(s.fragment_ids) for d in latest.describe_indices() for s in d.segments] + assert segments == [[0, 1, 2, 3]] # old partial segment retired by the atomic swap + + +def test_backfill_different_column_rejected(tmp_path: Path) -> None: + """A different column under the same name is rejected by Lance's build API. + + daft-lance keeps no duplicated pre-check; the error surfaces from the + build API itself. + """ + uri = _make_dataset(tmp_path / "other_column.lance") + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0]) + with pytest.raises(Exception, match="(?i)different fields|already exists"): + create_scalar_index(uri, column="id", index_type="BTREE", name="name_idx", fragment_ids=[1]) + + +def test_backfill_mixed_index_type_rejected(tmp_path: Path) -> None: + """Appending segments of a different type is rejected by Lance's commit API. + + daft-lance keeps no duplicated pre-check; the manifest stays intact. + """ + uri = _make_dataset(tmp_path / "mixed_type.lance") + create_scalar_index(uri, column="name", index_type="INVERTED", name="shared_idx", fragment_ids=[0]) + with pytest.raises(ValueError, match="cannot change index 'shared_idx' from type"): + create_scalar_index(uri, column="name", index_type="BITMAP", name="shared_idx", fragment_ids=[2]) + # The manifest must still be describable and the original type intact. + ds = lance.dataset(uri) + assert [idx.name for idx in ds.describe_indices()] == ["shared_idx"] + + +def test_fragment_ids_rejected_for_unsupported_type(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "rtree.lance") + with pytest.raises(ValueError, match="Unsupported distributed index type 'RTREE'"): + create_scalar_index(uri, column="id", index_type="RTREE", fragment_ids=[0]) + assert lance.dataset(uri).describe_indices() == [] + + +def test_bitmap_partial_backfill_matches_full_build(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "bitmap_partial.lance") + full_uri = _make_dataset(tmp_path / "bitmap_full.lance") + + create_scalar_index(uri, column="id", index_type="BITMAP", fragment_ids=[0, 1]) + create_scalar_index(uri, column="id", index_type="BITMAP", fragment_ids=_fragment_ids(uri)) + create_scalar_index(full_uri, column="id", index_type="BITMAP") + + assert _covered_fragments(uri, "id_idx") == {0, 1, 2, 3} + for predicate in ["id = 42", "id in (0, 39, 79)", "id >= 60"]: + assert _sorted_rows(uri, predicate) == _sorted_rows(full_uri, predicate), predicate + + +def test_btree_partial_backfill_matches_full_build(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "btree_equiv.lance") + full_uri = _make_dataset(tmp_path / "btree_full.lance") + + create_scalar_index(uri, column="id", index_type="BTREE", fragment_ids=[0, 1]) + create_scalar_index(uri, column="id", index_type="BTREE", fragment_ids=_fragment_ids(uri)) + create_scalar_index(full_uri, column="id", index_type="BTREE") + + for predicate in ["id = 7", "id in (20, 59)", "id >= 40"]: + assert _sorted_rows(uri, predicate) == _sorted_rows(full_uri, predicate), predicate + + +def test_compacted_fragment_conflict_surfaces_as_error(tmp_path: Path) -> None: + """The coordinator must refuse to commit segments whose fragments died. + + If a compaction rewrites the fragments while workers are building, Lance + would still accept the commit and the index would reference dead fragment + IDs forever. + """ + import pickle + + from daft_lance.lance_scalar_index import _validate_segments_against_manifest + + uri = _make_dataset(tmp_path / "conflict.lance") + ds = lance.dataset(uri) + open_context = type( + "Ctx", + (), + { + "uri": uri, + "open_pinned": lambda self: ds, + "open_latest": lambda self: lance.dataset(uri), + }, + )() + + handler = SegmentedFragmentIndexHandler( + open_context=open_context, + column="name", + index_type="INVERTED", + name="name_idx", + ) + segment = pickle.loads(handler([0, 1])) + + # Compact away the fragments the segment was built against. + lance.dataset(uri).optimize.compact_files() + + latest = lance.dataset(uri) + with pytest.raises(RuntimeError, match="no longer exist"): + _validate_segments_against_manifest(latest, [segment]) + + +def test_label_list_backfill_and_rebuild_roundtrip(tmp_path: Path) -> None: + """Regression: pylance reports LABEL_LIST as 'LabelList'. + + The type guard must normalize separators so backfill and plain rebuild + both work. + """ + uri = str(tmp_path / "labels.lance") + lance.write_dataset( + pa.table({"tags": [[f"t{i % 4}", "common"] for i in range(80)]}), + uri, + mode="create", + max_rows_per_file=20, + ) + + create_scalar_index(uri, column="tags", index_type="LABEL_LIST", fragment_ids=[0, 1]) + create_scalar_index(uri, column="tags", index_type="LABEL_LIST", fragment_ids=[2, 3]) + assert _covered_fragments(uri, "tags_idx") == {0, 1, 2, 3} + + # Plain rebuild (default replace=True) also still works. + create_scalar_index(uri, column="tags", index_type="LABEL_LIST") + assert _covered_fragments(uri, "tags_idx") == {0, 1, 2, 3} + + +def test_full_rebuild_may_change_type(tmp_path: Path) -> None: + """A full rebuild may change the index type via the atomic swap. + + A different column under the same name is rejected by Lance's build API. + """ + uri = _make_dataset(tmp_path / "swap.lance") + create_scalar_index(uri, column="name", index_type="INVERTED", name="shared_idx", fragment_ids=[0]) + + create_scalar_index(uri, column="name", index_type="BTREE", name="shared_idx") + + described = lance.dataset(uri).describe_indices() + assert len(described) == 1 + assert described[0].index_type == "BTree" + assert lance.dataset(uri).scanner(filter="name = 'name-3'").to_table().num_rows == 10 + + with pytest.raises(Exception, match="(?i)different fields|already exists"): + create_scalar_index(uri, column="id", index_type="BTREE", name="shared_idx") + + +def test_replace_false_with_fragment_ids_still_backfills(tmp_path: Path) -> None: + """replace=False does not apply to backfill. + + Backfill appends and never replaces, so the flag is irrelevant on that + path (documented interaction). + """ + uri = _make_dataset(tmp_path / "rf.lance") + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 1]) + version_before = lance.dataset(uri).version + + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[2, 3], replace=False) + + latest = lance.dataset(uri) + assert latest.version == version_before + 1 + segments = [sorted(s.fragment_ids) for d in latest.describe_indices() for s in d.segments] + assert segments == [[0, 1], [2, 3]] + + +def test_backfill_with_multiple_segments_per_task(tmp_path: Path) -> None: + """fragment_group_size=1 splits the backfill across workers. + + Multiple new segments are appended alongside the preserved old one in a + single commit. + """ + uri = _make_dataset(tmp_path / "multi.lance") + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 1]) + + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[2, 3], fragment_group_size=1) + + segments = [sorted(s.fragment_ids) for d in lance.dataset(uri).describe_indices() for s in d.segments] + assert segments == [[0, 1], [2], [3]] + assert lance.dataset(uri).scanner(filter="name = 'name-3'").to_table().num_rows == 10 + + +def test_duplicate_fragment_ids_warn(caplog: pytest.LogCaptureFixture, tmp_path: Path) -> None: + """The de-duplication is user-visible via a log warning.""" + import logging + + uri = _make_dataset(tmp_path / "dupwarn.lance") + with caplog.at_level(logging.WARNING, logger="daft_lance.lance_scalar_index"): + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0, 0, 1]) + assert any("Duplicate fragment_ids" in r.message for r in caplog.records) + assert _covered_fragments(uri, "name_idx") == {0, 1} + + +def _guard_recorder(order: list[str], real_guard: Any, ds: Any, metas: Any, expected: Any) -> Any: + order.append("guard") + return real_guard(ds, metas, expected) + + +def _commit_recorder(order: list[str], real_commit: Any, ds: Any, *args: Any, **kwargs: Any) -> Any: + order.append("commit") + return real_commit(ds, *args, **kwargs) + + +def test_guard_runs_before_commit(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + """The manifest guard is on the commit path (before the commit call).""" + import daft_lance.lance_scalar_index as lsi + + order: list[str] = [] + real_guard = lsi._validate_segments_against_manifest + monkeypatch.setattr( + lsi, + "_validate_segments_against_manifest", + lambda ds, metas, expected=None: _guard_recorder(order, real_guard, ds, metas, expected), + ) + real_commit = lance.LanceDataset.commit_existing_index_segments + monkeypatch.setattr( + lance.LanceDataset, + "commit_existing_index_segments", + lambda self, *a, **k: _commit_recorder(order, real_commit, self, *a, **k), + ) + + uri = _make_dataset(tmp_path / "order.lance") + create_scalar_index(uri, column="name", index_type="INVERTED", fragment_ids=[0]) + + assert order == ["guard", "commit"] + + +def test_guard_rejects_overlapping_and_incomplete_coverage() -> None: + """The commit guard catches duplicate coverage and missing fragments.""" + from daft_lance.lance_scalar_index import _validate_segments_against_manifest + + class FakeSegment: + def __init__(self, fragment_ids: set[int]) -> None: + self.fragment_ids = fragment_ids + + class FakeFragments: + def __init__(self, ids: list[int]) -> None: + self._ids = ids + + def get_fragments(self) -> list[FakeFragment]: + return [FakeFragment(i) for i in self._ids] + + class FakeFragment: + def __init__(self, fragment_id: int) -> None: + self.fragment_id = fragment_id + + ds = FakeFragments([0, 1, 2, 3]) + + # Overlapping coverage: fragment 1 built by two workers. + with pytest.raises(RuntimeError, match="more than one segment"): + _validate_segments_against_manifest( + cast(Any, ds), [cast(Any, FakeSegment({0, 1})), cast(Any, FakeSegment({1, 2}))], [0, 1, 2] + ) + + # Incomplete coverage: fragment 3 scheduled but no segment covers it. + with pytest.raises(RuntimeError, match="not covered by any built segment"): + _validate_segments_against_manifest(cast(Any, ds), [cast(Any, FakeSegment({0, 1}))], [0, 1, 2, 3]) + + # Exact coverage passes. + _validate_segments_against_manifest( + cast(Any, ds), [cast(Any, FakeSegment({0, 1})), cast(Any, FakeSegment({2, 3}))], [0, 1, 2, 3] + ) diff --git a/tests/io/lance/test_lance_scalar_index.py b/tests/io/lance/test_lance_scalar_index.py index 54c5c34..bb24b6a 100644 --- a/tests/io/lance/test_lance_scalar_index.py +++ b/tests/io/lance/test_lance_scalar_index.py @@ -137,7 +137,7 @@ def test_build_distributed_index_search_functionality(self, multi_fragment_lance # populate Lance index details, so use list_indices() here. indices = updated_dataset.list_indices() index_names = [idx["name"] for idx in indices] - assert "text_inverted_idx" in index_names, f"Text index not found in {index_names}" + assert "text_idx" in index_names, f"Text index not found in {index_names}" # Test full-text search functionality search_term = "Python" @@ -546,17 +546,13 @@ def test_build_distributed_index_no_fragments(self, temp_dir): path = Path(temp_dir) / "empty_dataset.lance" dataset.write_lance(uri=path) - # Try to build index on empty dataset - create_scalar_index( - uri=path, - column="text", - index_type="INVERTED", - ) - - # Verify no index was created (since no data) - updated_dataset = lance.dataset(path) - indices = updated_dataset.list_indices() - assert len(indices) == 0, f"Expected no indices for empty dataset, got {len(indices)}" + # Building an index on an empty dataset fails loudly (no silent no-op). + with pytest.raises(ValueError, match="contains no fragments"): + create_scalar_index( + uri=path, + column="text", + index_type="INVERTED", + ) def test_build_distributed_index_zonemap_type(self, temp_dir): """Test building ZONEMAP index distributed on a numeric column.""" @@ -781,6 +777,8 @@ def __init__(self, fragment_ids: set[int]) -> None: class ExistingIndex: name = "flag_bitmap_idx" + index_type = "Bitmap" + field_names = ["flag"] segments = [FakeSegment({0, 1})] class FakeFragment: