diff --git a/burr/integrations/serde/pandas.py b/burr/integrations/serde/pandas.py index 22ee32b01..c917cb7f4 100644 --- a/burr/integrations/serde/pandas.py +++ b/burr/integrations/serde/pandas.py @@ -52,6 +52,12 @@ def serialize_pandas_df( hash_object = hashlib.sha256() hash_value = str(value.columns) + str(value.shape) + str(value.dtypes) hash_object.update(hash_value.encode()) + try: + content = pd.util.hash_pandas_object(value, index=True).values + except TypeError: + # cells such as lists cannot be hashed by pandas, so hash the parquet bytes instead + content = value.to_parquet() + hash_object.update(content) # Return the hexadecimal representation of the hash file_name = f"df_{hash_object.hexdigest()}.parquet" diff --git a/tests/integrations/serde/test_pandas.py b/tests/integrations/serde/test_pandas.py index 90ce9950b..55a2496bb 100644 --- a/tests/integrations/serde/test_pandas.py +++ b/tests/integrations/serde/test_pandas.py @@ -68,3 +68,20 @@ def test_deserialize_pandas_df_without_pandas_kwargs(tmp_path): serialized = og.serialize(pandas_kwargs={"path": tmp_path}) ng = state.State.deserialize(serialized) pd.testing.assert_frame_equal(ng["df"], df) + + +@pytest.mark.parametrize( + "df_1,df_2", + [ + (pd.DataFrame({"a": [1, 2]}), pd.DataFrame({"a": [9, 9]})), + (pd.DataFrame({"a": [[1, 2], [3]]}), pd.DataFrame({"a": [[9], [9, 9]]})), + ], + ids=["scalar_cells", "list_cells"], +) +def test_serde_of_same_shape_dataframes_keeps_each_dataframe(tmp_path, df_1, df_2): + serialized_1 = state.State({"df": df_1}).serialize(pandas_kwargs={"path": tmp_path}) + serialized_2 = state.State({"df": df_2}).serialize(pandas_kwargs={"path": tmp_path}) + assert serialized_1["df"]["path"] != serialized_2["df"]["path"] + + pd.testing.assert_frame_equal(state.State.deserialize(serialized_1)["df"], df_1) + pd.testing.assert_frame_equal(state.State.deserialize(serialized_2)["df"], df_2)