Pyspark dataframe.limit很慢

Question

我正在尝试使用大型数据集，但只使用其中的一小部分。每个操作都花费很长时间，我想看看数据帧的head或limit。

例如，我调用UDF（用户定义的函数）添加一列，但是我只关心第一个行（例如10行）。

sum_cols = F.udf(lambda x:x[0] + x[1], IntegerType())
df_with_sum = df.limit(10).withColumn('C',sum_cols(F.array('A','B')))

但是，如果我不使用limit，这仍然需要花费很长时间。

Answer 1

如果您首先处理10行，我认为最好创建一个新的df并将其缓存

df2 = df.limit(10).cache()
df_with_sum = df2.withColumn('C',sum_cols(F.array('A','B')))

Answer 2

limit首先将尝试从单个分区获取所需的数据。如果它没有在一个分区中获取全部数据，那么它将从下一分区中获取剩余数据。

因此，请使用df.rdd.getNumPartition检查您有多少个分区

为了证明这一点，我建议首先将df合并到一个分区并进行限制。您可以看到此时间限制更快，因为它从一个分区中过滤数据