This is an automated email from the ASF dual-hosted git repository.
zhengruifeng pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.0 by this push:
new fcf427772c68 [SPARK-60608][ML][CONNECT] Apply the class allowlist to
nested model readers
fcf427772c68 is described below
commit fcf427772c685425375cb380a9755f31fe772ea9
Author: Ruifeng Zheng <[email protected]>
AuthorDate: Thu Oct 29 08:03:03 2026 +0800
[SPARK-60608][ML][CONNECT] Apply the class allowlist to nested model readers
### What changes were proposed in this pull request?
Apply the Spark Connect class allowlist in the shared reader used to load
nested ML stages. Explicitly allow RFormulaModel's built-in ColumnPruner and
VectorAttributeRewriter stages, which cannot be registered through
ServiceLoader because they lack no-argument constructors.
Add MLClassLoadingSuite with generated temporary fixtures, including a
fitted RFormulaModel round trip and a harmless test-only model omitted from the
allowlist. Keep the fixture helper in RFormula.scala so RFormulaModel retains
its private[feature] constructor.
### Why are the changes needed?
An allowlisted model can load nested stages selected by saved metadata.
Those stages need the same allowlist enforcement as the outer model. The
explicit entries for RFormulaModel's internal stages preserve normal model
loading.
### Does this PR introduce _any_ user-facing change?
Yes. Spark Connect rejects non-allowlisted nested ML classes. Supported
fitted RFormulaModel instances continue to load, including their internal
preprocessing stages.
### How was this patch tested?
MLClassLoadingSuite covers an allowed nested stage, a fitted RFormulaModel
saving and loading through Connect with identical transformed features and
labels, an unregistered model loading outside Connect, and Connect rejecting
that unregistered nested model.
Attempted `connect/testOnly
org.apache.spark.sql.connect.ml.MLClassLoadingSuite
org.apache.spark.sql.connect.ml.MLSuite` with Java 17. SBT dependency
resolution failed before compilation with `UnknownHostException` for Maven
repositories, so the updated suites have not run locally.
`git diff --check`, ASCII and line-length checks, and source validation
passed.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: Codex (GPT-6)
Closes #59854 from zhengruifeng/ml-connect-nested-class-allowlist.
Authored-by: Ruifeng Zheng <[email protected]>
Signed-off-by: Ruifeng Zheng <[email protected]>
(cherry picked from commit 44ec3726bc088d0d8ec2e5c3aee3dffe16fa8d16)
Signed-off-by: Ruifeng Zheng <[email protected]>
---
.../org/apache/spark/ml/feature/RFormula.scala | 10 ++
.../scala/org/apache/spark/ml/util/ReadWrite.scala | 2 +-
.../apache/spark/sql/connect/ml/MLHandler.scala | 9 +-
.../spark/sql/connect/ml/MLClassLoadingSuite.scala | 135 +++++++++++++++++++++
4 files changed, 154 insertions(+), 2 deletions(-)
diff --git a/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
b/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
index d2191185dddd..f157366c1c1a 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/feature/RFormula.scala
@@ -475,6 +475,16 @@ object RFormulaModel extends MLReadable[RFormulaModel] {
}
}
+/** Creates minimal RFormulaModel containers for nested class-loading tests. */
+private[spark] object RFormulaModelTestHelper {
+
+ def withStage(stage: Transformer): RFormulaModel = {
+ val pipeline = new PipelineModel("pipelineFixture", Array(stage))
+ val formula = ResolvedRFormula("y", Seq(Seq("x")), hasIntercept = true)
+ new RFormulaModel("rFormulaFixture", formula, pipeline)
+ }
+}
+
/**
* Utility transformer for removing temporary columns from a DataFrame.
* TODO(ekl) make this a public transformer
diff --git a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
index a3e7b0c15a5f..f0668f900a34 100644
--- a/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
+++ b/mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala
@@ -740,7 +740,7 @@ private[ml] object DefaultParamsReader {
def loadParamsInstanceReader[T](path: String, spark: SparkSession):
MLReader[T] = {
val metadata = DefaultParamsReader.loadMetadata(path, spark)
- val cls = Utils.classForName(metadata.className)
+ val cls = MLAllowListedLoader.load(metadata.className)
cls.getMethod("read").invoke(null).asInstanceOf[MLReader[T]]
}
}
diff --git
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
index ca3e013a093f..1d619f31de5c 100644
---
a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
+++
b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/ml/MLHandler.scala
@@ -123,7 +123,14 @@ private[connect] object MLHandler extends Logging {
val transformerClasses = MLUtils.loadOperators(classOf[Transformer])
val estimatorClasses = MLUtils.loadOperators(classOf[Estimator[_]])
val evaluatorClasses = MLUtils.loadOperators(classOf[Evaluator])
- transformerClasses ++ estimatorClasses ++ evaluatorClasses ++ Map(
+ // RFormulaModel saves these internal stages, which lack the no-argument
constructors
+ // required by ServiceLoader.
+ val internalTransformerClasses = Seq(
+ "org.apache.spark.ml.feature.ColumnPruner",
+ "org.apache.spark.ml.feature.VectorAttributeRewriter")
+ .map(name => name -> Utils.classForName(name))
+ .toMap
+ transformerClasses ++ estimatorClasses ++ evaluatorClasses ++
internalTransformerClasses ++ Map(
"org.apache.spark.ml.clustering.PowerIterationClustering" ->
classOf[org.apache.spark.ml.clustering.PowerIterationClustering])
}
diff --git
a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLClassLoadingSuite.scala
b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLClassLoadingSuite.scala
new file mode 100644
index 000000000000..d51997bad213
--- /dev/null
+++
b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ml/MLClassLoadingSuite.scala
@@ -0,0 +1,135 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connect.ml
+
+import java.io.File
+
+import org.apache.spark.connect.proto
+import org.apache.spark.ml.{Model, Transformer}
+import org.apache.spark.ml.feature.{RFormula, RFormulaModel,
RFormulaModelTestHelper, VectorAssembler}
+import org.apache.spark.ml.param.ParamMap
+import org.apache.spark.ml.util.{DefaultParamsReadable, DefaultParamsWritable}
+import org.apache.spark.sql.{DataFrame, Dataset}
+import org.apache.spark.sql.connect.SparkConnectTestUtils
+import org.apache.spark.sql.types.StructType
+
+// A harmless ML-readable model deliberately omitted from the ServiceLoader
registrations.
+class UnregisteredModel(override val uid: String)
+ extends Model[UnregisteredModel]
+ with DefaultParamsWritable {
+
+ override def transform(dataset: Dataset[_]): DataFrame = dataset.toDF()
+
+ override def transformSchema(schema: StructType): StructType = schema
+
+ override def copy(extra: ParamMap): UnregisteredModel = defaultCopy(extra)
+}
+
+object UnregisteredModel extends DefaultParamsReadable[UnregisteredModel]
+
+class MLClassLoadingSuite extends MLHelper {
+
+ override protected def afterEach(): Unit = {
+ try {
+ MLHandler.currentSessionHolder.remove()
+ } finally {
+ super.afterEach()
+ }
+ }
+
+ private def withSavedModel(stage: Transformer)(f: String => Unit): Unit = {
+ withTempDir { dir =>
+ val path = new File(dir, "model").getAbsolutePath
+ RFormulaModelTestHelper.withStage(stage).write.save(path)
+ f(path)
+ }
+ }
+
+ private def readModel(path: String): RFormulaModel = {
+ val command = proto.MlCommand
+ .newBuilder()
+ .setRead(
+ proto.MlCommand.Read
+ .newBuilder()
+ .setOperator(
+ proto.MlOperator
+ .newBuilder()
+ .setName("org.apache.spark.ml.feature.RFormulaModel")
+ .setType(proto.MlOperator.OperatorType.OPERATOR_TYPE_MODEL))
+ .setPath(path))
+ .build()
+ val sessionHolder = SparkConnectTestUtils.createDummySessionHolder(spark)
+ val result = MLHandler.handleMlCommand(sessionHolder, command)
+ val modelId = result.getOperatorInfo.getObjRef.getId
+ assert(modelId.nonEmpty)
+ sessionHolder.mlCache.get(modelId).asInstanceOf[RFormulaModel]
+ }
+
+ test("load an RFormulaModel with an allowlisted nested stage") {
+ val stage = new VectorAssembler("stageFixture")
+ .setInputCols(Array("x"))
+ .setOutputCol("features")
+ withSavedModel(stage) { path =>
+ val model = readModel(path)
+ assert(model.uid == "rFormulaFixture")
+ }
+ }
+
+ test("load a fitted RFormulaModel with internal nested stages") {
+ val dataset = spark
+ .createDataFrame(Seq((1.0, 1.0, "a"), (0.0, 2.0, "b"), (1.0, 0.0, "a")))
+ .toDF("y", "x", "s")
+ val model = new RFormula()
+ .setFormula("y ~ x + s")
+ .setStringIndexerOrderType("alphabetDesc")
+ .fit(dataset)
+ val expected = model.transform(dataset).select("features",
"label").collect().toSeq
+
+ withTempDir { dir =>
+ val path = new File(dir, "model").getAbsolutePath
+ model.write.save(path)
+ val loadedModel = readModel(path)
+ assert(loadedModel.uid == model.uid)
+ assert(loadedModel.getFormula == model.getFormula)
+ assert(
+ loadedModel.transform(dataset).select("features",
"label").collect().toSeq == expected)
+ }
+ }
+
+ test("the unregistered nested model is readable outside Spark Connect") {
+ val stage = new UnregisteredModel("stageFixture")
+ withSavedModel(stage) { path =>
+ val nestedPath = new File(path,
s"pipelineModel/stages/0_${stage.uid}").getAbsolutePath
+ MLHandler.currentSessionHolder.remove()
+ val model = UnregisteredModel.load(nestedPath)
+ assert(model.uid == stage.uid)
+ val transformers = MLUtils.loadOperators(classOf[Transformer])
+ assert(!transformers.contains(classOf[UnregisteredModel].getName))
+ }
+ }
+
+ test("reject an RFormulaModel with a non-allowlisted nested stage") {
+ withSavedModel(new UnregisteredModel("stageFixture")) { path =>
+ val error = intercept[MlUnsupportedException] {
+ readModel(path)
+ }
+ assert(error.getMessage.contains(classOf[UnregisteredModel].getName))
+ assert(error.getMessage.contains("not in the allowlist"))
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
