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,25 @@ case class AddIndexExec(
     val btreeBuildMode = IndexUtils.btreeBuildMode(indexType, args)
     val scalarSegmentIndexType = IndexUtils.scalarSegmentIndexType(method)
 
-    val (fragmentWorkloads, canonicalColumns) = {
+    // Plan and build against a single pinned version. Tasks open the dataset themselves, so without
+    // pinning each one resolves the latest version independently. The driver's batches fix which
+    // fragments a segment covers either way, but not the version each segment records, and that
+    // version is what core checks against. It revalidates a segment's coverage of a fragment only
+    // when the recorded version predates the commit and the fragment already existed at it, so a
+    // task opening after a concurrent rewrite of the indexed column stamps the current version over
+    // keys it read earlier and those keys are trusted unchecked. Coverage the commit cannot
+    // establish is accounted for at commit time; see IndexUtils.establishedCoverage.
+    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 +170,31 @@ case class AddIndexExec(
     val (nsImpl, nsProps, tableId, initialStorageOpts) =
       extractNamespaceInfo(lanceDataset, readOptions)
 
-    // Range-mode BTree uses preprocessed data from Spark and keeps its dedicated path.
+    // Range-mode BTree uses preprocessed data from Spark and keeps its dedicated path: its coverage
+    // follows the fragment ids in the scanned rows rather than the fragment list planned above. The
+    // pinned options drive both halves of that path, the producer scan and the builder, so the rows a
+    // segment holds and the version it is stamped with describe the same snapshot; the job explains
+    // why both are needed.
     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 +203,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 +242,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()
     }
@@ -276,6 +302,13 @@ case class AddIndexExec(
  * covering exactly those fragments, so the resulting segments have disjoint fragment coverage and
  * can be committed directly as a single logical index.
  *
+ * Unlike the segmented path, coverage here is derived from the fragment ids present in the scanned
+ * rows rather than from a fragment list fixed at planning time, so the producer scan is pinned to the
+ * same version as the build. Both halves have to describe one snapshot: the segment records the
+ * version the builder opened, and core validates a segment's coverage against that version only for
+ * a fragment that existed at it. A scan left to resolve its own version could hand the builder rows
+ * from a fragment appended afterwards, whose keys would then be committed unvalidated.
+ *
  * @param addIndexExec       The AddIndexExec instance that initiated this job
  * @param readOptions        Configuration options for reading the Lance dataset
  * @param numFragments       Number of fragments in the dataset, used to bound shuffle partitions
@@ -295,6 +328,26 @@ class RangeBasedBTreeIndexJob(
 
   private val VALUE_COLUMN_NAME = "value"
 
+  /**
+   * The version [[readOptions]] is pinned to.
+   *
+   * The caller pins these options with [[IndexUtils.pinVersion]], and CREATE INDEX only runs against
+   * a writable table, so the ref names a version on main; a branch or tag target is rejected before
+   * this point by `LanceDataset.ensureWritable`. Anything else cannot be expressed as a scan option
+   * at all, since branch and version are mutually exclusive there, so it would leave the scan
+   * unpinned and reading a different snapshot than the build stamps. That is the defect this pin
+   * exists to prevent, and it is silent, so refuse rather than fall through to it.
+   */
+  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 +369,19 @@ class RangeBasedBTreeIndexJob(
     }
     val fullTableName = parts.mkString(".")
 
-    // Read the indexed column with the row id and fragment id metadata columns.
+    // Read the indexed column with the row id and fragment id metadata columns, at the version the
+    // build is pinned to. The scan otherwise resolves its own version through the catalog, which is
+    // at or after the pinned one, and coverage here comes from the fragment ids in the scanned rows:
+    // a fragment appended after planning would be declared covered by a segment stamped with a
+    // version that predates it, and core skips staleness validation for a fragment absent at the
+    // stamped version, so its keys would be trusted however they changed afterwards. Pinning the
+    // scan keeps the rows and the stamp describing one snapshot. Fragments appended after planning
+    // are simply not part of this build, exactly as in the segmented path, whose fragment list is
+    // fixed at planning time too; Dataset.optimizeIndices covers them incrementally.
     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 +682,115 @@ object IndexUtils extends Logging {
       .get(method.toLowerCase(Locale.ROOT))
       .filter(scalarSegmentIndexTypes.contains)
 
+  /**
+   * Pins `readOptions` to the version `dataset` is open at.
+   *
+   * Distributed index builds hand read options to tasks that open the dataset themselves. Pinning
+   * makes every task observe the fragment set the driver planned over instead of resolving the
+   * latest version independently.
+   */
+  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 to publish a segment set that would establish no coverage at all.
+   *
+   * Commit intersects each segment's declared coverage with the dataset's live fragments, so a
+   * fragment retired while the build ran contributes nothing. Lance accepts such a set: an empty
+   * fragment bitmap is valid metadata, and an existing segment is trivially disjoint from it and so
+   * survives. Nothing is corrupted, but the transaction would publish segments that index no data and
+   * report a build that achieved nothing, and this is the last point at which it can still be
+   * declined.
+   *
+   * Partial loss is not refused. A fragment leaves the manifest either because its rows moved
+   * (compaction, an in-place rewrite) or because they were all deleted; only the first leaves data
+   * unindexed, and either way the segments covering what remains are correct. Discarding a finished
+   * distributed build over that would be the wrong response, and it would make a routine concurrent
+   * DELETE fatal on a long one. What the commit actually established is reported afterwards, by
+   * [[establishedCoverage]].
+   */
+  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.")
+    }
+  }
+
+  /**
+   * The coverage a commit established, read back from the metadata the commit returned.
+   *
+   * A check taken before the commit cannot answer this, for two reasons. It reads the manifest its
+   * handle was opened at, while the commit lands on whatever version is current by then. And Lance
+   * prunes an incoming segment's coverage not only for a fragment that is gone but also for one whose
+   * indexed field was rewritten under the same id, which no comparison of fragment ids can see. The
+   * returned metadata is post-pruning, so it is the only truthful account of what was indexed.
+   *
+   * Only the segments this build produced are counted. Existing segments that were disjoint from them
+   * survive the commit, and their coverage is not this command's to report.
+   *
+   * @return the fragment ids the commit covered with these segments
+   */
+  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,28 @@ public void testBranchAndVersionOptionsFail() {
     Assertions.assertTrue(conflictMessages.contains("version"));
   }
 
+  /**
+   * The version option has to pin the scan itself, not merely be accepted. Distributed index builds
+   * rely on it to read the same snapshot the build is stamped with, and a silently ignored option
+   * would put the scan back on the latest version without any visible failure.
+   */
+  @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();