FeatureVectorQuantiles.scala

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import org.apache.spark.sql.DataFrame
import org.apache.spark.ml.linalg.Vector
import org.apache.spark.sql.expressions.UserDefinedFunction


object FeatureVectorQuantiles {

  // Simple helper to convert vector to array<double>
  val vecToSeq: UserDefinedFunction = udf((v: Vector) => v.toArray)
  
  def calculateBuckets (n: Integer): Array[Double] = {
    (0 until n + 1).map {x => (1.0/n) * x}.toArray
  }

  def Calculate (df: DataFrame, featureColName: String = "features",
                 numBuckets: Int = 2,   relativeError: Double = 0.001): Array[Array[Double]] = {

    val spark: SparkSession = SparkSession.builder().getOrCreate()
    import spark.implicits._

    val nFeatures = df.select(featureColName).first.getAs[Vector](0).size

    val exprs = (0 until nFeatures).map(i => $"_tmp".getItem(i).alias(s"feature_$i"))

    val unpackedDF = df.select(vecToSeq(col(featureColName)).alias("_tmp")).select(exprs:_*)

    unpackedDF.stat.approxQuantile(unpackedDF.columns, calculateBuckets(numBuckets), relativeError)

  }
}