From c5728ecc39d68a2c6d4821088a40b711a3fc332c Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Wed, 26 Aug 2026 15:34:40 +0800 Subject: [PATCH 1/3] feat: report segment index build progress --- docs/src/operations/ddl/create-index.md | 6 + .../datasources/v2/AddIndexExec.scala | 140 ++++++++++++++++-- .../lance/spark/update/BaseAddIndexTest.java | 15 ++ .../datasources/v2/IndexUtilsTest.scala | 80 ++++++++++ 4 files changed, 232 insertions(+), 9 deletions(-) diff --git a/docs/src/operations/ddl/create-index.md b/docs/src/operations/ddl/create-index.md index 03d1a1191..15cc8690e 100755 --- a/docs/src/operations/ddl/create-index.md +++ b/docs/src/operations/ddl/create-index.md @@ -284,6 +284,12 @@ 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 named `index build completed segments` and `index build total segments`. 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: diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala index 085090996..e5ae62ee4 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala @@ -23,6 +23,7 @@ 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.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 +39,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 +84,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 +198,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 +211,14 @@ case class AddIndexExec( throw new UnsupportedOperationException(s"Unsupported index type: $indexType") } + private def createIndexSegmentProgress(): SparkIndexSegmentProgress = + new SparkIndexSegmentProgress( + indexName, + completedDelta => metrics(AddIndexExec.INDEX_BUILD_COMPLETED_SEGMENTS).add(completedDelta), + totalDelta => metrics(AddIndexExec.INDEX_BUILD_TOTAL_SEGMENTS).add(totalDelta), + 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 +284,98 @@ 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, + 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) + 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) + 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 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 +590,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 +623,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,16 +829,26 @@ 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) + 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]) => execute(taskIterator.next()), + (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) } diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java index 711292373..bbfb6330c 100755 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java @@ -32,6 +32,8 @@ import org.apache.spark.sql.Row; import org.apache.spark.sql.RowFactory; import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.execution.CommandResultExec; +import org.apache.spark.sql.execution.metric.SQLMetric; import org.apache.spark.sql.types.DataTypes; import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.MetadataBuilder; @@ -46,6 +48,7 @@ import org.junit.jupiter.params.provider.Arguments; import org.junit.jupiter.params.provider.MethodSource; import org.junit.jupiter.params.provider.ValueSource; +import scala.collection.JavaConverters; import java.io.IOException; import java.nio.file.FileSystems; @@ -400,6 +403,18 @@ public void testCreateZonemapIndexWithNumSegments() { fragmentCount, coveredFragments, "Expected committed segments to cover all fragments exactly once"); + + CommandResultExec commandResult = (CommandResultExec) result.queryExecution().executedPlan(); + Map progressMetrics = + JavaConverters.mapAsJavaMap(commandResult.commandPhysicalPlan().metrics()); + Assertions.assertEquals( + expectedSegmentCount, + progressMetrics.get("indexBuildTotalSegments").value(), + "Expected progress to report the planned segment count"); + Assertions.assertEquals( + expectedSegmentCount, + progressMetrics.get("indexBuildCompletedSegments").value(), + "Expected progress to report every successfully built segment"); } finally { lanceDataset.close(); } diff --git a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala index 9644e6f1c..d15b9ad9c 100644 --- a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala +++ b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala @@ -16,8 +16,11 @@ package org.apache.spark.sql.execution.datasources.v2 import org.apache.spark.sql.catalyst.plans.logical.LanceNamedArgument import org.junit.jupiter.api.Assertions._ import org.junit.jupiter.api.Test +import org.junit.jupiter.api.function.Executable import org.lance.index.IndexType +import scala.collection.mutable.ArrayBuffer + /** * Unit tests for [[IndexUtils]] helper methods. * @@ -257,4 +260,81 @@ class IndexUtilsTest { classOf[ArithmeticException], () => IndexUtils.batchFragments(fragmentWorkloads(Long.MaxValue, 1), Some(1), 1)) } + + @Test + def indexSegmentProgress_reportsSuccessfulPartitionsOnce(): Unit = { + var completed = 0L + var total = 0L + val logs = ArrayBuffer.empty[String] + val progress = new SparkIndexSegmentProgress( + "idx_text", + delta => completed += delta, + delta => total += delta, + message => logs += message, + (_, _) => ()) + + progress.start(3) + progress.segmentComplete(2) + progress.segmentComplete(0) + progress.segmentComplete(2) + progress.segmentComplete(1) + + assertEquals(3L, total) + assertEquals( + 3L, + completed, + "a retried or duplicate partition result must not increment progress twice") + assertTrue(logs.head.contains("started")) + assertTrue(logs.last.contains("3/3 segments completed")) + } + + @Test + def indexSegmentProgress_ignoresObservationFailures(): Unit = { + var warnings = 0 + val progress = new SparkIndexSegmentProgress( + "idx_text", + _ => throw new RuntimeException("metric update failed"), + _ => throw new RuntimeException("metric update failed"), + _ => throw new RuntimeException("log failed"), + (_, _) => { + warnings += 1 + throw new RuntimeException("warn failed") + }) + + assertProgressDoesNotThrow { + progress.start(2) + } + assertProgressDoesNotThrow { + progress.segmentComplete(0) + } + assertTrue(warnings >= 4) + } + + @Test + def indexSegmentMetricDefinitions_onlyReportsForEagerSegmentBuilds(): Unit = { + val ftsMetricNames = + AddIndexExec.indexSegmentMetricDefinitions("fts", Seq.empty).keySet + + assertEquals( + Set( + AddIndexExec.INDEX_BUILD_COMPLETED_SEGMENTS, + AddIndexExec.INDEX_BUILD_TOTAL_SEGMENTS), + ftsMetricNames) + assertFalse(AddIndexExec.indexSegmentMetricDefinitions("BTREE", Seq.empty).isEmpty) + assertTrue( + AddIndexExec.indexSegmentMetricDefinitions( + "btree", + Seq(LanceNamedArgument("build_mode", "range"))).isEmpty) + assertTrue( + AddIndexExec.indexSegmentMetricDefinitions( + "fts", + Seq(LanceNamedArgument("train", java.lang.Boolean.FALSE))).isEmpty) + assertTrue(AddIndexExec.indexSegmentMetricDefinitions("ivf_pq", Seq.empty).isEmpty) + } + + private def assertProgressDoesNotThrow(callback: => Unit): Unit = { + assertDoesNotThrow(new Executable { + override def execute(): Unit = callback + }) + } } From ca860d91c036de0b95b76edb8476155f74d6b517 Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Thu, 27 Aug 2026 10:45:47 +0800 Subject: [PATCH 2/3] fix: publish segment index progress metrics --- .../datasources/v2/AddIndexExec.scala | 28 ++++++- .../lance/spark/update/BaseAddIndexTest.java | 76 ++++++++++++++++++- .../datasources/v2/IndexUtilsTest.scala | 59 +++++++++++++- 3 files changed, 156 insertions(+), 7 deletions(-) diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala index e5ae62ee4..ed9cdadef 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala @@ -23,6 +23,7 @@ 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 @@ -211,13 +212,25 @@ case class AddIndexExec( throw new UnsupportedOperationException(s"Unsupported index type: $indexType") } - private def createIndexSegmentProgress(): SparkIndexSegmentProgress = + 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. + val executionId = sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY) + new SparkIndexSegmentProgress( indexName, - completedDelta => metrics(AddIndexExec.INDEX_BUILD_COMPLETED_SEGMENTS).add(completedDelta), - totalDelta => metrics(AddIndexExec.INDEX_BUILD_TOTAL_SEGMENTS).add(totalDelta), + 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( @@ -315,6 +328,7 @@ private[datasources] class SparkIndexSegmentProgress( indexName: String, addCompletedDelta: Long => Unit, addTotalDelta: Long => Unit, + publishMetricUpdates: () => Unit, logStatus: String => Unit, logWarningStatus: (String, Throwable) => Unit) { @@ -328,6 +342,7 @@ private[datasources] class SparkIndexSegmentProgress( addTotalDelta, "update index build total segments metric", total.toLong - previousTotal) + publishMetrics() logProgress(s"Index '$indexName' segment build started (0/$total segments)") } @@ -338,6 +353,7 @@ private[datasources] class SparkIndexSegmentProgress( addCompletedDelta, "update index build completed segments metric", 1L) + publishMetrics() logProgress( s"Index '$indexName' segment build progress: " + s"$completed/${totalSegments.get()} segments completed") @@ -358,6 +374,12 @@ private[datasources] class SparkIndexSegmentProgress( } } + private def publishMetrics(): Unit = { + observe("publish index build metrics") { + publishMetricUpdates() + } + } + private def observe(action: String)(callback: => Unit): Unit = { try { callback diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java index bbfb6330c..8070c9238 100755 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java @@ -28,12 +28,15 @@ import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.spark.SparkException; +import org.apache.spark.scheduler.SparkListener; +import org.apache.spark.scheduler.SparkListenerEvent; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; import org.apache.spark.sql.RowFactory; import org.apache.spark.sql.SparkSession; import org.apache.spark.sql.execution.CommandResultExec; import org.apache.spark.sql.execution.metric.SQLMetric; +import org.apache.spark.sql.execution.ui.SparkListenerDriverAccumUpdates; import org.apache.spark.sql.types.DataTypes; import org.apache.spark.sql.types.Metadata; import org.apache.spark.sql.types.MetadataBuilder; @@ -57,17 +60,45 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; import java.util.UUID; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; /** Base test for distributed CREATE INDEX. */ public abstract class BaseAddIndexTest { + private static class DriverMetricSnapshotListener extends SparkListener { + private final List> snapshots = new CopyOnWriteArrayList<>(); + + @Override + public void onOtherEvent(SparkListenerEvent event) { + if (event instanceof SparkListenerDriverAccumUpdates) { + SparkListenerDriverAccumUpdates updates = (SparkListenerDriverAccumUpdates) event; + Map snapshot = new HashMap<>(); + JavaConverters.seqAsJavaList(updates.accumUpdates()) + .forEach( + update -> + snapshot.put( + ((Number) update._1()).longValue(), ((Number) update._2()).longValue())); + snapshots.add(snapshot); + } + } + + List> snapshotsContaining(long firstMetricId, long secondMetricId) { + return snapshots.stream() + .filter( + snapshot -> + snapshot.containsKey(firstMetricId) && snapshot.containsKey(secondMetricId)) + .collect(Collectors.toList()); + } + } + protected String catalogName = "lance_test"; protected String tableName = "create_index_test"; protected String fullTable = catalogName + ".default." + tableName; @@ -364,9 +395,11 @@ public void testCreateZonemapIndex() { } @Test - public void testCreateZonemapIndexWithNumSegments() { + public void testCreateZonemapIndexWithNumSegments() throws Exception { prepareDataset(); + DriverMetricSnapshotListener metricListener = new DriverMetricSnapshotListener(); + spark.sparkContext().addSparkListener(metricListener); Dataset result = spark.sql( String.format( @@ -407,14 +440,51 @@ public void testCreateZonemapIndexWithNumSegments() { CommandResultExec commandResult = (CommandResultExec) result.queryExecution().executedPlan(); Map progressMetrics = JavaConverters.mapAsJavaMap(commandResult.commandPhysicalPlan().metrics()); + SQLMetric totalMetric = progressMetrics.get("indexBuildTotalSegments"); + SQLMetric completedMetric = progressMetrics.get("indexBuildCompletedSegments"); Assertions.assertEquals( expectedSegmentCount, - progressMetrics.get("indexBuildTotalSegments").value(), + totalMetric.value(), "Expected progress to report the planned segment count"); Assertions.assertEquals( expectedSegmentCount, - progressMetrics.get("indexBuildCompletedSegments").value(), + completedMetric.value(), "Expected progress to report every successfully built segment"); + + spark.sparkContext().listenerBus().waitUntilEmpty(5000); + List> progressSnapshots = + metricListener.snapshotsContaining(completedMetric.id(), totalMetric.id()); + Assertions.assertFalse( + progressSnapshots.isEmpty(), + "Expected CREATE INDEX driver metrics in SparkListenerDriverAccumUpdates"); + + long previousCompleted = -1L; + boolean observedStart = false; + boolean observedIntermediate = false; + boolean observedCompletion = false; + for (Map snapshot : progressSnapshots) { + long completed = snapshot.get(completedMetric.id()); + long total = snapshot.get(totalMetric.id()); + Assertions.assertEquals( + expectedSegmentCount, + total, + "Every published snapshot should retain the planned segment count"); + Assertions.assertTrue( + completed >= 0 && completed <= total, + "Published completed segments must remain within [0, total]"); + Assertions.assertTrue( + completed >= previousCompleted, + "Published completed segment snapshots must be monotonic"); + previousCompleted = completed; + observedStart |= completed == 0; + observedIntermediate |= completed > 0 && completed < total; + observedCompletion |= completed == total; + } + + Assertions.assertTrue(observedStart, "Expected an initial 0/total progress snapshot"); + Assertions.assertTrue( + observedIntermediate, "Expected a progress snapshot before every segment completed"); + Assertions.assertTrue(observedCompletion, "Expected a final total/total progress snapshot"); } finally { lanceDataset.close(); } diff --git a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala index d15b9ad9c..927ccb3bd 100644 --- a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala +++ b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala @@ -266,10 +266,12 @@ class IndexUtilsTest { var completed = 0L var total = 0L val logs = ArrayBuffer.empty[String] + val publishedSnapshots = ArrayBuffer.empty[(Long, Long)] val progress = new SparkIndexSegmentProgress( "idx_text", delta => completed += delta, delta => total += delta, + () => publishedSnapshots += completed -> total, message => logs += message, (_, _) => ()) @@ -284,10 +286,64 @@ class IndexUtilsTest { 3L, completed, "a retried or duplicate partition result must not increment progress twice") + assertEquals( + Seq(0L -> 3L, 1L -> 3L, 2L -> 3L, 3L -> 3L), + publishedSnapshots.toSeq, + "publish a start snapshot and one snapshot per unique successful partition") assertTrue(logs.head.contains("started")) assertTrue(logs.last.contains("3/3 segments completed")) } + @Test + def indexSegmentProgress_isolatesPublicationFailures(): Unit = { + var completed = 0L + var total = 0L + var logs = 0 + var warnings = 0 + val progress = new SparkIndexSegmentProgress( + "idx_text", + delta => completed += delta, + delta => total += delta, + () => throw new RuntimeException("publication failed"), + _ => logs += 1, + (_, _) => warnings += 1) + + assertProgressDoesNotThrow { + progress.start(2) + } + assertProgressDoesNotThrow { + progress.segmentComplete(0) + } + + assertEquals(2L, total) + assertEquals(1L, completed) + assertEquals(2, logs, "publication failures must not suppress progress logging") + assertEquals(2, warnings, "each failed publication should be reported once") + } + + @Test + def indexSegmentProgress_publishesWhenOtherObservationCallbacksFail(): Unit = { + var publications = 0 + var warnings = 0 + val progress = new SparkIndexSegmentProgress( + "idx_text", + _ => throw new RuntimeException("metric update failed"), + _ => throw new RuntimeException("metric update failed"), + () => publications += 1, + _ => throw new RuntimeException("log failed"), + (_, _) => warnings += 1) + + assertProgressDoesNotThrow { + progress.start(2) + } + assertProgressDoesNotThrow { + progress.segmentComplete(0) + } + + assertEquals(2, publications, "metric or log failures must not suppress publication attempts") + assertEquals(4, warnings) + } + @Test def indexSegmentProgress_ignoresObservationFailures(): Unit = { var warnings = 0 @@ -295,6 +351,7 @@ class IndexUtilsTest { "idx_text", _ => throw new RuntimeException("metric update failed"), _ => throw new RuntimeException("metric update failed"), + () => throw new RuntimeException("publication failed"), _ => throw new RuntimeException("log failed"), (_, _) => { warnings += 1 @@ -307,7 +364,7 @@ class IndexUtilsTest { assertProgressDoesNotThrow { progress.segmentComplete(0) } - assertTrue(warnings >= 4) + assertEquals(6, warnings) } @Test From e79d0605824f4296686ca7558db2e3a5635dfdaf Mon Sep 17 00:00:00 2001 From: "zhanghaobo@kanzhun.com" Date: Thu, 27 Aug 2026 15:52:39 +0800 Subject: [PATCH 3/3] fix: harden segment progress reporting --- docs/src/operations/ddl/create-index.md | 9 +++--- .../datasources/v2/AddIndexExec.scala | 19 +++++++++++- .../lance/spark/update/BaseAddIndexTest.java | 18 +++++++++-- .../datasources/v2/IndexUtilsTest.scala | 31 +++++++++++++++++++ 4 files changed, 70 insertions(+), 7 deletions(-) diff --git a/docs/src/operations/ddl/create-index.md b/docs/src/operations/ddl/create-index.md index 15cc8690e..b40f0fab5 100755 --- a/docs/src/operations/ddl/create-index.md +++ b/docs/src/operations/ddl/create-index.md @@ -285,10 +285,11 @@ The `CREATE INDEX` command returns the following information about the operation | `index_name` | String | The name of the created index. | For eager distributed segment builds, Lance Spark reports driver-side progress through Spark SQL -metrics named `index build completed segments` and `index build total segments`. 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. +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 diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala index ed9cdadef..872b4ab71 100755 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/AddIndexExec.scala @@ -217,6 +217,7 @@ case class AddIndexExec( 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( @@ -858,6 +859,8 @@ object IndexUtils extends Logging { } else { try { 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 @@ -865,7 +868,7 @@ object IndexUtils extends Logging { // exposing live progress before the full job result is available. sc.runJob( taskRdd, - (taskIterator: Iterator[T]) => execute(taskIterator.next()), + (taskIterator: Iterator[T]) => executeSinglePartitionTask(taskIterator)(execute), (partitionId: Int, encoded: String) => { encodedResults(partitionId) = encoded progress.segmentComplete(partitionId) @@ -877,6 +880,20 @@ object IndexUtils extends Logging { } } + 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], diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java index 8070c9238..0b11a6490 100755 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/update/BaseAddIndexTest.java @@ -67,6 +67,7 @@ import java.util.Set; import java.util.UUID; import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeoutException; import java.util.stream.Collectors; import java.util.stream.IntStream; import java.util.stream.Stream; @@ -104,6 +105,7 @@ List> snapshotsContaining(long firstMetricId, long secondMetricI protected String fullTable = catalogName + ".default." + tableName; protected SparkSession spark; + private DriverMetricSnapshotListener metricListener; @TempDir Path tempDir; protected String tableDir; @@ -135,6 +137,9 @@ public void setup() throws IOException { @AfterEach public void tearDown() throws IOException { if (spark != null) { + if (metricListener != null) { + spark.sparkContext().removeSparkListener(metricListener); + } spark.close(); } } @@ -398,7 +403,7 @@ public void testCreateZonemapIndex() { public void testCreateZonemapIndexWithNumSegments() throws Exception { prepareDataset(); - DriverMetricSnapshotListener metricListener = new DriverMetricSnapshotListener(); + metricListener = new DriverMetricSnapshotListener(); spark.sparkContext().addSparkListener(metricListener); Dataset result = spark.sql( @@ -451,7 +456,16 @@ public void testCreateZonemapIndexWithNumSegments() throws Exception { completedMetric.value(), "Expected progress to report every successfully built segment"); - spark.sparkContext().listenerBus().waitUntilEmpty(5000); + try { + spark.sparkContext().listenerBus().waitUntilEmpty(10000); + } catch (TimeoutException timeout) { + List> receivedSnapshots = + metricListener.snapshotsContaining(completedMetric.id(), totalMetric.id()); + Assertions.fail( + "Timed out waiting for Spark listener events; received progress snapshots: " + + receivedSnapshots, + timeout); + } List> progressSnapshots = metricListener.snapshotsContaining(completedMetric.id(), totalMetric.id()); Assertions.assertFalse( diff --git a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala index 927ccb3bd..5615bab3b 100644 --- a/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala +++ b/lance-spark-base_2.12/src/test/scala/org/apache/spark/sql/execution/datasources/v2/IndexUtilsTest.scala @@ -261,6 +261,37 @@ class IndexUtilsTest { () => IndexUtils.batchFragments(fragmentWorkloads(Long.MaxValue, 1), Some(1), 1)) } + @Test + def executeSinglePartitionTask_executesTheOnlyTask(): Unit = { + assertEquals( + "segment-7", + IndexUtils.executeSinglePartitionTask(Iterator(7))(task => s"segment-$task")) + } + + @Test + def executeSinglePartitionTask_rejectsEmptyPartitions(): Unit = { + val error = assertThrows( + classOf[IllegalStateException], + () => IndexUtils.executeSinglePartitionTask[Int, String](Iterator.empty)(_.toString)) + + assertTrue(error.getMessage.contains("partition was empty")) + } + + @Test + def executeSinglePartitionTask_rejectsMultipleTasksBeforeExecution(): Unit = { + var executed = false + val error = assertThrows( + classOf[IllegalStateException], + () => + IndexUtils.executeSinglePartitionTask(Iterator(1, 2)) { task => + executed = true + task.toString + }) + + assertFalse(executed, "do not build any segment when the partition invariant is violated") + assertTrue(error.getMessage.contains("partition contained multiple tasks")) + } + @Test def indexSegmentProgress_reportsSuccessfulPartitionsOnce(): Unit = { var completed = 0L