Может ли pd.DataFrame.set_index поддерживать dtype?

Question

Я пытаюсь вызвать df.set_index таким образом, чтобы dtype столбца, для которого я установлен set_index, был новым 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

Можно ли избежать такого поведения? Моя версия для панд 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