Skip to content

Commit d1d5ef5

Browse files
author
Daria Tikhonovich
committed
fixed extra cols
1 parent 2e594b3 commit d1d5ef5

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

rectools/dataset/interactions.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -195,13 +195,17 @@ def to_external(
195195
Columns.Item: item_id_map.convert_to_external(self.df[Columns.Item].values),
196196
}
197197
)
198+
cols_to_add = []
198199

199200
if include_weight:
200-
res[Columns.Weight] = self.df[Columns.Weight]
201+
cols_to_add.append(Columns.Weight)
201202
if include_datetime:
202-
res[Columns.Datetime] = self.df[Columns.Datetime]
203+
cols_to_add.append(Columns.Datetime)
203204
if include_extra_cols:
204-
target_cols = [col for col in self.df if col not in [Columns.Weight, Columns.Datetime]]
205-
self._add_extra_cols(res, self.df[target_cols])
205+
cols_not_to_add = [Columns.User, Columns.Item, Columns.Weight, Columns.Datetime]
206+
extra_cols = [col for col in self.df if col not in cols_not_to_add]
207+
cols_to_add.extend(extra_cols)
206208

209+
for col in cols_to_add:
210+
res[col] = self.df[col]
207211
return res

0 commit comments

Comments
 (0)