[SPARK-50941][ML][PYTHON][CONNECT] add supports for TrainValidationSplit

### What changes were proposed in this pull request?

This PR adds support for TrainValidationSplit and TrainValidationSplitModel on Connect

### Why are the changes needed?
new feature parity

### Does this PR introduce _any_ user-facing change?
Yes

### How was this patch tested?
The CI passes

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

Closes #49688 from wbo4958/train_validation_split.

Lead-authored-by: Bobby Wang <wbo4958@gmail.com>
Co-authored-by: Ruifeng Zheng <ruifengz@foxmail.com>
Signed-off-by: Ruifeng Zheng <ruifengz@apache.org>
(cherry picked from commit 9d0e888)
Signed-off-by: Ruifeng Zheng <ruifengz@apache.org>

Author: wbo4958

python/pyspark/ml/connect/readwrite.py

@@ -19,7 +19,12 @@
 
 import pyspark.sql.connect.proto as pb2
 from pyspark.ml.connect.serialize import serialize_ml_params, deserialize, deserialize_param
-from pyspark.ml.tuning import CrossValidatorModelWriter, CrossValidatorModel
+from pyspark.ml.tuning import (
+    CrossValidatorModelWriter,
+    CrossValidatorModel,
+    TrainValidationSplitModel,
+    TrainValidationSplitModelWriter,
+)
 from pyspark.ml.util import MLWriter, MLReader, RL
 from pyspark.ml.wrapper import JavaWrapper
 
@@ -42,6 +47,19 @@ def __init__(
         self.session(session)  # type: ignore[arg-type]
 
 
+class RemoteTrainValidationSplitModelWriter(TrainValidationSplitModelWriter):
+    def __init__(
+        self,
+        instance: "TrainValidationSplitModel",
+        optionMap: Dict[str, Any] = {},
+        session: Optional["SparkSession"] = None,
+    ):
+        super(RemoteTrainValidationSplitModelWriter, self).__init__(instance)
+        self.instance = instance
+        self.optionMap = optionMap
+        self.session(session)  # type: ignore[arg-type]
+
+
 class RemoteMLWriter(MLWriter):
     def __init__(self, instance: "JavaMLWritable") -> None:
         super().__init__()
@@ -76,7 +94,7 @@ def saveInstance(
         from pyspark.ml.wrapper import JavaModel, JavaEstimator, JavaTransformer
         from pyspark.ml.evaluation import JavaEvaluator
         from pyspark.ml.pipeline import Pipeline, PipelineModel
-        from pyspark.ml.tuning import CrossValidator
+        from pyspark.ml.tuning import CrossValidator, TrainValidationSplit
 
         # Spark Connect ML is built on scala Spark.ML, that means we're only
         # supporting JavaModel or JavaEstimator or JavaEvaluator
@@ -155,6 +173,20 @@ def saveInstance(
                 warnings.warn("Overwrite doesn't take effect for CrossValidatorModel")
             cvm_writer = RemoteCrossValidatorModelWriter(instance, optionMap, session)
             cvm_writer.save(path)
+        elif isinstance(instance, TrainValidationSplit):
+            from pyspark.ml.tuning import TrainValidationSplitWriter
+
+            if shouldOverwrite:
+                # TODO(SPARK-50954): Support client side model path overwrite
+                warnings.warn("Overwrite doesn't take effect for TrainValidationSplit")
+            tvs_writer = TrainValidationSplitWriter(instance)
+            tvs_writer.save(path)
+        elif isinstance(instance, TrainValidationSplitModel):
+            if shouldOverwrite:
+                # TODO(SPARK-50954): Support client side model path overwrite
+                warnings.warn("Overwrite doesn't take effect for TrainValidationSplitModel")
+            tvsm_writer = RemoteTrainValidationSplitModelWriter(instance, optionMap, session)
+            tvsm_writer.save(path)
         else:
             raise NotImplementedError(f"Unsupported write for {instance.__class__}")
 
@@ -182,7 +214,7 @@ def loadInstance(
         from pyspark.ml.wrapper import JavaModel, JavaEstimator, JavaTransformer
         from pyspark.ml.evaluation import JavaEvaluator
         from pyspark.ml.pipeline import Pipeline, PipelineModel
-        from pyspark.ml.tuning import CrossValidator
+        from pyspark.ml.tuning import CrossValidator, TrainValidationSplit
 
         if (
             issubclass(clazz, JavaModel)
@@ -261,5 +293,19 @@ def _get_class() -> Type[RL]:
             cvm_reader.session(session)
             return cvm_reader.load(path)
 
+        elif issubclass(clazz, TrainValidationSplit):
+            from pyspark.ml.tuning import TrainValidationSplitReader
+
+            tvs_reader = TrainValidationSplitReader(TrainValidationSplit)
+            tvs_reader.session(session)
+            return tvs_reader.load(path)
+
+        elif issubclass(clazz, TrainValidationSplitModel):
+            from pyspark.ml.tuning import TrainValidationSplitModelReader
+
+            tvs_reader = TrainValidationSplitModelReader(TrainValidationSplitModel)
+            tvs_reader.session(session)
+            return tvs_reader.load(path)
+
         else:
             raise RuntimeError(f"Unsupported read for {clazz}")

python/pyspark/ml/tests/test_tuning.py

@@ -24,11 +24,69 @@
 from pyspark.ml.evaluation import BinaryClassificationEvaluator
 from pyspark.ml.linalg import Vectors
 from pyspark.ml.classification import LogisticRegression
-from pyspark.ml.tuning import ParamGridBuilder, CrossValidator, CrossValidatorModel
+from pyspark.ml.tuning import (
+    ParamGridBuilder,
+    CrossValidator,
+    CrossValidatorModel,
+    TrainValidationSplit,
+    TrainValidationSplitModel,
+)
 from pyspark.testing.sqlutils import ReusedSQLTestCase
 
 
 class TuningTestsMixin:
+    def test_train_validation_split(self):
+        dataset = self.spark.createDataFrame(
+            [
+                (Vectors.dense([0.0]), 0.0),
+                (Vectors.dense([0.4]), 1.0),
+                (Vectors.dense([0.5]), 0.0),
+                (Vectors.dense([0.6]), 1.0),
+                (Vectors.dense([1.0]), 1.0),
+            ]
+            * 10,  # Repeat the data 10 times
+            ["features", "label"],
+        ).repartition(1)
+
+        lr = LogisticRegression()
+        grid = ParamGridBuilder().addGrid(lr.maxIter, [0, 1]).build()
+        evaluator = BinaryClassificationEvaluator()
+
+        tvs = TrainValidationSplit(
+            estimator=lr, estimatorParamMaps=grid, evaluator=evaluator, parallelism=1, seed=42
+        )
+        self.assertEqual(tvs.getEstimator(), lr)
+        self.assertEqual(tvs.getEvaluator(), evaluator)
+        self.assertEqual(tvs.getParallelism(), 1)
+        self.assertEqual(tvs.getEstimatorParamMaps(), grid)
+
+        tvs_model = tvs.fit(dataset)
+
+        # Access the train ratio
+        self.assertEqual(tvs_model.getTrainRatio(), 0.75)
+        print("----------- ", tvs_model.validationMetrics)
+        self.assertTrue(np.isclose(tvs_model.validationMetrics[0], 0.5, atol=1e-4))
+        self.assertTrue(np.isclose(tvs_model.validationMetrics[1], 0.8857142857142857, atol=1e-4))
+
+        evaluation_score = evaluator.evaluate(tvs_model.transform(dataset))
+        self.assertTrue(np.isclose(evaluation_score, 0.8333333333333333, atol=1e-4))
+
+        # save & load
+        with tempfile.TemporaryDirectory(prefix="train_validation_split") as d:
+            path1 = os.path.join(d, "cv")
+            tvs.write().save(path1)
+            tvs2 = TrainValidationSplit.load(path1)
+            self.assertEqual(str(tvs), str(tvs2))
+            self.assertEqual(str(tvs.getEstimator()), str(tvs2.getEstimator()))
+            self.assertEqual(str(tvs.getEvaluator()), str(tvs2.getEvaluator()))
+
+            path2 = os.path.join(d, "cv_model")
+            tvs_model.write().save(path2)
+            model2 = TrainValidationSplitModel.load(path2)
+            self.assertEqual(str(tvs_model), str(model2))
+            self.assertEqual(str(tvs_model.getEstimator()), str(model2.getEstimator()))
+            self.assertEqual(str(tvs_model.getEvaluator()), str(model2.getEvaluator()))
+
     def test_cross_validator(self):
         dataset = self.spark.createDataFrame(
             [

python/pyspark/ml/tuning.py

@@ -1186,12 +1186,12 @@ def __init__(self, cls: Type["TrainValidationSplit"]):
         self.cls = cls
 
     def load(self, path: str) -> "TrainValidationSplit":
-        metadata = DefaultParamsReader.loadMetadata(path, self.sc)
+        metadata = DefaultParamsReader.loadMetadata(path, self.sparkSession)
         if not DefaultParamsReader.isPythonParamsInstance(metadata):
             return JavaMLReader(self.cls).load(path)  # type: ignore[arg-type]
         else:
             metadata, estimator, evaluator, estimatorParamMaps = _ValidatorSharedReadWrite.load(
-                path, self.sc, metadata
+                path, self.sparkSession, metadata
             )
             tvs = TrainValidationSplit(
                 estimator=estimator, estimatorParamMaps=estimatorParamMaps, evaluator=evaluator
@@ -1209,7 +1209,7 @@ def __init__(self, instance: "TrainValidationSplit"):
 
     def saveImpl(self, path: str) -> None:
         _ValidatorSharedReadWrite.validateParams(self.instance)
-        _ValidatorSharedReadWrite.saveImpl(path, self.instance, self.sc)
+        _ValidatorSharedReadWrite.saveImpl(path, self.instance, self.sparkSession)
 
 
 @inherit_doc
@@ -1219,15 +1219,17 @@ def __init__(self, cls: Type["TrainValidationSplitModel"]):
         self.cls = cls
 
     def load(self, path: str) -> "TrainValidationSplitModel":
-        metadata = DefaultParamsReader.loadMetadata(path, self.sc)
+        metadata = DefaultParamsReader.loadMetadata(path, self.sparkSession)
         if not DefaultParamsReader.isPythonParamsInstance(metadata):
             return JavaMLReader(self.cls).load(path)  # type: ignore[arg-type]
         else:
             metadata, estimator, evaluator, estimatorParamMaps = _ValidatorSharedReadWrite.load(
-                path, self.sc, metadata
+                path, self.sparkSession, metadata
             )
             bestModelPath = os.path.join(path, "bestModel")
-            bestModel: Model = DefaultParamsReader.loadParamsInstance(bestModelPath, self.sc)
+            bestModel: Model = DefaultParamsReader.loadParamsInstance(
+                bestModelPath, self.sparkSession
+            )
             validationMetrics = metadata["validationMetrics"]
             persistSubModels = ("persistSubModels" in metadata) and metadata["persistSubModels"]
 
@@ -1236,7 +1238,7 @@ def load(self, path: str) -> "TrainValidationSplitModel":
                 for paramIndex in range(len(estimatorParamMaps)):
                     modelPath = os.path.join(path, "subModels", f"{paramIndex}")
                     subModels[paramIndex] = DefaultParamsReader.loadParamsInstance(
-                        modelPath, self.sc
+                        modelPath, self.sparkSession
                     )
             else:
                 subModels = None
@@ -1273,7 +1275,9 @@ def saveImpl(self, path: str) -> None:
             "validationMetrics": instance.validationMetrics,
             "persistSubModels": persistSubModels,
         }
-        _ValidatorSharedReadWrite.saveImpl(path, instance, self.sc, extraMetadata=extraMetadata)
+        _ValidatorSharedReadWrite.saveImpl(
+            path, instance, self.sparkSession, extraMetadata=extraMetadata
+        )
         bestModelPath = os.path.join(path, "bestModel")
         cast(MLWritable, instance.bestModel).save(bestModelPath)
         if persistSubModels:
@@ -1473,7 +1477,7 @@ def _fit(self, dataset: DataFrame) -> "TrainValidationSplitModel":
             subModels = [None for i in range(numModels)]
 
         tasks = map(
-            inheritable_thread_target,
+            inheritable_thread_target(dataset.sparkSession),
             _parallelFitTasks(est, train, eva, validation, epm, collectSubModelsParam),
         )
         pool = ThreadPool(processes=min(self.getParallelism(), numModels))
@@ -1529,6 +1533,7 @@ def copy(self, extra: Optional["ParamMap"] = None) -> "TrainValidationSplit":
         return newTVS
 
     @since("2.3.0")
+    @try_remote_write
     def write(self) -> MLWriter:
         """Returns an MLWriter instance for this ML instance."""
         if _ValidatorSharedReadWrite.is_java_convertible(self):
@@ -1537,6 +1542,7 @@ def write(self) -> MLWriter:
 
     @classmethod
     @since("2.3.0")
+    @try_remote_read
     def read(cls) -> TrainValidationSplitReader:
         """Returns an MLReader instance for this class."""
         return TrainValidationSplitReader(cls)
@@ -1649,6 +1655,7 @@ def copy(self, extra: Optional["ParamMap"] = None) -> "TrainValidationSplitModel
         )
 
     @since("2.3.0")
+    @try_remote_write
     def write(self) -> MLWriter:
         """Returns an MLWriter instance for this ML instance."""
         if _ValidatorSharedReadWrite.is_java_convertible(self):
@@ -1657,6 +1664,7 @@ def write(self) -> MLWriter:
 
     @classmethod
     @since("2.3.0")
+    @try_remote_read
     def read(cls) -> TrainValidationSplitModelReader:
         """Returns an MLReader instance for this class."""
         return TrainValidationSplitModelReader(cls)