spark scala实现 pandas的melt

import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions.{array, col, explode, lit, struct}

object SparkPandas {
    def sparkMelt(frame:DataFrame, idVars:List[String], valueVars:List[String],
                  varName:String, valueName:String):DataFrame={
        val _idVars:List[String] = idVars.filter(x=>frame.columns.contains(x))
        val _valueVars:Seq[String] = valueVars.filter(x=>frame.columns.contains(x))
        var _df:DataFrame = frame
        for (v <- _valueVars){
            _df = _df.withColumn(v,struct(lit(v).alias("key"),col(v).alias("value")))
        }
        _df = _df.withColumn("_zip",array(_valueVars.map(x=>col(x)):_*))
            .withColumn("_keyValue",explode(col("_zip")))
            .withColumn(varName,col("_keyValue").getField("key"))
            .withColumn(valueName,col("_keyValue").getField("value"))
        val oCol:List[String] = _idVars:+varName:+valueName
        _df.selectExpr(oCol:_*).where(s"${valueName} is not null")
    }
}

参考https://www.yyshao.com/2021/07/28/spark-melt/ 里的python代码

后面发现也不必这么复杂,这样就行了。

val cols = spark.table(tableName).columns.filter(x=> x!="msisdn" && x!="dt").map(x=>s"'$x',$x")
spark.sql(s"""
               select msisdn,value,label,'20240428' dt from
               (
              select msisdn,stack(${cols.length},${cols.mkString(",")}) as (label,value) 
              from $tableName
              where dt='20240428'
              ) where value is not null
              """).show()