diff --git a/README.md b/README.md index 57d5634..7517ef7 100644 --- a/README.md +++ b/README.md @@ -24,12 +24,32 @@ compact_files("s3://bucket/my_dataset") ### Scalar Indexing +Every supported index type is built distributed: Daft workers build +independent index segments and the coordinator commits them atomically with +complete metadata. Supported types: `BITMAP`, `BTREE`, `INVERTED`, `FTS`, +`ZONEMAP`, `NGRAM`, `LABEL_LIST`, `BLOOMFILTER`. + ```python from daft_lance import create_scalar_index create_scalar_index("s3://bucket/my_dataset", column="name", index_type="INVERTED") +create_scalar_index("s3://bucket/my_dataset", column="ts", index_type="ZONEMAP") ``` +`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. + +> **Breaking change (from 0.5.0):** the `segmented` parameter was removed — +> the distributed segment-index workflow is now the only code path, and the +> single-node fallbacks are gone. `replace` now defaults to `True` (matching +> pylance). Indexes created by older versions of +> `create_scalar_index` on `INVERTED` columns may carry empty index metadata +> (see #69); rebuilding them with `replace=True` records full metadata. + ### Column Merging ```python diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 15659d3..f6dbe87 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -400,7 +400,7 @@ def create_scalar_index( column: str, index_type: str = "INVERTED", name: str | None = None, - replace: bool = False, + replace: bool = True, storage_options: dict[str, Any] | None = None, version: int | str | None = None, asof: str | None = None, @@ -412,14 +412,13 @@ def create_scalar_index( fragment_group_size: int | None = None, num_partitions: int | None = None, max_concurrency: int | None = None, - segmented: bool = False, **kwargs: Any, ) -> None: """Build a distributed scalar index using Daft's distributed execution. This function distributes the index building process across multiple Daft workers, - with each worker building indices for a subset of fragments. The indices are then - merged and committed as a single index. + with each worker building independent index segments for a subset of fragments + that the coordinator commits atomically, recording complete index metadata. Args: uri: The URI of the Lance table (supports remote URLs to object stores such as `s3://` or `gs://`) @@ -430,20 +429,17 @@ def create_scalar_index( namespace_properties: Properties for connecting to the namespace, e.g. {"root": "/data"} for "dir" or {"uri": "http://host:port"} for "rest". column: Column name to index - index_type: Type of index to build. - For distributed segmented execution this supports "BITMAP", "BTREE", "INVERTED", and "FTS". - Other scalar index types supported by Lance (for example "NGRAM", "ZONEMAP", - "LABEL_LIST", "BLOOMFILTER") are passed directly to - ``LanceDataset.create_scalar_index(...)``. + index_type: Type of index to build. Built distributed for + "BITMAP", "BTREE", "INVERTED", "FTS", "ZONEMAP", "NGRAM", + "LABEL_LIST", and "BLOOMFILTER". Other types raise ``ValueError``; + for those call pylance directly + (``lance.dataset(uri).create_scalar_index(...)``). name: Name of the index (generated if None). - replace: Whether to replace an existing index with the same name. Defaults to False. - This is only supported by scalar index types that are passed directly to - ``LanceDataset.create_scalar_index(...)``. Segmented BITMAP/BTREE/INVERTED/FTS - indexes use Lance's public segmented-index commit API, which does not - currently expose atomic replacement, so existing index names are rejected. - For BITMAP indexing with the default ``segmented=False`` behavior, - ``replace=True`` on an existing index falls back to Lance's scalar - replacement path to preserve existing API behavior. + 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. storage_options: Storage options for the dataset. version: Version of the dataset to use. asof: Timestamp to use for time travel queries. @@ -458,18 +454,14 @@ 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. - segmented: If True, force the segmented index workflow where each worker builds - a fully independent index segment and the coordinator commits them via - ``commit_existing_index_segments``. ``"FTS"`` is normalized to Lance's - inverted full-text index. If False, scalar index creation uses the legacy - partitioned workflow or Lance's direct ``create_scalar_index`` path. - **kwargs: Additional keyword arguments forwarded to the selected Lance index creation API. + **kwargs: Additional keyword arguments forwarded to Lance's index segment creation API. Returns: None Raises: - ValueError: If input parameters are invalid (e.g., empty column name, non-existent column, invalid index type, etc.) + ValueError: If input parameters are invalid (e.g., empty column name, non-existent + column, unsupported index type, or an existing index name with ``replace=False``) TypeError: If column type is incompatible with the chosen ``index_type`` RuntimeError: If index building fails (e.g., version compatibility issues, commit failures) ImportError: If lance package is not available @@ -493,17 +485,16 @@ def create_scalar_index( ... "s3://my-bucket/dataset/", column="price", index_type="BTREE", name="price_idx" ... ) - Create a segmented BTREE index (supports describe_indices): - >>> daft_lance.create_scalar_index( - ... "s3://my-bucket/dataset/", column="price", index_type="BTREE", segmented=True - ... ) + Create a distributed ZONEMAP or NGRAM index (newly distributed): + >>> daft_lance.create_scalar_index("s3://my-bucket/dataset/", column="ts", index_type="ZONEMAP") + >>> daft_lance.create_scalar_index("s3://my-bucket/dataset/", column="doc", index_type="NGRAM") Create an index with custom fragment grouping and partitioning: >>> daft_lance.create_scalar_index( ... "s3://my-bucket/dataset/", column="description", fragment_group_size=8, num_partitions=16 ... ) - Create an index without replacing existing ones: + Refuse to overwrite an existing index: >>> daft_lance.create_scalar_index("s3://my-bucket/dataset/", column="title", replace=False) """ io_config = context.get_context().daft_planning_config.default_io_config if io_config is None else io_config @@ -534,7 +525,6 @@ def create_scalar_index( fragment_group_size=fragment_group_size, num_partitions=num_partitions, max_concurrency=max_concurrency, - segmented=segmented, **kwargs, ) diff --git a/daft_lance/lance_scalar_index.py b/daft_lance/lance_scalar_index.py index d29fd23..47eec95 100644 --- a/daft_lance/lance_scalar_index.py +++ b/daft_lance/lance_scalar_index.py @@ -2,7 +2,6 @@ import logging import pickle -import uuid from typing import TYPE_CHECKING, Any, cast import daft @@ -18,54 +17,11 @@ logger = logging.getLogger(__name__) -# Segmented index types whose worker-built segments must be merged before commit. -MERGED_SEGMENTED_INDEX_TYPES = {"BITMAP", "INVERTED"} - - -class FragmentIndexHandler: - """Handler for distributed scalar index creation on fragment batches.""" - - def __init__( - self, - open_context: DatasetOpenContext, - column: str, - index_type: str, - name: str, - fragment_uuid: str, - replace: bool, - **kwargs: Any, - ) -> None: - self.open_context = open_context - self.column = column - self.index_type = index_type - self.name = name - self.fragment_uuid = fragment_uuid - self.replace = replace - self.kwargs = kwargs - 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_ids: list[int]) -> bool: - """Process a batch of fragment IDs for scalar index creation.""" - logger.info( - "Building distributed scalar index for fragments %s using create_scalar_index", - fragment_ids, - ) - - self._dataset().create_scalar_index( - column=self.column, - index_type=self.index_type, # type: ignore[arg-type] - name=self.name, - replace=self.replace, - index_uuid=self.fragment_uuid, - fragment_ids=fragment_ids, - **self.kwargs, - ) - return True +# Scalar index types built with the distributed segment workflow. pylance 11 +# reports every scalar index type as segment-native; RTREE additionally +# requires GeoArrow extension columns and stays unsupported until the test +# suite can create those columns. +DISTRIBUTED_INDEX_TYPES = frozenset({"BTREE", "BITMAP", "INVERTED", "ZONEMAP", "NGRAM", "LABEL_LIST", "BLOOMFILTER"}) class SegmentedFragmentIndexHandler: @@ -84,12 +40,14 @@ def __init__( column: str, index_type: str, name: str, + replace: bool = False, **kwargs: Any, ) -> None: self.open_context = open_context self.column = column self.index_type = index_type self.name = name + self.replace = replace self.kwargs = kwargs self._lance_ds: lance.LanceDataset | None = None @@ -98,7 +56,7 @@ def _dataset(self) -> lance.LanceDataset: self._lance_ds = self.open_context.open_pinned() return self._lance_ds - def __call__(self, fragment_ids: list[int], shard_id: int | None = None) -> bytes: + def __call__(self, fragment_ids: list[int]) -> bytes: """Build an independent index segment and return its pickled metadata.""" logger.info( "Building segmented index segment for fragments %s (column=%s, type=%s)", @@ -107,20 +65,22 @@ def __call__(self, fragment_ids: list[int], shard_id: int | None = None) -> byte self.index_type, ) segment_kwargs = self.kwargs.copy() - if self.index_type == "BITMAP" and shard_id is not None: - # Lance's BITMAP segment builder needs a stable shard number to - # distinguish independently-built bitmap segments before merge/commit. - segment_kwargs["shard_id"] = shard_id - - # Create one uncommitted index segment. ``pylance 8.0.0`` supports - # scalar index segments through this public API. Segment creation always - # uses ``replace=False`` because replacement, if supported, must happen - # in the final manifest commit rather than independently in each worker. + + # Create one uncommitted index segment. Segment creation normally + # uses ``replace=False`` because replacement must happen in the final + # manifest commit rather than independently in each worker. The one + # exception is a replace=True rebuild: the worker's pinned snapshot + # still contains the same-named index (whether or not the driver + # dropped it), so Lance rejects building against that name with + # ``replace=False``. The driver opts workers into ``replace=True`` + # in exactly that case, and the coordinator's commit lands the new + # segments atomically — retiring overlapped old segments, or creating + # the index fresh on a post-drop manifest. index_meta = self._dataset().create_index_uncommitted( column=self.column, index_type=self.index_type, name=self.name, - replace=False, + replace=self.replace, train=True, fragment_ids=fragment_ids, **segment_kwargs, @@ -149,11 +109,10 @@ def create_scalar_index_internal( column: str, index_type: str = "INVERTED", name: str | None = None, - replace: bool = False, + replace: bool = True, fragment_group_size: int | None = None, num_partitions: int | None = None, max_concurrency: int | None = None, - segmented: bool = False, **kwargs: Any, ) -> None: """Internal implementation of distributed scalar index creation. @@ -162,16 +121,35 @@ def create_scalar_index_internal( ``open_context`` is the serializable handle workers reopen from and the single source of uri, storage options and namespace kwargs. - When ``segmented=True``, ``BITMAP``, ``BTREE``, and ``INVERTED`` use - Lance's public segment-index workflow: each worker builds a fully - independent index segment, and the coordinator commits them atomically with - ``commit_existing_index_segments``. ``FTS`` is normalized to ``INVERTED`` - (same Lance index); see Lance Rust/Python bindings: ``INVERTED`` and - ``FTS`` map to the same inverted full-text index type. + Every supported index type is built with the distributed segment-index + workflow: each worker builds a fully independent index segment with + ``create_index_uncommitted``, and the coordinator commits them atomically + with ``commit_existing_index_segments``, which records complete index + metadata (no more empty ``index_details``). ``FTS`` is normalized to + ``INVERTED`` (same Lance index). + + ``replace=True`` (the default) relies on Lance core's atomic overlap + replacement: ``commit_existing_index_segments`` removes committed segments + whose fragments overlap the incoming ones in the same CreateIndex + transaction, so a full-coverage rebuild swaps the old index in one + transaction. ``replace=False`` refuses to touch an existing index; the + default matches pylance's own ``replace`` default. Types without a distributed + path raise ``ValueError`` instead of silently falling back to single-node + Lance indexing — callers wanting single-node execution should call pylance + directly. """ if not column: raise ValueError("Column name cannot be empty") + if "segmented" in kwargs: + # Removed parameter: **kwargs would otherwise forward it to Lance, + # which fails deep inside a worker with a confusing index-parameter + # error instead of at the API boundary. + raise TypeError( + "The 'segmented' parameter was removed: the distributed segment-index " + "workflow is now the only code path. Remove the argument." + ) + index_type = index_type.upper() if index_type == "FTS": logger.info( @@ -179,6 +157,13 @@ def create_scalar_index_internal( ) index_type = "INVERTED" + if index_type not in DISTRIBUTED_INDEX_TYPES: + raise ValueError( + f"Unsupported distributed index type '{index_type}'. Supported types: " + f"{sorted(DISTRIBUTED_INDEX_TYPES)} (plus 'FTS'). For other types call pylance " + f"directly: lance.dataset().create_scalar_index(...)." + ) + # Validate column exists and has correct type try: field = lance_ds.schema.field(column) @@ -186,7 +171,8 @@ def create_scalar_index_internal( available_columns = [field.name for field in lance_ds.schema] raise ValueError(f"Column '{column}' not found. Available: {available_columns}") from e - # Check column type + # Check column type for the types with an obvious Python-side rule; the + # rest are validated by Lance during the distributed build. value_type = field.type if pa.types.is_list(field.type) or pa.types.is_large_list(field.type): value_type = field.type.value_type @@ -202,70 +188,29 @@ def create_scalar_index_internal( and not pa.types.is_string(value_type) ): raise TypeError(f"Column {column} must be numeric or string type for BTREE index, got {value_type}") - case "BITMAP": - # BITMAP supports multiple physical Arrow types depending on the - # Lance release. Leave final type validation to Lance rather than - # duplicating a narrower Python-side allowlist here. - pass case _: - logger.warning( - "Distributed indexing currently only supports 'BITMAP', 'INVERTED', and 'BTREE' index types, not '%s'. So we are falling back to single-threaded index creation.", - index_type, - ) - lance_ds.create_scalar_index( - column=column, - index_type=index_type, # type: ignore[arg-type] - name=name, - replace=replace, - **kwargs, - ) - return - - if index_type == "BITMAP" and not segmented: - logger.info( - "Falling back to Lance scalar index creation for non-segmented BITMAP index %s.", - name, - ) - lance_ds.create_scalar_index( - column=column, - index_type=index_type, - name=name, - replace=replace, - **kwargs, - ) - return + pass # Generate index name if not provided if name is None: name = f"{column}_{index_type.lower()}_idx" - # Handle replace parameter - check for existing index with same name - if not replace or segmented: - existing_names = _existing_index_names(lance_ds) - if name in existing_names and segmented: - raise ValueError( - f"Index with name '{name}' already exists and cannot atomically replace existing index " - "with Lance's public segmented index API. Drop the existing index first or use a different name." - ) - if name in existing_names: + # Replacement rides on Lance core's atomic overlap replacement: + # commit_existing_index_segments retires committed segments whose fragments + # overlap the incoming ones in the same CreateIndex transaction, so a + # full-coverage rebuild swaps the old index atomically. Segments that no + # 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 - if index_type == "BTREE" and not segmented: - logger.info( - "Falling back to Lance scalar index creation for non-segmented %s index %s.", - index_type, - name, - ) - lance_ds.create_scalar_index( - column=column, - index_type=index_type, - name=name, - replace=replace, - **kwargs, - ) - return - - # Get available fragment IDs to use fragments = lance_ds.get_fragments() fragment_ids_to_use = [fragment.fragment_id for fragment in fragments] @@ -291,43 +236,26 @@ def create_scalar_index_internal( return logger.info( - "Starting distributed scalar index creation: column=%s, type=%s, name=%s, fragment_group_size=%s, max_concurrency=%s, segmented=%s", + "Starting distributed scalar index creation: column=%s, type=%s, name=%s, fragment_group_size=%s, max_concurrency=%s", column, index_type, name, fragment_group_size, max_concurrency, - segmented, ) - # Use segment-index creation for Lance scalar index types that expose the - # public uncommitted segment API. The legacy path is kept as a fallback for - # older/unsupported distributed scalar index types. - if segmented: - _create_segmented_index( - open_context=open_context, - column=column, - index_type=index_type, - name=name, - fragment_data=fragment_data, - fragment_ids_to_use=fragment_ids_to_use, - num_partitions=num_partitions, - max_concurrency=max_concurrency, - **kwargs, - ) - else: - _create_partitioned_index( - open_context=open_context, - column=column, - index_type=index_type, - name=name, - replace=replace, - fragment_data=fragment_data, - fragment_ids_to_use=fragment_ids_to_use, - num_partitions=num_partitions, - max_concurrency=max_concurrency, - **kwargs, - ) + _create_segmented_index( + open_context=open_context, + column=column, + index_type=index_type, + name=name, + fragment_data=fragment_data, + fragment_ids_to_use=fragment_ids_to_use, + num_partitions=num_partitions, + max_concurrency=max_concurrency, + handler_replace=handler_replace, + **kwargs, + ) def _create_segmented_index( @@ -340,6 +268,7 @@ def _create_segmented_index( fragment_ids_to_use: list[int], num_partitions: int | None, max_concurrency: int | None, + handler_replace: bool = False, **kwargs: Any, ) -> None: """Segmented index workflow: each worker builds an independent segment. @@ -347,7 +276,15 @@ def _create_segmented_index( Workers call Lance's uncommitted index segment API, pickle the returned ``lance.Index`` metadata so it can traverse Daft serialisation boundaries, and return it. The coordinator unpickles all segments and commits them - atomically via ``commit_existing_index_segments``. + as-is via ``commit_existing_index_segments``: committed segments whose + fragments overlap the incoming ones are retired in the same transaction + (atomic replacement), and non-overlapping ones are appended. Multi-segment + indexes are fully functional without merging (verified: split segments are + loaded and pruned at query time with identical scores), so no merge is + performed; compaction is left to ``optimize_indices``. When + ``handler_replace`` is set the workers' pinned snapshot still contains a + same-named index; they must build with ``replace=True`` for Lance to + accept the name. """ handler_cls = daft.cls( SegmentedFragmentIndexHandler, @@ -358,34 +295,17 @@ def _create_segmented_index( column=column, index_type=index_type, name=name, + replace=handler_replace, **kwargs, ) - segment_data: list[dict[str, Any]] - if index_type == "BITMAP": - # Give each worker-built BITMAP segment a stable Lance shard id. This - # is separate from Lance fragment ids: shard_id identifies the bitmap - # segment, while fragment_ids identify the rows covered by that segment. - segment_data = [ - { - **group, - "shard_id": shard_id, - } - for shard_id, group in enumerate(fragment_data) - ] - else: - segment_data = fragment_data - with execution_config_ctx(maintain_order=False): if num_partitions is not None and num_partitions > 1: - df = from_pylist(segment_data).repartition(num_partitions) + df = from_pylist(fragment_data).repartition(num_partitions) else: - df = from_pylist(segment_data) + df = from_pylist(fragment_data) - if index_type == "BITMAP": - df = df.select(handler(df["fragment_ids"], df["shard_id"]).alias("index_meta")) - else: - df = df.select(handler(df["fragment_ids"]).alias("index_meta")) + df = df.select(handler(df["fragment_ids"]).alias("index_meta")) collected = df.collect() # Deserialise the Index metadata returned by each worker. @@ -396,7 +316,6 @@ 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() - index_metas = _prepare_index_segments_for_commit(lance_ds, index_type, index_metas) logger.info( "Collected %d index segments; committing as segmented index %s", @@ -406,114 +325,3 @@ def _create_segmented_index( lance_ds.commit_existing_index_segments(name, column, index_metas) logger.info("Segmented index %s committed successfully", name) - - -def _prepare_index_segments_for_commit( - lance_ds: lance.LanceDataset, - index_type: str, - index_metas: list[lance.Index | lance.indices.IndexSegment], -) -> list[lance.Index | lance.indices.IndexSegment]: - """Prepare worker-built segments for the final manifest commit.""" - if index_type not in MERGED_SEGMENTED_INDEX_TYPES or len(index_metas) <= 1: - return index_metas - - merged = lance_ds.merge_existing_index_segments([cast(lance.Index, segment) for segment in index_metas]) - return [merged] - - -def _create_partitioned_index( - open_context: DatasetOpenContext, - *, - column: str, - index_type: str, - name: str, - replace: bool, - fragment_data: list[dict[str, list[int]]], - fragment_ids_to_use: list[int], - num_partitions: int | None, - max_concurrency: int | None, - **kwargs: Any, -) -> None: - """Legacy partitioned-and-merged index workflow. - - Workers build partial index files sharing the same UUID, then the - coordinator merges them with ``merge_index_metadata`` and commits via a - manual ``CreateIndex`` transaction. - """ - # Generate unique index ID (shared across all partitions) - index_id = str(uuid.uuid4()) - - handler_cls = daft.cls( - FragmentIndexHandler, - max_concurrency=max_concurrency, - ) - handler = handler_cls( - open_context=open_context, - column=column, - index_type=index_type, - name=name, - fragment_uuid=index_id, - replace=replace, - **kwargs, - ) - - with execution_config_ctx(maintain_order=False): - if num_partitions is not None and num_partitions > 1: - df = from_pylist(fragment_data).repartition(num_partitions) - else: - df = from_pylist(fragment_data) - - df = df.select(handler(df["fragment_ids"])) - df.collect() - - logger.info("Starting index metadata merging by reloading dataset to get latest state") - lance_ds = open_context.open_latest() - lance_ds.merge_index_metadata(index_id, index_type) - - logger.info("Starting atomic index creation and commit") - field_id = lance_ds.schema.get_field_index(column) - index = lance.Index( - uuid=index_id, - name=name, - fields=[field_id], - dataset_version=lance_ds.version, - fragment_ids=set(fragment_ids_to_use), - index_version=0, - ) - removed_indices = [] - if replace: - # NOTE: kept on list_indices() until the distributed-index commit path is - # rewritten to populate index_details (e.g. via commit_existing_index_segments). - # describe_indices() raises on indices produced by this flow because their - # index_details field is empty. - for idx_info_raw in lance_ds.list_indices(): - idx_info = cast(dict[str, Any], idx_info_raw) - if idx_info["name"] == name: - field_ids = [lance_ds.schema.get_field_index(f) for f in idx_info["fields"]] - removed_indices.append( - lance.Index( - uuid=idx_info["uuid"], - name=idx_info["name"], - fields=field_ids, - dataset_version=lance_ds.version, - fragment_ids=idx_info["fragment_ids"], - index_version=idx_info["version"], - base_id=idx_info.get("base_id"), - ) - ) - - create_index_op = lance.LanceOperation.CreateIndex( - new_indices=[index], - removed_indices=removed_indices, - ) - - # Commit the index operation atomically - lance.LanceDataset.commit( - open_context.uri, - create_index_op, - read_version=lance_ds.version, - storage_options=open_context.storage_options, - **open_context.commit_kwargs, - ) - - logger.info("Index %s created successfully with ID %s", name, index_id) diff --git a/tests/io/lance/test_lance_distributed_index_types.py b/tests/io/lance/test_lance_distributed_index_types.py new file mode 100644 index 0000000..2855801 --- /dev/null +++ b/tests/io/lance/test_lance_distributed_index_types.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import warnings +from pathlib import Path + +import lance +import pytest + +from daft.dependencies import pa +from daft_lance import create_scalar_index + +warnings.filterwarnings("ignore", category=DeprecationWarning, module="lance") + + +def _make_dataset(path: Path) -> str: + table = pa.table( + { + "id": list(range(80)), + "doc": [f"document number {i % 8}" for i in range(80)], + "tags": [[f"tag-{i % 4}", "common"] for i in range(80)], + } + ) + lance.write_dataset(table, str(path), mode="create", max_rows_per_file=20) + return str(path) + + +@pytest.mark.parametrize( + ("index_type", "column", "expected_type", "predicate", "expected_rows"), + [ + ("ZONEMAP", "id", "ZoneMap", "id >= 40", 40), + ("NGRAM", "doc", "NGram", "contains(doc, 'number 3')", 10), + ("LABEL_LIST", "tags", "LabelList", "array_contains(tags, 'tag-1')", 20), + ("BLOOMFILTER", "id", "BloomFilter", "id = 42", 1), + ], +) +def test_new_type_builds_distributed_with_full_metadata( + tmp_path: Path, + index_type: str, + column: str, + expected_type: str, + predicate: str, + expected_rows: int, +) -> None: + """ZONEMAP/NGRAM/LABEL_LIST/BLOOMFILTER now use the distributed segment workflow.""" + uri = _make_dataset(tmp_path / f"{index_type.lower()}.lance") + + # fragment_group_size=1 forces one segment per worker build. + create_scalar_index(uri, column=column, index_type=index_type, fragment_group_size=1) + + described = lance.dataset(uri).describe_indices() + assert len(described) == 1 + desc = described[0] + assert desc.name == f"{column}_{index_type.lower()}_idx" + # The whole point of the segment workflow: metadata must be complete, not "Unknown". + assert desc.index_type == expected_type + assert desc.type_url != "" + assert desc.num_rows_indexed == 80 + assert len(desc.segments) >= 1 + + results = lance.dataset(uri).scanner(filter=predicate).to_table() + assert results.num_rows == expected_rows + + +def test_new_types_replace_rebuilds_in_place(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "replace.lance") + create_scalar_index(uri, column="id", index_type="ZONEMAP", name="z_idx") + create_scalar_index(uri, column="id", index_type="ZONEMAP", name="z_idx") + + described = lance.dataset(uri).describe_indices() + assert len(described) == 1 + assert described[0].name == "z_idx" + assert lance.dataset(uri).scanner(filter="id < 10").to_table().num_rows == 10 + + +def test_new_types_reject_existing_name_without_replace(tmp_path: Path) -> None: + uri = _make_dataset(tmp_path / "noreplace.lance") + create_scalar_index(uri, column="id", index_type="ZONEMAP", name="z_idx") + with pytest.raises(ValueError, match="already exists. Set replace=True"): + create_scalar_index(uri, column="id", index_type="ZONEMAP", name="z_idx", replace=False) + + +def test_all_supported_types_match_pylance_segment_native_set(tmp_path: Path) -> None: + """DISTRIBUTED_INDEX_TYPES must stay in sync with what Lance can build as segments.""" + from daft_lance.lance_scalar_index import DISTRIBUTED_INDEX_TYPES + + ds = lance.dataset(_make_dataset(tmp_path / "probe.lance")) + segment_native = { + t + for t in ("BTREE", "BITMAP", "INVERTED", "NGRAM", "ZONEMAP", "LABEL_LIST", "BLOOMFILTER", "RTREE") + if ds._is_segment_native_scalar_index_type(t) + } + # Everything we route must be segment-native; RTREE stays excluded until + # GeoArrow columns are testable. + assert DISTRIBUTED_INDEX_TYPES <= segment_native + assert "RTREE" not in DISTRIBUTED_INDEX_TYPES diff --git a/tests/io/lance/test_lance_scalar_index.py b/tests/io/lance/test_lance_scalar_index.py index faa2c13..54c5c34 100644 --- a/tests/io/lance/test_lance_scalar_index.py +++ b/tests/io/lance/test_lance_scalar_index.py @@ -4,6 +4,7 @@ import tempfile from inspect import signature from pathlib import Path +from typing import Any, cast import lance import pytest @@ -14,7 +15,6 @@ from daft_lance.lance_scalar_index import ( SegmentedFragmentIndexHandler, _existing_index_names, - _prepare_index_segments_for_commit, create_scalar_index_internal, ) from daft_lance.namespace import DatasetOpenContext @@ -105,10 +105,21 @@ def generate_multi_fragment_dataset(tmp_path, num_fragments=4, rows_per_fragment class TestDistributedIndexing: """Test cases for distributed indexing functionality.""" - def test_replace_defaults_to_false(self): - """Test that scalar index replacement is opt-in.""" - assert signature(create_scalar_index).parameters["replace"].default is False - assert signature(create_scalar_index_internal).parameters["replace"].default is False + def test_replace_defaults_to_true(self) -> None: + """Replacement is opt-out, matching pylance's default.""" + assert signature(create_scalar_index).parameters["replace"].default is True + assert signature(create_scalar_index_internal).parameters["replace"].default is True + + def test_segmented_kwarg_is_rejected_loudly(self) -> None: + """The removed segmented parameter must fail at the API boundary.""" + with pytest.raises(TypeError, match="'segmented' parameter was removed"): + create_scalar_index_internal( + lance_ds=cast(Any, None), + open_context=cast(Any, None), + column="a", + index_type="INVERTED", + segmented=True, + ) def test_build_distributed_index_search_functionality(self, multi_fragment_lance_dataset): """Test that the built index actually works for searching.""" @@ -548,7 +559,7 @@ def test_build_distributed_index_no_fragments(self, temp_dir): assert len(indices) == 0, f"Expected no indices for empty dataset, got {len(indices)}" def test_build_distributed_index_zonemap_type(self, temp_dir): - """Test building ZONEMAP index on numeric column (falls back to single-threaded).""" + """Test building ZONEMAP index distributed on a numeric column.""" data = { "id": [1, 2, 3, 4, 5, 6, 7, 8], "price": [10.5, 20.75, 30.0, 40.25, 50.5, 60.75, 70.0, 80.25], @@ -558,8 +569,6 @@ def test_build_distributed_index_zonemap_type(self, temp_dir): path = Path(temp_dir) / "zonemap_test.lance" dataset.write_lance(uri=path, max_rows_per_file=2) - # ZONEMAP is not supported by merge_index_metadata, so it falls back - # to single-threaded creation via Lance's create_scalar_index. create_scalar_index( uri=path, column="price", @@ -581,7 +590,7 @@ def test_build_distributed_index_zonemap_type(self, temp_dir): assert results.num_rows > 0, "No results found for ZONEMAP index query" def test_build_distributed_index_zonemap_integer_column(self, temp_dir): - """Test building ZONEMAP index on integer column (falls back to single-threaded).""" + """Test building ZONEMAP index distributed on an integer column.""" data = { "id": [1, 2, 3, 4, 5, 6, 7, 8], "score": [100, 200, 300, 400, 500, 600, 700, 800], @@ -666,37 +675,16 @@ def create_index_uncommitted(self, **kwargs): } ] - def test_segmented_inverted_segments_are_merged_before_commit(self): - """Test that INVERTED segments are merged into one physical segment before commit.""" - - class FakeLanceDataset: - def __init__(self): - self.calls = [] - - def merge_existing_index_segments(self, segments): - self.calls.append(segments) - return {"segment": "merged"} - - fake_ds = FakeLanceDataset() - segments = [{"segment": "a"}, {"segment": "b"}] - - prepared = _prepare_index_segments_for_commit(fake_ds, "INVERTED", segments) - - assert prepared == [{"segment": "merged"}] - assert fake_ds.calls == [segments] - - def test_segmented_btree_segments_are_committed_without_merge(self): - """Test that non-merged segmented index types keep their physical segments.""" - - class FakeLanceDataset: - def merge_existing_index_segments(self, segments): - raise AssertionError("BTREE segments must not be merged") - - segments = [{"segment": "a"}, {"segment": "b"}] + def test_segments_are_committed_without_merge(self): + """Worker-built segments commit as-is, with no merge step. - prepared = _prepare_index_segments_for_commit(FakeLanceDataset(), "BTREE", segments) - - assert prepared is segments + Multi-segment indexes are fully functional (verified: split segments + load and prune at query time with identical scores); compaction is + left to optimize_indices. + """ + # The merge helper is gone from the module entirely. + assert not hasattr(lance_scalar_index, "_prepare_index_segments_for_commit") + assert not hasattr(lance_scalar_index, "MERGED_SEGMENTED_INDEX_TYPES") def test_existing_index_names_falls_back_to_list_indices(self): """Test that existing-name checks still work for legacy indexes with bad details.""" @@ -710,17 +698,13 @@ def list_indices(self): assert _existing_index_names(FakeLanceDataset()) == {"legacy_idx"} - def test_segmented_bitmap_handler_forwards_shard_id(self): - """Test that BITMAP segment creation forwards the required shard id.""" + def test_segmented_bitmap_handler_builds_without_shard_id(self): + """BITMAP segment creation needs no shard id; segments commit as-is.""" class FakeLanceDataset: - def __init__(self): + def __init__(self) -> None: self.calls = [] - @property - def _ds(self): - raise AssertionError("public BITMAP segment creation should not use fallback") - def create_index_uncommitted(self, **kwargs): self.calls.append(kwargs) return {"segment": "bitmap-metadata"} @@ -733,7 +717,7 @@ def create_index_uncommitted(self, **kwargs): name="flag_idx", ) - raw_segment = handler([1, 2], shard_id=7) + raw_segment = handler([1, 2]) assert pickle.loads(raw_segment) == {"segment": "bitmap-metadata"} assert fake_ds.calls == [ @@ -744,37 +728,17 @@ def create_index_uncommitted(self, **kwargs): "replace": False, "train": True, "fragment_ids": [1, 2], - "shard_id": 7, } ] - def test_segmented_bitmap_segments_are_merged_before_commit(self): - """Test that BITMAP segments are merged to avoid one physical segment per fragment.""" - - class FakeLanceDataset: - def __init__(self): - self.calls = [] - - def merge_existing_index_segments(self, segments): - self.calls.append(segments) - return {"segment": "merged-bitmap"} - - fake_ds = FakeLanceDataset() - segments = [{"segment": "a"}, {"segment": "b"}] - - prepared = _prepare_index_segments_for_commit(fake_ds, "BITMAP", segments) - - assert prepared == [{"segment": "merged-bitmap"}] - assert fake_ds.calls == [segments] - def test_segmented_bitmap_respects_fragment_group_size(self, monkeypatch): """Test that segmented BITMAP can group multiple fragments per segment.""" class FakeFragment: - def __init__(self, fragment_id): + def __init__(self, fragment_id: int) -> None: self.fragment_id = fragment_id - def count_rows(self): + def count_rows(self) -> int: return 1 class FakeLanceDataset: @@ -799,34 +763,54 @@ def fake_create_segmented_index(**kwargs): column="flag", index_type="BITMAP", name="flag_bitmap_idx", - segmented=True, fragment_group_size=2, ) assert [len(group["fragment_ids"]) for group in calls[0]["fragment_data"]] == [2, 2] - def test_bitmap_replace_true_existing_index_preserves_lance_replacement(self): - """Test that default BITMAP indexing keeps replace=True behavior for existing indexes.""" + def test_replace_true_uses_atomic_overlap_replacement(self, monkeypatch): + """replace=True replaces atomically, without dropping the old index. + + The segment commit retires the overlapped segments in the same + transaction. + """ + + class FakeSegment: + def __init__(self, fragment_ids: set[int]) -> None: + self.fragment_ids = fragment_ids class ExistingIndex: name = "flag_bitmap_idx" + segments = [FakeSegment({0, 1})] + + class FakeFragment: + def __init__(self, fragment_id: int) -> None: + self.fragment_id = fragment_id + + def count_rows(self) -> int: + return 1 class FakeLanceDataset: schema = pa.schema([("flag", pa.int64())]) - def __init__(self): - self.calls = [] + def __init__(self) -> None: + self.dropped: list[str] = [] - def describe_indices(self): + def describe_indices(self) -> list[Any]: return [ExistingIndex()] - def create_scalar_index(self, **kwargs): - self.calls.append(kwargs) + def drop_index(self, name: str) -> None: + self.dropped.append(name) + + def get_fragments(self) -> list[Any]: + return [FakeFragment(0), FakeFragment(1)] fake_ds = FakeLanceDataset() + calls = [] + monkeypatch.setattr(lance_scalar_index, "_create_segmented_index", lambda **kwargs: calls.append(kwargs)) create_scalar_index_internal( - lance_ds=fake_ds, + lance_ds=cast(Any, fake_ds), open_context=DatasetOpenContext(uri="memory://bitmap", version=1), column="flag", index_type="BITMAP", @@ -834,39 +818,32 @@ def create_scalar_index(self, **kwargs): replace=True, ) - assert fake_ds.calls == [ - { - "column": "flag", - "index_type": "BITMAP", - "name": "flag_bitmap_idx", - "replace": True, - } - ] + # No drop: overlap replacement retires the old segments atomically. + assert fake_ds.dropped == [] + # Workers still build with replace=True: the pinned snapshot has the name. + assert calls[0]["handler_replace"] is True - def test_bitmap_replace_true_default_name_preserves_lance_replacement(self, temp_dir): - """Test that default BITMAP replacement preserves Lance's default index name.""" + def test_replace_true_rebuilds_single_index_with_same_name(self, temp_dir): + """Replacing an index must leave exactly one index behind, same name.""" path = Path(temp_dir) / "bitmap_default_replace.lance" table = pa.table( { - "flag": pa.array([1, 2, 1, 3], type=pa.int64()), + "flag": pa.array([1, 2, 1, 3, 2, 1, 3, 2], type=pa.int64()), } ) lance.write_dataset(table, str(path), max_rows_per_file=2) - initial_dataset = lance.dataset(str(path)) - initial_dataset.create_scalar_index(column="flag", index_type="BITMAP") - initial_names = [idx["name"] for idx in lance.dataset(str(path)).list_indices()] - assert len(initial_names) == 1 + create_scalar_index(uri=path, column="flag", index_type="BITMAP", name="flag_idx") + assert len(lance.dataset(str(path)).describe_indices()) == 1 - create_scalar_index( - uri=path, - column="flag", - index_type="BITMAP", - replace=True, - ) + create_scalar_index(uri=path, column="flag", index_type="BITMAP", name="flag_idx") - final_names = [idx["name"] for idx in lance.dataset(str(path)).list_indices()] - assert final_names == initial_names + described = lance.dataset(str(path)).describe_indices() + assert len(described) == 1 + assert described[0].name == "flag_idx" + assert described[0].num_rows_indexed == 8 + results = lance.dataset(str(path)).scanner(filter="flag = 1").to_table() + assert results.num_rows == 3 def test_segmented_btree_basic(self, temp_dir): """Test basic segmented BTree index creation and query.""" @@ -884,7 +861,6 @@ def test_segmented_btree_basic(self, temp_dir): column="price", index_type="BTREE", name="price_seg_idx", - segmented=True, max_concurrency=2, ) @@ -920,7 +896,6 @@ def test_segmented_btree_multiple_segments(self, temp_dir): column="score", index_type="BTREE", name="score_seg_idx", - segmented=True, fragment_group_size=2, max_concurrency=2, ) @@ -957,7 +932,6 @@ def test_segmented_btree_describe_indices_works(self, temp_dir): column="value", index_type="BTREE", name="value_idx", - segmented=True, ) updated_dataset = lance.dataset(path) @@ -970,8 +944,8 @@ def test_segmented_btree_describe_indices_works(self, temp_dir): assert desc.type_url == "/lance.table.BTreeIndexDetails" assert desc.num_rows_indexed == 4 - def test_segmented_btree_replace_existing_is_rejected(self, temp_dir): - """Test that replacing an existing segmented BTree index fails safely.""" + def test_btree_replace_semantics(self, temp_dir): + """replace=False rejects an existing name; replace=True rebuilds it.""" data = { "id": [1, 2, 3, 4, 5, 6, 7, 8], "price": [10.5, 20.75, 30.0, 40.25, 50.5, 60.75, 70.0, 80.25], @@ -980,37 +954,93 @@ def test_segmented_btree_replace_existing_is_rejected(self, temp_dir): path = Path(temp_dir) / "segmented_btree_replace.lance" dataset.write_lance(uri=path, max_rows_per_file=2) - # Create initial index create_scalar_index( uri=path, column="price", index_type="BTREE", name="price_idx", - segmented=True, ) + assert len(lance.dataset(path).describe_indices()) == 1 - ds1 = lance.dataset(path) - assert len(ds1.describe_indices()) == 1 - - with pytest.raises(ValueError, match="cannot atomically replace existing index"): + # replace=False refuses to touch the existing index. + with pytest.raises(ValueError, match="already exists. Set replace=True"): create_scalar_index( uri=path, column="price", index_type="BTREE", name="price_idx", - segmented=True, - replace=True, + replace=False, ) + # Default (replace=True) drops and rebuilds; still exactly one index. + create_scalar_index( + uri=path, + column="price", + index_type="BTREE", + name="price_idx", + ) ds2 = lance.dataset(path) described = ds2.describe_indices() assert len(described) == 1 assert described[0].name == "price_idx" - # Query still works after the rejected replacement results = ds2.scanner(filter="price > 50.0", columns=["id", "price"]).to_table() assert results.num_rows == 4 + def test_replace_rebuild_advances_version_exactly_once(self, temp_dir): + """Atomic overlap replacement lands as a single new dataset version.""" + data = { + "id": [1, 2, 3, 4, 5, 6, 7, 8], + "price": [10.5, 20.75, 30.0, 40.25, 50.5, 60.75, 70.0, 80.25], + } + dataset = daft.from_pydict(data) + path = Path(temp_dir) / "atomic_replace.lance" + dataset.write_lance(uri=path, max_rows_per_file=2) + + create_scalar_index(uri=path, column="price", index_type="BTREE", name="atomic_idx") + version_before = lance.dataset(path).version + + create_scalar_index(uri=path, column="price", index_type="BTREE", name="atomic_idx") + + latest = lance.dataset(path) + # One transaction: no intermediate drop version, no append duplication. + assert latest.version == version_before + 1 + described = latest.describe_indices() + assert len(described) == 1 + assert described[0].name == "atomic_idx" + assert len(described[0].segments) == 1 + assert latest.scanner(filter="price > 50.0", columns=["id", "price"]).to_table().num_rows == 4 + + def test_replace_with_stale_coverage_rebuilds_atomically(self, temp_dir) -> None: + """Stale coverage from a fully deleted fragment retires atomically. + + A fully deleted fragment inside a mixed segment is the only stale + coverage normal operations produce; the rebuild retires it without a + drop, in exactly one new version. + """ + data = { + "id": list(range(80)), + "name": [f"name-{i % 8}" for i in range(80)], + } + path = Path(temp_dir) / "stale_coverage.lance" + lance.write_dataset(daft.from_pydict(data).to_arrow(), str(path), max_rows_per_file=20) + + create_scalar_index(uri=path, column="name", index_type="INVERTED", name="stale_idx") + # Delete every row of fragment 0: the committed segment keeps its + # (now partially dead) coverage {0,1,2,3}. + lance.dataset(str(path)).delete("id < 20") + + version_before = lance.dataset(str(path)).version + create_scalar_index(uri=path, column="name", index_type="INVERTED", name="stale_idx") + + latest = lance.dataset(str(path)) + assert latest.version == version_before + 1 + described = latest.describe_indices()[0] + assert len(described.segments) == 1 + assert sorted(described.segments[0].fragment_ids) == [1, 2, 3] + # ids 3, 11, 19 (name-3) were among the 20 deleted rows: 10 - 3 remain. + assert latest.scanner(filter="name = 'name-3'").to_table().num_rows == 7 + def test_segmented_btree_string_column(self, temp_dir): """Test segmented BTree index on a string column.""" import pyarrow as pa @@ -1030,7 +1060,6 @@ def test_segmented_btree_string_column(self, temp_dir): column="category", index_type="BTREE", name="cat_idx", - segmented=True, ) updated_dataset = lance.dataset(path) @@ -1054,7 +1083,6 @@ def test_segmented_btree_integer_column(self, temp_dir): column="count", index_type="BTREE", name="count_idx", - segmented=True, ) updated_dataset = lance.dataset(path) @@ -1065,42 +1093,18 @@ def test_segmented_btree_integer_column(self, temp_dir): results = updated_dataset.scanner(filter="count > 500", columns=["id", "count"]).to_table() assert results.num_rows == 3 # 600, 700, 800 - def test_segmented_false_uses_lance_scalar_flow_for_btree(self, monkeypatch): - """Test that segmented=False opts out of distributed BTREE workflows.""" - - class FakeLanceDataset: - schema = pa.schema([("price", pa.float64())]) - - def describe_indices(self): - return [] - - def create_scalar_index(self, **kwargs): - calls.append(("lance_scalar", kwargs)) - - calls = [] - - def fake_create_segmented_index(**kwargs): - calls.append(("segmented", kwargs)) - - def fake_create_partitioned_index(**kwargs): - calls.append(("partitioned", kwargs)) - - monkeypatch.setattr(lance_scalar_index, "_create_segmented_index", fake_create_segmented_index) - monkeypatch.setattr(lance_scalar_index, "_create_partitioned_index", fake_create_partitioned_index) - - create_scalar_index_internal( - lance_ds=FakeLanceDataset(), - open_context=DatasetOpenContext(uri="memory://btree", version=1), - column="price", - index_type="BTREE", - name="price_btree_idx", - segmented=False, - ) - - assert [call[0] for call in calls] == ["lance_scalar"] + def test_unsupported_type_raises_instead_of_single_node_fallback(self) -> None: + """Types without a distributed path fail loudly; no silent fallback.""" + with pytest.raises(ValueError, match="Unsupported distributed index type 'RTREE'"): + create_scalar_index_internal( + lance_ds=cast(Any, None), + open_context=cast(Any, None), + column="geom", + index_type="RTREE", + ) def test_segmented_inverted_creates_index(self, multi_fragment_lance_dataset): - """Test that segmented=True with INVERTED creates an index.""" + """INVERTED creates an index through the distributed segment workflow.""" dataset_uri = multi_fragment_lance_dataset create_scalar_index( @@ -1108,7 +1112,6 @@ def test_segmented_inverted_creates_index(self, multi_fragment_lance_dataset): column="text", index_type="INVERTED", name="text_inv_idx", - segmented=True, ) updated_dataset = lance.dataset(dataset_uri) diff --git a/tests/io/lance/test_namespace.py b/tests/io/lance/test_namespace.py index cb73b12..80e3b5b 100644 --- a/tests/io/lance/test_namespace.py +++ b/tests/io/lance/test_namespace.py @@ -248,10 +248,9 @@ def test_namespace_create_scalar_index(tmp_path: Path) -> None: assert any(idx["fields"] == ["price"] for idx in indices) -@pytest.mark.parametrize("segmented", [False, True]) -def test_namespace_create_distributed_inverted_index(tmp_path: Path, segmented: bool) -> None: +def test_namespace_create_distributed_inverted_index(tmp_path: Path) -> None: ns = _dir_ns(tmp_path) - table_id = [f"inverted_{segmented}"] + table_id = ["inverted_dist"] daft_lance.write_lance( daft.from_pydict({"id": list(range(20)), "text": [f"document {i}" for i in range(20)]}), @@ -265,7 +264,6 @@ def test_namespace_create_distributed_inverted_index(tmp_path: Path, segmented: column="text", index_type="INVERTED", name="text_idx", - segmented=segmented, **ns, ) @@ -799,16 +797,12 @@ def test_maintenance_udfs_hold_a_context_not_a_dataset(tmp_path: Path) -> None: FragmentHandler, GroupFragmentMergeUDF, ) - from daft_lance.lance_scalar_index import ( - FragmentIndexHandler, - SegmentedFragmentIndexHandler, - ) + from daft_lance.lance_scalar_index import SegmentedFragmentIndexHandler context = _ns_handle(tmp_path, "udf_tbl").worker_open_context() plain = [ CompactionTaskUDF(context), - FragmentIndexHandler(context, "score", "BTREE", "idx", "uuid", False), SegmentedFragmentIndexHandler(context, "score", "BTREE", "idx"), ] # daft.cls wraps these, so reach through to the instance it actually holds.