hiframes_typed.py 文件源码

def _handle_df_col_filter(self, lhs_name, rhs, assign):
        # find df['col2'] = df['col1'][arr]
        # since columns should have the same size, output is filled with NaNs
        # TODO: check for float, make sure col1 and col2 are in the same df
        if (rhs.op=='getitem'
                and rhs.value.name in self.df_cols
                and lhs_name in self.df_cols
                and self.is_bool_arr(rhs.index.name)):
            lhs = assign.target
            in_arr = rhs.value
            index_var = rhs.index
            f_blocks = compile_to_numba_ir(_column_filter_impl_float,
                    {'numba': numba, 'np': np}, self.typingctx,
                    (self.typemap[lhs.name], self.typemap[in_arr.name],
                    self.typemap[index_var.name]),
                    self.typemap, self.calltypes).blocks
            first_block = min(f_blocks.keys())
            replace_arg_nodes(f_blocks[first_block], [lhs, in_arr, index_var])
            alloc_nodes = gen_np_call('empty_like', np.empty_like, lhs, [in_arr],
                        self.typingctx, self.typemap, self.calltypes)
            f_blocks[first_block].body = alloc_nodes + f_blocks[first_block].body
            return f_blocks