DECROMAX

Add UI on top of a dataframe to let viewers filter columns

Aug 11th, 2023
1,832
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 2.71 KB | None | 0 0
  1. def filter_dataframe(df: pd.DataFrame) -> pd.DataFrame:
  2.     """
  3.    Adds a UI on top of a dataframe to let viewers filter columns
  4.  
  5.    Args:
  6.        df (pd.DataFrame): Original dataframe
  7.  
  8.    Returns:
  9.        pd.DataFrame: Filtered dataframe
  10.    """
  11.     modify = st.checkbox("Add filters")
  12.  
  13.     if not modify:
  14.         return df
  15.  
  16.     df = df.copy()
  17.  
  18.     # Try to convert datetimes into a standard format (datetime, no timezone)
  19.     for col in df.columns:
  20.         if is_object_dtype(df[col]):
  21.             try:
  22.                 df[col] = pd.to_datetime(df[col])
  23.             except Exception:
  24.                 pass
  25.  
  26.         if is_datetime64_any_dtype(df[col]):
  27.             df[col] = df[col].dt.tz_localize(None)
  28.  
  29.     modification_container = st.container()
  30.  
  31.     with modification_container:
  32.         to_filter_columns = st.multiselect("Filter dataframe on", df.columns)
  33.         for column in to_filter_columns:
  34.             left, right = st.columns((1, 20))
  35.             # Treat columns with < 10 unique values as categorical
  36.             if is_categorical_dtype(df[column]) or df[column].nunique() < 10:
  37.                 user_cat_input = right.multiselect(
  38.                     f"Values for {column}",
  39.                     df[column].unique(),
  40.                     default=list(df[column].unique()),
  41.                 )
  42.                 df = df[df[column].isin(user_cat_input)]
  43.             elif is_numeric_dtype(df[column]):
  44.                 _min = float(df[column].min())
  45.                 _max = float(df[column].max())
  46.                 step = (_max - _min) / 100
  47.                 user_num_input = right.slider(
  48.                     f"Values for {column}",
  49.                     min_value=_min,
  50.                     max_value=_max,
  51.                     value=(_min, _max),
  52.                     step=step,
  53.                 )
  54.                 df = df[df[column].between(*user_num_input)]
  55.             elif is_datetime64_any_dtype(df[column]):
  56.                 user_date_input = right.date_input(
  57.                     f"Values for {column}",
  58.                     value=(
  59.                         df[column].min(),
  60.                         df[column].max(),
  61.                     ),
  62.                 )
  63.                 if len(user_date_input) == 2:
  64.                     user_date_input = tuple(map(pd.to_datetime, user_date_input))
  65.                     start_date, end_date = user_date_input
  66.                     df = df.loc[df[column].between(start_date, end_date)]
  67.             else:
  68.                 user_text_input = right.text_input(
  69.                     f"Substring or regex in {column}",
  70.                 )
  71.                 if user_text_input:
  72.                     df = df[df[column].astype(str).str.contains(user_text_input)]
  73.  
  74.     return df
Tags: streamlit
Advertisement
Add Comment
Please, Sign In to add comment