擴充套件 PyArrow

使用 PyCapsule 介面控制向 (Py)Arrow 的轉換

Arrow C 資料介面 允許在不同的 Arrow 實現之間傳輸 Arrow 資料。這是一個通用的、跨語言的介面，並非 Python 專用。但對於 Python 庫，該介面透過一個 Python 特有的層進行了擴充套件：Arrow PyCapsule 介面。

此 Python 介面確保支援 C 資料介面的不同庫能夠以標準方式匯出 Arrow 資料結構，並識別彼此的物件。

如果您擁有一個提供底層持有 Arrow 相容資料的 Python 庫，則可以在這些物件上實現以下方法：

• __arrow_c_schema__，用於 Schema（模式）或類型別物件。

• __arrow_c_array__，用於陣列和記錄批次（連續表）。

• __arrow_c_stream__，用於分塊陣列、表和資料流。

這些方法返回 PyCapsule 物件，關於確切語義的更多詳細資訊，請參見 規範。

當您的資料結構定義了這些方法時，PyArrow 的建構函式（見下文）將識別這些物件支援此協議，並將它們零複製地轉換為 PyArrow 資料結構。對於任何其他支援此攝取資料協議的庫，情況也是如此。

同樣，如果您的庫具有接受使用者提供資料的函式，您可以透過檢查這些方法是否存在來新增對該協議的支援，從而接受任何 Arrow 資料（而不是硬編碼對特定 Arrow 生產者（如 PyArrow）的支援）。

如需透過此協議使用 PyArrow 消費資料，可以使用以下建構函式來建立各種 PyArrow 物件：

結果類	PyArrow 建構函式	支援的協議
陣列	pyarrow.array()	陣列
ChunkedArray	pyarrow.chunked_array()	array（陣列）, stream（流）
RecordBatch	pyarrow.record_batch()	陣列
表	pyarrow.table()	array（陣列）, stream（流）
RecordBatchReader	pyarrow.RecordBatchReader.from_stream()	stream（流）
Field	pyarrow.field()	模式
Schema	pyarrow.schema()	模式

DataType 可以透過使用 pyarrow.field() 消費 Schema 相容物件，然後訪問結果 Field 的 .type 來建立。

使用 __arrow_array__ 協議控制向 pyarrow.Array 的轉換

pyarrow.array() 函式內建了對 Python 序列、numpy 陣列和 pandas 一維物件（Series、Index、Categorical 等）的支援，可將它們轉換為 Arrow 陣列。透過實現 __arrow_array__ 方法（類似於 numpy 的 __array__ 協議），可以將其擴充套件到其他類陣列物件。

例如，要支援將您的 duck 陣列類轉換為 Arrow 陣列，請定義 __arrow_array__ 方法以返回一個 Arrow 陣列。

>>> class MyDuckArray:
...
...     def __arrow_array__(self, type=None):
...         # convert the underlying array values to a PyArrow Array
...         import pyarrow
...         return pyarrow.array(..., type=type)

__arrow_array__ 方法採用一個可選的 type 關鍵字，該關鍵字由 pyarrow.array() 傳入。該方法允許返回 Array 或 ChunkedArray。

注意

若要以更通用的方式控制 Python 物件到 Arrow 資料的轉換，請考慮 Arrow PyCapsule 介面。它不特定於 PyArrow，並且支援轉換其他物件，如表（tables）和 Schema（模式）。

定義擴充套件型別（“使用者定義型別”）

Arrow 提供了一種擴充套件型別概念，允許使用者使用額外的語義來註釋資料型別。這使得開發者既可以指定自定義的序列化和反序列化程式（例如，用於 Python 標量 和 pandas），也可以更輕鬆地解釋資料。

在 Arrow 中，擴充套件型別 是透過使用自定義型別名稱（以及可選的位元組串，該位元組串可用於提供額外元資料，在本規範中稱為“引數”）來註釋任何內建 Arrow 資料型別（即“儲存型別”）來指定的。這些出現在 Field 的 custom_metadata 中的 ARROW:extension:name 和 ARROW:extension:metadata 鍵中。

請注意，由於這些註釋是 Arrow 規範的一部分，它們可能被其他（非 Python）Arrow 消費者（如 PySpark）識別。

PyArrow 允許您透過繼承 ExtensionType 並賦予派生類其自己的副檔名及 (反)序列化任何引數的機制，來從 Python 定義擴充套件型別。例如，我們可以定義一個自定義的有理數型別，表示為一對整數。

>>> import pyarrow as pa
>>> class RationalType(pa.ExtensionType):
...
...     def __init__(self, data_type: pa.DataType):
...         if not pa.types.is_integer(data_type):
...             raise TypeError(f"data_type must be an integer type not {data_type}")
...
...         super().__init__(
...             pa.struct(
...                 [
...                     ("numer", data_type),
...                     ("denom", data_type),
...                 ],
...             ),
...             "my_package.rational",
...         )
...
...     def __arrow_ext_serialize__(self) -> bytes:
...         # No parameters are necessary
...         return b""
...
...     @classmethod
...     def __arrow_ext_deserialize__(cls, storage_type, serialized):
...         # Sanity checks, not required but illustrate the method signature.
...         assert pa.types.is_struct(storage_type)
...         assert pa.types.is_integer(storage_type[0].type)
...         assert storage_type[0].type == storage_type[1].type
...         assert serialized == b""
...
...         # return an instance of this subclass
...         return RationalType(storage_type[0].type)

特殊方法 __arrow_ext_serialize__ 和 __arrow_ext_deserialize__ 定義了擴充套件型別例項的序列化和反序列化。

這現在可以用於建立持有該擴充套件型別的陣列和表。

>>> rational_type = RationalType(pa.int32())
>>> rational_type.extension_name
'my_package.rational'
>>> rational_type.storage_type
StructType(struct<numer: int32, denom: int32>)
>>> storage_array = pa.array(
...     [
...         {"numer": 10, "denom": 17},
...         {"numer": 20, "denom": 13},
...     ],
...     type=rational_type.storage_type,
... )
>>> arr = rational_type.wrap_array(storage_array)
>>> # or equivalently
>>> arr = pa.ExtensionArray.from_storage(rational_type, storage_array)
>>> arr
<pyarrow.lib.ExtensionArray object at ...>
-- is_valid: all not null
-- child 0 type: int32
  [
    10,
    20
  ]
-- child 1 type: int32
  [
    17,
    13
  ]

此陣列可以包含在 RecordBatches 中，透過 IPC 傳送並在另一個 Python 程序中接收。接收程序必須顯式註冊該擴充套件型別以進行反序列化，否則它將回退到儲存型別。

>>> pa.register_extension_type(RationalType(pa.int32()))

例如，建立一個 RecordBatch 並使用 IPC 協議將其寫入流：

>>> batch = pa.RecordBatch.from_arrays([arr], ["ext"])
>>> sink = pa.BufferOutputStream()
>>> with pa.RecordBatchStreamWriter(sink, batch.schema) as writer:
...     writer.write_batch(batch)
>>> buf = sink.getvalue()

然後讀回它會得到正確的型別。

>>> with pa.ipc.open_stream(buf) as reader:
...     result = reader.read_all()
>>> result.column("ext").type
RationalType(StructType(struct<numer: int32, denom: int32>))

此外，請注意，雖然我們註冊了具體型別 RationalType(pa.int32())，但 RationalType(integer_type) 對 *所有* Arrow 整數型別都使用相同的副檔名（"my_package.rational"）。因此，上述程式碼也允許使用者 (反)序列化這些資料型別。

>>> big_rational_type = RationalType(pa.int64())
>>> storage_array = pa.array(
...     [
...         {"numer": 10, "denom": 17},
...         {"numer": 20, "denom": 13},
...     ],
...     type=big_rational_type.storage_type,
... )
>>> arr = big_rational_type.wrap_array(storage_array)
>>> batch = pa.RecordBatch.from_arrays([arr], ["ext"])
>>> sink = pa.BufferOutputStream()
>>> with pa.RecordBatchStreamWriter(sink, batch.schema) as writer:
...     writer.write_batch(batch)
>>> buf = sink.getvalue()
>>> with pa.ipc.open_stream(buf) as reader:
...     result = reader.read_all()
>>> result.column("ext").type
RationalType(StructType(struct<numer: int64, denom: int64>))

接收應用程式不必是 Python，但如果它已實現自己的擴充套件型別來接收它，它仍然可以識別該擴充套件型別為“my_package.rational”型別。如果該型別未在接收應用程式中註冊，它將回退到儲存型別。

引數化擴充套件型別

上面的例子展示瞭如何構建一個除了儲存型別外不需要額外元資料的擴充套件型別。但 Arrow 也提供了更靈活的引數化擴充套件型別。

此處給出的示例實現了 pandas “period” 資料型別 的擴充套件型別，表示時間跨度（例如，頻率為天、月、季度等）。它儲存為 int64 陣列，被解釋為自 1970 年以來的給定頻率的時間跨度數。

>>> class PeriodType(pa.ExtensionType):
...
...     def __init__(self, freq):
...         # attributes need to be set first before calling
...         # super init (as that calls serialize)
...         self._freq = freq
...         super().__init__(pa.int64(), "my_package.period")
...
...     @property
...     def freq(self):
...         return self._freq
...
...     def __arrow_ext_serialize__(self):
...         return "freq={}".format(self.freq).encode()
...
...     @classmethod
...     def __arrow_ext_deserialize__(cls, storage_type, serialized):
...         # Return an instance of this subclass given the serialized
...         # metadata.
...         serialized = serialized.decode()
...         assert serialized.startswith("freq=")
...         freq = serialized.split("=")[1]
...         return PeriodType(freq)

在這裡，我們確保將重建例項（在 __arrow_ext_deserialize__ 類方法中）所需的所有資訊儲存在序列化元資料中，本例中即頻率字串。

請注意，一旦建立，資料型別例項就被視為不可變的。因此，在上面的例子中，freq 引數被儲存在一個私有屬性中，並帶有一個用於訪問它的公共只讀屬性。

自定義擴充套件陣列類

預設情況下，所有具有擴充套件型別的陣列都被構建或反序列化為內建的 ExtensionArray 物件。然而，人們可能希望繼承 ExtensionArray 以新增特定於擴充套件型別的自定義邏輯。Arrow 允許透過向擴充套件型別的定義新增特殊方法 __arrow_ext_class__ 來實現這一點。

例如，讓我們考慮 Numpy 快速入門 中 3D 空間點的示例。我們可以將它們儲存為定長列表，我們希望能夠以 (N, 3) 的二維 Numpy 陣列形式提取資料，且無需任何複製。

>>> class Point3DArray(pa.ExtensionArray):
...     def to_numpy_array(self):
...         return self.storage.flatten().to_numpy().reshape((-1, 3))
>>> class Point3DType(pa.ExtensionType):
...     def __init__(self):
...         super().__init__(pa.list_(pa.float32(), 3), "my_package.Point3DType")
...
...     def __arrow_ext_serialize__(self):
...         return b""
...
...     @classmethod
...     def __arrow_ext_deserialize__(cls, storage_type, serialized):
...         return Point3DType()
...
...     def __arrow_ext_class__(self):
...         return Point3DArray

使用此擴充套件型別構建的陣列現在擁有預期的自定義陣列類。

>>> storage = pa.array([[1, 2, 3], [4, 5, 6]], pa.list_(pa.float32(), 3))
>>> arr = pa.ExtensionArray.from_storage(Point3DType(), storage)
>>> arr
<__main__.Point3DArray object at ...>
[
  [
    1,
    2,
    3
  ],
  [
    4,
    5,
    6
  ]
]

擴充套件類中的附加方法隨後可供使用者使用。

>>> arr.to_numpy_array()
array([[1., 2., 3.],
       [4., 5., 6.]], dtype=float32)

此陣列可以透過 IPC 傳送，在另一個 Python 程序中接收，並且自定義擴充套件陣列類將被保留（只要接收程序在讀取 IPC 資料之前使用 register_extension_type() 註冊了該擴充套件型別）。

自定義標量轉換

如果您希望當呼叫 ExtensionScalar.as_py() 時，您的自定義擴充套件型別的標量能轉換為自定義型別，您可以透過繼承 ExtensionScalar 來重寫 ExtensionScalar.as_py() 方法。例如，如果我們希望上面的 3D 點型別示例返回一個自定義的 3D 點類而不是列表，我們將實現：

>>> from collections import namedtuple
>>> Point3D = namedtuple("Point3D", ["x", "y", "z"])
>>> class Point3DScalar(pa.ExtensionScalar):
...     def as_py(self, **kwargs) -> Point3D:
...         return Point3D(*self.value.as_py(**kwargs))
>>> class Point3DType(pa.ExtensionType):
...     def __init__(self):
...         super().__init__(pa.list_(pa.float32(), 3), "my_package.Point3DType")
...
...     def __arrow_ext_serialize__(self):
...         return b""
...
...     @classmethod
...     def __arrow_ext_deserialize__(cls, storage_type, serialized):
...         return Point3DType()
...
...     def __arrow_ext_scalar_class__(self):
...         return Point3DScalar

使用此擴充套件型別構建的陣列現在提供的標量會轉換為我們的 Point3D 類。

>>> storage = pa.array([[1, 2, 3], [4, 5, 6]], pa.list_(pa.float32(), 3))
>>> arr = pa.ExtensionArray.from_storage(Point3DType(), storage)
>>> arr[0].as_py()
Point3D(x=1.0, y=2.0, z=3.0)
>>> arr.to_pylist()
[Point3D(x=1.0, y=2.0, z=3.0), Point3D(x=4.0, y=5.0, z=6.0)]

向 pandas 的轉換

如果您的擴充套件型別有對應的 pandas 擴充套件陣列，則帶有擴充套件型別的列向 pandas 的轉換（在 Table.to_pandas() 中）是可以控制的。

為此，需要實現 ExtensionType.to_pandas_dtype() 方法，該方法應返回一個 pandas.api.extensions.ExtensionDtype 子類例項。

以使用上述 pandas period 型別為例，程式碼如下：

>>> class PeriodType(pa.ExtensionType):
...
...     def to_pandas_dtype(self):
...         import pandas as pd
...         return pd.PeriodDtype(freq=self.freq)

其次，pandas 的 ExtensionDtype 反過來也需要實現 __from_arrow__ 方法：該方法給定擴充套件型別的 PyArrow Array 或 ChunkedArray，可以構建對應的 pandas ExtensionArray。此方法應具有以下簽名：

>>> import pandas as pd
>>> class MyExtensionDtype(pd.api.extensions.ExtensionDtype):
...
...     def __from_arrow__(self, array):  # pyarrow.Array/ChunkedArray -> pandas.ExtensionArray
...         pass

透過這種方式，您可以控制 PyArrow 擴充套件型別的 PyArrow Array 到可以儲存在 DataFrame 中的 pandas ExtensionArray 的轉換。

規範擴充套件型別

您可以在 規範擴充套件型別 部分找到規範擴充套件型別的官方列表。在這裡，我們添加了關於如何在 PyArrow 中使用它們的示例。

定形張量

要建立具有相同形狀的張量陣列（定形張量陣列），我們首先需要定義一個具有值型別和形狀的定形張量擴充套件型別。

>>> tensor_type = pa.fixed_shape_tensor(pa.int32(), (2, 2))

然後我們需要 pyarrow.list_() 型別的儲存陣列，其中 value_type 是定形張量的值型別，列表大小是 tensor_type 形狀元素的乘積。然後我們可以使用 pa.ExtensionArray.from_storage() 方法建立張量陣列。

>>> arr = [[1, 2, 3, 4], [10, 20, 30, 40], [100, 200, 300, 400]]
>>> storage = pa.array(arr, pa.list_(pa.int32(), 4))
>>> tensor_array = pa.ExtensionArray.from_storage(tensor_type, storage)

我們也可以建立另一個具有不同值型別的張量陣列。

>>> tensor_type_2 = pa.fixed_shape_tensor(pa.float32(), (2, 2))
>>> storage_2 = pa.array(arr, pa.list_(pa.float32(), 4))
>>> tensor_array_2 = pa.ExtensionArray.from_storage(tensor_type_2, storage_2)

擴充套件陣列可用作 pyarrow.Table 或 pyarrow.RecordBatch 中的列。

>>> data = [
...     pa.array([1, 2, 3]),
...     pa.array(["foo", "bar", None]),
...     pa.array([True, None, True]),
...     tensor_array,
...     tensor_array_2
... ]
>>> my_schema = pa.schema([("f0", pa.int8()),
...                        ("f1", pa.string()),
...                        ("f2", pa.bool_()),
...                        ("tensors_int", tensor_type),
...                        ("tensors_float", tensor_type_2)])
>>> table = pa.Table.from_arrays(data, schema=my_schema)
>>> table
pyarrow.Table
f0: int8
f1: string
f2: bool
tensors_int: extension<arrow.fixed_shape_tensor[value_type=int32, shape=[2,2]]>
tensors_float: extension<arrow.fixed_shape_tensor[value_type=float, shape=[2,2]]>
----
f0: [[1,2,3]]
f1: [["foo","bar",null]]
f2: [[true,null,true]]
tensors_int: [[[1,2,3,4],[10,20,30,40],[100,200,300,400]]]
tensors_float: [[[1,2,3,4],[10,20,30,40],[100,200,300,400]]]

我們還可以將張量陣列轉換為單個多維 numpy ndarray。透過這種轉換，Arrow 陣列的長度成為 numpy ndarray 的第一個維度。

>>> numpy_tensor = tensor_array_2.to_numpy_ndarray()
>>> numpy_tensor
array([[[  1.,   2.],
        [  3.,   4.]],

       [[ 10.,  20.],
        [ 30.,  40.]],

       [[100., 200.],
        [300., 400.]]], dtype=float32)
 >>> numpy_tensor.shape
 (3, 2, 2)

注意

兩個可選引數 permutation 和 dim_names 旨在為使用者提供關於資料邏輯佈局與物理佈局相比的資訊。

向 numpy ndarray 的轉換僅對平凡置換（None 或 [0, 1, ... N-1]，其中 N 是張量維數）是可能的。

反之亦然，我們可以將 numpy ndarray 轉換為定形張量陣列。

>>> pa.FixedShapeTensorArray.from_numpy_ndarray(numpy_tensor)
<pyarrow.lib.FixedShapeTensorArray object at ...>
[
  [
    1,
    2,
    3,
    4
  ],
  [
    10,
    20,
    30,
    40
  ],
  [
    100,
    200,
    300,
    400
  ]
]

透過轉換，ndarray 的第一個維度成為 PyArrow 擴充套件陣列的長度。我們在示例中可以看到，形狀為 (3, 2, 2) 的 ndarray 變成了長度為 3 的 Arrow 陣列，其張量元素形狀為 (2, 2)。

# ndarray of shape (3, 2, 2)
>>> numpy_tensor.shape
(3, 2, 2)

# arrow array of length 3 with tensor elements of shape (2, 2)
>>> pyarrow_tensor_array = pa.FixedShapeTensorArray.from_numpy_ndarray(numpy_tensor)
>>> len(pyarrow_tensor_array)
3
>>> pyarrow_tensor_array.type.shape
[2, 2]

擴充套件型別還可以定義 permutation 和 dim_names。例如：

>>> tensor_type = pa.fixed_shape_tensor(pa.float64(), [2, 2, 3], permutation=[0, 2, 1])

或

>>> tensor_type = pa.fixed_shape_tensor(pa.bool_(), [2, 2, 3], dim_names=["C", "H", "W"])

對於 NCHW 格式，其中：

• N：影像數量，在我們的例子中是陣列的長度，總是位於第一個維度

• C：影像的通道數

• H：影像的高度

• W：影像的寬度

UUID

UUID 擴充套件型別（arrow.uuid）將通用唯一識別符號表示為 16 位元組的定長二進位制值。PyArrow 提供了與 Python 內建 uuid 模組的整合，包括自動型別推斷。

建立 UUID 標量和陣列

PyArrow 從 Python 的 uuid.UUID 物件中推斷 UUID 型別，因此您可以直接將它們傳遞給 pyarrow.scalar() 和 pyarrow.array()。

>>> import uuid
>>> import pyarrow as pa

>>> pa.scalar(uuid.uuid4())
<pyarrow.UuidScalar: UUID('...')>

>>> uuids = [uuid.uuid4() for _ in range(3)]
>>> arr = pa.array(uuids)
>>> arr.type
UuidType(extension<arrow.uuid>)

您也可以使用 pyarrow.uuid() 顯式指定 UUID 型別。

>>> pa.array([uuid.uuid4(), uuid.uuid4()], type=pa.uuid())
<pyarrow.lib.UuidArray object at ...>
[
  ...,
  ...
]