You are viewing a plain text version of this content. The canonical link for it is here.

Posted to reviews@spark.apache.org by "maoli67660 (via GitHub)" <gi...@apache.org> on 2026/07/01 23:35:39 UTC

[PR] [SPARK-57860][ML] Add HasIntermediateStorageLevel shared param and apply to KMeans [spark]

maoli67660 opened a new pull request, #56949:
URL: https://github.com/apache/spark/pull/56949

### 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)

--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org