Pandas API on Spark使用详解

在上一篇文章中我们介绍了《PySpark DataFrame使用详解》，本篇文章我们继续介绍PySpark系列的第二个重要内容——Pandas API on Spark。

PySpark DataFrame虽然已经很大程度上方便了代码开发，并且支持pandas udf，但是Python开发者仍然需要学习相关的API，这对于习惯使用Pandas的用户而言仍然不够友好。基于此，PySpark祭出了核弹级大杀器——Pandas API on Spark，用户几乎可以完全按照Pandas的API写PySpark程序，这也意味着已有的Pandas程序几乎不需要修改就可以在Spark上执行，从而突破Pandas的串行计算瓶颈。

本文涉及的关于Pandas API on Spark的API都可在官方文档中查阅。

1. 环境配置

关于环境配置，版本要求，请参考前一篇文章《PySpark DataFrame使用详解》。针对本文的内容，安装PySpark的时候务必要安装安装Pandas依赖。

2. Pandas API on Spark使用

Pandas API on Spark的用法和Pandas API几乎一致，只需要修改一下导包即可：

import pandas as pd
import numpy as np
import pyspark.pandas as ps
from pyspark.sql import SparkSession

s = ps.Series([1, 3, 5, np.nan, 6, 8])
print(type(s)
# 

psdf = ps.DataFrame(
    {'a': [1, 2, 3, 4, 5, 6],
     'b': [100, 200, 300, 400, 500, 600],
     'c': ["one", "two", "three", "four", "five", "six"]},
    index=[10, 20, 30, 40, 50, 60])
print(type(psdf)
# 

• 创建pyspark.pandas

# 直接创建
psdf = ps.DataFrame(
    {'a': [1, 2, 3, 4, 5, 6],
     'b': [100, 200, 300, 400, 500, 600],
     'c': ["one", "two", "three", "four", "five", "six"]},
    index=[10, 20, 30, 40, 50, 60])

# 基于pandas.DataFrame创建
psdf = ps.from_pandas(pd.DataFrame())

# Pandas API on Spark转换为pandas.DataFrame。此操作需要把所有数据搜集到client，注意可能会出现oom。尽可能使用Pandas API on Spark
pdf = psdf.to_pandas()

# 基于Spark DataFrame创建
spark = SparkSession.builder.getOrCreate()
sdf = spark.createDataFrame(pdf)
psdf = sdf.pandas_api()

# Pandas API on Spark转为Spark DataFrame
sdf = psdf.to_spark().filter("id > 5")

# 当从Spark DataFrame创建pandas-on-Spark DataFrame时，会创建一个新的默认索引。为了避免这种开销，请尽可能指定要用作索引的列
# Create a pandas-on-Spark DataFrame with an explicit index.
psdf = ps.DataFrame({'id': range(10)}, index=range(10))
# Keep the explicit index.
sdf = psdf.to_spark(index_col='index')
# Call Spark APIs
sdf = sdf.filter("id > 5")
# Uses the explicit index to avoid to create default index.
sdf.pandas_api(index_col='index')

• 使用pandas api

psdf.to_csv('foo.csv')
ps.read_csv('foo.csv')
psdf.to_spark_io('zoo.orc', format="orc")
ps.read_spark_io('zoo.orc', format="orc")
psdf.index
psdf.columns
psdf.to_numpy()
psdf.describe()
psdf.sort_index(ascending=False)
psdf.sort_values(by='B')
psdf.dropna(how='any')
psdf.fillna(value=5)
psdf.groupby(['A', 'B']).sum()

# 绘图
pdf = pd.DataFrame(np.random.randn(1000, 4), index=pser.index, columns=['A', 'B', 'C', 'D'])
psdf = ps.from_pandas(pdf)
psdf = psdf.cummax()
psdf.plot()

注意：因为数据是分布式计算的，所以尽量不要使用没有分区的跨数据段操作，例如 psdf[‘B’].shift(1)，这会带来极大的计算性能开销，而且因为数据是分批计算的，不能保证数据有序。

• 配置项

# 设置最大显示行数，配置项都可在 pyspark.pandas.config.py 下查看，或者参考[官方文档](http://spark.incubator.apache.org/docs/latest/api/python/user_guide/pandas_on_spark/options.html#available-options)
ps.options.display.max_rows = 10		# 值形式
ps.options.display.max_rows
ps.set_option("display.max_rows", 101)	# API形式
ps.get_option("display.max_rows")
ps.reset_option("display.max_rows")

# 通过option_context设置，在with外无效
with ps.option_context("display.max_rows", 10, "compute.max_rows", 5):
	print(ps.get_option("display.max_rows"))
	print(ps.get_option("compute.max_rows"))

# 启用arrow优化
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", True)

# 使用默认索引防止多余开销
ps.set_option("compute.default_index_type", "distributed")

# 允许不同dataframe直接计算，应慎重考虑（效率低，而且必须保证索引相同）
ps.set_option('compute.ops_on_diff_frames', True)
psdf1 - psdf2
ps.reset_option('compute.ops_on_diff_frames')

# 设置索引类型
ps.set_option('compute.default_index_type', 'sequence')	# 全局唯一连续递增
ps.set_option('compute.default_index_type', 'distributed-sequence')	# 默认值，全局唯一连续递增，以分布式形式生成
ps.set_option('compute.default_index_type', 'distributed')	# 全局唯一递增，以分布式形式生成，但是不连续，性能更好

# 设置数据有序
# compute.ordered_head

• transform and apply

transform 和 apply在pandas中是使用非常高频的API。在Pandas API on Spark中，两者之间具有一些和pandas中不同的区别。

# apply中的axis参数可接受0或者1，transform只能接受0，也就意味着transform只能基于列计算
psdf = ps.DataFrame({'a': [1,2,3], 'b':[4,5,6]})
def pandas_plus(pser: pd.Series):
    return pser + 1

psdf.transform(pandas_plus)		# transform虽然有axis参数，但是只能为0
psdf.apply(pandas_plus, axis=0)
psdf.apply(pandas_plus, axis=1)

# apply可接受聚合计算，transform只能接受map关系计算，即输入和输出长度必须相同
def pandas_plus(pser: pd.Series):
    return sum(pser)  # allows an arbitrary length

psdf.apply(pandas_plus, axis='columns')	# 不能使用transform

# batch api，输入和输出都是pandas.DataFrame
# 除apply和transform外，Pandas API on Spark还提供了DataFrame.pandas_on_spark.transform_batch(), DataFrame.pandas_on_spark.apply_batch(), Series.pandas_on_spark.transform_batch()三个API用于分批处理
# batch api会对DataFrame或Series进行切片，然后对每个数据片分别应用func
psdf = ps.DataFrame({'a': [1,2,3], 'b':[4,5,6]})
def pandas_plus(pdf: pd.DataFrame):
    return pdf + 1  # should always return the same length as input.

psdf.pandas_on_spark.transform_batch(pandas_plus)
psdf.a.pandas_on_spark.transform_batch(pandas_plus)

psdf = ps.DataFrame({'a': [1,2,3], 'b':[4,5,6]})
def pandas_plus(pdf: pd.DataFrame):
    return pdf.query("a > 1")  # allow arbitrary length

psdf.pandas_on_spark.apply_batch(pandas_plus) # 不能使用transform_batch

• 数据类型

Pandas API on Spark目前暂不支持下面的Pandas特有类型：

pd.Timedelta
pd.Categorical
pd.CategoricalDtype
# 前三种计划支持

pd.SparseDtype
pd.DatetimeTZDtype
pd.UInt*Dtype
pd.BooleanDtype
pd.StringDtype
# 后五种没有支持计划

Python和PySpark类型映射关系如下：

Python	PySpark
bytes	BinaryType
int	LongType
float	DoubleType
str	StringType
bool	BooleanType
datetime.datetime	TimestampType
datetime.date	DateType
decimal.Decimal	DecimalType(38, 18)

可以使用as_spark_type查看映射关系：

import typing
import numpy as np
from pyspark.pandas.typedef import as_spark_type

as_spark_type(int)
# LongType
as_spark_type(np.int32)
# IntegerType
as_spark_type(typing.List[float])
# ArrayType(DoubleType,true)

也可以使用Spark accessor来检查底层的PySpark数据类型：

ps.Series([0.3, 0.1, 0.8]).spark.data_type
# DoubleType
ps.DataFrame({"d": [0.3, 0.1, 0.8], "s": ["welcome", "to", "pandas-on-Spark"], "b": [False, True, False]}).spark.print_schema()
# root
#  |-- d: double (nullable = false)
#  |-- s: string (nullable = false)
#  |-- b: boolean (nullable = false)

注意：Pandas API on Spark目前不支持Pandas中单列存储多种数据类型。

Numpy和PySpark的类型映射关系如下：

NumPy	PySpark
np.character	BinaryType
np.bytes_	BinaryType
np.string_	BinaryType
np.int8	ByteType
np.byte	ByteType
np.int16	ShortType
np.int32	IntegerType
np.int64	LongType
np.float32	FloatType
np.float64	DoubleType
np.unicode_	StringType
np.datetime64	TimestampType
np.ndarray	ArrayType(StringType())

• 类型声明

不管是PySpark还是PyFlink都有一个明显的特点就是在写udf函数的时候需要指定返回值类型，这也是和Java、Scala语言相比一个非常不便的点。在PyFlink中，如果没有写Function函数的schema，那么一般是按照pickle序列化的形式返回。

和PyFlink不同，Pandas API on Spark通过从输出中获取一些顶层输出来推断schema，然而这可能会导致一些高昂的性能开销，尤其是一些shuffle操作，如groupby，这是因为Pandas API on Spark会执行两次Spark作业，一次用于schema推断，然后使用得到的schema用于处理实际数据。为了消除这种影响，所以强烈建议在开发过程中显示的声明返回值schema来避免schema推断。

# pd.DataFrame schema声明
def pandas_div(pdf) -> pd.DataFrame[float, float]:
   # pdf is a pandas DataFrame.
   return pdf[['B', 'C']] / pdf[['B', 'C']]

df = ps.DataFrame({'A': ['a', 'a', 'b'], 'B': [1, 2, 3], 'C': [4, 6, 5]})
df.groupby('A').apply(pandas_div)

# pd.Series schema声明
def sqrt(x) -> pd.Series[float]:
    return np.sqrt(x)

df = ps.DataFrame([[4, 9]] * 3, columns=['A', 'B'])
df.apply(sqrt, axis=0)

# 有名schema声明（可以指定列名）
def transform(pdf) -> pd.DataFrame["id": int, "A": int]:
    pdf['A'] = pdf.id + 1
    return pdf

ps.range(5).pandas_on_spark.apply_batch(transform)

#动态声明schema
def transform(pdf) -> pd.DataFrame[
   return pdf + 1

psdf.pandas_on_spark.apply_batch(transform)

如果在schema声明中没有指定索引的类型，Pandas API on Spark会添加默认索引（原始索引会丢失，在操作中需要特别注意，这一点和Pandas不同），默认索引类型由compute.default_index_type设置决定。默认索引在shuffle过程中需要消耗额外的性能，因此最好同时指定索引类型。

pdf = pd.DataFrame({'id': range(5)})
sample = pdf.copy()
sample["a"] = sample.id + 1

def transform(pdf) -> pd.DataFrame[int, [int, int]]:
    pdf["a"] = pdf.id + 1
    return pdf

def transform(pdf) -> pd.DataFrame[sample.index.dtype, sample.dtypes]:
    pdf["a"] = pdf.id + 1
    return pdf

def transform(pdf) -> pd.DataFrame[("idxA", int), [("id", int), ("a", int)]]:
    pdf["a"] = pdf.id + 1
    return pdf

def transform(pdf) -> pd.DataFrame[(sample.index.name, sample.index.dtype), zip(sample.columns, sample.dtypes)]:
    pdf["a"] = pdf.id + 1
    return pdf

ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform)

如果是组合索引：

midx = pd.MultiIndex.from_arrays([(1, 1, 2), (1.5, 4.5, 7.5)], names=("int", "float"))
pdf = pd.DataFrame(range(3), index=midx, columns=["id"])
sample = pdf.copy()
sample["a"] = sample.id + 1

def transform(pdf) -> pd.DataFrame[[int, float], [int, int]]:
    pdf["a"] = pdf.id + 1
    return pdf

def transform(pdf) -> pd.DataFrame[sample.index.dtypes, sample.dtypes]:
    pdf["a"] = pdf.id + 1
    return pdf

def transform(pdf) -> pd.DataFrame[[("int", int), ("float", float)], [("id", int), ("a", int)]]:
    pdf["a"] = pdf.id + 1
    return pdf
    
def transform(pdf) -> pd.DataFrame[zip(sample.index.names, sample.index.dtypes), zip(sample.columns, sample.dtypes)]:
    pdf["A"] = pdf.id + 1
    return pdf
ps.from_pandas(pdf).pandas_on_spark.apply_batch(transform)

• 读写数据库

和Pandas一样，Pandas API on Spark也提供了读写数据库的API，但是又和Pandas有一些不同。Pandas API on Spark同样提供了下面三个API：

ps.read_sql_table(table_name, con[, schema, …])
ps.read_sql_query(sql, con[, index_col])
ps.read_sql(sql, con[, index_col, columns])

不同之处在于con参数需要的是标准jdbc连接，所以在使用之前必须把对应的数据库驱动包放到Spark类库中，或者通过SparkSession配置驱动包路径：

import os

from pyspark.sql import SparkSession

(SparkSession.builder
    .master("local")
    .appName("SQLite JDBC")
    .config(
        "spark.jars",
        "{}/sqlite-jdbc-3.34.0.jar".format(os.getcwd()))
    .config(
        "spark.driver.extraClassPath",
        "{}/sqlite-jdbc-3.34.0.jar".format(os.getcwd()))
    .getOrCreate())

读写数据方式如下：

df = ps.read_sql("stocks", con="jdbc:sqlite:{}/example.db".format(os.getcwd()))

# Pandas API on Spark没有直接提供to_sql API
df.spark.to_spark_io(format="jdbc", mode="append", dbtable="stocks", url="jdbc:sqlite:{}/example.db".format(os.getcwd()))

• SparkSession管理

和PySpark DataFrame不同，在上面的案例中，我们都没有定义SparkSession。这是因为在Pandas API on Spark中，SparkContext/SparkSession是开箱即用的(参考spark shell中的sc)，一旦创建了SparkContext/SparkSession，Pandas API on Spark就会自动使用显式创建的SparkContext/SparkSession。

• 查看执行计划

Pandas API on Spark是懒加载执行的，可以通过explain查看执行计划，从而确定复杂计算。

import pyspark.pandas as ps

psdf = ps.DataFrame({'id': range(10)})
psdf = psdf[psdf.id > 5]
psdf.spark.explain()

• 避免使用 __前缀/后缀列名

__ (双下划线)开头或者结尾的列名在Pandas API on Spark中是保留列，在其内部计算中会用到保留列列名，因此在使用中应避免使用这种列名，如果使用也可能会失效。

此外，Pandas API on Spark在定义DataFrame的时候列名如果只有大小写不同也是不允许的，例如：ps.DataFrame({‘a’: [1, 2], ‘A’:[3, 4]}) 这种定义就是错误的。当然也可以设置spark.sql.caseSensitive允许这种操作，但是非常不推荐：

import pyspark.pandas as ps
from pyspark.sql import SparkSession

builder = SparkSession.builder.appName("pandas-on-spark")
builder = builder.config("spark.sql.caseSensitive", "true")
builder.getOrCreate()

psdf = ps.DataFrame({'a': [1, 2], 'A':[3, 4]})

• 尽可能避免使用Python内置操作

Python中内置了很多便捷易用的函数和操作，如max、列表推导式等，然而这些操作在Pandas API on Spark中可能会失效或者带来一些潜在的风险。因为Python处理数据时数据都是一个节点中，所以都很便捷，但是Pandas API on Spark中的数据分布在不同的节点中，直接使用Python内置函数极有可能导致oom。例如下面的对比操作：

##########################################
import pandas as pd
max(pd.Series([1, 2, 3]))

import pyspark.pandas as ps
ps.Series([1, 2, 3]).max()

##########################################
import pandas as pd
data = []
countries = ['London', 'New York', 'Helsinki']
pser = pd.Series([20., 21., 12.], index=countries)
for temperature in pser:
    assert temperature > 0
    if temperature > 1000:
        temperature = None
    data.append(temperature ** 2)

pd.Series(data, index=countries)

import pyspark.pandas as ps
import numpy as np
countries = ['London', 'New York', 'Helsinki']
psser = ps.Series([20., 21., 12.], index=countries)
def square(temperature) -> np.float64:
    assert temperature > 0
    if temperature > 1000:
        temperature = None
    return temperature ** 2

psser.apply(square)