[SPARK-48883][ML][R] Replace RDD read / write API invocation with Dataframe read / write API

[WeichenXu123]

What changes were proposed in this pull request?

Replace RDD read / write API invocation with Dataframe read / write API

Why are the changes needed?

In certain storage configuration cases, RDD read / write API has some issue for certain storage types that requires the account key, but Dataframe read / write API works.

Does this PR introduce any user-facing change?

No.

How was this patch tested?

Unit tests

Was this patch authored or co-authored using generative AI tooling?

No.

[dongjoon-hyun]

Hi, @WeichenXu123 and @zhengruifeng . This is Apache Software Foundation repository. Please avoid to mention a specific commercial vender in the PR description and commit logs.

In databricks runtime, RDD read / write API has some issue for certain storage types that requires the account key, but Dataframe read / write API works.

[dongjoon-hyun]

This is a regression because SparkSession is heavier than SparkContext.

[WeichenXu123]

I think it is not a regression, we are using active spark session but not creating a new one.

[zhengruifeng]

this is what ML scala using for model save and load:

spark/mllib/src/main/scala/org/apache/spark/ml/regression/AFTSurvivalRegression.scala

Lines 493 to 523 in 3ae78c4

	private case class Data(coefficients: Vector, intercept: Double, scale: Double)
	override protected def saveImpl(path: String): Unit = {
	// Save metadata and Params
	DefaultParamsWriter.saveMetadata(instance, path, sc)
	// Save model data: coefficients, intercept, scale
	val data = Data(instance.coefficients, instance.intercept, instance.scale)
	val dataPath = new Path(path, "data").toString
	sparkSession.createDataFrame(Seq(data)).repartition(1).write.parquet(dataPath)
	}
	}
	private class AFTSurvivalRegressionModelReader extends MLReader[AFTSurvivalRegressionModel] {
	/** Checked against metadata when loading model */
	private val className = classOf[AFTSurvivalRegressionModel].getName
	override def load(path: String): AFTSurvivalRegressionModel = {
	val metadata = DefaultParamsReader.loadMetadata(path, sc, className)
	val dataPath = new Path(path, "data").toString
	val data = sparkSession.read.parquet(dataPath)
	val Row(coefficients: Vector, intercept: Double, scale: Double) =
	MLUtils.convertVectorColumnsToML(data, "coefficients")
	.select("coefficients", "intercept", "scale")
	.head()
	val model = new AFTSurvivalRegressionModel(metadata.uid, coefficients, intercept, scale)
	metadata.getAndSetParams(model)
	model
	}

for this R wrapper, I think it is fine to do it in the same way

[WeichenXu123]

Hi, @WeichenXu123 and @zhengruifeng . This is Apache Software Foundation repository. Please avoid to mention a specific commercial vender in the PR description and commit logs.

In databricks runtime, RDD read / write API has some issue for certain storage types that requires the account key, but Dataframe read / write API works.

Oh my bad. I updated PR description. Thanks for your reminder.

[HyukjinKwon]

Can we avoid repartition?

[zhengruifeng]

No, the single partition is needed to avoid creating empty files, even though there is only one row:

spark.range(1).select(sf.lit("123")).write.text("/tmp/test-test")

(base) ➜  runtime git:(master) ls /tmp/test-test
_SUCCESS                                                 part-00011-e548682e-21a5-41ea-93bc-0089efccbf4f-c000.txt
part-00000-e548682e-21a5-41ea-93bc-0089efccbf4f-c000.txt
(base) ➜  runtime git:(master) wc -l /tmp/test-test/*
       0 /tmp/test-test/_SUCCESS
       0 /tmp/test-test/part-00000-e548682e-21a5-41ea-93bc-0089efccbf4f-c000.txt
       1 /tmp/test-test/part-00011-e548682e-21a5-41ea-93bc-0089efccbf4f-c000.txt
       1 total

But I think we can use coalesce(1) to avoid this small shuffle

[HyukjinKwon]

This is single row, and will produce one partition, see LocalTableScanExec.

spark.createDataFrame(Seq(Tuple1("a"))).write.text("/tmp/abc")

write single file.

[HyukjinKwon]

Made a followup: #47341

[zhengruifeng]

ah, I see, this LocalTableScanExec ganrantees single row -> single partition.

It will be better if we can specify the numPartitions somewhere

[HyukjinKwon]

Yeah I think so but removing it should work for now.

[HyukjinKwon]

Reverted per the dev mailing list discussion

[dongjoon-hyun]

Thank you for reverting this, Hyukjin.