feat(naive_bayes): dataframe-agnostic MultinomialNB via narwhals

docs/releases/unreleased.md

@@ -1,5 +1,8 @@
 # Unreleased
 
+## naive_bayes
+
+- `MultinomialNB.learn_many`, `predict_many`, and `predict_proba_many` now accept any [narwhals](https://github.com/narwhals-dev/narwhals)-supported eager backend (pandas, polars, pyarrow, ...) instead of being pandas-only, preserving the input backend (including the pandas index) on output. A backend-agnostic `BaseNB.predict_many` was added so the argmax-over-probabilities logic is shared by all Naive Bayes variants.
 ## stream
 
 - `stream.Cache` now writes a pass to a temporary file and renames it into place once the stream is exhausted. An interrupted first pass (a `break`, an exception, an abandoned generator) used to leave a truncated file behind, which every later pass then read back as if it were the whole dataset.

river/naive_bayes/base.py

@@ -8,7 +8,7 @@
 from scipy import sparse, special
 
 from river import base, utils
-from river.utils.dataframe import into_frame, to_native_frame, to_numpy
+from river.utils.dataframe import into_frame, to_native_frame, to_native_series, to_numpy
 
 if typing.TYPE_CHECKING:
     import pandas as pd
@@ -55,6 +55,15 @@ def predict_proba_many(self, X: IntoDataFrame) -> IntoDataFrame:
         result = np.exp(jll_np - lse[:, np.newaxis])
         return to_native_frame({col: result[:, i] for i, col in enumerate(columns)}, like=jll_nw)
 
+    def predict_many(self, X: IntoDataFrame):
+        y_pred = self.predict_proba_many(X)
+        y_pred_nw = into_frame(y_pred)
+        if y_pred_nw.is_empty() or not y_pred_nw.columns:
+            return y_pred
+        values = to_numpy(y_pred_nw)
+        labels = np.asarray(y_pred_nw.columns)
+        return to_native_series(labels[np.argmax(values, axis=1)], name=None, like=y_pred_nw)
+
     @property
     def _multiclass(self):
         return True

river/naive_bayes/multinomial.py

@@ -13,7 +13,11 @@
 from . import base
 
 if typing.TYPE_CHECKING:
+    from typing import Any
+
     import pandas as pd
+    from narwhals.stable.v2.typing import IntoDataFrame, IntoSeries
+    from numpy.typing import NDArray
 
 __all__ = ["MultinomialNB"]
 
@@ -185,7 +189,25 @@ def joint_log_likelihood(self, x):
             for c in self.classes_
         }
 
-    def learn_many(self, X: pd.DataFrame, y: pd.Series):
+    @staticmethod
+    def _one_hot_targets(y: IntoSeries) -> tuple[Any, NDArray[np.object_]]:
+        y = utils.dataframe.into_series(y)
+        y_np = np.asarray(y.to_numpy())
+        raw_classes = np.unique(y_np)
+        classes = np.asarray([str(c) for c in raw_classes], dtype=object)
+        indices = np.searchsorted(raw_classes, y_np)
+        rows = np.arange(len(y_np))
+        data: NDArray[np.int64] = np.ones(len(y_np), dtype=np.int64)
+        return sparse.csr_matrix((data, (indices, rows)), shape=(len(classes), len(y_np))), classes
+
+    @staticmethod
+    def _as_sparse_matrix(X):
+        native = X.to_native()
+        if hasattr(native, "sparse"):
+            return sparse.csr_matrix(native.sparse.to_coo())
+        return sparse.csr_matrix(utils.dataframe.to_numpy(X))
+
+    def learn_many(self, X: IntoDataFrame, y: IntoSeries):
         """Learn from a batch of count vectors.
 
         Parameters
@@ -196,16 +218,17 @@ def learn_many(self, X: pd.DataFrame, y: pd.Series):
             Target classes.
 
         """
-        y = base.one_hot_encode(y)
-        columns, classes = X.columns, y.columns
-        y = sparse.csc_matrix(y.sparse.to_coo()).T
+        X = utils.dataframe.into_frame(X)
+        y_one_hot, classes = self._one_hot_targets(y)
+        columns = X.columns
 
-        self.class_counts.update({c: count.item() for c, count in zip(classes, y.sum(axis=1))})
+        self.class_counts.update(
+            {c: int(count.item()) for c, count in zip(classes, y_one_hot.sum(axis=1))}
+        )
 
-        if hasattr(X, "sparse"):
-            X = sparse.csr_matrix(X.sparse.to_coo())
+        X = self._as_sparse_matrix(X)
 
-        fc = y @ X
+        fc = y_one_hot @ X
 
         self.class_totals.update({c: count.item() for c, count in zip(classes, fc.sum(axis=1))})
 
@@ -223,31 +246,22 @@ def learn_many(self, X: pd.DataFrame, y: pd.Series):
                 for f, count in dict_count.items():
                     self.feature_counts[f].update(count)
 
-    def _feature_log_prob(self, columns: list, known: list, unknown: list) -> pd.DataFrame:
-        """Compute log probabilities of input features.
-
-        Parameters
-        ----------
-        columns
-            List of input features.
-        known
-            List of input features that are part of the vocabulary.
-        unknown
-            List of input features that are not part the vocabulary.
-
-        Returns
-        -------
-        Log probabilities of input features.
-
-        """
-        smooth_fc = np.log(base.from_dict(self.feature_counts).fillna(0).T[known] + self.alpha)
-        smooth_fc[unknown] = np.log(self.alpha)
-
-        smooth_cc = np.log(base.from_dict(self.class_totals) + self.alpha * self.n_terms)
-
-        return smooth_fc.subtract(smooth_cc.values, axis="rows")[columns].T
+    def _feature_log_prob(self, columns: list) -> np.ndarray:
+        classes = self.classes_
+        smooth_cc = np.array(
+            [self.class_totals[c] + self.alpha * self.n_terms for c in classes],
+            dtype=float,
+        )
+        feature_log_prob: NDArray[np.float64] = np.empty((len(columns), len(classes)), dtype=float)
+        for i, f in enumerate(columns):
+            smooth_fc = np.array(
+                [self.feature_counts.get(f, {}).get(c, 0.0) + self.alpha for c in classes],
+                dtype=float,
+            )
+            feature_log_prob[i] = np.log(smooth_fc) - np.log(smooth_cc)
+        return feature_log_prob
 
-    def joint_log_likelihood_many(self, X: pd.DataFrame) -> pd.DataFrame:
+    def joint_log_likelihood_many(self, X: IntoDataFrame) -> IntoDataFrame:
         """Computes the joint log likelihood of input features.
 
         Parameters
@@ -260,25 +274,15 @@ def joint_log_likelihood_many(self, X: pd.DataFrame) -> pd.DataFrame:
         Input samples joint log likelihood.
 
         """
-        pd = utils.pandas.import_pandas()
-        index, columns = X.index, X.columns
-        known, unknown = [], []
+        X = utils.dataframe.into_frame(X)
+        columns = X.columns
 
         if not self.class_counts or not self.feature_counts:
-            return pd.DataFrame(index=index)
-
-        for f in columns:
-            if f in self.feature_counts:
-                known.append(f)
-            else:
-                unknown.append(f)
-
-        if hasattr(X, "sparse"):
-            X = sparse.csr_matrix(X.sparse.to_coo())
-
-        return pd.DataFrame(
-            X @ self._feature_log_prob(columns=columns, known=known, unknown=unknown)
-            + np.log(self.p_class_many()).values,
-            index=index,
-            columns=self.class_totals.keys(),
-        )
+            return utils.dataframe.to_native_frame(np.empty((len(X), 0)), columns=[], like=X)
+
+        X_matrix = self._as_sparse_matrix(X)
+        classes = self.classes_
+        jll = X_matrix @ self._feature_log_prob(columns=columns)
+        jll = np.asarray(jll) + np.log(np.array([self.p_class(c) for c in classes], dtype=float))
+
+        return utils.dataframe.to_native_frame(jll, columns=classes, like=X)

tests/naive_bayes/test_naive_bayes.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
+import narwhals.stable.v2 as nw
 import numpy as np
 import pandas as pd
 import pytest
 from sklearn import naive_bayes as sk_naive_bayes
 
 from river import compose, feature_extraction, naive_bayes
+from tests.frames import FRAME_BACKENDS, FrameBackend
 
 
 def river_models():
@@ -241,6 +243,71 @@ def test_river_vs_sklearn(model, sk_model, bag):
             assert river_pred == pytest.approx(1 - sk_pred) or river_pred == pytest.approx(sk_pred)
 
 
+MULTINOMIAL_X = {
+    "Chinese": [2, 2, 1, 1],
+    "Beijing": [1, 0, 0, 0],
+    "Shanghai": [0, 1, 0, 0],
+    "Macao": [0, 0, 1, 0],
+    "Tokyo": [0, 0, 0, 1],
+    "Japan": [0, 0, 0, 1],
+}
+MULTINOMIAL_Y = ["yes", "yes", "maybe", "no"]
+MULTINOMIAL_TEST = {
+    "Chinese": [0, 1, 2],
+    "Beijing": [0, 0, 1],
+    "Shanghai": [0, 1, 0],
+    "Macao": [0, 0, 0],
+    "Tokyo": [1, 0, 0],
+    "Japan": [1, 0, 0],
+    "Taipei": [1, 1, 0],
+}
+
+
+def test_multinomial_predict_many_backend_agnostic(frame_backend: FrameBackend):
+    pandas = FRAME_BACKENDS["pandas"]()
+
+    reference = naive_bayes.MultinomialNB(alpha=1)
+    reference.learn_many(pandas.frame(MULTINOMIAL_X), pandas.series(MULTINOMIAL_Y))
+    expected_proba = nw.from_native(
+        reference.predict_proba_many(pandas.frame(MULTINOMIAL_TEST)),
+        eager_only=True,
+    )
+    expected_pred = nw.from_native(
+        reference.predict_many(pandas.frame(MULTINOMIAL_TEST)),
+        series_only=True,
+    ).to_list()
+
+    model = naive_bayes.MultinomialNB(alpha=1)
+    model.learn_many(frame_backend.frame(MULTINOMIAL_X), frame_backend.series(MULTINOMIAL_Y))
+
+    native_test = frame_backend.frame(MULTINOMIAL_TEST)
+    got_proba_native = model.predict_proba_many(native_test)
+    got_pred_native = model.predict_many(native_test)
+
+    assert type(got_proba_native) is type(native_test)
+    assert type(got_pred_native) is type(frame_backend.series(MULTINOMIAL_Y))
+
+    got_proba = nw.from_native(got_proba_native, eager_only=True)
+    got_pred = nw.from_native(got_pred_native, series_only=True).to_list()
+
+    assert got_proba.columns == expected_proba.columns
+    np.testing.assert_allclose(got_proba.to_numpy(), expected_proba.to_numpy())
+    assert got_pred == expected_pred
+
+
+def test_multinomial_predict_many_not_fit_backend_agnostic(frame_backend: FrameBackend):
+    X = frame_backend.frame({"Chinese": [1, 0], "Tokyo": [0, 1]})
+    model = naive_bayes.MultinomialNB()
+
+    proba = model.predict_proba_many(X)
+    pred = model.predict_many(X)
+
+    assert type(proba) is type(X)
+    assert type(pred) is type(X)
+    assert nw.from_native(proba, eager_only=True).columns == []
+    assert nw.from_native(pred, eager_only=True).columns == []
+
+
 def test_gaussian_learn_many_vs_learn_one():
     X = pd.DataFrame(
         [