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 数据导出到此类胶囊对象中，以便任何希望接受 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

生命周期语义#

导出的 PyCapsules 应该有一个析构函数，如果 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。转换为该模式是尽力而为的。请参见 模式请求。

返回值:

包含 C ArrowSchema 和 ArrowArray 的 PyCapsules 对。模式胶囊的名称应为 "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 的默认值时，才会接受其他关键字参数，以允许将来添加新关键字。有关更多详细信息，请参见 设备支持。

返回值:

包含 C ArrowSchema 和 ArrowDeviceArray 的 PyCapsules 对。模式胶囊的名称应为 "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 的代码类似。

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

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 的代码类似。

#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");
}

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)

与其他协议的比较#

与数据帧交换协议的比较#

数据帧交换协议 是 Python 中的另一个协议，它允许在库之间共享数据。此协议与数据帧交换协议互补。实现此协议的许多对象也将实现数据帧交换协议。

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

此协议也比数据帧交换协议要小得多。它只处理数据导出，而不是定义像行数或列数这样的详细信息的访问器。

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

与 __arrow_array__ 协议的比较#

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

本页内容
在 GitHub 上编辑