Refactor javadoc and UT

Signed-off-by: Chen Dai <[email protected]>
opensearch-project · Jan 8, 2024 · 1242d8c · 1242d8c
1 parent 17b3803
commit 1242d8c
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 13 deletions.
diff --git a/flint-spark-integration/src/main/antlr4/FlintSparkSqlExtensions.g4 b/flint-spark-integration/src/main/antlr4/FlintSparkSqlExtensions.g4
@@ -154,10 +154,10 @@ indexColTypeList
 
 indexColType
     : identifier skipType=(PARTITION | VALUE_SET | MIN_MAX)
-        (LEFT_PAREN propertyValues RIGHT_PAREN)?
+        (LEFT_PAREN skipParams RIGHT_PAREN)?
     ;
 
-propertyValues
+skipParams
     : propertyValue (COMMA propertyValue)*
     ;
 

diff --git a/...ation/src/main/scala/org/opensearch/flint/spark/skipping/FlintSparkSkippingStrategy.scala b/...ation/src/main/scala/org/opensearch/flint/spark/skipping/FlintSparkSkippingStrategy.scala
@@ -32,7 +32,7 @@ trait FlintSparkSkippingStrategy {
   val columnType: String
 
   /**
-   * Skipping algorithm parameters.
+   * Skipping algorithm named parameters.
    */
   val parameters: Map[String, String] = Map.empty
 

diff --git a/...ain/scala/org/opensearch/flint/spark/sql/skipping/FlintSparkSkippingIndexAstBuilder.scala b/...ain/scala/org/opensearch/flint/spark/sql/skipping/FlintSparkSkippingIndexAstBuilder.scala
@@ -46,11 +46,12 @@ trait FlintSparkSkippingIndexAstBuilder extends FlintSparkSqlExtensionsVisitor[A
       ctx.indexColTypeList().indexColType().forEach { colTypeCtx =>
         val colName = colTypeCtx.identifier().getText
         val skipType = SkippingKind.withName(colTypeCtx.skipType.getText)
-        val paramValues = visitPropertyValues(colTypeCtx.propertyValues())
+        val skipParams = visitSkipParams(colTypeCtx.skipParams())
         skipType match {
           case PARTITION => indexBuilder.addPartitions(colName)
           case VALUE_SET =>
-            indexBuilder.addValueSet(colName, (Seq(VALUE_SET_MAX_SIZE_KEY) zip paramValues).toMap)
+            val valueSetParams = (Seq(VALUE_SET_MAX_SIZE_KEY) zip skipParams).toMap
+            indexBuilder.addValueSet(colName, valueSetParams)
           case MIN_MAX => indexBuilder.addMinMax(colName)
         }
       }
@@ -112,7 +113,7 @@ trait FlintSparkSkippingIndexAstBuilder extends FlintSparkSqlExtensionsVisitor[A
     }
   }
 
-  override def visitPropertyValues(ctx: PropertyValuesContext): Seq[String] = {
+  override def visitSkipParams(ctx: SkipParamsContext): Seq[String] = {
     if (ctx == null) {
       Seq.empty
     } else {

diff --git a/...st/scala/org/opensearch/flint/spark/skipping/valueset/ValueSetSkippingStrategySuite.scala b/...st/scala/org/opensearch/flint/spark/skipping/valueset/ValueSetSkippingStrategySuite.scala
@@ -7,28 +7,46 @@ package org.opensearch.flint.spark.skipping.valueset
 
 import org.opensearch.flint.spark.skipping.{FlintSparkSkippingStrategy, FlintSparkSkippingStrategySuite}
 import org.opensearch.flint.spark.skipping.valueset.ValueSetSkippingStrategy.{DEFAULT_VALUE_SET_MAX_SIZE, VALUE_SET_MAX_SIZE_KEY}
-import org.scalatest.matchers.should.Matchers
+import org.scalatest.matchers.should.Matchers.convertToAnyShouldWrapper
 
 import org.apache.spark.SparkFunSuite
 import org.apache.spark.sql.catalyst.expressions.{Abs, AttributeReference, EqualTo, Literal}
-import org.apache.spark.sql.functions.{col, isnull}
+import org.apache.spark.sql.functions._
 import org.apache.spark.sql.types.StringType
 
-class ValueSetSkippingStrategySuite
-    extends SparkFunSuite
-    with FlintSparkSkippingStrategySuite
-    with Matchers {
+class ValueSetSkippingStrategySuite extends SparkFunSuite with FlintSparkSkippingStrategySuite {
 
   override val strategy: FlintSparkSkippingStrategy =
     ValueSetSkippingStrategy(columnName = "name", columnType = "string")
 
   private val name = AttributeReference("name", StringType, nullable = false)()
 
-  test("should return properties with default value") {
+  test("should return parameters with default value") {
     strategy.parameters shouldBe Map(
       VALUE_SET_MAX_SIZE_KEY -> DEFAULT_VALUE_SET_MAX_SIZE.toString)
   }
 
+  test("should build aggregator with default parameter") {
+    strategy.getAggregators.head.semanticEquals(
+      when(size(collect_set("name")) > DEFAULT_VALUE_SET_MAX_SIZE, lit(null))
+        .otherwise(collect_set("name"))
+        .expr) shouldBe true
+  }
+
+  test("should use given parameter value") {
+    val strategy =
+      ValueSetSkippingStrategy(
+        columnName = "name",
+        columnType = "string",
+        params = Map(VALUE_SET_MAX_SIZE_KEY -> "5"))
+
+    strategy.parameters shouldBe Map(VALUE_SET_MAX_SIZE_KEY -> "5")
+    strategy.getAggregators.head.semanticEquals(
+      when(size(collect_set("name")) > 5, lit(null))
+        .otherwise(collect_set("name"))
+        .expr) shouldBe true
+  }
+
   test("should rewrite EqualTo(<indexCol>, <value>)") {
     EqualTo(name, Literal("hello")) shouldRewriteTo
       (isnull(col("name")) || col("name") === "hello")