pyspark从数据帧到每次执行迭代N行

Question

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]

输入数据帧为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|
+----+----+-----+ 

下面是一段代码片段，它可以工作，但需要很长时间

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，内存有问题，一次最多只能处理50000行。我希望使用funudf = udf(fun_1)。 因此，如何将PySpark DF拆分为50000行的段，调用funudf ->fun_1。 输出有两个列，分别来自输入的“col1”和“funudf返回值”

Answer 1

通过使用RDDAPI中公开的groupByKey方法，可以实现强制PySpark对固定批次的行进行操作的预期结果。使用groupByKey将强制PySpark将单个密钥的所有数据洗牌到单个执行器

注意：出于同样的原因，由于网络成本，通常不鼓励使用groupByKey

战略：

1. 添加一列，将数据分组到所需的批大小和groupByKey
2. 定义一个复制UDF逻辑的函数（并返回一个id供以后加入）。这在pyspark.resultiterable.ResultIterable上运行，这是groupByKey的结果。使用mapValues将函数应用于您的组
3. 将生成的RDD转换为数据帧并重新加入

例如：

# 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|
+ -+  -+  +    +