Pandas join with categorical indices

Viewed 114

I'm getting unexpected results with attempting to join Pandas DataFrame objects with categorical indices. Here is the minimum reproducible example (boiled down from my real-life use case):

shape_categories=['square', 'circle']
color_categories=['red', 'blue', 'green']

test_a = pd.DataFrame({
    'shape': pd.Categorical(['square', 'circle'], categories=shape_categories, ordered=True),
    'color': pd.Categorical(['red', 'blue'], categories=color_categories, ordered=True),
    'value_a': [1.0, 2.0]
})
test_a.set_index(['shape', 'color'], inplace=True)

test_b = pd.DataFrame({
    'shape': pd.Categorical(['square', 'square', 'circle', 'circle'], categories=shape_categories, ordered=True),
    'color': pd.Categorical(['red', 'blue', 'red', 'blue'], categories=color_categories, ordered=True),
    'value_b': [10.0, np.nan, np.nan, 40.0]
})
test_b.set_index(['shape', 'color'], inplace=True)

test_a.join(test_b, how='left')

I expect

shape color value_a value_b
square red 1.0 10.0
circle blue 2.0 40.0

but instead I get

shape color value_a value_b
square red 1.0 10.0
circle blue 2.0 NaN

What am I missing? I've tried to be careful to keep the dtypes of the categorical variables exactly the same.

1 Answers

You need to specify all of the indices to merge on, otherwise pandas will only use the first one. So in your case:

test_a.join(test_b, how='left', on = ['shape', 'color'])

should produce the expected result.

Related