Export and clean notebook

nixtlats/nixtla_client.py

@@ -749,12 +749,15 @@ def partition_by_uid(func):
     def wrapper(self, num_partitions, **kwargs):
         if num_partitions is None or num_partitions == 1:
             return func(self, **kwargs, num_partitions=1)
+        main_logger.info(f"Number of partitions: {num_partitions}")
         df = kwargs.pop("df")
         X_df = kwargs.pop("X_df", None)
         id_col = kwargs["id_col"]
         uids = df["unique_id"].unique()
         results_df = []
+        split_index = 1
         for uids_split in np.array_split(uids, num_partitions):
+            main_logger.info(f"Partition {split_index} of {num_partitions}")
             df_uids = df.query("unique_id in @uids_split")
             if X_df is not None:
                 X_df_uids = X_df.query("unique_id in @uids_split")
@@ -767,6 +770,7 @@ def wrapper(self, num_partitions, **kwargs):
                 kwargs_uids["X_df"] = X_df_uids
             results_uids = func(self, **kwargs_uids, num_partitions=1)
             results_df.append(results_uids)
+            split_index += 1
         results_df = pd.concat(results_df).reset_index(drop=True)
         return results_df