Data_editor behaves different when using button to load data

Using Streamlit
1

I’m running this locally. version 1.37.1

I have created this simple code that reproduce the problem.
Editor 1 works as it should, and always updates the values every time

When using editor 2, that loads data to df after a press of a button.
The data_editor always need to input data twice to get it to be stored.

First time edit a row = WORKS
second time edit a row = flashes and resets the values
third time edit a row = WORKS
etc…

I can’t figure out what i’m doing wrong. If i remove the button and just preload it everything works. but that is not what i need.

import streamlit as st
import pandas as pd

# Data for df1 (always loaded)
st.subheader("EDITOR 1")
st.session_state.df1 = pd.DataFrame({
    'A': [1, 2, 3],
    'B': [4, 5, 6]})

st.session_state.df1 = st.data_editor(st.session_state.df1, key="ED1")
st.write("DataFrame 1 res:", st.session_state.df1)

# Data for df2 (loaded only if button is pressed)
# Initialize session state for df2 if it doesn't exist
if 'df2' not in st.session_state:
    st.session_state.df2 = pd.DataFrame()  # Start with an empty DataFrame

st.subheader("EDITOR 2")
if st.button("Load Data for Editor 2"):
    st.session_state.df2 = pd.DataFrame({
        'X': ['a', 'b', 'c'],
        'Y': ['d', 'e', 'f']})

st.session_state.df2 = st.data_editor(st.session_state.df2, key="ED2")
st.write("DataFrame 2 res:", st.session_state.df2)
2

So this actually solves the problem. but seems hacky…

if 'df2' not in st.session_state:
    st.session_state.df2 = pd.DataFrame()  # Start with an empty DataFrame

st.subheader("EDITOR 2")


def update():
    for idx, change in st.session_state.ED2["edited_rows"].items():
        for label, value in change.items():
            st.session_state.df2.loc[idx, label] = value


if st.button("Load Data for Editor 2"):
    st.session_state.df2 = pd.DataFrame({
        'X': ['a', 'b', 'c'],
        'Y': ['d', 'e', 'f']})


st.session_state.df2 = st.data_editor(st.session_state.df2, key="ED2", on_change=update)
st.write("DataFrame 2 res:", st.session_state.df2)
st.write(st.session_state.ED2)