python - Pandas 分层列和 csv 函数

Question

是否有可能以尊重分层列结构的方式通过 csv 来回传输 DataFrame？换句话说，如果我有以下 DataFrame:

>>> cols = pd.MultiIndex.from_arrays([["foo", "foo", "bar", "bar"],
                                      ["a", "b", "c", "d"]])
>>> df = pd.DataFrame(np.random.randn(5, 4), index=range(5), columns=cols)

执行以下操作失败:

>>> df.to_csv("df.csv", index_label="index")
>>> df_new = pd.read_csv("df.csv", index_col="index")
>>> assert df.columns == df_new.columns

我是否遗漏了 csv 保存/读取步骤中的某些选项？

Answer 1

在特殊情况下，您有一个列式 MultiIndex，但索引很简单，您可以转置 DataFrame 并使用 index_label 和 index_col，如下所示:

import numpy as np
import pandas as pd

cols = pd.MultiIndex.from_arrays([["foo", "foo", "bar", "bar"],
                                  ["a", "b", "c", "d"]])

df = pd.DataFrame(np.random.randn(5, 4), index=range(5), columns=cols)

(df.T).to_csv('/tmp/df.csv', index_label=['first','second'])
df_new = pd.read_csv('/tmp/df.csv', index_col=['first','second']).T
assert np.all(df.columns.values == df_new.columns.values)

但不幸的是，这引出了一个问题，如果索引和列都是多索引该怎么办？

---

这是一个 hacky 解决方法:

import numpy as np
import pandas as pd
import ast

cols = pd.MultiIndex.from_arrays([["foo", "foo", "bar", "bar"],
                                  ["a", "b", "c", "d"]])

df = pd.DataFrame(np.random.randn(5, 4), index=range(5), columns=cols)
print(df)

df.to_csv('/tmp/df.csv', index_label='index')
df_new = pd.read_csv('/tmp/df.csv', index_col='index')

columns = pd.MultiIndex.from_tuples([ast.literal_eval(item) for item in df_new.columns])
df_new.columns = columns
df_new.index.name = None
print(df_new)
assert np.all(df.columns.values == df_new.columns.values)

---

当然，如果你只是想将DataFrame存储在任何格式的文件中，那么df.save和pd.load提供了一个更愉快的解决方案:

import numpy as np
import pandas as pd

cols = pd.MultiIndex.from_arrays([["foo", "foo", "bar", "bar"],
                                  ["a", "b", "c", "d"]])

df = pd.DataFrame(np.random.randn(5, 4), index=range(5), columns=cols)

df.save('/tmp/df.df')
df_new = pd.load('/tmp/df.df')
assert np.all(df.columns.values == df_new.columns.values)