Pandas 2.2.1中float64[pyarrow]拆分合并DataFrame无额外内存占用原因咨询

在Pandas 2.2.1中，使用float64[pyarrow]数据类型时，拆分DataFrame后再合并几乎不会占用额外内存；而使用常规float64时，内存占用会达到原DataFrame的3倍，符合数据被多次复制的预期。

测试代码 #

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()

常规float64输出 #

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

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)

背后原因 #

核心差异来自两种数据类型的底层存储机制和Pandas的处理逻辑：

• 常规float64（NumPy后端）：
  NumPy的ndarray采用行优先存储，每一行的数据是连续的内存块。当通过loc选择列时，这些列在内存中是非连续的，NumPy无法创建视图，必须复制数据生成新的ndarray。因此拆分出的split1和split2各自占用和原DataFrame一半相当的内存，加上原DataFrame，内存翻倍。后续pd.concat合并时，又需要把两个副本的数据复制到新的ndarray中，最终内存达到原DataFrame的3倍。

• float64[pyarrow]（PyArrow后端）：
  PyArrow采用列优先存储，每一列是独立的连续内存块。Pandas对PyArrow数据的列选择操作，只是创建原PyArrow Table的视图引用，完全不复制底层数据。拆分出的split1和split2本质上是指向原数据列的指针集合，几乎不占用额外内存。而pd.concat合并时，只是将这些列视图重新组合成一个新的PyArrow Table，依然共享原数据的内存块，没有任何数据复制。只有当对视图中的数据进行修改时，才会触发写时复制（Copy-on-Write）机制生成新数据，而测试中仅进行拆分合并操作，因此内存几乎没有变化。

内容的提问来源于stack exchange，提问作者Neil