Как сохранить идентификатор dataframe pandas при извлечении одной строки

Question

Я извлекаю подмножество своей DataFrame по индексу, используя либо .xs, либо .loc (они, похоже, ведут себя одинаково). Когда мое условие извлекает несколько строк, результат остается в кадре данных. Когда извлекается только одна строка, она автоматически преобразуется в серию. Я не хочу этого поведения, так как это означает, что мне нужно обрабатывать несколько случаев downstream (разные наборы методов, доступные для серии vs dataframe).

In [1]: df = pd.DataFrame({'a':range(7), 'b':['one']*4 + ['two'] + ['three']*2, 
'c':range(10,17)}) 

In [2]: df.set_index('b', inplace=True) 

In [3]: df.xs('one') 
Out[3]: 
    a c 
b 
one 0 10 
one 1 11 
one 2 12 
one 3 13 

In [4]: df.xs('two') 
Out[4]: 
a  4 
c 14 
Name: two, dtype: int64 

In [5]: type(df.xs('two')) 
Out [5]: pandas.core.series.Series 

можно вручную преобразовать эту серию обратно в dataframe, но это, кажется громоздким и также потребует случай тестирования, чтобы увидеть, если я должен сделать это. Есть ли более чистый способ просто получить DataFrame для начала?

Answer 1

IIUC, вы можете просто добавить фигурные скобки, [] и использовать .loc:

>>> df.loc["two"] 
a  4 
c 14 
Name: two, dtype: int64 
>>> type(_) 
<class 'pandas.core.series.Series'> 
>>> df.loc[["two"]] 
    a c 
b   
two 4 14 

[1 rows x 2 columns] 
>>> type(_) 
<class 'pandas.core.frame.DataFrame'> 

Это может напомнить вам о том, как работает numpy расширенные индексации:

>>> a = np.arange(9).reshape(3,3) 
>>> a[1] 
array([3, 4, 5]) 
>>> a[[1]] 
array([[3, 4, 5]]) 

Теперь это, вероятно, потребуется некоторый рефакторинг кода, так что вы всегда получаете доступ к списку, даже если в списке есть только один элемент, но он работает хорошо для меня на практике.