返回映射函数使用的UDF中的所有列+更多列

Question

我正在使用map函数来生成一个新列，其值取决于数据框中已经存在的列的结果。

def computeTechFields(row):
    if row.col1!=VALUE_TO_COMPARE:
            tech1=0
    else:
            tech1=1
    return (row.col1, row.col2, row.col3, tech1)
delta2rdd = delta.map(computeTechFields)

问题是，我的主数据框架有超过150列，我必须返回与地图函数，所以最后我有这样的东西:

return (row.col1, row.col2, row.col3, row.col4, row.col5, row.col6, row.col7, row.col8, row.col9, row.col10, row.col11, row.col12, row.col13, row.col14, row.col15, row.col16, row.col17, row.col18 ..... row.col149, row.col150, row.col151, tech1)

正如你所看到的，它写起来很长，读起来很困难。所以我试着这样做:

return (row.*, tech1)

当然，它没有工作。

我知道"withColumn"函数存在，但我不太了解它的性能，无论如何也无法使它工作。

编辑(withColumn函数发生了什么):

def computeTech1(row):
    if row.col1!=VALUE_TO_COMPARE:
            tech1=0
    else:
            tech1=1
    return tech1
delta2 = delta.withColumn("tech1", computeTech1)

它给了我这个错误:

AssertionError: col should be Column

我试着这样做:

return col(tech1)

错误是相同的

我也试过:

delta2 = delta.withColumn("tech1", col(computeTech1))

这个错误是:

AttributeError: 'function' object has no attribute '_get_object_id'

编辑结束

所以我的问题是，我怎么能返回所有的列+几个在我的UDF中使用的映射函数?

谢谢!

Answer 1

对于Python来说不是很坚定，所以人们可能会纠正我这里的语法，但一般的想法是让你的函数成为一个UDF，其中一列作为输入，然后在withColumn中调用它。我在这里使用了lambda，但经过一些调整，它应该也可以与函数一起工作。

from pyspark.sql.functions import udf
computeTech1UDF = udf(
   lambda col: 0 if col != VALUE_TO_COMPARE else 1, IntegerType())
delta2 = delta.withColumn("tech1", computeTech1UDF(col1))

由于没有为withColumn提供列表达式(请参阅http://spark.apache.org/docs/1.6.0/api/python/pyspark.sql.html#pyspark.sql.DataFrame.withColumn)，所以您所尝试的操作不起作用。使用UDF包装器可以实现这一点。