How to filter out some specific rows from a PySpark dataframe depending on an interable (Set)?

user_11891255 0
Viewed 109

I have a dataframe like this:

FLAG DATA
S D1
S D2
S D3
S SP
N D5
N D2
N D6
N SP

Now I need to follow some instructions. They are:

  1. Get dataframe where FLAG is S.
  2. Get dataframe where FLAG is N.
  3. On a new dataframe, first add all data from dataframe1 (data with FLAG S). Then, take only those rows from the dataframe2 (data with FLAG N), whose DATA value is not in dataframe1.
  4. The only exception of point 3 is where DATA value is SP. If datafram2 has SP, it will be added to the new dataframe. For SP, no need to check the dataframe1 for addition.

Illustrations:

  1. DataFrame 1
FLAG DATA
S D1
S D2
S D3
S SP
  1. DataFrame 2
FLAG DATA
N D5
N D2
N D6
N SP
  1. DataFrame 3
FLAG DATA
S D1
S D2
S D3
S SP
N D5
N D6
N SP

N D2 is eliminated because DATA D2 is already present for FLAG S. But, SP is present for both N and S.

I tried this code snippet:

import pyspark.sql.functions as F    

all_data = final_df.select("FLAG", "DATA")
sf_DATA_set = set()
for data in all_data:
    if data["FLAG"] == 'S' and data["DATA"] != 'SP':
        sf_DATA_set.add(data["DATA"]) # Make a set of all DATA where FLAG is S except SP
    
df1 = final_df.filter("FLAG='S'")
df2 = final_df.filter("FLAG='N'")
   
df4 = df2.filter(~F.col("DATA").isin(tuple(sf_DATA_set)) # Getting only those fields from df2 whose DATA value isn't present in sf_DATA_set.
final_df = df4.unionByName(df1)
display(final_df.select(['FLAG', 'DATA']))

But this part isn't working properly.

df4 = df2.filter(~F.col("DATA").isin(tuple(sf_DATA_set))

How can I filter out the data from datafram2?

pyspark python-3.x
1 Answers

You can create a list of DATA values where FLAG == 'S' and DATA != 'SP' and then filter those rows out by negating .isin

import pyspark.sql.functions as F

def my_function(df):
  """
  Input: PySpark DataFrame
  Output: three PySpark DataFrames
  """
  
  df1 = df.filter(F.col('FLAG') == 'S')
  df2 = df.filter(F.col('FLAG') == 'N')
  
  # make a list of all DATA where FLAG is S except DATA is SP
  sf_data_list = all_data \
    .filter((F.col('FLAG') == 'S') & (F.col('DATA') != 'SP')) \
    .select('DATA') \
    .distinct() \
    .rdd.flatMap(lambda x: x) \
    .collect()
  
  df3 = df2 \
    .filter(~F.col('DATA').isin(sf_data_list)) \
    .unionByName(df1)
  
  return df1, df2, df3

Output

my_function(all_data)[2].show()

+----+----+
|FLAG|DATA|
+----+----+
|   N|  D5|
|   N|  D6|
|   N|  SP|
|   S|  D1|
|   S|  D2|
|   S|  D3|
|   S|  SP|
+----+----+
user_74654620
Related