Add IF NOT EXISTS for skipping and covering index with IT

Signed-off-by: Chen Dai <[email protected]>
opensearch-project · Sep 25, 2023 · 0884216 · 0884216
1 parent e72f054
commit 0884216
Show file tree

Hide file tree

Showing 9 changed files with 71 additions and 15 deletions.
diff --git a/docs/index.md b/docs/index.md
@@ -121,7 +121,7 @@ High level API is dependent on query engine implementation. Please see Query Eng
 #### Skipping Index
 
 ```sql
-CREATE SKIPPING INDEX
+CREATE SKIPPING INDEX [IF NOT EXISTS]
 ON <object>
 ( column <index_type> [, ...] )
 WHERE <filter_predicate>
@@ -161,7 +161,7 @@ DROP SKIPPING INDEX ON alb_logs
 #### Covering Index
 
 ```sql
-CREATE INDEX name ON <object>
+CREATE INDEX [IF NOT EXISTS] name ON <object>
 ( column [, ...] )
 WHERE <filter_predicate>
 WITH (auto_refresh = (true|false))

diff --git a/flint-spark-integration/src/main/antlr4/FlintSparkSqlExtensions.g4 b/flint-spark-integration/src/main/antlr4/FlintSparkSqlExtensions.g4
@@ -27,7 +27,8 @@ skippingIndexStatement
     ;
 
 createSkippingIndexStatement
-    : CREATE SKIPPING INDEX ON tableName=multipartIdentifier
+    : CREATE SKIPPING INDEX (IF NOT EXISTS)?
+        ON tableName=multipartIdentifier
         LEFT_PAREN indexColTypeList RIGHT_PAREN
         (WITH LEFT_PAREN propertyList RIGHT_PAREN)?
     ;
@@ -51,7 +52,8 @@ coveringIndexStatement
     ;
 
 createCoveringIndexStatement
-    : CREATE INDEX indexName=identifier ON tableName=multipartIdentifier
+    : CREATE INDEX (IF NOT EXISTS)? indexName=identifier
+        ON tableName=multipartIdentifier
         LEFT_PAREN indexColumns=multipartIdentifierPropertyList RIGHT_PAREN
         (WITH LEFT_PAREN propertyList RIGHT_PAREN)?
     ;

diff --git a/flint-spark-integration/src/main/antlr4/SparkSqlBase.g4 b/flint-spark-integration/src/main/antlr4/SparkSqlBase.g4
@@ -158,8 +158,11 @@ CREATE: 'CREATE';
 DESC: 'DESC';
 DESCRIBE: 'DESCRIBE';
 DROP: 'DROP';
+EXISTS: 'EXISTS';
 FALSE: 'FALSE';
+IF: 'IF';
 INDEX: 'INDEX';
+NOT: 'NOT';
 ON: 'ON';
 PARTITION: 'PARTITION';
 REFRESH: 'REFRESH';

diff --git a/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSpark.scala b/flint-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSpark.scala
@@ -74,14 +74,18 @@ class FlintSpark(val spark: SparkSession) {
    *
    * @param index
    *   Flint index to create
+   * @param ignoreIfExists
+   *   Ignore existing index
    */
-  def createIndex(index: FlintSparkIndex): Unit = {
+  def createIndex(index: FlintSparkIndex, ignoreIfExists: Boolean = false): Unit = {
     val indexName = index.name()
     if (flintClient.exists(indexName)) {
-      throw new IllegalStateException(
-        s"A table can only have one Flint skipping index: Flint index $indexName is found")
+      if (!ignoreIfExists) {
+        throw new IllegalStateException(s"Flint index $indexName already exists")
+      }
+    } else {
+      flintClient.createIndex(indexName, index.metadata())
     }
-    flintClient.createIndex(indexName, index.metadata())
   }
 
   /**

diff --git a/...-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexBuilder.scala b/...-spark-integration/src/main/scala/org/opensearch/flint/spark/FlintSparkIndexBuilder.scala
@@ -31,8 +31,12 @@ abstract class FlintSparkIndexBuilder(flint: FlintSpark) {
 
   /**
    * Create Flint index.
+   *
+   * @param ignoreIfExists
+   *   ignore existing index
    */
-  def create(): Unit = flint.createIndex(buildIndex())
+  def create(ignoreIfExists: Boolean = false): Unit =
+    flint.createIndex(buildIndex(), ignoreIfExists)
 
   /**
    * Build method for concrete builder class to implement

diff --git a/...ain/scala/org/opensearch/flint/spark/sql/covering/FlintSparkCoveringIndexAstBuilder.scala b/...ain/scala/org/opensearch/flint/spark/sql/covering/FlintSparkCoveringIndexAstBuilder.scala
@@ -35,7 +35,9 @@ trait FlintSparkCoveringIndexAstBuilder extends FlintSparkSqlExtensionsVisitor[C
         val colName = indexColCtx.multipartIdentifier().getText
         indexBuilder.addIndexColumns(colName)
       }
-      indexBuilder.create()
+
+      val ignoreIfExists = ctx.EXISTS() != null
+      indexBuilder.create(ignoreIfExists)
 
       // Trigger auto refresh if enabled
       if (isAutoRefreshEnabled(ctx.propertyList())) {

diff --git a/...ain/scala/org/opensearch/flint/spark/sql/skipping/FlintSparkSkippingIndexAstBuilder.scala b/...ain/scala/org/opensearch/flint/spark/sql/skipping/FlintSparkSkippingIndexAstBuilder.scala
@@ -42,7 +42,9 @@ trait FlintSparkSkippingIndexAstBuilder extends FlintSparkSqlExtensionsVisitor[C
           case MIN_MAX => indexBuilder.addMinMax(colName)
         }
       }
-      indexBuilder.create()
+
+      val ignoreIfExists = ctx.EXISTS() != null
+      indexBuilder.create(ignoreIfExists)
 
       // Trigger auto refresh if enabled
       if (isAutoRefreshEnabled(ctx.propertyList())) {

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
@@ -64,6 +64,26 @@ class FlintSparkCoveringIndexSqlITSuite extends FlintSparkSuite {
     indexData.count() shouldBe 2
   }
 
+  test("create covering index if not exists") {
+    sql(s"""
+           | CREATE INDEX IF NOT EXISTS $testIndex
+           | ON $testTable (name, age)
+           |""".stripMargin)
+    flint.describeIndex(testFlintIndex) shouldBe defined
+
+    // Expect error without IF NOT EXISTS, otherwise success
+    assertThrows[IllegalStateException] {
+      sql(s"""
+             | CREATE INDEX $testIndex
+             | ON $testTable (name, age)
+             |""".stripMargin)
+    }
+    sql(s"""
+           | CREATE INDEX IF NOT EXISTS $testIndex
+           | ON $testTable (name, age)
+           |""".stripMargin)
+  }
+
   test("drop covering index") {
     flint
       .coveringIndex()

diff --git a/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexSqlITSuite.scala b/integ-test/src/test/scala/org/opensearch/flint/spark/FlintSparkSkippingIndexSqlITSuite.scala
@@ -12,11 +12,8 @@ import org.opensearch.flint.spark.skipping.FlintSparkSkippingIndex.getSkippingIn
 import org.scalatest.matchers.must.Matchers.defined
 import org.scalatest.matchers.should.Matchers.convertToAnyShouldWrapper
 
-import org.apache.spark.FlintSuite
-import org.apache.spark.sql.{QueryTest, Row}
+import org.apache.spark.sql.Row
 import org.apache.spark.sql.flint.FlintDataSourceV2.FLINT_DATASOURCE
-import org.apache.spark.sql.flint.config.FlintSparkConf.{HOST_ENDPOINT, HOST_PORT, REFRESH_POLICY}
-import org.apache.spark.sql.streaming.StreamTest
 
 class FlintSparkSkippingIndexSqlITSuite extends FlintSparkSuite {
 
@@ -78,6 +75,28 @@ class FlintSparkSkippingIndexSqlITSuite extends FlintSparkSuite {
     indexData.count() shouldBe 2
   }
 
+  test("create skipping index if not exists") {
+    sql(s"""
+           | CREATE SKIPPING INDEX
+           | IF NOT EXISTS
+           | ON $testTable ( year PARTITION )
+           | """.stripMargin)
+    flint.describeIndex(testIndex) shouldBe defined
+
+    // Expect error without IF NOT EXISTS, otherwise success
+    assertThrows[IllegalStateException] {
+      sql(s"""
+             | CREATE SKIPPING INDEX
+             | ON $testTable ( year PARTITION )
+             | """.stripMargin)
+    }
+    sql(s"""
+           | CREATE SKIPPING INDEX
+           | IF NOT EXISTS
+           | ON $testTable ( year PARTITION )
+           | """.stripMargin)
+  }
+
   test("describe skipping index") {
     flint
       .skippingIndex()