pyspark 数据框所有列中唯一元素的数量

Question

如何计算 pyspark 数据帧的每一列中唯一元素的数量：

import pandas as pd
from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()
df = pd.DataFrame([[1, 100], [1, 200], [2, 300], [3, 100], [4, 100], [4, 300]], columns=['col1', 'col2'])
df_spark = spark.createDataFrame(df)
print(df_spark.show())
# +----+----+
# |col1|col2|
# +----+----+
# |   1| 100|
# |   1| 200|
# |   2| 300|
# |   3| 100|
# |   4| 100|
# |   4| 300|
# +----+----+

# Some transformations on df_spark here

# How to get a number of unique elements (just a number) in each columns?

我只知道以下非常慢的解决方案，这两行都是在相同的时间内计算出来的：

col1_num_unique = df_spark.select('col1').distinct().count()
col2_num_unique = df_spark.select('col2').distinct().count()

中大约有 1000 万行df_spark。

Answer 1

尝试这个：

from pyspark.sql.functions import col, countDistinct

df_spark.agg(*(countDistinct(col(c)).alias(c) for c in df_spark.columns))

编辑： 正如@pault 建议的那样，这是一项昂贵的操作，您可以使用approx_count_distinct()他建议的操作目前已弃用（spark 版本 >= 2.1）

Answer 2

@Manrique 解决了这个问题，但只有稍微修改的解决方案对我有用：

expression = [countDistinct(c).alias(c) for c in df.columns]
df.select(*expression).show()

Answer 3

这要快得多：

df_spark.select(F.countDistinct("col1")).show()