Add refresh and drop statement support

Signed-off-by: Chen Dai <[email protected]>
dai-chen · Sep 6, 2023 · 88ca8dd · 88ca8dd
1 parent c967577
commit 88ca8dd
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 1 deletion.
diff --git a/flint-spark-integration/src/main/antlr4/FlintSparkSqlExtensions.g4 b/flint-spark-integration/src/main/antlr4/FlintSparkSqlExtensions.g4
@@ -46,6 +46,8 @@ dropSkippingIndexStatement
 
 coveringIndexStatement
     : createCoveringIndexStatement
+    | refreshCoveringIndexStatement
+    | dropCoveringIndexStatement
     ;
 
 createCoveringIndexStatement
@@ -54,6 +56,14 @@ createCoveringIndexStatement
         (WITH LEFT_PAREN propertyList RIGHT_PAREN)?
     ;
 
+refreshCoveringIndexStatement
+    : REFRESH INDEX indexName=identifier ON tableName=multipartIdentifier
+    ;
+
+dropCoveringIndexStatement
+    : DROP INDEX indexName=identifier ON tableName=multipartIdentifier
+    ;
+
 indexColTypeList
     : indexColType (COMMA indexColType)*
     ;

diff --git a/...ain/scala/org/opensearch/flint/spark/sql/covering/FlintSparkCoveringIndexAstBuilder.scala b/...ain/scala/org/opensearch/flint/spark/sql/covering/FlintSparkCoveringIndexAstBuilder.scala
@@ -9,7 +9,7 @@ import org.opensearch.flint.spark.FlintSpark.RefreshMode
 import org.opensearch.flint.spark.covering.FlintSparkCoveringIndex.getFlintIndexName
 import org.opensearch.flint.spark.sql.{FlintSparkSqlCommand, FlintSparkSqlExtensionsVisitor}
 import org.opensearch.flint.spark.sql.FlintSparkSqlAstBuilder.{getFullTableName, isAutoRefreshEnabled}
-import org.opensearch.flint.spark.sql.FlintSparkSqlExtensionsParser.CreateCoveringIndexStatementContext
+import org.opensearch.flint.spark.sql.FlintSparkSqlExtensionsParser.{CreateCoveringIndexStatementContext, DropCoveringIndexStatementContext, RefreshCoveringIndexStatementContext}
 
 import org.apache.spark.sql.catalyst.plans.logical.Command
 
@@ -43,4 +43,28 @@ trait FlintSparkCoveringIndexAstBuilder extends FlintSparkSqlExtensionsVisitor[C
       Seq.empty
     }
   }
+
+  override def visitRefreshCoveringIndexStatement(
+      ctx: RefreshCoveringIndexStatementContext): Command = {
+    FlintSparkSqlCommand() { flint =>
+      val indexName = ctx.indexName.getText
+      val tableName = ctx.tableName.getText
+      val flintIndexName = getFlintIndexName(indexName, tableName)
+
+      flint.refreshIndex(flintIndexName, RefreshMode.FULL)
+      Seq.empty
+    }
+  }
+
+  override def visitDropCoveringIndexStatement(
+      ctx: DropCoveringIndexStatementContext): Command = {
+    FlintSparkSqlCommand() { flint =>
+      val indexName = ctx.indexName.getText
+      val tableName = getFullTableName(flint, ctx.tableName)
+      val flintIndexName = getFlintIndexName(indexName, tableName)
+
+      flint.deleteIndex(flintIndexName)
+      Seq.empty
+    }
+  }
 }
diff --git a/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkCoveringIndexSqlITSuite.scala b/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkCoveringIndexSqlITSuite.scala
@@ -5,6 +5,8 @@
 
 package org.opensearch.flint.spark
 
+import scala.Option.empty
+
 import org.opensearch.flint.spark.covering.FlintSparkCoveringIndex.getFlintIndexName
 import org.scalatest.matchers.must.Matchers.defined
 import org.scalatest.matchers.should.Matchers.convertToAnyShouldWrapper
@@ -46,4 +48,32 @@ class FlintSparkCoveringIndexSqlITSuite extends FlintSparkSuite {
     val indexData = flint.queryIndex(testFlintIndex)
     indexData.count() shouldBe 2
   }
+
+  test("create covering index with manual refresh") {
+    sql(s"""
+           | CREATE INDEX $testIndex ON $testTable
+           | (name, age)
+           |""".stripMargin)
+
+    val indexData = flint.queryIndex(testFlintIndex)
+
+    flint.describeIndex(testFlintIndex) shouldBe defined
+    indexData.count() shouldBe 0
+
+    sql(s"REFRESH INDEX $testIndex ON $testTable")
+    indexData.count() shouldBe 2
+  }
+
+  test("drop covering index") {
+    flint
+      .coveringIndex()
+      .name(testIndex)
+      .onTable(testTable)
+      .addIndexColumns("name", "age")
+      .create()
+
+    sql(s"DROP INDEX $testIndex ON $testTable")
+
+    flint.describeIndex(testFlintIndex) shouldBe empty
+  }
 }