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

Commit 85f6e15

Browse files
authored
Re-implement df.rolling methods based on new structure (#852)
1 parent 401313c commit 85f6e15

File tree

2 files changed

+327
-28
lines changed

2 files changed

+327
-28
lines changed

sdc/datatypes/hpat_pandas_dataframe_rolling_functions.py

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -123,19 +123,25 @@ def df_rolling_method_other_df_codegen(method_name, self, other, args=None, kws=
123123
f' raise ValueError("Method rolling.{method_name}(). The object pairwise\\n expected: False, None")'
124124
]
125125

126-
data_length = 'len(self._data._data[0])' if data_columns else '0'
127-
other_length = 'len(other._data[0])' if other_columns else '0'
126+
data_length = 'len(self._data._data[0][0])' if data_columns else '0'
127+
other_length = 'len(other._data[0][0])' if other_columns else '0'
128128
func_lines += [f' length = max([{data_length}, {other_length}])']
129129

130130
for col in all_columns:
131131
res_data = f'result_data_{col}'
132132
if col in common_columns:
133+
col_loc = self.data.column_loc[col]
134+
type_id, col_id = col_loc.type_id, col_loc.col_id
135+
other_col_loc = other.column_loc[col]
136+
other_type_id = other_col_loc.type_id
137+
other_col_id = other_col_loc.col_id
138+
133139
other_series = f'other_series_{col}'
134140
method_kws['other'] = other_series
135141
method_params = ', '.join(args + kwsparams2list(method_kws))
136142
func_lines += [
137-
f' data_{col} = self._data._data[{data_columns[col]}]',
138-
f' other_data_{col} = other._data[{other_columns[col]}]',
143+
f' data_{col} = self._data._data[{type_id}][{col_id}]',
144+
f' other_data_{col} = other._data[{other_type_id}][{other_col_id}]',
139145
f' series_{col} = pandas.Series(data_{col})',
140146
f' {other_series} = pandas.Series(other_data_{col})',
141147
f' rolling_{col} = series_{col}.rolling({rolling_params})',
@@ -158,16 +164,18 @@ def df_rolling_method_other_df_codegen(method_name, self, other, args=None, kws=
158164
return func_text, global_vars
159165

160166

161-
def df_rolling_method_main_codegen(method_params, df_columns, method_name):
167+
def df_rolling_method_main_codegen(method_params, df_columns, column_loc, method_name):
162168
rolling_params = df_rolling_params_codegen()
163169
method_params_as_str = ', '.join(method_params)
164170

165171
results = []
166172
func_lines = []
167173
for idx, col in enumerate(df_columns):
174+
col_loc = column_loc[col]
175+
type_id, col_id = col_loc.type_id, col_loc.col_id
168176
res_data = f'result_data_{col}'
169177
func_lines += [
170-
f' data_{col} = self._data._data[{idx}]',
178+
f' data_{col} = self._data._data[{type_id}][{col_id}]',
171179
f' series_{col} = pandas.Series(data_{col})',
172180
f' rolling_{col} = series_{col}.rolling({rolling_params})',
173181
f' result_{col} = rolling_{col}.{method_name}({method_params_as_str})',
@@ -229,7 +237,8 @@ def df_rolling_method_codegen(method_name, self, args=None, kws=None):
229237
func_lines = [f'def {impl_name}({impl_params_as_str}):']
230238

231239
method_params = args + ['{}={}'.format(k, k) for k in kwargs]
232-
func_lines += df_rolling_method_main_codegen(method_params, self.data.columns, method_name)
240+
func_lines += df_rolling_method_main_codegen(method_params, self.data.columns,
241+
self.data.column_loc, method_name)
233242
func_text = '\n'.join(func_lines)
234243

235244
global_vars = {'pandas': pandas}

0 commit comments

Comments
 (0)