From 05a91e2bfe1b0af9c505466b661a1e9a13b3f285 Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Wed, 29 Jul 2026 11:39:12 +0800 Subject: [PATCH 1/7] feat(index): support distributed LABEL_LIST index --- docs/src/distributed-indexing.md | 4 +- lance_ray/index.py | 10 ++++- tests/test_distributed_indexing.py | 67 ++++++++++++++++++++++++++++++ tests/test_vector_index_options.py | 40 +++++++++++++++--- 4 files changed, 112 insertions(+), 9 deletions(-) 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..50bc7702 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: @@ -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/tests/test_distributed_indexing.py b/tests/test_distributed_indexing.py index 2eeeac12..c4e7787a 100755 --- a/tests/test_distributed_indexing.py +++ b/tests/test_distributed_indexing.py @@ -1514,6 +1514,73 @@ def test_distributed_rtree_index_matches_baseline(self, temp_dir): assert "point_rtree_idx" in explain +class TestDistributedLabelListIndexing: + """Distributed LABEL_LIST indexing tests.""" + + def test_distributed_large_list_index_matches_baseline(self, temp_dir): + """Build one segment per fragment and verify a list-membership query.""" + rows_per_fragment = 8 + num_fragments = 3 + num_rows = rows_per_fragment * num_fragments + dataset = lance.write_dataset( + pa.table( + { + "id": pa.array(range(num_rows), type=pa.int32()), + "labels": pa.array( + [ + ["distributed"] if row_id % 2 == 0 else ["other"] + for row_id in range(num_rows) + ], + type=pa.large_list(pa.string()), + ), + } + ), + str(Path(temp_dir) / "distributed_label_list.lance"), + max_rows_per_file=rows_per_fragment, + ) + assert len(dataset.get_fragments()) == num_fragments + + index_name = "labels_idx" + indexed_dataset = lr.create_scalar_index( + uri=dataset.uri, + column="labels", + index_type="LABEL_LIST", + name=index_name, + replace=False, + num_workers=num_fragments, + num_segments=num_fragments, + ) + + index = next( + index + for index in indexed_dataset.describe_indices() + if index.name == index_name + ) + assert index.index_type == "LabelList" + assert len(index.segments) == num_fragments + + filter_expr = "array_has_any(labels, ['distributed'])" + indexed = indexed_dataset.scanner( + filter=filter_expr, + columns=["id", "labels"], + use_scalar_index=True, + ).to_table() + baseline = indexed_dataset.scanner( + filter=filter_expr, + columns=["id", "labels"], + use_scalar_index=False, + ).to_table() + + assert indexed.equals(baseline) + plan = indexed_dataset.scanner( + filter=filter_expr, + columns=["id"], + use_scalar_index=True, + ).explain_plan() + assert "ScalarIndexQuery" in plan + assert index_name in plan + + class TestOptimizeIndices: """Test cases for optimize_indices (incremental index optimization).""" diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index 8d3b4ae9..711ca88e 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())), ] ) @@ -447,7 +450,16 @@ def test_create_index_rejects_invalid_num_segments(monkeypatch): @pytest.mark.parametrize( "index_type", - ["BTREE", "BITMAP", "INVERTED", "FTS", "NGRAM", "BLOOMFILTER", "RTREE"], + [ + "BTREE", + "BITMAP", + "INVERTED", + "FTS", + "NGRAM", + "BLOOMFILTER", + "RTREE", + "LABEL_LIST", + ], ) def test_create_scalar_index_uses_segment_path(monkeypatch, index_type): """Migrated scalar indexes should use Lance's segment workflow.""" @@ -486,7 +498,12 @@ 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" + if index_type in {"INVERTED", "FTS", "NGRAM"}: + column = "text" + elif index_type == "LABEL_LIST": + column = "labels" + else: + column = "value" updated_dataset = index_mod.create_scalar_index( uri="memory://fake", column=column, @@ -502,6 +519,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.""" @@ -594,9 +622,9 @@ def fake_lance_dataset(*args, **kwargs): scalar_handler = index_mod._handle_fragment_index( dataset_uri="memory://fake", - column="value", - index_type="LABEL_LIST", - name="value_idx", + column="text", + index_type="NGRAM", + name="text_idx", index_uuid="scalar-index", replace=False, train=True, From bf8fd630b0b05def938cfbe32a982c7a7b6339f7 Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Wed, 29 Jul 2026 13:18:44 +0800 Subject: [PATCH 2/7] fix(index): align LABEL_LIST segment validation --- lance_ray/index.py | 18 +++++++++--------- tests/test_distributed_indexing.py | 22 ++++++++++++++++------ tests/test_vector_index_options.py | 24 +++++++++++++++++++++--- 3 files changed, 46 insertions(+), 18 deletions(-) diff --git a/lance_ray/index.py b/lance_ray/index.py index 50bc7702..52890cc5 100755 --- a/lance_ray/index.py +++ b/lance_ray/index.py @@ -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,18 +596,18 @@ 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 + if _scalar_index_type_name(index_type) == "LABEL_LIST" and 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}" + ) + use_segment_workflow = ( _scalar_index_type_name(index_type) in _SCALAR_SEGMENT_INDEX_TYPES ) diff --git a/tests/test_distributed_indexing.py b/tests/test_distributed_indexing.py index c4e7787a..ee068ce9 100755 --- a/tests/test_distributed_indexing.py +++ b/tests/test_distributed_indexing.py @@ -1517,20 +1517,29 @@ def test_distributed_rtree_index_matches_baseline(self, temp_dir): class TestDistributedLabelListIndexing: """Distributed LABEL_LIST indexing tests.""" - def test_distributed_large_list_index_matches_baseline(self, temp_dir): - """Build one segment per fragment and verify a list-membership query.""" + def test_distributed_large_list_index_preserves_nullable_query_semantics( + self, temp_dir + ): + """Build one segment per fragment and preserve nullable list semantics.""" rows_per_fragment = 8 num_fragments = 3 num_rows = rows_per_fragment * num_fragments + labels_per_fragment = [ + ["distributed", "shared"], + ["other", None], + None, + [], + ["distributed"], + ["shared", "other"], + [None], + ["other"], + ] dataset = lance.write_dataset( pa.table( { "id": pa.array(range(num_rows), type=pa.int32()), "labels": pa.array( - [ - ["distributed"] if row_id % 2 == 0 else ["other"] - for row_id in range(num_rows) - ], + labels_per_fragment * num_fragments, type=pa.large_list(pa.string()), ), } @@ -1572,6 +1581,7 @@ def test_distributed_large_list_index_matches_baseline(self, temp_dir): ).to_table() assert indexed.equals(baseline) + assert indexed.column("id").to_pylist() == [0, 4, 8, 12, 16, 20] plan = indexed_dataset.scanner( filter=filter_expr, columns=["id"], diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index 711ca88e..b3f0c7f1 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") @@ -519,14 +525,26 @@ def fake_map_async_with_pool(**kwargs): assert fake_dataset.commit_kwargs["segments"] == ["segment"] -def test_create_label_list_index_rejects_non_list_column(): +@pytest.mark.parametrize( + "index_type", + ["LABEL_LIST", index_mod.IndexConfig("LABEL_LIST", {})], + ids=["string", "index-config"], +) +def test_create_label_list_index_rejects_non_list_column(monkeypatch, index_type): """LABEL_LIST should reject invalid columns before Ray workers start.""" + def fail_if_workers_are_scheduled(**kwargs): + raise AssertionError("Ray workers should not be scheduled") + + monkeypatch.setattr( + index_mod, "_map_async_with_pool", fail_if_workers_are_scheduled + ) + 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", + index_type=index_type, ) From 73b089e6fbfe9a6b8bb76c181395e66065e654f4 Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Wed, 29 Jul 2026 13:31:56 +0800 Subject: [PATCH 3/7] fix(index): keep LABEL_LIST validation string-only --- lance_ray/index.py | 16 ++++++++-------- tests/test_vector_index_options.py | 24 +++--------------------- 2 files changed, 11 insertions(+), 29 deletions(-) diff --git a/lance_ray/index.py b/lance_ray/index.py index 52890cc5..b36b4e3a 100755 --- a/lance_ray/index.py +++ b/lance_ray/index.py @@ -596,18 +596,18 @@ 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 - if _scalar_index_type_name(index_type) == "LABEL_LIST" and 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}" - ) - use_segment_workflow = ( _scalar_index_type_name(index_type) in _SCALAR_SEGMENT_INDEX_TYPES ) diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index b3f0c7f1..711ca88e 100644 --- a/tests/test_vector_index_options.py +++ b/tests/test_vector_index_options.py @@ -20,13 +20,7 @@ def _load_index_module_with_stubs(): lance_dataset = ModuleType("lance.dataset") lance_dataset.Index = type("Index", (), {}) - - class IndexConfig: - def __init__(self, index_type, parameters): - self.index_type = index_type - self.parameters = parameters - - lance_dataset.IndexConfig = IndexConfig + lance_dataset.IndexConfig = type("IndexConfig", (), {}) lance_dataset.LanceDataset = object lance_indices = ModuleType("lance.indices") @@ -525,26 +519,14 @@ def fake_map_async_with_pool(**kwargs): assert fake_dataset.commit_kwargs["segments"] == ["segment"] -@pytest.mark.parametrize( - "index_type", - ["LABEL_LIST", index_mod.IndexConfig("LABEL_LIST", {})], - ids=["string", "index-config"], -) -def test_create_label_list_index_rejects_non_list_column(monkeypatch, index_type): +def test_create_label_list_index_rejects_non_list_column(): """LABEL_LIST should reject invalid columns before Ray workers start.""" - def fail_if_workers_are_scheduled(**kwargs): - raise AssertionError("Ray workers should not be scheduled") - - monkeypatch.setattr( - index_mod, "_map_async_with_pool", fail_if_workers_are_scheduled - ) - with pytest.raises(TypeError, match="must be list or large list type"): index_mod.create_scalar_index( uri=_FakeDataset(), column="value", - index_type=index_type, + index_type="LABEL_LIST", ) From f39a394aca2d7cc047ae50fcadb6bbbd40104f67 Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Wed, 29 Jul 2026 13:58:50 +0800 Subject: [PATCH 4/7] test(index): restore block size handler fixture --- tests/test_vector_index_options.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index 711ca88e..edc288a8 100644 --- a/tests/test_vector_index_options.py +++ b/tests/test_vector_index_options.py @@ -622,9 +622,9 @@ def fake_lance_dataset(*args, **kwargs): scalar_handler = index_mod._handle_fragment_index( dataset_uri="memory://fake", - column="text", - index_type="NGRAM", - name="text_idx", + column="value", + index_type="LABEL_LIST", + name="value_idx", index_uuid="scalar-index", replace=False, train=True, From 0f3946898457a18c0cae34fd1be7135822532ac8 Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Thu, 30 Jul 2026 11:00:25 +0800 Subject: [PATCH 5/7] chore: update pylance to v10.0.0-beta.7 --- pyproject.toml | 2 +- uv.lock | 16 ++++++++-------- 2 files changed, 9 insertions(+), 9 deletions(-) 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/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]] From 09eb31a1c5ae8d5789d2693f05db47e2cf91c059 Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Thu, 30 Jul 2026 14:04:44 +0800 Subject: [PATCH 6/7] test(index): expect minimal field path quoting --- tests/test_distributed_indexing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_distributed_indexing.py b/tests/test_distributed_indexing.py index ee068ce9..d208f4d6 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", From 62799fb35694417da763b37f504560877d886bc8 Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Thu, 30 Jul 2026 15:29:11 +0800 Subject: [PATCH 7/7] test(index): merge LABEL_LIST segment coverage --- tests/test_distributed_indexing.py | 100 +++++++---------------------- tests/test_vector_index_options.py | 26 +++----- 2 files changed, 33 insertions(+), 93 deletions(-) diff --git a/tests/test_distributed_indexing.py b/tests/test_distributed_indexing.py index d208f4d6..7cc0686c 100755 --- a/tests/test_distributed_indexing.py +++ b/tests/test_distributed_indexing.py @@ -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( @@ -1514,83 +1537,6 @@ def test_distributed_rtree_index_matches_baseline(self, temp_dir): assert "point_rtree_idx" in explain -class TestDistributedLabelListIndexing: - """Distributed LABEL_LIST indexing tests.""" - - def test_distributed_large_list_index_preserves_nullable_query_semantics( - self, temp_dir - ): - """Build one segment per fragment and preserve nullable list semantics.""" - rows_per_fragment = 8 - num_fragments = 3 - num_rows = rows_per_fragment * num_fragments - labels_per_fragment = [ - ["distributed", "shared"], - ["other", None], - None, - [], - ["distributed"], - ["shared", "other"], - [None], - ["other"], - ] - dataset = lance.write_dataset( - pa.table( - { - "id": pa.array(range(num_rows), type=pa.int32()), - "labels": pa.array( - labels_per_fragment * num_fragments, - type=pa.large_list(pa.string()), - ), - } - ), - str(Path(temp_dir) / "distributed_label_list.lance"), - max_rows_per_file=rows_per_fragment, - ) - assert len(dataset.get_fragments()) == num_fragments - - index_name = "labels_idx" - indexed_dataset = lr.create_scalar_index( - uri=dataset.uri, - column="labels", - index_type="LABEL_LIST", - name=index_name, - replace=False, - num_workers=num_fragments, - num_segments=num_fragments, - ) - - index = next( - index - for index in indexed_dataset.describe_indices() - if index.name == index_name - ) - assert index.index_type == "LabelList" - assert len(index.segments) == num_fragments - - filter_expr = "array_has_any(labels, ['distributed'])" - indexed = indexed_dataset.scanner( - filter=filter_expr, - columns=["id", "labels"], - use_scalar_index=True, - ).to_table() - baseline = indexed_dataset.scanner( - filter=filter_expr, - columns=["id", "labels"], - use_scalar_index=False, - ).to_table() - - assert indexed.equals(baseline) - assert indexed.column("id").to_pylist() == [0, 4, 8, 12, 16, 20] - plan = indexed_dataset.scanner( - filter=filter_expr, - columns=["id"], - use_scalar_index=True, - ).explain_plan() - assert "ScalarIndexQuery" in plan - assert index_name in plan - - class TestOptimizeIndices: """Test cases for optimize_indices (incremental index optimization).""" diff --git a/tests/test_vector_index_options.py b/tests/test_vector_index_options.py index edc288a8..b114ef70 100644 --- a/tests/test_vector_index_options.py +++ b/tests/test_vector_index_options.py @@ -449,19 +449,19 @@ def test_create_index_rejects_invalid_num_segments(monkeypatch): @pytest.mark.parametrize( - "index_type", + ("index_type", "column"), [ - "BTREE", - "BITMAP", - "INVERTED", - "FTS", - "NGRAM", - "BLOOMFILTER", - "RTREE", - "LABEL_LIST", + ("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": []} @@ -498,12 +498,6 @@ def fake_map_async_with_pool(**kwargs): ) monkeypatch.setattr(index_mod, "_map_async_with_pool", fake_map_async_with_pool) - if index_type in {"INVERTED", "FTS", "NGRAM"}: - column = "text" - elif index_type == "LABEL_LIST": - column = "labels" - else: - column = "value" updated_dataset = index_mod.create_scalar_index( uri="memory://fake", column=column,