[SPARK-48843] Adding more tests that would have failed had the fix apache#47271 not been sumitted. Prevents regression. Hardens previous tests to catch infinite loop caused by head() and first() in Connect. Adds a test that would have caused infinite loop with BindParameters in Classic.

Author: nemanjapetr-db

python/pyspark/sql/tests/connect/test_connect_basic.py

@@ -22,6 +22,7 @@
 import tempfile
 import io
 from contextlib import redirect_stdout
+import datetime
 
 from pyspark.util import is_remote_only
 from pyspark.errors import PySparkTypeError, PySparkValueError
@@ -637,6 +638,16 @@ def test_namedargs_with_global_limit(self):
         df2 = self.spark.sql(sqlText, args={"val": 1})
         self.assert_eq(df.toPandas(), df2.toPandas())
 
+        self.assert_eq(df.first()[0], datetime.datetime(2022, 12, 25, 10, 30))
+        self.assert_eq(df.first().date, datetime.datetime(2022, 12, 25, 10, 30))
+        self.assert_eq(df.first()[1], 1)
+        self.assert_eq(df.first().val, 1)
+
+        self.assert_eq(df.head()[0], datetime.datetime(2022, 12, 25, 10, 30))
+        self.assert_eq(df.head().date, datetime.datetime(2022, 12, 25, 10, 30))
+        self.assert_eq(df.head()[1], 1)
+        self.assert_eq(df.head().val, 1)
+
     def test_sql_with_pos_args(self):
         sqlText = "SELECT *, element_at(?, 1) FROM range(10) WHERE id > ?"
         df = self.connect.sql(sqlText, args=[CF.array(CF.lit(1)), 7])

sql/core/src/test/scala/org/apache/spark/sql/ParametersSuite.scala

@@ -22,6 +22,7 @@ import java.time.{Instant, LocalDate, LocalDateTime, ZoneId}
 import org.apache.spark.sql.catalyst.expressions.Literal
 import org.apache.spark.sql.catalyst.parser.ParseException
 import org.apache.spark.sql.catalyst.plans.PlanTest
+import org.apache.spark.sql.catalyst.plans.logical.Limit
 import org.apache.spark.sql.functions.{array, call_function, lit, map, map_from_arrays, map_from_entries, str_to_map, struct}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SharedSparkSession
@@ -633,4 +634,11 @@ class ParametersSuite extends QueryTest with SharedSparkSession with PlanTest {
         |""".stripMargin)
     checkAnswer(query, Row("ABC"))
   }
+
+  test("SPARK-48843: Prevent infinite loop with BindParameters") {
+    val df = sql("EXECUTE IMMEDIATE 'SELECT SUM(c1) num_sum FROM VALUES (?), (?) AS t(c1) ' USING 5, 6;")
+    val analyzedPlan = Limit(Literal.create(100), df.queryExecution.initialParsedPlan)
+    spark.sessionState.analyzer.executeAndCheck(analyzedPlan, df.queryExecution.tracker)
+    checkAnswer(df, Row(11))
+  }
 }