fix(index): preserve contiguous fragment coverage for compaction

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

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

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(

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)

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]):