如何在scala中将withcolumn编写为可重用代码

Posted 2023-03-23

【中文标题】如何在scala中将withcolumn编写为可重用代码

【英文标题】：How to write withcolumn as reusable code in scala

【发布时间】：2022-01-22 07:41:03

【问题描述】：

以下是我的代码，我如何将 withcolumn 的最佳代码编写为可重用代码，因为条件非常相似，我也可能在其他数据帧中使用 withcolumn。

    val sample = df1.alias("a").join(df2.alias("b")
      ,Seq("ReportingSetID","ProductLine","PlanID","ReportType","NumeratorID","MemberID","ServDate"))
      .withColumn("IsNumeratorSupplemental",
        when (col("a.IsNumerator")===1 && col("b.IsNumerator")===0,1).otherwise(0))
      .withColumn("IsExclusionSupplemental",
        when (col("a.IsExclusion")===1 && col("b.IsExclusion")===0,1).otherwise(0))
      .withColumn("IsSubExclusionSupplemental",
        when (col("a.IsSubExclusion")===1 && col("b.IsSubExclusion")===0,1).otherwise(0))
      .withColumn("IsRequiredExclusionSupplemental",
        when (col("a.IsRequiredExclusion")===1 && col("b.IsRequiredExclusion")===0,1).otherwise(0))
      .filter(!col("b.MeasureID").isin("44","70"))

【参考方案1】：

您可以编写一个简单的方法，该方法接受 DataFrame 并像这样添加列：

def withSupplemental(columnName : String)(df: DataFrame) = 
    df.withColumn(s"$columnNameSupplemental",
        when (col(s"a.$columnName")===1 && col(s"b.$columnName")===0,1).otherwise(0))

如果您觉得需要，这也可以成为一个扩展方法。