Fix vis for spark df with timestamp (#865)

closes #861
evidentlyai · Nov 30, 2023 · 3c38a9d · 3c38a9d
1 parent 728d74c
commit 3c38a9d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/evidently/spark/visualizations.py b/src/evidently/spark/visualizations.py
@@ -36,16 +36,16 @@ def prepare_df_for_time_index_plot(
             period_col = sf.date_format(date_col, pattern).alias(PERIOD_COL)
         plot_df = (
             df.select(column_name, period_col)
-            .groupby(period_col)
+            .groupby(PERIOD_COL)
             .agg(sf.mean(column_name).alias("mean"), sf.stddev_pop(column_name).alias("std"))
         )
         if pattern == "week":
             split = sf.split(PERIOD_COL, "-")
             week = split.getItem(1)
             year = sf.to_timestamp(split.getItem(0), "y")
             week_start_diff = sf.date_format(year, "F")
-            plot_df = plot_df.select("*", sf.date_add(year, week * 7 - week_start_diff).alias(PERIOD_COL)).toPandas()
-        return plot_df, prefix
+            plot_df = plot_df.select("*", sf.date_add(year, week * 7 - week_start_diff).alias(PERIOD_COL))
+        return plot_df.toPandas(), prefix
     ptp = df.count() - 1
 
     schema = StructType(fields=[StructField("_1", dataType=df.schema), StructField("_2", dataType=LongType())])