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、表格結構或流。