在 Spark 数据框中创建一个新列，该列是其他列值的列表

Question

我有一个名为“df”的数据框，其结构如下

ID	姓名	lv1	lv2
abb	名称1	40.34	21.56
粑粑	名称2	21.30	67.45
巴巴	名称3	32.45	45.44

在 Pandas 中，我可以使用以下代码创建一个包含 lv1 和 lv2 值列表的新列

cols = ['lv1', 'lv2']
df['new_col'] = df[cols].values.tolist()

由于数据大小导致的内存问题，我现在使用 Databricks 代替（我以前从未使用过）并且需要复制上述内容。 我通过安装我的数据的位置然后加载成功创建了一个 Spark 数据框

file_location = 'dbfs:/mnt/<mountname>/filename.csv'
file_type = "csv"
   
infer_schema = "false"
first_row_is_header = "true"
delimiter = ","

df = spark.read.format(file_type)
  .option("inferSchema", infer_schema)
  .option("header", first_row_is_header)
  .option("sep", delimiter)
  .load(file_location)

display(df)

这会加载数据，但是，我不知道如何完成必要的下一步。 我在 Spark 中找到了一个名为struct的函数，但是，我似乎在 PySpark 中找不到相应的函数。 有什么建议么？

Answer 1

它可能是您正在寻找的array函数。

from pyspark.sql import functions as F
df = spark.createDataFrame(
    [('abb', 'name1', 40.34, 21.56),
     ('bab', 'name2', 21.30, 67.45),
     ('bba', 'name3', 32.45, 45.44)],
    ['ID', 'name', 'lv1', 'lv2'])

df = df.withColumn('new_col', F.array('lv1', 'lv2'))

df.show()
# +---+-----+-----+-----+--------------+
# | ID| name|  lv1|  lv2|       new_col|
# +---+-----+-----+-----+--------------+
# |abb|name1|40.34|21.56|[40.34, 21.56]|
# |bab|name2| 21.3|67.45| [21.3, 67.45]|
# |bba|name3|32.45|45.44|[32.45, 45.44]|
# +---+-----+-----+-----+--------------+