PySpark passing Dataframe as extra parameter to map - scala

Question

I want to parallelize a python list, use a map on that list, and pass a Dataframe to the mapper function also
def output_age_split(df):
ages= [18, 19, 20, 21, 22]
age_dfs= spark.sparkContext.parallelize(ages).map(lambda x: test(x, df)
# Unsure of type of age_dfs, but should be able to split into the smaller dfs like this somehow
return age_dfs[0], age_dfs[1] ...
def test(age, df):
return df.where(col("age")==age)
This results in a pickling error
raise pickle.PicklingError(msg)
_pickle.PicklingError: Could not serialize object: TypeError: can't pickle _thread.RLock objects
How should i parallelize this operation so that i get returned a collection of Dataframes?
EDIT: Sample of df
|age|name|salary|
|---|----|------|
|18 |John|40000 |
|22 |Joseph|60000 |

Answer 1

The issue is ages_dfs is not a dataframe, it's an RDD. Now, when you are applying a map with test function in it (which returns the dataframe), we end up getting into a weird situation where ages_dfs is actually an RDD of type PipelinedRDD which is neither a dataframe nor iterable.
TypeError: 'PipelinedRDD' object is not iterable
You can try the below workaround, where you just iterate on the list instead and create a collection of dataframe and iterate on them at will.
from pyspark.sql.functions import *
def output_age_split(df):
ages = [18, 19, 20, 21, 22]
result = []
for age in ages:
temp_df = test(age, df)
if(not len(temp_df.head(1)) == 0):
result.append(temp_df)
return result
def test(age, df):
return df.where(col("age")==age)
# +---+------+------+
# |age|name |salary|
# +---+------+------+
# |18 |John |40000 |
# |22 |Joseph|60000 |
# +---+------+------+
df = spark.sparkContext.parallelize(
[
(18, "John", 40000),
(22, "Jpseph", 60000)
]
).toDF(["age", "name", "salary"])
df.show()
result = output_age_split(df)
# Output type is: <class 'list'>
print(f"Output type is: {type(result)}")
for r in result:
r.show()
# +---+----+------+
# |age|name|salary|
# +---+----+------+
# | 18|John| 40000|
# +---+----+------+
# +---+------+------+
# |age| name|salary|
# +---+------+------+
# | 22|Jpseph| 60000|
# +---+------+------+
I am also attaching screenshot from my workspace for your reference.
Problem:
Solution: