有一个缺少值的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|
+----------+---------+
我想创建一个新列,如下所示:
我可以构建一个简单的函数:
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
值的首选方法?
答案 0 :(得分:6)
一如既往,最好直接在本机表示上操作,而不是将数据提取到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")).
答案 1 :(得分:3)
您只需使用接收两个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|
#+----------+---------+------------+