使用scala将数据框列名作为参数传递给函数？答案

【问题标题】：Pass dataframe column name as parameter to the function using scala?使用scala将数据框列名作为参数传递给函数？
【发布时间】：2021-03-10 16:50:07
【问题描述】：


I have a function that adds 2 columns:

def sum_num (num1: Int, num2: Int): Int = {
    return num1 + num2
}

我有一个具有以下值的数据框 df

+----+----+----+
|col1|col2|col3|
+----+----+----+
|1   |2   |5   |
|7   |4   |4   |
+----+----+----+

我想添加一列并将列名传递给函数，但下面的代码不起作用。它给出了错误发现所需的列是 Int

val newdf = df.withColumn("sum_of_cols1", sum_num($col1, $ col2))
              .withColumn("sum_of_cols2", sum_num($col1, $ col3))

【问题讨论】：

this 有帮助吗？
@GuruStron 我见过这个，但不确定如何使用多列创建 udf。我还读到 udf 可能会对性能产生一些影响，因为我正在对十亿条记录进行计算，试图找到其他解决方案！

标签： scala apache-spark databricks

【解决方案1】：

将您的代码更改为：

import spark.implicits._

def sum_num (num1: Column, num2: Column): Column = {
  return num1 + num2
}

val newdf = df.withColumn("sum_of_cols1", sum_num($"col1", $"col2"))
  .withColumn("sum_of_cols2", sum_num($"col1", $"col3"))

您必须对 Spark SQL 列进行操作。您可以对它们进行算术运算。看看能用的operators

【讨论】：

感谢您。对于函数 sum_num，输出是否需要是一列，或者我可以在那里输出 Int，因为我需要该值来进行其他计算
您可以使用可用的运算符直接对列进行操作。不用担心，如果您需要其他特定功能，您可以添加一个 udf，它也可以作为输入列。