diff --git a/docs/src/distributed-indexing.md b/docs/src/distributed-indexing.md index ba54e9bd..5865c44c 100755 --- a/docs/src/distributed-indexing.md +++ b/docs/src/distributed-indexing.md @@ -22,10 +22,8 @@ def create_scalar_index( index_type: Union[ Literal["BTREE"], Literal["BITMAP"], - Literal["LABEL_LIST"], Literal["INVERTED"], Literal["FTS"], - Literal["NGRAM"], Literal["ZONEMAP"], IndexConfig, ], @@ -53,7 +51,7 @@ def create_scalar_index( |-----------|------|-------------| | `uri` | `str`, optional | The URI of the Lance dataset. Either `uri` OR (`namespace_impl` + `table_id`) must be provided. | | `column` | `str` | Column name to index | -| `index_type` | `str` or `IndexConfig` | Index type, can be `"INVERTED"`, `"FTS"`, `"BTREE"`, `"BITMAP"`, `"LABEL_LIST"`, `"NGRAM"`, `"ZONEMAP"`, or `IndexConfig` object | +| `index_type` | `str` or `IndexConfig` | Index type, can be `"INVERTED"`, `"FTS"`, `"BTREE"`, `"BITMAP"`, `"ZONEMAP"`, or `IndexConfig` object. `"FTS"` is an alias for `"INVERTED"`. | | `table_id` | `list[str]`, optional | The table identifier as a list of strings. | | `name` | `str`, optional | Index name, auto-generated if not provided | | `replace` | `bool`, optional | Whether to replace existing index with the same name, default is `True` | @@ -486,7 +484,7 @@ except ValueError as e: ### Important Notes -- **Index Type Support**: For distributed indexing, currently only `"INVERTED"`/`"FTS"`/`"BTREE"`/`"BITMAP"`/`"ZONEMAP"` index types are supported, even though the function signature accepts other index types. +- **Index Type Support**: Distributed scalar indexing supports only `"INVERTED"`/`"FTS"`/`"BTREE"`/`"BITMAP"`/`"ZONEMAP"`. - **Default Behavior**: The `replace` parameter defaults to `True`, meaning existing indices with the same name will be replaced without warning. Set `replace=False` to prevent accidental overwrites. - **Fragment Selection**: Use `fragment_ids` parameter to build indices on specific fragments only. This is useful for incremental index building or testing. - **Error Handling**: When `replace=False` and an index with the same name exists, a `ValueError` or `RuntimeError` will be raised depending on the execution context. diff --git a/lance_ray/index.py b/lance_ray/index.py index ea01bf33..9b6400c2 100755 --- a/lance_ray/index.py +++ b/lance_ray/index.py @@ -215,15 +215,91 @@ def _build_rabitq_model(*, dimension: int, num_bits: int = 1) -> str: return indices.build_rq_model(dimension=dimension, num_bits=num_bits) -_SCALAR_SEGMENT_INDEX_TYPES = {"BTREE", "BITMAP", "INVERTED", "FTS", "ZONEMAP"} +_DistributedScalarIndexType: TypeAlias = Literal[ + "BTREE", "BITMAP", "INVERTED", "FTS", "ZONEMAP" +] +_SCALAR_SEGMENT_INDEX_TYPES = frozenset( + {"BTREE", "BITMAP", "INVERTED", "FTS", "ZONEMAP"} +) + +def _normalize_distributed_scalar_index_type( + index_type: _DistributedScalarIndexType | str | IndexConfig, +) -> tuple[str, str | IndexConfig]: + """Return the logical type and worker input for a distributed scalar index. -def _scalar_index_type_name(index_type: str | IndexConfig) -> str | None: + Keep the caller-owned configuration immutable while normalizing its type for the + worker. ``IndexConfig`` plugin names are normalized to lowercase; FTS is a + public alias for Core's ``inverted`` plugin. + """ if isinstance(index_type, str): - return index_type.upper() - if isinstance(index_type, IndexConfig): - return index_type.index_type.upper() - return None + logical_index_type = index_type.upper() + worker_index_type: str | IndexConfig = logical_index_type + elif isinstance(index_type, IndexConfig): + logical_index_type = index_type.index_type.upper() + core_index_type = ( + "inverted" if logical_index_type == "FTS" else logical_index_type.lower() + ) + worker_index_type = IndexConfig(core_index_type, index_type.parameters) + else: + raise ValueError( + "index_type must be a string literal or IndexConfig object, got " + f"{type(index_type)}" + ) + + if logical_index_type not in _SCALAR_SEGMENT_INDEX_TYPES: + raise ValueError( + f"Distributed indexing does not support index type '{logical_index_type}'" + ) + + return logical_index_type, worker_index_type + + +def _validate_distributed_scalar_index_field( + field: pa.Field, column: str, logical_index_type: str +) -> None: + """Apply the Core scalar-index field contract before scheduling Ray workers.""" + field_type = field.type + field_metadata = field.metadata or {} + if hasattr(field_type, "storage_type"): + field_type = field_type.storage_type + + if logical_index_type in {"BTREE", "BITMAP", "ZONEMAP"}: + is_supported = ( + pa.types.is_integer(field_type) + or pa.types.is_floating(field_type) + or pa.types.is_boolean(field_type) + or pa.types.is_string(field_type) + or pa.types.is_large_string(field_type) + or pa.types.is_temporal(field_type) + or pa.types.is_fixed_size_binary(field_type) + ) + if not is_supported: + raise TypeError( + f"Column {column} must be int, float, bool, string, large string, " + f"fixed-size binary, or temporal for {logical_index_type} index, " + f"got {field_type}" + ) + else: + 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 + is_json = ( + pa.types.is_large_binary(value_type) + and field_metadata.get(b"ARROW:extension:name") == b"lance.json" + ) + if not ( + pa.types.is_string(value_type) + or pa.types.is_large_string(value_type) + or is_json + ): + raise TypeError( + f"Column {column} must be string, large string, list of strings, " + f"or json for {logical_index_type} index, got {value_type}" + ) + + if pa.types.is_duration(field_type): + raise TypeError(f"Scalar index column {column} cannot currently be a duration") def _handle_scalar_segment_index( @@ -411,14 +487,7 @@ def create_scalar_index( uri: Optional[Union[str, "lance.LanceDataset"]] = None, *, column: str, - index_type: Literal["BTREE"] - | Literal["BITMAP"] - | Literal["LABEL_LIST"] - | Literal["INVERTED"] - | Literal["FTS"] - | Literal["NGRAM"] - | Literal["ZONEMAP"] - | IndexConfig, + index_type: _DistributedScalarIndexType | IndexConfig, table_id: Optional[list[str]] = None, name: Optional[str] = None, replace: bool = True, @@ -442,8 +511,8 @@ def create_scalar_index( ``LanceDataset``, its dataset URI is retained so distributed workers and the final commit stay on the same Lance branch. column: Column name to index. - index_type: Type of index to build ("BTREE", "BITMAP", "LABEL_LIST", - "INVERTED", "FTS", "NGRAM", "ZONEMAP") or IndexConfig object. + index_type: Type of index to build ("BTREE", "BITMAP", "INVERTED", "FTS", + or "ZONEMAP") or IndexConfig object. FTS is an alias for INVERTED. table_id: The table identifier as a list of strings. Must be provided together with namespace_impl. name: Name of the index (generated if None). @@ -469,7 +538,7 @@ def create_scalar_index( Raises: ValueError: If input parameters are invalid. - TypeError: If column type is not string. + TypeError: If the column type is unsupported by the selected index type. RuntimeError: If index building fails or pylance version is incompatible. """ # Check pylance version compatibility @@ -510,33 +579,9 @@ def create_scalar_index( if block_size is not None and block_size <= 0: raise ValueError(f"block_size must be positive, got {block_size}") - if isinstance(index_type, str): - valid_index_types = [ - "BTREE", - "BITMAP", - "LABEL_LIST", - "INVERTED", - "FTS", - "NGRAM", - "ZONEMAP", - ] - if index_type not in valid_index_types: - raise ValueError( - f"Index type must be one of {valid_index_types}, not '{index_type}'" - ) - - supported_distributed_types = {"INVERTED", "FTS", "BTREE", "BITMAP", "ZONEMAP"} - if index_type not in supported_distributed_types: - raise ValueError( - "Distributed indexing currently supports " - f"{sorted(supported_distributed_types)} index types, " - f"not '{index_type}'" - ) - elif not isinstance(index_type, IndexConfig): - raise ValueError( - "index_type must be a string literal or IndexConfig object, got " - f"{type(index_type)}" - ) + logical_index_type, index_type = _normalize_distributed_scalar_index_type( + index_type + ) # Note: Ray initialization is now handled by the Pool, following the pattern from io.py # This removes the need for explicit ray.init() calls @@ -573,41 +618,9 @@ def create_scalar_index( column = resolved_column.path field = resolved_column.field - # Check column type according to index type - 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 + _validate_distributed_scalar_index_field(field, column, logical_index_type) - if isinstance(index_type, str): - match index_type: - case "INVERTED" | "FTS": - if not ( - pa.types.is_string(value_type) - or pa.types.is_large_string(value_type) - ): - raise TypeError( - f"Column {column} must be string type for {index_type} " - f"index, got {value_type}" - ) - case "BTREE" | "ZONEMAP": - is_supported = ( - pa.types.is_integer(value_type) - or pa.types.is_floating(value_type) - or pa.types.is_string(value_type) - or pa.types.is_large_string(value_type) - ) - if not is_supported: - raise TypeError( - f"Column {column} must be numeric or string type for " - f"{index_type} index, got {value_type}" - ) - case _: - # For other index types, skip strict validation to maintain compatibility - pass - - use_segment_workflow = ( - _scalar_index_type_name(index_type) in _SCALAR_SEGMENT_INDEX_TYPES - ) + use_segment_workflow = logical_index_type in _SCALAR_SEGMENT_INDEX_TYPES if name is None: name = f"{column}_idx" diff --git a/tests/test_distributed_indexing.py b/tests/test_distributed_indexing.py index 81d9854d..df0b061e 100755 --- a/tests/test_distributed_indexing.py +++ b/tests/test_distributed_indexing.py @@ -404,7 +404,7 @@ def test_build_distributed_index_invalid_index_type( with pytest.raises( ValueError, - match=r"Index type must be one of \['BTREE', 'BITMAP', 'LABEL_LIST', 'INVERTED', 'FTS', 'NGRAM', 'ZONEMAP'\], not 'INVALID'", + match="Distributed indexing does not support index type 'INVALID'", ): lr.create_scalar_index( uri=dataset_uri, @@ -453,7 +453,10 @@ def test_build_distributed_index_non_string_column(self, temp_dir): path = Path(temp_dir) / "non_string_test.lance" lr.write_lance(dataset, str(path), min_rows_per_file=2, max_rows_per_file=2) - with pytest.raises(TypeError, match="must be string type"): + with pytest.raises( + TypeError, + match="must be string, large string, list of strings, or json", + ): lr.create_scalar_index( uri=str(path), column="numeric_col", @@ -925,7 +928,7 @@ def test_distributed_index_error_handling_new_api(self, temp_dir): # Test with invalid index type with pytest.raises( ValueError, - match=r"Index type must be one of \['BTREE', 'BITMAP', 'LABEL_LIST', 'INVERTED', 'FTS', 'NGRAM', 'ZONEMAP'\], not 'INVALID_TYPE'", + match="Distributed indexing does not support index type 'INVALID_TYPE'", ): lr.create_scalar_index( uri=ds.uri, diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index b521df50..1c9115b7 100644 --- a/tests/test_vector_index_options.py +++ b/tests/test_vector_index_options.py @@ -20,7 +20,13 @@ def _load_index_module_with_stubs(): lance_dataset = ModuleType("lance.dataset") lance_dataset.Index = type("Index", (), {}) - lance_dataset.IndexConfig = type("IndexConfig", (), {}) + + class IndexConfig: + def __init__(self, index_type, parameters): + self.index_type = index_type + self.parameters = parameters + + lance_dataset.IndexConfig = IndexConfig lance_dataset.LanceDataset = object lance_indices = ModuleType("lance.indices") @@ -60,6 +66,7 @@ class _FakeField: def __init__(self, name, field_type=None): self.name = name self.type = field_type or index_mod.pa.float32() + self.metadata = None class _FakeLanceField: @@ -69,7 +76,7 @@ def id(self): class _FakeLanceSchema: def field(self, column): - if column not in {"value", "text"}: + if column not in {"value", "text", "flag", "tags"}: raise KeyError(column) return _FakeLanceField() @@ -82,6 +89,10 @@ def field(self, column): return _FakeField(column, index_mod.pa.int64()) if column == "text": return _FakeField(column, index_mod.pa.string()) + if column == "flag": + return _FakeField(column, index_mod.pa.bool_()) + if column == "tags": + return _FakeField(column, index_mod.pa.list_(index_mod.pa.string())) else: raise KeyError(column) @@ -91,6 +102,8 @@ def __iter__(self): _FakeField("vector"), _FakeField("value", index_mod.pa.int64()), _FakeField("text", index_mod.pa.string()), + _FakeField("flag", index_mod.pa.bool_()), + _FakeField("tags", index_mod.pa.list_(index_mod.pa.string())), ] ) @@ -445,10 +458,7 @@ def test_create_index_rejects_invalid_num_segments(monkeypatch): ) -@pytest.mark.parametrize("index_type", ["BTREE", "BITMAP", "INVERTED", "FTS"]) -def test_create_scalar_index_uses_segment_path(monkeypatch, index_type): - """Migrated scalar indexes should use Lance's segment workflow.""" - +def _patch_scalar_segment_workflow(monkeypatch): captured = {"loads": []} fake_dataset = _FakeDataset() @@ -482,6 +492,16 @@ def fake_map_async_with_pool(**kwargs): fake_handle_scalar_segment_index, ) monkeypatch.setattr(index_mod, "_map_async_with_pool", fake_map_async_with_pool) + return captured, fake_dataset + + +@pytest.mark.parametrize( + "index_type", ["BTREE", "BITMAP", "INVERTED", "FTS", "ZONEMAP"] +) +def test_create_scalar_index_uses_segment_path(monkeypatch, index_type): + """Migrated scalar indexes should use Lance's segment workflow.""" + + captured, fake_dataset = _patch_scalar_segment_workflow(monkeypatch) column = "text" if index_type in {"INVERTED", "FTS"} else "value" updated_dataset = index_mod.create_scalar_index( @@ -499,6 +519,153 @@ def fake_map_async_with_pool(**kwargs): assert fake_dataset.commit_kwargs["segments"] == ["segment"] +@pytest.mark.parametrize( + ("config_type", "column", "worker_type"), + [ + ("btree", "value", "btree"), + ("bitmap", "value", "bitmap"), + ("inverted", "text", "inverted"), + ("fts", "text", "inverted"), + ("zonemap", "value", "zonemap"), + ], +) +def test_scalar_index_config_uses_segment_path( + monkeypatch, config_type, column, worker_type +): + """IndexConfig input uses the same segment path without mutating the caller.""" + + captured, fake_dataset = _patch_scalar_segment_workflow(monkeypatch) + parameters = {"test_parameter": config_type} + index_config = index_mod.IndexConfig(config_type, parameters) + + updated_dataset = index_mod.create_scalar_index( + uri="memory://fake", + column=column, + index_type=index_config, + num_workers=2, + ) + + worker_config = captured["fragment_handler_kwargs"]["index_type"] + assert updated_dataset is fake_dataset + assert worker_config is not index_config + assert worker_config.index_type == worker_type + assert worker_config.parameters == parameters + assert index_config.index_type == config_type + assert index_config.parameters == parameters + + +@pytest.mark.parametrize( + "make_index_type", + [ + lambda: "ngram", + lambda: index_mod.IndexConfig("ngram", {}), + lambda: "LABEL_LIST", + lambda: index_mod.IndexConfig("label_list", {}), + ], + ids=[ + "string-ngram", + "config-ngram", + "string-label-list", + "config-label-list", + ], +) +def test_scalar_index_rejects_unsupported_types( + monkeypatch, make_index_type +): + """Unsupported input forms fail before a Ray worker is created.""" + + captured, _ = _patch_scalar_segment_workflow(monkeypatch) + + with pytest.raises(ValueError, match="Distributed indexing does not support"): + index_mod.create_scalar_index( + uri="memory://fake", + column="text", + index_type=make_index_type(), + num_workers=2, + ) + + assert "map_kwargs" not in captured + + +@pytest.mark.parametrize( + "make_index_type", + [lambda: "INVERTED", lambda: index_mod.IndexConfig("inverted", {})], + ids=["string", "config"], +) +def test_scalar_index_validates_text_column( + monkeypatch, make_index_type +): + """String and IndexConfig forms reject invalid FTS fields on the driver.""" + + captured, _ = _patch_scalar_segment_workflow(monkeypatch) + + with pytest.raises(TypeError, match="must be string, large string"): + index_mod.create_scalar_index( + uri="memory://fake", + column="value", + index_type=make_index_type(), + num_workers=2, + ) + + assert "map_kwargs" not in captured + + +@pytest.mark.parametrize( + "make_index_type", + [ + lambda: "BTREE", + lambda: index_mod.IndexConfig("btree", {}), + lambda: "ZONEMAP", + lambda: index_mod.IndexConfig("zonemap", {}), + ], + ids=["btree-string", "btree-config", "zonemap-string", "zonemap-config"], +) +def test_scalar_index_accepts_ordered_field( + monkeypatch, make_index_type +): + """Core-supported boolean fields are accepted for ordered index forms.""" + + captured, fake_dataset = _patch_scalar_segment_workflow(monkeypatch) + + updated_dataset = index_mod.create_scalar_index( + uri="memory://fake", + column="flag", + index_type=make_index_type(), + num_workers=2, + ) + + assert updated_dataset is fake_dataset + assert "map_kwargs" in captured + + +@pytest.mark.parametrize( + "make_index_type", + [ + lambda: "BITMAP", + lambda: index_mod.IndexConfig("bitmap", {}), + lambda: "ZONEMAP", + lambda: index_mod.IndexConfig("zonemap", {}), + ], + ids=["bitmap-string", "bitmap-config", "zonemap-string", "zonemap-config"], +) +def test_scalar_index_validates_ordered_field( + monkeypatch, make_index_type +): + """Invalid ordered-index fields fail on the driver for both input forms.""" + + captured, _ = _patch_scalar_segment_workflow(monkeypatch) + + with pytest.raises(TypeError, match="must be int, float, bool, string"): + index_mod.create_scalar_index( + uri="memory://fake", + column="tags", + index_type=make_index_type(), + num_workers=2, + ) + + assert "map_kwargs" not in captured + + def test_create_index_passes_block_size_to_loads_and_handler(monkeypatch): """The vector index path should use block_size for driver and worker loads."""