KeyError when joining on index level

[andyvanyperenAM]

When using the index level (multi-index) as join key, an KeyError is raised.
The multi-index originates for example from a groupby operation in previous steps.
In pandas, it is allowed to join on index keys.

example.

toy_pd = pd.DataFrame(columns = ['day','item','size'], data = [[5, 0, 500],[5, 0, 550],[5, 1, 1500],[5, 1, 700],[5, 1, 900],
                                                               [6, 0, 400],[6, 0, 300],[6, 0, 600], [6, 1, 800],[6, 1, 200],
                                                               [7, 0, 600],[7, 1, 700],[7, 1, 700], [7, 2, 750],[7, 2, 500]])

toy_ks1 = ks.from_pandas(toy_pd).groupby(['day','item']).agg({'size':'mean'})
toy_ks2 = ks.from_pandas(toy_pd).groupby(['day','item']).agg({'size':'mean'})
toy_ks1.join(toy_ks2, on = ['day','item'], rsuffix='r')

results in the error

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-36-43ac3f9768b4> in <module>
----> 1 toy_ks1.join(toy_ks2, on = ['day','item'], rsuffix='r')

~\AppData\Local\Continuum\anaconda3\envs\spark_env\lib\site-packages\databricks\koalas\frame.py in join(self, right, on, how, lsuffix, rsuffix)
   6928             )
   6929         if on:
-> 6930             self = self.set_index(on)
   6931             join_kdf = self.merge(
   6932                 right, left_index=True, right_index=True, how=how, suffixes=(lsuffix, rsuffix)

~\AppData\Local\Continuum\anaconda3\envs\spark_env\lib\site-packages\databricks\koalas\frame.py in set_index(self, keys, drop, append, inplace)
   3182         for key in keys:
   3183             if key not in columns:
-> 3184                 raise KeyError(key)
   3185         keys = [key if isinstance(key, tuple) else (key,) for key in keys]
   3186 

KeyError: 'day'

Anyone else with this problem? Is it fundamendal or a bug? Possible workarounds in the first case?

Koalas version 1.2.0
OS: Windows 10
python 3.7.6

[itholic]

Thanks for the reporting, @andyvanyperenAM !
It seems like bug, let me check !

[andyvanyperenAM]

Thanks @itholic, I tried the bug fix in your PR and it resolves my issue!

[itholic]

Nice, @andyvanyperenAM .
Good to hear that!! 👍