如何有效地比较两个 1x1 Spark DataFrames？

Question

I have a use case where I need to efficiently compare the average values of two columns. More specifically, I want to find the percent change between two values (which involves some algebra and comparisons between the numbers).

To do this, I start by grouping and aggregating the average over the column that I want, which gives me a DataFrame with a single float in it (ie a DataFrame with one cell). Now, what I originally did was grab this value from the DataFrame using:

my_df.head()[0]

but it turns out that this is very slow (several seconds to bring this DataFrame to the driver, I believe). I am unsure how else to get this value, or how to compare it with another average value (which is aggregated/obtained in the same way). Side note, .collect()[0][0] also has this speed issue.

Is there a way to get this average value without such a slow runtime, or otherwise compare the two average values in these separate DataFrames?

Answer 1

If your two dataframes have a common key, you can join the average column on the key. If they don't and you just need the average across the complete dataframe you can add a key column with a constant value to both dataframes using F.lit .

df = df.withColumn("key", F.lit(1))

However, as soon as you want to actually see the results, they need to be collected anyway. This will take some time even for really small dataframes. Spark will be able to compute an average for 1 billion rows very quickly compared to tools like Pandas. This is because it builds up an infrastructure that allows it to compute subtasks of the problem you give it in a distributed fashion. Building this infrastructure takes Spark some time. If you want to compute the average of 3 rows it is not worth to build such a complex infrastructure. That means don't use Spark for things like that or live with the fact that Spark is slower on small data sets than Pandas. To speed things up a little bit for small dataframes in Spark cache them (that is load them into memory) before you work with them.

df = df.cache()

dfA = spark.createDataFrame([
    [1, 1],
    [1, 2],
    [1, 3]
], ["key", "a"])

dfB = spark.createDataFrame([
    [1, 2],
    [1, 3],
    [1, 4]
], ["key", "b"])

dfC = dfA.join(dfB, on=["key"], how="inner")
(
    dfC
    .groupBy("key")
    .agg(
        F.avg("a").alias("avg_a"), 
        F.avg("b").alias("avg_b")
    )
    .withColumn("avg_is_equal", F.expr("avg_a = avg_b"))
    .show()
)

Output

+---+-----+-----+------------+
|key|avg_a|avg_b|avg_is_equal|
+---+-----+-----+------------+
|  1|  2.0|  3.0|       false|
+---+-----+-----+------------+