SIENTIAPDE-1273

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.

laborious/activities/gates.py

@@ -15,6 +15,7 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
 
     from laborious import metrics
+    from laborious.utils.dataframe_utils import ensure_dataframe
     from laborious.utils.filters.conditional_filters import (
         filter_empty_data,
         filter_specific_variables_null_values,
@@ -149,7 +150,7 @@ class Gates(SientiaMonitoring):
         self.info('Performing input gate...', metadata)
 
         filters = input_data['filters']
-        data = DataFrame(input_data['data'])
+        data = ensure_dataframe(input_data['data'])
         path_priority = input_data['path_priority']
 
         filter_output = []
@@ -404,7 +405,7 @@ class Gates(SientiaMonitoring):
 
 
     @activity.defn(name='format_transformed_data')
-    async def format_transformed_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    async def format_transformed_data(self, input_data: dict[str, Any]) -> DataFrame:
         """
         Format transformed data according to configured storage policies.
         """
@@ -414,7 +415,7 @@ class Gates(SientiaMonitoring):
 
         self.info('Formatting transformed data...', metadata)
 
-        data = DataFrame(input_data['data'])
+        data = ensure_dataframe(input_data['data'])
 
         data['timestamp'] = data.index
         data = data.reset_index(drop=True)
@@ -422,10 +423,10 @@ class Gates(SientiaMonitoring):
         data = data.melt(id_vars='timestamp', var_name='variable', value_name='value')
         data['model_id'] = model_id
 
-        return data.to_dict()
+        return data
 
     @activity.defn(name='format_prediction')
-    async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    async def format_prediction(self, input_data: dict[str, Any]) -> DataFrame:
         """
         Format prediction data according to configured storage policies.
 
@@ -453,7 +454,7 @@ class Gates(SientiaMonitoring):
         prediction_store_policy = input_data['prediction_store_policy']
         self.info('Formatting prediction...', metadata)
 
-        data = DataFrame(input_data['data'])
+        data = ensure_dataframe(input_data['data'])
 
         # Create timestamp column from index and reset index
         data['timestamp'] = data.index
@@ -495,10 +496,10 @@ class Gates(SientiaMonitoring):
         self.info(f'Prediction formatted: {len(data)} rows', metadata)
         self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
 
-        return data.to_dict()
+        return data
 
     @activity.defn(name='format_default_prediction')
-    async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    async def format_default_prediction(self, input_data: dict[str, Any]) -> DataFrame:
         """
         Create and format default prediction data for error conditions.
 
@@ -540,10 +541,10 @@ class Gates(SientiaMonitoring):
         )
 
         self.info(f'Default prediction formatted: {data.size} rows', metadata)
-        return data.to_dict()
+        return data
 
     @activity.defn(name='format_retrain_report')
-    async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[Any, Any]:
+    async def format_retrain_report(self, input_data: dict[str, Any]) -> DataFrame:
         """
         Format retrain report data according to configured storage policies.
         """
@@ -572,7 +573,7 @@ class Gates(SientiaMonitoring):
 
         self.debug(f'Retrain report: {report.to_csv()}', metadata)
 
-        return report.to_dict()
+        return report
 
     @activity.defn(name='get_last_timestamp')
     async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
@@ -601,7 +602,7 @@ class Gates(SientiaMonitoring):
 
         self.info('Getting last timestamp...', metadata)
 
-        data = DataFrame(input_data['data'])
+        data = ensure_dataframe(input_data['data'])
 
         self.debug(f'Input data: {data.head(5).to_string()}', metadata)
 

laborious/activities/mlflow.py

@@ -20,6 +20,7 @@ with workflow.unsafe.imports_passed_through():
         now,
     )
 
+    from laborious.utils.dataframe_utils import ensure_dataframe
     from laborious.utils.repository.minio_repository import MinioRepository
     from laborious.utils.repository.model_repository import MLFlowRepository
 
@@ -139,7 +140,7 @@ class MLFlow(SientiaMonitoring):
         """
         metadata = input_data['metadata']
         self.info('Transforming data...', metadata)
-        data = DataFrame(input_data['data'])
+        data = ensure_dataframe(input_data['data'])
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
@@ -211,7 +212,7 @@ class MLFlow(SientiaMonitoring):
         """
         metadata = input_data['metadata']
         self.info('Predicting data...', metadata)
-        data = DataFrame(input_data['data'])
+        data = ensure_dataframe(input_data['data'])
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
@@ -422,7 +423,7 @@ class MLFlow(SientiaMonitoring):
 
 
     @activity.defn(name='get_reference_data')
-    async def get_reference_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any] | None:
+    async def get_reference_data(self, input_data: dict[str, Any]) -> DataFrame | None:
         """
         Get reference data from the MLflow Model Registry.
 
@@ -451,4 +452,4 @@ class MLFlow(SientiaMonitoring):
         reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
         reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
 
-        return reference_data.to_dict()
+        return reference_data

laborious/activities/model_metrics.py

@@ -138,7 +138,7 @@ class ModelMetrics(SientiaMonitoring):
 
 
     @activity.defn(name='calculate_drift')
-    async def calculate_drift(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
+    async def calculate_drift(self, input_data: dict[str, Any]) -> DataFrame | dict:
         """
         Calculate drift metrics for a model.
 
@@ -262,9 +262,9 @@ class ModelMetrics(SientiaMonitoring):
         self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata)
 
         
-        return drift_df.to_dict()
+        return drift_df
 
-    async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
+    async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> DataFrame:
         """
         Calculate simple metrics for a model. Metrics available are:
             - rmse
@@ -342,6 +342,6 @@ class ModelMetrics(SientiaMonitoring):
 
         self.debug(f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata)
 
-        return data.to_dict()
+        return data
 
 

laborious/activities/opc.py

@@ -1,3 +1,4 @@
+from typing import Hashable
 from temporalio import activity, workflow
 
 with workflow.unsafe.imports_passed_through():
@@ -11,6 +12,7 @@ with workflow.unsafe.imports_passed_through():
     from sientia_do.observability.metrics_controller import MetricsController
     from sientia_do.observability.sientia_monitoring import SientiaMonitoring
 
+    from laborious.utils.dataframe_utils import ensure_dataframe
     from laborious.utils.repository.opc_repository import OpcRepository
 
 OPC_WRITTING_ERROR_CONFIDENCE = 12
@@ -275,7 +277,7 @@ class OPC(SientiaMonitoring):
     @activity.defn(name='write_opc_data')
     async def write_opc_data(
         self, input_data: dict[str, Any]
-    ) -> tuple[dict[Any, Any], dict[str, dict[str, float | None]]]:
+    ) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]:
         """
         Write prediction and confidence data to OPC servers. The two writing
         operations are optional and independent of each other.
@@ -295,7 +297,7 @@ class OPC(SientiaMonitoring):
         """
         metadata = input_data['metadata']
         self.info('Writing data to OPC servers...', metadata)
-        data = DataFrame(input_data['data'])
+        data = ensure_dataframe(input_data['data'])
         opc_output_config = input_data['opc_output_config']
         self.info(f'Data to write: {data.size} rows', metadata)
 
@@ -324,7 +326,7 @@ class OPC(SientiaMonitoring):
 
     def process_confidence(
         self, data: DataFrame, success: bool, metadata: dict[str, Any]
-    ) -> dict[Any, Any]:
+    ) -> dict[Hashable, Any]:
         """
         Process prediction confidence based on OPC write operation success.
 

laborious/utils/dataframe_utils.py

@@ -0,0 +1,31 @@
+"""
+DataFrame utility functions for handling serialized DataFrames.
+
+This module provides helper functions to work with DataFrames that may
+come from Temporal serialization (already as DataFrame) or from legacy
+code (as dict).
+"""
+
+from typing import Any
+
+from pandas import DataFrame
+
+
+def ensure_dataframe(data: Any) -> DataFrame:
+    """
+    Ensure that data is a DataFrame, converting from dict if necessary.
+
+    This function handles both cases:
+    - Data already deserialized as DataFrame (from Temporal codec)
+    - Data as dict (legacy format or non-DataFrame serialization)
+
+    Args:
+        data: Data that should be a DataFrame (can be DataFrame or dict)
+
+    Returns:
+        DataFrame: The data as a pandas DataFrame
+    """
+    if isinstance(data, DataFrame):
+        return data
+    return DataFrame(data)
+

laborious/utils/temporal_codec.py

@@ -0,0 +1,149 @@
+"""
+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
+    )
+

laborious/worker/worker.py

@@ -48,6 +48,7 @@ with workflow.unsafe.imports_passed_through():
         build_opc_config,
         build_postgres_config,
     )
+    from laborious.utils.temporal_codec import create_dataframe_data_converter
     from laborious.workflows.minimal_retrain import MinimalRetrain
     from laborious.workflows.predictions_batch import PredictionsBatch
     from laborious.workflows.drift import Drift
@@ -132,10 +133,14 @@ async def main():
 
     logger.custom_info(f'Starting Temporal Client at {host}...', metadata)
 
+    # Create custom data converter with DataFrame support
+    data_converter = create_dataframe_data_converter()
+
     temporal_client = await client.Client.connect(
         target_host=host,
         namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
         runtime=new_runtime,
+        data_converter=data_converter,
     )
 
     logger.custom_info('Starting Workers...', metadata)

laborious/workflows/drift.py

@@ -93,7 +93,7 @@ class Drift:
                 Activities.export_data_to_postgres,
                 {
                     **metadata,
-                    'data': drift_data,
+                    'data': drift_data.to_dict(orient='records'),
                     'schema': input_data['schema'],
                     'table_name': input_data['target_table_name'],
                     'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},

laborious/workflows/minimal_retrain.py

@@ -127,7 +127,7 @@ class MinimalRetrain:
             Activities.export_data_to_postgres,
             {
                 **metadata,
-                'data': report,
+                'data': report.to_dict(orient='records'),
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
             },

laborious/workflows/simple_metrics.py

@@ -82,7 +82,7 @@ class SimpleMetrics:
             Activities.export_data_to_postgres,
             {
                 **metadata,
-                'data': simple_metrics,
+                'data': simple_metrics.to_dict(orient='records'),
                 'schema': input_data['schema'],
                 'table_name': input_data['target_table_name'],
                 'timestamp_conversion': {

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -107,7 +107,7 @@ class FormatAndExportPrediction:
                         **metadata,
                         'schema': input_data['schema'],
                         'table_name': input_data['transform_table_name'],
-                        'data': transformed,
+                        'data': transformed.to_dict(orient='records'),
                         'timestamp_conversion': {
                             'column': 'timestamp',
                             'format': DATETIME_FORMAT_WITH_TZ,
@@ -156,7 +156,7 @@ class FormatAndExportPrediction:
                 **metadata,
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
-                'data': prediction,
+                'data': prediction.to_dict(orient='records'),
                 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
             },
             retry_policy=retry_policy,

tests.ipynb

@@ -588,6 +588,157 @@
     "\n",
     "display(prediction_data)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "486b95b3",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>a</th>\n",
+       "      <th>b</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>2025-01-01</th>\n",
+       "      <td>1</td>\n",
+       "      <td>4</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2025-01-02</th>\n",
+       "      <td>2</td>\n",
+       "      <td>5</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2025-01-03</th>\n",
+       "      <td>3</td>\n",
+       "      <td>6</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "            a  b\n",
+       "2025-01-01  1  4\n",
+       "2025-01-02  2  5\n",
+       "2025-01-03  3  6"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'index': ['2025-01-01', '2025-01-02', '2025-01-03'],\n",
+       " 'columns': ['a', 'b'],\n",
+       " 'data': [[1, 4], [2, 5], [3, 6]]}"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>0</th>\n",
+       "      <th>1</th>\n",
+       "      <th>2</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>index</th>\n",
+       "      <td>2025-01-01</td>\n",
+       "      <td>2025-01-02</td>\n",
+       "      <td>2025-01-03</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>columns</th>\n",
+       "      <td>a</td>\n",
+       "      <td>b</td>\n",
+       "      <td>None</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>data</th>\n",
+       "      <td>[1, 4]</td>\n",
+       "      <td>[2, 5]</td>\n",
+       "      <td>[3, 6]</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "                  0           1           2\n",
+       "index    2025-01-01  2025-01-02  2025-01-03\n",
+       "columns           a           b        None\n",
+       "data         [1, 4]      [2, 5]      [3, 6]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from pandas import DataFrame\n",
+    "\n",
+    "data = DataFrame({\n",
+    "    \"a\": {\"2025-01-01\": 1, \"2025-01-02\": 2, \"2025-01-03\": 3},\n",
+    "    \"b\": {\"2025-01-01\": 4, \"2025-01-02\": 5, \"2025-01-03\": 6},\n",
+    "})\n",
+    "\n",
+    "display(data)\n",
+    "\n",
+    "data_list = data.to_dict('split')\n",
+    "\n",
+    "display(data_list)\n",
+    "\n",
+    "data_rec = DataFrame.from_dict(data_list, orient='index')\n",
+    "\n",
+    "display(data_rec)"
+   ]
   }
  ],
  "metadata": {