内存和 IO 接口

内存和 IO 接口#

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

  • 缓冲区

  • 内存池

  • 类文件和流式对象

引用和分配内存#

pyarrow.Buffer#

Buffer 对象封装了 C++ 的 arrow::Buffer 类型，这是 Apache Arrow 在 C++ 中进行内存管理的主要工具。它允许更高级别的数组类安全地与它们可能拥有或可能不拥有的内存进行交互。arrow::Buffer 可以进行零拷贝切片，以允许 Buffers 廉价地引用其他 Buffers，同时保持内存生命周期和清晰的父子关系。

arrow::Buffer 有许多实现，但它们都提供了标准接口：数据指针和长度。这与 Python 的内置缓冲区协议和 memoryview 对象类似。

可以通过调用 py_buffer() 函数从任何实现了缓冲区协议的 Python 对象创建 Buffer。

考虑字节对象：

import pyarrow as pa

data = b'abcdefghijklmnopqrstuvwxyz'

buf = pa.py_buffer(data)

buf.size, buf

(26,
 <pyarrow.Buffer address=0x7f2eb4158d10 size=26 is_cpu=True is_mutable=False>)

以这种方式创建 Buffer 不会分配任何内存；它是对从数据字节对象导出的内存的零拷贝视图。

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

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

memoryview(buf)

<memory at 0x7f2e8c053700>

Buffer 的 to_pybytes() 方法将 Buffer 的数据转换为 Python 字节字符串（因此会复制数据）：

buf.to_pybytes()

b'abcdefghijklmnopqrstuvwxyz'

内存池#

所有内存分配和释放（如 C 中的 malloc 和 free）都在 pyarrow.MemoryPool 的一个实例中进行跟踪。这意味着我们可以精确地跟踪已分配的内存量：

pa.total_allocated_bytes()

0

从默认池分配一个可调整大小的 pyarrow.Buffer：

buf = pa.allocate_buffer(1024, resizable=True)
pa.total_allocated_bytes()

1024

buf.resize(2048)

pa.total_allocated_bytes()

2048

默认分配器以最小 64 字节的增量请求内存。如果缓冲区被垃圾回收，所有内存都会被释放：

buf = None

pa.total_allocated_bytes()

0

除了默认的内置内存池外，根据 Arrow 是如何构建的，还可以选择其他内存池（如 mimalloc）。可以获取内存池的后端名称：

pa.default_memory_pool().backend_name

'jemalloc'

参见

输入和输出#

Arrow C++ 库为不同种类的 IO 对象提供了几个抽象接口：

  • 只读流

  • 支持随机访问的只读文件

  • 只写流

  • 支持随机访问的只写文件

  • 支持读写和随机访问的文件

为了使得这些对象的行为更像 Python 的内置 file 对象，我们定义了 NativeFile 基类，它实现了与常规 Python 文件对象相同的 API。

NativeFile 具有一些重要特性，这使得在可能的情况下，使用它比使用 Python file 与 PyArrow 更为可取：

  • 其他 Arrow 类可以原生地访问内部 C++ IO 对象，无需获取 Python GIL

  • Native C++ IO 可能能够进行零拷贝 IO，例如使用内存映射

有多种 NativeFile 选项可供选择：

还有一些高级 API，使得实例化常见类型的流更加容易。

高级 API#

输入流#

input_stream() 函数允许从各种来源创建可读的 NativeFile。

buf = memoryview(b"some data")

stream = pa.input_stream(buf)

stream.read(4)

b'some'

  • 如果传入字符串或文件路径，它将打开给定的文件进行读取，创建 OSFile。可选地，文件可以是压缩的：如果其文件名以诸如 .gz 这样的可识别扩展名结尾，其内容将在读取时自动解压缩。

import gzip

with gzip.open('example.gz', 'wb') as f:
    f.write(b'some data\n' * 3)
stream = pa.input_stream('example.gz')
stream.read()

b'some data\nsome data\nsome data\n'

  • 如果传入 Python 文件对象，它将被包装在 PythonFile 中，以便 Arrow C++ 库可以从中读取数据（这会带来一些开销）。

输出流#

output_stream() 是用于输出流的等效函数，允许创建可写的 NativeFile。它具有与上述 input_stream() 相同的特性，例如能够写入缓冲区或进行实时压缩。

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


f = open('example1.dat', 'rb')

f.read()

b'some data'

磁盘上的文件和内存映射文件#

PyArrow 包括两种与磁盘上的数据交互的方式：标准操作系统级别的文件API，以及内存映射文件。在常规Python中，我们可以编写：

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

使用 pyarrow.OSFile 类，你可以编写：

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

对于读取文件，你可以使用 pyarrow.OSFile 或 pyarrow.MemoryMappedFile。它们之间的区别在于，pyarrow.OSFile 在每次读取时都会分配新的内存，就像 Python 文件对象一样。在从 pyarrow.MemoryMappedFile 中读取时，库会构造引用映射内存的缓冲区，而不进行任何内存分配或复制：

file_obj = pa.OSFile('example2.dat')

mmap = pa.memory_map('example3.dat')

file_obj.read(4)

b'some'

mmap.read(4)

b'some'

read 方法实现了标准的 Python file 读取 API。要读取到 Arrow Buffer 对象中，使用 read_buffer：

mmap.seek(0)

0

buf = mmap.read_buffer(4)

print(buf)

buf.to_pybytes()

<pyarrow.Buffer address=0x7f2eba7f6000 size=4 is_cpu=True is_mutable=False>

b'some'

PyArrow 中的许多工具，特别是 Apache Parquet 接口以及文件和流消息工具，在使用这些 NativeFile 类型时比使用常规 Python 文件对象更高效。

内存中读写#

为了帮助内存数据的序列化和反序列化，我们提供了可以读写 Arrow Buffers 的文件接口。

writer = pa.BufferOutputStream()

writer.write(b'hello, friends')

buf = writer.getvalue()

print(buf)

print(buf.size)

reader = pa.BufferReader(buf)

<pyarrow.Buffer address=0x7f2ea4210000 size=14 is_cpu=True is_mutable=True>
14

reader.seek(7)

7

reader.read(7)

b'friends'

这些接口与 Python 内置的 io.BytesIO 具有相似的语义。