内存和 IO 接口#

本节将向您介绍 PyArrow 内存管理和 IO 系统中的主要概念

缓冲区
内存池
类文件和类流对象

引用和分配内存#

pyarrow.Buffer#

Buffer 对象封装了 C++ 的 arrow::Buffer 类型，后者是 Apache Arrow 在 C++ 中进行内存管理的主要工具。它允许更高级别的数组类安全地与它们可能拥有或不拥有的内存进行交互。arrow::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

以这种方式创建缓冲区不会分配任何内存；它是从 data bytes 对象导出的内存的零拷贝视图。

也可以使用 foreign_buffer() 函数引用外部内存，形式为原始指针和大小。

在需要 Python 缓冲区或 memoryview 的情况下可以使用缓冲区，并且这种转换是零拷贝的

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

Buffer 的 to_pybytes() 方法将缓冲区的数据转换为 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 具有一些重要的特性，使其在可能的情况下比使用 PyArrow 的 Python 文件更可取

其他 Arrow 类可以原生访问内部 C++ IO 对象，而无需获取 Python GIL
原生 C++ IO 可能能够进行零拷贝 IO，例如使用内存映射

有几种 NativeFile 选项可用

OSFile，一个使用操作系统文件描述符的本机文件
MemoryMappedFile，用于使用内存映射进行读取（零拷贝）和写入
BufferReader，用于将 Buffer 对象作为文件读取
BufferOutputStream，用于在内存中写入数据，并在最后生成一个缓冲区
FixedSizeBufferWriter，用于将数据写入已分配的缓冲区
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 Buffers 的文件接口。

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 具有相似的语义。