Я хочу рекурсивно добавлять столбцы в переменное количество кадров данных PySpark, пока все они не будут использовать одни и те же столбцы. (Добавленные столбцы будут заполнены нулевыми значениями). Ниже функция работает для 2 фреймов данных, и мой вопрос: как это обобщить для любого количества фреймов данных (2, 3 и т. д.)?
def add_missing_col(df_1, df_2):
"""Compare column names in df_1 and df_2
and insert missing columns in df_2 with null values.
"""
missing_cols = [i for i in df_1.schema.names if i not in df_2.schema.names]
for i in missing_cols:
df_2 = df_2.withColumn(i, lit(None).cast(StringType()))
return df_2
Я пробовал с functools.reduce
и определял сигнатуру функции как *dfs
, но не знаю, что делать дальше:
def add_missing_col_r(*dfs):
"""Compare column names in dfs and insert missing columns with null values recursively."""
return reduce(DataFrame.withColumn(lambda i : i for i in DataFrame.schema.names), dfs)
Стоит ли использовать здесь лямбда-функцию или есть лучший подход?
Тестовые кадры данных, которые я использую:
# Test dataframes
df1 = spark.createDataFrame([(1, "foo1", "qux1"),
(2, "bar1", "quux1"),
(3, "baz1", "quuz1")],
("a", "b", "c"))
df2 = spark.createDataFrame([(4, "foo2"), (5, "baz2")], ("a", "c"))
df3 = spark.createDataFrame([("bar3", "bar3", "bar3", "bar3"),
("qux3", "quux3", "quuz3", "corge3"),
("grault3", "garply3", "waldo3", "fred3")
],
("b", "d", "e", "f")
)
Я не уверен, что здесь уместно сокращение. Просто использование обычного python будет в порядке. Если вы хотите, чтобы результирующие столбцы были в правильном порядке, ознакомьтесь с моим предыдущим ответом на другой ваш вопрос.
dfs = [df1, df2, df3]
all_cols = set(sum([i.columns for i in dfs], []))
def add_missing_col_r(dfs):
return_dfs = []
for df in dfs:
missing_cols = all_cols - set(df.columns)
for i in sorted(missing_cols):
df = df.withColumn(i, lit(None).cast(StringType()))
return_dfs.append(df)
return return_dfs
new_dfs = add_missing_col_r(dfs)
[x.show() for x in new_dfs]
+---+----+-----+----+----+----+
| a| b| c| d| e| f|
+---+----+-----+----+----+----+
| 1|foo1| qux1|null|null|null|
| 2|bar1|quux1|null|null|null|
| 3|baz1|quuz1|null|null|null|
+---+----+-----+----+----+----+
+---+----+----+----+----+----+
| a| c| b| d| e| f|
+---+----+----+----+----+----+
| 4|foo2|null|null|null|null|
| 5|baz2|null|null|null|null|
+---+----+----+----+----+----+
+-------+-------+------+------+----+----+
| b| d| e| f| a| c|
+-------+-------+------+------+----+----+
| bar3| bar3| bar3| bar3|null|null|
| qux3| quux3| quuz3|corge3|null|null|
|grault3|garply3|waldo3| fred3|null|null|
+-------+-------+------+------+----+----+