small fix: removing repartition as it has no effect before groupby

databricks-industry-solutions · Dec 1, 2023 · c5a6c4e · c5a6c4e
1 parent 2838a51
commit c5a6c4e
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/forecasting_sa/Forecaster.py b/forecasting_sa/Forecaster.py
@@ -467,8 +467,7 @@ def evaluate_local_model(self, model_conf):
         )
 
         res_sdf = (
-            src_df.repartition(100)
-            .groupby(self.conf["group_id"])
+            src_df.groupby(self.conf["group_id"])
             .applyInPandas(evaluate_one_model_fn, schema=output_schema_intm)
         )
 
@@ -477,8 +476,7 @@ def evaluate_local_model(self, model_conf):
         )
 
         res_sdf = (
-            res_sdf.repartition(100)
-            .groupby(self.conf["group_id"])
+            res_sdf.groupby(self.conf["group_id"])
             .applyInPandas(unpack_fn, schema=output_schema_final)
         )
 
@@ -610,8 +608,7 @@ def run_scoring_for_local_model(self, model_conf):
         score_one_model_fn = functools.partial(Forecaster.score_one_model, model=model)
 
         res_sdf = (
-            src_df.repartition(100)
-            .groupby(self.conf["group_id"])
+            src_df.groupby(self.conf["group_id"])
             .applyInPandas(score_one_model_fn, schema=output_schema)
         )
         if not isinstance(res_sdf.schema[self.conf["group_id"]].dataType, StringType):