python - pd.DataFrame.set_index 可以维护数据类型吗？

Question

我正在尝试调用 df.set_index，使我设置索引的列的 dtype 是新的 index.dtype。不幸的是，在下面的示例中，set_index 更改了 dtype。

df = pd.DataFrame({'a': pd.Series(np.array([-1, 0, 1, 2], dtype=np.int8))})
df['ignore'] = df['a']
assert (df.dtypes == np.int8).all() # fine
df2=  df.set_index('a')
assert df2.index.dtype == df['a'].dtype, df2.index.dtype

是否可以避免这种行为？我的 pandas 版本是 0.23.3

同样，

new_idx = pd.Index(np.array([-1, 0, 1, 2]), dtype=np.dtype('int8'))
assert new_idx.dtype == np.dtype('int64')

尽管 dtype 参数的文档说:“如果提供了实际的 dtype，我们会在安全的情况下强制使用该 dtype。 否则，将引发错误。”

Answer 1

尽管我在上面的评论中夸大其词，但这可能足以获得一个适当的索引，该索引既低内存又从 -1 开始。

pandas.RangeIndex

采用开始和停止参数，如 range

df = df.set_index(pd.RangeIndex(-1, len(df) - 1))

print(df.index, df.index.dtype, sep='\n')

这应该非常节省内存。

尽管它仍然是 dtype int64(您应该想要的)，但它占用的内存非常少。

pd.RangeIndex(-1, 4000000).memory_usage()

84

和

for i in range(1, 1000000, 100000):
  print(pd.RangeIndex(-1, i).memory_usage())

84
84
84
84
84
84
84
84
84
84