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

您可以通过使用 RDD API 中公开的groupByKey 方法来实现强制 PySpark 对固定批次的行进行操作的预期结果。使用 groupByKey 将强制 PySpark 将单个密钥的所有数据随机分配给单个执行程序。

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

策略：

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

例子：

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