Stacking options - st.bar_chart (streamlit#8945)

PR extends flexibility of `st.bar_chart` by adding a new parameter `stack` to allow explicit control over multiple stacking options.

e2e_playwright/st_bar_chart.py

@@ -15,6 +15,7 @@
 
 import numpy as np
 import pandas as pd
+from vega_datasets import data as vega_data
 
 import streamlit as st
 
@@ -71,3 +72,11 @@
 st.bar_chart(df, x_label="X Axis Label", y_label="Y Axis Label")
 st.bar_chart(df, horizontal=True)
 st.bar_chart(df, horizontal=True, x_label="X Label", y_label="Y Label")
+
+# Additional tests for stacking options
+source = vega_data.barley()
+st.bar_chart(source, x="variety", y="yield", color="site", stack=True)
+st.bar_chart(source, x="variety", y="yield", color="site", stack=False)
+st.bar_chart(source, x="variety", y="yield", color="site", stack="normalize")
+st.bar_chart(source, x="variety", y="yield", color="site", stack="center")
+st.bar_chart(source, x="variety", y="yield", color="site", stack="layered")

e2e_playwright/st_bar_chart_test.py

@@ -16,7 +16,7 @@
 
 from e2e_playwright.conftest import ImageCompareFunction
 
-TOTAL_BAR_CHARTS = 13
+TOTAL_BAR_CHARTS = 18
 
 
 def test_bar_chart_rendering(app: Page, assert_snapshot: ImageCompareFunction):

lib/streamlit/elements/lib/built_in_chart_utils.py

@@ -45,7 +45,7 @@
     import pandas as pd
 
     from streamlit.elements.arrow import Data
-    from streamlit.type_util import DataFrameCompatible
+    from streamlit.type_util import ChartStackType, DataFrameCompatible
 
 
 class PrepDataColumns(TypedDict):
@@ -123,6 +123,8 @@ def generate_chart(
     size_from_user: str | float | None = None,
     width: int | None = None,
     height: int | None = None,
+    # Bar charts only:
+    stack: bool | ChartStackType | None = None,
 ) -> tuple[alt.Chart, AddRowsMetadata]:
     """Function to use the chart's type, data columns and indices to figure out the chart's spec."""
     import altair as alt
@@ -166,21 +168,18 @@ def generate_chart(
 
     # At this point, x_column is only None if user did not provide one AND df is empty.
 
-    if chart_type == ChartType.HORIZONTAL_BAR:
-        # Handle horizontal bar chart - switches x and y data:
-        x_encoding = _get_x_encoding(
-            df, y_column, y_from_user, x_axis_label, chart_type
-        )
-        y_encoding = _get_y_encoding(
-            df, x_column, x_from_user, y_axis_label, chart_type
-        )
-    else:
-        x_encoding = _get_x_encoding(
-            df, x_column, x_from_user, x_axis_label, chart_type
-        )
-        y_encoding = _get_y_encoding(
-            df, y_column, y_from_user, y_axis_label, chart_type
-        )
+    # Get x and y encodings
+    x_encoding, y_encoding = _get_axis_encodings(
+        df,
+        chart_type,
+        x_column,
+        y_column,
+        x_from_user,
+        y_from_user,
+        x_axis_label,
+        y_axis_label,
+        stack,
+    )
 
     # Create a Chart with x and y encodings.
     chart = alt.Chart(
@@ -193,8 +192,21 @@ def generate_chart(
         y=y_encoding,
     )
 
+    # Offset encoding only works for Altair >= 5.0.0
+    is_altair_version_offset_compatible = not type_util.is_altair_version_less_than(
+        "5.0.0"
+    )
+    # Set up offset encoding (creates grouped/non-stacked bar charts, so only applicable when stack=False).
+    if (
+        is_altair_version_offset_compatible
+        and stack is False
+        and color_column is not None
+    ):
+        x_offset, y_offset = _get_offset_encoding(chart_type, color_column)
+        chart = chart.encode(xOffset=x_offset, yOffset=y_offset)
+
     # Set up opacity encoding.
-    opacity_enc = _get_opacity_encoding(chart_type, color_column)
+    opacity_enc = _get_opacity_encoding(chart_type, stack, color_column)
     if opacity_enc is not None:
         chart = chart.encode(opacity=opacity_enc)
 
@@ -577,14 +589,38 @@ def _parse_y_columns(
     return y_column_list
 
 
+def _get_offset_encoding(
+    chart_type: ChartType,
+    color_column: str | None,
+) -> tuple[alt.XOffset, alt.YOffset]:
+    # Vega's Offset encoding channel is used to create grouped/non-stacked bar charts
+    import altair as alt
+
+    x_offset = alt.XOffset()
+    y_offset = alt.YOffset()
+
+    if chart_type is ChartType.VERTICAL_BAR:
+        x_offset = alt.XOffset(field=color_column)
+    elif chart_type is ChartType.HORIZONTAL_BAR:
+        y_offset = alt.YOffset(field=color_column)
+
+    return x_offset, y_offset
+
+
 def _get_opacity_encoding(
-    chart_type: ChartType, color_column: str | None
+    chart_type: ChartType,
+    stack: bool | ChartStackType | None,
+    color_column: str | None,
 ) -> alt.OpacityValue | None:
     import altair as alt
 
     if color_column and chart_type == ChartType.AREA:
         return alt.OpacityValue(0.7)
 
+    # Layered bar chart
+    if color_column and stack == "layered":
+        return alt.OpacityValue(0.7)
+
     return None
 
 
@@ -634,6 +670,42 @@ def _maybe_melt(
     return df, y_column, color_column
 
 
+def _get_axis_encodings(
+    df: pd.DataFrame,
+    chart_type: ChartType,
+    x_column: str | None,
+    y_column: str | None,
+    x_from_user: str | None,
+    y_from_user: str | Sequence[str] | None,
+    x_axis_label: str | None,
+    y_axis_label: str | None,
+    stack: bool | ChartStackType | None,
+) -> tuple[alt.X, alt.Y]:
+    stack_encoding: alt.X | alt.Y
+    if chart_type == ChartType.HORIZONTAL_BAR:
+        # Handle horizontal bar chart - switches x and y data:
+        x_encoding = _get_x_encoding(
+            df, y_column, y_from_user, x_axis_label, chart_type
+        )
+        y_encoding = _get_y_encoding(
+            df, x_column, x_from_user, y_axis_label, chart_type
+        )
+        stack_encoding = x_encoding
+    else:
+        x_encoding = _get_x_encoding(
+            df, x_column, x_from_user, x_axis_label, chart_type
+        )
+        y_encoding = _get_y_encoding(
+            df, y_column, y_from_user, y_axis_label, chart_type
+        )
+        stack_encoding = y_encoding
+
+    # Handle stacking - only relevant for bar charts
+    _update_encoding_with_stack(stack, stack_encoding)
+
+    return x_encoding, y_encoding
+
+
 def _get_x_encoding(
     df: pd.DataFrame,
     x_column: str | None,
@@ -730,6 +802,19 @@ def _get_y_encoding(
     )
 
 
+def _update_encoding_with_stack(
+    stack: bool | ChartStackType | None,
+    encoding: alt.X | alt.Y,
+) -> None:
+    if stack is None:
+        return None
+    # Our layered option maps to vega's stack=False option
+    elif stack == "layered":
+        stack = False
+
+    encoding["stack"] = stack
+
+
 def _get_color_encoding(
     df: pd.DataFrame,
     color_value: Color | None,

lib/streamlit/elements/vega_charts.py

@@ -57,7 +57,7 @@
 from streamlit.runtime.scriptrunner import get_script_run_ctx
 from streamlit.runtime.state import register_widget
 from streamlit.runtime.state.common import compute_widget_id
-from streamlit.type_util import Key, to_key
+from streamlit.type_util import ChartStackType, Key, to_key
 from streamlit.util import HASHLIB_KWARGS
 
 if TYPE_CHECKING:
@@ -957,6 +957,7 @@ def bar_chart(
         y_label: str | None = None,
         color: str | Color | list[Color] | None = None,
         horizontal: bool = False,
+        stack: bool | ChartStackType | None = None,
         width: int | None = None,
         height: int | None = None,
         use_container_width: bool = True,
@@ -1046,6 +1047,14 @@ def bar_chart(
             Streamlit swaps the x-axis and y-axis and the bars display
             horizontally.
 
+        stack : bool, "normalize", "center", "layered", or None
+            Whether to stack the bars. If this is ``None`` (default), uses Vega's
+            default. If this is ``True``, the bars are stacked on top of each other.
+            If this is ``False``, the bars are displayed side by side. If "normalize",
+            the bars are stacked and normalized to 100%. If "center", the bars are
+            stacked around a central axis. If "layered", the bars are stacked on top
+            of one another.
+
         width : int or None
             Desired width of the chart expressed in pixels. If ``width`` is
             ``None`` (default), Streamlit sets the width of the chart to fit
@@ -1139,6 +1148,20 @@ def bar_chart(
 
         """
 
+        # Offset encodings (used for non-stacked/grouped bar charts) are not supported in Altair < 5.0.0
+        if type_util.is_altair_version_less_than("5.0.0") and stack is False:
+            raise StreamlitAPIException(
+                "Streamlit does not support non-stacked (grouped) bar charts with Altair 4.x. Please upgrade to Version 5."
+            )
+
+        # Check that the stack parameter is valid, raise more informative error message if not
+        VALID_STACK_TYPES = (None, True, False, "normalize", "center", "layered")
+        if stack not in VALID_STACK_TYPES:
+            raise StreamlitAPIException(
+                f'Invalid value for stack parameter: {stack}. Stack must be one of True, False, "normalize", "center", "layered" or None. '
+                "See documentation for `st.bar_chart` [here](https://docs.streamlit.io/develop/api-reference/charts/st.bar_chart) for more information."
+            )
+
         bar_chart_type = (
             ChartType.HORIZONTAL_BAR if horizontal else ChartType.VERTICAL_BAR
         )
@@ -1154,6 +1177,7 @@ def bar_chart(
             size_from_user=None,
             width=width,
             height=height,
+            stack=stack,
         )
         return cast(
             "DeltaGenerator",

lib/streamlit/type_util.py

@@ -142,6 +142,8 @@ def iloc(self) -> _iLocIndexer: ...
 
 VegaLiteType = Literal["quantitative", "ordinal", "temporal", "nominal"]
 
+ChartStackType = Literal["normalize", "center", "layered"]
+
 
 class SupportsStr(Protocol):
     def __str__(self) -> str: ...