[SPARK-48376][SQL][FOLLOWUP] Add ITERATE statement fixes

### What changes were proposed in this pull request?
Previous [pull request](#47973) introduced new logic for `LEAVE` and `ITERATE` statements, but I missed to introduce minor part of `ITERATE` statement handling.
While `ITERATE` statement is not allowed for compounds (i.e. user cannot use `LEAVE <label>` where label belongs to compound) it's still possible to use `ITERATE` statement within the compound (i.e. when using label of the surrounding `WHILE` loop for example). In such cases, it's required to drop variables defined in the compound (the same way as for `ITERATE` statement).
We don't yet have execution logic (I will work on POC and design doc during the next week), but this is setting up stuff now, so we don't miss anything in future.

### Why are the changes needed?
Changes are minor, they improve consistency among test suites for SQL scripting.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
Added small test, but not completely related to this change. Since we don't have execution logic yet, it's hard to test. I will add more tests together with the execution logic.

### Was this patch authored or co-authored using generative AI tooling?
No.

Closes #48022 from davidm-db/missing_logic_for_leave_statement.

Authored-by: David Milicevic <[email protected]>
Signed-off-by: Max Gekk <[email protected]>

Author: davidm-db

sql/core/src/main/scala/org/apache/spark/sql/scripting/SqlScriptingExecutionNode.scala

@@ -176,6 +176,10 @@ abstract class CompoundNestedStatementIteratorExec(
             handleLeaveStatement(leaveStatement)
             curr = None
             leaveStatement
+          case Some(iterateStatement: IterateStatementExec) =>
+            handleIterateStatement(iterateStatement)
+            curr = None
+            iterateStatement
           case Some(statement: LeafStatementExec) =>
             curr = if (localIterator.hasNext) Some(localIterator.next()) else None
             statement
@@ -185,6 +189,9 @@ abstract class CompoundNestedStatementIteratorExec(
                 case leaveStatement: LeaveStatementExec =>
                   handleLeaveStatement(leaveStatement)
                   leaveStatement
+                case iterateStatement: IterateStatementExec =>
+                  handleIterateStatement(iterateStatement)
+                  iterateStatement
                 case other => other
               }
             } else {
@@ -206,18 +213,35 @@ abstract class CompoundNestedStatementIteratorExec(
     stopIteration = false
   }
 
-  /** Actions to do when LEAVE statement is encountered to stop the execution of this compound. */
+  /** Actions to do when LEAVE statement is encountered, to stop the execution of this compound. */
   private def handleLeaveStatement(leaveStatement: LeaveStatementExec): Unit = {
     if (!leaveStatement.hasBeenMatched) {
       // Stop the iteration.
       stopIteration = true
 
       // TODO: Variable cleanup (once we add SQL script execution logic).
+      // TODO: Add interpreter tests as well.
 
       // Check if label has been matched.
       leaveStatement.hasBeenMatched = label.isDefined && label.get.equals(leaveStatement.label)
     }
   }
+
+  /**
+   * Actions to do when ITERATE statement is encountered, to stop the execution of this compound.
+   */
+  private def handleIterateStatement(iterateStatement: IterateStatementExec): Unit = {
+    if (!iterateStatement.hasBeenMatched) {
+      // Stop the iteration.
+      stopIteration = true
+
+      // TODO: Variable cleanup (once we add SQL script execution logic).
+      // TODO: Add interpreter tests as well.
+
+      // No need to check if label has been matched, since ITERATE statement is already
+      //   not allowed in CompoundBody.
+    }
+  }
 }
 
 /**

sql/core/src/test/scala/org/apache/spark/sql/scripting/SqlScriptingInterpreterSuite.scala

@@ -678,4 +678,38 @@ class SqlScriptingInterpreterSuite extends QueryTest with SharedSparkSession {
     )
     verifySqlScriptResult(sqlScriptText, expected)
   }
+
+  test("nested compounds in loop - leave in inner compound") {
+    val sqlScriptText =
+      """
+        |BEGIN
+        |  DECLARE x INT;
+        |  SET x = 0;
+        |  lbl: WHILE x < 2 DO
+        |    SET x = x + 1;
+        |    BEGIN
+        |      SELECT 1;
+        |      lbl2: BEGIN
+        |        SELECT 2;
+        |        LEAVE lbl2;
+        |        SELECT 3;
+        |      END;
+        |    END;
+        |  END WHILE;
+        |  SELECT x;
+        |END""".stripMargin
+    val expected = Seq(
+      Seq.empty[Row], // declare
+      Seq.empty[Row], // set x = 0
+      Seq.empty[Row], // set x = 1
+      Seq(Row(1)), // select 1
+      Seq(Row(2)), // select 2
+      Seq.empty[Row], // set x = 2
+      Seq(Row(1)), // select 1
+      Seq(Row(2)), // select 2
+      Seq(Row(2)), // select x
+      Seq.empty[Row] // drop
+    )
+    verifySqlScriptResult(sqlScriptText, expected)
+  }
 }