SIENTIAPDE-1222

Refactor datetime index handling in MLFlow and MLFlowRepository

- Moved the detect_and_parse_datetime_index method from MLFlow to MLFlowRepository for better organization and reusability.
- Updated the method to include enhanced logging and error handling for invalid datetime formats.
- Adjusted the transform method in MLFlowRepository to utilize the new datetime index parsing logic.
- Added unit tests for both valid and invalid datetime index cases to ensure robustness.
This commit is contained in:
vitor-aignosi
2025-09-17 08:59:12 -03:00
parent ee2ac5a365
commit 20fc938cc0
3 changed files with 107 additions and 76 deletions

View File

@@ -63,44 +63,6 @@ class MLFlow(BaseActivity):
f"{mlflow_host}:{mlflow_port}", mlflow_username, mlflow_password, logger.base_logger
)
def detect_and_parse_datetime_index(self, data: DataFrame, metadata: dict) -> DataFrame:
"""
Detect and parse datetime index from data. index must be a timestamp like column.
This function must detect the timestamp type (pandas Timestamp or datetime) and convert it to DATETIME_FORMAT_WITH_TZ.
If the index is a string, must be in format DATETIME_FORMAT_WITH_TZ.
If another type or format, must raise an error.
"""
index = data.index
# Get type of first element of index
index_type = type(index[0])
self.info(f"Index type: {index_type}", metadata)
message = f"Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}"
# Check if all in index are of the same type
if not all(isinstance(i, index_type) for i in index):
raise ValueError(
f"{message}")
# Check type and converts to DATETIME_FORMAT_WITH_TZ
if index_type == str:
# Validate format of string and return error if not valid
try:
to_datetime(data.index)
except ValueError:
raise ValueError(
f"{message}")
elif index_type == datetime or index_type == Timestamp:
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ)
else:
raise ValueError(
f"{message}")
return data
@activity.defn(name="request_transform")
async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
@@ -164,35 +126,6 @@ class MLFlow(BaseActivity):
self.debug(
f"Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
if response_data['success']:
response_dataframe = DataFrame(response_data['content'])
if len(response_dataframe) == 0:
return response_data
try:
response_dataframe = self.detect_and_parse_datetime_index(
response_dataframe, metadata)
response_dataframe['timestamp'] = to_datetime(
response_dataframe.index, format=DATETIME_FORMAT_WITH_TZ)
response_dataframe['timestamp'] = response_dataframe['timestamp'].dt.strftime(
DATETIME_FORMAT)
except ValueError as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id='TRANSFORM_DATA_INDEX_ERROR',
message=f'Error parsing trasnformed data index: {e}',
block='transform',
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.error(trace, metadata=metadata)
raise e
response_dataframe.to_csv('response_data.csv')
response_data['content'] = response_dataframe.to_dict()
self.debug(
f"Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)