Arrow PyCapsule 接口

警告

Arrow PyCapsule 接口应被视为实验性的

基本原理

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

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

目标

• 标准化表示 ArrowSchema、ArrowArray、ArrowArrayStream、ArrowDeviceArray 和 ArrowDeviceArrayStream 的 PyCapsule 对象。

• 定义将 Arrow 数据导出到此类胶囊对象的标准方法，以便任何希望接受 Arrow 数据作为输入的 Python 库都可以调用相应的方法，而不是硬编码对特定 Arrow 生产者的支持。

非目标

• 标准化应该使用哪些公共 API 进行导入。这取决于各个库。

PyCapsule 标准

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

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

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

生命周期语义

如果导出的 PyCapsule 的析构函数尚未为空，则它应该调用 Arrow 结构的 释放回调。这可以防止在胶囊从未传递给其他使用者的情况下发生内存泄漏。

如果胶囊已传递给使用者，则使用者应已移动数据并将释放回调标记为空，因此不存在释放使用者正在使用的数据的风险。在 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 导出

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

__arrow_c_schema__(self)

将对象导出为 ArrowSchema。

返回:

包含对象的 C ArrowSchema 表示形式的 PyCapsule。胶囊的名称必须为 "arrow_schema"。

ArrowArray 导出

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

__arrow_c_array__(self, requested_schema=None)

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

参数:

requested_schema (PyCapsule 或 None) – 包含请求模式的 C ArrowSchema 表示形式的 PyCapsule。转换为此模式是尽力而为的。请参阅模式请求。

返回:

一对 PyCapsule，分别包含 C ArrowSchema 和 ArrowArray。模式胶囊的名称应为 "arrow_schema"，数组胶囊的名称应为 "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) – 包含请求模式的 C ArrowSchema 表示形式的 PyCapsule。转换为此模式是尽力而为的。请参阅模式请求。

• kwargs – 仅当附加关键字参数具有默认值 None 时才应接受它们，以便将来添加新的关键字。有关更多详细信息，请参阅设备支持。

返回:

一对 PyCapsule，分别包含 C ArrowSchema 和 ArrowDeviceArray。模式胶囊的名称应为 "arrow_schema"，数组胶囊的名称应为 "arrow_device_array"。

ArrowStream 导出

表/数据帧和流可以实现 __arrow_c_stream__ 方法。

__arrow_c_stream__(self, requested_schema=None)

将对象导出为 ArrowArrayStream。

参数:

requested_schema (PyCapsule 或 None) – 包含请求模式的 C ArrowSchema 表示形式的 PyCapsule。转换为此模式是尽力而为的。请参阅模式请求。

返回:

一个包含对象的 C ArrowArrayStream 表示形式的 PyCapsule。该胶囊必须具有名称 "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) – 包含请求模式的 C ArrowSchema 表示形式的 PyCapsule。转换为此模式是尽力而为的。请参阅模式请求。

• kwargs – 仅当附加关键字参数具有默认值 None 时才应接受它们，以便将来添加新的关键字。有关更多详细信息，请参阅设备支持。

返回:

一个包含对象的 C ArrowDeviceArrayStream 表示形式的 PyCapsule。该胶囊必须具有名称 "arrow_device_array_stream"。

模式请求

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

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

被调用方应尝试以请求的模式提供数据。但是，如果被调用方无法以请求的模式提供数据，则它们可以返回与将 None 传递给 requested_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 函数。必须将析构函数传递给该函数，该函数将被调用以释放胶囊指向的数据。它必须首先调用释放回调（如果它不为空），然后释放结构体。

以下是为 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 数据接口导入函数导入结构。只有在此之后才能释放胶囊。

以下示例显示了如何使用 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)

您可以通过将对象传递给 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 交换协议 是 Python 中的另一个协议，允许在库之间共享数据。此协议是对 DataFrame 交换协议的补充。许多实现此协议的对象也将实现 DataFrame 交换协议。

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

此协议也比 DataFrame 交换协议更小。它只处理数据导出，而不是定义诸如行数或列数之类的详细信息的访问器。

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

与 __arrow_array__ 协议的比较

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