pyspark iterate N rows from Data Frame to each execution

user_2458922 0
Viewed 296
def fun_1(csv):
     # returns int[] of length = Number of New Lines in String csv

def fun_2(csv): # My WorkArround to Pass one CSV Line at One Time
     return fun_1(csv)[0]

The Input Data Frame is df

+----+----+-----+
|col1|col2|CSVs |
+----+----+-----+
|   1|   a|2,0,1|
|   2|   b|2,0,2|
|   3|   c|2,0,3|
|   4|   a|2,0,1|
|   5|   b|2,0,2|
|   6|   c|2,0,3|
|   7|   a|2,0,1|
+----+----+-----+

Below is a Code Snippet which works but takes long time

from pyspark.sql.functions import udf
from pyspark.sql import functions as sf
funudf = udf(fun_2) # wish it could be fun_1
df=df.withColumn( 'pred' , funudf(sf.col('csv')))

fun_1 , has a memory issue and could only handle 50000 max rows at a time. I wish to use funudf = udf(fun_1) . Hence, how can I split the PySpark DF into segments of 50000 rows , call funudf ->fun_1 . Output has two colunms, 'col1' from the input and 'funudf return value' .

pyspark-dataframes pyspark user-defined-functions python
1 Answers

You can achieve the desired result of forcing PySpark to operate on fixed batches of rows by using the groupByKey method exposed in the RDD API. Using groupByKey will force PySpark to shuffle all the data for a single key to a single executor.

NOTE: for this very same reason using groupByKey is often discouraged because of the network cost.

Strategy:

  1. Add a column that groups your data into the desired size of batches and groupByKey
  2. Define a function that reproduces the logic of your UDF (and also returns an id for joining later on). This operates on pyspark.resultiterable.ResultIterable, the result of groupByKey. Apply function to your groups using mapValues
  3. Convert the resulting RDD into a DataFrame and join back in.

Example:

# Synthesize DF
data = {'_id': range(9), 'group': ['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b', 'c'], 'vals': [2.0*i for i in range(9)]}
df = spark.createDataFrame(pd.DataFrame(data))

df.show()

##
# Step - 1 Convert to rdd and groupByKey to force each group to separate executor
##
kv = df.rdd.map(lambda r: (r.group, [r._id, r.group, r.vals]))
groups = kv.groupByKey()

##
# Step 2 - Calulate function
##

# Dummy function taking 
def mult3(ditr):
    data = ditr.data
    ids = [v[0] for v in data]
    vals = [3*v[2] for v in data]
    return zip(ids, vals)

# run mult3 and flaten results
mv = groups.mapValues(mult3).map(lambda r: r[1]).flatMap(lambda r: r) # rdd[(id, val)]

## 
# Step 3 - Join results back into base DF
## 

# convert results into a DF and join back in
schema = t.StructType([t.StructField('_id', t.LongType()), t.StructField('vals_x_3', t.FloatType())])
df_vals = spark.createDataFrame(mv, schema)
joined = df.join(df_vals, '_id')

joined.show()

>>>

+---+-----+----+
|_id|group|vals|
+---+-----+----+
|  0|    a| 0.0|
|  1|    b| 2.0|
|  2|    c| 4.0|
|  3|    a| 6.0|
|  4|    b| 8.0|
|  5|    c|10.0|
|  6|    a|12.0|
|  7|    b|14.0|
|  8|    c|16.0|
+---+-----+----+

+---+-----+----+--------+
|_id|group|vals|vals_x_3|
+---+-----+----+--------+
|  0|    a| 0.0|     0.0|
|  7|    b|14.0|    42.0|
|  6|    a|12.0|    36.0|
|  5|    c|10.0|    30.0|
|  1|    b| 2.0|     6.0|
|  3|    a| 6.0|    18.0|
|  8|    c|16.0|    48.0|
|  2|    c| 4.0|    12.0|
|  4|    b| 8.0|    24.0|
+---+-----+----+--------+
user_18841710
Related