This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 6313ea6f804b [SPARK-57860][ML][PYTHON] Add HasIntermediateStorageLevel
shared param and apply to KMeans
6313ea6f804b is described below
commit 6313ea6f804b981338c0da0218ca3480574cfb15
Author: Mao Li <[email protected]>
AuthorDate: Thu Jul 2 09:40:47 2026 +0800
[SPARK-57860][ML][PYTHON] Add HasIntermediateStorageLevel shared param and
apply to KMeans
### What changes were proposed in this pull request?
This is the first sub-task of
[SPARK-47103](https://issues.apache.org/jira/browse/SPARK-47103), which aims to
make the default storage level of MLlib's intermediate datasets configurable.
This PR:
1. Adds a new shared param `HasIntermediateStorageLevel` (default
`"MEMORY_AND_DISK"`, cannot be `"NONE"`) by extending the code generators on
both sides:
- Scala: `SharedParamsCodeGen.scala` -> regenerated `sharedParams.scala`
- Python: `_shared_params_code_gen.py` -> regenerated `shared.py`
2. Applies it to `KMeans` (Scala and PySpark): mixes in the trait, adds
`setIntermediateStorageLevel`, and uses `$(intermediateStorageLevel)` at the
`persist` call site instead of the hardcoded `StorageLevel.MEMORY_AND_DISK`.
3. Adds test coverage in `KMeansSuite`.
The design follows the suggestion from zhengruifeng on the earlier PR
#45182 (a per-estimator param via a shared `HasIntermediateStorageLevel` trait,
consistent with ALS's existing `intermediateStorageLevel`), rather than the
global SQL config explored there. The remaining estimators are tracked as
sibling sub-tasks under SPARK-47103.
### Why are the changes needed?
MLlib persists *intermediate* datasets internally during training (e.g.
blockified instances), with the storage level hardcoded to `MEMORY_AND_DISK`.
These datasets are created inside the algorithm and are not the user's input
`DataFrame`, so users currently have **no way** to change their storage level
-- unlike the input, which they can already cache themselves.
Making this configurable (e.g. `DISK_ONLY`) improves resilience to executor
loss: since SPARK-27677, the External Shuffle Service can serve disk-persisted
cached blocks, so disk-based intermediate storage survives executor failures.
`ALS` already exposes exactly this via `intermediateStorageLevel`; this PR
starts extending the same capability to the rest of MLlib.
### Does this PR introduce _any_ user-facing change?
Yes. `KMeans` gains a new expert param `intermediateStorageLevel` and a
`setIntermediateStorageLevel` setter.
The default is `"MEMORY_AND_DISK"`, so **behavior is unchanged unless the
user sets it**.
Before (no way to change intermediate storage level):
```python
kmeans = KMeans(k=3) # intermediate data always MEMORY_AND_DISK
```
After:
```python
kmeans = KMeans(k=3).setIntermediateStorageLevel("DISK_ONLY")
```
### How was this patch tested?
- Extended `KMeansSuite` to assert the new param's default value, that it
can be set, and that invalid values (`"NONE"` and non-existent levels) are
rejected. `KMeansSuite` passes (15/15).
- PySpark param parity is covered by the existing
`pyspark.ml.tests.test_param.test_java_params`, which checks that Python params
match their Scala counterparts.
- `dev/mima` (mllib `mimaReportBinaryIssues`) reports no binary
compatibility problems; no `MimaExcludes` entries were needed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Claude Opus 4.8)
Closes #56949 from maoli67660/SPARK-57860.
Lead-authored-by: Mao Li <[email protected]>
Co-authored-by: Mao Li <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../org/apache/spark/ml/clustering/KMeans.scala | 10 +++++++---
.../ml/param/shared/SharedParamsCodeGen.scala | 11 ++++++++++-
.../spark/ml/param/shared/sharedParams.scala | 21 ++++++++++++++++++++
.../apache/spark/ml/clustering/KMeansSuite.scala | 9 +++++++++
python/pyspark/ml/clustering.py | 9 +++++++++
python/pyspark/ml/param/_shared_params_code_gen.py | 6 ++++++
python/pyspark/ml/param/shared.py | 23 ++++++++++++++++++++++
7 files changed, 85 insertions(+), 4 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala
b/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala
index ad6ca0924064..2de3eb531f56 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/clustering/KMeans.scala
@@ -49,7 +49,7 @@ import org.apache.spark.util.VersionUtils.majorVersion
*/
private[clustering] trait KMeansParams extends Params with HasMaxIter with
HasFeaturesCol
with HasSeed with HasPredictionCol with HasTol with HasDistanceMeasure with
HasWeightCol
- with HasSolver with HasMaxBlockSizeInMB {
+ with HasSolver with HasMaxBlockSizeInMB with HasIntermediateStorageLevel {
import KMeans._
/**
@@ -434,6 +434,10 @@ class KMeans @Since("1.5.0") (
@Since("3.4.0")
def setMaxBlockSizeInMB(value: Double): this.type = set(maxBlockSizeInMB,
value)
+ /** @group expertSetParam */
+ @Since("5.0.0")
+ def setIntermediateStorageLevel(value: String): this.type =
set(intermediateStorageLevel, value)
+
@Since("2.0.0")
override def fit(dataset: Dataset[_]): KMeansModel = instrumented { instr =>
transformSchema(dataset.schema, logging = true)
@@ -441,7 +445,7 @@ class KMeans @Since("1.5.0") (
instr.logPipelineStage(this)
instr.logDataset(dataset)
instr.logParams(this, featuresCol, predictionCol, k, initMode, initSteps,
distanceMeasure,
- maxIter, seed, tol, weightCol, solver, maxBlockSizeInMB)
+ maxIter, seed, tol, weightCol, solver, maxBlockSizeInMB,
intermediateStorageLevel)
val oldModel = if (preferBlockSolver(dataset)) {
trainWithBlock(dataset, instr)
@@ -547,7 +551,7 @@ class KMeans @Since("1.5.0") (
}
val maxMemUsage = (actualBlockSizeInMB * 1024L * 1024L).ceil.toLong
val blocks = InstanceBlock.blokifyWithMaxMemUsage(instances, maxMemUsage)
- .persist(StorageLevel.MEMORY_AND_DISK)
+ .persist(StorageLevel.fromString($(intermediateStorageLevel)))
.setName(s"$uid: training blocks (blockSizeInMB=$actualBlockSizeInMB)")
val distanceFunction = getDistanceFunction
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/param/shared/SharedParamsCodeGen.scala
b/mllib/src/main/scala/org/apache/spark/ml/param/shared/SharedParamsCodeGen.scala
index c61aa14edca2..a689ab37d678 100644
---
a/mllib/src/main/scala/org/apache/spark/ml/param/shared/SharedParamsCodeGen.scala
+++
b/mllib/src/main/scala/org/apache/spark/ml/param/shared/SharedParamsCodeGen.scala
@@ -113,7 +113,13 @@ private[shared] object SharedParamsCodeGen {
"into blocks. Data is stacked within partitions. If more than
remaining data size in a " +
"partition then it is adjusted to the data size. Default 0.0
represents choosing " +
"optimal value, depends on specific algorithm. Must be >= 0.",
- Some("0.0"), isValid = "ParamValidators.gtEq(0.0)", isExpertParam =
true)
+ Some("0.0"), isValid = "ParamValidators.gtEq(0.0)", isExpertParam =
true),
+ ParamDesc[String]("intermediateStorageLevel",
+ "StorageLevel for intermediate datasets. Pass in a string
representation of " +
+ "`StorageLevel`. Cannot be 'NONE'.",
+ Some("\"MEMORY_AND_DISK\""),
+ isValid = "(s: String) => Try(StorageLevel.fromString(s)).isSuccess &&
s != \"NONE\"",
+ isExpertParam = true)
)
val code = genSharedParams(params)
@@ -248,7 +254,10 @@ private[shared] object SharedParamsCodeGen {
|
|package org.apache.spark.ml.param.shared
|
+ |import scala.util.Try
+ |
|import org.apache.spark.ml.param._
+ |import org.apache.spark.storage.StorageLevel
|
|// DO NOT MODIFY THIS FILE! It was generated by SharedParamsCodeGen.
|
diff --git
a/mllib/src/main/scala/org/apache/spark/ml/param/shared/sharedParams.scala
b/mllib/src/main/scala/org/apache/spark/ml/param/shared/sharedParams.scala
index 425bf91fd00b..0df0567e4692 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/param/shared/sharedParams.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/param/shared/sharedParams.scala
@@ -17,7 +17,10 @@
package org.apache.spark.ml.param.shared
+import scala.util.Try
+
import org.apache.spark.ml.param._
+import org.apache.spark.storage.StorageLevel
// DO NOT MODIFY THIS FILE! It was generated by SharedParamsCodeGen.
@@ -580,4 +583,22 @@ trait HasMaxBlockSizeInMB extends Params {
/** @group expertGetParam */
final def getMaxBlockSizeInMB: Double = $(maxBlockSizeInMB)
}
+
+/**
+ * Trait for shared param intermediateStorageLevel (default:
"MEMORY_AND_DISK"). This trait may be changed or
+ * removed between minor versions.
+ */
+trait HasIntermediateStorageLevel extends Params {
+
+ /**
+ * Param for StorageLevel for intermediate datasets. Pass in a string
representation of `StorageLevel`. Cannot be 'NONE'..
+ * @group expertParam
+ */
+ final val intermediateStorageLevel: Param[String] = new Param[String](this,
"intermediateStorageLevel", "StorageLevel for intermediate datasets. Pass in a
string representation of `StorageLevel`. Cannot be 'NONE'.", (s: String) =>
Try(StorageLevel.fromString(s)).isSuccess && s != "NONE")
+
+ setDefault(intermediateStorageLevel, "MEMORY_AND_DISK")
+
+ /** @group expertGetParam */
+ final def getIntermediateStorageLevel: String = $(intermediateStorageLevel)
+}
// scalastyle:on
diff --git
a/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala
b/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala
index 9e5be94a9ccb..852cce6645aa 100644
--- a/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala
+++ b/mllib/src/test/scala/org/apache/spark/ml/clustering/KMeansSuite.scala
@@ -58,6 +58,7 @@ class KMeansSuite extends MLTest with DefaultReadWriteTest
with PMMLReadWriteTes
assert(kmeans.getTol === 1e-4)
assert(kmeans.getSolver === KMeans.AUTO)
assert(kmeans.getDistanceMeasure === KMeans.EUCLIDEAN)
+ assert(kmeans.getIntermediateStorageLevel === "MEMORY_AND_DISK")
val model = kmeans.setMaxIter(1).fit(dataset)
val transformed = model.transform(dataset)
@@ -85,6 +86,7 @@ class KMeansSuite extends MLTest with DefaultReadWriteTest
with PMMLReadWriteTes
.setSeed(123)
.setTol(1e-3)
.setDistanceMeasure(KMeans.COSINE)
+ .setIntermediateStorageLevel("MEMORY_ONLY")
assert(kmeans.getK === 9)
assert(kmeans.getFeaturesCol === "test_feature")
@@ -95,6 +97,7 @@ class KMeansSuite extends MLTest with DefaultReadWriteTest
with PMMLReadWriteTes
assert(kmeans.getSeed === 123)
assert(kmeans.getTol === 1e-3)
assert(kmeans.getDistanceMeasure === KMeans.COSINE)
+ assert(kmeans.getIntermediateStorageLevel === "MEMORY_ONLY")
}
test("parameters validation") {
@@ -110,6 +113,12 @@ class KMeansSuite extends MLTest with DefaultReadWriteTest
with PMMLReadWriteTes
intercept[IllegalArgumentException] {
new KMeans().setDistanceMeasure("no_such_a_measure")
}
+ intercept[IllegalArgumentException] {
+ new KMeans().setIntermediateStorageLevel("NONE")
+ }
+ intercept[IllegalArgumentException] {
+ new KMeans().setIntermediateStorageLevel("no_such_a_level")
+ }
}
test("fit, transform and summary") {
diff --git a/python/pyspark/ml/clustering.py b/python/pyspark/ml/clustering.py
index c3964724dd3a..ba3e486cb4ef 100644
--- a/python/pyspark/ml/clustering.py
+++ b/python/pyspark/ml/clustering.py
@@ -36,6 +36,7 @@ from pyspark.ml.param.shared import (
HasCheckpointInterval,
HasSolver,
HasMaxBlockSizeInMB,
+ HasIntermediateStorageLevel,
Param,
Params,
TypeConverters,
@@ -582,6 +583,7 @@ class _KMeansParams(
HasWeightCol,
HasSolver,
HasMaxBlockSizeInMB,
+ HasIntermediateStorageLevel,
):
"""
Params for :py:class:`KMeans` and :py:class:`KMeansModel`.
@@ -921,6 +923,13 @@ class KMeans(JavaEstimator[KMeansModel], _KMeansParams,
JavaMLWritable, JavaMLRe
"""
return self._set(maxBlockSizeInMB=value)
+ @since("5.0.0")
+ def setIntermediateStorageLevel(self, value: str) -> "KMeans":
+ """
+ Sets the value of :py:attr:`intermediateStorageLevel`.
+ """
+ return self._set(intermediateStorageLevel=value)
+
@inherit_doc
class _BisectingKMeansParams(
diff --git a/python/pyspark/ml/param/_shared_params_code_gen.py
b/python/pyspark/ml/param/_shared_params_code_gen.py
index bbcaa208a39d..a885d6f1be5e 100644
--- a/python/pyspark/ml/param/_shared_params_code_gen.py
+++ b/python/pyspark/ml/param/_shared_params_code_gen.py
@@ -333,6 +333,12 @@ if __name__ == "__main__":
"0.0",
"TypeConverters.toFloat",
),
+ (
+ "intermediateStorageLevel",
+ "StorageLevel for intermediate datasets. Cannot be 'NONE'.",
+ '"MEMORY_AND_DISK"',
+ "TypeConverters.toString",
+ ),
(
"numTrainWorkers",
"number of training workers",
diff --git a/python/pyspark/ml/param/shared.py
b/python/pyspark/ml/param/shared.py
index e60f2a7432f7..ae8100b155e1 100644
--- a/python/pyspark/ml/param/shared.py
+++ b/python/pyspark/ml/param/shared.py
@@ -789,6 +789,29 @@ class HasMaxBlockSizeInMB(Params):
return self.getOrDefault(self.maxBlockSizeInMB)
+class HasIntermediateStorageLevel(Params):
+ """
+ Mixin for param intermediateStorageLevel: StorageLevel for intermediate
datasets. Cannot be 'NONE'.
+ """
+
+ intermediateStorageLevel: "Param[str]" = Param(
+ Params._dummy(),
+ "intermediateStorageLevel",
+ "StorageLevel for intermediate datasets. Cannot be 'NONE'.",
+ typeConverter=TypeConverters.toString,
+ )
+
+ def __init__(self) -> None:
+ super().__init__()
+ self._setDefault(intermediateStorageLevel="MEMORY_AND_DISK")
+
+ def getIntermediateStorageLevel(self) -> str:
+ """
+ Gets the value of intermediateStorageLevel or its default value.
+ """
+ return self.getOrDefault(self.intermediateStorageLevel)
+
+
class HasNumTrainWorkers(Params):
"""
Mixin for param numTrainWorkers: number of training workers
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
