test_distributed_indexing.py

"""Test cases for lance_ray.indexing module."""

import random
import tempfile
from pathlib import Path

import lance
import lance_ray as lr
import numpy as np
import pyarrow as pa
import pytest
import ray
from lance_ray.search import _scanner_accepts_index_segments
from packaging import version

import pandas as pd


def check_lance_version_compatibility():
    """Check if lance version supports distributed indexing."""
    try:
        lance_version = version.parse(lance.__version__)
        min_required_version = version.parse("0.36.0")
        return lance_version >= min_required_version
    except (AttributeError, Exception):
        return False


# Skip all distributed indexing tests if lance version is incompatible
pytestmark = pytest.mark.skipif(
    not check_lance_version_compatibility(),
    reason="Distributed indexing requires pylance >= 0.36.0. Current version: {}".format(
        getattr(lance, "__version__", "unknown")
    ),
)


@pytest.fixture
def text_data():
    """Create sample text data for indexing tests."""
    return pd.DataFrame(
        {
            "id": [1, 2, 3, 4, 5, 6, 7, 8],
            "text": [
                "The quick brown fox jumps over the lazy dog",
                "Python is a powerful programming language",
                "Machine learning algorithms are fascinating",
                "Data science requires statistical knowledge",
                "Natural language processing uses text analysis",
                "Distributed computing scales horizontally",
                "Ray framework enables parallel processing",
                "Lance format provides efficient storage",
            ],
            "category": [
                "animals",
                "tech",
                "ml",
                "data",
                "nlp",
                "distributed",
                "ray",
                "storage",
            ],
        }
    )


@pytest.fixture
def temp_dir():
    """Create a temporary directory for testing."""
    with tempfile.TemporaryDirectory() as temp_dir:
        yield temp_dir


@pytest.fixture
def text_dataset(text_data):
    """Create a Ray Dataset from text data."""
    return ray.data.from_pandas(text_data)


@pytest.fixture
def multi_fragment_lance_dataset(text_dataset, temp_dir):
    """Create a Lance dataset with multiple fragments for testing."""
    path = Path(temp_dir) / "multi_fragment_text.lance"
    # Create dataset with multiple fragments (2 rows per fragment)
    lr.write_lance(text_dataset, str(path), min_rows_per_file=2, max_rows_per_file=2)
    return str(path)


def generate_multi_fragment_dataset(tmp_path, num_fragments=4, rows_per_fragment=250):
    """Generate a test dataset with multiple fragments."""
    all_data = []
    for frag_idx in range(num_fragments):
        for row_idx in range(rows_per_fragment):
            row_id = frag_idx * rows_per_fragment + row_idx
            all_data.append(
                {
                    "id": row_id,
                    "text": f"This is test document {row_id} with some sample text content for fragment {frag_idx}",
                    "fragment_id": frag_idx,
                }
            )

    df = pd.DataFrame(all_data)
    dataset = ray.data.from_pandas(df)

    path = Path(tmp_path) / "large_multi_fragment.lance"
    lr.write_lance(
        dataset,
        str(path),
        min_rows_per_file=rows_per_fragment,
        max_rows_per_file=rows_per_fragment,
    )

    return lance.dataset(str(path))


def generate_mixed_schema_dataset(
    tmp_path,
    num_rows: int = 200,
    vector_dim: int = 8,
    rows_per_fragment: int = 50,
):
    """Generate a Lance dataset with both scalar and vector columns.

    Schema: id (int64), vector (fixed-size list float32), label (int64), score (float64).
    Used to test creating a scalar index on a dataset that also has a vector column.
    """
    ids = pa.array(range(num_rows), type=pa.int64())
    vectors = np.random.randn(num_rows, vector_dim).astype(np.float32)
    vector_values = pa.array(vectors.ravel(), type=pa.float32())
    vector_array = pa.FixedSizeListArray.from_arrays(vector_values, vector_dim)
    labels = pa.array(
        np.random.randint(0, 10, size=num_rows),
        type=pa.int64(),
    )
    scores = pa.array(
        np.random.uniform(0, 100, size=num_rows),
        type=pa.float64(),
    )
    tbl = pa.table(
        {
            "id": ids,
            "vector": vector_array,
            "label": labels,
            "score": scores,
        }
    )
    dataset = ray.data.from_arrow(tbl)
    path = Path(tmp_path) / "mixed_schema.lance"
    lr.write_lance(
        dataset,
        str(path),
        min_rows_per_file=rows_per_fragment,
        max_rows_per_file=rows_per_fragment,
    )
    return str(path)


def generate_nested_contract_dataset(tmp_path, rows_per_fragment: int = 2):
    """Generate a multi-fragment dataset with nested field-path edge cases."""
    schema = pa.schema(
        [
            pa.field("id", pa.int64()),
            pa.field(
                "meta",
                pa.struct(
                    [
                        pa.field("text", pa.string()),
                        pa.field("a.b", pa.string()),
                    ]
                ),
            ),
            pa.field(
                "meta-data",
                pa.struct([pa.field("user-id", pa.int64())]),
            ),
            pa.field("outer", pa.struct([pa.field("leaf", pa.int64())])),
            pa.field("other", pa.struct([pa.field("leaf", pa.int64())])),
        ]
    )
    table = pa.Table.from_arrays(
        [
            pa.array([1, 2, 3, 4], type=pa.int64()),
            pa.array(
                [
                    {"text": "nestedone", "a.b": "literalone"},
                    {"text": "nestedtwo", "a.b": "literaltwo"},
                    {"text": "nestedthree", "a.b": "literalthree"},
                    {"text": "nestedfour", "a.b": "literalfour"},
                ],
                type=schema.field("meta").type,
            ),
            pa.array(
                [
                    {"user-id": 101},
                    {"user-id": 102},
                    {"user-id": 103},
                    {"user-id": 104},
                ],
                type=schema.field("meta-data").type,
            ),
            pa.array(
                [{"leaf": 10}, {"leaf": 20}, {"leaf": 30}, {"leaf": 40}],
                type=schema.field("outer").type,
            ),
            pa.array(
                [{"leaf": 40}, {"leaf": 30}, {"leaf": 20}, {"leaf": 10}],
                type=schema.field("other").type,
            ),
        ],
        schema=schema,
    )

    path = Path(tmp_path) / "nested_contract.lance"
    lr.write_lance(
        ray.data.from_arrow(table),
        str(path),
        min_rows_per_file=rows_per_fragment,
        max_rows_per_file=rows_per_fragment,
    )
    return str(path)


class TestDistributedIndexing:
    """Test cases for distributed indexing functionality."""

    def test_build_distributed_fts_index_basic(self, multi_fragment_lance_dataset):
        """Test basic distributed FTS index building."""
        dataset_uri = multi_fragment_lance_dataset

        # Build distributed index
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            num_workers=2,
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after building"

        # Find our index
        text_index = None
        for idx in indices:
            if "text" in idx.name:
                text_index = idx
                break

        assert text_index is not None, "Text index not found"
        assert text_index.index_type == "Inverted", (
            f"Expected Inverted index, got {text_index.index_type}"
        )

    def test_build_distributed_fts_index_with_name(self, multi_fragment_lance_dataset):
        """Test building distributed index with custom name."""
        dataset_uri = multi_fragment_lance_dataset
        custom_name = "custom_text_index"

        # Build distributed index with custom name
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            name=custom_name,
            num_workers=2,
        )

        # Verify the index was created with correct name
        indices = updated_dataset.describe_indices()
        index_names = [idx.name for idx in indices]
        assert custom_name in index_names, (
            f"Custom index name '{custom_name}' not found in {index_names}"
        )

    def test_build_distributed_fts_index_search_functionality(
        self, multi_fragment_lance_dataset
    ):
        """Test that the built index actually works for searching."""
        dataset_uri = multi_fragment_lance_dataset

        # Build distributed index
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            num_workers=2,
        )

        # Test full-text search functionality
        search_term = "Python"
        results = updated_dataset.scanner(
            full_text_query=search_term,
            columns=["id", "text"],
        ).to_table()

        # Should find at least one result containing "Python"
        assert results.num_rows > 0, f"No results found for search term '{search_term}'"

        # Verify results contain the search term
        text_results = results.column("text").to_pylist()
        assert any(search_term in text for text in text_results), (
            "Search results don't contain the search term"
        )

    def test_build_distributed_fts_index_fts_type(self, multi_fragment_lance_dataset):
        """Test building distributed FTS index."""
        dataset_uri = multi_fragment_lance_dataset

        # Build distributed FTS index
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            num_workers=2,
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after building"

    def test_build_distributed_fts_index_list_large_utf8(self, temp_dir):
        """Test distributed FTS index building on list<large_utf8> columns."""
        search_term = "needlelarge"
        table = pa.table(
            {
                "id": pa.array([1, 2, 3, 4], type=pa.int64()),
                "tags": pa.array(
                    [
                        ["alpha", "beta"],
                        ["distributed", search_term],
                        ["search", "fts"],
                        ["other", "tokens"],
                    ],
                    type=pa.list_(pa.large_string()),
                ),
            }
        )
        dataset = ray.data.from_arrow(table)
        path = Path(temp_dir) / "list_large_utf8_text.lance"
        lr.write_lance(dataset, str(path), min_rows_per_file=2, max_rows_per_file=2)

        updated_dataset = lr.create_scalar_index(
            uri=str(path),
            column="tags",
            index_type="INVERTED",
            num_workers=2,
        )

        results = updated_dataset.scanner(
            full_text_query=search_term,
            columns=["id", "tags"],
        ).to_table()

        assert results.num_rows == 1
        assert results.column("id").to_pylist() == [2]

    def test_build_distributed_index_large_dataset(self, temp_dir):
        """Test distributed indexing on a larger dataset with multiple fragments."""
        # Generate larger dataset
        dataset = generate_multi_fragment_dataset(
            temp_dir, num_fragments=4, rows_per_fragment=50
        )

        # Build distributed index
        updated_dataset = lr.create_scalar_index(
            uri=dataset.uri,
            column="text",
            index_type="INVERTED",
            num_workers=4,
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after building"

        # Test search functionality
        search_term = "test"
        results = updated_dataset.scanner(
            full_text_query=search_term,
            columns=["id", "text"],
        ).to_table()

        assert results.num_rows > 0, f"No results found for search term '{search_term}'"

    def test_build_distributed_index_invalid_column(self, multi_fragment_lance_dataset):
        """Test error handling for invalid column."""
        dataset_uri = multi_fragment_lance_dataset

        with pytest.raises(ValueError, match="Column 'nonexistent' not found"):
            lr.create_scalar_index(
                uri=dataset_uri,
                column="nonexistent",
                index_type="INVERTED",
                num_workers=2,
            )

    def test_build_distributed_index_invalid_index_type(
        self, multi_fragment_lance_dataset
    ):
        """Test error handling for invalid index type."""
        dataset_uri = multi_fragment_lance_dataset

        with pytest.raises(
            ValueError,
            match="Distributed indexing does not support index type 'INVALID'",
        ):
            lr.create_scalar_index(
                uri=dataset_uri,
                column="text",
                index_type="INVALID",
                num_workers=2,
            )

    def test_build_distributed_index_invalid_num_workers(
        self, multi_fragment_lance_dataset
    ):
        """Test error handling for invalid num_workers."""
        dataset_uri = multi_fragment_lance_dataset

        with pytest.raises(ValueError, match="num_workers must be positive"):
            lr.create_scalar_index(
                uri=dataset_uri,
                column="text",
                index_type="INVERTED",
                num_workers=0,
            )

    def test_build_distributed_index_empty_column(self, multi_fragment_lance_dataset):
        """Test error handling for empty column name."""
        dataset_uri = multi_fragment_lance_dataset

        with pytest.raises(ValueError, match="Column name cannot be empty"):
            lr.create_scalar_index(
                uri=dataset_uri,
                column="",
                index_type="INVERTED",
                num_workers=2,
            )

    def test_build_distributed_index_non_string_column(self, temp_dir):
        """Test error handling for non-string column."""
        # Create dataset with non-string column
        data = pd.DataFrame(
            {
                "id": [1, 2, 3, 4],
                "numeric_col": [10, 20, 30, 40],
                "text": ["text1", "text2", "text3", "text4"],
            }
        )
        dataset = ray.data.from_pandas(data)
        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, large string, list of strings, or json",
        ):
            lr.create_scalar_index(
                uri=str(path),
                column="numeric_col",
                index_type="INVERTED",
                num_workers=2,
            )

    def test_build_distributed_index_with_ray_remote_args(
        self, multi_fragment_lance_dataset
    ):
        """Test building distributed index with Ray options."""
        dataset_uri = multi_fragment_lance_dataset

        # Build distributed index with Ray options
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            num_workers=2,
            ray_remote_args={"num_cpus": 1},
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after building"

    def test_build_distributed_index_with_storage_options(
        self, multi_fragment_lance_dataset
    ):
        """Test building distributed index with storage options."""
        dataset_uri = multi_fragment_lance_dataset

        # Build distributed index with storage options
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            num_workers=2,
            storage_options={},  # Empty storage options should work
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after building"

    def test_build_distributed_index_with_kwargs(self, multi_fragment_lance_dataset):
        """Test building distributed index with additional kwargs."""
        dataset_uri = multi_fragment_lance_dataset

        # Build distributed index with additional kwargs
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            num_workers=2,
            remove_stop_words=False,  # Additional kwarg for create_scalar_index
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after building"

    def test_build_distributed_index_dataset_object(self, multi_fragment_lance_dataset):
        """Test building distributed index with Lance dataset object instead of URI."""
        dataset = lance.dataset(multi_fragment_lance_dataset)

        # Build distributed index using dataset object
        updated_dataset = lr.create_scalar_index(
            uri=dataset.uri,
            column="text",
            index_type="INVERTED",
            num_workers=2,
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after building"

    def test_build_distributed_nested_scalar_indexes(self, temp_dir):
        """Nested field paths should pass driver validation and reach workers."""
        dataset_uri = generate_nested_contract_dataset(temp_dir)

        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="`meta`.`text`",
            index_type="INVERTED",
            name="nested_text_idx",
            num_workers=2,
        )
        updated_dataset = lr.create_scalar_index(
            uri=updated_dataset.uri,
            column="meta.`a.b`",
            index_type="INVERTED",
            name="literal_dot_text_idx",
            num_workers=2,
        )
        updated_dataset = lr.create_scalar_index(
            uri=updated_dataset.uri,
            column="`meta-data`.`user-id`",
            index_type="BTREE",
            name="hyphen_user_id_idx",
            num_workers=2,
        )

        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`"]

        nested_results = updated_dataset.scanner(
            full_text_query="nestedthree",
            columns=["id", "meta.text"],
        ).to_table()
        literal_dot_results = updated_dataset.scanner(
            full_text_query="literaltwo",
            columns=["id", "meta.`a.b`"],
        ).to_table()

        assert nested_results.column("id").to_pylist() == [3]
        assert literal_dot_results.column("id").to_pylist() == [2]

    def test_build_distributed_nested_same_leaf_scalar_indexes(self, temp_dir):
        """Same leaf names must resolve through their full nested paths."""
        dataset_uri = generate_nested_contract_dataset(temp_dir)

        with pytest.raises(ValueError, match="Column 'leaf' not found"):
            lr.create_scalar_index(
                uri=dataset_uri,
                column="leaf",
                index_type="BTREE",
                name="ambiguous_leaf_idx",
                num_workers=2,
            )

        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="outer.leaf",
            index_type="BTREE",
            name="outer_leaf_idx",
            num_workers=2,
        )
        updated_dataset = lr.create_scalar_index(
            uri=updated_dataset.uri,
            column="other.leaf",
            index_type="BTREE",
            name="other_leaf_idx",
            num_workers=2,
        )

        indices = {idx.name: idx for idx in updated_dataset.describe_indices()}
        assert indices["outer_leaf_idx"].field_names == ["outer.leaf"]
        assert indices["other_leaf_idx"].field_names == ["other.leaf"]

        outer_results = updated_dataset.scanner(
            filter="outer.leaf = 20",
            columns=["id", "outer.leaf"],
        ).to_table()
        other_results = updated_dataset.scanner(
            filter="other.leaf = 20",
            columns=["id", "other.leaf"],
        ).to_table()

        assert outer_results.column("id").to_pylist() == [2]
        assert other_results.column("id").to_pylist() == [3]

    def test_scalar_index_on_mixed_schema_describe_indices(self, temp_dir):
        """Create scalar index on schema with both scalar and vector columns; verify describe_indices."""
        dataset_uri = generate_mixed_schema_dataset(
            temp_dir,
            num_rows=200,
            vector_dim=8,
            rows_per_fragment=50,
        )
        index_name = "label_btree_idx"

        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="label",
            index_type="BTREE",
            name=index_name,
            num_workers=2,
        )

        indices = updated_dataset.describe_indices()
        assert len(indices) >= 1, (
            "describe_indices should return at least the new scalar index"
        )
        names = [idx.name for idx in indices]
        assert index_name in names, (
            f"Expected index name {index_name!r} in describe_indices: {names}"
        )

        label_index = next(idx for idx in indices if idx.name == index_name)
        assert label_index.index_type == "BTree", (
            f"Expected BTree type for scalar index, got {label_index.index_type!r}"
        )

        # Schema should still have both scalar and vector columns
        schema = updated_dataset.schema
        assert schema.field("id") is not None
        assert schema.field("vector") is not None
        assert schema.field("label") is not None
        assert schema.field("score") is not None

    def test_build_distributed_index_replace_false_existing_index(
        self, multi_fragment_lance_dataset
    ):
        """Test that replace=False raises error when trying to create index with existing name."""
        dataset_uri = multi_fragment_lance_dataset
        index_name = "test_replace_false_index"

        # First, create an index
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            name=index_name,
            num_workers=2,
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "Initial index creation failed"

        # Now try to create another index with the same name but replace=False
        # The error might be raised as RuntimeError during distributed processing
        with pytest.raises((ValueError, RuntimeError)) as exc_info:
            lr.create_scalar_index(
                uri=dataset_uri,
                column="text",
                index_type="INVERTED",
                name=index_name,
                replace=False,
                num_workers=2,
            )

        # Verify the error message contains information about existing index
        error_msg = str(exc_info.value)
        assert "already exists" in error_msg and index_name in error_msg

    def test_build_distributed_index_replace_true_overwrite_existing(
        self, multi_fragment_lance_dataset
    ):
        """Test that replace=True successfully overwrites existing index."""
        dataset_uri = multi_fragment_lance_dataset
        index_name = "test_replace_true_index"

        # First, create an index
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            name=index_name,
            num_workers=2,
        )

        # Verify the index was created
        initial_indices = updated_dataset.describe_indices()
        assert len(initial_indices) > 0, "Initial index creation failed"

        # Find our initial index
        initial_index = None
        for idx in initial_indices:
            if idx.name == index_name:
                initial_index = idx
                break
        assert initial_index is not None, "Initial index not found"

        # Now create another index with the same name but replace=True
        updated_dataset = lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            name=index_name,
            replace=True,
            num_workers=2,
        )

        # Verify the index still exists (should have been replaced)
        final_indices = updated_dataset.describe_indices()
        final_index = None
        for idx in final_indices:
            if idx.name == index_name:
                final_index = idx
                break

        assert final_index is not None, "Index should still exist after replacement"
        assert final_index.index_type == "Inverted", "Index type should remain Inverted"

        # Test that the replaced index still works for searching
        search_term = "Python"
        results = updated_dataset.scanner(
            full_text_query=search_term,
            columns=["id", "text"],
        ).to_table()

        assert results.num_rows > 0, (
            f"No results found for search term '{search_term}' after index replacement"
        )

    def test_failed_replace_keeps_existing_index(
        self, multi_fragment_lance_dataset, monkeypatch
    ):
        """A failed replacement must not remove the previously committed index."""
        index_name = "test_replace_failure_keeps_existing_index"

        lr.create_scalar_index(
            uri=multi_fragment_lance_dataset,
            column="text",
            index_type="INVERTED",
            name=index_name,
            num_workers=2,
        )

        monkeypatch.setattr(
            "lance_ray.index._map_async_with_pool",
            lambda **_: [
                {
                    "status": "error",
                    "fragment_ids": [0],
                    "error": "injected worker failure",
                }
            ],
        )

        with pytest.raises(RuntimeError, match="injected worker failure"):
            lr.create_scalar_index(
                uri=multi_fragment_lance_dataset,
                column="text",
                index_type="INVERTED",
                name=index_name,
                replace=True,
                num_workers=2,
            )

        indices = lance.dataset(multi_fragment_lance_dataset).describe_indices()
        assert index_name in [index.name for index in indices]

    def test_build_distributed_index_auto_adjust_workers(self, temp_dir):
        """Test that num_workers is automatically adjusted if it exceeds fragment count."""
        # Create dataset with only 2 fragments
        data = pd.DataFrame(
            {
                "id": [1, 2, 3, 4],
                "text": ["text1", "text2", "text3", "text4"],
            }
        )
        dataset = ray.data.from_pandas(data)
        path = Path(temp_dir) / "small_dataset.lance"
        lr.write_lance(dataset, str(path), min_rows_per_file=2, max_rows_per_file=2)

        # Request more workers than fragments
        updated_dataset = lr.create_scalar_index(
            uri=str(path),
            column="text",
            index_type="INVERTED",
            num_workers=10,  # More than the 2 fragments
        )

        # Should still work and create the index
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after building"

    def test_distributed_fts_index_new_api(self, temp_dir):
        """
        Test distributed FTS index building with the segment workflow.
        """
        # Generate test dataset with multiple fragments
        ds = generate_multi_fragment_dataset(
            temp_dir, num_fragments=4, rows_per_fragment=250
        )

        # Test with the new distributed index building function
        updated_dataset = lr.create_scalar_index(
            uri=ds.uri,
            column="text",
            index_type="INVERTED",
            name="new_api_test_idx",
            num_workers=2,
            num_segments=4,
            remove_stop_words=False,
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after distributed index creation"

        # Find our index
        our_index = None
        for idx in indices:
            if idx.name == "new_api_test_idx":
                our_index = idx
                break

        assert our_index is not None, (
            "Index 'new_api_test_idx' not found in indices list"
        )
        assert our_index.index_type == "Inverted", (
            f"Expected Inverted index, got {our_index.index_type}"
        )
        assert len(our_index.segments) == 4

        # Test that the index works for searching
        sample_data = updated_dataset.take([0], columns=["text"])
        sample_text = sample_data.column(0)[0].as_py()
        search_word = sample_text.split()[0] if sample_text.split() else "test"

        # Perform a full-text search to verify the index works
        results = updated_dataset.scanner(
            full_text_query=search_word,
            columns=["id", "text"],
        ).to_table()

        print(f"Search for '{search_word}' returned {results.num_rows} results")
        assert results.num_rows > 0, f"No results found for search term '{search_word}'"

    def test_distributed_index_with_index_uuid(self, temp_dir):
        """
        Test distributed FTS index creation records the requested index name.
        """
        # Generate test dataset
        ds = generate_multi_fragment_dataset(
            temp_dir, num_fragments=3, rows_per_fragment=100
        )

        # Test with explicit fragment UUID handling
        updated_dataset = lr.create_scalar_index(
            uri=ds.uri,
            column="text",
            index_type="INVERTED",
            name="index_uuid_test_idx",
            num_workers=2,
        )

        # Verify the index was created
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after index creation"

        # Find our index
        our_index = None
        for idx in indices:
            if idx.name == "index_uuid_test_idx":
                our_index = idx
                break

        assert our_index is not None, "Index 'index_uuid_test_idx' not found"
        assert our_index.index_type == "Inverted", (
            f"Expected Inverted index, got {our_index.index_type}"
        )

    def test_distributed_index_error_handling_new_api(self, temp_dir):
        """
        Test error handling in the distributed indexing API.
        """
        # Generate test dataset
        ds = generate_multi_fragment_dataset(
            temp_dir, num_fragments=2, rows_per_fragment=50
        )

        # Test with invalid parameters that should be caught by the new API
        with pytest.raises(ValueError, match="Column name cannot be empty"):
            lr.create_scalar_index(
                uri=ds.uri,
                column="",
                index_type="INVERTED",
                num_workers=2,
            )

        # Test with invalid index type
        with pytest.raises(
            ValueError,
            match="Distributed indexing does not support index type 'INVALID_TYPE'",
        ):
            lr.create_scalar_index(
                uri=ds.uri,
                column="text",
                index_type="INVALID_TYPE",
                num_workers=2,
            )


def check_btree_version_compatibility():
    """Check if lance version supports distributed B-tree indexing (>= 0.37.0)."""
    try:
        lance_version = version.parse(lance.__version__)
        btree_min_version = version.parse("0.37.0")
        return lance_version >= btree_min_version
    except (AttributeError, Exception):
        return False


@pytest.mark.skipif(
    not check_btree_version_compatibility(),
    reason="B-tree indexing requires pylance >= 0.37.0. Current version: {}".format(
        getattr(lance, "__version__", "unknown")
    ),
)
class TestDistributedBTreeIndexing:
    """Distributed BTREE indexing tests using the unified lr.create_scalar_index entrypoint."""

    def test_distributed_btree_index_basic(self, temp_dir):
        """Build a distributed BTREE index and verify search works and type is BTree."""
        ds = generate_multi_fragment_dataset(
            temp_dir, num_fragments=3, rows_per_fragment=500
        )

        updated_dataset = lr.create_scalar_index(
            uri=ds.uri,
            column="id",
            index_type="BTREE",
            name="btree_multiple_fragment_idx",
            replace=False,
            num_workers=3,
        )

        # Verify index
        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after distributed BTREE build"

        our_index = None
        for idx in indices:
            if idx.name == "btree_multiple_fragment_idx":
                our_index = idx
                break
        assert our_index is not None, "BTREE index not found by name"
        assert our_index.index_type == "BTree", (
            f"Expected BTree index, got {our_index.index_type}"
        )

        # Spot-check equality and range queries
        eq_id = 100
        eq_tbl = updated_dataset.scanner(
            filter=f"id = {eq_id}", columns=["id", "text"]
        ).to_table()
        assert eq_tbl.num_rows == 1
        eq_plan = updated_dataset.scanner(
            filter=f"id = {eq_id}",
            columns=["id"],
            use_scalar_index=True,
        ).explain_plan()
        assert "ScalarIndexQuery" in eq_plan
        assert "btree_multiple_fragment_idx" in eq_plan

        rg_tbl = updated_dataset.scanner(
            filter="id >= 200 AND id < 800",
            columns=["id", "text"],
        ).to_table()
        assert rg_tbl.num_rows > 0

    def test_distributed_btree_index_on_branch(self, temp_dir):
        """Build a distributed BTREE index on a branch without indexing main."""
        main_dataset = generate_multi_fragment_dataset(
            temp_dir, num_fragments=3, rows_per_fragment=500
        )
        branch_dataset = main_dataset.create_branch("distributed-btree-index")

        indexed_branch = lr.create_scalar_index(
            uri=branch_dataset,
            column="id",
            index_type="BTREE",
            name="branch_btree_idx",
            replace=False,
            num_workers=3,
        )

        assert "branch_btree_idx" in {
            index.name for index in indexed_branch.describe_indices()
        }
        assert "branch_btree_idx" not in {
            index.name for index in lance.dataset(main_dataset.uri).describe_indices()
        }

        plan = indexed_branch.scanner(
            filter="id = 100", columns=["id"], use_scalar_index=True
        ).explain_plan()
        assert "ScalarIndexQuery" in plan
        assert "branch_btree_idx" in plan

    @pytest.fixture
    def btree_comp_datasets(self, tmp_path):
        """Build two datasets: one with a distributed BTREE index and one without index as baseline."""
        with_index = generate_multi_fragment_dataset(
            tmp_path / "with_index", num_fragments=3, rows_per_fragment=500
        )
        without_index = generate_multi_fragment_dataset(
            tmp_path / "without_index", num_fragments=3, rows_per_fragment=500
        )

        # Build BTREE index on the first dataset using unified API
        with_index = lr.create_scalar_index(
            uri=with_index.uri,
            column="id",
            index_type="BTREE",
            name="btree_comp_idx",
            replace=True,
            num_workers=2,
        )

        return {"with_index": with_index, "without_index": without_index}

    @pytest.mark.parametrize(
        "test_name,filter_expr",
        [
            ("First value", "id = 0"),
            ("Fragment 0 last value", "id = 499"),
            ("Fragment 1 first value", "id = 500"),
            ("Fragment 1 last value", "id = 999"),
            ("Fragment 2 first value", "id = 1000"),
            ("Last value", "id = 1499"),
            ("Fragment 0 middle", "id = 250"),
            ("Fragment 1 middle", "id = 750"),
            ("Fragment 2 middle", "id = 1250"),
            ("Range within fragment 0", "id >= 10 AND id < 20"),
            ("Range within fragment 1", "id >= 510 AND id < 520"),
            ("Range within fragment 2", "id >= 1010 AND id < 1020"),
            ("Cross fragment 0-1", "id >= 495 AND id < 505"),
            ("Cross fragment 1-2", "id >= 995 AND id < 1005"),
            ("Cross all fragments", "id >= 250 AND id < 1250"),
            ("Non-existent small value", "id = -1"),
            ("Non-existent large value", "id = 2000"),
            ("Large range", "id >= 0 AND id < 1500"),
            ("Less than boundary", "id < 500"),
            ("Greater than boundary", "id > 999"),
            ("Less than or equal", "id <= 505"),
            ("Greater than or equal", "id >= 995"),
        ],
    )
    def test_btree_query_results_match_baseline(
        self, btree_comp_datasets, test_name, filter_expr
    ):
        """Compare query results between an indexed dataset and an identical baseline dataset without index."""
        with_index = btree_comp_datasets["with_index"]
        without_index = btree_comp_datasets["without_index"]

        res_idx = with_index.scanner(
            filter=filter_expr, columns=["id", "text"]
        ).to_table()
        res_base = without_index.scanner(
            filter=filter_expr, columns=["id", "text"]
        ).to_table()

        assert res_idx.num_rows == res_base.num_rows, (
            f"Test '{test_name}' failed: indexed returned {res_idx.num_rows} rows, "
            f"baseline returned {res_base.num_rows} rows for filter: {filter_expr}"
        )

        if res_idx.num_rows > 0:
            ids_idx = sorted(res_idx.column("id").to_pylist())
            ids_base = sorted(res_base.column("id").to_pylist())
            assert ids_idx == ids_base, (
                f"Test '{test_name}' failed: indexed and baseline results differ for filter: {filter_expr}"
            )

    def test_distributed_btree_index_many_fragments_many_workers(self, temp_dir):
        """
        Test distributed BTREE index building with many fragments and many workers.

        This test reproduces the scenario reported in the bug where:
        - 5000 fragments with 16 workers resulted in "No partition metadata files found"
        - The test uses a smaller scale (50 fragments, 16 workers) for faster execution
        """
        available_cpus = int(ray.cluster_resources().get("CPU", 4))
        num_workers = min(16, available_cpus)

        ds = generate_multi_fragment_dataset(
            temp_dir, num_fragments=50, rows_per_fragment=100
        )

        fragments = list(ds.get_fragments())
        assert len(fragments) == 50, f"Expected 50 fragments, got {len(fragments)}"

        updated_dataset = lr.create_scalar_index(
            uri=ds.uri,
            column="id",
            index_type="BTREE",
            name="btree_many_workers_idx",
            replace=False,
            num_workers=num_workers,
        )

        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after distributed BTREE build"

        our_index = None
        for idx in indices:
            if idx.name == "btree_many_workers_idx":
                our_index = idx
                break
        assert our_index is not None, "BTREE index not found by name"
        assert our_index.index_type == "BTree", (
            f"Expected BTree index, got {our_index.index_type}"
        )

        eq_tbl = updated_dataset.scanner(filter="id = 2500", columns=["id"]).to_table()
        assert eq_tbl.num_rows == 1, "Exact match query failed"

        rg_tbl = updated_dataset.scanner(
            filter="id >= 1000 AND id < 4000",
            columns=["id"],
        ).to_table()
        assert rg_tbl.num_rows == 3000, (
            f"Range query returned {rg_tbl.num_rows} rows, expected 3000"
        )

    def test_distributed_btree_index_string_column(self, temp_dir):
        """Test distributed BTREE index on string column (like video_uuid in the bug report)."""
        import uuid as uuid_module

        available_cpus = int(ray.cluster_resources().get("CPU", 4))
        num_workers = min(8, available_cpus)

        all_data = []
        num_fragments = 20
        rows_per_fragment = 100
        for frag_idx in range(num_fragments):
            for row_idx in range(rows_per_fragment):
                row_id = frag_idx * rows_per_fragment + row_idx
                all_data.append(
                    {
                        "id": row_id,
                        "video_uuid": str(uuid_module.uuid4()),
                        "fragment_id": frag_idx,
                    }
                )

        df = pd.DataFrame(all_data)
        dataset = ray.data.from_pandas(df)

        path = Path(temp_dir) / "string_btree_test.lance"
        lr.write_lance(
            dataset,
            str(path),
            min_rows_per_file=rows_per_fragment,
            max_rows_per_file=rows_per_fragment,
        )

        ds = lance.dataset(str(path))
        assert len(list(ds.get_fragments())) == num_fragments

        updated_dataset = lr.create_scalar_index(
            uri=ds.uri,
            column="video_uuid",
            index_type="BTREE",
            name="video_uuid_idx",
            replace=False,
            num_workers=num_workers,
        )

        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after distributed BTREE build"

        our_index = None
        for idx in indices:
            if idx.name == "video_uuid_idx":
                our_index = idx
                break
        assert our_index is not None, "String BTREE index not found"
        assert our_index.index_type == "BTree"

        sample_uuid = all_data[500]["video_uuid"]
        result = updated_dataset.scanner(
            filter=f"video_uuid = '{sample_uuid}'",
            columns=["id", "video_uuid"],
        ).to_table()
        assert result.num_rows == 1, (
            f"Expected 1 row for exact UUID match, got {result.num_rows}"
        )


class TestDistributedZoneMapIndexing:
    """Distributed ZONEMAP indexing tests."""

    def test_distributed_zonemap_index_basic(self, temp_dir):
        """Build a distributed ZONEMAP index on an int column and verify index type."""
        ds = generate_multi_fragment_dataset(
            temp_dir, num_fragments=3, rows_per_fragment=500
        )

        updated_dataset = lr.create_scalar_index(
            uri=ds.uri,
            column="id",
            index_type="ZONEMAP",
            name="zonemap_id_idx",
            replace=False,
            num_workers=3,
        )

        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after distributed ZONEMAP build"

        our_index = next((idx for idx in indices if idx.name == "zonemap_id_idx"), None)
        assert our_index is not None, "ZONEMAP index not found by name"
        assert our_index.index_type == "ZoneMap", (
            f"Expected ZoneMap index, got {our_index.index_type}"
        )

    def test_zonemap_query_results_match_baseline(self, temp_dir):
        """ZONEMAP-indexed range queries must return the same rows as a non-indexed scan."""
        with_index = generate_multi_fragment_dataset(
            Path(temp_dir) / "with_zonemap",
            num_fragments=3,
            rows_per_fragment=500,
        )
        without_index = generate_multi_fragment_dataset(
            Path(temp_dir) / "without_zonemap",
            num_fragments=3,
            rows_per_fragment=500,
        )

        updated_dataset = lr.create_scalar_index(
            uri=with_index.uri,
            column="id",
            index_type="ZONEMAP",
            name="zonemap_range_idx",
            replace=False,
            num_workers=3,
        )

        for filter_expr in [
            "id = 250",
            "id >= 200 AND id < 800",
            "id < 500",
            "id > 999",
        ]:
            indexed = updated_dataset.scanner(
                filter=filter_expr, columns=["id"]
            ).to_table()
            baseline = without_index.scanner(
                filter=filter_expr, columns=["id"]
            ).to_table()
            assert indexed.num_rows == baseline.num_rows, (
                f"Row count mismatch for '{filter_expr}': "
                f"indexed={indexed.num_rows}, baseline={baseline.num_rows}"
            )
            if indexed.num_rows > 0:
                assert sorted(indexed.column("id").to_pylist()) == sorted(
                    baseline.column("id").to_pylist()
                ), f"Result mismatch for '{filter_expr}'"

    def test_distributed_zonemap_index_string_column(self, temp_dir):
        """Build a distributed ZONEMAP index on a string column."""
        ds = generate_multi_fragment_dataset(
            temp_dir, num_fragments=3, rows_per_fragment=200
        )

        updated_dataset = lr.create_scalar_index(
            uri=ds.uri,
            column="text",
            index_type="ZONEMAP",
            name="zonemap_text_idx",
            replace=False,
            num_workers=3,
        )

        indices = updated_dataset.describe_indices()
        our_index = next(
            (idx for idx in indices if idx.name == "zonemap_text_idx"), None
        )
        assert our_index is not None, "ZONEMAP string index not found by name"
        assert our_index.index_type == "ZoneMap"


class TestDistributedBitmapIndexing:
    """Distributed BITMAP indexing tests."""

    def test_distributed_bitmap_index_matches_baseline(self, temp_dir):
        """Build a distributed BITMAP index and verify query results."""
        with_index = generate_multi_fragment_dataset(
            Path(temp_dir) / "with_bitmap",
            num_fragments=3,
            rows_per_fragment=250,
        )
        without_index = generate_multi_fragment_dataset(
            Path(temp_dir) / "without_bitmap",
            num_fragments=3,
            rows_per_fragment=250,
        )

        updated_dataset = lr.create_scalar_index(
            uri=with_index.uri,
            column="fragment_id",
            index_type="BITMAP",
            name="fragment_bitmap_idx",
            replace=False,
            num_workers=3,
        )

        indices = updated_dataset.describe_indices()
        our_index = next(
            (idx for idx in indices if idx.name == "fragment_bitmap_idx"),
            None,
        )

        assert our_index is not None, "BITMAP index not found by name"
        assert our_index.index_type == "Bitmap"

        indexed = updated_dataset.scanner(
            filter="fragment_id = 1",
            columns=["id", "fragment_id"],
        ).to_table()
        baseline = without_index.scanner(
            filter="fragment_id = 1",
            columns=["id", "fragment_id"],
        ).to_table()

        assert indexed.num_rows == baseline.num_rows
        assert sorted(indexed.column("id").to_pylist()) == sorted(
            baseline.column("id").to_pylist()
        )

        plan = updated_dataset.scanner(
            filter="fragment_id = 1",
            columns=["id"],
            use_scalar_index=True,
        ).explain_plan()
        assert "ScalarIndexQuery" in plan
        assert "fragment_bitmap_idx" in plan


class TestOptimizeIndices:
    """Test cases for optimize_indices (incremental index optimization)."""

    def test_optimize_indices_uri_required(self):
        """optimize_indices raises ValueError when neither uri nor namespace provided."""
        with pytest.raises(ValueError, match="Must provide either"):
            lr.optimize_indices()

    def test_optimize_indices_uri_and_namespace_exclusive(self):
        """optimize_indices raises ValueError when both uri and namespace provided."""
        with pytest.raises(ValueError, match="Cannot provide both"):
            lr.optimize_indices(
                uri="/some/path.lance",
                namespace_impl="dir",
                namespace_properties={"root": "/tmp"},
                table_id=["t1"],
            )

    def test_optimize_indices_success_with_uri(self, multi_fragment_lance_dataset):
        """optimize_indices returns LanceDataset and describe_indices is consistent when API is available."""
        dataset_uri = multi_fragment_lance_dataset
        lr.create_scalar_index(
            uri=dataset_uri,
            column="text",
            index_type="INVERTED",
            name="text_idx",
            num_workers=2,
        )

        ds = lance.LanceDataset(dataset_uri)
        has_optimize = (
            getattr(ds, "optimize_indices", None) is not None
            or getattr(ds, "optimize", None) is not None
        )
        if not has_optimize:
            pytest.skip(
                "Lance dataset does not expose optimize_indices or optimize; "
                "skipping optimize_indices integration test."
            )

        result = lr.optimize_indices(uri=dataset_uri)
        assert result is not None
        assert isinstance(result, lance.LanceDataset)
        assert result.count_rows() == lance.LanceDataset(dataset_uri).count_rows()

        indices = result.describe_indices()
        assert len(indices) >= 1, (
            "describe_indices should include at least the existing index"
        )
        names = [idx.name for idx in indices]
        assert "text_idx" in names, f"Expected 'text_idx' in describe_indices: {names}"

    def test_optimize_indices_runtime_error_when_api_missing(self, temp_dir):
        """optimize_indices raises RuntimeError when dataset has no optimize API."""
        path = Path(temp_dir) / "no_optimize.lance"
        df = pd.DataFrame({"id": [1, 2], "t": ["a", "b"]})
        lr.write_lance(ray.data.from_pandas(df), str(path))
        ds = lance.LanceDataset(str(path))

        if (
            getattr(ds, "optimize_indices", None) is not None
            or getattr(ds, "optimize", None) is not None
        ):
            pytest.skip(
                "This lance version exposes optimize_indices/optimize; "
                "cannot test RuntimeError path."
            )

        with pytest.raises(RuntimeError, match="optimize_indices or optimize"):
            lr.optimize_indices(uri=str(path))


class TestNamespaceIndexing:
    """Test cases for distributed indexing with DirectoryNamespace."""

    def test_distributed_fts_index_with_directory_namespace(self, temp_dir):
        """Test distributed FTS index building using DirectoryNamespace."""
        table_id = ["fts_index_test_table"]

        data = pd.DataFrame(
            {
                "id": range(100),
                "text": [
                    f"This is document {i} with searchable content" for i in range(100)
                ],
            }
        )
        dataset = ray.data.from_pandas(data)
        lr.write_lance(
            dataset,
            namespace_impl="dir",
            namespace_properties={"root": temp_dir},
            table_id=table_id,
            min_rows_per_file=25,
            max_rows_per_file=25,
        )

        # Use namespace params only - create_scalar_index will resolve URI internally
        updated_dataset = lr.create_scalar_index(
            column="text",
            index_type="INVERTED",
            name="fts_namespace_idx",
            num_workers=2,
            namespace_impl="dir",
            namespace_properties={"root": temp_dir},
            table_id=table_id,
        )

        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after distributed FTS build"

        our_index = None
        for idx in indices:
            if idx.name == "fts_namespace_idx":
                our_index = idx
                break
        assert our_index is not None, "FTS index not found"
        assert our_index.index_type == "Inverted"

        results = updated_dataset.scanner(
            full_text_query="document",
            columns=["id", "text"],
        ).to_table()
        assert results.num_rows > 0, "FTS search should return results"

    def test_distributed_btree_index_with_directory_namespace(self, temp_dir):
        """Test distributed BTREE index building using DirectoryNamespace."""
        table_id = ["btree_index_test_table"]

        data = pd.DataFrame(
            {
                "id": range(200),
                "value": [f"value_{i}" for i in range(200)],
            }
        )
        dataset = ray.data.from_pandas(data)
        lr.write_lance(
            dataset,
            namespace_impl="dir",
            namespace_properties={"root": temp_dir},
            table_id=table_id,
            min_rows_per_file=50,
            max_rows_per_file=50,
        )

        # Use namespace params only - create_scalar_index will resolve URI internally
        updated_dataset = lr.create_scalar_index(
            column="id",
            index_type="BTREE",
            name="btree_namespace_idx",
            num_workers=2,
            namespace_impl="dir",
            namespace_properties={"root": temp_dir},
            table_id=table_id,
        )

        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after distributed BTREE build"

        our_index = None
        for idx in indices:
            if idx.name == "btree_namespace_idx":
                our_index = idx
                break
        assert our_index is not None, "BTREE index not found"
        assert our_index.index_type == "BTree"

        result = updated_dataset.scanner(filter="id = 100", columns=["id"]).to_table()
        assert result.num_rows == 1, "BTREE index query should return 1 row"

    def test_distributed_vector_index_with_directory_namespace(self, temp_dir):
        """Test distributed vector index building using DirectoryNamespace.

        Verifies that create_index() correctly resolves the dataset URI and
        passes a storage_options_provider to workers when namespace params
        are supplied, mirroring the behaviour of create_scalar_index().
        """
        table_id = ["vector_index_namespace_test"]
        dim = 32
        num_rows = 1024  # needs >= num_partitions * sample_rate (4 * 256)

        # Build a vector dataset and register it under the dir namespace.
        values = pa.array(
            [random.gauss(0, 1) for _ in range(num_rows * dim)], type=pa.float32()
        )
        vector_array = pa.FixedSizeListArray.from_arrays(values, dim)
        tbl = pa.Table.from_arrays(
            [vector_array, pa.array(range(num_rows), type=pa.int64())],
            names=["vector", "id"],
        )
        dataset = ray.data.from_arrow(tbl)
        lr.write_lance(
            dataset,
            namespace_impl="dir",
            namespace_properties={"root": temp_dir},
            table_id=table_id,
            min_rows_per_file=64,
            max_rows_per_file=64,
        )

        # Build vector index via namespace params — no URI needed.
        # sample_rate=4: PQ requires 256 * sample_rate <= num_rows (256*4=1024 ✓)
        try:
            updated_dataset = lr.create_index(
                column="vector",
                index_type="IVF_PQ",
                name="vec_namespace_idx",
                num_workers=2,
                num_partitions=4,
                num_sub_vectors=4,
                sample_rate=4,
                namespace_impl="dir",
                namespace_properties={"root": temp_dir},
                table_id=table_id,
            )
        except RuntimeError as exc:
            if "not yet implemented" in str(exc):
                pytest.skip(f"Skipping: lance version limitation: {exc}")
            raise

        indices = updated_dataset.describe_indices()
        assert len(indices) > 0, "No indices found after distributed vector build"

        our_index = next(
            (idx for idx in indices if idx.name == "vec_namespace_idx"), None
        )
        assert our_index is not None, "vector index not found in dataset"

        # Sanity-check that the index is usable for ANN search.
        query = [random.gauss(0, 1) for _ in range(dim)]
        results = updated_dataset.to_table(
            nearest={"column": "vector", "q": query, "k": 5}
        )
        assert results.num_rows == 5, "ANN search should return 5 results"

    def test_create_index_namespace_uri_mutual_exclusion(self, temp_dir):
        """create_index raises ValueError when both uri and namespace params are given."""
        with pytest.raises(ValueError, match="Cannot provide both"):
            lr.create_index(
                uri="/some/path.lance",
                column="vector",
                index_type="IVF_PQ",
                namespace_impl="dir",
                namespace_properties={"root": temp_dir},
                table_id=["some_table"],
            )

    def test_create_index_namespace_requires_uri_or_namespace(self):
        """create_index raises ValueError when neither uri nor namespace params are given."""
        with pytest.raises(ValueError, match="Must provide either"):
            lr.create_index(
                column="vector",
                index_type="IVF_PQ",
            )


def generate_multi_fragment_vector_dataset(
    tmp_path, num_fragments: int = 4, rows_per_fragment: int = 64, dim: int = 128
) -> str:
    """Generate a Lance dataset with a vector column and multiple fragments.

    The dataset is written via lance-ray so that each fragment has the same
    number of rows, which makes it suitable for distributed vector index
    building tests.
    """
    num_rows = num_fragments * rows_per_fragment
    # Generate random vectors
    data = [random.gauss(0, 1) for _ in range(num_rows * dim)]

    # Manually construct a FixedSizeList "vector" column compatible with
    # Lance's vector index requirements.
    values = pa.array(data, type=pa.float32())
    vector_array = pa.FixedSizeListArray.from_arrays(values, dim)
    tbl = pa.Table.from_arrays([vector_array], names=["vector"])
    tbl = tbl.append_column("id", pa.array(range(num_rows), type=pa.int64()))

    path = Path(tmp_path) / "multi_fragment_vector.lance"
    lance.write_dataset(
        tbl,
        str(path),
        max_rows_per_file=rows_per_fragment,
    )

    return str(path)


def generate_nested_vector_dataset(
    tmp_path, num_fragments: int = 2, rows_per_fragment: int = 256, dim: int = 8
) -> str:
    """Generate a Lance dataset with a nested vector column."""
    num_rows = num_fragments * rows_per_fragment
    values = pa.array(
        [
            float(row_idx) + (dim_idx / 100.0)
            for row_idx in range(num_rows)
            for dim_idx in range(dim)
        ],
        type=pa.float32(),
    )
    vector_array = pa.FixedSizeListArray.from_arrays(values, dim)
    meta_type = pa.struct([pa.field("vector", vector_array.type)])
    meta_array = pa.StructArray.from_arrays([vector_array], fields=list(meta_type))
    tbl = pa.Table.from_arrays(
        [pa.array(range(num_rows), type=pa.int64()), meta_array],
        names=["id", "meta"],
    )

    path = Path(tmp_path) / "nested_vector.lance"
    lance.write_dataset(
        tbl,
        str(path),
        max_rows_per_file=rows_per_fragment,
    )

    return str(path)


@pytest.mark.parametrize("index_type", ["IVF_FLAT", "IVF_SQ", "IVF_PQ"])
def test_build_distributed_vector_index(tmp_path, index_type):
    """Build a distributed vector index and verify nearest search works."""
    dataset_uri = generate_multi_fragment_vector_dataset(
        tmp_path, num_fragments=4, rows_per_fragment=1024, dim=128
    )
    build_kwargs = {"num_sub_vectors": 16, "sample_rate": 16}
    if index_type == "IVF_PQ":
        build_kwargs.update(num_bits=4, sample_rate=4)
    index_name = f"idx_{index_type}"

    # Build distributed vector index using the high-level Ray entrypoint.
    try:
        updated_dataset = lr.create_index(
            uri=dataset_uri,
            column="vector",
            index_type=index_type,
            name=index_name,
            num_workers=2,
            num_segments=4,
            num_partitions=4,
            **build_kwargs,
        )
    except RuntimeError as exc:
        # Older pylance builds may not yet support creating empty distributed
        # vector indices with train=False. In that case we skip the functional
        # verification while still ensuring the Ray entrypoint is wired
        # correctly.
        msg = str(exc)
        if (
            "Creating empty vector indices with train=False is not yet implemented"
            in msg
        ):
            pytest.skip(
                "Current pylance build does not yet support distributed vector "
                "indices with train=False; skipping functional test."
            )
        raise

    indices = updated_dataset.describe_indices()
    assert len(indices) > 0, "No indices found after distributed vector index build"

    # Find the index with the name we specified
    vec_index = next((idx for idx in indices if idx.name == index_name), None)
    assert vec_index is not None, f"Index with name {index_name} not found"
    assert vec_index.index_type == index_type, (
        f"Expected {index_type} vector index, got {vec_index.index_type}"
    )
    assert len(vec_index.segments) == 4

    # Run a simple nearest-neighbor query to ensure the index is usable.
    q = [random.gauss(0, 1) for _ in range(128)]
    result = updated_dataset.to_table(
        nearest={"column": "vector", "q": q, "k": 5},
        columns=["id", "vector"],
    )

    assert result.num_rows == 5

    plan = updated_dataset.scanner(
        nearest={"column": "vector", "q": q, "k": 5},
        columns=["id"],
    ).explain_plan()
    assert "ANNSubIndex" in plan
    assert index_name in plan

    if index_type == "IVF_PQ":
        stats = updated_dataset.stats.index_stats(index_name)
        assert stats["indices"]
        assert all(index["sub_index"]["nbits"] == 4 for index in stats["indices"])


@pytest.mark.parametrize("index_type", ["IVF_FLAT", "IVF_PQ"])
def test_distributed_nested_vector_index_and_search(tmp_path, index_type):
    """Distributed vector index and search should accept nested canonical paths."""
    dataset_uri = generate_nested_vector_dataset(tmp_path)
    query = [dim_idx / 100.0 for dim_idx in range(8)]
    index_kwargs = {"num_sub_vectors": 2} if index_type == "IVF_PQ" else {}

    try:
        updated_dataset = lr.create_index(
            uri=dataset_uri,
            column="meta.vector",
            index_type=index_type,
            name=f"nested_vector_{index_type.lower()}_idx",
            num_workers=2,
            num_partitions=1,
            sample_rate=2,
            **index_kwargs,
        )
    except RuntimeError as exc:
        msg = str(exc)
        if (
            "Creating empty vector indices with train=False is not yet implemented"
            in msg
        ):
            pytest.skip(f"Skipping: lance version limitation: {exc}")
        raise

    indices = {idx.name: idx for idx in updated_dataset.describe_indices()}
    index_name = f"nested_vector_{index_type.lower()}_idx"
    assert indices[index_name].field_names == ["meta.vector"]

    if not _scanner_accepts_index_segments(updated_dataset):
        with pytest.raises(RuntimeError, match="does not support index_segments"):
            lr.vector_search(
                uri=updated_dataset.uri,
                nearest={"column": "`meta`.`vector`", "q": query, "k": 5},
                index_name=index_name,
                columns=["id"],
                num_workers=2,
                fast_search=True,
            )
        return

    result = lr.vector_search(
        uri=updated_dataset.uri,
        nearest={"column": "`meta`.`vector`", "q": query, "k": 5},
        index_name=index_name,
        columns=["id"],
        num_workers=2,
        fast_search=True,
    )

    assert result.num_rows == 5
    assert result.column("id").to_pylist()[0] == 0


@pytest.fixture
def large_string_dataset(temp_dir):
    """Create a multi-fragment dataset whose indexed column is large_string."""
    path = Path(temp_dir) / "large_string.lance"
    table = pa.table(
        {
            "id": pa.array(range(8), type=pa.int64()),
            "large_text": pa.array(
                [f"value-{i:02d}" for i in range(8)], type=pa.large_string()
            ),
        }
    )
    lance.write_dataset(table, str(path), max_rows_per_file=2)
    return str(path)


def test_build_distributed_btree_index_on_large_string(large_string_dataset):
    """BTREE indexes must accept large_string columns.

    This exercises the real pylance build (lance-format/lance#7525) rather than
    a mocked worker, so it fails if the pinned pylance cannot index large_string.
    """
    updated_dataset = lr.create_scalar_index(
        uri=large_string_dataset,
        column="large_text",
        index_type="BTREE",
        num_workers=2,
    )

    indices = {idx.name: idx for idx in updated_dataset.describe_indices()}
    assert "large_text_idx" in indices
    assert indices["large_text_idx"].field_names == ["large_text"]

    results = updated_dataset.scanner(
        filter="large_text = 'value-03'", columns=["id"]
    ).to_table()
    assert results.column("id").to_pylist() == [3]