Allow selectExpr to accept multiple columns

data-catering · Dec 4, 2023 · 88ebb73 · 88ebb73
1 parent 9eb092f
commit 88ebb73
Show file tree

Hide file tree

Showing 7 changed files with 12 additions and 10 deletions.
diff --git a/api/src/main/scala/com/github/pflooky/datacaterer/api/ValidationBuilder.scala b/api/src/main/scala/com/github/pflooky/datacaterer/api/ValidationBuilder.scala
@@ -137,16 +137,16 @@ case class ValidationBuilder(validation: Validation = ExpressionValidation()) {
    * SQL expression used to apply to columns before running validations.
    *
    * For example,
-   * {{{validation.selectExpr("PERCENTILE(amount, 0.5) AS median_amount, *")}}}
+   * {{{validation.selectExpr("PERCENTILE(amount, 0.5) AS median_amount", "*")}}}
    *
-   * @param expr SQL expression
+   * @param expr SQL expressions
    * @return ValidationBuilder
    * @see <a href="https://spark.apache.org/docs/latest/api/sql/">SQL expressions</a>
    */
-  def selectExpr(expr: String): ValidationBuilder = {
+  @varargs def selectExpr(expr: String*): ValidationBuilder = {
     validation match {
       case expressionValidation: ExpressionValidation =>
-        val withExpr = expressionValidation.modify(_.selectExpr).setTo(expr)
+        val withExpr = expressionValidation.modify(_.selectExpr).setTo(expr.toList)
         copyWithDescAndThreshold(withExpr)
     }
   }

diff --git a/api/src/main/scala/com/github/pflooky/datacaterer/api/model/ValidationModels.scala b/api/src/main/scala/com/github/pflooky/datacaterer/api/model/ValidationModels.scala
@@ -19,7 +19,7 @@ trait Validation {
 
 case class ExpressionValidation(
                                  whereExpr: String = "true",
-                                 selectExpr: String = "*"
+                                 selectExpr: List[String] = List("*")
                                ) extends Validation
 
 case class GroupByValidation(

diff --git a/api/src/main/scala/com/github/pflooky/datacaterer/api/parser/ValidationIdResolver.scala b/api/src/main/scala/com/github/pflooky/datacaterer/api/parser/ValidationIdResolver.scala
@@ -32,8 +32,10 @@ class ValidationBuilderSerializer extends JsonSerializer[ValidationBuilder] {
     gen.writeStartObject()
     validation match {
       case ExpressionValidation(expr, selectExpr) =>
+        gen.writeArrayFieldStart("selectExpr")
+        selectExpr.foreach(gen.writeObject)
+        gen.writeEndArray()
         gen.writeStringField("whereExpr", expr)
-        gen.writeStringField("selectExpr", selectExpr)
       case GroupByValidation(groupByCols, aggCol, aggType, expr) =>
         gen.writeArrayFieldStart("groupByCols")
         groupByCols.foreach(gen.writeObject)

diff --git a/app/src/main/scala/com/github/pflooky/datagen/core/generator/result/ResultHtmlWriter.scala b/app/src/main/scala/com/github/pflooky/datagen/core/generator/result/ResultHtmlWriter.scala
@@ -679,7 +679,7 @@ class ResultHtmlWriter {
     val options = validation match {
       case ExpressionValidation(expr, selectExpr) =>
         List(
-          List("selectExpr", selectExpr),
+          List("selectExpr", selectExpr.mkString(", ")),
           List("whereExpr", expr),
           List("errorThreshold", validation.errorThreshold.getOrElse(0.0).toString)
         )

diff --git a/app/src/main/scala/com/github/pflooky/datagen/core/validator/ValidationOperations.scala b/app/src/main/scala/com/github/pflooky/datagen/core/validator/ValidationOperations.scala
@@ -47,7 +47,7 @@ object ValidationHelper {
 
 class ExpressionValidationOps(expressionValidation: ExpressionValidation) extends ValidationOps(expressionValidation) {
   override def validate(df: DataFrame, dfCount: Long): ValidationResult = {
-    val dfWithSelectExpr = df.selectExpr(expressionValidation.selectExpr)
+    val dfWithSelectExpr = df.selectExpr(expressionValidation.selectExpr: _*)
     validateWithExpression(dfWithSelectExpr, dfCount, expressionValidation.whereExpr)
   }
 }

diff --git a/app/src/test/scala/com/github/pflooky/datagen/core/model/ValidationOperationsTest.scala b/app/src/test/scala/com/github/pflooky/datagen/core/model/ValidationOperationsTest.scala
@@ -29,7 +29,7 @@ class ValidationOperationsTest extends SparkSuite {
   }
 
   test("Can define select expression to run before where expression") {
-    val validation = ExpressionValidation("median_amount < 1000", "PERCENTILE(amount, 0.5) AS median_amount")
+    val validation = ExpressionValidation("median_amount < 1000", List("PERCENTILE(amount, 0.5) AS median_amount"))
     val result = new ExpressionValidationOps(validation).validate(df, 4)
 
     assert(result.isSuccess)

diff --git a/gradle.properties b/gradle.properties
@@ -1,5 +1,5 @@
 groupId=io.github.data-catering
-version=0.5.4
+version=0.5.5
 
 scalaVersion=2.12
 scalaSpecificVersion=2.12.15