Github user sethah commented on a diff in the pull request:
https://github.com/apache/spark/pull/19876#discussion_r160461644
--- Diff: mllib/src/main/scala/org/apache/spark/ml/util/ReadWrite.scala ---
@@ -126,15 +180,69 @@ abstract class MLWriter extends BaseReadWrite with
Logging {
this
}
+ // override for Java compatibility
+ override def session(sparkSession: SparkSession): this.type =
super.session(sparkSession)
+
+ // override for Java compatibility
+ override def context(sqlContext: SQLContext): this.type =
super.session(sqlContext.sparkSession)
+}
+
+/**
+ * A ML Writer which delegates based on the requested format.
+ */
+class GeneralMLWriter(stage: PipelineStage) extends MLWriter with Logging {
+ private var source: String = "internal"
+
/**
- * Overwrites if the output path already exists.
+ * Specifies the format of ML export (e.g. PMML, internal, or
+ * the fully qualified class name for export).
*/
- @Since("1.6.0")
- def overwrite(): this.type = {
- shouldOverwrite = true
+ @Since("2.3.0")
+ def format(source: String): this.type = {
+ this.source = source
this
}
+ /**
+ * Dispatches the save to the correct MLFormat.
+ */
+ @Since("2.3.0")
+ @throws[IOException]("If the input path already exists but overwrite is
not enabled.")
+ @throws[SparkException]("If multiple sources for a given short name
format are found.")
+ override protected def saveImpl(path: String) = {
+ val loader = Utils.getContextOrSparkClassLoader
+ val serviceLoader = ServiceLoader.load(classOf[MLFormatRegister],
loader)
+ val stageName = stage.getClass.getName
+ val targetName = s"${source}+${stageName}"
--- End diff --
don't need brackets
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
