Pandas 集成#

为了与 pandas 交互,PyArrow 提供了各种转换例程来使用 pandas 结构并将其转换回它们。

注意

虽然 pandas 使用 NumPy 作为后端,但它有足够的特殊性(例如不同的类型系统和对空值的支持),因此这是与 NumPy 集成 不同的主题。

要按照本文件中示例操作,请确保运行

In [1]: import pandas as pd

In [2]: import pyarrow as pa

数据帧#

Arrow 中与 pandas DataFrame 等效的是 Table。两者都包含一组长度相等的命名列。虽然 pandas 仅支持扁平列,但 Table 还提供嵌套列,因此它可以表示比 DataFrame 更多的数据,因此并非总是可以进行完整的转换。

从 Table 转换为 DataFrame 是通过调用 pyarrow.Table.to_pandas() 完成的。反之,则通过使用 pyarrow.Table.from_pandas() 来实现。

import pyarrow as pa
import pandas as pd

df = pd.DataFrame({"a": [1, 2, 3]})
# Convert from pandas to Arrow
table = pa.Table.from_pandas(df)
# Convert back to pandas
df_new = table.to_pandas()

# Infer Arrow schema from pandas
schema = pa.Schema.from_pandas(df)

默认情况下,pyarrow 尝试尽可能准确地保留和恢复 .index 数据。有关此内容以及如何禁用此逻辑的更多信息,请参见下面的部分。

系列#

在 Arrow 中,与 pandas Series 最相似的结构是 Array。它是一个向量,包含与线性内存中相同类型的數據。可以使用 pyarrow.Array.from_pandas() 将 pandas Series 转换为 Arrow Array。由于 Arrow Array 总是可为空的,可以使用 mask 参数提供一个可选的掩码,以标记所有空条目。

处理 pandas 索引#

pyarrow.Table.from_pandas() 这样的方法有一个 preserve_index 选项,它定义了如何保留(存储)或不保留(不存储)对应 pandas 对象的 index 成员中的数据。这些数据使用内部 arrow::Schema 对象中的模式级元数据进行跟踪。

preserve_index 的默认值为 None,其行为如下

  • RangeIndex 仅作为元数据存储,不需要任何额外存储。

  • 其他索引类型将作为结果 Table 中的一个或多个物理数据列存储。

要完全不存储索引,请传递 preserve_index=False。由于在某些有限情况下存储 RangeIndex 会导致问题(例如,将多个 DataFrame 对象存储在 Parquet 文件中),因此要强制将所有索引数据序列化到结果表中,请传递 preserve_index=True

类型差异#

在 pandas 和 Arrow 的当前设计中,无法将所有列类型进行未修改的转换。这里的主要问题之一是 pandas 不支持任意类型的可为空列。此外,datetime64 目前固定为纳秒分辨率。另一方面,Arrow 可能仍然缺少对某些类型的支持。

pandas -> Arrow 转换#

源类型 (pandas)

目标类型 (Arrow)

bool

BOOL

(u)int{8,16,32,64}

(U)INT{8,16,32,64}

float32

FLOAT

float64

DOUBLE

str / unicode

STRING

pd.Categorical

DICTIONARY

pd.Timestamp

TIMESTAMP(unit=ns)

datetime.date

DATE

datetime.time

TIME64

Arrow -> pandas 转换#

源类型 (Arrow)

目标类型 (pandas)

BOOL

bool

BOOL with nulls

object (with values True, False, None)

(U)INT{8,16,32,64}

(u)int{8,16,32,64}

(U)INT{8,16,32,64} with nulls

float64

FLOAT

float32

DOUBLE

float64

STRING

str

DICTIONARY

pd.Categorical

TIMESTAMP(unit=*)

pd.Timestamp (np.datetime64[ns])

DATE

object (with datetime.date objects)

TIME64

object (with datetime.time objects)

分类类型#

Pandas 分类 列被转换为 Arrow 字典数组,这是一种专门的数组类型,经过优化以处理重复和有限数量的可能值。

In [3]: df = pd.DataFrame({"cat": pd.Categorical(["a", "b", "c", "a", "b", "c"])})

In [4]: df.cat.dtype.categories
Out[4]: Index(['a', 'b', 'c'], dtype='object')

In [5]: df
Out[5]: 
  cat
0   a
1   b
2   c
3   a
4   b
5   c

In [6]: table = pa.Table.from_pandas(df)

In [7]: table
Out[7]: 
pyarrow.Table
cat: dictionary<values=string, indices=int8, ordered=0>
----
cat: [  -- dictionary:
["a","b","c"]  -- indices:
[0,1,2,0,1,2]]

我们可以检查创建的表的 ChunkedArray,并看到 Pandas DataFrame 的相同类别。

In [8]: column = table[0]

In [9]: chunk = column.chunk(0)

In [10]: chunk.dictionary
Out[10]: 
<pyarrow.lib.StringArray object at 0x7f15601f3400>
[
  "a",
  "b",
  "c"
]

In [11]: chunk.indices
Out[11]: 
<pyarrow.lib.Int8Array object at 0x7f15601f3280>
[
  0,
  1,
  2,
  0,
  1,
  2
]

日期时间 (时间戳) 类型#

Pandas 时间戳 在 Pandas 中使用 datetime64[ns] 类型,并被转换为 Arrow TimestampArray

In [12]: df = pd.DataFrame({"datetime": pd.date_range("2020-01-01T00:00:00Z", freq="h", periods=3)})

In [13]: df.dtypes
Out[13]: 
datetime    datetime64[ns, UTC]
dtype: object

In [14]: df
Out[14]: 
                   datetime
0 2020-01-01 00:00:00+00:00
1 2020-01-01 01:00:00+00:00
2 2020-01-01 02:00:00+00:00

In [15]: table = pa.Table.from_pandas(df)

In [16]: table
Out[16]: 
pyarrow.Table
datetime: timestamp[ns, tz=UTC]
----
datetime: [[2020-01-01 00:00:00.000000000Z,2020-01-01 01:00:00.000000000Z,2020-01-01 02:00:00.000000000Z]]

在这个例子中,Pandas 时间戳是时区感知的(在本例中为 UTC),此信息用于创建 Arrow TimestampArray

日期类型#

虽然可以使用 pandas 中的 datetime64[ns] 类型来处理日期,但一些系统使用 Python 内置 datetime.date 对象的对象数组。

In [17]: from datetime import date

In [18]: s = pd.Series([date(2018, 12, 31), None, date(2000, 1, 1)])

In [19]: s
Out[19]: 
0    2018-12-31
1          None
2    2000-01-01
dtype: object

转换为 Arrow 数组时,默认情况下将使用 date32 类型。

In [20]: arr = pa.array(s)

In [21]: arr.type
Out[21]: DataType(date32[day])

In [22]: arr[0]
Out[22]: <pyarrow.Date32Scalar: datetime.date(2018, 12, 31)>

要使用 64 位 date64,请明确指定。

In [23]: arr = pa.array(s, type='date64')

In [24]: arr.type
Out[24]: DataType(date64[ms])

使用 to_pandas 转换回来时,将返回 datetime.date 对象的对象数组。

In [25]: arr.to_pandas()
Out[25]: 
0    2018-12-31
1          None
2    2000-01-01
dtype: object

如果你想使用 NumPy 的 datetime64 dtype,请传递 date_as_object=False

In [26]: s2 = pd.Series(arr.to_pandas(date_as_object=False))

In [27]: s2.dtype
Out[27]: dtype('<M8[ms]')

警告

从 Arrow 0.13 开始,参数 date_as_object 默认情况下为 True。较旧的版本必须传递 date_as_object=True 才能获得此行为。

时间类型#

Pandas 数据结构中的内置 datetime.time 对象将分别转换为 Arrow time64Time64Array

In [28]: from datetime import time

In [29]: s = pd.Series([time(1, 1, 1), time(2, 2, 2)])

In [30]: s
Out[30]: 
0    01:01:01
1    02:02:02
dtype: object

In [31]: arr = pa.array(s)

In [32]: arr.type
Out[32]: Time64Type(time64[us])

In [33]: arr
Out[33]: 
<pyarrow.lib.Time64Array object at 0x7f15ee735e40>
[
  01:01:01.000000,
  02:02:02.000000
]

转换为 pandas 时,将返回 datetime.time 对象的数组。

In [34]: arr.to_pandas()
Out[34]: 
0    01:01:01
1    02:02:02
dtype: object

可为空类型#

在 Arrow 中,所有数据类型都是可为空的,这意味着它们支持存储缺失值。但是,在 pandas 中,并非所有数据类型都支持缺失数据。最值得注意的是,默认的整数数据类型不支持,并且在引入缺失值时将转换为浮点数。因此,当 Arrow 数组或表格转换为 pandas 时,整数列将在存在缺失值时变为浮点数。

>>> arr = pa.array([1, 2, None])
>>> arr
<pyarrow.lib.Int64Array object at 0x7f07d467c640>
[
  1,
  2,
  null
]
>>> arr.to_pandas()
0    1.0
1    2.0
2    NaN
dtype: float64

Pandas 具有实验性的可为空数据类型 (https://pandas.ac.cn/docs/user_guide/integer_na.html)。Arrows 支持这些类型的双向转换。

>>> df = pd.DataFrame({'a': pd.Series([1, 2, None], dtype="Int64")})
>>> df
      a
0     1
1     2
2  <NA>

>>> table = pa.table(df)
>>> table
Out[32]:
pyarrow.Table
a: int64
----
a: [[1,2,null]]

>>> table.to_pandas()
      a
0     1
1     2
2  <NA>

>>> table.to_pandas().dtypes
a    Int64
dtype: object

这种双向转换之所以有效,是因为有关原始 pandas DataFrame 的元数据存储在 Arrow 表格中。但是,如果你有非源自具有可为空数据类型的 pandas DataFrame 的 Arrow 数据(或例如 Parquet 文件),则 pandas 的默认转换将不会使用这些可为空数据类型。

pyarrow.Table.to_pandas() 方法有一个 types_mapper 关键字,可以用来覆盖用于结果 pandas DataFrame 的默认数据类型。这样,你可以指示 Arrow 使用可为空数据类型来创建 pandas DataFrame。

>>> table = pa.table({"a": [1, 2, None]})
>>> table.to_pandas()
     a
0  1.0
1  2.0
2  NaN
>>> table.to_pandas(types_mapper={pa.int64(): pd.Int64Dtype()}.get)
      a
0     1
1     2
2  <NA>

types_mapper 关键字期望一个函数,该函数将在给定 pyarrow 数据类型的情况下返回要使用的 pandas 数据类型。通过使用 dict.get 方法,我们可以使用字典创建这样的函数。

如果你想使用 pandas 支持的所有当前可为空数据类型,这个字典将变为

dtype_mapping = {
    pa.int8(): pd.Int8Dtype(),
    pa.int16(): pd.Int16Dtype(),
    pa.int32(): pd.Int32Dtype(),
    pa.int64(): pd.Int64Dtype(),
    pa.uint8(): pd.UInt8Dtype(),
    pa.uint16(): pd.UInt16Dtype(),
    pa.uint32(): pd.UInt32Dtype(),
    pa.uint64(): pd.UInt64Dtype(),
    pa.bool_(): pd.BooleanDtype(),
    pa.float32(): pd.Float32Dtype(),
    pa.float64(): pd.Float64Dtype(),
    pa.string(): pd.StringDtype(),
}

df = table.to_pandas(types_mapper=dtype_mapping.get)

当使用 pandas API 读取 Parquet 文件 (pd.read_parquet(..)) 时,也可以通过传递 use_nullable_dtypes 来实现。

df = pd.read_parquet(path, use_nullable_dtypes=True)

内存使用情况和零拷贝#

当使用各种 to_pandas 方法从 Arrow 数据结构转换为 pandas 对象时,有时必须注意与性能和内存使用情况相关的问题。

由于 pandas 的内部数据表示通常不同于 Arrow 的列式格式,因此只有在某些有限情况下才能实现零拷贝转换(不需要内存分配或计算)。

在最坏的情况下,调用 to_pandas 将导致内存中存在数据的两个版本,一个用于 Arrow,一个用于 pandas,从而产生大约两倍的内存占用。我们已经为这种情况实施了一些缓解措施,特别是当创建大型 DataFrame 对象时,我们将在下面描述这些措施。

零拷贝系列转换#

ArrayChunkedArray 到 NumPy 数组或 pandas Series 的零拷贝转换在某些狭窄的情况下是可能的。

  • Arrow 数据存储在整数(有符号或无符号 int8int64)或浮点类型 (float16float64) 中。这包括许多数值类型以及时间戳。

  • Arrow 数据没有空值(因为这些值使用位图表示,而位图不受 pandas 支持)。

  • 对于 ChunkedArray,数据由单个块组成,即 arr.num_chunks == 1。多个块总是需要复制,因为 pandas 要求数据连续。

在这些情况下,to_pandasto_numpy 将是零拷贝。在所有其他情况下,都需要进行复制。

减少 Table.to_pandas 中的内存使用#

截至撰写本文时,pandas 采用了一种称为“合并”的数据管理策略,将类型相同的 DataFrame 列收集到二维 NumPy 数组中,在内部称为“块”。我们付出了巨大的努力来构建精确的“合并”块,以便在我们将数据传递给 pandas.DataFrame 后,pandas 不会执行任何进一步的分配或复制。这种合并策略的明显缺点是它会导致“内存翻倍”。

为了尝试限制在 Table.to_pandas 期间“内存翻倍”的潜在影响,我们提供了一些选项

  • split_blocks=True,启用后 Table.to_pandas 为每一列生成一个内部 DataFrame “块”,跳过“合并”步骤。请注意,许多 pandas 操作会触发合并,但峰值内存使用可能低于完全内存翻倍的最坏情况。由于有了这个选项,我们能够在与 ArrayChunkedArray 执行零拷贝转换相同的情况下,对列进行零拷贝转换。

  • self_destruct=True,这会在将每列 Table 对象转换为与 pandas 兼容的表示形式时,销毁其内部 Arrow 内存缓冲区,从而可能在列转换后立即将内存释放给操作系统。请注意,这会使调用的 Table 对象不安全,无法进一步使用,任何进一步调用的方法都会导致您的 Python 进程崩溃。

共同使用时,调用

df = table.to_pandas(split_blocks=True, self_destruct=True)
del table  # not necessary, but a good practice

在某些情况下将产生明显更低的内存使用量。如果没有这些选项,to_pandas 将始终使内存翻倍。

请注意,self_destruct=True 并不保证节省内存。由于转换是逐列进行的,因此内存也是逐列释放的。但是,如果多个列共享一个基础缓冲区,那么在所有这些列都转换完成之前,将不会释放任何内存。特别是,由于实现细节,来自 IPC 或 Flight 的数据容易出现这种情况,因为内存将按以下方式布局

Record Batch 0: Allocation 0: array 0 chunk 0, array 1 chunk 0, ...
Record Batch 1: Allocation 1: array 0 chunk 1, array 1 chunk 1, ...
...

在这种情况下,即使使用 self_destruct=True,在整个表格转换完成之前,也无法释放任何内存。