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 資料匯出至此類 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 匯出#

結構描述、欄位和資料類型可以實作 __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) – 包含請求的結構描述的 C ArrowSchema 表示法的 PyCapsule。盡力轉換為此結構描述。請參閱結構描述請求。

傳回值:

一對 PyCapsule，分別包含 C ArrowSchema 和 ArrowArray。結構描述 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) – 包含請求的結構描述的 C ArrowSchema 表示法的 PyCapsule。盡力轉換為此結構描述。請參閱結構描述請求。

  • kwargs – 額外的關鍵字引數應僅在它們具有 None 的預設值時才接受，以允許未來新增關鍵字。有關更多詳細資訊，請參閱裝置支援。

傳回值:

一對 PyCapsule，分別包含 C ArrowSchema 和 ArrowDeviceArray。結構描述 capsule 應具有名稱 "arrow_schema"，而陣列 capsule 應具有名稱 "arrow_device_array"。

ArrowStream 匯出#

表格 / DataFrame 和串流可以實作 __arrow_c_stream__ 方法。

__arrow_c_stream__(self, requested_schema=None)#

將物件匯出為 ArrowArrayStream。

參數:

requested_schema (PyCapsule 或 None) – 包含請求的結構描述的 C ArrowSchema 表示法的 PyCapsule。盡力轉換為此結構描述。請參閱結構描述請求。

傳回值:

包含物件的 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) – 包含請求的結構描述的 C ArrowSchema 表示法的 PyCapsule。盡力轉換為此結構描述。請參閱結構描述請求。

  • kwargs – 額外的關鍵字引數應僅在它們具有 None 的預設值時才接受，以允許未來新增關鍵字。有關更多詳細資訊，請參閱裝置支援。

傳回值:

包含物件的 C ArrowDeviceArrayStream 表示法的 PyCapsule。 capsule 必須具有名稱 "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 函數。該函數必須傳遞一個解構子函數，該函數將被呼叫以釋放 capsule 指向的資料。它必須先呼叫釋放回呼（如果不是空值），然後釋放結構。

以下是為 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 資料介面匯入函數匯入結構。只有在那之後才應釋放 capsule。

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

與其他協定的比較#

與 DataFrame 交換協定的比較#

DataFrame 交換協定是 Python 中的另一個協定，允許在程式庫之間共用資料。此協定是 DataFrame 交換協定的補充。許多實作此協定的物件也將實作 DataFrame 交換協定。

此協定專用於基於 Arrow 的資料結構，而 DataFrame 交換協定也允許共用非 Arrow 資料框架和陣列。因此，這些 PyCapsule 可以支援 Arrow 特定的功能，例如巢狀欄。

此協定也比 DataFrame 交換協定更精簡。它僅處理資料匯出，而不是定義用於詳細資訊（如列數或欄數）的存取器。

總之，如果您要實作此協定，也應考慮實作 DataFrame 交換協定。

與 __arrow_array__ 協定的比較#

使用 __arrow_array__ 協定控制轉換為 pyarrow.Array 協定是一種 dunder 方法，用於定義 PyArrow 應如何將物件匯入為 Arrow 陣列。與此協定不同，它專用於 PyArrow，並且不被其他程式庫使用。它也僅限於陣列，並且不支援結構描述、表格結構或串流。