How to remove rows in a spark dataset on the basis of count of a specific group

user_9031512 0
Viewed 1217

I have a dataframe in Spark- enter image description here

I want to filter the rows on the basic of count of Person for a particular State. What I mean by that is, I only want to keep the rows for those State which have Person count of >1.

So the desired output would look like-enter image description here

The desired output would have rows for State S2 and S4 removed because they had Person count of less than 2.

It would be great if someone could help me with this. Thanks in advance : )

pyspark apache-spark-sql apache-spark
2 Answers

You can use Count of Person over the window and filter out the rows as below.

from pyspark.sql.window import Window
from pyspark.sql import functions as F

df = spark.createDataFrame([
    ("S1", "S1_P1"),
    ("S1", "S1_P2"),
    ("S1", "S1_P3"),
    ("S2", "S2_P1"),
    ("S3", "S3_P1"),
    ("S3", "S3_P2"),
    ("S4", "S4_P1")
],["State", "Person"])

window = Window.partitionBy("State")

df.withColumn("count", F.count("Person").over(window))\
    .filter(F.col("count") > 1).drop("count").show()

Output:

+-----+------+
|State|Person|
+-----+------+
|   S3| S3_P1|
|   S3| S3_P2|
|   S1| S1_P1|
|   S1| S1_P2|
|   S1| S1_P3|
+-----+------+
user_65514260

One alternative to the response of @koiralo would be to use a groupby and an inner join. Which approach Window vs groupby depends on the cardinality of the grouping columns. The rule of thumb: If the groupby results in a small number of rows it should be used. Otherwise if the cardinality of the grouping column is high a window function is the better approach.

from pyspark.sql import functions as F

df = spark.createDataFrame([
    ("S1", "S1_P1"),
    ("S1", "S1_P2"),
    ("S1", "S1_P3"),
    ("S2", "S2_P1"),
    ("S3", "S3_P1"),
    ("S3", "S3_P2"),
    ("S4", "S4_P1")
],["State", "Person"])

states_grouped = df.groupBy(["State"]).agg(F.count("State").alias("count"))
states_relevant = states_grouped.filter(F.col("count") > 1).drop("count")

df_filtered = df.join(states_relevant, "State", "inner")
user_143387160
Related