pyspark融入pandas的优势，真香！

以你之姓@ 2022-11-03 11:50 103阅读 0赞

近期，在使用spark的时候，发现spark在python下的使用，pyspark还挺好用的。

你甚至可以把它当作pandas来使用，众所周知，pandas在数据处理方面是很强大的，不谈性能，它提供了许多的内置方法，非常的方便，极大的减少我们的开发时间。

下面，简答来得展示一下它的具体使用。

**首先，我们初设一个SparkSession，并开启pandas的支持：Pandas with Apache Arrow；**

接着，简单的新建一个spark的DataFrame对象：df。

import pandas as pd
    
    from pyspark.sql import SparkSession, Window
    from pyspark.sql.functions import col, pandas_udf
    from pyspark.sql.types import *
    
    spark = SparkSession \
        .builder \
        .appName("Python Spark SQL Hive integration example") \
        .config("spark.sql.warehouse.dir", warehouse_location) \
        .enableHiveSupport() \
        .getOrCreate()
        
    # Enable Arrow-based columnar data transfers
    spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
    
    df = spark.createDataFrame(
            [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)],
            ("id", "v"))

# 列的自定义处理 #

@pandas_udf(“double”)
    def multiply_func(a: pd.Series, b: pd.Series) -> pd.Series:
        return a * b
     
    df.select(multiply_func(col("id"), col("v")))

上面，实现了一个两列的相乘。每一列在处理的都是pandas的Series，这个时候你就get到了吧，它能用上所有pandas的技巧。

# 列的聚合 #

@pandas_udf("double")
    def mean_udf(v: pd.Series) -> float:
    		return v.mean()
    
    # 对字段'v'进行求均值
    df.select(mean_udf(df['v'])).show()
    # +-----------+
    # |mean_udf(v)|
    # +-----------+
    # | 4.2|
    # +-----------+
    
    # 按照‘id’分组，求'v'的均值
    df.groupby("id").agg(mean_udf(df['v'])).show()
    # +---+-----------+
    # | id|mean_udf(v)|
    # +---+-----------+
    # | 1| 1.5|
    # | 2| 6.0|
    # +---+-----------+
    
    # 按照‘id’分组，求'v'的均值，并赋值给新的一列，但保持df的形式不变
    w = Window \
    		.partitionBy('id') \
    		.rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
    df.withColumn('mean_v', mean_udf(df['v']).over(w)).show()
    # +---+----+------+
    # | id| v|mean_v|
    # +---+----+------+
    # | 1| 1.0| 1.5|
    # | 1| 2.0| 1.5|
    # | 2| 3.0| 6.0|
    # | 2| 5.0| 6.0|
    # | 2|10.0| 6.0|
    # +---+----+------+

这里的聚合的单位是一列（pd.Series）：即要么只能是某一列如`mean_udf(df['v'])`，要么是拆开的某几列如`mean_udf(df['v'], df['id'])`，

所以返回的只能int、float、str这种非迭代的类型。

# 整体聚合 #

def subtract_mean(pdf):
        # pdf is a pandas.DataFrame
        # 这里根据id分组，那么每次传入的就是对应id的dataframe数据
        v = pdf.v
        pdf['v1'] = pdf['v'] - v.mean()
        pdf['v2'] = pdf['v'] + v.mean()
        return pdf
    
    df.groupby("id").applyInPandas(subtract_mean, schema="id long, v double, v1 double, v2 double").show()
    # +---+----+----+----+
    # | id | v | v1 | v2 |
    # +---+----+----+----+
    # | 1 | 1.0 | -0.5 | 2.5 |
    # | 1 | 2.0 | 0.5 | 3.5 |
    # | 2 | 3.0 | -3.0 | 9.0 |
    # | 2 | 5.0 | -1.0 | 11.0 |
    # | 2 | 10.0 | 4.0 | 16.0 |
    # +---+----+----+----+

这里的聚合跟上面的就不一样了，它处理的单位是pandas的dataframe，所以它的自由度就非常的高了。

在这里，你可以自由的选择返回多行数据，甚至可以新增字段。

# join #

df1 = spark.createDataFrame(
        [(20000101, 1, 1.0), (20000101, 2, 2.0), (20000102, 1, 3.0), (20000102, 2, 4.0)],
        ("time", "id", "v1"))
    
    df2 = spark.createDataFrame(
        [(20000101, 1, "x"), (20000101, 2, "y")],
        ("time", "id", "v2"))
    
    def asof_join(l, r):
      	# l、r is a pandas.DataFrame
      	# 这里是按照id分组
      	# 那么，l和r分别是对应id的df1和df2数据
        return pd.merge_asof(l, r, on="time", by="id")
    
    df1.groupby("id").cogroup(df2.groupby("id")).applyInPandas(
        asof_join, schema="time int, id int, v1 double, v2 string").show()
    # +--------+---+---+---+
    # | time| id| v1| v2|
    # +--------+---+---+---+
    # |20000101| 1|1.0| x|
    # |20000102| 1|3.0| x|
    # |20000101| 2|2.0| y|
    # |20000102| 2|4.0| y|
    # +--------+---+---+---+

# map #

df = spark.createDataFrame([(1, 21), (2, 30)], ("id", "age"))
    
    def filter_func(iterator):
      	# 传入的是一个pdf迭代器
        # 每一个pdf都是pandas的dataframe
        # 即按照df，一行数据对应一个pdf
        for pdf in iterator:
            id = int(pdf.id)
            age = int(pdf.age)
            pdf = pdf.append([{ "id": id*10, "age": age*10}])
            yield pdf
    
    df.mapInPandas(filter_func, schema=df.schema).show()
    # +---+---+
    # | id | age |
    # +---+---+
    # | 1 | 21 |
    # | 10 | 210 |
    # | 2 | 30 |
    # | 20 | 300 |
    # +---+---+