vitor-aignosi
d2b365a34d
Refactor data handling in various modules to ensure DataFrame consistency - Replaced direct DataFrame instantiation with `ensure_dataframe` utility in Gates, MLFlow, OPC, and ModelMetrics classes to standardize data handling. - Updated return types in several asynchronous methods to return DataFrames instead of dictionaries for improved usability. - Adjusted data export processes in workflows to convert DataFrames to dictionaries with `to_dict(orient='records')` for compatibility with downstream systems.
150 lines
5.3 KiB
Python
150 lines
5.3 KiB
Python
"""
|
|
Temporal Codec for DataFrame Serialization
|
|
|
|
This module provides a custom Temporal DataConverter that automatically
|
|
serializes pandas DataFrames to Parquet format and deserializes them back.
|
|
|
|
The codec only handles DataFrames, leaving all other types to the default
|
|
Temporal serialization mechanism.
|
|
"""
|
|
|
|
import io
|
|
from collections.abc import Sequence
|
|
from typing import Any, List, Optional, Type
|
|
|
|
from temporalio.api.common.v1 import Payload
|
|
from temporalio.converter import DataConverter, PayloadConverter
|
|
from pandas import DataFrame, read_parquet
|
|
|
|
|
|
class DataFramePayloadConverter(PayloadConverter):
|
|
"""
|
|
Custom PayloadConverter that serializes pandas DataFrames to Parquet format.
|
|
|
|
Only DataFrames are handled by this converter. All other types are passed
|
|
to the default Temporal serialization mechanism.
|
|
"""
|
|
|
|
def __init__(self, default_payload_converter: PayloadConverter):
|
|
"""
|
|
Initialize the DataFrame payload converter.
|
|
|
|
Args:
|
|
default_payload_converter: The default Temporal payload converter
|
|
to use for non-DataFrame types
|
|
"""
|
|
self._default = default_payload_converter
|
|
|
|
def to_payloads(self, values: Sequence[Any]) -> List[Payload]:
|
|
"""
|
|
Convert values to Temporal Payloads.
|
|
|
|
If a value is a pandas DataFrame, it is serialized to Parquet format.
|
|
Otherwise, the default converter is used.
|
|
|
|
Args:
|
|
values: The values to serialize
|
|
|
|
Returns:
|
|
List[Payload]: The serialized payloads
|
|
"""
|
|
payloads = []
|
|
for value in values:
|
|
# Check if value is a DataFrame
|
|
try:
|
|
if isinstance(value, DataFrame):
|
|
buffer = io.BytesIO()
|
|
value.to_parquet(buffer, engine='pyarrow', index=True)
|
|
payloads.append(Payload(
|
|
metadata={"encoding": b"parquet-dataframe"},
|
|
data=buffer.getvalue()
|
|
))
|
|
continue
|
|
except ImportError:
|
|
# pandas not available, fall through to default
|
|
pass
|
|
except Exception:
|
|
# Error serializing DataFrame, fall through to default
|
|
pass
|
|
|
|
# Use default converter for all other types
|
|
default_payloads = self._default.to_payloads([value])
|
|
payloads.extend(default_payloads)
|
|
|
|
return payloads
|
|
|
|
def from_payloads(
|
|
self,
|
|
payloads: Sequence[Payload],
|
|
type_hints: Optional[List[Type]] = None,
|
|
) -> List[Any]:
|
|
"""
|
|
Convert Temporal Payloads back to Python values.
|
|
|
|
If a payload metadata indicates it's a Parquet-serialized DataFrame,
|
|
it is deserialized. Otherwise, the default converter is used.
|
|
|
|
Args:
|
|
payloads: The payloads to deserialize
|
|
type_hints: Optional type hints for the expected return types
|
|
|
|
Returns:
|
|
List[Any]: The deserialized values
|
|
"""
|
|
values = []
|
|
for i, payload in enumerate(payloads):
|
|
# Check if this is a Parquet-serialized DataFrame
|
|
if payload.metadata.get("encoding") == b"parquet-dataframe":
|
|
try:
|
|
buffer = io.BytesIO(payload.data)
|
|
values.append(read_parquet(buffer))
|
|
continue
|
|
except ImportError:
|
|
# pandas not available, fall through to default
|
|
pass
|
|
except Exception:
|
|
# Error deserializing DataFrame, fall through to default
|
|
pass
|
|
|
|
# Use default converter for all other types
|
|
# type_hints must have same length as payloads if provided
|
|
type_hint = type_hints[i] if type_hints and i < len(type_hints) else None
|
|
default_values = self._default.from_payloads([payload], [type_hint] if type_hint is not None else None)
|
|
values.extend(default_values)
|
|
|
|
return values
|
|
|
|
|
|
def create_dataframe_data_converter() -> DataConverter:
|
|
"""
|
|
Create a DataConverter with DataFrame serialization support.
|
|
|
|
Returns:
|
|
DataConverter: A DataConverter that handles DataFrames automatically
|
|
"""
|
|
# Get default converter to use as fallback
|
|
# DataConverter.default is an attribute, not a method
|
|
default_converter = DataConverter.default
|
|
|
|
# Create a factory class that extends PayloadConverter
|
|
class DataFramePayloadConverterFactory(PayloadConverter):
|
|
def __init__(self):
|
|
super().__init__()
|
|
# Create instance of default payload converter to use as fallback
|
|
self._default_converter = default_converter.payload_converter_class()
|
|
|
|
def to_payloads(self, values: Sequence[Any]) -> List[Payload]:
|
|
return DataFramePayloadConverter(self._default_converter).to_payloads(values)
|
|
|
|
def from_payloads(
|
|
self,
|
|
payloads: Sequence[Payload],
|
|
type_hints: Optional[List[Type]] = None,
|
|
) -> List[Any]:
|
|
return DataFramePayloadConverter(self._default_converter).from_payloads(payloads, type_hints)
|
|
|
|
return DataConverter(
|
|
payload_converter_class=DataFramePayloadConverterFactory
|
|
)
|
|
|