Arrow PyCapsule 接口

警告

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 允许将 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

生命周期语义

导出的 PyCapsule 应具有一个析构函数，如果 Arrow 结构不为空，则调用其释放回调。这可以防止在 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。

返回:

一个 PyCapsule，其中包含对象的 C ArrowSchema 表示。该 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 请求。

返回:

一对 PyCapsule，分别包含 C ArrowSchema 和 ArrowArray。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 时才应接受其他关键字参数，以允许将来添加新关键字。有关更多详细信息，请参阅设备支持。

返回:

一对 PyCapsule，分别包含 C ArrowSchema 和 ArrowDeviceArray。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 请求。

返回:

一个 PyCapsule，其中包含对象的 C ArrowArrayStream 表示。该 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 时才应接受其他关键字参数，以允许将来添加新关键字。有关更多详细信息，请参阅设备支持。

返回:

一个 PyCapsule，其中包含对象的 C ArrowDeviceArrayStream 表示。该 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 转换。

设备支持

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__()）应接受附加关键字参数（**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 指向的数据。它必须首先调用释放回调（如果它不为空），然后释放结构。

下面是为 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 版本。这可以通过鸭子类型来实现。

例如，如果您的库有导入方法，例如

# 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)

您可以重写此函数以使用 PyCapsule 接口，方法是将您的对象传递给 array() 构造函数，该构造函数接受任何实现该协议的对象。检查 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 的比较

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

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

此协议也比 DataFrame Interchange Protocol 更加精简。它只处理数据导出，而不定义访问行数或列数等详细信息的访问器。

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

与 __arrow_array__ 协议的比较

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