[SPARK-51873][ML] For OneVsRest algorithm, allow using save / load to replace cache

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

For OneVsRest algorithm, allow using save / load to replace cache

### Why are the changes needed?

Dataframe persisting is not well supported in certain cases, so we need a replacement.

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

No.

### How was this patch tested?

Unit tests.

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

No.

Closes #50672 from WeichenXu123/one-vs-rest-cache.

Authored-by: Weichen Xu <[email protected]>
Signed-off-by: Weichen Xu <[email protected]>

Author: WeichenXu123

python/pyspark/ml/classification.py

@@ -89,6 +89,7 @@
     try_remote_read,
     try_remote_write,
     try_remote_attribute_relation,
+    _cache_spark_dataset,
 )
 from pyspark.ml.wrapper import JavaParams, JavaPredictor, JavaPredictionModel, JavaWrapper
 from pyspark.ml.common import inherit_doc
@@ -3603,46 +3604,47 @@ def _fit(self, dataset: DataFrame) -> "OneVsRestModel":
 
         # persist if underlying dataset is not persistent.
         handlePersistence = dataset.storageLevel == StorageLevel(False, False, False, False)
-        if handlePersistence:
-            multiclassLabeled.persist(StorageLevel.MEMORY_AND_DISK)
 
-        def _oneClassFitTasks(numClasses: int) -> List[Callable[[], Tuple[int, CM]]]:
-            indices = iter(range(numClasses))
-
-            def trainSingleClass() -> Tuple[int, CM]:
-                index = next(indices)
-
-                binaryLabelCol = "mc2b$" + str(index)
-                trainingDataset = multiclassLabeled.withColumn(
-                    binaryLabelCol,
-                    F.when(multiclassLabeled[labelCol] == float(index), 1.0).otherwise(0.0),
-                )
-                paramMap = dict(
-                    [
-                        (classifier.labelCol, binaryLabelCol),
-                        (classifier.featuresCol, featuresCol),
-                        (classifier.predictionCol, predictionCol),
-                    ]
-                )
-                if weightCol:
-                    paramMap[cast(HasWeightCol, classifier).weightCol] = weightCol
-                return index, classifier.fit(trainingDataset, paramMap)
-
-            return [trainSingleClass] * numClasses
-
-        tasks = map(
-            inheritable_thread_target(dataset.sparkSession),
-            _oneClassFitTasks(numClasses),
-        )
-        pool = ThreadPool(processes=min(self.getParallelism(), numClasses))
-
-        subModels = [None] * numClasses
-        for j, subModel in pool.imap_unordered(lambda f: f(), tasks):
-            assert subModels is not None
-            subModels[j] = subModel
+        with _cache_spark_dataset(
+            multiclassLabeled,
+            storageLevel=StorageLevel.MEMORY_AND_DISK,
+            enable=handlePersistence,
+        ) as multiclassLabeled:
+
+            def _oneClassFitTasks(numClasses: int) -> List[Callable[[], Tuple[int, CM]]]:
+                indices = iter(range(numClasses))
+
+                def trainSingleClass() -> Tuple[int, CM]:
+                    index = next(indices)
+
+                    binaryLabelCol = "mc2b$" + str(index)
+                    trainingDataset = multiclassLabeled.withColumn(
+                        binaryLabelCol,
+                        F.when(multiclassLabeled[labelCol] == float(index), 1.0).otherwise(0.0),
+                    )
+                    paramMap = dict(
+                        [
+                            (classifier.labelCol, binaryLabelCol),
+                            (classifier.featuresCol, featuresCol),
+                            (classifier.predictionCol, predictionCol),
+                        ]
+                    )
+                    if weightCol:
+                        paramMap[cast(HasWeightCol, classifier).weightCol] = weightCol
+                    return index, classifier.fit(trainingDataset, paramMap)
+
+                return [trainSingleClass] * numClasses
+
+            tasks = map(
+                inheritable_thread_target(dataset.sparkSession),
+                _oneClassFitTasks(numClasses),
+            )
+            pool = ThreadPool(processes=min(self.getParallelism(), numClasses))
 
-        if handlePersistence:
-            multiclassLabeled.unpersist()
+            subModels = [None] * numClasses
+            for j, subModel in pool.imap_unordered(lambda f: f(), tasks):
+                assert subModels is not None
+                subModels[j] = subModel
 
         return self._copyValues(OneVsRestModel(models=cast(List[ClassificationModel], subModels)))
 
@@ -3868,32 +3870,31 @@ def _transform(self, dataset: DataFrame) -> DataFrame:
 
         # persist if underlying dataset is not persistent.
         handlePersistence = dataset.storageLevel == StorageLevel(False, False, False, False)
-        if handlePersistence:
-            newDataset.persist(StorageLevel.MEMORY_AND_DISK)
-
-        # update the accumulator column with the result of prediction of models
-        aggregatedDataset = newDataset
-        for index, model in enumerate(self.models):
-            rawPredictionCol = self.getRawPredictionCol()
-
-            columns = origCols + [rawPredictionCol, accColName]
-
-            # add temporary column to store intermediate scores and update
-            tmpColName = "mbc$tmp" + str(uuid.uuid4())
-            transformedDataset = model.transform(aggregatedDataset).select(*columns)
-            updatedDataset = transformedDataset.withColumn(
-                tmpColName,
-                F.array_append(accColName, SF.vector_get(F.col(rawPredictionCol), F.lit(1))),
-            )
-            newColumns = origCols + [tmpColName]
-
-            # switch out the intermediate column with the accumulator column
-            aggregatedDataset = updatedDataset.select(*newColumns).withColumnRenamed(
-                tmpColName, accColName
-            )
+        with _cache_spark_dataset(
+            newDataset,
+            storageLevel=StorageLevel.MEMORY_AND_DISK,
+            enable=handlePersistence,
+        ) as newDataset:
+            # update the accumulator column with the result of prediction of models
+            aggregatedDataset = newDataset
+            for index, model in enumerate(self.models):
+                rawPredictionCol = self.getRawPredictionCol()
+
+                columns = origCols + [rawPredictionCol, accColName]
+
+                # add temporary column to store intermediate scores and update
+                tmpColName = "mbc$tmp" + str(uuid.uuid4())
+                transformedDataset = model.transform(aggregatedDataset).select(*columns)
+                updatedDataset = transformedDataset.withColumn(
+                    tmpColName,
+                    F.array_append(accColName, SF.vector_get(F.col(rawPredictionCol), F.lit(1))),
+                )
+                newColumns = origCols + [tmpColName]
 
-        if handlePersistence:
-            newDataset.unpersist()
+                # switch out the intermediate column with the accumulator column
+                aggregatedDataset = updatedDataset.select(*newColumns).withColumnRenamed(
+                    tmpColName, accColName
+                )
 
         if self.getRawPredictionCol():
             aggregatedDataset = aggregatedDataset.withColumn(

python/pyspark/ml/tests/test_algorithms.py

@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-
+import os
 from shutil import rmtree
 import tempfile
 import unittest
@@ -154,6 +154,28 @@ def test_support_for_weightCol(self):
         ovr2 = OneVsRest(classifier=dt, weightCol="weight")
         self.assertIsNotNone(ovr2.fit(df))
 
+    def test_tmp_dfs_cache(self):
+        from pyspark.ml.util import _SPARKML_TEMP_DFS_PATH
+
+        with tempfile.TemporaryDirectory(prefix="ml_tmp_dir") as d:
+            os.environ[_SPARKML_TEMP_DFS_PATH] = d
+            try:
+                df = self.spark.createDataFrame(
+                    [
+                        (0.0, Vectors.dense(1.0, 0.8)),
+                        (1.0, Vectors.sparse(2, [], [])),
+                        (2.0, Vectors.dense(0.5, 0.5)),
+                    ],
+                    ["label", "features"],
+                )
+                lr = LogisticRegression(maxIter=5, regParam=0.01)
+                ovr = OneVsRest(classifier=lr, parallelism=1)
+                model = ovr.fit(df)
+                model.transform(df)
+                assert len(os.listdir(d)) == 0
+            finally:
+                os.environ.pop(_SPARKML_TEMP_DFS_PATH, None)
+
 
 class KMeansTests(SparkSessionTestCase):
     def test_kmeans_cosine_distance(self):

python/pyspark/ml/tests/test_tuning.py

@@ -73,12 +73,15 @@ def test_train_validation_split(self):
 
         with tempfile.TemporaryDirectory(prefix="ml_tmp_dir") as d:
             os.environ[_SPARKML_TEMP_DFS_PATH] = d
-            tvs_model2 = tvs.fit(dataset)
-            assert len(os.listdir(d)) == 0
-            self.assertTrue(np.isclose(tvs_model2.validationMetrics[0], 0.5, atol=1e-4))
-            self.assertTrue(
-                np.isclose(tvs_model2.validationMetrics[1], 0.8857142857142857, atol=1e-4)
-            )
+            try:
+                tvs_model2 = tvs.fit(dataset)
+                assert len(os.listdir(d)) == 0
+                self.assertTrue(np.isclose(tvs_model2.validationMetrics[0], 0.5, atol=1e-4))
+                self.assertTrue(
+                    np.isclose(tvs_model2.validationMetrics[1], 0.8857142857142857, atol=1e-4)
+                )
+            finally:
+                os.environ.pop(_SPARKML_TEMP_DFS_PATH, None)
 
         # save & load
         with tempfile.TemporaryDirectory(prefix="train_validation_split") as d:
@@ -131,9 +134,12 @@ def test_cross_validator(self):
 
         with tempfile.TemporaryDirectory(prefix="ml_tmp_dir") as d:
             os.environ[_SPARKML_TEMP_DFS_PATH] = d
-            model2 = cv.fit(dataset)
-            assert len(os.listdir(d)) == 0
-            self.assertTrue(np.isclose(model2.avgMetrics[0], 0.5, atol=1e-4))
+            try:
+                model2 = cv.fit(dataset)
+                assert len(os.listdir(d)) == 0
+                self.assertTrue(np.isclose(model2.avgMetrics[0], 0.5, atol=1e-4))
+            finally:
+                os.environ.pop(_SPARKML_TEMP_DFS_PATH, None)
 
         output = model.transform(dataset)
         self.assertEqual(

python/pyspark/ml/util.py

@@ -41,6 +41,7 @@
 from pyspark.ml.common import inherit_doc
 from pyspark.sql import SparkSession
 from pyspark.sql.utils import is_remote
+from pyspark.storagelevel import StorageLevel
 from pyspark.util import VersionUtils
 
 if TYPE_CHECKING:
@@ -1138,7 +1139,15 @@ def _remove_dfs_dir(path: str, spark_session: "SparkSession") -> None:
 
 
 @contextmanager
-def _cache_spark_dataset(dataset: "DataFrame") -> Iterator[Any]:
+def _cache_spark_dataset(
+    dataset: "DataFrame",
+    storageLevel: "StorageLevel" = StorageLevel.MEMORY_AND_DISK_DESER,
+    enable: bool = True,
+) -> Iterator[Any]:
+    if not enable:
+        yield dataset
+        return
+
     spark_session = dataset._session
     tmp_dfs_path = os.environ.get(_SPARKML_TEMP_DFS_PATH)
 
@@ -1150,7 +1159,7 @@ def _cache_spark_dataset(dataset: "DataFrame") -> Iterator[Any]:
         finally:
             _remove_dfs_dir(tmp_cache_path, spark_session)
     else:
-        dataset.cache()
+        dataset.persist(storageLevel)
         try:
             yield dataset
         finally: