How to remove all the subset from a column except few based on the other column in Pyspark?

user_17363735 0
Viewed 47

I have a pyspark data frame with a column of lists(column a) and another column with numbers(column b), I want to retain all are supersets rows, and subsets that have greater value in column b than their supersets.

For Example,

Input data frame:

Column a = ([A,B,C],[A,C],[B,C],[J,S,K],[J,S],[J,K])
Column b = (10,15,7,8,9,8)

Expected Outcome:

Column a = ([A,B,C],[A,C],[J,S,K],[J,S])
Column b = (10,15,8,9)

Here [B,C] and [A,C] are subsets of [A,B,C] but we only retain [A,C] because this subset has 15 in column b which is greater than 10 the supersets([A,B,C]) value in column b. Similarly, the superset [J,S,K] is retained along with its subset [J,S] because its value in column b is greater than the superset column b value.

pyspark apache-spark-sql apache-spark
1 Answers

You can use a self left_anti join to filter rows that satisfies that condition.

You'll need to have an ID column in your dataframe, here I'm using monotonically_increasing_id function to generate an ID for each row:

import pyspark.sql.functions as F

df = df.withColumn("ID", F.monotonically_increasing_id())
df.show() 

#+---------+---+-----------+
#|        a|  b|         ID|
#+---------+---+-----------+
#|[A, B, C]| 10| 8589934592|
#|   [A, C]| 15|17179869184|
#|   [B, C]|  7|25769803776|
#|[J, S, K]|  8|42949672960|
#|   [J, S]|  9|51539607552|
#|   [J, K]|  8|60129542144|
#+---------+---+-----------+

Now, to verify an array arr1 is a subset of another array arr2, you can use array_intersect and size functions size(array_intersect(arr1, arr2)) = size(arr1):

df_result = df.alias("df1").join(
    df.alias("df2"),
    (
        (F.size(F.array_intersect("df1.a", "df2.a")) == F.size("df1.a"))
        & (F.col("df1.b") <= F.col("df2.b"))
        & (F.col("df1.ID") != F.col("df2.ID")) # not the same row
    ),
    "left_anti"
).drop("ID")

df_result.show()

#+---------+---+
#|        a|  b|
#+---------+---+
#|[A, B, C]| 10|
#|   [A, C]| 15|
#|[J, S, K]|  8|
#|   [J, S]|  9|
#+---------+---+
user_13865510
Related