Arrow PyCapsule 接口#

原理#

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

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

目标#

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

非目标#

标准化应该使用哪些公共 API 进行导入。这由各个库自行决定。

PyCapsule 标准#

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

PyCapsule 允许为 capsule 关联一个 name，从而允许消费者验证该 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 结构体的释放回调 (release callback) 不为空时调用它。这可以防止在 capsule 从未传递给其他消费者的情况下发生内存泄漏。

如果 capsule 已传递给消费者，消费者应已移动数据并将释放回调标记为空，因此不存在释放消费者正在使用的数据的风险。请在 C 数据接口规范中阅读更多信息。

对于设备结构体，上述释放回调是嵌入式 ArrowArray 结构体的 release 成员。请在 C 设备接口规范中阅读更多信息。

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

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

导出协议#

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

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

设备接口另外定义了两个协议

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

ArrowSchema 导出#

Schema、字段和数据类型可以实现 __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 或 None) – 一个包含所请求 schema 的 C ArrowSchema 表示形式的 PyCapsule。转换为该 schema 是尽力而为的。请参阅 Schema 请求。
返回：: 一对分别包含 C ArrowSchema 和 ArrowArray 的 PyCapsules。Schema capsule 的名称应为 "arrow_schema"，数组 capsule 的名称应为 "arrow_array"。

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

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

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

参数:

requested_schema (PyCapsule 或 None) – 一个包含所请求 schema 的 C ArrowSchema 表示形式的 PyCapsule。转换为该 schema 是尽力而为的。请参阅 Schema 请求。
kwargs – 仅当其他关键字参数的默认值为 None 时，才应接受它们，以便将来添加新关键字。有关更多详细信息，请参阅设备支持。

返回：

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

ArrowStream 导出#

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

__arrow_c_stream__(self, requested_schema=None)#

将对象导出为 ArrowArrayStream。

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

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

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

将对象导出为 ArrowDeviceArrayStream。

参数:

requested_schema (PyCapsule 或 None) – 一个包含所请求 schema 的 C ArrowSchema 表示形式的 PyCapsule。转换为该 schema 是尽力而为的。请参阅 Schema 请求。
kwargs – 仅当其他关键字参数的默认值为 None 时，才应接受它们，以便将来添加新关键字。有关更多详细信息，请参阅设备支持。

返回：

一个包含对象 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 时相同的 schema。

如果调用者请求的 schema 与数据不兼容（例如请求具有不同字段数量的 schema），则被调用者应引发异常。所请求的 schema 机制仅旨在协调同一数据的不同表示形式，而不是允许任意的 schema 转换。

设备支持#

PyCapsule 接口通过使用 C 设备接口提供了跨硬件支持。这意味着可以在非 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__()）的默认值为 None，则它们应接受额外的关键字参数（**kwargs）。这允许将来添加新的可选关键字，其中此类新关键字的默认值始终为 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}"
        )

    ...

协议类型提示 (Typehints)#

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

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 指向的数据。它必须首先调用释放回调（如果它不为空），然后释放该结构体。

以下是为 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 交换协议的比较#

DataFrame 交换协议 (DataFrame Interchange Protocol) 是 Python 中的另一种允许库之间共享数据的协议。此协议是 DataFrame 交换协议的补充。许多实现此协议的对象也将实现 DataFrame 交换协议。

此协议特定于基于 Arrow 的数据结构，而 DataFrame 交换协议也允许将非 Arrow 数据帧和数组进行共享。因此，这些 PyCapsules 可以支持 Arrow 特有的功能，例如嵌套列。

此协议也比 DataFrame 交换协议极简得多。它只处理数据导出，而不是定义行数或列数等详细信息的访问器。

总之，如果您正在实现此协议，也应该考虑实现 DataFrame 交换协议。

与 `__arrow_array__` 协议的比较#

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