Left anti join in groups - apache-spark

Question

I have a dataframe that has two columns a and b where the values in the b column are a subset of the values in the a column. For instance:
df
+---+---+
| a| b|
+---+---+
| 1| 2|
| 1| 3|
| 2| 1|
| 3| 2|
+---+---+
I’d like to produce a dataframe with columns a and anti_b where the values in the anti_b column are any values from the a column such that a!=anti_b and the row (a,anti_b) does not appear in the original dataframe. So in the above dataframe, the result should be:
anti df
+---+------+
| a|anti_b|
+---+------+
| 3| 1|
| 2| 3|
+---+------+
This can be accomplished with a crossJoin and a call to array_contains, but it’s extremely slow and inefficient. Does anyone know of a better spark idiom to accomplish this, something like anti_join?
Here is the inefficient example using a small dataframe so you can see what I’m after:
df = spark.createDataFrame(pandas.DataFrame(numpy.array(
[[1,2],[1,3],[2,1],[3,2]]),columns=['a','b']))
crossed_df = df.select('a').withColumnRenamed('a','_a').distinct().crossJoin(df.select('a').withColumnRenamed('a','anti_b').distinct()).where(pyspark.sql.functions.col('_a')!=pyspark.sql.functions.col('anti_b'))
anti_df = df.groupBy(
'a'
).agg(
pyspark.sql.functions.collect_list('b').alias('bs')
).join(
crossed_df,
on=((pyspark.sql.functions.col('a')==pyspark.sql.functions.col('_a'))&(~pyspark.sql.functions.expr('array_contains(bs,anti_b)'))),
how='inner'
).select(
'a','anti_b'
)
print('df')
df.show()
print('anti df')
anti_df.show()
Edit: This also works, but it's not much faster:
df = spark.createDataFrame(pandas.DataFrame(numpy.array(
[[1,2],[1,3],[2,1],[3,2]]),columns=['a','b']))
crossed_df = df.select('a').distinct().crossJoin(df.select('a').withColumnRenamed('a','b').distinct()).where(pyspark.sql.functions.col('a')!=pyspark.sql.functions.col('b'))
anti_df = crossed_df.join(
df,
on=['a','b'],
how='left_anti'
)

Answer 1

This should be better than what you have:
from pyspark.sql.functions import collect_set, expr
anti_df = df.groupBy("a").agg(collect_set("b").alias("bs")).alias("l")\
.join(df.alias("r"), on=expr("NOT array_contains(l.bs, r.b)"))\
.where("l.a != r.b")\
.selectExpr("l.a", "r.b AS anti_b")\
anti_df.show()
#+---+------+
#| a|anti_b|
#+---+------+
#| 3| 1|
#| 2| 3|
#+---+------+
If you compare this execution plan with your methods, you'll see that is better (because you can swap out distinct for collect_set), but it still has a Cartesian product.
anti_df.explain()
#== Physical Plan ==
#*(3) Project [a#0, b#294 AS anti_b#308]
#+- CartesianProduct (NOT (a#0 = b#294) && NOT array_contains(bs#288, b#294))
# :- *(1) Filter isnotnull(a#0)
# : +- ObjectHashAggregate(keys=[a#0], functions=[collect_set(b#1, 0, 0)])
# : +- Exchange hashpartitioning(a#0, 200)
# : +- ObjectHashAggregate(keys=[a#0], functions=[partial_collect_set(b#1, 0, 0)])
# : +- Scan ExistingRDD[a#0,b#1]
# +- *(2) Project [b#294]
# +- *(2) Filter isnotnull(b#294)
# +- Scan ExistingRDD[a#293,b#294]
However, I don't think there's any way to avoid the Cartesian product for this particular problem without more information.