@@ -60,7 +60,7 @@ def run(self, data):
60
60
61
61
"""
62
62
clean_df = self ._remove_trailing_whitespace (data )
63
- if self .dataset_info . horizon :
63
+ if hasattr ( self .dataset_info , ' horizon' ) :
64
64
clean_df = self ._clean_column_names (clean_df )
65
65
if self .name == "historical_data" :
66
66
self ._check_historical_dataset (clean_df )
@@ -113,15 +113,18 @@ def _clean_column_names(self, df):
113
113
col : col .replace (" " , "" ) for col in df .columns if " " in col
114
114
}
115
115
116
- self .target_column_name = self .raw_column_names .get (
117
- self .target_column_name , self .target_column_name
118
- )
116
+ if self .target_column_name :
117
+ self .target_column_name = self .raw_column_names .get (
118
+ self .target_column_name , self .target_column_name
119
+ )
119
120
self .dt_column_name = self .raw_column_names .get (
120
121
self .dt_column_name , self .dt_column_name
121
122
)
122
- self .target_category_columns = [
123
- self .raw_column_names .get (col , col ) for col in self .target_category_columns
124
- ]
123
+
124
+ if self .target_category_columns :
125
+ self .target_category_columns = [
126
+ self .raw_column_names .get (col , col ) for col in self .target_category_columns
127
+ ]
125
128
df .columns = df .columns .str .replace (" " , "" )
126
129
return df
127
130
0 commit comments