Filter dataframe with a window which contains at least one particular value

user_3955022 0
Viewed 491

Suppose you have a pyspark dataframe with two columns, id and val.
You want to filter rows, where a single id-window has at least one val entry with a certain value.

E.g. if we want to get rows, where an id-window has at least one 5 value in the val column:

# Input:

df = spark.createDataFrame(
    [(1, '0'), 
     (1, '5'),
     (2, '2'),
     (2, '5'),    
     (2, '5'),
     (3, '1'),
     (3, '0'),],
     ['id', 'val']
)

# Desired output:

+---+----+
| id| val|
+---+----+
|  1|   0|
|  1|   5|
|  2|   2|
|  2|   5|
|  2|   5|
+---+----+

I thought about using the window function somehow?

pyspark apache-spark
2 Answers

Use array_contains() function from Spark-2.4 then join to df to get only id's that has 5 val in them.

Example:

from pyspark.sql.functions import *
from pyspark.sql.types import *

df1=df.groupBy("id").\
agg(array_contains(collect_set(col("val")).cast("array<int>"),5).alias("has_5")).\
filter(col("has_5")).\
drop('has_5')

df.join(df1,['id'],'inner').show()
#+---+---+
#| id|val|
#+---+---+
#|  1|  0|
#|  1|  5|
#|  2|  2|
#|  2|  5|
#|  2|  5|
#+---+---+

Another way Using Window function:

import sys
from pyspark.sql import *

w=Window.partitionBy("id").orderBy("val").rowsBetween(-sys.maxsize,sys.maxsize)

df.withColumn("has_5",array_contains(collect_set(col("val")).over(w).cast("array<int>"),5)).\
filter(col("has_5")).\
drop("has_5").\
show()
#+---+---+
#| id|val|
#+---+---+
#|  1|  0|
#|  1|  5|
#|  2|  2|
#|  2|  5|
#|  2|  5|
#+---+---+
user_76326950

Try this.

import pyspark.sql.functions as f

value = '5'

df2 = df.groupBy('id').agg(f.collect_list('val').alias('vals')).filter(f.array_contains(f.col('vals'), value)).select('id')
df.join(df2, ['id'], 'inner').show(10, False)

+---+---+
|id |val|
+---+---+
|1  |0  |
|1  |5  |
|2  |2  |
|2  |5  |
|2  |5  |
+---+---+
user_118415710
Related