[SPARK-51843][PYTHON][ML][TESTS] Avoid per-test classic session start & stop

python/pyspark/ml/tests/connect/test_legacy_mode_classification.py

@@ -22,9 +22,9 @@
 import numpy as np
 
 from pyspark.util import is_remote_only
-from pyspark.sql import SparkSession
 from pyspark.testing.connectutils import should_test_connect, connect_requirement_message
 from pyspark.testing.utils import have_torch, torch_requirement_message
+from pyspark.testing.sqlutils import ReusedSQLTestCase
 
 if should_test_connect:
     from pyspark.ml.connect.classification import (
@@ -231,12 +231,10 @@ def test_save_load(self):
     or torch_requirement_message
     or "pyspark-connect cannot test classic Spark",
 )
-class ClassificationTests(ClassificationTestsMixin, unittest.TestCase):
-    def setUp(self) -> None:
-        self.spark = SparkSession.builder.master("local[2]").getOrCreate()
-
-    def tearDown(self) -> None:
-        self.spark.stop()
+class ClassificationTests(ClassificationTestsMixin, ReusedSQLTestCase):
+    @classmethod
+    def master(cls):
+        return "local[2]"
 
 
 if __name__ == "__main__":

python/pyspark/ml/tests/connect/test_legacy_mode_evaluation.py

@@ -21,9 +21,9 @@
 import numpy as np
 
 from pyspark.util import is_remote_only
-from pyspark.sql import SparkSession
 from pyspark.testing.connectutils import should_test_connect, connect_requirement_message
 from pyspark.testing.utils import have_torcheval, torcheval_requirement_message
+from pyspark.testing.sqlutils import ReusedSQLTestCase
 
 if should_test_connect:
     from pyspark.ml.connect.evaluation import (
@@ -178,12 +178,10 @@ def test_multiclass_classifier_evaluator(self):
     or torcheval_requirement_message
     or "pyspark-connect cannot test classic Spark",
 )
-class EvaluationTests(EvaluationTestsMixin, unittest.TestCase):
-    def setUp(self) -> None:
-        self.spark = SparkSession.builder.master("local[2]").getOrCreate()
-
-    def tearDown(self) -> None:
-        self.spark.stop()
+class EvaluationTests(EvaluationTestsMixin, ReusedSQLTestCase):
+    @classmethod
+    def master(cls):
+        return "local[2]"
 
 
 if __name__ == "__main__":

python/pyspark/ml/tests/connect/test_legacy_mode_feature.py

@@ -24,9 +24,9 @@
 import numpy as np
 
 from pyspark.util import is_remote_only
-from pyspark.sql import SparkSession
 from pyspark.testing.connectutils import should_test_connect, connect_requirement_message
 from pyspark.testing.utils import have_torch, torch_requirement_message
+from pyspark.testing.sqlutils import ReusedSQLTestCase
 
 if should_test_connect:
     from pyspark.ml.connect.feature import (
@@ -201,12 +201,10 @@ def test_array_assembler(self):
     or torch_requirement_message
     or "pyspark-connect cannot test classic Spark",
 )
-class FeatureTests(FeatureTestsMixin, unittest.TestCase):
-    def setUp(self) -> None:
-        self.spark = SparkSession.builder.master("local[2]").getOrCreate()
-
-    def tearDown(self) -> None:
-        self.spark.stop()
+class FeatureTests(FeatureTestsMixin, ReusedSQLTestCase):
+    @classmethod
+    def master(cls):
+        return "local[2]"
 
 
 if __name__ == "__main__":

python/pyspark/ml/tests/connect/test_legacy_mode_pipeline.py

@@ -22,9 +22,9 @@
 import numpy as np
 
 from pyspark.util import is_remote_only
-from pyspark.sql import SparkSession
 from pyspark.testing.connectutils import should_test_connect, connect_requirement_message
 from pyspark.testing.utils import have_torch, torch_requirement_message
+from pyspark.testing.sqlutils import ReusedSQLTestCase
 
 if should_test_connect:
     from pyspark.ml.connect.feature import StandardScaler
@@ -175,12 +175,10 @@ def test_pipeline_copy():
     or torch_requirement_message
     or "pyspark-connect cannot test classic Spark",
 )
-class PipelineTests(PipelineTestsMixin, unittest.TestCase):
-    def setUp(self) -> None:
-        self.spark = SparkSession.builder.master("local[2]").getOrCreate()
-
-    def tearDown(self) -> None:
-        self.spark.stop()
+class PipelineTests(PipelineTestsMixin, ReusedSQLTestCase):
+    @classmethod
+    def master(cls):
+        return "local[2]"
 
 
 if __name__ == "__main__":

python/pyspark/ml/tests/connect/test_legacy_mode_summarizer.py

@@ -21,8 +21,8 @@
 import numpy as np
 
 from pyspark.util import is_remote_only
-from pyspark.sql import SparkSession
 from pyspark.testing.connectutils import should_test_connect, connect_requirement_message
+from pyspark.testing.sqlutils import ReusedSQLTestCase
 
 if should_test_connect:
     from pyspark.ml.connect.summarizer import summarize_dataframe
@@ -67,12 +67,10 @@ def assert_dict_allclose(dict1, dict2):
     not should_test_connect or is_remote_only(),
     connect_requirement_message or "pyspark-connect cannot test classic Spark",
 )
-class SummarizerTests(SummarizerTestsMixin, unittest.TestCase):
-    def setUp(self) -> None:
-        self.spark = SparkSession.builder.master("local[2]").getOrCreate()
-
-    def tearDown(self) -> None:
-        self.spark.stop()
+class SummarizerTests(SummarizerTestsMixin, ReusedSQLTestCase):
+    @classmethod
+    def master(cls):
+        return "local[2]"
 
 
 if __name__ == "__main__":

python/pyspark/ml/tests/connect/test_legacy_mode_tuning.py

@@ -25,7 +25,6 @@
 from pyspark.util import is_remote_only
 from pyspark.ml.param import Param, Params
 from pyspark.ml.tuning import ParamGridBuilder
-from pyspark.sql import SparkSession
 from pyspark.sql.functions import rand
 from pyspark.testing.connectutils import should_test_connect, connect_requirement_message
 from pyspark.testing.utils import (
@@ -36,6 +35,7 @@
     have_torcheval,
     torcheval_requirement_message,
 )
+from pyspark.testing.sqlutils import ReusedSQLTestCase
 
 if should_test_connect:
     import pandas as pd
@@ -294,12 +294,10 @@ def test_crossvalidator_with_fold_col(self):
     or torcheval_requirement_message
     or "pyspark-connect cannot test classic Spark",
 )
-class CrossValidatorTests(CrossValidatorTestsMixin, unittest.TestCase):
-    def setUp(self) -> None:
-        self.spark = SparkSession.builder.master("local[2]").getOrCreate()
-
-    def tearDown(self) -> None:
-        self.spark.stop()
+class CrossValidatorTests(CrossValidatorTestsMixin, ReusedSQLTestCase):
+    @classmethod
+    def master(cls):
+        return "local[2]"
 
 
 if __name__ == "__main__":

python/pyspark/testing/utils.py

@@ -246,7 +246,11 @@ def conf(cls):
     def setUpClass(cls):
         from pyspark import SparkContext
 
-        cls.sc = SparkContext("local[4]", cls.__name__, conf=cls.conf())
+        cls.sc = SparkContext(cls.master(), cls.__name__, conf=cls.conf())
+
+    @classmethod
+    def master(cls):
+        return "local[4]"
 
     @classmethod
     def tearDownClass(cls):