Pyarrow 是否存在内存泄漏问题？

Question

为了解析一个较大的文件，我需要循环连续写入大量的parquet文件。然而，每次迭代时这个任务消耗的内存都会增加，而我希望它保持不变（因为在内存中不应该添加任何东西）。这使得扩展变得棘手。

我已经添加了一个最小可重现示例，它创建了 10,000 个parquet，并对其进行循环追加。

import resource
import random
import string
import pyarrow as pa
import pyarrow.parquet as pq
import pandas as pd


def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
    return ''.join(random.choice(chars) for _ in range(size))

schema = pa.schema([
                        pa.field('test', pa.string()),
                    ])

resource.setrlimit(resource.RLIMIT_NOFILE, (1000000, 1000000))
number_files = 10000
number_rows_increment = 1000
number_iterations = 100

writers = [pq.ParquetWriter('test_'+id_generator()+'.parquet', schema) for i in range(number_files)]

for i in range(number_iterations):
    for writer in writers:
        table_to_write = pa.Table.from_pandas(
                            pd.DataFrame({'test': [id_generator() for i in range(number_rows_increment)]}),
                            preserve_index=False,
                            schema = schema,
                            nthreads = 1)
        table_to_write = table_to_write.replace_schema_metadata(None)
        writer.write_table(table_to_write)
    print(i)

for writer in writers:
    writer.close()

有人知道是什么原因导致这个漏洞并如何预防吗？

- Abel Riboulot

---

你能说明一下你的pandas版本吗？ - Uwe L. Korn

Pandas: 0.22.0 PyArrow: 0.10.0 - Abel Riboulot

请更新到pandas>=0.23。Pandas存在泄漏问题，也会影响到pyarrow。 - Uwe L. Korn

我尝试过了，但仍然存在内存泄漏问题。 - Abel Riboulot

更新到pyarrow==0.15.0对我很有帮助。 - Dima Fomin

Answer 1

我们不确定出了什么问题，但是一些其他用户报告还未诊断出的内存泄漏。我将您的示例添加到其中一个跟踪 JIRA 问题中 https://issues.apache.org/jira/browse/ARROW-3324

- Wes McKinney

Answer 2

2022年更新：

我花了几天时间解决来自pyarrow的内存泄漏问题。请参见这里以更好地理解。我将在下面贴出关键点。基本上，他们说这不是一个库内存泄漏问题，而是一种常见行为。 Pyarrow使用jemalloc，一个自定义内存分配器，尽力保留从操作系统分配的内存（因为这可能是一个昂贵的操作）。不幸的是，这使得使用像memory_profiler这样的工具逐行跟踪内存使用变得困难。有几个选项：

1. 您可以使用这个库函数pyarrow.total_allocated_bytes来跟踪分配，而不是使用memory_profiler。
2. 您还可以在脚本顶部放置以下行，这将配置jemalloc立即释放内存，而不是保留它（这可能会产生一些性能影响）。但是，我使用它，但没有效果。

import pyarrow as pa
pa.jemalloc_set_decay_ms(0)

你看到的行为是jemalloc相当典型的。如果想进一步了解，你也可以查看其他问题，以获取更多有关jemalloc行为的讨论和示例：

https://issues.apache.org/jira/browse/ARROW-6910

https://issues.apache.org/jira/browse/ARROW-7305

- user3503711

Answer 3

与@user3503711类似，我们在PyArrow中遇到了严重的内存分配问题。PyArrow在macOS上使用其“mimalloc”后端而不是“jemalloc”。PyArrow为自身分配的内存以千兆字节为单位，导致我们的并行pytest测试运行器崩溃（每个工作进程都分配了大量内存，导致Github Actions默默失败）。请注意，这主要是在大量读取和少量写入的情况下。

我们无法完全消除PyArrow的内存问题，并怀疑存在内部泄漏。但是，我们设法在一定程度上减轻了问题，并降低了内存消耗，以便测试运行能够完成。

尝试设置参数给`pyarrow.parquet.read_table`函数。你可能想要调整的一些参数有`read_table(f, use_threads=False, pre_buffer=False, memory_map=True)`。

如果你正在使用`pytest`测试库，请确保不要在测试固定装置中使用Arrow/DataFrame，因为pytest可能会在幕后缓存这些对象。

使用Python的`del`关键字显式地删除你的对象。

尝试强制`pyarrow`进行垃圾回收。

    import gc
    gc.collect()
    import pyarrow
    pool = pyarrow.default_memory_pool()
    pool.release_unused()

    p = psutil.Process()
    rss = p.memory_info().rss

    print(f"Pool bytes {pool.bytes_allocated():,} max memory {pool.max_memory():,}, RSS after cleaning is {rss:,}")

此外，

• 如果可能的话，请尝试使用FastParquet库代替PyArrow。可能会有兼容性问题阻止您这样做。

- Mikko Ohtamaa