如果数据框基于列值的过滤器，则从字典中提取行数据

Question

The dictionary dict_set has dataframes as the value for their keys.

I'm trying to extract data from a dictionary of dataframes based on a filter on 'A' column in the dataframe based on the value in column.

dict_set={}
dict_set['a']=pd.DataFrame({'A':[1,2,3],'B':[1,2,3]})
dict_set['b']=pd.DataFrame({'A':[1,4,5],'B':[1,5,6]})
    
df=pd.concat([dict_set[x][dict_set[x]['A']==1] for x in dict_set.keys()],axis=0)

output being the below.

   A  B
0  1  1
0  1  1

But I would want the output to be

   A  B  x
0  1  1  a
0  1  1  b

Basically, I want the value of x to be present in the new dataframe formed as a column, say column x in the dataframe formed such that df[x] would give me the x values. Is there a simple way to do this?

Answer 1

Try this:

pd.concat([df.query("A == 1") for df in dict_set.values()], keys=dict_set.keys())\
  .reset_index(level=0)\
  .rename(columns={'level_0':'x'})

Output:

   x  A  B
0  a  1  1
0  b  1  1

Details:

Let's get the dataframes from the dictionary using list comprehension and filter the datafames. Here, I choose to use query , but you could use boolean index with df[df['A'] == 1] also, then pd.concat with the keys parameter set to the dictionary keys. Lastly, reset_index level=0 and rename .