diff --git a/graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_label_transformation.py b/graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_label_transformation.py index 4076dcdbf..a69c10561 100644 --- a/graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_label_transformation.py +++ b/graphstorm-processing/graphstorm_processing/data_transformations/dist_transformations/dist_label_transformation.py @@ -62,7 +62,7 @@ def apply(self, input_df: DataFrame) -> DataFrame: processed_col_name, self.label_column, ) - input_df = input_df.orderBy(COLUMN_ORDER_FLAG) + input_df = input_df.orderBy(COLUMN_ORDER_FLAG).drop(COLUMN_ORDER_FLAG) # Labels that were missing and were assigned the value numLabels by the StringIndexer # are converted to None