diff --git a/python/treelite/sklearn/exporter.py b/python/treelite/sklearn/exporter.py index b8c9cd30..8d8b510d 100644 --- a/python/treelite/sklearn/exporter.py +++ b/python/treelite/sklearn/exporter.py @@ -26,41 +26,27 @@ def _ensure_numpy(x: Any) -> np.ndarray: raise ValueError(f"x is not a valid NumPy array. {x.type=}") -BITSET_LENGTH = 8 - -_node_dtype_old = np.dtype( +# Fields of scikit-learn's private Node struct that this exporter knows how to +# handle. If scikit-learn grows a field outside this set we cannot populate it, +# and since we now allocate with scikit-learn's own NODE_DTYPE the dtype check in +# Tree.__setstate__ will no longer catch it -- the field would silently keep its +# zero value. `left_cat_bitset` is knowingly left zeroed: trees with categorical +# splits are rejected before we allocate. +_KNOWN_NODE_FIELDS = frozenset( { - "names": [ - "left_child", - "right_child", - "feature", - "threshold", - "impurity", - "n_node_samples", - "weighted_n_node_samples", - "missing_go_to_left", - ], - "formats": ["= parse_version("1.10.0.dev0"): n_categories = np.full(n_features, -1, dtype=np.intp) tree = SKLearnTree(n_features, n_classes, n_targets, n_categories) - nodes = np.empty(n_nodes, dtype=_node_dtype_sklearn1_10) else: tree = SKLearnTree(n_features, n_classes, n_targets) - nodes = np.empty(n_nodes, dtype=_node_dtype_old) + # Use scikit-learn's runtime node layout, since its private ABI can change + # independently of the package version. + nodes = np.zeros(n_nodes, dtype=NODE_DTYPE) nodes["left_child"] = tree_accessor.get_field("cleft") nodes["right_child"] = tree_accessor.get_field("cright") @@ -106,6 +102,11 @@ def _export_tree( nodes["n_node_samples"] = -1 nodes["weighted_n_node_samples"] = np.nan nodes["missing_go_to_left"] = tree_accessor.get_field("default_left") + if "split_kind" in nodes.dtype.names: + from sklearn.tree._utils import SPLIT_LEAF, SPLIT_NUMERIC + + nodes["split_kind"] = SPLIT_NUMERIC + nodes["split_kind"][nodes["left_child"] == TREE_LEAF] = SPLIT_LEAF if n_targets == 1 and n_classes[0] == 1: leaf_value = (