fix: make distributed index builds version-consistent and report real coverage

lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala

@@ -88,18 +88,18 @@ case class AddIndexExec(
     val btreeBuildMode = IndexUtils.btreeBuildMode(indexType, args)
     val scalarSegmentIndexType = IndexUtils.scalarSegmentIndexType(method)
 
-    val (fragmentWorkloads, canonicalColumns) = {
+    // Plan and build at one pinned version; the commit opens the live dataset.
+    val (fragmentWorkloads, canonicalColumns, buildReadOptions) = {
       val ds = Utils.openDatasetBuilder(readOptions).build()
       try {
         val canonical = columns.map { column =>
           val field = IndexUtils.resolveIndexField(ds.getLanceSchema, indexType, column)
           FieldPathUtils.pathByFieldId(ds.getLanceSchema, field.getId)
         }
         (
-          ds.getFragments.asScala
-            .map(fragment => FragmentWorkload(fragment.getId, fragment.metadata().getNumRows))
-            .toList,
-          canonical)
+          IndexUtils.fragmentWorkloads(ds),
+          canonical,
+          IndexUtils.pinVersion(readOptions, ds))
       } finally {
         ds.close()
       }
@@ -163,27 +163,26 @@ case class AddIndexExec(
     val (nsImpl, nsProps, tableId, initialStorageOpts) =
       extractNamespaceInfo(lanceDataset, readOptions)
 
-    // Range-mode BTree uses preprocessed data from Spark and keeps its dedicated path.
     if (btreeBuildMode.contains("range")) {
       val segments = new RangeBasedBTreeIndexJob(
         this.copy(columns = canonicalColumns),
-        readOptions,
+        buildReadOptions,
         fragmentIds.size,
         nsImpl,
         nsProps,
         tableId,
         initialStorageOpts).run()
-      commitIndexSegments(readOptions, canonicalColumns.head, segments)
+      val indexed = commitIndexSegments(readOptions, canonicalColumns.head, segments)
       return Seq(new GenericInternalRow(Array[Any](
-        fragmentIds.size.toLong,
+        indexed.toLong,
         UTF8String.fromString(indexName))))
     }
 
     // Scalar segment indexes use the logical segment commit path.
     if (scalarSegmentIndexType.isDefined) {
       val segmentJob = new ScalarSegmentIndexJob(
         this.copy(columns = canonicalColumns),
-        readOptions,
+        buildReadOptions,
         fragmentWorkloads,
         validatedNumSegments,
         nsImpl,
@@ -192,9 +191,9 @@ case class AddIndexExec(
         initialStorageOpts)
       val segments = segmentJob.run()
       // Atomic add+remove via Lance core; see commitIndexSegments
-      commitIndexSegments(readOptions, canonicalColumns.head, segments)
+      val indexed = commitIndexSegments(readOptions, canonicalColumns.head, segments)
       return Seq(new GenericInternalRow(Array[Any](
-        fragmentIds.size.toLong,
+        indexed.toLong,
         UTF8String.fromString(indexName))))
     }
 
@@ -231,16 +230,31 @@ case class AddIndexExec(
   // Lance core's commitExistingIndexSegments handles atomic replacement:
   // it finds existing segments whose fragments overlap with incoming ones
   // and removes them in the same CreateIndex transaction.
+  //
+  // Returns the number of fragments the commit actually covers, which is what the command reports.
   private def commitIndexSegments(
       readOptions: LanceSparkReadOptions,
       column: String,
-      segments: Seq[Index]): Unit = {
+      segments: Seq[Index]): Int = {
     val dataset = Utils.openDatasetBuilder(readOptions).build()
     try {
-      dataset.commitExistingIndexSegments(
+      IndexUtils.requireCommittableCoverage(
+        IndexUtils.liveFragmentIds(dataset),
+        segments,
+        indexName)
+      val committed = dataset.commitExistingIndexSegments(
         indexName,
         column,
         segments.toList.asJava)
+      // The commit advances this handle to the manifest it wrote, so both the returned metadata and
+      // the fragment list below describe the committed state rather than the one validated above.
+      IndexUtils
+        .establishedCoverage(
+          segments,
+          committed.asScala.toSeq,
+          IndexUtils.liveFragmentIds(dataset),
+          indexName)
+        .size
     } finally {
       dataset.close()
     }
@@ -295,6 +309,17 @@ class RangeBasedBTreeIndexJob(
 
   private val VALUE_COLUMN_NAME = "value"
 
+  /** Version `readOptions` is pinned to. Throws if the ref is not a version on main. */
+  private def pinnedVersion: Long = {
+    val ref = readOptions.getRef
+    if (ref == null || !ref.isMain || !ref.getVersionNumber.isPresent) {
+      throw new IllegalStateException(
+        "Range-mode BTree builds need read options pinned to a version on main so the scan and the " +
+          "segment stamp describe one snapshot; got a ref that names none")
+    }
+    ref.getVersionNumber.get.longValue()
+  }
+
   def run(): Seq[Index] = {
     if (addIndexExec.columns.size != 1) {
       throw new UnsupportedOperationException(
@@ -316,9 +341,10 @@ class RangeBasedBTreeIndexJob(
     }
     val fullTableName = parts.mkString(".")
 
-    // Read the indexed column with the row id and fragment id metadata columns.
     val fragmentColumn = LanceDataset.FRAGMENT_ID_COLUMN.name
-    val df = session.table(fullTableName)
+    val df = session.read
+      .option(LanceSparkReadOptions.CONFIG_VERSION, pinnedVersion)
+      .table(fullTableName)
     val selectDf = df.select(
       df.col(columns.head).as(VALUE_COLUMN_NAME),
       df.col(LanceDataset.ROW_ID_COLUMN.name),
@@ -619,6 +645,89 @@ object IndexUtils extends Logging {
       .get(method.toLowerCase(Locale.ROOT))
       .filter(scalarSegmentIndexTypes.contains)
 
+  /** Pins `readOptions` to the version `dataset` is open at. */
+  def pinVersion(
+      readOptions: LanceSparkReadOptions,
+      dataset: Dataset): LanceSparkReadOptions =
+    readOptions.withRef(Utils.pinOpenedRef(dataset, readOptions.getRef))
+
+  /**
+   * Fragment ids and live row counts of `dataset`, in manifest order.
+   *
+   * Reads the primitive fragment-statistics view rather than [[Dataset#getFragments]], which
+   * materializes a Java object per fragment and per data file. Both commands enumerate fragments on
+   * the driver, once to plan and once to check the commit, so on a large table that difference is
+   * the bulk of planning cost.
+   */
+  def fragmentWorkloads(dataset: Dataset): List[FragmentWorkload] = {
+    val stats = dataset.getFragmentStatistics
+    val ids = stats.getIds
+    val rowCounts = stats.getRowCounts
+    List.tabulate(ids.length)(index =>
+      FragmentWorkload(Integer.valueOf(ids(index)), rowCounts(index)))
+  }
+
+  /** Fragment ids live in `dataset`. See [[fragmentWorkloads]] for why this avoids getFragments. */
+  def liveFragmentIds(dataset: Dataset): Set[Int] =
+    dataset.getFragmentStatistics.getIds.toSet
+
+  /** Fragment ids the given segments declare coverage of. */
+  def declaredCoverage(segments: Seq[Index]): Set[Int] =
+    segments.iterator
+      .flatMap(_.fragments().orElse(Collections.emptyList[Integer]()).asScala)
+      .map(_.intValue)
+      .toSet
+
+  /**
+   * Refuses a segment set that declares coverage but intersects no live fragment. Partial loss is
+   * accepted: fragments can be retired by compaction or deletion during the build, and the segments
+   * covering what remains are still correct. [[establishedCoverage]] reports what was actually
+   * committed.
+   */
+  def requireCommittableCoverage(
+      liveFragmentIds: Set[Int],
+      segments: Seq[Index],
+      indexName: String): Unit = {
+    val declared = declaredCoverage(segments)
+    if (declared.nonEmpty && declared.intersect(liveFragmentIds).isEmpty) {
+      throw new IllegalStateException(
+        s"Index '$indexName' build raced a concurrent operation: every fragment it covers " +
+          s"(${describeFragmentIds(declared)}) was retired while the build ran, so the segments " +
+          "would cover nothing. No index change was committed; re-run the command.")
+    }
+  }
+
+  /**
+   * Committed fragment ids from the segments this build produced, intersected with what is still
+   * live. Only segments whose UUID matches `builtSegments` are counted; pre-existing survivors are
+   * not this command's coverage.
+   */
+  def establishedCoverage(
+      builtSegments: Seq[Index],
+      committedSegments: Seq[Index],
+      liveFragmentIds: Set[Int],
+      indexName: String): Set[Int] = {
+    val builtUuids = builtSegments.map(_.uuid).toSet
+    val established =
+      declaredCoverage(committedSegments.filter(segment => builtUuids.contains(segment.uuid)))
+        .intersect(liveFragmentIds)
+    val uncovered = declaredCoverage(builtSegments).diff(established)
+    if (uncovered.nonEmpty) {
+      logWarning(
+        s"Index '$indexName' build raced a concurrent operation: fragments " +
+          s"${describeFragmentIds(uncovered)} are not covered by this commit, because they were " +
+          "retired or because their indexed field was rewritten while the build ran. The segments " +
+          "for the remaining fragments are committed; re-run the command to cover them.")
+    }
+    established
+  }
+
+  private def describeFragmentIds(ids: Set[Int]): String = {
+    val ordered = ids.toSeq.sorted
+    val shown = ordered.take(10).mkString(", ")
+    if (ordered.size > 10) s"$shown, ... (${ordered.size} total)" else shown
+  }
+
   def resolveIndexField(
       schema: LanceSchema,
       indexType: IndexType,

lance-spark-base_2.12/src/test/java/org/lance/spark/branch/BaseBranchDDLTest.java

@@ -590,6 +590,23 @@ public void testBranchAndVersionOptionsFail() {
     Assertions.assertTrue(conflictMessages.contains("version"));
   }
 
+  @Test
+  public void testVersionOptionPinsTheScanToThatVersion() {
+    DatasetVersions versions = prepareDatasetWithHistory();
+    Assertions.assertNotEquals(versions.firstInsertVersion, versions.latestVersion);
+
+    Assertions.assertEquals(
+        5,
+        spark
+            .read()
+            .option("version", Long.toString(versions.firstInsertVersion))
+            .table(fullTable)
+            .count(),
+        "reading with the version option must see only the rows present at that version");
+    Assertions.assertEquals(
+        10, spark.table(fullTable).count(), "the unpinned read still sees the latest version");
+  }
+
   @Test
   public void testBranchIdentifierRejectsVersionAsOf() {
     DatasetVersions versions = prepareDatasetWithHistory();

lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java

@@ -17,8 +17,11 @@
 import org.lance.index.Index;
 import org.lance.index.IndexCriteria;
 import org.lance.index.IndexDescription;
+import org.lance.index.IndexOptions;
+import org.lance.index.IndexParams;
 import org.lance.index.IndexType;
 import org.lance.index.OptimizeOptions;
+import org.lance.index.scalar.ScalarIndexParams;
 import org.lance.ipc.FullTextQuery;
 import org.lance.ipc.LanceScanner;
 import org.lance.ipc.ScanOptions;
@@ -105,6 +108,64 @@ public void tearDown() throws IOException {
     }
   }
 
+  @Test
+  public void testSegmentCommitReportsCoverageAsCommitted() {
+    spark.sql(String.format("create table %s (id int) using lance", fullTable));
+    spark.sql(String.format("insert into %s values (0), (1), (2)", fullTable));
+    spark.sql(String.format("insert into %s values (3), (4), (5)", fullTable));
+
+    try (org.lance.Dataset committer =
+        Utils.openDatasetBuilder(LanceSparkReadOptions.from(tableDir)).build()) {
+      List<Fragment> fragments = committer.getFragments();
+      int coveredFragmentId = fragments.get(fragments.size() - 1).getId();
+
+      IndexParams indexParams =
+          IndexParams.builder()
+              .setScalarIndexParams(ScalarIndexParams.create("zonemap", "{}"))
+              .build();
+      Index built =
+          committer.createIndex(
+              IndexOptions.builder(Collections.singletonList("id"), IndexType.ZONEMAP, indexParams)
+                  .withIndexName("idx_committed_coverage")
+                  .replace(true)
+                  .withFragmentIds(Collections.singletonList(coveredFragmentId))
+                  .build());
+      Assertions.assertEquals(
+          Collections.singletonList(coveredFragmentId),
+          built.fragments().orElse(Collections.emptyList()),
+          "the uncommitted segment should declare the fragment it was built for");
+
+      // Retire every fragment from another handle, leaving the committer on a stale manifest.
+      spark.sql(String.format("delete from %s where id >= 0", fullTable));
+
+      List<Index> committed =
+          committer.commitExistingIndexSegments(
+              "idx_committed_coverage", "id", Collections.singletonList(built));
+
+      Set<UUID> ours = Collections.singleton(built.uuid());
+      Assertions.assertTrue(
+          committed.stream().anyMatch(index -> ours.contains(index.uuid())),
+          "the commit must return the metadata of the segments it was handed");
+
+      Set<Integer> liveAfter =
+          committer.getFragments().stream().map(Fragment::getId).collect(Collectors.toSet());
+      Assertions.assertFalse(
+          liveAfter.contains(coveredFragmentId),
+          "the committing handle must advance to the manifest the commit wrote");
+
+      Set<Integer> established =
+          committed.stream()
+              .filter(index -> ours.contains(index.uuid()))
+              .flatMap(index -> index.fragments().orElse(Collections.emptyList()).stream())
+              .filter(liveAfter::contains)
+              .collect(Collectors.toSet());
+      Assertions.assertEquals(
+          Collections.emptySet(),
+          established,
+          "coverage read from the committed state must not count a fragment retired in between");
+    }
+  }
+
   private void prepareDataset() {
     spark.sql(String.format("create table %s (id int, text string) using lance;", fullTable));
     // First insert to create initial fragments