From 4afb71ef3bd5598ef2f628c67f9c45f896fd5a1c Mon Sep 17 00:00:00 2001 From: wangzheyan Date: Sun, 20 Sep 2026 12:29:40 +0800 Subject: [PATCH] refactor: distributed-only scalar indexing with complete index metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make the distributed segment-index workflow (create_index_uncommitted per worker + commit_existing_index_segments at the coordinator) the only code path for create_scalar_index. - Remove the public 'segmented' parameter (it leaked an implementation mechanism into the API; passing it now fails loudly with a TypeError). - Remove the legacy partitioned-and-merged INVERTED path whose manual CreateIndex commit cannot record index_details — the source of the 'Unknown' index type metadata (references #69). - Remove all single-node fallbacks: types without a distributed path raise ValueError with a pointer to pylance instead of silently degrading to single-node execution. - Unlock distributed builds for ZONEMAP, NGRAM, LABEL_LIST and BLOOMFILTER (pylance 11 builds all of them as segments; RTREE needs GeoArrow columns and stays excluded until testable). - replace now defaults to True (matching pylance) and is atomic: the segment commit retires the existing index's overlapped segments in the same transaction as the new ones. replace=False rejects an existing name. Stale segments that no longer overlap any live fragment cannot be retired by the swap, but normal operations never produce them and any that appear are healed by optimize_indices (verified). - Commit worker-built segments as-is: no merge step and no BITMAP shard ids. Multi-segment INVERTED/BITMAP indexes are fully functional (verified via scan statistics: split segments are loaded and pruned at query time, rows_scanned=1, with BM25 scores identical to a merged build), so compaction is left to optimize_indices. This also removes the single-node merge bottleneck from the build pipeline. --- README.md | 20 + daft_lance/_lance.py | 50 +-- daft_lance/lance_scalar_index.py | 386 +++++------------- .../test_lance_distributed_index_types.py | 95 +++++ tests/io/lance/test_lance_scalar_index.py | 307 +++++++------- tests/io/lance/test_namespace.py | 12 +- 6 files changed, 390 insertions(+), 480 deletions(-) create mode 100644 tests/io/lance/test_lance_distributed_index_types.py 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.