Add scheduler_mode index option (opensearch-project#415)

* Add scheduler_mode index option Signed-off-by: Louis Chu <[email protected]> * Add integ test Signed-off-by: Louis Chu <[email protected]> * Delete flint-spark-integration/src/main/scala/org/opensearch/flint/spark/skipping/worksheet.sc Signed-off-by: Louis Chu <[email protected]> * Fix CI Signed-off-by: Louis Chu <[email protected]> * Resolve comment Signed-off-by: Louis Chu <[email protected]> * Fix CI Signed-off-by: Louis Chu <[email protected]> --------- Signed-off-by: Louis Chu <[email protected]>
seankao-az · Jul 15, 2024 · 890930c · 890930c
1 parent 978c028
commit 890930c
Show file tree

Hide file tree

Showing 17 changed files with 291 additions and 21 deletions.
diff --git a/docs/index.md b/docs/index.md
@@ -30,7 +30,13 @@ Please see the following example in which Index Building Logic and Query Rewrite
 ### Flint Index Refresh
 
 - **Auto Refresh:**
-  - This feature allows the Flint Index to automatically refresh. Users can configure such as frequency of auto-refresh based on their preferences.
+  - This feature allows the Flint Index to automatically refresh. Users can configure such as frequency of auto-refresh based on their preferences. There are two modes available for scheduling the auto-refresh:
+  - **Internal Scheduler:**
+    - Description: The data refresh is executed in micro-batch mode using the internal scheduler. 
+    - Recommended Use-Case: This mode is ideal for low-latency use-cases where data needs to be refreshed frequently and quickly.    
+  - **External Scheduler:**
+    - Description: The data refresh is executed using an external scheduler. 
+    - Recommended Use-Case: This mode is suitable for scenarios where data responsiveness is less critical, helping to reduce the cost of maintaining a long-running Spark cluster.
 - **Manual Refresh:**
   - Users have the option to manually trigger a refresh for the Flint Index. This provides flexibility and control over when the refresh occurs.
     - **Full Refresh:**
@@ -369,7 +375,8 @@ fetched rows / total rows = 5/5
 
 User can provide the following options in `WITH` clause of create statement:
 
-+ `auto_refresh`: default value is false. Automatically refresh the index if set to true. Otherwise, user has to trigger refresh by `REFRESH` statement manually.
++ `auto_refresh`: default value is false. Automatically refresh the index if set to true. Otherwise, user has to trigger refresh by `REFRESH` statement manually. 
++ `scheduler_mode`: A mode string (`internal` or `external`) that describes how `auto_refresh` is scheduled. `checkpoint_location` is required for the external scheduler.
 + `refresh_interval`: a string as the time interval for incremental refresh, e.g. 1 minute, 10 seconds. This is only applicable when auto refresh enabled. Please check `org.apache.spark.unsafe.types.CalendarInterval` for valid duration identifiers. By default, next micro batch will be generated as soon as the previous one complete processing.
 + `incremental_refresh`: default value is false. incrementally refresh the index if set to true. Otherwise, fully refresh the entire index. This only applicable when auto refresh disabled.
 + `checkpoint_location`: a string as the location path for refresh job checkpoint (auto or incremental). The location has to be a path in an HDFS compatible file system and only applicable when auto refresh enabled. If unspecified, temporary checkpoint directory will be used and may result in checkpoint data lost upon restart.

diff --git a/...-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexOptions.scala b/...-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexOptions.scala
@@ -8,8 +8,9 @@ package org.opensearch.flint.spark
 import org.json4s.{Formats, NoTypeHints}
 import org.json4s.native.JsonMethods._
 import org.json4s.native.Serialization
-import org.opensearch.flint.spark.FlintSparkIndexOptions.OptionName.{AUTO_REFRESH, CHECKPOINT_LOCATION, EXTRA_OPTIONS, INCREMENTAL_REFRESH, INDEX_SETTINGS, OptionName, OUTPUT_MODE, REFRESH_INTERVAL, WATERMARK_DELAY}
+import org.opensearch.flint.spark.FlintSparkIndexOptions.OptionName.{AUTO_REFRESH, CHECKPOINT_LOCATION, EXTRA_OPTIONS, INCREMENTAL_REFRESH, INDEX_SETTINGS, OptionName, OUTPUT_MODE, REFRESH_INTERVAL, SCHEDULER_MODE, WATERMARK_DELAY}
 import org.opensearch.flint.spark.FlintSparkIndexOptions.validateOptionNames
+import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.SchedulerMode
 
 /**
  * Flint Spark index configurable options.
@@ -22,6 +23,7 @@ case class FlintSparkIndexOptions(options: Map[String, String]) {
   implicit val formats: Formats = Serialization.formats(NoTypeHints)
 
   validateOptionNames(options)
+  validateOptionSchedulerModeValue()
 
   /**
    * Is Flint index auto refreshed or manual refreshed.
@@ -31,6 +33,19 @@ case class FlintSparkIndexOptions(options: Map[String, String]) {
    */
   def autoRefresh(): Boolean = getOptionValue(AUTO_REFRESH).getOrElse("false").toBoolean
 
+  /**
+   * The scheduler mode for the Flint index refresh.
+   *
+   * @return
+   *   scheduler mode option value
+   */
+  def schedulerMode(): SchedulerMode.Value = {
+    // TODO: Change default value to external once the external scheduler is enabled
+    val defaultMode = "internal"
+    val modeStr = getOptionValue(SCHEDULER_MODE).getOrElse(defaultMode)
+    SchedulerMode.fromString(modeStr)
+  }
+
   /**
    * The refresh interval (only valid if auto refresh enabled).
    *
@@ -112,6 +127,21 @@ case class FlintSparkIndexOptions(options: Map[String, String]) {
     if (!options.contains(AUTO_REFRESH.toString)) {
       map += (AUTO_REFRESH.toString -> autoRefresh().toString)
     }
+
+    // Add default option only when auto refresh is TRUE
+    if (autoRefresh() == true) {
+      if (!options.contains(SCHEDULER_MODE.toString)) {
+        map += (SCHEDULER_MODE.toString -> schedulerMode().toString)
+      }
+
+      // The query will be executed in micro-batch mode using the internal scheduler
+      // The default interval for the external scheduler is 15 minutes.
+      if (SchedulerMode.EXTERNAL == schedulerMode() && !options.contains(
+          REFRESH_INTERVAL.toString)) {
+        map += (REFRESH_INTERVAL.toString -> "15 minutes")
+      }
+    }
+
     if (!options.contains(INCREMENTAL_REFRESH.toString)) {
       map += (INCREMENTAL_REFRESH.toString -> incrementalRefresh().toString)
     }
@@ -127,6 +157,14 @@ case class FlintSparkIndexOptions(options: Map[String, String]) {
       .map(opt => (parse(opt) \ key).extract[Map[String, String]])
       .getOrElse(Map.empty)
   }
+
+  private def validateOptionSchedulerModeValue(): Unit = {
+    getOptionValue(SCHEDULER_MODE) match {
+      case Some(modeStr) =>
+        SchedulerMode.fromString(modeStr) // Will throw an exception if the mode is invalid
+      case None => // no action needed if modeStr is empty
+    }
+  }
 }
 
 object FlintSparkIndexOptions {
@@ -142,6 +180,7 @@ object FlintSparkIndexOptions {
   object OptionName extends Enumeration {
     type OptionName = Value
     val AUTO_REFRESH: OptionName.Value = Value("auto_refresh")
+    val SCHEDULER_MODE: OptionName.Value = Value("scheduler_mode")
     val REFRESH_INTERVAL: OptionName.Value = Value("refresh_interval")
     val INCREMENTAL_REFRESH: OptionName.Value = Value("incremental_refresh")
     val CHECKPOINT_LOCATION: OptionName.Value = Value("checkpoint_location")

diff --git a/...park-integration/src/main/scala/org/opensearch/flint/spark/refresh/AutoIndexRefresh.scala b/...park-integration/src/main/scala/org/opensearch/flint/spark/refresh/AutoIndexRefresh.scala
@@ -10,6 +10,7 @@ import java.util.Collections
 import org.opensearch.flint.spark.{FlintSparkIndex, FlintSparkIndexOptions, FlintSparkValidationHelper}
 import org.opensearch.flint.spark.FlintSparkIndex.{quotedTableName, StreamingRefresh}
 import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.RefreshMode.{AUTO, RefreshMode}
+import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.SchedulerMode
 
 import org.apache.spark.sql.{DataFrame, Row, SparkSession}
 import org.apache.spark.sql.flint.FlintDataSourceV2.FLINT_DATASOURCE
@@ -43,7 +44,7 @@ class AutoIndexRefresh(indexName: String, index: FlintSparkIndex)
       !isTableProviderSupported(spark, index),
       "Index auto refresh doesn't support Hive table")
 
-    // Checkpoint location is required if mandatory option set
+    // Checkpoint location is required if mandatory option set or external scheduler is used
     val flintSparkConf = new FlintSparkConf(Collections.emptyMap[String, String])
     val checkpointLocation = options.checkpointLocation()
     if (flintSparkConf.isCheckpointMandatory) {
@@ -90,7 +91,7 @@ class AutoIndexRefresh(indexName: String, index: FlintSparkIndex)
           .foreachBatch { (batchDF: DataFrame, _: Long) =>
             new FullIndexRefresh(indexName, index, Some(batchDF))
               .start(spark, flintSparkConf)
-            () // discard return value above and return unit to use right overridden method
+            () // discard return value above and return unit to use the right overridden method
           }
           .start()
         Some(job.id.toString)
@@ -103,10 +104,12 @@ class AutoIndexRefresh(indexName: String, index: FlintSparkIndex)
     def addSinkOptions(
         options: FlintSparkIndexOptions,
         flintSparkConf: FlintSparkConf): DataStreamWriter[Row] = {
+      // For incremental refresh, the refresh_interval option is overridden by Trigger.AvailableNow().
       dataStream
         .addCheckpointLocation(options.checkpointLocation(), flintSparkConf.isCheckpointMandatory)
         .addRefreshInterval(options.refreshInterval())
-        .addAvailableNowTrigger(options.incrementalRefresh())
+        .addAvailableNowTrigger(
+          SchedulerMode.EXTERNAL == options.schedulerMode() || options.incrementalRefresh())
         .addOutputMode(options.outputMode())
         .options(options.extraSinkOptions())
     }
@@ -129,8 +132,8 @@ class AutoIndexRefresh(indexName: String, index: FlintSparkIndex)
         .getOrElse(dataStream)
     }
 
-    def addAvailableNowTrigger(incrementalRefresh: Boolean): DataStreamWriter[Row] = {
-      if (incrementalRefresh) {
+    def addAvailableNowTrigger(setAvailableNow: Boolean): DataStreamWriter[Row] = {
+      if (setAvailableNow) {
         dataStream.trigger(Trigger.AvailableNow())
       } else {
         dataStream

diff --git a/...ntegration/src/main/scala/org/opensearch/flint/spark/refresh/FlintSparkIndexRefresh.scala b/...ntegration/src/main/scala/org/opensearch/flint/spark/refresh/FlintSparkIndexRefresh.scala
@@ -59,6 +59,20 @@ object FlintSparkIndexRefresh {
     val AUTO, FULL, INCREMENTAL = Value
   }
 
+  /** Index scheduler mode for auto refresh */
+  object SchedulerMode extends Enumeration {
+    type SchedulerMode = Value
+    val INTERNAL: SchedulerMode.Value = Value("internal")
+    val EXTERNAL: SchedulerMode.Value = Value("external")
+
+    def fromString(s: String): SchedulerMode.Value = {
+      require(
+        SchedulerMode.values.exists(_.toString.equalsIgnoreCase(s)),
+        s"Scheduler mode $s is invalid. Must be 'internal' or 'external'.")
+      SchedulerMode.values.find(_.toString.equalsIgnoreCase(s)).get
+    }
+  }
+
   /**
    * Create concrete index refresh implementation for the given index.
    *
@@ -71,7 +85,10 @@ object FlintSparkIndexRefresh {
    */
   def create(indexName: String, index: FlintSparkIndex): FlintSparkIndexRefresh = {
     val options = index.options
-    if (options.autoRefresh()) {
+    // TODO: Refactor the refresh class names and RefreshMode enum to distinguish Spark side physical trigger setup from the logical Flint index option.
+    if (SchedulerMode.EXTERNAL == options.schedulerMode()) {
+      new IncrementalIndexRefresh(indexName, index)
+    } else if (options.autoRefresh()) {
       new AutoIndexRefresh(indexName, index)
     } else if (options.incrementalRefresh()) {
       new IncrementalIndexRefresh(indexName, index)

diff --git a/...tegration/src/main/scala/org/opensearch/flint/spark/refresh/IncrementalIndexRefresh.scala b/...tegration/src/main/scala/org/opensearch/flint/spark/refresh/IncrementalIndexRefresh.scala
@@ -34,9 +34,7 @@ class IncrementalIndexRefresh(indexName: String, index: FlintSparkIndex)
     // Checkpoint location is required regardless of mandatory option
     val options = index.options
     val checkpointLocation = options.checkpointLocation()
-    require(
-      options.checkpointLocation().nonEmpty,
-      "Checkpoint location is required by incremental refresh")
+    require(options.checkpointLocation().nonEmpty, "Checkpoint location is required")
     require(
       isCheckpointLocationAccessible(spark, checkpointLocation.get),
       s"No sufficient permission to access the checkpoint location ${checkpointLocation.get}")
@@ -50,6 +48,7 @@ class IncrementalIndexRefresh(indexName: String, index: FlintSparkIndex)
       new AutoIndexRefresh(indexName, index)
         .start(spark, flintSparkConf)
 
+    // Blocks the calling thread until the streaming query finishes
     spark.streams
       .get(jobId.get)
       .awaitTermination()

diff --git a/...ain/scala/org/opensearch/flint/spark/sql/covering/FlintSparkCoveringIndexAstBuilder.scala b/...ain/scala/org/opensearch/flint/spark/sql/covering/FlintSparkCoveringIndexAstBuilder.scala
@@ -8,6 +8,7 @@ package org.opensearch.flint.spark.sql.covering
 import org.antlr.v4.runtime.tree.RuleNode
 import org.opensearch.flint.spark.FlintSpark
 import org.opensearch.flint.spark.covering.FlintSparkCoveringIndex
+import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.SchedulerMode
 import org.opensearch.flint.spark.sql.{FlintSparkSqlCommand, FlintSparkSqlExtensionsVisitor, SparkSqlAstBuilder}
 import org.opensearch.flint.spark.sql.FlintSparkSqlAstBuilder.{getFullTableName, getSqlText}
 import org.opensearch.flint.spark.sql.FlintSparkSqlExtensionsParser._
@@ -49,8 +50,9 @@ trait FlintSparkCoveringIndexAstBuilder extends FlintSparkSqlExtensionsVisitor[A
         .options(indexOptions)
         .create(ignoreIfExists)
 
-      // Trigger auto refresh if enabled
-      if (indexOptions.autoRefresh()) {
+      // Trigger auto refresh if enabled and not using external scheduler
+      if (indexOptions
+          .autoRefresh() && SchedulerMode.INTERNAL == indexOptions.schedulerMode()) {
         val flintIndexName = getFlintIndexName(flint, ctx.indexName, ctx.tableName)
         flint.refreshIndex(flintIndexName)
       }

diff --git a/...c/main/scala/org/opensearch/flint/spark/sql/mv/FlintSparkMaterializedViewAstBuilder.scala b/...c/main/scala/org/opensearch/flint/spark/sql/mv/FlintSparkMaterializedViewAstBuilder.scala
@@ -10,6 +10,7 @@ import scala.collection.convert.ImplicitConversions.`collection AsScalaIterable`
 import org.antlr.v4.runtime.tree.RuleNode
 import org.opensearch.flint.spark.FlintSpark
 import org.opensearch.flint.spark.mv.FlintSparkMaterializedView
+import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.SchedulerMode
 import org.opensearch.flint.spark.sql.{FlintSparkSqlCommand, FlintSparkSqlExtensionsVisitor, SparkSqlAstBuilder}
 import org.opensearch.flint.spark.sql.FlintSparkSqlAstBuilder.{getFullTableName, getSqlText, IndexBelongsTo}
 import org.opensearch.flint.spark.sql.FlintSparkSqlExtensionsParser._
@@ -42,8 +43,9 @@ trait FlintSparkMaterializedViewAstBuilder extends FlintSparkSqlExtensionsVisito
         .options(indexOptions)
         .create(ignoreIfExists)
 
-      // Trigger auto refresh if enabled
-      if (indexOptions.autoRefresh()) {
+      // Trigger auto refresh if enabled and not using external scheduler
+      if (indexOptions
+          .autoRefresh() && SchedulerMode.INTERNAL == indexOptions.schedulerMode()) {
         val flintIndexName = getFlintIndexName(flint, ctx.mvName)
         flint.refreshIndex(flintIndexName)
       }

diff --git a/...ain/scala/org/opensearch/flint/spark/sql/skipping/FlintSparkSkippingIndexAstBuilder.scala b/...ain/scala/org/opensearch/flint/spark/sql/skipping/FlintSparkSkippingIndexAstBuilder.scala
@@ -10,6 +10,7 @@ import scala.collection.JavaConverters.collectionAsScalaIterableConverter
 import org.antlr.v4.runtime.tree.RuleNode
 import org.opensearch.flint.core.field.bloomfilter.BloomFilterFactory._
 import org.opensearch.flint.spark.FlintSpark
+import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.SchedulerMode
 import org.opensearch.flint.spark.skipping.FlintSparkSkippingIndex
 import org.opensearch.flint.spark.skipping.FlintSparkSkippingStrategy.SkippingKind
 import org.opensearch.flint.spark.skipping.FlintSparkSkippingStrategy.SkippingKind.{BLOOM_FILTER, MIN_MAX, PARTITION, VALUE_SET}
@@ -75,8 +76,9 @@ trait FlintSparkSkippingIndexAstBuilder extends FlintSparkSqlExtensionsVisitor[A
         .options(indexOptions)
         .create(ignoreIfExists)
 
-      // Trigger auto refresh if enabled
-      if (indexOptions.autoRefresh()) {
+      // Trigger auto refresh if enabled and not using external scheduler
+      if (indexOptions
+          .autoRefresh() && SchedulerMode.INTERNAL == indexOptions.schedulerMode()) {
         val indexName = getSkippingIndexName(flint, ctx.tableName)
         flint.refreshIndex(indexName)
       }

diff --git a/...k-integration/src/test/scala/org/opensearch/flint/spark/FlintSparkIndexOptionsSuite.scala b/...k-integration/src/test/scala/org/opensearch/flint/spark/FlintSparkIndexOptionsSuite.scala
@@ -6,6 +6,7 @@
 package org.opensearch.flint.spark
 
 import org.opensearch.flint.spark.FlintSparkIndexOptions.OptionName._
+import org.opensearch.flint.spark.refresh.FlintSparkIndexRefresh.SchedulerMode
 import org.scalatest.matchers.should.Matchers
 
 import org.apache.spark.FlintSuite
@@ -14,6 +15,7 @@ class FlintSparkIndexOptionsSuite extends FlintSuite with Matchers {
 
   test("should return lowercase name as option name") {
     AUTO_REFRESH.toString shouldBe "auto_refresh"
+    SCHEDULER_MODE.toString shouldBe "scheduler_mode"
     REFRESH_INTERVAL.toString shouldBe "refresh_interval"
     INCREMENTAL_REFRESH.toString shouldBe "incremental_refresh"
     CHECKPOINT_LOCATION.toString shouldBe "checkpoint_location"
@@ -27,6 +29,7 @@ class FlintSparkIndexOptionsSuite extends FlintSuite with Matchers {
     val options = FlintSparkIndexOptions(
       Map(
         "auto_refresh" -> "true",
+        "scheduler_mode" -> "external",
         "refresh_interval" -> "1 Minute",
         "incremental_refresh" -> "true",
         "checkpoint_location" -> "s3://test/",
@@ -45,6 +48,7 @@ class FlintSparkIndexOptionsSuite extends FlintSuite with Matchers {
             | }""".stripMargin))
 
     options.autoRefresh() shouldBe true
+    options.schedulerMode() shouldBe SchedulerMode.EXTERNAL
     options.refreshInterval() shouldBe Some("1 Minute")
     options.incrementalRefresh() shouldBe true
     options.checkpointLocation() shouldBe Some("s3://test/")
@@ -73,6 +77,7 @@ class FlintSparkIndexOptionsSuite extends FlintSuite with Matchers {
     val options = FlintSparkIndexOptions(Map.empty)
 
     options.autoRefresh() shouldBe false
+    options.schedulerMode() shouldBe SchedulerMode.INTERNAL
     options.refreshInterval() shouldBe empty
     options.checkpointLocation() shouldBe empty
     options.watermarkDelay() shouldBe empty
@@ -92,6 +97,27 @@ class FlintSparkIndexOptionsSuite extends FlintSuite with Matchers {
       "refresh_interval" -> "1 Minute")
   }
 
+  test("should return include default scheduler_mode option when auto refresh is set to true") {
+    val options = FlintSparkIndexOptions(Map("auto_refresh" -> "true"))
+
+    options.optionsWithDefault shouldBe Map(
+      "auto_refresh" -> "true",
+      "scheduler_mode" -> "internal",
+      "incremental_refresh" -> "false")
+  }
+
+  test(
+    "should return include default refresh_interval option with auto_refresh=true and scheduler_mode=external") {
+    val options =
+      FlintSparkIndexOptions(Map("auto_refresh" -> "true", "scheduler_mode" -> "external"))
+
+    options.optionsWithDefault shouldBe Map(
+      "auto_refresh" -> "true",
+      "scheduler_mode" -> "external",
+      "refresh_interval" -> "15 minutes",
+      "incremental_refresh" -> "false")
+  }
+
   test("should report error if any unknown option name") {
     the[IllegalArgumentException] thrownBy
       FlintSparkIndexOptions(Map("autoRefresh" -> "true"))
@@ -102,5 +128,9 @@ class FlintSparkIndexOptionsSuite extends FlintSuite with Matchers {
     the[IllegalArgumentException] thrownBy {
       FlintSparkIndexOptions(Map("auto_refresh" -> "true", "indexSetting" -> "test"))
     } should have message "requirement failed: option name indexSetting is invalid"
+
+    the[IllegalArgumentException] thrownBy {
+      FlintSparkIndexOptions(Map("scheduler_mode" -> "invalid_mode"))
+    } should have message "requirement failed: Scheduler mode invalid_mode is invalid. Must be 'internal' or 'external'."
   }
 }
diff --git a/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkCoveringIndexITSuite.scala b/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkCoveringIndexITSuite.scala
@@ -118,6 +118,29 @@ class FlintSparkCoveringIndexITSuite extends FlintSparkSuite {
     checkAnswer(indexData, Seq(Row("Hello", 30), Row("World", 25)))
   }
 
+  test("auto refresh covering index successfully with external scheduler") {
+    withTempDir { checkpointDir =>
+      flint
+        .coveringIndex()
+        .name(testIndex)
+        .onTable(testTable)
+        .addIndexColumns("name", "age")
+        .options(
+          FlintSparkIndexOptions(
+            Map(
+              "auto_refresh" -> "true",
+              "scheduler_mode" -> "external",
+              "checkpoint_location" -> checkpointDir.getAbsolutePath)))
+        .create()
+
+      val jobId = flint.refreshIndex(testFlintIndex)
+      jobId shouldBe None
+
+      val indexData = flint.queryIndex(testFlintIndex)
+      checkAnswer(indexData, Seq(Row("Hello", 30), Row("World", 25)))
+    }
+  }
+
   test("update covering index successfully") {
     // Create full refresh Flint index
     flint