Arrow PyCapsule 接口#

警告

Arrow PyCapsule 接口应被视为实验性功能

基本原理#

C 数据接口、C 流接口和C device 接口允许在 Arrow 的不同实现之间移动 Arrow 数据。然而，这些接口没有指定 Python 库应如何向其他库公开这些结构体。在此之前，许多库仅提供导出到 PyArrow 数据结构的功能，使用 _import_from_c 和 _export_to_c 方法。但是，这始终要求安装 PyArrow。此外，如果处理不当，这些 API 可能会导致内存泄漏。

此接口允许任何库将其 Arrow 数据结构导出到理解相同协议的其他库。

目标#

标准化代表 ArrowSchema、ArrowArray、ArrowArrayStream、ArrowDeviceArray 和 ArrowDeviceArrayStream 的 PyCapsule 对象。
定义将 Arrow 数据导出到此类 capsule 对象的标准方法，以便任何想要接受 Arrow 数据作为输入的 Python 库都可以调用相应的方法，而不是硬编码对特定 Arrow 生成者的支持。

非目标#

标准化用于导入的公共 API。这取决于各个库。

PyCapsule 标准#

通过 Python 导出 Arrow 数据时，C 数据接口 / C 流接口结构体应被包装在 capsules 中。Capsules 通过为指针附加名称来避免无效访问，并通过附加析构函数来避免内存泄漏。因此，它们比将指针作为整数传递要安全得多。

PyCapsule 允许将 name 与 capsule 关联，使消费者能够验证 capsule 是否包含预期类型的数据。为确保 Arrow 结构体被识别，必须使用以下名称：

C 接口类型	PyCapsule 名称
ArrowSchema	`arrow_schema`
ArrowArray	`arrow_array`
ArrowArrayStream	`arrow_array_stream`
ArrowDeviceArray	`arrow_device_array`
ArrowDeviceArrayStream	`arrow_device_array_stream`

生命周期语义#

导出的 PyCapsules 应具有一个析构函数，如果 Arrow 结构体的释放回调不为 null，则调用它。这可以防止 capsule 未被传递给其他消费者时发生内存泄漏。

如果 capsule 已被传递给消费者，则消费者应已移动数据并将释放回调标记为 null，这样就不会有释放消费者正在使用的数据的风险。在 C 数据接口规范中阅读更多内容。

对于 device 结构体，上述释放回调是嵌入的 ArrowArray 结构体的 release 成员。在 C Device 接口规范中阅读更多内容。

就像在 C 数据接口中一样，此处定义的 PyCapsule 对象只能被消费一次。

有关带有析构函数的 PyCapsule 示例，请参见创建 PyCapsule。

导出协议#

该接口由三个独立的协议组成：

ArrowSchemaExportable，定义了 __arrow_c_schema__ 方法。
ArrowArrayExportable，定义了 __arrow_c_array__ 方法。
ArrowStreamExportable，定义了 __arrow_c_stream__ 方法。

为 Device 接口定义了另外两个协议：

ArrowDeviceArrayExportable，定义了 __arrow_c_device_array__ 方法。
ArrowDeviceStreamExportable，定义了 __arrow_c_device_stream__ 方法。

ArrowSchema 导出#

Schemas、字段和数据类型可以实现方法 __arrow_c_schema__。

__arrow_c_schema__(self)#

将对象导出为 ArrowSchema。

返回值：: 包含对象 C ArrowSchema 表示形式的 PyCapsule。该 capsule 的名称必须是 "arrow_schema"。

ArrowArray 导出#

数组和记录批次（连续表）可以实现方法 __arrow_c_array__。

__arrow_c_array__(self, requested_schema=None)#

将对象导出为一对 ArrowSchema 和 ArrowArray 结构体。

参数：: requested_schema (PyCapsule or None) – 一个 PyCapsule，包含请求 schema 的 C ArrowSchema 表示形式。转换为此 schema 是尽力而为的。参见schema 请求。
返回值：: 一对 PyCapsules，分别包含 C ArrowSchema 和 ArrowArray。schema capsule 的名称应为 "arrow_schema"，array capsule 的名称应为 "arrow_array"。

支持 Device 接口的库可以在这些对象上实现 __arrow_c_device_array__ 方法，该方法的工作方式与 __arrow_c_array__ 相同，只是返回的是 ArrowDeviceArray 结构体而不是 ArrowArray 结构体

__arrow_c_device_array__(self, requested_schema=None, **kwargs)#

将对象导出为一对 ArrowSchema 和 ArrowDeviceArray 结构体。

参数：

requested_schema (PyCapsule or None) – 一个 PyCapsule，包含请求 schema 的 C ArrowSchema 表示形式。转换为此 schema 是尽力而为的。参见schema 请求。
kwargs – 额外的关键字参数应仅在默认值为 None 时接受，以允许将来添加新关键字。更多详细信息请参见Device 支持。

返回值：

一对 PyCapsules，分别包含 C ArrowSchema 和 ArrowDeviceArray。schema capsule 的名称应为 "arrow_schema"，array capsule 的名称应为 "arrow_device_array"。

ArrowStream 导出#

表 / DataFrames 和流可以实现方法 __arrow_c_stream__。

__arrow_c_stream__(self, requested_schema=None)#

将对象导出为 ArrowArrayStream。

参数：: requested_schema (PyCapsule or None) – 一个 PyCapsule，包含请求 schema 的 C ArrowSchema 表示形式。转换为此 schema 是尽力而为的。参见schema 请求。
返回值：: 包含对象 C ArrowArrayStream 表示形式的 PyCapsule。该 capsule 的名称必须是 "arrow_array_stream"。

支持 Device 接口的库可以在这些对象上实现 __arrow_c_device_stream__ 方法，该方法的工作方式与 __arrow_c_stream__ 相同，只是返回的是 ArrowDeviceArrayStream 结构体而不是 ArrowArrayStream 结构体

__arrow_c_device_stream__(self, requested_schema=None, **kwargs)#

将对象导出为 ArrowDeviceArrayStream。

参数：

requested_schema (PyCapsule or None) – 一个 PyCapsule，包含请求 schema 的 C ArrowSchema 表示形式。转换为此 schema 是尽力而为的。参见schema 请求。
kwargs – 额外的关键字参数应仅在默认值为 None 时接受，以允许将来添加新关键字。更多详细信息请参见Device 支持。

返回值：

包含对象 C ArrowDeviceArrayStream 表示形式的 PyCapsule。该 capsule 的名称必须是 "arrow_device_array_stream"。

schema 请求#

在某些情况下，相同的数据可能存在多种 Arrow 表示形式。例如，一个库可能只有一个整数类型，但 Arrow 有多种不同大小和符号的整数类型。再例如，Arrow 对字符串数组有几种可能的编码：32位偏移量、64位偏移量、字符串视图和字典编码。字符串序列可以导出为这些 Arrow 表示形式中的任何一种。

为了允许调用者请求特定的表示形式，__arrow_c_array__() 和 __arrow_c_stream__() 方法接受一个可选的 requested_schema 参数。此参数是一个包含 ArrowSchema 的 PyCapsule。

被调用者应尝试按照请求的 schema 提供数据。但是，如果被调用者无法按请求的 schema 提供数据，则可以返回与将 None 传递给 requested_schema 时相同的 schema。

如果调用者请求的 schema 与数据不兼容，例如请求具有不同字段数量的 schema，被调用者应抛出异常。请求 schema 机制仅用于在同一数据的不同表示形式之间进行协商，而不允许任意的 schema 转换。

Device 支持#

PyCapsule 接口通过使用C device 接口支持跨硬件操作。这意味着可以在非 CPU 设备（例如 CUDA GPU）上交换数据，并检查交换的数据位于哪个设备上。

为了交换数据结构，此接口有两组协议方法：标准的仅 CPU 版本（__arrow_c_array__() 和 __arrow_c_stream__()）和等效的设备感知版本（__arrow_c_device_array__() 和 __arrow_c_device_stream__()）。

对于仅 CPU 的生产者，可以只实现标准的仅 CPU 协议方法，或者同时实现仅 CPU 和设备感知方法。缺少设备版本方法意味着数据仅位于 CPU 上。对于仅 CPU 的消费者，鼓励能够消费这两种版本的协议。

对于其数据结构只能驻留在非 CPU 内存中的设备感知生产者，建议仅实现协议的设备版本（例如，只添加 __arrow_c_device_array__，而不添加 __arrow_c_array__）。数据结构既可以在 CPU 上也可以在非 CPU 设备上驻留的生产者可以实现协议的两个版本，但应保证仅 CPU 版本（__arrow_c_array__() 和 __arrow_c_stream__()）包含指向 CPU 内存的有效指针（因此，尝试导出非 CPU 数据时，要么引发错误，要么复制到 CPU 内存）。

生成 ArrowDeviceArray 和 ArrowDeviceArrayStream 结构体预计不会涉及任何跨设备的复制数据。

设备感知方法（__arrow_c_device_array__() 和 __arrow_c_device_stream__()）应接受额外的关键字参数（**kwargs），如果它们的默认值为 None。这允许将来添加新的可选关键字，此类新关键字的默认值将始终为 None。实现者负责对用户传递的任何未识别的额外关键字抛出 NotImplementedError 异常。例如：

def __arrow_c_device_array__(self, requested_schema=None, **kwargs):

    non_default_kwargs = [
        name for name, value in kwargs.items() if value is not None
    ]
    if non_default_kwargs:
        raise NotImplementedError(
            f"Received unsupported keyword argument(s): {non_default_kwargs}"
        )

    ...

协议类型提示#

以下类型提示可以复制到您的库中，以标注函数接受实现这些协议之一的对象。

from typing import Tuple, Protocol

class ArrowSchemaExportable(Protocol):
    def __arrow_c_schema__(self) -> object: ...

class ArrowArrayExportable(Protocol):
    def __arrow_c_array__(
        self,
        requested_schema: object | None = None
    ) -> Tuple[object, object]:
        ...

class ArrowStreamExportable(Protocol):
    def __arrow_c_stream__(
        self,
        requested_schema: object | None = None
    ) -> object:
        ...

class ArrowDeviceArrayExportable(Protocol):
    def __arrow_c_device_array__(
        self,
        requested_schema: object | None = None,
        **kwargs,
    ) -> Tuple[object, object]:
        ...

class ArrowDeviceStreamExportable(Protocol):
    def __arrow_c_device_stream__(
        self,
        requested_schema: object | None = None,
        **kwargs,
    ) -> object:
        ...

示例#

创建 PyCapsule#

要创建 PyCapsule，请使用 PyCapsule_New 函数。必须向函数传递一个析构函数，该函数将在释放 capsule 指向的数据时被调用。它必须首先调用释放回调（如果不为 null），然后释放结构体。

下面是为 ArrowSchema 创建 PyCapsule 的代码。ArrowArray 和 ArrowArrayStream 的代码类似。

C

#include <Python.h>

void ReleaseArrowSchemaPyCapsule(PyObject* capsule) {
    struct ArrowSchema* schema =
        (struct ArrowSchema*)PyCapsule_GetPointer(capsule, "arrow_schema");
    if (schema->release != NULL) {
        schema->release(schema);
    }
    free(schema);
}

PyObject* ExportArrowSchemaPyCapsule() {
    struct ArrowSchema* schema =
        (struct ArrowSchema*)malloc(sizeof(struct ArrowSchema));
    // Fill in ArrowSchema fields
    // ...
    return PyCapsule_New(schema, "arrow_schema", ReleaseArrowSchemaPyCapsule);
}

Cython

cimport cpython
from libc.stdlib cimport malloc, free

cdef void release_arrow_schema_py_capsule(object schema_capsule):
    cdef ArrowSchema* schema = <ArrowSchema*>cpython.PyCapsule_GetPointer(
        schema_capsule, 'arrow_schema'
    )
    if schema.release != NULL:
        schema.release(schema)

    free(schema)

cdef object export_arrow_schema_py_capsule():
    cdef ArrowSchema* schema = <ArrowSchema*>malloc(sizeof(ArrowSchema))
    # It's recommended to immediately wrap the struct in a capsule, so
    # if subsequent lines raise an exception memory will not be leaked.
    schema.release = NULL
    capsule = cpython.PyCapsule_New(
        <void*>schema, 'arrow_schema', release_arrow_schema_py_capsule
    )
    # Fill in ArrowSchema fields:
    # schema.format = ...
    # ...
    return capsule

消费 PyCapsule#

要消费 PyCapsule，请使用 PyCapsule_GetPointer 函数获取底层结构体的指针。使用您的系统的 Arrow C 数据接口导入函数导入结构体。之后才能释放 capsule。

下面的示例展示了如何消费一个 ArrowSchema 的 PyCapsule。ArrowArray 和 ArrowArrayStream 的代码类似。

C

#include <Python.h>

// If the capsule is not an ArrowSchema, will return NULL and set an exception.
struct ArrowSchema* GetArrowSchemaPyCapsule(PyObject* capsule) {
  return PyCapsule_GetPointer(capsule, "arrow_schema");
}

Cython

cimport cpython

cdef ArrowSchema* get_arrow_schema_py_capsule(object capsule) except NULL:
    return <ArrowSchema*>cpython.PyCapsule_GetPointer(capsule, 'arrow_schema')

与 PyArrow 的向后兼容性#

与 PyArrow 交互时，应优先使用 PyCapsule 接口而非 _export_to_c 和 _import_from_c 方法。然而，许多库会希望支持一系列 PyArrow 版本。这可以通过 Duck typing 实现。

例如，如果您的库有一个导入方法，如下所示：

# OLD METHOD
def from_arrow(arr: pa.Array)
    array_import_ptr = make_array_import_ptr()
    schema_import_ptr = make_schema_import_ptr()
    arr._export_to_c(array_import_ptr, schema_import_ptr)
    return import_c_data(array_import_ptr, schema_import_ptr)

您可以重写此方法以同时支持 PyArrow 和其他实现了 PyCapsule 接口的库

# NEW METHOD
def from_arrow(arr)
    # Newer versions of PyArrow as well as other libraries with Arrow data
    # implement this method, so prefer it over _export_to_c.
    if hasattr(arr, "__arrow_c_array__"):
         schema_ptr, array_ptr = arr.__arrow_c_array__()
         return import_c_capsule_data(schema_ptr, array_ptr)
    elif isinstance(arr, pa.Array):
         # Deprecated method, used for older versions of PyArrow
         array_import_ptr = make_array_import_ptr()
         schema_import_ptr = make_schema_import_ptr()
         arr._export_to_c(array_import_ptr, schema_import_ptr)
         return import_c_data(array_import_ptr, schema_import_ptr)
    else:
        raise TypeError(f"Cannot import {type(arr)} as Arrow array data.")

您可能还希望在您的构造函数中接受实现该协议的对象。例如，在 PyArrow 中，array() 和 record_batch() 构造函数接受任何实现了 __arrow_c_array__() 方法协议的对象。类似地，PyArrow 的 schema() 构造函数接受任何实现了 __arrow_c_schema__() 方法的对象。

现在，如果您的库有一个导出到 PyArrow 的函数，如下所示：

# OLD METHOD
def to_arrow(self) -> pa.Array:
    array_export_ptr = make_array_export_ptr()
    schema_export_ptr = make_schema_export_ptr()
    self.export_c_data(array_export_ptr, schema_export_ptr)
    return pa.Array._import_from_c(array_export_ptr, schema_export_ptr)

您可以通过将您的对象传递给 array() 构造函数来重写此函数以使用 PyCapsule 接口，该构造函数接受任何实现了该协议的对象。检查 PyArrow 版本是否足够新以支持此功能的一个简单方法是检查 pa.Array 是否具有 __arrow_c_array__ 方法。

import warnings

# NEW METHOD
def to_arrow(self) -> pa.Array:
    # PyArrow added support for constructing arrays from objects implementing
    # __arrow_c_array__ in the same version it added the method for it's own
    # arrays. So we can use hasattr to check if the method is available as
    # a proxy for checking the PyArrow version.
    if hasattr(pa.Array, "__arrow_c_array__"):
        return pa.array(self)
    else:
        array_export_ptr = make_array_export_ptr()
        schema_export_ptr = make_schema_export_ptr()
        self.export_c_data(array_export_ptr, schema_export_ptr)
        return pa.Array._import_from_c(array_export_ptr, schema_export_ptr)

与其他协议的比较#

与 DataFrame Interchange Protocol 的比较#

The DataFrame Interchange Protocol 是 Python 中另一个允许在库之间共享数据的协议。此协议与 DataFrame Interchange Protocol 互补。许多实现了此协议的对象也将实现 DataFrame Interchange Protocol。

此协议专门针对基于 Arrow 的数据结构，而 DataFrame Interchange Protocol 也允许共享非 Arrow 的数据帧和数组。因此，这些 PyCapsules 可以支持 Arrow 特有的功能，例如嵌套列。

此协议也比 DataFrame Interchange Protocol 精简得多。它只处理数据导出，而不是定义访问器来获取诸如行数或列数之类的详细信息。

总而言之，如果您正在实现此协议，也应考虑实现 DataFrame Interchange Protocol。

与 `__arrow_array__` 协议的比较#

使用 __arrow_array__ 协议控制到 pyarrow.Array 的转换协议是一个双下划线方法，定义了 PyArrow 应如何将对象导入为 Arrow 数组。与此协议不同，它专用于 PyArrow，不被其他库使用。它也仅限于数组，不支持 schemas、表格结构或流。