Update FlintSpark to use progress trackers

Signed-off-by: Simeon Widdis <[email protected]>

Author: Swiddis

flint-commons/src/main/scala/org/opensearch/flint/common/metadata/FlintMetadata.scala

@@ -41,7 +41,9 @@ case class FlintMetadata(
      */
     latestLogEntry: Option[FlintMetadataLogEntry] = None,
     /** Optional Flint index settings. TODO: move elsewhere? */
-    indexSettings: Option[String]) {
+    indexSettings: Option[String],
+    /** Information on the current progress of the index job */
+    currentProgress: Option[Map[String, Double]]) {
 
   require(version != null, "version is required")
   require(name != null, "name is required")
@@ -69,6 +71,7 @@ object FlintMetadata {
     private var latestId: Option[String] = None
     private var latestLogEntry: Option[FlintMetadataLogEntry] = None
     private var indexSettings: Option[String] = None
+    private var currentProgress: Option[Map[String, Double]] = None
 
     def version(version: FlintVersion): this.type = {
       this.version = version
@@ -131,6 +134,11 @@ object FlintMetadata {
       this
     }
 
+    def currentProgress(currentProgress: Map[String, Double]): this.type = {
+      this.currentProgress = Some(currentProgress)
+      this
+    }
+
     // Build method to create the FlintMetadata instance
     def build(): FlintMetadata = {
       FlintMetadata(
@@ -144,7 +152,8 @@ object FlintMetadata {
         schema = schema,
         indexSettings = indexSettings,
         latestId = latestId,
-        latestLogEntry = latestLogEntry)
+        latestLogEntry = latestLogEntry,
+        currentProgress = currentProgress)
     }
   }
 }

flint-core/src/main/scala/org/opensearch/flint/core/metrics/ProgressListener.scala

@@ -7,7 +7,17 @@ package org.opensearch.flint.core.metrics
 
 import org.apache.spark.scheduler.{SparkListenerStageSubmitted, SparkListenerTaskEnd}
 
-case class Progress(bytesCompleted: Long, estimatedBytesTotal: Double, estimatedBytesPerSecond: Double) {}
+case class Progress(
+    bytesCompleted: Long,
+    estimatedBytesTotal: Double,
+    estimatedBytesPerSecond: Double) {
+  def asMap(): Map[String, Double] = {
+    Map(
+      ("bytes_read", bytesCompleted.asInstanceOf[Double]),
+      ("est_bytes_total", estimatedBytesTotal),
+      ("est_bytes_per_second", estimatedBytesPerSecond))
+  }
+}
 
 /**
  * Collect and emit metrics by listening spark events

flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSpark.scala

@@ -556,11 +556,16 @@ class FlintSpark(val spark: SparkSession) extends FlintSparkTransactionSupport w
       .transientLog(latest => {
         val currentTime = System.currentTimeMillis()
         val updatedLatest = latest
-          .copy(state = REFRESHING, createTime = currentTime, lastRefreshStartTime = currentTime, progress = indexRefresh.progress())
+          .copy(state = REFRESHING, createTime = currentTime, lastRefreshStartTime = currentTime)
         flintMetadataCacheWriter
           .updateMetadataCache(
             indexName,
-            index.metadata.copy(latestLogEntry = Some(updatedLatest)))
+            index.metadata.copy(
+              latestLogEntry = Some(updatedLatest),
+              currentProgress = indexRefresh.progress() match {
+                case Some(p) => Some(p.asMap())
+                case None => None
+              }))
         updatedLatest
       })
       .finalLog(latest => {

flint-spark-integration/src/main/scala/org/opensearch/flint/spark/metadatacache/FlintMetadataCache.scala flint-spark-integration/src/main/scala/org/opensearch/flint/spark/metadatacache/ExportedFlintMetadata.scala

@@ -9,15 +9,17 @@ import scala.collection.JavaConverters.mapAsScalaMapConverter
 
 import org.opensearch.flint.common.metadata.FlintMetadata
 import org.opensearch.flint.common.metadata.log.FlintMetadataLogEntry
+import org.opensearch.flint.core.metrics.Progress
 import org.opensearch.flint.spark.FlintSparkIndexOptions
 import org.opensearch.flint.spark.mv.FlintSparkMaterializedView.{getSourceTablesFromMetadata, MV_INDEX_TYPE}
 import org.opensearch.flint.spark.scheduler.util.IntervalSchedulerParser
 
 /**
- * Flint metadata cache defines metadata required to store in read cache for frontend user to
- * access.
+ * Select fields that are exported as part of the mapping `_meta` object under `properties`.
+ * Useful for providing front-end features that need specific data without needing to go through
+ * the full async query system.
  */
-case class FlintMetadataCache(
+case class ExportedFlintMetadata(
     metadataCacheVersion: String,
     /** Refresh interval for Flint index with auto refresh. Unit: seconds */
     refreshInterval: Option[Int],
@@ -26,7 +28,9 @@ case class FlintMetadataCache(
     /** Source query for MV */
     sourceQuery: Option[String],
     /** Timestamp when Flint index is last refreshed. Unit: milliseconds */
-    lastRefreshTime: Option[Long]) {
+    lastRefreshTime: Option[Long],
+    /** Information on the current progress of the index job */
+    currentProgress: Option[Map[String, Double]]) {
 
   /**
    * Convert FlintMetadataCache to a map. Skips a field if its value is not defined.
@@ -47,11 +51,11 @@ case class FlintMetadataCache(
   }
 }
 
-object FlintMetadataCache {
+object ExportedFlintMetadata {
 
   val metadataCacheVersion = "1.0"
 
-  def apply(metadata: FlintMetadata): FlintMetadataCache = {
+  def apply(metadata: FlintMetadata): ExportedFlintMetadata = {
     val indexOptions = FlintSparkIndexOptions(
       metadata.options.asScala.mapValues(_.asInstanceOf[String]).toMap)
     val refreshInterval = if (indexOptions.autoRefresh()) {
@@ -76,11 +80,12 @@ object FlintMetadataCache {
         case timestamp => Some(timestamp)
       }
     }
-    FlintMetadataCache(
+    ExportedFlintMetadata(
       metadataCacheVersion,
       refreshInterval,
       sourceTables,
       sourceQuery,
-      lastRefreshTime)
+      lastRefreshTime,
+      None)
   }
 }

flint-spark-integration/src/main/scala/org/opensearch/flint/spark/metadatacache/FlintMetadataCacheWriter.scala

@@ -8,7 +8,7 @@ package org.opensearch.flint.spark.metadatacache
 import org.opensearch.flint.common.metadata.{FlintIndexMetadataService, FlintMetadata}
 
 /**
- * Writes {@link FlintMetadataCache} to a storage of choice. This is different from {@link
+ * Writes {@link ExportedFlintMetadata} to a storage of choice. This is different from {@link
  * FlintIndexMetadataService} which persists the full index metadata to a storage for single
  * source of truth.
  */

flint-spark-integration/src/main/scala/org/opensearch/flint/spark/metadatacache/FlintOpenSearchMetadataCacheWriter.scala

@@ -21,7 +21,8 @@ import org.opensearch.flint.core.storage.OpenSearchClientUtils
 import org.apache.spark.internal.Logging
 
 /**
- * Writes {@link FlintMetadataCache} to index mappings `_meta` field for frontend user to access.
+ * Writes {@link ExportedFlintMetadata} to index mappings `_meta` field for frontend user to
+ * access.
  */
 class FlintOpenSearchMetadataCacheWriter(options: FlintOptions)
     extends FlintMetadataCacheWriter
@@ -34,7 +35,7 @@ class FlintOpenSearchMetadataCacheWriter(options: FlintOptions)
     try {
       client = OpenSearchClientUtils.createClient(options)
       val indexMapping = getIndexMapping(client, osIndexName)
-      val metadataCacheProperties = FlintMetadataCache(metadata).toMap.asJava
+      val metadataCacheProperties = ExportedFlintMetadata(metadata).toMap.asJava
       mergeMetadataCacheProperties(indexMapping, metadataCacheProperties)
       val serialized = buildJson(builder => {
         builder.field("_meta", indexMapping.get("_meta"))

flint-spark-integration/src/main/scala/org/opensearch/flint/spark/refresh/AutoIndexRefresh.scala

@@ -6,10 +6,12 @@
 package org.opensearch.flint.spark.refresh
 
 import java.util.Collections
-import org.opensearch.flint.core.metrics.{MetricsSparkListener, ProgressListener, WithSparkListeners}
+
+import org.opensearch.flint.core.metrics.{MetricsSparkListener, Progress, ProgressListener, WithSparkListeners}
 import org.opensearch.flint.spark.{FlintSparkIndex, FlintSparkIndexOptions, FlintSparkValidationHelper}
-import org.opensearch.flint.spark.FlintSparkIndex.{StreamingRefresh, quotedTableName}
+import org.opensearch.flint.spark.FlintSparkIndex.{quotedTableName, StreamingRefresh}
 import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.RefreshMode.{AUTO, RefreshMode}
+
 import org.apache.spark.sql.{DataFrame, Row, SparkSession}
 import org.apache.spark.sql.flint.FlintDataSourceV2.FLINT_DATASOURCE
 import org.apache.spark.sql.flint.config.FlintSparkConf
@@ -29,6 +31,7 @@ class AutoIndexRefresh(indexName: String, index: FlintSparkIndex)
     with FlintSparkValidationHelper {
 
   override def refreshMode: RefreshMode = AUTO
+  val progressTracker: ProgressListener = ProgressListener()
 
   override def validate(spark: SparkSession): Unit = {
     // Incremental refresh cannot enabled at the same time
@@ -66,18 +69,16 @@ class AutoIndexRefresh(indexName: String, index: FlintSparkIndex)
       // Flint index has specialized logic and capability for incremental refresh
       case refresh: StreamingRefresh =>
         logInfo("Start refreshing index in streaming style")
-        val jobContext = WithSparkListeners(spark, List(MetricsSparkListener(), ProgressListener()))
-
-        val job = jobContext.run(() =>
-          refresh
-            .buildStream(spark)
-            .writeStream
-            .queryName(indexName)
-            .format(FLINT_DATASOURCE)
-            .options(flintSparkConf.properties)
-            .addSinkOptions(options, flintSparkConf)
-            .start(indexName))
-        val progress = jobContext.listeners(1).asInstanceOf[ProgressListener]
+        val job =
+          WithSparkListeners(spark, List(MetricsSparkListener(), progressTracker)).run(() =>
+            refresh
+              .buildStream(spark)
+              .writeStream
+              .queryName(indexName)
+              .format(FLINT_DATASOURCE)
+              .options(flintSparkConf.properties)
+              .addSinkOptions(options, flintSparkConf)
+              .start(indexName))
         Some(job.id.toString)
 
       // Otherwise, fall back to foreachBatch + batch refresh
@@ -99,6 +100,8 @@ class AutoIndexRefresh(indexName: String, index: FlintSparkIndex)
     }
   }
 
+  override def progress(): Option[Progress] = Some(progressTracker.currentProgress())
+
   // Using Scala implicit class to avoid breaking method chaining of Spark data frame fluent API
   private implicit class FlintDataStreamWriter(val dataStream: DataStreamWriter[Row]) {
 

flint-spark-integration/src/main/scala/org/opensearch/flint/spark/refresh/FlintSparkIndexRefresh.scala

@@ -5,12 +5,13 @@
 
 package org.opensearch.flint.spark.refresh
 
+import org.opensearch.flint.core.metrics.Progress
 import org.opensearch.flint.spark.FlintSparkIndex
 import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.RefreshMode.RefreshMode
+
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.flint.config.FlintSparkConf
-import org.opensearch.flint.core.metrics.Progress
 
 /**
  * Flint Spark index refresh that sync index data with source in style defined by concrete
@@ -53,9 +54,9 @@ trait FlintSparkIndexRefresh extends Logging {
   /**
    * For refresh types supporting progress recording, return the progress.
    *
-   * Either always or never set depending on the specific implementation. If the job hasn't started yet, the progress
-   * should be filled with zeroes, not None. This is necessary for the front-end to determine whether to render a
-   * progress bar at all.
+   * Either always or never set depending on the specific implementation. If the job hasn't
+   * started yet, the progress should be Some(zeroes), not None. This is necessary for the
+   * front-end to determine whether to render a progress bar at all.
    */
   def progress(): Option[Progress]
 }

flint-spark-integration/src/main/scala/org/opensearch/flint/spark/refresh/FullIndexRefresh.scala

@@ -5,6 +5,7 @@
 
 package org.opensearch.flint.spark.refresh
 
+import org.opensearch.flint.core.metrics.Progress
 import org.opensearch.flint.spark.FlintSparkIndex
 import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.RefreshMode.{FULL, RefreshMode}
 
@@ -47,4 +48,6 @@ class FullIndexRefresh(
       .save(indexName)
     None
   }
+
+  override def progress(): Option[Progress] = None
 }

flint-spark-integration/src/main/scala/org/opensearch/flint/spark/refresh/IncrementalIndexRefresh.scala

@@ -5,6 +5,7 @@
 
 package org.opensearch.flint.spark.refresh
 
+import org.opensearch.flint.core.metrics.Progress
 import org.opensearch.flint.spark.{FlintSparkIndex, FlintSparkValidationHelper}
 import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.RefreshMode.{INCREMENTAL, RefreshMode}
 import org.opensearch.flint.spark.refresh.util.RefreshMetricsAspect
@@ -62,4 +63,6 @@ class IncrementalIndexRefresh(val indexName: String, index: FlintSparkIndex)
       None
     }
   }
+
+  override def progress(): Option[Progress] = None
 }

flint-spark-integration/src/test/scala/org/opensearch/flint/spark/metadatacache/FlintMetadataCacheSuite.scala

@@ -46,8 +46,8 @@ class FlintMetadataCacheSuite extends AnyFlatSpec with Matchers {
       .deserialize(content)
       .copy(latestLogEntry = Some(flintMetadataLogEntry))
 
-    val metadataCache = FlintMetadataCache(metadata)
-    metadataCache.metadataCacheVersion shouldBe FlintMetadataCache.metadataCacheVersion
+    val metadataCache = ExportedFlintMetadata(metadata)
+    metadataCache.metadataCacheVersion shouldBe ExportedFlintMetadata.metadataCacheVersion
     metadataCache.refreshInterval.get shouldBe 600
     metadataCache.sourceTables shouldBe Array("spark_catalog.default.test_table")
     metadataCache.lastRefreshTime.get shouldBe 1234567890123L
@@ -75,8 +75,8 @@ class FlintMetadataCacheSuite extends AnyFlatSpec with Matchers {
       .deserialize(content)
       .copy(latestLogEntry = Some(flintMetadataLogEntry))
 
-    val metadataCache = FlintMetadataCache(metadata)
-    metadataCache.metadataCacheVersion shouldBe FlintMetadataCache.metadataCacheVersion
+    val metadataCache = ExportedFlintMetadata(metadata)
+    metadataCache.metadataCacheVersion shouldBe ExportedFlintMetadata.metadataCacheVersion
     metadataCache.refreshInterval.get shouldBe 600
     metadataCache.sourceTables shouldBe Array("spark_catalog.default.test_table")
     metadataCache.lastRefreshTime.get shouldBe 1234567890123L
@@ -112,8 +112,8 @@ class FlintMetadataCacheSuite extends AnyFlatSpec with Matchers {
       .deserialize(content)
       .copy(latestLogEntry = Some(flintMetadataLogEntry))
 
-    val metadataCache = FlintMetadataCache(metadata)
-    metadataCache.metadataCacheVersion shouldBe FlintMetadataCache.metadataCacheVersion
+    val metadataCache = ExportedFlintMetadata(metadata)
+    metadataCache.metadataCacheVersion shouldBe ExportedFlintMetadata.metadataCacheVersion
     metadataCache.refreshInterval.get shouldBe 600
     metadataCache.sourceTables shouldBe Array(
       "spark_catalog.default.test_table",
@@ -144,8 +144,8 @@ class FlintMetadataCacheSuite extends AnyFlatSpec with Matchers {
       .deserialize(content)
       .copy(latestLogEntry = Some(flintMetadataLogEntry.copy(lastRefreshCompleteTime = 0L)))
 
-    val metadataCache = FlintMetadataCache(metadata)
-    metadataCache.metadataCacheVersion shouldBe FlintMetadataCache.metadataCacheVersion
+    val metadataCache = ExportedFlintMetadata(metadata)
+    metadataCache.metadataCacheVersion shouldBe ExportedFlintMetadata.metadataCacheVersion
     metadataCache.refreshInterval shouldBe empty
     metadataCache.sourceTables shouldBe Array("spark_catalog.default.test_table")
     metadataCache.sourceQuery shouldBe empty

spark-sql-application/src/main/scala/org/apache/spark/sql/JobOperator.scala

@@ -7,13 +7,16 @@ package org.apache.spark.sql
 
 import java.util.concurrent.{ThreadPoolExecutor, TimeUnit}
 import java.util.concurrent.atomic.AtomicInteger
+
 import scala.concurrent.{ExecutionContext, Future, TimeoutException}
 import scala.concurrent.duration.{Duration, MINUTES}
 import scala.util.{Failure, Success, Try}
+
 import org.opensearch.flint.common.model.FlintStatement
 import org.opensearch.flint.common.scheduler.model.LangType
 import org.opensearch.flint.core.metrics.{MetricConstants, MetricsSparkListener, MetricsUtil, ProgressListener}
 import org.opensearch.flint.spark.FlintSpark
+
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.flint.config.FlintSparkConf
 import org.apache.spark.sql.util.ShuffleCleaner