Overload df.getitem with idx of tuple#573
Overload df.getitem with idx of tuple#573AlexanderKalistratov merged 2 commits intoIntelPython:masterfrom
Conversation
| # SDC pd.df.getitem does not support idx as a list | ||
| sdc_func = self.jit(lambda df: df[('A', 'C')]) |
There was a problem hiding this comment.
I think we should state this as a limitation in our documentation too.
Also it might look better if you write it this way:
test_impl = lambda df: df[['A', 'C']]
sdc_func = self.jit(lambda df: df[('A', 'C')])
or use do_jit parameter, e.g.
# Pandas implementation does not support idx as a tuple
# but SDC implementation only supports tuple, so adjust tested code
def test_impl(df, do_jit):
if do_jit == True: # noqa
return df[('A', 'C')]
else:
return df[['A', 'C']]
sdc_func = self.jit(test_impl)
pd.testing.assert_frame_equal(sdc_func(df, True), test_impl(df, False))
| func_lines += df_getitem_tuple_idx_main_codelines(self, literal_idx) | ||
| else: | ||
| # raise KeyError if input DF is empty or idx is invalid | ||
| func_lines += [' raise KeyError'] |
There was a problem hiding this comment.
The exception should really have a message, maybe something like "Column is not in the DataFrame" would fit?
| return func_text, global_vars | ||
|
|
||
|
|
||
| def gen_df_getitem_impl_generator(codegen, impl_name): |
There was a problem hiding this comment.
This looks as a very generic function, that we could use everywhere else, so it might be useful to move it to sdc_typing_utils. So it's responsibility would be to call exec on provided func_code and return back created func object.
|
|
||
| return _df_getitem_unicode_idx_impl | ||
|
|
||
| if isinstance(idx, types.Tuple): |
There was a problem hiding this comment.
This should probably be a more strict check, e.g.
| if isinstance(idx, types.Tuple): | |
| if (isinstance(idx, types.Tuple) | |
| and all([isinstance(item, types.StringLiteral) for item in idx])): |
kozlov-alexey
left a comment
There was a problem hiding this comment.
Looks very good! No major comments from me, only minor ones, that might be applied later on.
No description provided.