diff --git a/docs/src/distributed-indexing.md b/docs/src/distributed-indexing.md index 382dd25c..c2146d3c 100755 --- a/docs/src/distributed-indexing.md +++ b/docs/src/distributed-indexing.md @@ -7,7 +7,7 @@ Lance-Ray provides distributed index building functionality that leverages Ray's ### Scalar Indexing -`create_scalar_index()` - Distributedly create scalar index using ray. Currently only Inverted/FTS/BTREE/BITMAP/NGRAM/ZONEMAP/BLOOMFILTER/RTREE are supported. Will add more index type support in the future. +`create_scalar_index()` - Distributedly create scalar index using ray. Currently Inverted/FTS/BTREE/BITMAP/LABEL_LIST/NGRAM/ZONEMAP/BLOOMFILTER/RTREE are supported. Will add more index type support in the future. To construct GeoArrow data for an RTREE index, install the PyLance geo extra: @@ -77,7 +77,7 @@ def create_scalar_index( | `ray_remote_args` | `Dict[str, Any]`, optional | Ray task options (e.g., `num_cpus`, `resources`) | | `**kwargs` | `Any` | Additional arguments passed to `create_scalar_index` | -**Note:** For distributed scalar indexing, currently only `"INVERTED"`, `"FTS"`, `"BTREE"`, `"BITMAP"`, `"NGRAM"`, `"ZONEMAP"`, `"BLOOMFILTER"` and `"RTREE"` index types are supported. +**Note:** For distributed scalar indexing, currently `"INVERTED"`, `"FTS"`, `"BTREE"`, `"BITMAP"`, `"LABEL_LIST"`, `"NGRAM"`, `"ZONEMAP"`, `"BLOOMFILTER"`, and `"RTREE"` index types are supported. #### Return Value diff --git a/lance_ray/index.py b/lance_ray/index.py index fc3d8c8c..b36b4e3a 100755 --- a/lance_ray/index.py +++ b/lance_ray/index.py @@ -225,7 +225,7 @@ def _build_rabitq_model(*, dimension: int, num_bits: int = 1) -> str: "RTREE", ] _SCALAR_INDEX_TYPES = get_args(_ScalarIndexType) -_SCALAR_SEGMENT_INDEX_TYPES = frozenset(_SCALAR_INDEX_TYPES) - {"LABEL_LIST"} +_SCALAR_SEGMENT_INDEX_TYPES = frozenset(_SCALAR_INDEX_TYPES) def _scalar_index_type_name(index_type: str | IndexConfig) -> str | None: @@ -473,7 +473,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 incompatible with the index type. RuntimeError: If index building fails or pylance version is incompatible. """ # Check pylance version compatibility @@ -596,6 +596,14 @@ def create_scalar_index( f"Column {column} must be numeric or string type for " f"{index_type} index, got {value_type}" ) + case "LABEL_LIST": + if not ( + pa.types.is_list(field.type) or pa.types.is_large_list(field.type) + ): + raise TypeError( + f"Column {column} must be list or large list type for " + f"LABEL_LIST index, got {field.type}" + ) case _: # For other index types, skip strict validation to maintain compatibility pass diff --git a/pyproject.toml b/pyproject.toml index 48177918..f40fac68 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ dependencies = [ "ray[data]>=2.41.0", - "pylance>=10.0.0b5", + "pylance>=10.0.0b7", "lance-namespace", "packaging", "pyarrow>=17.0.0", diff --git a/tests/test_distributed_indexing.py b/tests/test_distributed_indexing.py index 2eeeac12..7cc0686c 100755 --- a/tests/test_distributed_indexing.py +++ b/tests/test_distributed_indexing.py @@ -583,7 +583,7 @@ def test_build_distributed_nested_scalar_indexes(self, temp_dir): indices = {idx.name: idx for idx in updated_dataset.describe_indices()} assert indices["nested_text_idx"].field_names == ["meta.text"] assert indices["literal_dot_text_idx"].field_names == ["meta.`a.b`"] - assert indices["hyphen_user_id_idx"].field_names == ["`meta-data`.`user-id`"] + assert indices["hyphen_user_id_idx"].field_names == ["meta-data.user-id"] nested_results = updated_dataset.scanner( full_text_query="nestedthree", @@ -1395,6 +1395,29 @@ class TestDistributedScalarSegmentIndexes: ["value = 4", "value IN (1, 6, 11)"], id="bloomfilter", ), + pytest.param( + "LABEL_LIST", + "labels", + [ + ["distributed", "shared"], + ["other", None], + None, + [], + ["distributed"], + ["shared", "other"], + [None], + ["other"], + ["distributed", "shared"], + ["other", None], + None, + [], + ], + pa.large_list(pa.string()), + "labels_idx", + "LabelList", + ["array_has_any(labels, ['distributed'])"], + id="label-list", + ), ], ) def test_filter_index_matches_baseline( diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index 8d3b4ae9..b114ef70 100644 --- a/tests/test_vector_index_options.py +++ b/tests/test_vector_index_options.py @@ -69,7 +69,7 @@ def id(self): class _FakeLanceSchema: def field(self, column): - if column not in {"value", "text"}: + if column not in {"value", "text", "labels"}: raise KeyError(column) return _FakeLanceField() @@ -82,6 +82,8 @@ def field(self, column): return _FakeField(column, index_mod.pa.int64()) if column == "text": return _FakeField(column, index_mod.pa.string()) + if column == "labels": + return _FakeField(column, index_mod.pa.list_(index_mod.pa.string())) else: raise KeyError(column) @@ -91,6 +93,7 @@ def __iter__(self): _FakeField("vector"), _FakeField("value", index_mod.pa.int64()), _FakeField("text", index_mod.pa.string()), + _FakeField("labels", index_mod.pa.list_(index_mod.pa.string())), ] ) @@ -446,10 +449,19 @@ def test_create_index_rejects_invalid_num_segments(monkeypatch): @pytest.mark.parametrize( - "index_type", - ["BTREE", "BITMAP", "INVERTED", "FTS", "NGRAM", "BLOOMFILTER", "RTREE"], + ("index_type", "column"), + [ + ("BTREE", "value"), + ("BITMAP", "value"), + ("INVERTED", "text"), + ("FTS", "text"), + ("NGRAM", "text"), + ("BLOOMFILTER", "value"), + ("RTREE", "value"), + ("LABEL_LIST", "labels"), + ], ) -def test_create_scalar_index_uses_segment_path(monkeypatch, index_type): +def test_create_scalar_index_uses_segment_path(monkeypatch, index_type, column): """Migrated scalar indexes should use Lance's segment workflow.""" captured = {"loads": []} @@ -486,7 +498,6 @@ def fake_map_async_with_pool(**kwargs): ) monkeypatch.setattr(index_mod, "_map_async_with_pool", fake_map_async_with_pool) - column = "text" if index_type in {"INVERTED", "FTS", "NGRAM"} else "value" updated_dataset = index_mod.create_scalar_index( uri="memory://fake", column=column, @@ -502,6 +513,17 @@ def fake_map_async_with_pool(**kwargs): assert fake_dataset.commit_kwargs["segments"] == ["segment"] +def test_create_label_list_index_rejects_non_list_column(): + """LABEL_LIST should reject invalid columns before Ray workers start.""" + + with pytest.raises(TypeError, match="must be list or large list type"): + index_mod.create_scalar_index( + uri=_FakeDataset(), + column="value", + index_type="LABEL_LIST", + ) + + 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.""" diff --git a/uv.lock b/uv.lock index 2f2f8ccd..d6bc1032 100644 --- a/uv.lock +++ b/uv.lock @@ -501,7 +501,7 @@ requires-dist = [ { name = "more-itertools", marker = "python_full_version < '3.12'", specifier = ">=2.6.0" }, { name = "packaging" }, { name = "pyarrow", specifier = ">=17.0.0" }, - { name = "pylance", specifier = ">=10.0.0b5" }, + { name = "pylance", specifier = ">=10.0.0b7" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.4.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=1.0.0" }, { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=5.0.0" }, @@ -1138,7 +1138,7 @@ wheels = [ [[package]] name = "pylance" -version = "10.0.0b6" +version = "10.0.0b7" source = { registry = "https://pypi.fury.io/lance-format" } dependencies = [ { name = "lance-namespace" }, @@ -1147,12 +1147,12 @@ dependencies = [ { name = "pyarrow" }, ] wheels = [ - { url = "https://pypi.fury.io/lance-format/-/ver_1S1urO/pylance-10.0.0b6-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:0529e466b3545eedef1dc55acfd41014f9a051eccc129b1688045dbe142e45a7" }, - { url = "https://pypi.fury.io/lance-format/-/ver_2h9KXm/pylance-10.0.0b6-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f741e3ae2b8c928675477ca7e012a74be280009178f22b80c2ec8ee8fbc5b32e" }, - { url = "https://pypi.fury.io/lance-format/-/ver_PFSbW/pylance-10.0.0b6-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c623e0ae1220d769827b4b0aa64f1bcd00cb18e4a81ecf8f7617bc00e359fbe6" }, - { url = "https://pypi.fury.io/lance-format/-/ver_vww2D/pylance-10.0.0b6-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:dfb232cc641f06799659c1a9d57969e1f920be8f9208ad2c29fce8edf31b8f16" }, - { url = "https://pypi.fury.io/lance-format/-/ver_ECyWk/pylance-10.0.0b6-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b683832edd828691bbbe55d97c36d0859ea9734c673219c12c740f817f9442fa" }, - { url = "https://pypi.fury.io/lance-format/-/ver_ub84x/pylance-10.0.0b6-cp310-abi3-win_amd64.whl", hash = "sha256:831ee58fd98c1d1306e93b874878bc0d5da320e7e04f5b60d94535816b46b490" }, + { url = "https://pypi.fury.io/lance-format/-/ver_F25tX/pylance-10.0.0b7-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:3c91750f59df60157ae2a784195358ac6e87c312348bab6aa0d3074daf570b8c" }, + { url = "https://pypi.fury.io/lance-format/-/ver_285dmG/pylance-10.0.0b7-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:248ecb3d46323c2c9363e030b82f9973daa6d9df1044fc4dafe513fc83780ad8" }, + { url = "https://pypi.fury.io/lance-format/-/ver_kRD7S/pylance-10.0.0b7-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e4f60a028e8a2bbddfae1cbf288a47375e1c2bfd7a1bb26e0a1da53f7f4f166a" }, + { url = "https://pypi.fury.io/lance-format/-/ver_RywJJ/pylance-10.0.0b7-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0cd91ad4a1d240c6a8cf399233a5f9c1c050e035fc960766ebc584dfadcc817e" }, + { url = "https://pypi.fury.io/lance-format/-/ver_UyK0D/pylance-10.0.0b7-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1fbd4a40a9dc17104abe467870aa89521a1bdeed62c80547e326cc353f2ef0a3" }, + { url = "https://pypi.fury.io/lance-format/-/ver_1UP4Pu/pylance-10.0.0b7-cp310-abi3-win_amd64.whl", hash = "sha256:23b44b64270f7ac11c75d8a084e94867708541722da5405531a6caa61783680e" }, ] [[package]]