You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2024/08/05 07:32:13 UTC
(spark) branch master updated: [SPARK-48346][SQL] Support for IF ELSE statements in SQL scripts
This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new f99291ac864d [SPARK-48346][SQL] Support for IF ELSE statements in SQL scripts
f99291ac864d is described below
commit f99291ac864d8902aa131fa93d46725a47d88263
Author: David Milicevic <da...@databricks.com>
AuthorDate: Mon Aug 5 15:31:59 2024 +0800
[SPARK-48346][SQL] Support for IF ELSE statements in SQL scripts
### What changes were proposed in this pull request?
This PR proposes introduction of IF/ELSE statement to SQL scripting language.
To evaluate conditions in IF or ELSE IF clauses, introduction of boolean statement evaluator is required as well.
Changes summary:
- Grammar/parser changes:
- `ifElseStatement` grammar rule
- `visitIfElseStatement` rule visitor
- `IfElseStatement` logical operator
- `IfElseStatementExec` execution node:
- Internal states - `Condition` and `Body`
- Iterator implementation - iterate over conditions until the one that evaluates to `true` is found
- Use `StatementBooleanEvaluator` implementation to evaluate conditions
- `DataFrameEvaluator`:
- Implementation of `StatementBooleanEvaluator`
- Evaluates results to `true` if it is single row, single column of boolean type with value `true`
- `SqlScriptingInterpreter` - add logic to transform `IfElseStatement` to `IfElseStatementExec`
### Why are the changes needed?
We are gradually introducing SQL Scripting to Spark, and IF/ELSE is one of the basic control flow constructs in the SQL language. For more details, check [JIRA item](https://issues.apache.org/jira/browse/SPARK-48346).
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
New tests are introduced to all of the three scripting test suites: `SqlScriptingParserSuite`, `SqlScriptingExecutionNodeSuite` and `SqlScriptingInterpreterSuite`.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #47442 from davidm-db/sql_scripting_if_else.
Authored-by: David Milicevic <da...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../spark/sql/catalyst/parser/SqlBaseParser.g4 | 7 +
.../spark/sql/catalyst/parser/AstBuilder.scala | 15 +-
.../parser/SqlScriptingLogicalOperators.scala | 16 ++
.../catalyst/parser/SqlScriptingParserSuite.scala | 178 ++++++++++++++++++
.../sql/scripting/SqlScriptingExecutionNode.scala | 105 +++++++++++
.../sql/scripting/SqlScriptingInterpreter.scala | 28 ++-
.../scripting/SqlScriptingExecutionNodeSuite.scala | 201 ++++++++++++++++++---
.../scripting/SqlScriptingInterpreterSuite.scala | 150 ++++++++++++++-
8 files changed, 663 insertions(+), 37 deletions(-)
diff --git a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4 b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4
index c7aa56cf920a..841c94d7868f 100644
--- a/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4
+++ b/sql/api/src/main/antlr4/org/apache/spark/sql/catalyst/parser/SqlBaseParser.g4
@@ -63,6 +63,7 @@ compoundStatement
: statement
| setStatementWithOptionalVarKeyword
| beginEndCompoundBlock
+ | ifElseStatement
;
setStatementWithOptionalVarKeyword
@@ -71,6 +72,12 @@ setStatementWithOptionalVarKeyword
LEFT_PAREN query RIGHT_PAREN #setVariableWithOptionalKeyword
;
+ifElseStatement
+ : IF booleanExpression THEN conditionalBodies+=compoundBody
+ (ELSE IF booleanExpression THEN conditionalBodies+=compoundBody)*
+ (ELSE elseBody=compoundBody)? END IF
+ ;
+
singleStatement
: (statement|setResetStatement) SEMICOLON* EOF
;
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala
index feb3ef4e7155..baeda8dd80be 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/AstBuilder.scala
@@ -205,10 +205,23 @@ class AstBuilder extends DataTypeAstBuilder
.map { s =>
SingleStatement(parsedPlan = visit(s).asInstanceOf[LogicalPlan])
}.getOrElse {
- visit(ctx.beginEndCompoundBlock()).asInstanceOf[CompoundPlanStatement]
+ visitChildren(ctx).asInstanceOf[CompoundPlanStatement]
}
}
+ override def visitIfElseStatement(ctx: IfElseStatementContext): IfElseStatement = {
+ IfElseStatement(
+ conditions = ctx.booleanExpression().asScala.toList.map(boolExpr => withOrigin(boolExpr) {
+ SingleStatement(
+ Project(
+ Seq(Alias(expression(boolExpr), "condition")()),
+ OneRowRelation()))
+ }),
+ conditionalBodies = ctx.conditionalBodies.asScala.toList.map(body => visitCompoundBody(body)),
+ elseBody = Option(ctx.elseBody).map(body => visitCompoundBody(body))
+ )
+ }
+
override def visitSingleStatement(ctx: SingleStatementContext): LogicalPlan = withOrigin(ctx) {
Option(ctx.statement().asInstanceOf[ParserRuleContext])
.orElse(Option(ctx.setResetStatement().asInstanceOf[ParserRuleContext]))
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingLogicalOperators.scala
index a68898d151e9..716129e58a42 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingLogicalOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingLogicalOperators.scala
@@ -58,3 +58,19 @@ case class SingleStatement(parsedPlan: LogicalPlan)
case class CompoundBody(
collection: Seq[CompoundPlanStatement],
label: Option[String]) extends CompoundPlanStatement
+
+/**
+ * Logical operator for IF ELSE statement.
+ * @param conditions Collection of conditions. First condition corresponds to IF clause,
+ * while others (if any) correspond to following ELSE IF clauses.
+ * @param conditionalBodies Collection of bodies that have a corresponding condition,
+ * in IF or ELSE IF branches.
+ * @param elseBody Body that is executed if none of the conditions are met,
+ * i.e. ELSE branch.
+ */
+case class IfElseStatement(
+ conditions: Seq[SingleStatement],
+ conditionalBodies: Seq[CompoundBody],
+ elseBody: Option[CompoundBody]) extends CompoundPlanStatement {
+ assert(conditions.length == conditionalBodies.length)
+}
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala
index 0b7a8d53612d..b108564aa406 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/parser/SqlScriptingParserSuite.scala
@@ -348,6 +348,184 @@ class SqlScriptingParserSuite extends SparkFunSuite with SQLHelper {
assert(e.getMessage.contains("Syntax error"))
}
+ test("if") {
+ val sqlScriptText =
+ """
+ |BEGIN
+ | IF 1=1 THEN
+ | SELECT 42;
+ | END IF;
+ |END
+ |""".stripMargin
+ val tree = parseScript(sqlScriptText)
+ assert(tree.collection.length == 1)
+ assert(tree.collection.head.isInstanceOf[IfElseStatement])
+ val ifStmt = tree.collection.head.asInstanceOf[IfElseStatement]
+ assert(ifStmt.conditions.length == 1)
+ assert(ifStmt.conditions.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditions.head.getText == "1=1")
+ }
+
+ test("if else") {
+ val sqlScriptText =
+ """BEGIN
+ |IF 1 = 1 THEN
+ | SELECT 1;
+ |ELSE
+ | SELECT 2;
+ |END IF;
+ |END
+ """.stripMargin
+ val tree = parseScript(sqlScriptText)
+ assert(tree.collection.length == 1)
+ assert(tree.collection.head.isInstanceOf[IfElseStatement])
+
+ val ifStmt = tree.collection.head.asInstanceOf[IfElseStatement]
+ assert(ifStmt.conditions.length == 1)
+ assert(ifStmt.conditionalBodies.length == 1)
+ assert(ifStmt.elseBody.isDefined)
+
+ assert(ifStmt.conditions.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditions.head.getText == "1 = 1")
+
+ assert(ifStmt.conditionalBodies.head.collection.length == 1)
+ assert(ifStmt.conditionalBodies.head.collection.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditionalBodies.head.collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 1")
+
+ assert(ifStmt.elseBody.get.collection.length == 1)
+ assert(ifStmt.elseBody.get.collection.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.elseBody.get.collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 2")
+ }
+
+ test("if else if") {
+ val sqlScriptText =
+ """BEGIN
+ |IF 1 = 1 THEN
+ | SELECT 1;
+ |ELSE IF 2 = 2 THEN
+ | SELECT 2;
+ |ELSE
+ | SELECT 3;
+ |END IF;
+ |END
+ """.stripMargin
+ val tree = parseScript(sqlScriptText)
+ assert(tree.collection.length == 1)
+ assert(tree.collection.head.isInstanceOf[IfElseStatement])
+
+ val ifStmt = tree.collection.head.asInstanceOf[IfElseStatement]
+ assert(ifStmt.conditions.length == 2)
+ assert(ifStmt.conditionalBodies.length == 2)
+ assert(ifStmt.elseBody.isDefined)
+
+ assert(ifStmt.conditions.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditions.head.getText == "1 = 1")
+
+ assert(ifStmt.conditionalBodies.head.collection.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditionalBodies.head.collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 1")
+
+ assert(ifStmt.conditions(1).isInstanceOf[SingleStatement])
+ assert(ifStmt.conditions(1).getText == "2 = 2")
+
+ assert(ifStmt.conditionalBodies(1).collection.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditionalBodies(1).collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 2")
+
+ assert(ifStmt.elseBody.get.collection.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.elseBody.get.collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 3")
+ }
+
+ test("if multi else if") {
+ val sqlScriptText =
+ """BEGIN
+ |IF 1 = 1 THEN
+ | SELECT 1;
+ |ELSE IF 2 = 2 THEN
+ | SELECT 2;
+ |ELSE IF 3 = 3 THEN
+ | SELECT 3;
+ |END IF;
+ |END
+ """.stripMargin
+ val tree = parseScript(sqlScriptText)
+ assert(tree.collection.length == 1)
+ assert(tree.collection.head.isInstanceOf[IfElseStatement])
+
+ val ifStmt = tree.collection.head.asInstanceOf[IfElseStatement]
+ assert(ifStmt.conditions.length == 3)
+ assert(ifStmt.conditionalBodies.length == 3)
+ assert(ifStmt.elseBody.isEmpty)
+
+ assert(ifStmt.conditions.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditions.head.getText == "1 = 1")
+
+ assert(ifStmt.conditionalBodies.head.collection.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditionalBodies.head.collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 1")
+
+ assert(ifStmt.conditions(1).isInstanceOf[SingleStatement])
+ assert(ifStmt.conditions(1).getText == "2 = 2")
+
+ assert(ifStmt.conditionalBodies(1).collection.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditionalBodies(1).collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 2")
+
+ assert(ifStmt.conditions(2).isInstanceOf[SingleStatement])
+ assert(ifStmt.conditions(2).getText == "3 = 3")
+
+ assert(ifStmt.conditionalBodies(2).collection.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditionalBodies(2).collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 3")
+ }
+
+ test("if nested") {
+ val sqlScriptText =
+ """
+ |BEGIN
+ | IF 1=1 THEN
+ | IF 2=1 THEN
+ | SELECT 41;
+ | ELSE
+ | SELECT 42;
+ | END IF;
+ | END IF;
+ |END
+ |""".stripMargin
+ val tree = parseScript(sqlScriptText)
+ assert(tree.collection.length == 1)
+ assert(tree.collection.head.isInstanceOf[IfElseStatement])
+
+ val ifStmt = tree.collection.head.asInstanceOf[IfElseStatement]
+ assert(ifStmt.conditions.length == 1)
+ assert(ifStmt.conditionalBodies.length == 1)
+ assert(ifStmt.elseBody.isEmpty)
+
+ assert(ifStmt.conditions.head.isInstanceOf[SingleStatement])
+ assert(ifStmt.conditions.head.getText == "1=1")
+
+ assert(ifStmt.conditionalBodies.head.collection.head.isInstanceOf[IfElseStatement])
+ val nestedIfStmt = ifStmt.conditionalBodies.head.collection.head.asInstanceOf[IfElseStatement]
+
+ assert(nestedIfStmt.conditions.length == 1)
+ assert(nestedIfStmt.conditionalBodies.length == 1)
+ assert(nestedIfStmt.elseBody.isDefined)
+
+ assert(nestedIfStmt.conditions.head.isInstanceOf[SingleStatement])
+ assert(nestedIfStmt.conditions.head.getText == "2=1")
+
+ assert(nestedIfStmt.conditionalBodies.head.collection.head.isInstanceOf[SingleStatement])
+ assert(nestedIfStmt.conditionalBodies.head.collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 41")
+
+ assert(nestedIfStmt.elseBody.get.collection.head.isInstanceOf[SingleStatement])
+ assert(nestedIfStmt.elseBody.get.collection.head.asInstanceOf[SingleStatement]
+ .getText == "SELECT 42")
+ }
+
// Helper methods
def cleanupStatementString(statementStr: String): String = {
statementStr
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala
index 51e9304297f4..993d08878b22 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala
@@ -19,8 +19,10 @@ package org.apache.spark.sql.scripting
import org.apache.spark.SparkException
import org.apache.spark.internal.Logging
+import org.apache.spark.sql.{Dataset, SparkSession}
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
import org.apache.spark.sql.catalyst.trees.{Origin, WithOrigin}
+import org.apache.spark.sql.types.BooleanType
/**
* Trait for all SQL scripting execution nodes used during interpretation phase.
@@ -55,6 +57,33 @@ trait NonLeafStatementExec extends CompoundStatementExec {
* Tree iterator.
*/
def getTreeIterator: Iterator[CompoundStatementExec]
+
+ /**
+ * Evaluate the boolean condition represented by the statement.
+ * @param session SparkSession that SQL script is executed within.
+ * @param statement Statement representing the boolean condition to evaluate.
+ * @return Whether the condition evaluates to True.
+ */
+ protected def evaluateBooleanCondition(
+ session: SparkSession,
+ statement: LeafStatementExec): Boolean = statement match {
+ case statement: SingleStatementExec =>
+ assert(!statement.isExecuted)
+ statement.isExecuted = true
+
+ // DataFrame evaluates to True if it is single row, single column
+ // of boolean type with value True.
+ val df = Dataset.ofRows(session, statement.parsedPlan)
+ df.schema.fields match {
+ case Array(field) if field.dataType == BooleanType =>
+ df.limit(2).collect() match {
+ case Array(row) => row.getBoolean(0)
+ case _ => false
+ }
+ case _ => false
+ }
+ case _ => false
+ }
}
/**
@@ -155,3 +184,79 @@ abstract class CompoundNestedStatementIteratorExec(collection: Seq[CompoundState
*/
class CompoundBodyExec(statements: Seq[CompoundStatementExec])
extends CompoundNestedStatementIteratorExec(statements)
+
+/**
+ * Executable node for IfElseStatement.
+ * @param conditions Collection of executable conditions. First condition corresponds to IF clause,
+ * while others (if any) correspond to following ELSE IF clauses.
+ * @param conditionalBodies Collection of executable bodies that have a corresponding condition,
+* in IF or ELSE IF branches.
+ * @param elseBody Body that is executed if none of the conditions are met,
+ * i.e. ELSE branch.
+ * @param session Spark session that SQL script is executed within.
+ */
+class IfElseStatementExec(
+ conditions: Seq[SingleStatementExec],
+ conditionalBodies: Seq[CompoundBodyExec],
+ elseBody: Option[CompoundBodyExec],
+ session: SparkSession) extends NonLeafStatementExec {
+ private object IfElseState extends Enumeration {
+ val Condition, Body = Value
+ }
+
+ private var state = IfElseState.Condition
+ private var curr: Option[CompoundStatementExec] = Some(conditions.head)
+
+ private var clauseIdx: Int = 0
+ private val conditionsCount = conditions.length
+ assert(conditionsCount == conditionalBodies.length)
+
+ private lazy val treeIterator: Iterator[CompoundStatementExec] =
+ new Iterator[CompoundStatementExec] {
+ override def hasNext: Boolean = curr.nonEmpty
+
+ override def next(): CompoundStatementExec = state match {
+ case IfElseState.Condition =>
+ assert(curr.get.isInstanceOf[SingleStatementExec])
+ val condition = curr.get.asInstanceOf[SingleStatementExec]
+ if (evaluateBooleanCondition(session, condition)) {
+ state = IfElseState.Body
+ curr = Some(conditionalBodies(clauseIdx))
+ } else {
+ clauseIdx += 1
+ if (clauseIdx < conditionsCount) {
+ // There are ELSE IF clauses remaining.
+ state = IfElseState.Condition
+ curr = Some(conditions(clauseIdx))
+ } else if (elseBody.isDefined) {
+ // ELSE clause exists.
+ state = IfElseState.Body
+ curr = Some(elseBody.get)
+ } else {
+ // No remaining clauses.
+ curr = None
+ }
+ }
+ condition
+ case IfElseState.Body =>
+ assert(curr.get.isInstanceOf[CompoundBodyExec])
+ val currBody = curr.get.asInstanceOf[CompoundBodyExec]
+ val retStmt = currBody.getTreeIterator.next()
+ if (!currBody.getTreeIterator.hasNext) {
+ curr = None
+ }
+ retStmt
+ }
+ }
+
+ override def getTreeIterator: Iterator[CompoundStatementExec] = treeIterator
+
+ override def reset(): Unit = {
+ state = IfElseState.Condition
+ curr = Some(conditions.head)
+ clauseIdx = 0
+ conditions.foreach(c => c.reset())
+ conditionalBodies.foreach(b => b.reset())
+ elseBody.foreach(b => b.reset())
+ }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala
index 6453b204e162..e553ff10f0a2 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreter.scala
@@ -17,8 +17,9 @@
package org.apache.spark.sql.scripting
+import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.analysis.UnresolvedIdentifier
-import org.apache.spark.sql.catalyst.parser.{CompoundBody, CompoundPlanStatement, SingleStatement}
+import org.apache.spark.sql.catalyst.parser.{CompoundBody, CompoundPlanStatement, IfElseStatement, SingleStatement}
import org.apache.spark.sql.catalyst.plans.logical.{CreateVariable, DropVariable, LogicalPlan}
import org.apache.spark.sql.catalyst.trees.Origin
@@ -33,11 +34,15 @@ case class SqlScriptingInterpreter() {
*
* @param compound
* CompoundBody for which to build the plan.
+ * @param session
+ * Spark session that SQL script is executed within.
* @return
* Iterator through collection of statements to be executed.
*/
- def buildExecutionPlan(compound: CompoundBody): Iterator[CompoundStatementExec] = {
- transformTreeIntoExecutable(compound).asInstanceOf[CompoundBodyExec].getTreeIterator
+ def buildExecutionPlan(
+ compound: CompoundBody,
+ session: SparkSession): Iterator[CompoundStatementExec] = {
+ transformTreeIntoExecutable(compound, session).asInstanceOf[CompoundBodyExec].getTreeIterator
}
/**
@@ -55,12 +60,16 @@ case class SqlScriptingInterpreter() {
/**
* Transform the parsed tree to the executable node.
+ *
* @param node
* Root node of the parsed tree.
+ * @param session
+ * Spark session that SQL script is executed within.
* @return
* Executable statement.
*/
- private def transformTreeIntoExecutable(node: CompoundPlanStatement): CompoundStatementExec =
+ private def transformTreeIntoExecutable(
+ node: CompoundPlanStatement, session: SparkSession): CompoundStatementExec =
node match {
case body: CompoundBody =>
// TODO [SPARK-48530]: Current logic doesn't support scoped variables and shadowing.
@@ -73,7 +82,16 @@ case class SqlScriptingInterpreter() {
.map(new SingleStatementExec(_, Origin(), isInternal = true))
.reverse
new CompoundBodyExec(
- body.collection.map(st => transformTreeIntoExecutable(st)) ++ dropVariables)
+ body.collection.map(st => transformTreeIntoExecutable(st, session)) ++ dropVariables)
+ case IfElseStatement(conditions, conditionalBodies, elseBody) =>
+ val conditionsExec = conditions.map(condition =>
+ new SingleStatementExec(condition.parsedPlan, condition.origin, isInternal = false))
+ val conditionalBodiesExec = conditionalBodies.map(body =>
+ transformTreeIntoExecutable(body, session).asInstanceOf[CompoundBodyExec])
+ val unconditionalBodiesExec = elseBody.map(body =>
+ transformTreeIntoExecutable(body, session).asInstanceOf[CompoundBodyExec])
+ new IfElseStatementExec(
+ conditionsExec, conditionalBodiesExec, unconditionalBodiesExec, session)
case sparkStatement: SingleStatement =>
new SingleStatementExec(
sparkStatement.parsedPlan,
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNodeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNodeSuite.scala
index 003a8061a604..03b1ccfb7a3e 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNodeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNodeSuite.scala
@@ -21,67 +21,208 @@ import org.apache.spark.SparkFunSuite
import org.apache.spark.sql.catalyst.expressions.{Alias, Literal}
import org.apache.spark.sql.catalyst.plans.logical.{OneRowRelation, Project}
import org.apache.spark.sql.catalyst.trees.Origin
+import org.apache.spark.sql.test.SharedSparkSession
/**
* Unit tests for execution nodes from SqlScriptingExecutionNode.scala.
* Execution nodes are constructed manually and iterated through.
* It is then checked if the leaf statements have been iterated in the expected order.
*/
-class SqlScriptingExecutionNodeSuite extends SparkFunSuite {
+class SqlScriptingExecutionNodeSuite extends SparkFunSuite with SharedSparkSession {
// Helpers
case class TestLeafStatement(testVal: String) extends LeafStatementExec {
override def reset(): Unit = ()
}
- case class TestNestedStatementIterator(statements: Seq[CompoundStatementExec])
- extends CompoundNestedStatementIteratorExec(statements)
-
- case class TestBody(statements: Seq[CompoundStatementExec]) extends CompoundBodyExec(statements)
-
- case class TestSparkStatementWithPlan(testVal: String)
+ case class TestIfElseCondition(condVal: Boolean, description: String)
extends SingleStatementExec(
- parsedPlan = Project(Seq(Alias(Literal(testVal), "condition")()), OneRowRelation()),
- Origin(startIndex = Some(0), stopIndex = Some(testVal.length)),
+ parsedPlan = Project(Seq(Alias(Literal(condVal), description)()), OneRowRelation()),
+ Origin(startIndex = Some(0), stopIndex = Some(description.length)),
isInternal = false)
- // Tests
- test("test body - single statement") {
- val iter = TestNestedStatementIterator(Seq(TestLeafStatement("one"))).getTreeIterator
- val statements = iter.map {
- case TestLeafStatement(v) => v
+ private def extractStatementValue(statement: CompoundStatementExec): String =
+ statement match {
+ case TestLeafStatement(testVal) => testVal
+ case TestIfElseCondition(_, description) => description
case _ => fail("Unexpected statement type")
- }.toList
+ }
- assert(statements === List("one"))
+ // Tests
+ test("test body - single statement") {
+ val iter = new CompoundBodyExec(Seq(TestLeafStatement("one"))).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq("one"))
}
test("test body - no nesting") {
- val iter = TestNestedStatementIterator(
+ val iter = new CompoundBodyExec(
Seq(
TestLeafStatement("one"),
TestLeafStatement("two"),
TestLeafStatement("three")))
.getTreeIterator
- val statements = iter.map {
- case TestLeafStatement(v) => v
- case _ => fail("Unexpected statement type")
- }.toList
-
+ val statements = iter.map(extractStatementValue).toSeq
assert(statements === Seq("one", "two", "three"))
}
test("test body - nesting") {
- val iter = TestNestedStatementIterator(
+ val iter = new CompoundBodyExec(
Seq(
- TestNestedStatementIterator(Seq(TestLeafStatement("one"), TestLeafStatement("two"))),
+ new CompoundBodyExec(Seq(TestLeafStatement("one"), TestLeafStatement("two"))),
TestLeafStatement("three"),
- TestNestedStatementIterator(Seq(TestLeafStatement("four"), TestLeafStatement("five")))))
+ new CompoundBodyExec(Seq(TestLeafStatement("four"), TestLeafStatement("five")))))
.getTreeIterator
- val statements = iter.map {
- case TestLeafStatement(v) => v
- case _ => fail("Unexpected statement type")
- }.toList
-
+ val statements = iter.map(extractStatementValue).toSeq
assert(statements === Seq("one", "two", "three", "four", "five"))
}
+
+ test("if else - enter body of the IF clause") {
+ val iter = new CompoundBodyExec(Seq(
+ new IfElseStatementExec(
+ conditions = Seq(
+ TestIfElseCondition(condVal = true, description = "con1")
+ ),
+ conditionalBodies = Seq(
+ new CompoundBodyExec(Seq(TestLeafStatement("body1")))
+ ),
+ elseBody = Some(new CompoundBodyExec(Seq(TestLeafStatement("body2")))),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq("con1", "body1"))
+ }
+
+ test("if else - enter body of the ELSE clause") {
+ val iter = new CompoundBodyExec(Seq(
+ new IfElseStatementExec(
+ conditions = Seq(
+ TestIfElseCondition(condVal = false, description = "con1")
+ ),
+ conditionalBodies = Seq(
+ new CompoundBodyExec(Seq(TestLeafStatement("body1")))
+ ),
+ elseBody = Some(new CompoundBodyExec(Seq(TestLeafStatement("body2")))),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq("con1", "body2"))
+ }
+
+ test("if else if - enter body of the IF clause") {
+ val iter = new CompoundBodyExec(Seq(
+ new IfElseStatementExec(
+ conditions = Seq(
+ TestIfElseCondition(condVal = true, description = "con1"),
+ TestIfElseCondition(condVal = false, description = "con2")
+ ),
+ conditionalBodies = Seq(
+ new CompoundBodyExec(Seq(TestLeafStatement("body1"))),
+ new CompoundBodyExec(Seq(TestLeafStatement("body2")))
+ ),
+ elseBody = Some(new CompoundBodyExec(Seq(TestLeafStatement("body3")))),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq("con1", "body1"))
+ }
+
+ test("if else if - enter body of the ELSE IF clause") {
+ val iter = new CompoundBodyExec(Seq(
+ new IfElseStatementExec(
+ conditions = Seq(
+ TestIfElseCondition(condVal = false, description = "con1"),
+ TestIfElseCondition(condVal = true, description = "con2")
+ ),
+ conditionalBodies = Seq(
+ new CompoundBodyExec(Seq(TestLeafStatement("body1"))),
+ new CompoundBodyExec(Seq(TestLeafStatement("body2")))
+ ),
+ elseBody = Some(new CompoundBodyExec(Seq(TestLeafStatement("body3")))),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq("con1", "con2", "body2"))
+ }
+
+ test("if else if - enter body of the second ELSE IF clause") {
+ val iter = new CompoundBodyExec(Seq(
+ new IfElseStatementExec(
+ conditions = Seq(
+ TestIfElseCondition(condVal = false, description = "con1"),
+ TestIfElseCondition(condVal = false, description = "con2"),
+ TestIfElseCondition(condVal = true, description = "con3")
+ ),
+ conditionalBodies = Seq(
+ new CompoundBodyExec(Seq(TestLeafStatement("body1"))),
+ new CompoundBodyExec(Seq(TestLeafStatement("body2"))),
+ new CompoundBodyExec(Seq(TestLeafStatement("body3")))
+ ),
+ elseBody = Some(new CompoundBodyExec(Seq(TestLeafStatement("body4")))),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq("con1", "con2", "con3", "body3"))
+ }
+
+ test("if else if - enter body of the ELSE clause") {
+ val iter = new CompoundBodyExec(Seq(
+ new IfElseStatementExec(
+ conditions = Seq(
+ TestIfElseCondition(condVal = false, description = "con1"),
+ TestIfElseCondition(condVal = false, description = "con2")
+ ),
+ conditionalBodies = Seq(
+ new CompoundBodyExec(Seq(TestLeafStatement("body1"))),
+ new CompoundBodyExec(Seq(TestLeafStatement("body2")))
+ ),
+ elseBody = Some(new CompoundBodyExec(Seq(TestLeafStatement("body3")))),
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq("con1", "con2", "body3"))
+ }
+
+ test("if else if - without else (successful check)") {
+ val iter = new CompoundBodyExec(Seq(
+ new IfElseStatementExec(
+ conditions = Seq(
+ TestIfElseCondition(condVal = false, description = "con1"),
+ TestIfElseCondition(condVal = true, description = "con2")
+ ),
+ conditionalBodies = Seq(
+ new CompoundBodyExec(Seq(TestLeafStatement("body1"))),
+ new CompoundBodyExec(Seq(TestLeafStatement("body2")))
+ ),
+ elseBody = None,
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq("con1", "con2", "body2"))
+ }
+
+ test("if else if - without else (unsuccessful checks)") {
+ val iter = new CompoundBodyExec(Seq(
+ new IfElseStatementExec(
+ conditions = Seq(
+ TestIfElseCondition(condVal = false, description = "con1"),
+ TestIfElseCondition(condVal = false, description = "con2")
+ ),
+ conditionalBodies = Seq(
+ new CompoundBodyExec(Seq(TestLeafStatement("body1"))),
+ new CompoundBodyExec(Seq(TestLeafStatement("body2")))
+ ),
+ elseBody = None,
+ session = spark
+ )
+ )).getTreeIterator
+ val statements = iter.map(extractStatementValue).toSeq
+ assert(statements === Seq("con1", "con2"))
+ }
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala
index 52dd7547c644..2d61f5660942 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala
@@ -32,7 +32,7 @@ class SqlScriptingInterpreterSuite extends QueryTest with SharedSparkSession {
private def verifySqlScriptResult(sqlText: String, expected: Seq[Seq[Row]]): Unit = {
val interpreter = SqlScriptingInterpreter()
val compoundBody = spark.sessionState.sqlParser.parseScript(sqlText)
- val executionPlan = interpreter.buildExecutionPlan(compoundBody)
+ val executionPlan = interpreter.buildExecutionPlan(compoundBody, spark)
val result = executionPlan.flatMap {
case statement: SingleStatementExec =>
if (statement.isExecuted) {
@@ -214,4 +214,152 @@ class SqlScriptingInterpreterSuite extends QueryTest with SharedSparkSession {
)
verifySqlScriptResult(sqlScript, expected)
}
+
+ test("if") {
+ val commands =
+ """
+ |BEGIN
+ | IF 1=1 THEN
+ | SELECT 42;
+ | END IF;
+ |END
+ |""".stripMargin
+ val expected = Seq(Seq(Row(42)))
+ verifySqlScriptResult(commands, expected)
+ }
+
+ test("if nested") {
+ val commands =
+ """
+ |BEGIN
+ | IF 1=1 THEN
+ | IF 2=1 THEN
+ | SELECT 41;
+ | ELSE
+ | SELECT 42;
+ | END IF;
+ | END IF;
+ |END
+ |""".stripMargin
+ val expected = Seq(Seq(Row(42)))
+ verifySqlScriptResult(commands, expected)
+ }
+
+ test("if else going in if") {
+ val commands =
+ """
+ |BEGIN
+ | IF 1=1
+ | THEN
+ | SELECT 42;
+ | ELSE
+ | SELECT 43;
+ | END IF;
+ |END
+ |""".stripMargin
+
+ val expected = Seq(Seq(Row(42)))
+ verifySqlScriptResult(commands, expected)
+ }
+
+ test("if else if going in else if") {
+ val commands =
+ """
+ |BEGIN
+ | IF 1=2
+ | THEN
+ | SELECT 42;
+ | ELSE IF 1=1
+ | THEN
+ | SELECT 43;
+ | ELSE
+ | SELECT 44;
+ | END IF;
+ |END
+ |""".stripMargin
+
+ val expected = Seq(Seq(Row(43)))
+ verifySqlScriptResult(commands, expected)
+ }
+
+ test("if else going in else") {
+ val commands =
+ """
+ |BEGIN
+ | IF 1=2
+ | THEN
+ | SELECT 42;
+ | ELSE
+ | SELECT 43;
+ | END IF;
+ |END
+ |""".stripMargin
+
+ val expected = Seq(Seq(Row(43)))
+ verifySqlScriptResult(commands, expected)
+ }
+
+ test("if else if going in else") {
+ val commands =
+ """
+ |BEGIN
+ | IF 1=2
+ | THEN
+ | SELECT 42;
+ | ELSE IF 1=3
+ | THEN
+ | SELECT 43;
+ | ELSE
+ | SELECT 44;
+ | END IF;
+ |END
+ |""".stripMargin
+
+ val expected = Seq(Seq(Row(44)))
+ verifySqlScriptResult(commands, expected)
+ }
+
+ test("if with count") {
+ withTable("t") {
+ val commands =
+ """
+ |BEGIN
+ |CREATE TABLE t (a INT, b STRING, c DOUBLE) USING parquet;
+ |INSERT INTO t VALUES (1, 'a', 1.0);
+ |INSERT INTO t VALUES (1, 'a', 1.0);
+ |IF (SELECT COUNT(*) > 2 FROM t) THEN
+ | SELECT 42;
+ | ELSE
+ | SELECT 43;
+ | END IF;
+ |END
+ |""".stripMargin
+
+ val expected = Seq(Seq.empty[Row], Seq.empty[Row], Seq.empty[Row], Seq(Row(43)))
+ verifySqlScriptResult(commands, expected)
+ }
+ }
+
+ test("if else if with count") {
+ withTable("t") {
+ val commands =
+ """
+ |BEGIN
+ | CREATE TABLE t (a INT, b STRING, c DOUBLE) USING parquet;
+ | INSERT INTO t VALUES (1, 'a', 1.0);
+ | INSERT INTO t VALUES (1, 'a', 1.0);
+ | IF (SELECT COUNT(*) > 2 FROM t) THEN
+ | SELECT 42;
+ | ELSE IF (SELECT COUNT(*) > 1 FROM t) THEN
+ | SELECT 43;
+ | ELSE
+ | SELECT 44;
+ | END IF;
+ |END
+ |""".stripMargin
+
+ val expected = Seq(Seq.empty[Row], Seq.empty[Row], Seq.empty[Row], Seq(Row(43)))
+ verifySqlScriptResult(commands, expected)
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org