From 6104c043f82236a6148df94ae9e80f9f4cf2b263 Mon Sep 17 00:00:00 2001 From: "Mao Li (Personal)" Date: Sun, 12 Jul 2026 01:20:06 +0100 Subject: [PATCH] [SPARK-57870][ML][PYTHON] Apply HasIntermediateStorageLevel to tree ensemble estimators (GBT, RandomForest) Co-Authored-By: Claude Fable 5 --- .../ml/classification/GBTClassifier.scala | 13 +++-- .../RandomForestClassifier.scala | 11 ++++- .../spark/ml/regression/GBTRegressor.scala | 12 +++-- .../ml/regression/RandomForestRegressor.scala | 10 +++- .../ml/tree/impl/GradientBoostedTrees.scala | 35 ++++++++------ .../spark/ml/tree/impl/RandomForest.scala | 13 +++-- .../org/apache/spark/ml/tree/treeParams.scala | 2 +- .../classification/GBTClassifierSuite.scala | 13 +++++ .../RandomForestClassifierSuite.scala | 43 +++++++++++++++++ .../ml/regression/GBTRegressorSuite.scala | 47 +++++++++++++++++++ .../RandomForestRegressorSuite.scala | 13 +++++ python/pyspark/ml/classification.py | 14 ++++++ python/pyspark/ml/regression.py | 14 ++++++ python/pyspark/ml/tree.py | 3 +- 14 files changed, 211 insertions(+), 32 deletions(-) diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala index 2c5c7e7740a33..b63d53f47684b 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/classification/GBTClassifier.scala @@ -36,6 +36,7 @@ import org.apache.spark.mllib.tree.model.{GradientBoostedTreesModel => OldGBTMod import org.apache.spark.sql._ import org.apache.spark.sql.functions._ import org.apache.spark.sql.types.StructType +import org.apache.spark.storage.StorageLevel /** * Gradient-Boosted Trees (GBTs) (http://en.wikipedia.org/wiki/Gradient_boosting) @@ -166,6 +167,10 @@ class GBTClassifier @Since("1.4.0") ( @Since("3.0.0") def setWeightCol(value: String): this.type = set(weightCol, value) + /** @group expertSetParam */ + @Since("5.0.0") + def setIntermediateStorageLevel(value: String): this.type = set(intermediateStorageLevel, value) + override protected def train( dataset: Dataset[_]): GBTClassificationModel = instrumented { instr => val withValidation = isDefined(validationIndicatorCol) && $(validationIndicatorCol).nonEmpty @@ -188,17 +193,19 @@ class GBTClassifier @Since("1.4.0") ( instr.logParams(this, labelCol, weightCol, featuresCol, predictionCol, leafCol, impurity, lossType, maxDepth, maxBins, maxIter, maxMemoryInMB, minInfoGain, minInstancesPerNode, minWeightFractionPerNode, seed, stepSize, subsamplingRate, cacheNodeIds, - checkpointInterval, featureSubsetStrategy, validationIndicatorCol, validationTol, thresholds) + checkpointInterval, featureSubsetStrategy, validationIndicatorCol, validationTol, thresholds, + intermediateStorageLevel) instr.logNumClasses(numClasses) val categoricalFeatures = MetadataUtils.getCategoricalFeatures(dataset.schema($(featuresCol))) val boostingStrategy = super.getOldBoostingStrategy(categoricalFeatures, OldAlgo.Classification) + val storageLevel = StorageLevel.fromString($(intermediateStorageLevel)) val (baseLearners, learnerWeights) = if (withValidation) { GradientBoostedTrees.runWithValidation(trainDataset, validationDataset, boostingStrategy, - $(seed), $(featureSubsetStrategy), Some(instr)) + $(seed), $(featureSubsetStrategy), Some(instr), storageLevel) } else { GradientBoostedTrees.run(trainDataset, boostingStrategy, $(seed), $(featureSubsetStrategy), - Some(instr)) + Some(instr), storageLevel) } baseLearners.foreach(copyValues(_)) diff --git a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala index 2c22ca5b42302..fd33ca7e4109a 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/classification/RandomForestClassifier.scala @@ -36,6 +36,7 @@ import org.apache.spark.mllib.tree.model.{RandomForestModel => OldRandomForestMo import org.apache.spark.sql._ import org.apache.spark.sql.functions.{col, udf} import org.apache.spark.sql.types.StructType +import org.apache.spark.storage.StorageLevel /** * Random Forest learning algorithm for @@ -139,6 +140,10 @@ class RandomForestClassifier @Since("1.4.0") ( @Since("3.0.0") def setWeightCol(value: String): this.type = set(weightCol, value) + /** @group expertSetParam */ + @Since("5.0.0") + def setIntermediateStorageLevel(value: String): this.type = set(intermediateStorageLevel, value) + override protected def train( dataset: Dataset[_]): RandomForestClassificationModel = instrumented { instr => instr.logPipelineStage(this) @@ -168,10 +173,12 @@ class RandomForestClassifier @Since("1.4.0") ( instr.logParams(this, labelCol, featuresCol, weightCol, predictionCol, probabilityCol, rawPredictionCol, leafCol, impurity, numTrees, featureSubsetStrategy, maxDepth, maxBins, maxMemoryInMB, minInfoGain, pruneTree, minInstancesPerNode, minWeightFractionPerNode, seed, - subsamplingRate, thresholds, cacheNodeIds, checkpointInterval, bootstrap) + subsamplingRate, thresholds, cacheNodeIds, checkpointInterval, bootstrap, + intermediateStorageLevel) val trees = RandomForest - .run(instances, strategy, getNumTrees, getFeatureSubsetStrategy, getSeed, Some(instr)) + .run(instances, strategy, getNumTrees, getFeatureSubsetStrategy, getSeed, Some(instr), + storageLevel = StorageLevel.fromString($(intermediateStorageLevel))) .map(_.asInstanceOf[DecisionTreeClassificationModel]) trees.foreach(copyValues(_)) diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala index 71436036d1ea6..91c8f0d17dff1 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/GBTRegressor.scala @@ -35,6 +35,7 @@ import org.apache.spark.mllib.tree.model.{GradientBoostedTreesModel => OldGBTMod import org.apache.spark.sql._ import org.apache.spark.sql.functions._ import org.apache.spark.sql.types.StructType +import org.apache.spark.storage.StorageLevel /** * Gradient-Boosted Trees (GBTs) @@ -164,6 +165,10 @@ class GBTRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: String) @Since("3.0.0") def setWeightCol(value: String): this.type = set(weightCol, value) + /** @group expertSetParam */ + @Since("5.0.0") + def setIntermediateStorageLevel(value: String): this.type = set(intermediateStorageLevel, value) + override protected def train(dataset: Dataset[_]): GBTRegressionModel = instrumented { instr => val withValidation = isDefined(validationIndicatorCol) && $(validationIndicatorCol).nonEmpty val (trainDataset, validationDataset) = if (withValidation) { @@ -178,16 +183,17 @@ class GBTRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: String) instr.logParams(this, labelCol, featuresCol, predictionCol, leafCol, weightCol, impurity, lossType, maxDepth, maxBins, maxIter, maxMemoryInMB, minInfoGain, minInstancesPerNode, minWeightFractionPerNode, seed, stepSize, subsamplingRate, cacheNodeIds, checkpointInterval, - featureSubsetStrategy, validationIndicatorCol, validationTol) + featureSubsetStrategy, validationIndicatorCol, validationTol, intermediateStorageLevel) val categoricalFeatures = MetadataUtils.getCategoricalFeatures(dataset.schema($(featuresCol))) val boostingStrategy = super.getOldBoostingStrategy(categoricalFeatures, OldAlgo.Regression) + val storageLevel = StorageLevel.fromString($(intermediateStorageLevel)) val (baseLearners, learnerWeights) = if (withValidation) { GradientBoostedTrees.runWithValidation(trainDataset, validationDataset, boostingStrategy, - $(seed), $(featureSubsetStrategy), Some(instr)) + $(seed), $(featureSubsetStrategy), Some(instr), storageLevel) } else { GradientBoostedTrees.run(trainDataset, boostingStrategy, - $(seed), $(featureSubsetStrategy), Some(instr)) + $(seed), $(featureSubsetStrategy), Some(instr), storageLevel) } baseLearners.foreach(copyValues(_)) diff --git a/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala b/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala index 8d9b4817833bc..1b64eaa3cba32 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/regression/RandomForestRegressor.scala @@ -36,6 +36,7 @@ import org.apache.spark.mllib.tree.model.{RandomForestModel => OldRandomForestMo import org.apache.spark.sql._ import org.apache.spark.sql.functions.{col, udf} import org.apache.spark.sql.types.StructType +import org.apache.spark.storage.StorageLevel /** * Random Forest @@ -133,6 +134,10 @@ class RandomForestRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: S @Since("3.0.0") def setWeightCol(value: String): this.type = set(weightCol, value) + /** @group expertSetParam */ + @Since("5.0.0") + def setIntermediateStorageLevel(value: String): this.type = set(intermediateStorageLevel, value) + override protected def train( dataset: Dataset[_]): RandomForestRegressionModel = instrumented { instr => val categoricalFeatures: Map[Int, Int] = @@ -154,10 +159,11 @@ class RandomForestRegressor @Since("1.4.0") (@Since("1.4.0") override val uid: S instr.logParams(this, labelCol, featuresCol, weightCol, predictionCol, leafCol, impurity, numTrees, featureSubsetStrategy, maxDepth, maxBins, maxMemoryInMB, minInfoGain, minInstancesPerNode, minWeightFractionPerNode, seed, subsamplingRate, cacheNodeIds, - checkpointInterval, bootstrap) + checkpointInterval, bootstrap, intermediateStorageLevel) val trees = RandomForest - .run(instances, strategy, getNumTrees, getFeatureSubsetStrategy, getSeed, Some(instr)) + .run(instances, strategy, getNumTrees, getFeatureSubsetStrategy, getSeed, Some(instr), + storageLevel = StorageLevel.fromString($(intermediateStorageLevel))) .map(_.asInstanceOf[DecisionTreeRegressionModel]) trees.foreach(copyValues(_)) diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/impl/GradientBoostedTrees.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/impl/GradientBoostedTrees.scala index 81ffa5c86a9fa..f9cfd4adfe31e 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/tree/impl/GradientBoostedTrees.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/tree/impl/GradientBoostedTrees.scala @@ -49,18 +49,19 @@ private[spark] object GradientBoostedTrees extends Logging { boostingStrategy: OldBoostingStrategy, seed: Long, featureSubsetStrategy: String, - instr: Option[Instrumentation] = None): + instr: Option[Instrumentation] = None, + storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK): (Array[DecisionTreeRegressionModel], Array[Double]) = { val algo = boostingStrategy.treeStrategy.algo algo match { case OldAlgo.Regression => GradientBoostedTrees.boost(input, input, boostingStrategy, validate = false, - seed, featureSubsetStrategy, instr) + seed, featureSubsetStrategy, instr, storageLevel) case OldAlgo.Classification => // Map labels to -1, +1 so binary classification can be treated as regression. val remappedInput = input.map(x => Instance((x.label * 2) - 1, x.weight, x.features)) GradientBoostedTrees.boost(remappedInput, remappedInput, boostingStrategy, validate = false, - seed, featureSubsetStrategy, instr) + seed, featureSubsetStrategy, instr, storageLevel) case _ => throw new IllegalArgumentException(s"$algo is not supported by gradient boosting.") } @@ -84,13 +85,14 @@ private[spark] object GradientBoostedTrees extends Logging { boostingStrategy: OldBoostingStrategy, seed: Long, featureSubsetStrategy: String, - instr: Option[Instrumentation] = None): + instr: Option[Instrumentation] = None, + storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK): (Array[DecisionTreeRegressionModel], Array[Double]) = { val algo = boostingStrategy.treeStrategy.algo algo match { case OldAlgo.Regression => GradientBoostedTrees.boost(input, validationInput, boostingStrategy, - validate = true, seed, featureSubsetStrategy, instr) + validate = true, seed, featureSubsetStrategy, instr, storageLevel) case OldAlgo.Classification => // Map labels to -1, +1 so binary classification can be treated as regression. val remappedInput = input.map( @@ -98,7 +100,7 @@ private[spark] object GradientBoostedTrees extends Logging { val remappedValidationInput = validationInput.map( x => Instance((x.label * 2) - 1, x.weight, x.features)) GradientBoostedTrees.boost(remappedInput, remappedValidationInput, boostingStrategy, - validate = true, seed, featureSubsetStrategy, instr) + validate = true, seed, featureSubsetStrategy, instr, storageLevel) case _ => throw new IllegalArgumentException(s"$algo is not supported by the gradient boosting.") } @@ -294,7 +296,8 @@ private[spark] object GradientBoostedTrees extends Logging { validate: Boolean, seed: Long, featureSubsetStrategy: String, - instr: Option[Instrumentation] = None): + instr: Option[Instrumentation] = None, + storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK): (Array[DecisionTreeRegressionModel], Array[Double]) = { val earlyStopModelSizeThresholdInBytes = TreeConfig.trainingEarlyStopModelSizeThresholdInBytes lastEarlyStoppedModelSize = 0 @@ -324,7 +327,7 @@ private[spark] object GradientBoostedTrees extends Logging { // Prepare periodic checkpointers // Note: this is checkpointing the unweighted training error val predErrorCheckpointer = new PeriodicRDDCheckpointer[(Double, Double)]( - treeStrategy.getCheckpointInterval(), sc, StorageLevel.MEMORY_AND_DISK) + treeStrategy.getCheckpointInterval(), sc, storageLevel) timer.stop("init") @@ -349,7 +352,7 @@ private[spark] object GradientBoostedTrees extends Logging { // Cache input RDD for speedup during multiple passes. val treePoints = TreePoint.convertToTreeRDD( retaggedInput, splits, metadata) - .persist(StorageLevel.MEMORY_AND_DISK) + .persist(storageLevel) .setName("binned tree points") val firstCounts = BaggedPoint @@ -359,7 +362,7 @@ private[spark] object GradientBoostedTrees extends Logging { require(bagged.subsampleCounts.length == 1) require(bagged.sampleWeight == bagged.datum.weight) bagged.subsampleCounts.head - }.persist(StorageLevel.MEMORY_AND_DISK) + }.persist(storageLevel) .setName("firstCounts at iter=0") val firstBagged = treePoints.zip(firstCounts) @@ -372,7 +375,8 @@ private[spark] object GradientBoostedTrees extends Logging { metadata = metadata, bcSplits = bcSplits, strategy = treeStrategy, numTrees = 1, featureSubsetStrategy = featureSubsetStrategy, seed = seed, instr = instr, parentUID = None, - earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes) + earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes, + storageLevel = storageLevel) .head.asInstanceOf[DecisionTreeRegressionModel] firstCounts.unpersist() @@ -397,11 +401,11 @@ private[spark] object GradientBoostedTrees extends Logging { timer.start("init validation") validationTreePoints = TreePoint.convertToTreeRDD( validationInput.retag(classOf[Instance]), splits, metadata) - .persist(StorageLevel.MEMORY_AND_DISK) + .persist(storageLevel) validatePredError = computeInitialPredictionAndError( validationTreePoints, firstTreeWeight, firstTreeModel, loss, bcSplits) validatePredErrorCheckpointer = new PeriodicRDDCheckpointer[(Double, Double)]( - treeStrategy.getCheckpointInterval(), sc, StorageLevel.MEMORY_AND_DISK) + treeStrategy.getCheckpointInterval(), sc, storageLevel) validatePredErrorCheckpointer.update(validatePredError) bestValidateError = computeWeightedError(validationTreePoints, validatePredError) timer.stop("init validation") @@ -437,7 +441,7 @@ private[spark] object GradientBoostedTrees extends Logging { // Update labels with pseudo-residuals val newLabel = -loss.gradient(pred, bagged.datum.label) (newLabel, bagged.subsampleCounts.head) - }.persist(StorageLevel.MEMORY_AND_DISK) + }.persist(storageLevel) .setName(s"labelWithCounts at iter=$m") val bagged = treePoints.zip(labelWithCounts) @@ -451,7 +455,8 @@ private[spark] object GradientBoostedTrees extends Logging { metadata = metadata, bcSplits = bcSplits, strategy = treeStrategy, numTrees = 1, featureSubsetStrategy = featureSubsetStrategy, seed = seed + m, instr = None, parentUID = None, - earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes - accTreeSize) + earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes - accTreeSize, + storageLevel = storageLevel) .head.asInstanceOf[DecisionTreeRegressionModel] labelWithCounts.unpersist() diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala index b3ca7f04c3de5..1a97406eaf1b5 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/tree/impl/RandomForest.scala @@ -125,7 +125,8 @@ private[spark] object RandomForest extends Logging with Serializable { seed: Long, instr: Option[Instrumentation], parentUID: Option[String] = None, - earlyStopModelSizeThresholdInBytes: Long = 0): Array[DecisionTreeModel] = { + earlyStopModelSizeThresholdInBytes: Long = 0, + storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK): Array[DecisionTreeModel] = { lastEarlyStoppedModelSize = 0 val timer = new TimeTracker() timer.start("total") @@ -173,7 +174,7 @@ private[spark] object RandomForest extends Logging with Serializable { // At first, all the rows belong to the root nodes (node Id == 1). nodeIds = baggedInput.map { _ => Array.fill(numTrees)(1) } nodeIdCheckpointer = new PeriodicRDDCheckpointer[Array[Int]]( - strategy.getCheckpointInterval(), sc, StorageLevel.MEMORY_AND_DISK) + strategy.getCheckpointInterval(), sc, storageLevel) nodeIdCheckpointer.update(nodeIds) } @@ -307,7 +308,8 @@ private[spark] object RandomForest extends Logging with Serializable { featureSubsetStrategy: String, seed: Long, instr: Option[Instrumentation], - parentUID: Option[String] = None): Array[DecisionTreeModel] = { + parentUID: Option[String] = None, + storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK): Array[DecisionTreeModel] = { val earlyStopModelSizeThresholdInBytes = TreeConfig.trainingEarlyStopModelSizeThresholdInBytes val timer = new TimeTracker() @@ -344,7 +346,7 @@ private[spark] object RandomForest extends Logging with Serializable { strategy.bootstrap, (tp: TreePoint) => tp.weight, seed = seed) - .persist(StorageLevel.MEMORY_AND_DISK) + .persist(storageLevel) .setName("bagged tree points") val trees = runBagged( @@ -357,7 +359,8 @@ private[spark] object RandomForest extends Logging with Serializable { seed = seed, instr = instr, parentUID = parentUID, - earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes) + earlyStopModelSizeThresholdInBytes = earlyStopModelSizeThresholdInBytes, + storageLevel = storageLevel) baggedInput.unpersist() bcSplits.destroy() diff --git a/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala b/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala index 2244d49b2a35f..08c6dc730f14b 100644 --- a/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala +++ b/mllib/src/main/scala/org/apache/spark/ml/tree/treeParams.scala @@ -340,7 +340,7 @@ private[spark] object TreeEnsembleParams { * * Note: Marked as private since this may be made public in the future. */ -private[ml] trait TreeEnsembleParams extends DecisionTreeParams { +private[ml] trait TreeEnsembleParams extends DecisionTreeParams with HasIntermediateStorageLevel { /** * Fraction of the training data used for learning each decision tree, in range (0, 1]. diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala index 6ce2108b1f7c8..75b06c2cb5ee9 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/classification/GBTClassifierSuite.scala @@ -82,6 +82,19 @@ class GBTClassifierSuite extends MLTest with DefaultReadWriteTest { ParamsSuite.checkParams(model) } + test("SPARK-57870: intermediateStorageLevel param") { + val gbt = new GBTClassifier() + assert(gbt.getIntermediateStorageLevel === "MEMORY_AND_DISK") + gbt.setIntermediateStorageLevel("MEMORY_ONLY") + assert(gbt.getIntermediateStorageLevel === "MEMORY_ONLY") + intercept[IllegalArgumentException] { + new GBTClassifier().setIntermediateStorageLevel("NONE") + } + intercept[IllegalArgumentException] { + new GBTClassifier().setIntermediateStorageLevel("no_such_a_level") + } + } + test("GBTClassifier: default params") { val gbt = new GBTClassifier assert(gbt.getLabelCol === "label") diff --git a/mllib/src/test/scala/org/apache/spark/ml/classification/RandomForestClassifierSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/classification/RandomForestClassifierSuite.scala index 562cccedeef4f..339bf1ca4fc04 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/classification/RandomForestClassifierSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/classification/RandomForestClassifierSuite.scala @@ -17,6 +17,8 @@ package org.apache.spark.ml.classification +import scala.collection.mutable + import org.apache.spark.SparkFunSuite import org.apache.spark.ml.classification.LinearSVCSuite.generateSVMInput import org.apache.spark.ml.feature.LabeledPoint @@ -30,8 +32,10 @@ import org.apache.spark.mllib.regression.{LabeledPoint => OldLabeledPoint} import org.apache.spark.mllib.tree.{EnsembleTestHelper, RandomForest => OldRandomForest} import org.apache.spark.mllib.tree.configuration.{Algo => OldAlgo} import org.apache.spark.rdd.RDD +import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.functions._ +import org.apache.spark.storage.StorageLevel import org.apache.spark.util.ArrayImplicits._ /** @@ -83,6 +87,45 @@ class RandomForestClassifierSuite extends MLTest with DefaultReadWriteTest { ParamsSuite.checkParams(model) } + test("SPARK-57870: intermediateStorageLevel param") { + val rf = new RandomForestClassifier() + assert(rf.getIntermediateStorageLevel === "MEMORY_AND_DISK") + rf.setIntermediateStorageLevel("MEMORY_ONLY") + assert(rf.getIntermediateStorageLevel === "MEMORY_ONLY") + intercept[IllegalArgumentException] { + new RandomForestClassifier().setIntermediateStorageLevel("NONE") + } + intercept[IllegalArgumentException] { + new RandomForestClassifier().setIntermediateStorageLevel("no_such_a_level") + } + } + + test("SPARK-57870: intermediateStorageLevel is applied to intermediate datasets") { + val df = TreeTests.setMetadata(orderedLabeledPoints5_20, Map.empty[Int, Int], 2) + val rf = new RandomForestClassifier() + .setNumTrees(2) + .setMaxDepth(2) + .setIntermediateStorageLevel("DISK_ONLY") + + val capturedLevels = mutable.ArrayBuffer.empty[StorageLevel] + val listener = new SparkListener { + override def onStageCompleted(stageCompleted: SparkListenerStageCompleted): Unit = { + capturedLevels ++= stageCompleted.stageInfo.rddInfos + .filter(_.name == "bagged tree points") + .map(_.storageLevel) + } + } + sc.addSparkListener(listener) + try { + rf.fit(df) + sc.listenerBus.waitUntilEmpty() + } finally { + sc.removeSparkListener(listener) + } + assert(capturedLevels.nonEmpty) + capturedLevels.foreach(level => assert(level === StorageLevel.DISK_ONLY)) + } + test("RandomForestClassifier validate input dataset") { testInvalidClassificationLabels(new RandomForestClassifier().fit(_), None) testInvalidWeights(new RandomForestClassifier().setWeightCol("weight").fit(_)) diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala index d7f15dc2cfe9e..35ae7cb59daad 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/regression/GBTRegressorSuite.scala @@ -17,6 +17,8 @@ package org.apache.spark.ml.regression +import scala.collection.mutable + import org.apache.spark.SparkFunSuite import org.apache.spark.ml.feature.LabeledPoint import org.apache.spark.ml.linalg.{Vector, Vectors} @@ -28,8 +30,10 @@ import org.apache.spark.mllib.tree.{EnsembleTestHelper, GradientBoostedTrees => import org.apache.spark.mllib.tree.configuration.{Algo => OldAlgo} import org.apache.spark.mllib.util.LinearDataGenerator import org.apache.spark.rdd.RDD +import org.apache.spark.scheduler.{SparkListener, SparkListenerStageCompleted} import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.functions.lit +import org.apache.spark.storage.StorageLevel import org.apache.spark.util.ArrayImplicits._ import org.apache.spark.util.Utils @@ -69,6 +73,49 @@ class GBTRegressorSuite extends MLTest with DefaultReadWriteTest { xVariance = Array(0.7, 1.2), nPoints = 1000, seed, eps = 0.5), 2).map(_.asML).toDF() } + test("SPARK-57870: intermediateStorageLevel param") { + val gbt = new GBTRegressor() + assert(gbt.getIntermediateStorageLevel === "MEMORY_AND_DISK") + gbt.setIntermediateStorageLevel("MEMORY_ONLY") + assert(gbt.getIntermediateStorageLevel === "MEMORY_ONLY") + intercept[IllegalArgumentException] { + new GBTRegressor().setIntermediateStorageLevel("NONE") + } + intercept[IllegalArgumentException] { + new GBTRegressor().setIntermediateStorageLevel("no_such_a_level") + } + } + + test("SPARK-57870: intermediateStorageLevel is applied to intermediate datasets") { + val df = trainData.toDF() + val gbt = new GBTRegressor() + .setMaxIter(2) + .setMaxDepth(2) + .setIntermediateStorageLevel("DISK_ONLY") + + val capturedLevels = mutable.ArrayBuffer.empty[StorageLevel] + val listener = new SparkListener { + override def onStageCompleted(stageCompleted: SparkListenerStageCompleted): Unit = { + capturedLevels ++= stageCompleted.stageInfo.rddInfos + .filter { info => + info.name == "binned tree points" || + info.name.startsWith("firstCounts at iter=") || + info.name.startsWith("labelWithCounts at iter=") + } + .map(_.storageLevel) + } + } + sc.addSparkListener(listener) + try { + gbt.fit(df) + sc.listenerBus.waitUntilEmpty() + } finally { + sc.removeSparkListener(listener) + } + assert(capturedLevels.nonEmpty) + capturedLevels.foreach(level => assert(level === StorageLevel.DISK_ONLY)) + } + test("Regression with continuous features") { val categoricalFeatures = Map.empty[Int, Int] GBTRegressor.supportedLossTypes.foreach { loss => diff --git a/mllib/src/test/scala/org/apache/spark/ml/regression/RandomForestRegressorSuite.scala b/mllib/src/test/scala/org/apache/spark/ml/regression/RandomForestRegressorSuite.scala index 15db8c5c22531..5265ddc52ddfc 100644 --- a/mllib/src/test/scala/org/apache/spark/ml/regression/RandomForestRegressorSuite.scala +++ b/mllib/src/test/scala/org/apache/spark/ml/regression/RandomForestRegressorSuite.scala @@ -58,6 +58,19 @@ class RandomForestRegressorSuite extends MLTest with DefaultReadWriteTest { // Tests calling train() ///////////////////////////////////////////////////////////////////////////// + test("SPARK-57870: intermediateStorageLevel param") { + val rf = new RandomForestRegressor() + assert(rf.getIntermediateStorageLevel === "MEMORY_AND_DISK") + rf.setIntermediateStorageLevel("MEMORY_ONLY") + assert(rf.getIntermediateStorageLevel === "MEMORY_ONLY") + intercept[IllegalArgumentException] { + new RandomForestRegressor().setIntermediateStorageLevel("NONE") + } + intercept[IllegalArgumentException] { + new RandomForestRegressor().setIntermediateStorageLevel("no_such_a_level") + } + } + test("RandomForestRegressor validate input dataset") { testInvalidRegressionLabels(new RandomForestRegressor().fit(_)) testInvalidWeights(new RandomForestRegressor().setWeightCol("weight").fit(_)) diff --git a/python/pyspark/ml/classification.py b/python/pyspark/ml/classification.py index 4b7f2e4da2090..030ef96f171e0 100644 --- a/python/pyspark/ml/classification.py +++ b/python/pyspark/ml/classification.py @@ -2268,6 +2268,13 @@ def setMinWeightFractionPerNode(self, value: float) -> "RandomForestClassifier": """ return self._set(minWeightFractionPerNode=value) + @since("5.0.0") + def setIntermediateStorageLevel(self, value: str) -> "RandomForestClassifier": + """ + Sets the value of :py:attr:`intermediateStorageLevel`. + """ + return self._set(intermediateStorageLevel=value) + class RandomForestClassificationModel( # type: ignore[misc] _TreeEnsembleModel, @@ -2757,6 +2764,13 @@ def setMinWeightFractionPerNode(self, value: float) -> "GBTClassifier": """ return self._set(minWeightFractionPerNode=value) + @since("5.0.0") + def setIntermediateStorageLevel(self, value: str) -> "GBTClassifier": + """ + Sets the value of :py:attr:`intermediateStorageLevel`. + """ + return self._set(intermediateStorageLevel=value) + class GBTClassificationModel( _TreeEnsembleModel, diff --git a/python/pyspark/ml/regression.py b/python/pyspark/ml/regression.py index d910db6c6d30b..a5a01974f2113 100644 --- a/python/pyspark/ml/regression.py +++ b/python/pyspark/ml/regression.py @@ -1584,6 +1584,13 @@ def setMinWeightFractionPerNode(self, value: float) -> "RandomForestRegressor": """ return self._set(minWeightFractionPerNode=value) + @since("5.0.0") + def setIntermediateStorageLevel(self, value: str) -> "RandomForestRegressor": + """ + Sets the value of :py:attr:`intermediateStorageLevel`. + """ + return self._set(intermediateStorageLevel=value) + class RandomForestRegressionModel( # type: ignore[misc] _JavaRegressionModel[Vector], @@ -1962,6 +1969,13 @@ def setMinWeightFractionPerNode(self, value: float) -> "GBTRegressor": """ return self._set(minWeightFractionPerNode=value) + @since("5.0.0") + def setIntermediateStorageLevel(self, value: str) -> "GBTRegressor": + """ + Sets the value of :py:attr:`intermediateStorageLevel`. + """ + return self._set(intermediateStorageLevel=value) + class GBTRegressionModel( _JavaRegressionModel[Vector], diff --git a/python/pyspark/ml/tree.py b/python/pyspark/ml/tree.py index 92692ec225a76..4220d9cdca9f6 100644 --- a/python/pyspark/ml/tree.py +++ b/python/pyspark/ml/tree.py @@ -21,6 +21,7 @@ from pyspark.ml.param import Params from pyspark.ml.param.shared import ( HasCheckpointInterval, + HasIntermediateStorageLevel, HasSeed, HasWeightCol, Param, @@ -256,7 +257,7 @@ def predictLeaf(self, value: Vector) -> float: return self._call_java("predictLeaf", value) -class _TreeEnsembleParams(_DecisionTreeParams): +class _TreeEnsembleParams(_DecisionTreeParams, HasIntermediateStorageLevel): """ Mixin for Decision Tree-based ensemble algorithms parameters. """