從PySpark DataFrame中的非空列中選擇值

Question

有一個缺少值的pyspark數據框：

tbl = sc.parallelize([
        Row(first_name='Alice', last_name='Cooper'),             
        Row(first_name='Prince', last_name=None),
        Row(first_name=None, last_name='Lenon')
    ]).toDF()
tbl.show()

這是表格：

  +----------+---------+
  |first_name|last_name|
  +----------+---------+
  |     Alice|   Cooper|
  |    Prince|     null|
  |      null|    Lenon|
  +----------+---------+

我想創建一個新列，如下所示：

• 如果名字是None，則取姓氏
• 如果姓氏為無，請取名字
• 如果他們都在場，請將它們連接起來
• 我們可以放心地假設它們中至少有一個存在

我可以構造一個簡單的函數：

def combine_data(row):
    if row.last_name is None:
        return row.first_name
    elif row.first_name is None:
        return row.last_name
    else:
        return '%s %s' % (row.first_name, row.last_name)
tbl.map(combine_data).collect()

我得到了正確的結果，但我無法將其作為列附加到表中： tbl.withColumn('new_col', tbl.map(combine_data))導致AssertionError: col should be Column

將map結果轉換為Column的最佳方法是什么？ 有沒有一種處理null值的首選方法？

Answer 1

一如既往，最好直接在本機表示上操作，而不是將數據提取到Python：

from pyspark.sql.functions import concat_ws, coalesce, lit, trim

def combine(*cols):
    return trim(concat_ws(" ", *[coalesce(c, lit("")) for c in cols]))

tbl.withColumn("foo", combine("first_name", "last_name")).

Answer 2

您只需要使用接收兩columns作為參數的UDF 。

from pyspark.sql.functions import *
from pyspark.sql import Row

tbl = sc.parallelize([
        Row(first_name='Alice', last_name='Cooper'),             
        Row(first_name='Prince', last_name=None),
        Row(first_name=None, last_name='Lenon')
    ]).toDF()

tbl.show()

def combine(c1, c2):
  if c1 != None and c2 != None:
    return c1 + " " + c2
  elif c1 == None:
    return c2
  else:
    return c1

combineUDF = udf(combine)

expr = [c for c in ["first_name", "last_name"]] + [combineUDF(col("first_name"), col("last_name")).alias("full_name")]

tbl.select(*expr).show()

#+----------+---------+------------+
#|first_name|last_name|   full_name|
#+----------+---------+------------+
#|     Alice|   Cooper|Alice Cooper|
#|    Prince|     null|      Prince|
#|      null|    Lenon|       Lenon|
#+----------+---------+------------+