diff --git a/README.md b/README.md index a5ff276..7ad621f 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,18 @@ transaction as the new ones (the index type may change on a full rebuild; the column may not — drop the index or use a different name for that). `replace=False` rejects an existing name. +Scalar and vector index builds partition the selected fragments in dataset +order into `ceil(fragment_count / fragment_group_size)` non-empty segments +(`fragment_group_size` defaults to 10). Within this fixed segment count, +contiguous groups minimize the largest group's row count; the parameter is +not a hard limit on fragments per group. Identical ordered fragments, row +counts, and group sizes produce identical boundaries across index builds. +This preserves opportunities to compact adjacent small fragments within a +segment, although compaction still cannot cross different index-coverage +boundaries. Partial scalar builds may leave gaps between selected fragments. +Existing indexes with interleaved coverage need a full rebuild with +`replace=True` to adopt the new layout. + #### Partial Builds and Incremental Backfill Pass `fragment_ids` to index only a subset of fragments today and backfill diff --git a/daft_lance/_lance.py b/daft_lance/_lance.py index 728defd..07b7d9e 100644 --- a/daft_lance/_lance.py +++ b/daft_lance/_lance.py @@ -720,8 +720,11 @@ def create_scalar_index( index_cache_size: Size of the index cache. default_scan_options: Default scan options for the dataset. metadata_cache_size_bytes: Size of the metadata cache in bytes. - fragment_group_size: Number of fragments to group together in each fragment batch. If None, - defaults to 10. Must be a positive integer. + fragment_group_size: Target number of fragments per index segment. Defaults to 10; + must be positive. Determines the segment count as ceil(fragment_count / fragment_group_size). + Selected fragments are grouped in dataset order, minimizing the largest group's row count; + this is not a hard limit on fragments per group. Contiguous coverage preserves compaction + opportunities within each segment. num_partitions: Optional number of partitions to use when repartitioning fragment batches before execution. Only values greater than 1 enable additional parallelism on distributed runners; values <= 1 or None will use the default partitioning. max_concurrency: Maximum number of concurrent tasks to use for processing fragment batches. @@ -915,8 +918,11 @@ def create_vector_index( index_cache_size: Size of the index cache. default_scan_options: Default scan options. metadata_cache_size_bytes: Size of the metadata cache in bytes. - fragment_group_size: Number of fragments grouped per worker task. If None, - defaults to 10. Each group becomes one independently built index segment. + fragment_group_size: Target number of fragments per index segment. Defaults to 10; + must be positive. Determines the segment count as ceil(fragment_count / fragment_group_size). + Contiguous groups in dataset order minimize the largest group's row count; + this is not a hard limit on fragments per group. Each group becomes one independently + built segment, preserving compaction opportunities within that segment. max_concurrency: Maximum number of concurrent tasks for fragment batches. If None, Daft uses its default concurrency setting. **kwargs: Additional keyword arguments forwarded to Lance's index segment diff --git a/daft_lance/utils.py b/daft_lance/utils.py index c3cf79d..151d0c0 100644 --- a/daft_lance/utils.py +++ b/daft_lance/utils.py @@ -89,7 +89,13 @@ def worker_open_context(self) -> DatasetOpenContext: def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int) -> list[dict[str, list[int]]]: - """Distribute fragments across workers using a balanced algorithm considering fragment sizes.""" + """Partition fragments in input order into contiguous, row-balanced groups. + + The group count is ceil(fragment_count / fragment_group_size). + Minimize the maximum row count within that count, preserving adjacency so + index segment coverage leaves small adjacent fragments mergeable. Identical + inputs produce identical boundaries, including across different columns. + """ if fragment_group_size <= 0: raise ValueError("fragment_group_size must be a positive integer") @@ -104,33 +110,44 @@ def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int row_count = fragment.count_rows() fragment_info.append({"id": fragment.fragment_id, "size": row_count}) except Exception as e: - # If we can't get size info, use fragment_id as a fallback logger.warning( - "Could not get size for fragment %s: %s. Using fragment_id as size estimate.", + "Could not get size for fragment %s: %s. Using 1 as size estimate.", fragment.fragment_id, e, ) - fragment_info.append({"id": fragment.fragment_id, "size": fragment.fragment_id}) - - # Sort fragments by size in descending order (largest first) - # This helps with better load balancing using the greedy algorithm - fragment_info.sort(key=lambda x: x["size"], reverse=True) - - # Initialize fragment groups for each worker - fragment_group_list: list[list[int]] = [[] for _ in range(num_groups)] - group_size_list = [0] * num_groups + fragment_info.append({"id": fragment.fragment_id, "size": 1}) - # Greedy assignment: assign each fragment to the worker with minimum workload - for frag_info in fragment_info: - # Find the worker with the minimum current workload - min_workload_idx = min(range(num_groups), key=lambda i: group_size_list[i]) - - # Assign fragment to this worker - fragment_group_list[min_workload_idx].append(frag_info["id"]) - group_size_list[min_workload_idx] += frag_info["size"] - - # Log distribution statistics for debugging total_size = sum(frag_info["size"] for frag_info in fragment_info) + # Binary search the smallest feasible maximum workload for contiguous groups. + lower = max(max(info["size"] for info in fragment_info), (total_size + num_groups - 1) // num_groups) + upper = total_size + while lower < upper: + limit = (lower + upper) // 2 + groups_needed = 1 + workload = 0 + for info in fragment_info: + if workload + info["size"] > limit: + groups_needed += 1 + workload = 0 + workload += info["size"] + if groups_needed <= num_groups: + upper = limit + else: + lower = limit + 1 + + fragment_group_list: list[list[int]] = [[]] + group_size_list = [0] + for i, info in enumerate(fragment_info): + # Reserve a fragment for each remaining group, even for zero-row inputs. + if fragment_group_list[-1] and ( + group_size_list[-1] + info["size"] > lower + or len(fragment_info) - i == num_groups - len(fragment_group_list) + ): + fragment_group_list.append([]) + group_size_list.append(0) + fragment_group_list[-1].append(info["id"]) + group_size_list[-1] += info["size"] + logger.info( "Fragment distribution statistics: Total fragments=%d, Total size=%d, Fragment group size=%d, the num of Fragment group=%d", len(fragment_info), @@ -149,10 +166,7 @@ def distribute_fragments_balanced(fragments: list[Any], fragment_group_size: int percentage, ) - # Filter out empty batches (shouldn't happen with proper input validation) - non_empty_batches = [{"fragment_ids": batch} for batch in fragment_group_list if batch] - - return non_empty_batches + return [{"fragment_ids": batch} for batch in fragment_group_list] def construct_lance_dataset_handle( diff --git a/tests/io/lance/test_fragment_distribution.py b/tests/io/lance/test_fragment_distribution.py new file mode 100644 index 0000000..1c980f6 --- /dev/null +++ b/tests/io/lance/test_fragment_distribution.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from dataclasses import dataclass +from itertools import combinations + +import pytest + +from daft_lance.utils import distribute_fragments_balanced + + +@dataclass +class _Fragment: + fragment_id: int + rows: int + + def count_rows(self) -> int: + return self.rows + + +def test_equal_fragments_form_contiguous_groups() -> None: + fragments = [_Fragment(i, 100) for i in range(8)] + + assert distribute_fragments_balanced(fragments, 4) == [ + {"fragment_ids": [0, 1, 2, 3]}, + {"fragment_ids": [4, 5, 6, 7]}, + ] + + +@pytest.mark.parametrize( + "rows", + [ + [], + [0], + [0] * 8, + [0, 10, 0, 0, 20, 0], + [100, 80, 60, 40, 20], + [10, 1, 1, 1, 1, 10], + [2, 3, 8, 1, 6], + [2**60, 2**60 + 1, 0, 1], + ], +) +@pytest.mark.parametrize("fragment_group_size", [1, 2, 3, 4, 8, 10]) +def test_contiguous_groups_minimize_maximum_rows(rows: list[int], fragment_group_size: int) -> None: + # Dataset order determines adjacency, not numerical fragment ID order. + fragments = [_Fragment(100 - i * 3, count) for i, count in enumerate(rows)] + groups = distribute_fragments_balanced(fragments, fragment_group_size) + batches = [group["fragment_ids"] for group in groups] + + assert len(batches) == (len(rows) + fragment_group_size - 1) // fragment_group_size + assert all(batches) + assert [fid for batch in batches for fid in batch] == [fragment.fragment_id for fragment in fragments] + assert groups == distribute_fragments_balanced(fragments, fragment_group_size) + if not rows: + return + + sizes = {fragment.fragment_id: fragment.rows for fragment in fragments} + actual = max(sum(sizes[fid] for fid in batch) for batch in batches) + # Independent oracle: enumerate every legal placement of contiguous cuts. + optimal = min( + max(sum(rows[start:end]) for start, end in zip((0, *cuts), (*cuts, len(rows)))) + for cuts in combinations(range(1, len(rows)), len(batches) - 1) + ) + assert actual == optimal + + +def test_unknown_fragment_sizes_preserve_order(caplog: pytest.LogCaptureFixture) -> None: + class UnknownSizeFragment: + def __init__(self, fragment_id: int) -> None: + self.fragment_id = fragment_id + + def count_rows(self) -> int: + raise OSError("row count unavailable") + + fragments = [UnknownSizeFragment(i) for i in [0, 3, 10, 100]] + assert distribute_fragments_balanced(fragments, 2) == [ + {"fragment_ids": [0, 3]}, + {"fragment_ids": [10, 100]}, + ] + assert "Using 1 as size estimate" in caplog.text + + +@pytest.mark.parametrize("fragment_group_size", [0, -1]) +def test_group_size_must_be_positive(fragment_group_size: int) -> None: + with pytest.raises(ValueError, match="fragment_group_size must be a positive integer"): + distribute_fragments_balanced([_Fragment(0, 100)], fragment_group_size) diff --git a/tests/io/lance/test_lance_compaction.py b/tests/io/lance/test_lance_compaction.py index 9f81fcf..3499ae3 100644 --- a/tests/io/lance/test_lance_compaction.py +++ b/tests/io/lance/test_lance_compaction.py @@ -1,13 +1,82 @@ from __future__ import annotations from pathlib import Path +from typing import cast import lance import pandas as pd import pyarrow as pa import pytest -from daft_lance import compact_files +from daft_lance import compact_files, create_scalar_index, create_vector_index + + +@pytest.mark.parametrize("index_type", ["BTREE", "IVF_FLAT"]) +def test_distributed_index_preserves_compaction(tmp_path: Path, index_type: str) -> None: + """Index segments leave adjacent small fragments mergeable.""" + uri = str(tmp_path / "indexed_compaction.lance") + table = pa.table( + { + "id": range(800), + "value": range(800), + "vector": pa.array([[float(i), float(i % 7)] for i in range(800)], type=pa.list_(pa.float32(), 2)), + } + ) + dataset = lance.write_dataset(table, uri, max_rows_per_file=100, max_rows_per_group=100) + assert len(dataset.get_fragments()) == 8 + + if index_type == "BTREE": + # Equivalent builds on different columns must use the same boundaries. + for column in ["id", "value"]: + create_scalar_index( + uri, + column=column, + index_type="BTREE", + name=f"{column}_idx", + fragment_group_size=4, + max_concurrency=2, + ) + else: + create_vector_index( + uri, + column="vector", + index_type="IVF_FLAT", + name="vector_idx", + fragment_group_size=4, + max_concurrency=2, + num_partitions=2, + ) + + dataset = lance.dataset(uri) + for index in dataset.describe_indices(): + assert sorted(sorted(segment.fragment_ids) for segment in index.segments) == [[0, 1, 2, 3], [4, 5, 6, 7]] + + def check_query(dataset: lance.LanceDataset) -> None: + if index_type == "BTREE": + scanner = dataset.scanner(filter="id >= 120 AND value < 130", columns=["id"]) + assert "ScalarIndexQuery" in scanner.explain_plan() + expected_ids = list(range(120, 130)) + else: + scanner = dataset.scanner(nearest={"column": "vector", "q": [123.0, 4.0], "k": 1}, columns=["id"]) + assert "ANNSubIndex" in scanner.explain_plan() + expected_ids = [123] + assert sorted(cast(list[int], scanner.to_table()["id"].to_pylist())) == expected_ids + + check_query(dataset) + metrics = compact_files( + uri, + compaction_options={"target_rows_per_fragment": 1000, "num_threads": 1}, + partition_num=2, + concurrency=2, + ) + assert metrics is not None + assert metrics.fragments_removed == 8 + assert metrics.fragments_added == 2 + + dataset = lance.dataset(uri) + assert [fragment.count_rows() for fragment in dataset.get_fragments()] == [400, 400] + assert dataset.to_table().sort_by("id") == table + check_query(dataset) def create_dataset_with_fragments(path: Path, fragment_data: list[pd.DataFrame]):