Skip to content

refactor: drop pylance version checks from indexing and search #5304

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
jonasdedden wants to merge 1 commit into
base: main
Choose a base branch
from
Open

refactor: drop pylance version checks from indexing and search #5304

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
93 changes: 6 additions & 87 deletions lance_ray/index.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
from lance.bitmap import Bitmap
from lance.dataset import Index, IndexConfig, LanceDataset
from lance.indices import IndicesBuilder
from packaging import version
from ray.exceptions import ActorUnavailableError

from .field_path import resolve_arrow_field_path, resolve_dataset_field_path
Expand Down Expand Up @@ -484,21 +483,6 @@ def func(fragment_ids: list[int]) -> dict[str, Any]:
return func


def merge_index_metadata_compat(
dataset: LanceDataset,
index_id: str,
index_type: str,
**kwargs: Any,
) -> Any:
"""Call ``merge_index_metadata`` with backwards compatible signature."""
try:
return dataset.merge_index_metadata(
index_id, index_type, batch_readhead=kwargs.get("batch_readhead")
)
except TypeError:
return dataset.merge_index_metadata(index_id) # type: ignore[call-arg]


def create_scalar_index(
uri: Optional[str | lance.LanceDataset] = None,
*,
Expand Down Expand Up @@ -556,30 +540,8 @@ def create_scalar_index(
Raises:
ValueError: If input parameters are invalid.
TypeError: If the column type is incompatible with the index type.
RuntimeError: If index building fails or pylance version is incompatible.
RuntimeError: If index building fails.
"""
# Check pylance version compatibility
try:
lance_version = version.parse(lance.__version__)
min_required_version = version.parse("0.36.0")

if lance_version < min_required_version:
raise RuntimeError(
"Distributed indexing requires pylance >= 0.36.0, but found "
f"{lance.__version__}. The distribute-related interfaces are "
"not available in older versions. Please upgrade pylance by "
"running: pip install --upgrade pylance"
)

logger.info("Pylance version check passed: %s >= 0.36.0", lance.__version__)

except AttributeError as err: # pragma: no cover - defensive
raise RuntimeError(
"Cannot determine pylance version. Distributed indexing requires "
"pylance >= 0.36.0. Please upgrade pylance by running: "
"pip install --upgrade pylance"
) from err

index_id = str(uuid.uuid4())
logger.info("Starting distributed scalar index build with ID: %s", index_id)

Expand Down Expand Up @@ -842,7 +804,9 @@ def create_fragment_handler() -> Any:
# Convert IndexConfig to string for merge_index_metadata which expects a string
# (lance's create_scalar_index converts IndexConfig to "scalar" internally)
index_type_str = "scalar" if isinstance(index_type, IndexConfig) else index_type
merge_index_metadata_compat(dataset, index_id, index_type=index_type_str, **kwargs)
dataset.merge_index_metadata(
index_id, index_type_str, batch_readhead=kwargs.get("batch_readhead")
)

logger.info("Phase 3: Creating and committing scalar index '%s'", name)

Expand Down Expand Up @@ -1163,31 +1127,6 @@ def _normalize_index_type(index_type: Any) -> str:
return index_type_name


def _check_pylance_version() -> None:
"""Ensure pylance (lance) provides distributed vector APIs."""

try:
lance_version = version.parse(lance.__version__)
min_required_version = version.parse("0.36.0")

if lance_version < min_required_version:
raise RuntimeError(
"Distributed vector indexing requires pylance >= 0.36.0, but found "
f"{lance.__version__}. The distributed vector interfaces are not "
"available in older versions. Please upgrade pylance by running: "
"pip install --upgrade pylance"
)

logger.info("Pylance version check passed: %s >= 0.36.0", lance.__version__)

except AttributeError as err: # pragma: no cover - defensive
raise RuntimeError(
"Cannot determine pylance version. Distributed vector indexing requires "
"pylance >= 0.36.0. Please upgrade pylance by running: "
"pip install --upgrade pylance"
) from err


def _validate_metric(metric: str) -> str:
"""Normalize and validate the distance metric string."""

Expand Down Expand Up @@ -1490,8 +1429,6 @@ def create_index(
Updated Lance dataset with the index created
"""

_check_pylance_version()

if not column:
raise ValueError("Column name cannot be empty")

Expand Down Expand Up @@ -1762,8 +1699,6 @@ def optimize_indices(

Raises:
ValueError: If input parameters are invalid.
RuntimeError: If optimize_indices is not supported by the current
lance version or if the operation fails.

Example:
>>> import lance_ray
Expand Down Expand Up @@ -1805,23 +1740,7 @@ def optimize_indices(
storage_options=merged_storage_options,
**namespace_kwargs,
)
logger.info(
"Loaded dataset: uri=%s, version=%s",
uri,
getattr(dataset, "version", "unknown"),
)

if not hasattr(dataset, "optimize"):
raise RuntimeError(
"LanceDataset has no 'optimize' property. Please ensure "
"lance is installed with a version that provides DatasetOptimizer."
)
optimizer = dataset.optimize
if not hasattr(optimizer, "optimize_indices"):
raise RuntimeError(
"optimize_indices is not available on DatasetOptimizer. Please ensure "
"lance is installed with a version that provides optimize_indices."
)
logger.info("Loaded dataset: uri=%s, version=%s", uri, dataset.version)

call_kw: dict[str, Any] = {
"num_indices_to_merge": num_indices_to_merge,
Expand All @@ -1835,7 +1754,7 @@ def optimize_indices(
"Calling DatasetOptimizer.optimize_indices with: %s",
{k: v for k, v in call_kw.items() if k != "storage_options"},
)
optimizer.optimize_indices(**call_kw)
dataset.optimize.optimize_indices(**call_kw)
logger.info(
"optimize_indices completed successfully for dataset uri=%s",
uri,
Expand Down
39 changes: 3 additions & 36 deletions lance_ray/search.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@

from __future__ import annotations

import inspect
import logging
import math
import pickle
Expand Down Expand Up @@ -59,20 +58,6 @@ def _dataset_load_kwargs(
return kwargs


def _get_dataset_storage_options(dataset: LanceDataset) -> dict[str, Any]:
try:
return dataset.initial_storage_options or {}
except AttributeError:
return getattr(dataset, "_storage_options", None) or {}


def _get_fragment_id(fragment: Any) -> int:
try:
return cast(int, fragment.fragment_id)
except AttributeError:
return cast(int, fragment.metadata.id)


def _index_value(index: Any, name: str, default: Any = None) -> Any:
if isinstance(index, dict):
return index.get(name, default)
Expand Down Expand Up @@ -133,13 +118,13 @@ def _plan_vector_search(
num_workers: int,
include_unindexed: bool,
) -> list[_SearchPlan]:
fragment_ids = {_get_fragment_id(fragment) for fragment in fragments}
fragment_ids = {fragment.fragment_id for fragment in fragments}
if not fragment_ids:
return []

fragment_weights: dict[int, int] = {}
for fragment in fragments:
fragment_id = _get_fragment_id(fragment)
fragment_id = fragment.fragment_id
try:
fragment_weights[fragment_id] = fragment.count_rows()
except Exception: # pragma: no cover - defensive fallback
Expand Down Expand Up @@ -268,13 +253,6 @@ def _execute_vector_search_plan(
analyze_plan=analyze_plan,
)

if not _scanner_accepts_index_segments(dataset):
raise RuntimeError(
"The installed pylance scanner does not support index_segments, "
"which is required for distributed indexed vector search plans. "
"Upgrade pylance or run without an indexed plan."
)

scanner_options = dict(base_scanner_options)
search_nearest = dict(nearest)
search_nearest["k"] = candidate_k
Expand All @@ -295,17 +273,6 @@ def _execute_vector_search_plan(
return scanner.to_table()


def _scanner_accepts_index_segments(dataset: LanceDataset) -> bool:
try:
parameters = inspect.signature(dataset.scanner).parameters
except (TypeError, ValueError): # pragma: no cover - defensive
return True
return "index_segments" in parameters or any(
parameter.kind == inspect.Parameter.VAR_KEYWORD
for parameter in parameters.values()
)


def _execute_flat_fallback_vector_search_plan(
dataset: LanceDataset,
*,
Expand Down Expand Up @@ -705,7 +672,7 @@ def vector_search(
else:
dataset = uri
if not merged_storage_options:
merged_storage_options.update(_get_dataset_storage_options(dataset))
merged_storage_options.update(dataset.initial_storage_options or {})

try:
resolved_column = resolve_arrow_field_path(dataset.schema, column)
Expand Down
Loading
Loading