如何在 Pyspark 中高效地创建具有每维度总计的多维交叉表？

Question

TL;DR：还有比这更好的方法吗？

columns = ['sex', 'class', 'survived'] # for many columns
grouped_crosstab = sdf.groupBy(*columns).count()

for column in columns:
    grouped_crosstab = grouped_crosstab.join(
        grouped_crosstab.groupBy(column).agg(F.sum('count').alias(f'{column}_total')),
        column,
        'left')

---

问题设置

在 Pyspark 中，您可以在

crosstab

上使用

DataFrame

方法来获取数据的二维交叉表。同样，

groupBy

方法可以返回数据的多维“交叉表”，尽管它采用又高又瘦的格式。

例如：

columns = ['x', 'y', 'z'] # columns are assumed to be rather low in cardinality, such as categorical values, not continuous values
two_dimensional_crosstab = df.crosstab(columns[0], columns[1]) # only compares 'x' and 'y'
multi_dimensional_view = df.groupBy(*columns).count() # compares 'x', 'y', and 'z'

让我们用一些示例数据来可视化这一点

import seaborn
df = seaborn.load_dataset('titanic')
sdf = spark.createDataFrame(df) # how to setup a spark context is outside the scope of this question

数据如下：

让我们在

sex

和

class

上创建一个二维交叉表，然后使用

crosstab

和

groupBy

函数来展示这两种方法的比较：

two_d_crosstab = sdf.crosstab('sex', 'class')
grouped_crosstab = sdf.groupBy('sex', 'class').count()

这些数据框看起来像这样：

与

crosstab

不同，

groupBy

方法可以很好地推广到多列，但必须注意表格的格式。

列和行总计

出于统计目的（例如调查排名），通常希望在交叉表上显示行和列的总和。在二维情况下，人们可以通过这种（诚然是复杂的）方法来获取此类信息：

index_column = two_d_crosstab.columns[0]
col_list = two_d_crosstab.columns[1:]
two_d_crosstab = two_d_crosstab.withColumn('column_total', sum([F.col(c) for c in col_list]))
transposed_df = two_d_crosstab.pandas_api()\
    .set_index(index_column)\
    .T.reset_index()\
    .rename(columns = {'index':index_column})\
    .to_spark()
col_list = transposed_df.columns[1:]
two_d_crosstab = transposed_df.withColumn('row_total', sum([F.col(c) for c in col_list]))

two_d_crosstab

看起来像这样：

多维度总计

如何在多维交叉表上进行这一计算？

这是我尝试过的：

sex_tot = grouped_crosstab.groupBy('sex').agg(F.sum('count').alias('sex_total'))
class_tot = grouped_crosstab.groupBy('class').agg(F.sum('count').alias('class_total'))
grouped_crosstab = grouped_crosstab.join(sex_tot, 'sex', 'left').join(class_tot, 'class', 'left')

输出如下所示：

让我们添加

survived

作为第三个维度：

columns = ['sex', 'class', 'survived']
grouped_crosstab = sdf.groupBy(*columns).count()

for column in columns:
    grouped_crosstab = grouped_crosstab.join(
        grouped_crosstab.groupBy(column).agg(F.sum('count').alias(f'{column}_total')),
        column,
        'left')

看起来像这样：

请注意，输出中有很多重复信息。随着列数的增加，分组和连接操作的数量也会增加，因此这将变得相当笨重，特别是在具有数百万行的大型数据帧上。

有更好的（更可扩展的）方法吗？

Answer 1

我尝试的解决方案（我在写这个问题时发现的）是使用

cube

方法而不是

groupBy

。从二维情况开始：

columns = ['sex', 'class']
cubed_crosstab = sdf.cube(*columns).count()

产品：

其中，当显示

null

时，计数被视为“不关心此列中的值”。因此，在屏幕截图的第三行，491 的计数适用于任何性别的三等舱乘客，第 5 行生成整个数据框的总数。请注意，这些计数与之前生成的计数一致。

推广到多个维度是微不足道的：

columns = ['sex', 'class', 'survived']
cubed_crosstab = sdf.cube(*columns).count()

此答案的局限性

当两列中都有空值时，答案就会变得不明确。例如：

columns = ['deck', 'embarked']
cubed_crosstab = sdf.cube(*columns).count()

存在数量不同的重复条目。 Null 可以表示“不关心”状态，也可以表示数据中的合法值。建议在

cubeing

（这是一个词吗？）之前将 null 重新编码为另一个值，以避免这种歧义。