Search code examples
pythonstreamlit

Remove rows in your dataframe with filling option in streamlit

I'm currently creating a streamlit app that shows the results of user's inputs. It may happen that the user accidentally submitted its results. So I want the user to allow removing its values in the dataframe. The most similar option I found was this Deleting rows in st.data_editor progmatically, which uses a callback function to remove the rows. Unfortunately, it only works on the first row of the dataframe but it removes the complete dataset. Also, it doesn't work on the latter rows. Here is some reproducible example:

import pandas as pd
import streamlit as st

st.sidebar.header("Submit results")

# default values
st.session_state["option"] = ""
st.session_state["number"] = 1

if 'data' not in st.session_state:
    data = pd.DataFrame(columns=["Track", "Result"])
    st.session_state.data = data

data = st.session_state.data

def onAddRow():
    row = pd.DataFrame({'Track':[st.session_state["option"]], 'Result':[st.session_state["number"]]})
    st.session_state.data = pd.concat([st.session_state.data, row])

def callback():
    edited_rows = st.session_state["data_editor"]["edited_rows"]
    rows_to_delete = []

    for idx, value in edited_rows.items():
        if value["x"] is True:
            rows_to_delete.append(idx)

    st.session_state["data"] = (
        st.session_state["data"].drop(rows_to_delete, axis=0).reset_index(drop=True)
    )

with st.form('Form1'):
    st.session_state["option"] = st.selectbox(
            "Select the tracked you played:",
            ("ds Mario Kart", "Toad Harbour", "Koopa Cape"))
    st.session_state["number"] = st.slider("Pick a number", 1, 12)
    st.session_state["submitted1"] = st.form_submit_button('Submit 1', on_click=onAddRow)


columns = st.session_state["data"].columns
column_config = {column: st.column_config.Column(disabled=True) for column in columns}

modified_df = st.session_state["data"].copy()
modified_df["x"] = False
# Make Delete be the first column
modified_df = modified_df[["x"] + modified_df.columns[:-1].tolist()]

st.data_editor(
    modified_df,
    key="data_editor",
    on_change=callback,
    hide_index=True,
    column_config=column_config,
)

Output:

enter image description here

So it looks like this. Imagine we want to remove the second row, we can click on the button in the first column. But this returns an error:

KeyError: '[1] not found in axis'

I don't understand why this happens. So I was wondering if anyone knows how to create an option to remove rows like above even when you have a filling option with a submit button?

Solution

  • Reset the index of your session state data in the callback. This will place an index from 0 to n. The drop method would then work not only for index 0 but also for other indices.

    st.session_state["data"] = st.session_state["data"].reset_index(drop=True)

    Full code.

    def callback():
    edited_rows = st.session_state["data_editor"]["edited_rows"]
    rows_to_delete = []

    for idx, value in edited_rows.items():
        if value["x"] is True:
            rows_to_delete.append(idx)

    st.session_state["data"] = st.session_state["data"].reset_index(drop=True)  # <--- this

    st.session_state["data"] = (
        st.session_state["data"].drop(rows_to_delete, axis=0).reset_index(drop=True)
    )