Github user BryanCutler commented on a diff in the pull request:
https://github.com/apache/spark/pull/21007#discussion_r180584207
--- Diff: python/pyspark/sql/tests.py ---
@@ -3062,6 +3062,73 @@ def
test_sparksession_with_stopped_sparkcontext(self):
sc.stop()
+class SQLTests3(unittest.TestCase):
+ # These tests are separate because it uses
'spark.sql.queryExecutionListeners' which is
+ # static and immutable. This can't be set or unset.
+
+ @classmethod
+ def setUpClass(cls):
+ import glob
+ from pyspark.find_spark_home import _find_spark_home
+
+ SPARK_HOME = _find_spark_home()
+ filename_pattern = (
+ "sql/core/target/scala-*/test-classes/org/apache/spark/sql/"
+ "TestQueryExecutionListener.class")
+ if not glob.glob(os.path.join(SPARK_HOME, filename_pattern)):
+ raise unittest.SkipTest(
+ "'org.apache.spark.sql.TestQueryExecutionListener' is not "
+ "available. Will skip the related tests.")
+
+ # Note that 'spark.sql.queryExecutionListeners' is a static
immutable configuration.
+ cls.spark = SparkSession.builder \
+ .master("local[4]") \
+ .appName(cls.__name__) \
+ .config(
+ "spark.sql.queryExecutionListeners",
+ "org.apache.spark.sql.TestQueryExecutionListener") \
+ .getOrCreate()
+
+ @classmethod
+ def tearDownClass(cls):
+ cls.spark.stop()
+
+ def tearDown(self):
+ self.spark._jvm.OnSuccessCall.clear()
+
+ def test_query_execution_listener_on_collect(self):
+ self.assertFalse(
+ self.spark._jvm.OnSuccessCall.isCalled(),
+ "The callback from the query execution listener should not be
called before 'collect'")
+ self.spark.sql("SELECT * FROM range(1)").collect()
+ self.assertTrue(
+ self.spark._jvm.OnSuccessCall.isCalled(),
+ "The callback from the query execution listener should be
called after 'collect'")
+
+ @unittest.skipIf(
+ not _have_pandas or not _have_pyarrow,
+ _pandas_requirement_message or _pyarrow_requirement_message)
+ def test_query_execution_listener_on_collect_with_arrow(self):
+ # Here, it deplicates codes in ReusedSQLTestCase.sql_conf context
manager.
--- End diff --
I think it would be fine refactor `sql_conf` a little to use it here, it
makes things much clearer
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
