Skip to content
This repository was archived by the owner on Feb 2, 2024. It is now read-only.

Commit 645b284

Browse files
authored
Re-implement df.drop based on new structure (#848)
1 parent f5dcbc1 commit 645b284

2 files changed

Lines changed: 12 additions & 19 deletions

File tree

‎sdc/datatypes/hpat_pandas_dataframe_functions.py‎

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1308,22 +1308,14 @@ def isna_overload(df):
13081308

13091309
def sdc_pandas_dataframe_drop_codegen(func_name, func_args, df, drop_cols):
13101310
"""
1311-
Input:
1312-
df.drop(columns='M', errors='ignore')
1313-
1314-
Func generated:
1315-
def sdc_pandas_dataframe_drop_impl(df, labels=None, axis=0, index=None, columns=None, level=None, inplace=False,
1316-
errors="raise"):
1317-
if errors == "raise":
1318-
raise ValueError("The label M is not found in the selected axis")
1319-
new_col_A_data_df = df._data[0]
1320-
new_col_B_data_df = df._data[1]
1321-
new_col_C_data_df = df._data[2]
1322-
return pandas.DataFrame({"A": new_col_A_data_df, "B": new_col_B_data_df, "C": new_col_C_data_df})
1323-
1311+
Example of generated implementation:
1312+
def sdc_pandas_dataframe_drop_impl(df, labels=None, axis=0, index=None, columns=None,
1313+
level=None, inplace=False, errors="raise"):
1314+
new_col_0_data_df = df._data[1][0]
1315+
new_col_1_data_df = df._data[0][1]
1316+
return pandas.DataFrame({"B": new_col_0_data_df, "C": new_col_1_data_df}, index=df.index)
13241317
"""
13251318
indent = 4 * ' '
1326-
df_columns_indx = {col_name: i for i, col_name in enumerate(df.columns)}
13271319
saved_df_columns = [column for column in df.columns if column not in drop_cols]
13281320
func_definition = [f'def sdc_pandas_dataframe_{func_name}_impl({", ".join(func_args)}):']
13291321
func_text = []
@@ -1336,8 +1328,9 @@ def sdc_pandas_dataframe_drop_impl(df, labels=None, axis=0, index=None, columns=
13361328
break
13371329

13381330
for column_id, column_name in enumerate(saved_df_columns):
1339-
func_text.append(f'new_col_{column_id}_data_{"df"} = get_dataframe_data({"df"}, '
1340-
f'{df_columns_indx[column_name]})')
1331+
col_loc = df.column_loc[column_name]
1332+
type_id, col_id = col_loc.type_id, col_loc.col_id
1333+
func_text.append(f'new_col_{column_id}_data_df = df._data[{type_id}][{col_id}]')
13411334
column_list.append((f'new_col_{column_id}_data_df', column_name))
13421335

13431336
data = ', '.join(f'"{column_name}": {column}' for column, column_name in column_list)
@@ -1346,7 +1339,7 @@ def sdc_pandas_dataframe_drop_impl(df, labels=None, axis=0, index=None, columns=
13461339
func_definition.extend([indent + func_line for func_line in func_text])
13471340
func_def = '\n'.join(func_definition)
13481341

1349-
global_vars = {'pandas': pandas, 'get_dataframe_data': sdc.hiframes.pd_dataframe_ext.get_dataframe_data}
1342+
global_vars = {'pandas': pandas}
13501343

13511344
return func_def, global_vars
13521345

‎sdc/tests/test_dataframe.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1616,7 +1616,7 @@ def test_impl(df):
16161616
index=index)
16171617
pd.testing.assert_frame_equal(sdc_func(df), test_impl(df))
16181618

1619-
@dfRefactoringNotImplemented
1619+
@dfRefactoringNotImplemented # required re-implementing DataFrame boxing
16201620
def test_df_drop_one_column(self):
16211621
def test_impl(index):
16221622
df = pd.DataFrame({'A': [1.0, 2.0, np.nan, 1.0], 'B': [4, 5, 6, 7], 'C': [1.0, 2.0, np.nan, 1.0]},
@@ -1658,7 +1658,7 @@ def test_impl(df):
16581658
index=index)
16591659
pd.testing.assert_frame_equal(sdc_func(df), test_impl(df))
16601660

1661-
@dfRefactoringNotImplemented
1661+
@dfRefactoringNotImplemented # required re-implementing DataFrame boxing
16621662
def test_df_drop_tuple_column(self):
16631663
def gen_test_impl(do_jit=False):
16641664
def test_impl(index):

0 commit comments

Comments
 (0)