Spark DataFrame 新增一列的四种方法

雨点打透心脏的1/2处 2023-03-01 06:10 58阅读 0赞

作为一个学习 Python spark 的菜鸟，刚开始学习 dataframe 的多样化处理有些新奇，对于新增一列的方法，经过多方查询学习，总结了如下四种常用方法，分享给大家。虽然以下语言是用Scala写的，但使用Python的原理也是一样，

*  **方法一：利用 createDataFrame 方法，新增列的过程包含在构建 rdd 和 schema 中**
 *  **方法二：利用 withColumn 方法，新增列的过程包含在 udf 函数中**
 *  **方法三：利用 SQL 代码，新增列的过程直接写入 SQL 代码中**
 *  **方法四：以上三种是增加一个有判断的列，如果想要增加一列唯一序号，可以使用 monotonically\_increasing\_id**

### 代码块 ###

//dataframe新增一列方法1，利用createDataFrame方法
    val trdd = input.select(targetColumns).rdd.map(x=>{
      if (x.get(0).toString().toDouble > critValueR || x.get(0).toString().toDouble < critValueL) 
        Row(x.get(0).toString().toDouble,"F")
      else Row(x.get(0).toString().toDouble,"T")      
      })      
    val schema = input.select(targetColumns).schema.add("flag", StringType, true)
    val sample3 = ss.createDataFrame(trdd, schema).distinct().withColumnRenamed(targetColumns, "idx")
    
    
    //dataframe新增一列方法2，withColumn
    val code :(Int => String) = (arg: Int) => {if (arg > critValueR || arg < critValueL) "F" else "T"}
    val addCol = udf(code)
    val sample3 = input.select(targetColumns).withColumn("flag", addCol(input(targetColumns)))
    .withColumnRenamed(targetColumns, "idx")
    
    
    //dataframe新增一列方法3，写SQL语句
    input.select(targetColumns).createOrReplaceTempView("tmp")
    val sample3 = ss.sqlContext.sql("select distinct "+targetColname+
        " as idx,case when "+targetColname+">"+critValueR+" then 'F'"+
        " when "+targetColname+"<"+critValueL+" then 'F' else 'T' end as flag from tmp")
    
    
    //dataframe新增一列方法3，添加序号列
    import org.apache.spark.sql.functions.monotonically_increasing_id
    val inputnew = input.withColumn("idx", monotonically_increasing_id)