From 0590fea33136265ec77f5a34b3eff98f6a9ac611 Mon Sep 17 00:00:00 2001
From: breken-ai <312387581+breken-ai@users.noreply.github.com>
Date: Fri, 25 Sep 2026 02:53:38 -0700
Subject: [PATCH] Count every boolean category in CategoryCount

CategoryCount(categories=[True, False]) only counted the first category, because the all-boolean branch read counts[categories[0]]. The per-category loop already handles boolean labels, so use it for every case.
---
 src/evidently/metrics/column_statistics.py  | 12 +++------
 tests/future/metrics/test_category_count.py | 28 +++++++++++++++++++++
 2 files changed, 32 insertions(+), 8 deletions(-)
 create mode 100644 tests/future/metrics/test_category_count.py

diff --git a/src/evidently/metrics/column_statistics.py b/src/evidently/metrics/column_statistics.py
index ddde369af1..ee2383b48f 100644
--- a/src/evidently/metrics/column_statistics.py
+++ b/src/evidently/metrics/column_statistics.py
@@ -322,14 +322,10 @@ def _calculate_value(self, dataset: Dataset):
         column = dataset.column(self.metric.column)
         try:
             counts = column.data.value_counts()
-            if all(isinstance(c, bool) for c in self.metric.categories):
-                #  only one boolean label is possible here
-                value = counts[self.metric.categories[0]]  # type: ignore[index]
-            else:
-                value = 0
-                for cat in self.metric.categories:
-                    if cat in counts:
-                        value += counts[cat]  # type: ignore[index]
+            value = 0
+            for cat in self.metric.categories:
+                if cat in counts:
+                    value += counts[cat]  # type: ignore[index]
         except KeyError:
             value = 0
         total = column.data.count()
diff --git a/tests/future/metrics/test_category_count.py b/tests/future/metrics/test_category_count.py
new file mode 100644
index 0000000000..dfd7796905
--- /dev/null
+++ b/tests/future/metrics/test_category_count.py
@@ -0,0 +1,28 @@
+import pandas as pd
+import pytest
+
+from evidently.core.datasets import Dataset
+from evidently.core.metric_types import CountValue
+from evidently.core.report import Report
+from evidently.metrics import CategoryCount
+
+
+@pytest.mark.parametrize(
+    "data,categories,expected_count,expected_share",
+    [
+        (["a", "b", "c", "a"], ["a", "b"], 3, 0.75),
+        (["a", "b", "c", "a"], ["a", "z"], 2, 0.5),
+        ([True, True, False, False], [True], 2, 0.5),
+        ([True, True, False, False], [True, False], 4, 1.0),
+        ([True, True, False, None], [False, True], 3, 1.0),
+        ([True, True, True], [True, False], 3, 1.0),
+    ],
+)
+def test_category_count_multiple_categories(data, categories, expected_count, expected_share):
+    dataset = Dataset.from_pandas(pd.DataFrame(data=dict(column=data)))
+    metric = CategoryCount("column", categories=categories)
+    snapshot = Report([metric]).run(dataset, None)
+    metric_result = snapshot._context.get_metric_result(metric.metric_id)
+    assert isinstance(metric_result, CountValue)
+    assert metric_result.count.value == expected_count
+    assert metric_result.share.value == pytest.approx(expected_share)