在pyspark中使用pandas udf而不进行循环

e1xvtsh3  于 2022-11-28  发布在  Spark
关注(0)|答案(2)|浏览(174)

假设我有一个很大的spark Dataframe ,我不知道有多少列。
(the解决方案必须在pyspark中使用panda udf。没有不同的方法)
我想在所有列上执行一个操作,所以可以在所有列上循环,但是我不想在行上循环,我想让它立即在列上执行。
我在网上找不到这样做的方法。
假设我有这个数据表

A   B    C
5   3    2
1   7    0

现在我想发送到Pandasudf得到每行的总和。

Sum 
 10
  8

列数未知。
我可以在udf中一次循环一行,但我不想这样做,我想让它在所有行上都不循环,如果需要的话,我允许在列中循环。
我尝试的一个方法是将所有列组合成数组列

ARR
[5,3,2]
[1,7,0]

但即使在这里,如果不循环,它也不起作用。我把这个列发送到udf,然后在里面,我需要循环遍历它的行,并对list-row的每个值求和。
如果我能将每一列作为一个单独的列,并同时对整个列执行操作,那就太好了
如何立即对列执行操作?而不循环遍历行?
如果我循环遍历这些行,我想它并不比普通的python udf好多少

jm2pwxwz

jm2pwxwz1#

我不会去Pandasudfs,诉诸udfs它不能在pyspark中完成。无论如何,代码为以下两者

df = spark.read.load('/databricks-datasets/asa/small/small.csv', header=True,format='csv')

sf = df.select(df.colRegex("`.*rrDelay$|.*pDelay$`"))

#sf.show()

columns = ["id","ArrDelay","DepDelay"]
data = [("a", 81.0,3),
    ("b", 36.2,5),
    ("c", 12.0,5),
    ("d", 81.0,5),
    ("e", 36.3,5),
    ("f", 12.0,5),
    ("g", 111.7,5)]

sf = spark.createDataFrame(data=data,schema=columns)

sf.show()

# Use aggregate function
new = (sf.withColumn('sums', array(*[x for x in ['ArrDelay','DepDelay'] ]))#Create an array of values per row on desired columns
       .withColumn('sums', expr("aggregate(sums,cast(0 as double), (c,i)-> c+i)"))# USE aggregate to sum
      ).show()

#use pandas udf
sch= sf.withColumn('v', lit(90.087654623)).schema
def sum_s(iterator: Iterator[pd.DataFrame]) -> Iterator[pd.DataFrame]:
    for pdf in iterator:
           
      yield pdf.assign(v=pdf.sum(1))

sf.mapInPandas(sum_s, schema=sch).show()
qvtsj1bj

qvtsj1bj2#

下面是一个简单方法

from pyspark.sql import functions as F
from pyspark.sql.types import *
from pyspark.sql import Window
from functools import reduce

df = spark.createDataFrame(
    [
        (5,3,2),
        (1,7,0),        
    ],
    ["A", "B", "C"],
)

cols = df.columns
calculate_sum = reduce(lambda a, x: a+x, map(col, cols))

df = (
    df
    .withColumn(
        "sum",calculate_sum
    )
)

df.show()

输出:

+---+---+---+---+
|  A|  B|  C|sum|
+---+---+---+---+
|  5|  3|  2| 10|
|  1|  7|  0|  8|
+---+---+---+---+

相关问题