diff --git a/optimus/dataframe/rows.py b/optimus/dataframe/rows.py index 1ef340d6..edd3d620 100644 --- a/optimus/dataframe/rows.py +++ b/optimus/dataframe/rows.py @@ -124,6 +124,8 @@ def sort(col_sort): """ # If a list of columns names are given order this by desc. If you need to specify the order of every # column use a list of tuples (col_name, "asc") + df = self + t = [] if is_list_of_str_or_int(col_sort): for col_name in col_sort: @@ -131,7 +133,6 @@ def sort(col_sort): col_sort = t func = [] - for cs in col_sort: col_name = one_list_to_val(cs[0]) order = cs[1] @@ -140,8 +141,13 @@ def sort(col_sort): sort_func = F.asc elif order == "desc": sort_func = F.desc + else: + RaiseIt.value_error(sort_func, ["asc", "desc"]) + func.append(sort_func(col_name)) - df = self.sort(*func) + df = df.preserve_meta(self, Actions.SORT_ROW, col_name) + + df = df.sort(*func) return df @add_attr(rows) diff --git a/optimus/helpers/constants.py b/optimus/helpers/constants.py index 8a7c53e1..045ee04d 100644 --- a/optimus/helpers/constants.py +++ b/optimus/helpers/constants.py @@ -104,12 +104,13 @@ class Actions(Enum): Z_SCORE = "z_score" NEST = "nest" UNNEST = "unnest" - DROP_ROW = "drop_row" VALUES_TO_COLS = "values_to_cols" SET = "set" STRING_TO_INDEX = "string_to_index" MIN_MAX_SCALER = "min_max_scaler" MAX_ABS_SCALER = "max_abs_scaler" + DROP_ROW = "drop_row" + SORT_ROW = "sort_row" @staticmethod def list():