將新列添加到 Python 列表中的 PySpark DataFrame

Question

我有一個清單：

dates = [2017, 2018, 2018, 2018, 2019, 2019, 2019, 2020, 2020, 2020]

我嘗試添加的 dataframe 長度相同（沒有問題）。

我試過了：

df = df.withColumn("YEARS", dates)
Error: Column needs to be col

我也試過：

df = df.withColumn("YEARS", f.lit(dates))

但這也行不通。

我看到了這個問題： How to add a constant column in a Spark DataFrame?

但是對於這種情況沒有任何用處。

更新：預期的結果是：

df_columns...   | dates_from_list
---------------------------------
original_df_data| 2017
original_df_data| 2018
original_df_data| 2018
original_df_data| 2018
original_df_data| 2019
original_df_data| 2019
original_df_data| 2019
original_df_data| 2020
original_df_data| 2020
original_df_data| 2020

Answer 1

您的錯誤來自您需要將Column object 傳遞給withColumn的事實。

這里有兩種方法可以將日期添加為 Spark DataFrame上的新列（使用每個記錄的順序進行連接），具體取決於日期數據的大小。

1）如果你操作一個小數據集

實現它的一種簡潔方法是將 UDF 應用於單調遞增的 id：

from pyspark.sql.functions import udf, monotonically_increasing_id

df = [...]  # 10 records

dates = [2017, 2018, 2018, 2018, 2019, 2019, 2019, 2020, 2020, 2020]

df = df.repartition(1).withColumn(
    "YEARS", 
    udf(lambda id: dates[id])(monotonically_increasing_id()))

df.show()

輸出：

+---+-----+
|...|YEARS|
+---+-----+
|...| 2017|
|...| 2018|
|...| 2018|
|...| 2018|
|...| 2019|
|...| 2019|
|...| 2019|
|...| 2020|
|...| 2020|
|...| 2020|
+---+-----+

注意： .repartition(1)確保生成的 id 是連續的。 如果您有另一種方法將 map 每個記錄轉換為dates值（如先前構建的 id 列），則可以避免重新分區到單個分區。 在這個用例中，正如我們預期的 Python 列表 object 很小，這意味着您的 DataFrame 也很小，所以這個重新分區不是什么大問題。

/!\ 如果 dataframe 和 python 列表太大，為什么它不會擴展：

• 需要對 dataframe 進行重新分區，從而導致昂貴的洗牌/交換
• .repartition(1)可能會導致生成一個非常大的分區，該分區的處理速度可能非常慢（因為它很大，而且如果它不適合執行 memory 它可能意味着許多額外的磁盤 I/O 會溢出 RDD塊到磁盤），或使用OutOfMemoryError使作業崩潰。
• python 列表由 udf 捕獲（通過 lambda 閉包），這意味着它將被廣播到集群的每個執行程序

2) 如果您操作的數據集大小 > 百萬行

這是另一種方法，可以通過使用 pandas 操作 id 和日期列來更好地處理數百萬行，並避免對 Spark DataFrame進行任何重新分區。

可以這樣做：

import pandas as pd
from pyspark.sql.functions import monotonically_increasing_id
from pyspark.sql.session import SparkSession

spark = SparkSession.builder.getOrCreate()

# some spark DataFrame of length N
df = [...]  

# generate monotically increasing ids (not consecutive) without repartitioning the Spark DataFrame.
df = df.withColumn("id", monotonically_increasing_id())

# get generated ids (not consecutive) as a mono-column pandas DataFrame
spark_df_ids = df.select("id").toPandas()

# some python list of length N
dates = [2017, 2018, 2018, 2018, 2019, ..., 2019, 2019, 2020, 2020, 2020]

# build pandas DataFrame from dates
dates_pandas_df = pd.DataFrame(dates, columns=["YEARS"])

# append the id column to the dates in pandas
dates_and_ids_pandas_df = dates_pandas_df.join(spark_df_ids)

# convert from pandas DataFrame to spark DataFrame
dates_and_ids_spark_df = spark.createDataFrame(dates_and_ids_pandas_df)

# Perform the final adding of the dates column to the Spark DataFrame with a join in Spark
df.join(dates_and_ids_spark_df, ["id"]).show()

重要提示： 使用 Apache 箭頭可以更快地從和到 pandas 的轉換

Answer 2

你可以試試這個：

dates = [2017, 2018, 2018, 2018, 2019, 2019, 2019, 2020, 2020, 2020]
df = spark.createDataFrame([Row(a=1)])
df = df.withColumn("YEARS",  array( [lit(x) for x in dates]  ))


df.show(truncate=False)
+---+------------------------------------------------------------+
|a  |YEARS                                                       |
+---+------------------------------------------------------------+
|1  |[2017, 2018, 2018, 2018, 2019, 2019, 2019, 2020, 2020, 2020]|
+---+------------------------------------------------------------+

df.select("a", explode("YEARS")).show()
+---+----+
|  a| col|
+---+----+
|  1|2017|
|  1|2018|
|  1|2018|
|  1|2018|
|  1|2019|
|  1|2019|
|  1|2019|
|  1|2020|
|  1|2020|
|  1|2020|
+---+----+