fix: bug fixed (#1601)

ydataai · Jun 17, 2024 · 9946037 · 9946037
1 parent 6d9e644
commit 9946037
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/src/ydata_profiling/visualisation/plot.py b/src/ydata_profiling/visualisation/plot.py
@@ -833,12 +833,12 @@ def _prepare_heatmap_data(
     )
 
     df = df.groupby([entity_column, "__bins"])[sortbykey].count()
-    df = df.reset_index().pivot_table(entity_column, "__bins", sortbykey).T
+    df = df.reset_index().pivot_table(values=sortbykey, index="__bins", columns=entity_column).T
 
     if selected_entities:
-        df = df[selected_entities].T
+        df = df[selected_entities]
     else:
-        df = df.T[:max_entities]
+        df = df[:max_entities]
 
     return df
 

diff --git a/tests/unit/test_plot.py b/tests/unit/test_plot.py
@@ -92,3 +92,9 @@ def test_timeseries_heatmap(dataframe: pd.DataFrame):
     df = _prepare_heatmap_data(dataframe, "entity")
     plot = _create_timeseries_heatmap(df)
     assert isinstance(plot, Axes)
+
+    dataframe["entity"] = dataframe["entity"].astype(str)
+    df = _prepare_heatmap_data(dataframe, "entity")
+    plot = _create_timeseries_heatmap(df)
+    assert isinstance(plot, Axes)
+