diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index df812dc..80d007d 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -531,23 +531,19 @@ class MLFlow(SientiaMonitoring): # Keep heavy model fitting outside MLflow run timing. prediction_data = wrapper.retrain(data).rename(columns={target: 'prediction'}) - # Merge prediction data with retrain data. Both sides are indexed by the pivoted - # timestamp, so one row per timestamp on each side; `validate` fails loudly instead - # of silently multiplying rows if the wrapper ever returns duplicated timestamps. - evaluation_data = pd.merge( - data, - prediction_data, - left_index=True, - right_index=True, - how='left', - validate='one_to_one', - ) - - # Rename target column to "target", then move the index into a "timestamp" column - evaluation_data = evaluation_data.rename(columns={target: 'target'}) - evaluation_data['timestamp'] = evaluation_data.index - evaluation_data = evaluation_data.reset_index(drop=True).sort_values( - by='timestamp', ascending=True + # Both frames are indexed by the pivoted timestamp: expose it as an explicit column + # and join on it. `validate` fails loudly instead of silently multiplying rows if the + # wrapper ever returns duplicated timestamps. + evaluation_data = ( + data.reset_index(names='timestamp') + .merge( + prediction_data.reset_index(names='timestamp'), + on='timestamp', + how='left', + validate='one_to_one', + ) + .rename(columns={target: 'target'}) + .sort_values(by='timestamp', ascending=True, ignore_index=True) ) run_name = f'{model_name}-retrain-{datetime.now().strftime("%Y%m%d%H%M%S")}'