使用pyarrow的Pandas在拆分数据框时不占用额外内存

Question

在使用 float64[pyarrow] 数据类型的 Pandas 2.2.1 时，发现把一个数据表分成两部分后再合并回来，似乎并没有使用额外的内存。

而如果使用普通的 float64 数据类型，内存的使用量是原数据表的三倍（这也是我直觉上认为的，如果数据表被复制的话）。

有没有人能解释一下这是为什么？这显然是个好事，但在 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), 33685.60 MB (VMS),  68.29 MB (Peak RSS)
Memory usage: 831.38 MB (RSS), 34448.54 MB (VMS),  831.38 MB (Peak RSS)
Memory usage: 1594.66 MB (RSS), 35211.49 MB (VMS),  1594.66 MB (Peak RSS)
Memory usage: 2346.45 MB (RSS), 35974.43 MB (VMS),  2346.45 MB (Peak 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默认就开启了写时复制？不过我在文档里找不到相关的信息。