Numpy test and null handling fixes

Author: mayagbarnes

e2e_playwright/e2e_utils.py

@@ -44,7 +44,6 @@
 INTERACTIONS_SCRIPT = Path(__file__).parent / "pivot_table_interactions_app.py"
 DATA_SCRIPT = Path(__file__).parent / "pivot_table_data_app.py"
 GOLDEN_SCRIPT = Path(__file__).parent / "pivot_table_golden_app.py"
-NUMPY_INPUT_SCRIPT = Path(__file__).parent / "pivot_table_numpy_input_app.py"
 
 PIVOT_KEYS = [
     "test_pivot",
@@ -82,6 +81,8 @@
     "test_pivot_subtotals",
     "test_pivot_cond_fmt",
     "test_pivot_scalar_roundtrip",
+    "test_pivot_numpy_array",
+    "test_pivot_numpy_list",
 ]
 
 # Must match the exact order of st_pivot_table(key=...) calls in
@@ -144,11 +145,6 @@
     "golden_export",
 ]
 
-NUMPY_INPUT_PIVOT_KEYS = [
-    "numpy_array_pivot",
-    "numpy_str_list_pivot",
-]
-
 APP_CONFIGS = {
     "default": {"script": SCRIPT, "pivot_keys": PIVOT_KEYS},
     "pivot_table_test.py": {"script": TOOLBAR_SCRIPT, "pivot_keys": TOOLBAR_PIVOT_KEYS},
@@ -161,10 +157,6 @@
         "script": GOLDEN_SCRIPT,
         "pivot_keys": GOLDEN_PIVOT_KEYS,
     },
-    "pivot_table_numpy_input_test.py": {
-        "script": NUMPY_INPUT_SCRIPT,
-        "pivot_keys": NUMPY_INPUT_PIVOT_KEYS,
-    },
 }
 
 

e2e_playwright/pivot_table_numpy_input_app.py

@@ -350,6 +350,39 @@ def test_scalar_aggregation_roundtrip_persists_and_python_override_wins(
     )
 
 
+def test_numpy_backed_configs_render_without_frontend_validation_errors(
+    page_at_app: Page,
+):
+    """NumPy-backed config inputs render like normal string-list configs."""
+    page = page_at_app
+
+    raw_array = get_pivot(page, "test_pivot_numpy_array")
+    expect(raw_array.get_by_test_id("pivot-table")).to_be_visible(timeout=15000)
+    expect(raw_array.get_by_test_id("pivot-data-cell").first).to_be_visible(
+        timeout=5000
+    )
+    expect(raw_array.get_by_test_id("toolbar-rows-chips")).to_contain_text("Region")
+    expect(raw_array.get_by_test_id("toolbar-columns-chips")).to_contain_text("Year")
+    expect(raw_array.get_by_test_id("toolbar-values-chips")).to_contain_text("Revenue")
+
+    numpy_str_list = get_pivot(page, "test_pivot_numpy_list")
+    expect(numpy_str_list.get_by_test_id("pivot-table")).to_be_visible(timeout=15000)
+    expect(numpy_str_list.get_by_test_id("pivot-data-cell").first).to_be_visible(
+        timeout=5000
+    )
+    expect(numpy_str_list.get_by_test_id("toolbar-rows-chips")).to_contain_text(
+        "Region"
+    )
+    expect(numpy_str_list.get_by_test_id("toolbar-columns-chips")).to_contain_text(
+        "Year"
+    )
+    expect(numpy_str_list.get_by_test_id("toolbar-values-chips")).to_contain_text(
+        "Revenue"
+    )
+
+    expect(page.locator("text=must contain only strings")).to_have_count(0)
+
+
 def test_toolbar_swap_rows_columns(page_at_app: Page):
     """Swapping rows and columns transposes the pivot layout."""
     page = page_at_app

e2e_playwright/pivot_table_numpy_input_test.py

@@ -17,6 +17,9 @@
 
 from __future__ import annotations
 
+from typing import Any, cast
+
+import numpy as np
 import streamlit as st
 
 from streamlit_pivot import st_pivot_table
@@ -124,6 +127,34 @@ def render_app(data):
         on_config_change=noop,
     )
 
+    raw_rows = cast(Any, np.array(["Region"]))
+    raw_columns = cast(Any, np.array(["Year"]))
+    raw_values = cast(Any, np.array(["Revenue"]))
+
+    st.subheader("NumPy ndarray Pivot")
+    st_pivot_table(
+        df,
+        key="test_pivot_numpy_array",
+        rows=raw_rows,
+        columns=raw_columns,
+        values=raw_values,
+        aggregation="sum",
+        interactive=True,
+        on_config_change=noop,
+    )
+
+    st.subheader("NumPy list Pivot")
+    st_pivot_table(
+        df,
+        key="test_pivot_numpy_list",
+        rows=list(np.array(["Region"])),
+        columns=list(np.array(["Year"])),
+        values=list(np.array(["Revenue"])),
+        aggregation="sum",
+        interactive=True,
+        on_config_change=noop,
+    )
+
 
 def main():
     init_page()

e2e_playwright/pivot_table_test.py

@@ -375,11 +375,14 @@ def _prepare_threshold_hybrid_frame(
     named: dict[str, pd.NamedAgg] = {}
     avg_fields: list[str] = []
     numeric_coerce_fields: list[str] = []
+    count_like_fields: list[str] = []
 
     for vf in value_fields:
         agg = aggregation.get(vf, "sum")
         if agg in _NUMERIC_COERCE_AGGS:
             numeric_coerce_fields.append(vf)
+        if agg in ("count", "count_distinct"):
+            count_like_fields.append(vf)
         if agg == "avg":
             avg_fields.append(vf)
             named[f"{vf}__sum"] = pd.NamedAgg(column=vf, aggfunc="sum")
@@ -412,6 +415,8 @@ def _prepare_threshold_hybrid_frame(
             ser = filtered_df[vf]
             if agg in _NUMERIC_COERCE_AGGS:
                 ser = _coerce_measure_series(ser, vf, null_handling)
+            elif agg in ("count", "count_distinct"):
+                ser = _resolve_count_series(ser, vf, null_handling)
             if agg == "avg":
                 cnt = int(ser.count())
                 row[vf] = float(ser.sum() / cnt) if cnt else float("nan")
@@ -452,6 +457,8 @@ def _prepare_threshold_hybrid_frame(
         working[dim] = _resolve_dim_value_series(working[dim], col_type, mode, grain)
     for vf in numeric_coerce_fields:
         working[vf] = _coerce_measure_series(working[vf], vf, null_handling)
+    for vf in count_like_fields:
+        working[vf] = _resolve_count_series(working[vf], vf, null_handling)
 
     out = (
         working.groupby(group_fields, dropna=False, observed=True, sort=False)
@@ -663,6 +670,13 @@ def _coerce_measure_series(series: Any, field: str, null_handling: Any) -> Any:
     return numeric
 
 
+def _resolve_count_series(series: Any, field: str, null_handling: Any) -> Any:
+    """Preserve values for count-like aggs, zero-filling nulls when requested."""
+    if _get_null_mode(field, null_handling) == "zero":
+        return series.fillna(0)
+    return series
+
+
 def _extract_styler_formats(
     styler: Any,
 ) -> tuple[dict[str, str], dict[str, str]]:

e2e_playwright/pivot_table_toolbar_app.py

@@ -1897,6 +1897,18 @@ describe("PivotData - hybrid agg remap", () => {
     expect(pd.getAggregator(["EU"], ["2023"]).value()).toBe(5);
   });
 
+  it("count remap preserves zero for empty intersections", () => {
+    const countData: DataRecord[] = [
+      { region: "US", year: "2023", revenue: 3 },
+      { region: "EU", year: "2024", revenue: 5 },
+    ];
+    const cfg = makeConfig({ aggregation: "count" });
+    const pd = new PivotData(countData, cfg, {
+      hybridAggRemap: { revenue: "sum" },
+    });
+    expect(pd.getAggregator(["US"], ["2024"]).value()).toBe(0);
+  });
+
   it("count_distinct leaf cells with remap show pre-computed value", () => {
     const data: DataRecord[] = [
       { region: "US", year: "2023", revenue: 7 },

streamlit_pivot/__init__.py

@@ -429,6 +429,12 @@ export class PivotData {
     );
   }
 
+  private _emptyAggregatorForField(valField: string): Aggregator {
+    return getAggregatorFactory(
+      getAggregationForField(valField, this._config),
+    ).create();
+  }
+
   // ---------------------------------------------------------------------------
   // Filtering
   // ---------------------------------------------------------------------------
@@ -1124,7 +1130,7 @@ export class PivotData {
     const keyStr = `${makeKeyString(rowKey)}\x01${makeKeyString(colKey)}\x01${field}`;
     const agg = this._cellAggs.get(keyStr);
     if (!agg) {
-      return this._factoryForField(field).create();
+      return this._emptyAggregatorForField(field);
     }
     return agg;
   }
@@ -1145,7 +1151,7 @@ export class PivotData {
     }
     const agg = this._rowTotalAggs.get(keyStr);
     if (!agg) {
-      return this._factoryForField(field).create();
+      return this._emptyAggregatorForField(field);
     }
     return agg;
   }
@@ -1166,7 +1172,7 @@ export class PivotData {
     }
     const agg = this._colTotalAggs.get(keyStr);
     if (!agg) {
-      return this._factoryForField(field).create();
+      return this._emptyAggregatorForField(field);
     }
     return agg;
   }
@@ -1187,7 +1193,7 @@ export class PivotData {
     }
     const agg = this._grandTotalAggs.get(field);
     if (!agg) {
-      return this._factoryForField(field).create();
+      return this._emptyAggregatorForField(field);
     }
     return agg;
   }

streamlit_pivot/frontend/src/engine/PivotData.test.ts

@@ -119,6 +119,14 @@ describe("CountAggregator", () => {
     expect(agg.value()).toBe(3);
   });
 
+  it("ignores NaN values", () => {
+    const agg = createAggregator("count");
+    agg.push(1);
+    agg.push(Number.NaN);
+    agg.push(0);
+    expect(agg.value()).toBe(2);
+  });
+
   it("returns 0 for all nulls", () => {
     const agg = createAggregator("count");
     agg.push(null);
@@ -249,6 +257,12 @@ describe("CountDistinctAggregator", () => {
     expect(agg.value()).toBe(3);
   });
 
+  it("ignores NaN values", () => {
+    const agg = createAggregator("count_distinct");
+    [1, Number.NaN, 1, 2].forEach((v) => agg.push(v));
+    expect(agg.value()).toBe(2);
+  });
+
   it("returns null for empty input", () => {
     const agg = createAggregator("count_distinct");
     expect(agg.value()).toBeNull();