Scale a column of SparseVectors without UDF - scala

Question

I have a dataframe with a feature column of SparseVector. I need to scale each row by a scalar. I have a working implementation below that employs a UDF. The following depicts the original and scaled feature columns:
+-------------------+-------+-------------------+
| features|weights| scaledFeatures|
+-------------------+-------+-------------------+
|(6,[0,1],[0.5,1.0])| 1.0|(6,[0,1],[0.5,1.0])|
|(6,[2,3],[1.5,2.0])| 2.0|(6,[2,3],[3.0,4.0])|
|(6,[4,5],[0.5,1.0])| 3.0|(6,[4,5],[1.5,3.0])|
+-------------------+-------+-------------------+
Is there a way to do this using Spark's native, and optimized, methods instead of a UDF?
Similarly, is there a Spark-native way to scale a SparseVector by a scalar? See the line below the "Scale the SparseVector" comment in the UDF defined below.
import org.apache.spark.ml.linalg.SparseVector
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.{col, udf}
// Scaling a SparseVector column
val data = Array(
(new SparseVector(6, Array(0,1), Array(0.5, 1.0)), 1.0),
(new SparseVector(6, Array(2,3), Array(1.5, 2.0)), 2.0),
(new SparseVector(6, Array(4,5), Array(0.5, 1.0)), 3.0)
)
val df = spark.createDataFrame(data).toDF("features", "weights")
val scaleUDF = udf((sv: SparseVector, w: Double) => {
// Scale the SparseVector
val unzipped = sv.indices.zip(sv.values).map(iv => (iv._1, iv._2*w)).unzip
new SparseVector(sv.size, unzipped._1, unzipped._2)
})
val scaledDF = df.withColumn("scaledFeatures", scaleUDF(col("features"), col("weights")))
scaledDF.show()
+-------------------+-------+-------------------+
| features|weights| scaledFeatures|
+-------------------+-------+-------------------+
|(6,[0,1],[0.5,1.0])| 1.0|(6,[0,1],[0.5,1.0])|
|(6,[2,3],[1.5,2.0])| 2.0|(6,[2,3],[3.0,4.0])|
|(6,[4,5],[0.5,1.0])| 3.0|(6,[4,5],[1.5,3.0])|
+-------------------+-------+-------------------+