Fix #1560 (ARF cornercase)

online-ml · Jun 28, 2024 · c8ea67f · c8ea67f
1 parent 6995f15
commit c8ea67f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/docs/releases/unreleased.md b/docs/releases/unreleased.md
@@ -5,3 +5,7 @@ This release makes Polars an optional dependency instead of a required one.
 ## cluster
 
 - Added `ODAC` (Online Divisive-Agglomerative Clustering) for clustering time series.
+
+## forest
+
+- Fix error in `forest.ARFClassifer` and `forest.ARFRegressor` where the algorithms would crash in case the number of features available for learning went below the value of the `max_features` parameter (#1560).
diff --git a/river/tree/nodes/arf_htc_nodes.py b/river/tree/nodes/arf_htc_nodes.py
@@ -44,7 +44,9 @@ def _iter_features(self, x) -> typing.Iterable:
                 yield att_id, x[att_id]
 
     def _sample_features(self, x, max_features):
-        return self.rng.sample(sorted(x.keys()), k=max_features)
+        if len(x) >= max_features:
+            return self.rng.sample(sorted(x.keys()), k=max_features)
+        return sorted(x.keys())
 
 
 class RandomLeafMajorityClass(BaseRandomLeaf, LeafMajorityClass):