Skip to content

fix(index): align distributed scalar index parameters #5250

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ddupg wants to merge 1 commit into
base: main
Choose a base branch
from
+272 −91
Open

fix(index): align distributed scalar index parameters #5250

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 2 additions & 4 deletions docs/src/distributed-indexing.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
],
Expand Down Expand Up @@ -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` |
Expand Down Expand Up @@ -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.
169 changes: 91 additions & 78 deletions lance_ray/index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand All @@ -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).
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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"
Expand Down
9 changes: 6 additions & 3 deletions tests/test_distributed_indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading