feat: report segment index build progress

docs/src/operations/ddl/create-index.md

@@ -284,6 +284,13 @@ The `CREATE INDEX` command returns the following information about the operation
 | `fragments_indexed` | Long   | The number of fragments that were indexed. |
 | `index_name`        | String | The name of the created index.         |
 
+For eager distributed segment builds, Lance Spark reports driver-side progress through Spark SQL
+metrics displayed as `index build completed segments` and `index build total segments`. Their
+executed-plan metric keys are `indexBuildCompletedSegments` and `indexBuildTotalSegments`,
+respectively. The driver also logs progress as successful segment tasks return. Task retries and
+speculative attempts are counted once per successful Spark partition. Progress reporting is
+informational and does not change the command output or the atomic segment commit.
+
 ## When to Use an Index
 
 Consider creating an index when:

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

@@ -23,6 +23,8 @@ import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.{Attribute, GenericInternalRow}
 import org.apache.spark.sql.catalyst.plans.logical.{AddIndexOutputType, LanceNamedArgument}
 import org.apache.spark.sql.connector.catalog.{Identifier, TableCatalog}
+import org.apache.spark.sql.execution.SQLExecution
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.LanceArrowUtils
 import org.apache.spark.sql.util.LanceSerializeUtil.{decode, encode}
@@ -38,10 +40,13 @@ import org.lance.spark.utils.{CloseableUtil, FieldPathUtils, Utils}
 import org.lance.spark.write.SingleBatchArrowReader
 
 import java.util.{Collections, Locale, UUID}
+import java.util.concurrent.ConcurrentHashMap
+import java.util.concurrent.atomic.AtomicLong
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.{ArrayBuffer, PriorityQueue}
 import scala.reflect.ClassTag
+import scala.util.control.NonFatal
 
 /**
  * Physical execution of distributed CREATE INDEX (ALTER TABLE ... CREATE INDEX ...) for Lance datasets.
@@ -80,6 +85,11 @@ case class AddIndexExec(
 
   override def output: Seq[Attribute] = AddIndexOutputType.SCHEMA
 
+  override lazy val metrics: Map[String, SQLMetric] =
+    AddIndexExec.indexSegmentMetricDefinitions(method, args).map { case (name, description) =>
+      name -> SQLMetrics.createMetric(sparkContext, description)
+    }.toMap
+
   override protected def run(): Seq[InternalRow] = {
     val lanceDataset = LanceDataset.requireWritable(catalog.loadTable(ident), "AddIndex")
 
@@ -189,7 +199,8 @@ case class AddIndexExec(
         nsImpl,
         nsProps,
         tableId,
-        initialStorageOpts)
+        initialStorageOpts,
+        createIndexSegmentProgress())
       val segments = segmentJob.run()
       // Atomic add+remove via Lance core; see commitIndexSegments
       commitIndexSegments(readOptions, canonicalColumns.head, segments)
@@ -201,6 +212,27 @@ case class AddIndexExec(
     throw new UnsupportedOperationException(s"Unsupported index type: $indexType")
   }
 
+  private def createIndexSegmentProgress(): SparkIndexSegmentProgress = {
+    val completedMetric = metrics(AddIndexExec.INDEX_BUILD_COMPLETED_SEGMENTS)
+    val totalMetric = metrics(AddIndexExec.INDEX_BUILD_TOTAL_SEGMENTS)
+    // runJob invokes its result handler on the DAGScheduler thread, which does not inherit the
+    // command thread's Spark local properties. Capture the SQL execution ID before entering it.
+    // Spark tolerates a null ID by skipping listener publication; driver logs remain available.
+    val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+
+    new SparkIndexSegmentProgress(
+      indexName,
+      completedDelta => completedMetric.add(completedDelta),
+      totalDelta => totalMetric.add(totalDelta),
+      () =>
+        SQLMetrics.postDriverMetricUpdates(
+          sparkContext,
+          executionId,
+          Seq(completedMetric, totalMetric)),
+      message => logInfo(message),
+      (message, cause) => logWarning(message, cause))
+  }
+
   /** Commits an empty (untrained) index on the driver, with an empty fragment bitmap. */
   private def commitEmptyIndex(
       dataset: Dataset,
@@ -266,6 +298,107 @@ case class AddIndexExec(
 
 }
 
+private[datasources] object AddIndexExec {
+  private[datasources] val INDEX_BUILD_COMPLETED_SEGMENTS = "indexBuildCompletedSegments"
+  private[datasources] val INDEX_BUILD_TOTAL_SEGMENTS = "indexBuildTotalSegments"
+
+  private[datasources] def indexSegmentMetricDefinitions(
+      method: String,
+      args: Seq[LanceNamedArgument]): Map[String, String] = {
+    val usesSegmentBuild =
+      try {
+        val indexType = IndexUtils.buildIndexType(method)
+        IndexUtils.extractTrain(args) &&
+        IndexUtils.scalarSegmentIndexType(method).isDefined &&
+        !IndexUtils.btreeBuildMode(indexType, args).contains("range")
+      } catch {
+        case NonFatal(_) => false
+      }
+
+    if (usesSegmentBuild) {
+      Map(
+        INDEX_BUILD_COMPLETED_SEGMENTS -> "index build completed segments",
+        INDEX_BUILD_TOTAL_SEGMENTS -> "index build total segments")
+    } else {
+      Map.empty
+    }
+  }
+}
+
+private[datasources] class SparkIndexSegmentProgress(
+    indexName: String,
+    addCompletedDelta: Long => Unit,
+    addTotalDelta: Long => Unit,
+    publishMetricUpdates: () => Unit,
+    logStatus: String => Unit,
+    logWarningStatus: (String, Throwable) => Unit) {
+
+  private val totalSegments = new AtomicLong(0L)
+  private val completedSegments = new AtomicLong(0L)
+  private val completedPartitions = ConcurrentHashMap.newKeySet[Integer]()
+
+  def start(total: Int): Unit = {
+    val previousTotal = totalSegments.getAndSet(total.toLong)
+    addDelta(
+      addTotalDelta,
+      "update index build total segments metric",
+      total.toLong - previousTotal)
+    publishMetrics()
+    logProgress(s"Index '$indexName' segment build started (0/$total segments)")
+  }
+
+  def segmentComplete(partitionId: Int): Unit = {
+    if (completedPartitions.add(Integer.valueOf(partitionId))) {
+      val completed = completedSegments.incrementAndGet()
+      addDelta(
+        addCompletedDelta,
+        "update index build completed segments metric",
+        1L)
+      publishMetrics()
+      logProgress(
+        s"Index '$indexName' segment build progress: " +
+          s"$completed/${totalSegments.get()} segments completed")
+    }
+  }
+
+  private def addDelta(add: Long => Unit, action: String, delta: Long): Unit = {
+    if (delta != 0L) {
+      observe(action) {
+        add(delta)
+      }
+    }
+  }
+
+  private def logProgress(message: String): Unit = {
+    observe("log index build progress") {
+      logStatus(message)
+    }
+  }
+
+  private def publishMetrics(): Unit = {
+    observe("publish index build metrics") {
+      publishMetricUpdates()
+    }
+  }
+
+  private def observe(action: String)(callback: => Unit): Unit = {
+    try {
+      callback
+    } catch {
+      case NonFatal(e) =>
+        warn(s"Ignoring failure to $action for index '$indexName'", e)
+    }
+  }
+
+  private def warn(message: String, cause: Throwable): Unit = {
+    try {
+      logWarningStatus(message, cause)
+    } catch {
+      case NonFatal(_) =>
+    }
+  }
+}
+
 /**
  * A job implementation for creating range-based BTree indexes using preprocessed, globally sorted data.
  * This approach distributes data across multiple partitions based on ranges of values and creates
@@ -480,7 +613,8 @@ class ScalarSegmentIndexJob(
     nsImpl: Option[String],
     nsProps: Option[Map[String, String]],
     tableId: Option[List[String]],
-    initialStorageOpts: Option[Map[String, String]]) {
+    initialStorageOpts: Option[Map[String, String]],
+    progress: SparkIndexSegmentProgress) {
 
   def run(): Seq[Index] = {
     val indexType = IndexUtils.scalarSegmentIndexType(addIndexExec.method).getOrElse {
@@ -512,7 +646,8 @@ class ScalarSegmentIndexJob(
       addIndexExec.session.sparkContext,
       tasks,
       s"${indexType.name()} index build failed. Uncommitted segments are not " +
-        "visible to readers and will not affect query correctness.")(_.execute())
+        "visible to readers and will not affect query correctness.",
+      progress)(_.execute())
   }
 }
 
@@ -717,22 +852,48 @@ object IndexUtils extends Logging {
   def runSegmentTasks[T <: Serializable: ClassTag](
       sc: org.apache.spark.SparkContext,
       tasks: Seq[T],
-      failureMessage: String)(execute: T => String): Seq[Index] = {
+      failureMessage: String,
+      progress: SparkIndexSegmentProgress)(execute: T => String): Seq[Index] = {
     if (tasks.isEmpty) {
       Seq.empty
     } else {
       try {
-        sc.parallelize(tasks, tasks.size)
-          .map(execute)
-          .collect()
-          .map(encoded => decode[Index](encoded))
-          .toSeq
+        val encodedResults = new Array[String](tasks.size)
+        // parallelize with one slice per task must produce exactly one task in every partition.
+        // Validate that invariant inside the job so empty or multi-task partitions fail clearly.
+        val taskRdd = sc.parallelize(tasks, tasks.size)
+        progress.start(tasks.size)
+        // Spark invokes the result handler on the driver once for each successful output
+        // partition. This preserves the independent-segment retry/speculation semantics while
+        // exposing live progress before the full job result is available.
+        sc.runJob(
+          taskRdd,
+          (taskIterator: Iterator[T]) => executeSinglePartitionTask(taskIterator)(execute),
+          (partitionId: Int, encoded: String) => {
+            encodedResults(partitionId) = encoded
+            progress.segmentComplete(partitionId)
+          })
+        encodedResults.map(encoded => decode[Index](encoded)).toSeq
       } catch {
         case e: Exception => throw new RuntimeException(failureMessage, e)
       }
     }
   }
 
+  private[datasources] def executeSinglePartitionTask[T, R](
+      taskIterator: Iterator[T])(execute: T => R): R = {
+    if (!taskIterator.hasNext) {
+      throw new IllegalStateException(
+        "Expected exactly one segment task per Spark partition, but partition was empty")
+    }
+    val task = taskIterator.next()
+    if (taskIterator.hasNext) {
+      throw new IllegalStateException(
+        "Expected exactly one segment task per Spark partition, but partition contained multiple tasks")
+    }
+    execute(task)
+  }
+
   def batchFragments(
       fragments: List[FragmentWorkload],
       numSegments: Option[Int],