記憶體與 IO 介面#

本節將向您介紹 PyArrow 記憶體管理和 IO 系統的主要概念

緩衝區
記憶體池
類檔案與類串流物件

參考與配置記憶體#

pyarrow.Buffer#

Buffer 物件封裝了 C++ arrow::Buffer 類型，它是 Apache Arrow 在 C++ 中用於記憶體管理的主要工具。它允許更高階的陣列類別安全地與它們可能擁有或不擁有的記憶體互動。arrow::Buffer 可以進行零複製切片，以允許 Buffer 以低成本參考其他 Buffer，同時保留記憶體生命週期和清晰的父子關係。

arrow::Buffer 有許多實作，但它們都提供標準介面：資料指標和長度。這類似於 Python 的內建 buffer 協定 和 memoryview 物件。

可以透過呼叫 py_buffer() 函數，從任何實作緩衝區協定的 Python 物件建立 Buffer。讓我們考慮一個 bytes 物件

In [1]: import pyarrow as pa

In [2]: data = b'abcdefghijklmnopqrstuvwxyz'

In [3]: buf = pa.py_buffer(data)

In [4]: buf
Out[4]: <pyarrow.Buffer address=0x7fe3eab1d850 size=26 is_cpu=True is_mutable=False>

In [5]: buf.size
Out[5]: 26

以這種方式建立 Buffer 不會配置任何記憶體；它是 data bytes 物件匯出的記憶體的零複製檢視。

外部記憶體，以原始指標和大小的形式，也可以使用 foreign_buffer() 函數來參考。

Buffer 可以用於需要 Python 緩衝區或 memoryview 的情況，並且這種轉換是零複製的

In [6]: memoryview(buf)
Out[6]: <memory at 0x7fe3e3ff5840>

Buffer 的 to_pybytes() 方法將 Buffer 的資料轉換為 Python 位元組字串（因此會複製資料）

In [7]: buf.to_pybytes()
Out[7]: b'abcdefghijklmnopqrstuvwxyz'

記憶體池#

所有記憶體配置和釋放（例如 C 中的 malloc 和 free）都在 MemoryPool 的實例中追蹤。這表示我們可以精確地追蹤已配置的記憶體量

In [8]: pa.total_allocated_bytes()
Out[8]: 56320

讓我們從預設池配置一個可調整大小的 Buffer

In [9]: buf = pa.allocate_buffer(1024, resizable=True)

In [10]: pa.total_allocated_bytes()
Out[10]: 57344

In [11]: buf.resize(2048)

In [12]: pa.total_allocated_bytes()
Out[12]: 58368

預設配置器以 64 位元組的最小增量請求記憶體。如果緩衝區被垃圾回收，則所有記憶體都會被釋放

In [13]: buf = None

In [14]: pa.total_allocated_bytes()
Out[14]: 56320

除了預設的內建記憶體池之外，可能還有其他記憶體池可供選擇（例如 jemalloc），具體取決於 Arrow 的建置方式。可以取得記憶體池的後端名稱

>>> pa.default_memory_pool().backend_name
'mimalloc'

另請參閱

記憶體池的 API 文件.

另請參閱

使用 Arrow 的可選 CUDA 整合的 GPU 上緩衝區。

輸入與輸出#

Arrow C++ 函式庫針對不同類型的 IO 物件有幾個抽象介面

唯讀串流
支援隨機存取的唯讀檔案
僅寫入串流
支援隨機存取的僅寫入檔案
支援讀取、寫入和隨機存取的檔案

為了使這些物件的行為更像 Python 的內建 file 物件，我們定義了一個 NativeFile 基底類別，它實作了與一般 Python 檔案物件相同的 API。

NativeFile 具有一些重要的功能，使其在可能的情況下比將 Python 檔案與 PyArrow 一起使用更為可取

其他 Arrow 類別可以原生存取內部 C++ IO 物件，而無需取得 Python GIL
原生 C++ IO 可能能夠進行零複製 IO，例如使用記憶體映射

有幾種可用的 NativeFile 選項

OSFile，一種使用您作業系統檔案描述器的原生檔案
MemoryMappedFile，用於使用記憶體映射進行讀取（零複製）和寫入
BufferReader，用於將 Buffer 物件作為檔案讀取
BufferOutputStream，用於在記憶體中寫入資料，並在最後產生一個 Buffer
FixedSizeBufferWriter，用於將資料寫入已配置的 Buffer 中
HdfsFile，用於讀取和寫入資料到 Hadoop 檔案系統
PythonFile，用於在 C++ 中與 Python 檔案物件介接
CompressedInputStream 和 CompressedOutputStream，用於對另一個串流進行即時壓縮或解壓縮

還有更高等級的 API，可以更輕鬆地實例化常見的串流類型。

高階 API#

輸入串流#

input_stream() 函數允許從各種來源建立可讀取的 NativeFile。

如果傳遞 Buffer 或 memoryview 物件，將會返回 BufferReader

In [15]: buf = memoryview(b"some data")

In [16]: stream = pa.input_stream(buf)

In [17]: stream.read(4)
Out[17]: b'some'

如果傳遞字串或檔案路徑，它將開啟磁碟上給定的檔案以進行讀取，並建立 OSFile。或者，檔案可以被壓縮：如果其檔名以可辨識的副檔名結尾，例如 .gz，則其內容將在讀取時自動解壓縮。

In [18]: import gzip

In [19]: with gzip.open('example.gz', 'wb') as f:
   ....:     f.write(b'some data\n' * 3)
   ....: 

In [20]: stream = pa.input_stream('example.gz')

In [21]: stream.read()
Out[21]: b'some data\nsome data\nsome data\n'

如果傳遞 Python 檔案物件，它將被封裝在 PythonFile 中，以便 Arrow C++ 函式庫可以從中讀取資料（但會產生輕微的效能負擔）。

輸出串流#

output_stream() 是輸出串流的等效函數，並允許建立可寫入的 NativeFile。它具有與上述 input_stream() 相同的功能，例如能夠寫入緩衝區或進行即時壓縮。

In [22]: with pa.output_stream('example1.dat') as stream:
   ....:     stream.write(b'some data')
   ....: 

In [23]: f = open('example1.dat', 'rb')

In [24]: f.read()
Out[24]: b'some data'

磁碟上和記憶體映射檔案#

PyArrow 包含兩種與磁碟上資料互動的方式：標準作業系統層級的檔案 API，以及記憶體映射檔案。在常規 Python 中，我們可以寫入

In [25]: with open('example2.dat', 'wb') as f:
   ....:     f.write(b'some example data')
   ....: 

使用 pyarrow 的 OSFile 類別，您可以寫入

In [26]: with pa.OSFile('example3.dat', 'wb') as f:
   ....:     f.write(b'some example data')
   ....: 

對於讀取檔案，您可以使用 OSFile 或 MemoryMappedFile。它們之間的區別在於 OSFile 在每次讀取時都會分配新的記憶體，就像 Python 檔案物件一樣。在從記憶體映射讀取時，函式庫會建構一個引用映射記憶體的緩衝區，而無需任何記憶體分配或複製

In [27]: file_obj = pa.OSFile('example2.dat')

In [28]: mmap = pa.memory_map('example3.dat')

In [29]: file_obj.read(4)
Out[29]: b'some'

In [30]: mmap.read(4)
Out[30]: b'some'

read 方法實作了標準 Python 檔案 read API。要讀取到 Arrow Buffer 物件中，請使用 read_buffer

In [31]: mmap.seek(0)
Out[31]: 0

In [32]: buf = mmap.read_buffer(4)

In [33]: print(buf)
<pyarrow.Buffer address=0x7fe4798a0000 size=4 is_cpu=True is_mutable=False>

In [34]: buf.to_pybytes()
Out[34]: b'some'

PyArrow 中的許多工具，尤其是 Apache Parquet 介面以及檔案和串流訊息傳遞工具，在使用這些 NativeFile 類型時，比使用普通的 Python 檔案物件更有效率。

記憶體內讀取和寫入#

為了協助記憶體內資料的序列化和反序列化，我們提供了可以讀取和寫入 Arrow Buffer 的檔案介面。

In [35]: writer = pa.BufferOutputStream()

In [36]: writer.write(b'hello, friends')
Out[36]: 14

In [37]: buf = writer.getvalue()

In [38]: buf
Out[38]: <pyarrow.Buffer address=0x200000601c0 size=14 is_cpu=True is_mutable=True>

In [39]: buf.size
Out[39]: 14

In [40]: reader = pa.BufferReader(buf)

In [41]: reader.seek(7)
Out[41]: 7

In [42]: reader.read(7)
Out[42]: b'friends'

這些具有與 Python 內建 io.BytesIO 相似的語義。