带有 pyarrow 的 Pandas 在分割数据帧时不使用额外的内存

Question

在 Pandas 2.2.1 中使用

float64[pyarrow]

dtype 时，将数据帧分成两部分，然后再次将其重新连接在一起时，似乎没有使用额外的内存。

当使用常规

float64

dtype 时，这会使用原始数据帧内存的 3 倍（如果复制数据帧，这就是我的直觉预期）。

谁能解释为什么会发生这种情况？这显然是一件好事，但这似乎并没有列在 pyarrow 的好处下，所以我想了解为什么会发生这种情况。

我正在运行的代码是：

import gc
import os

import numpy as np
import pandas as pd
import psutil


def log_memory_usage():
    gc.collect()
    pid = os.getpid()
    p = psutil.Process(pid)
    full_info = p.memory_info()
    print(
        f"Memory usage: {full_info.rss / 1024 / 1024:.2f} MB (RSS)"
    )


log_memory_usage()
df1 = pd.DataFrame(np.ones(shape=(10000000, 10)), columns=[f"col_{i}" for i in range(10)], dtype="float64")
log_memory_usage()
split1 = df1.loc[:, df1.columns[:5]]
split2 = df1.loc[:, df1.columns[5:]]
log_memory_usage()
joined_again = pd.concat([split1, split2], axis=1)
log_memory_usage()

打印：

Memory usage: 68.28 MB (RSS)
Memory usage: 831.38 MB (RSS)
Memory usage: 1594.66 MB (RSS)
Memory usage: 2346.45 MB (RSS)

因此每次拆分和连接数据帧都会使用额外的内存。

但是如果我将

dtype="float64"

更改为

dtype="float64[pyarrow]"

我会得到以下输出：

Memory usage: 68.14 MB (RSS)
Memory usage: 833.51 MB (RSS)
Memory usage: 833.84 MB (RSS)
Memory usage: 833.93 MB (RSS)

因此，数据帧的拆分和连接版本似乎只使用了很少的额外内存。

Answer 1

这似乎是一种“写时复制”优化。如果我将其添加到代码中： pd.set_option("mode.copy_on_write", True)

两种数据类型都会减少内存使用量。

那么也许 pyarrow 默认启用了写时复制？但我在任何地方的文档中都找不到这个。