#41 refactoring

Signed-off-by: Christian Munz <[email protected]>
amosproj · Nov 24, 2024 · 9508683 · 9508683
1 parent b477c1a
commit 9508683
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/sdk/python/rtdip_sdk/pipelines/data_wranglers/spark/data_quality/prediction/arima.py b/src/sdk/python/rtdip_sdk/pipelines/data_wranglers/spark/data_quality/prediction/arima.py
@@ -197,6 +197,8 @@ def filter(self) -> PySparkDataFrame:
 
             base_df = df.toPandas()
             base_df["EventTime"] = pd.to_datetime(base_df["EventTime"])
+            base_df["Value"] = base_df["Value"].astype("float64")
+            base_df["EventTime"] = base_df["EventTime"].astype("datetime64[ns]")
             last_event_time = base_df["EventTime"].iloc[-1]
             second_last_event_time = base_df["EventTime"].iloc[-2]
 
@@ -224,6 +226,10 @@ def filter(self) -> PySparkDataFrame:
                     "Value": prediction_series.values,
                 }
             )
+            predicted_df["EventTime"] = predicted_df["EventTime"].astype(
+                "datetime64[ns]"
+            )
+
             extended_df = pd.concat([base_df, predicted_df], ignore_index=True)
 
             # Workaround needed for PySpark versions <3.4