diff --git a/python/treelite/sklearn/exporter.py b/python/treelite/sklearn/exporter.py index ee0678a5..83b8a570 100644 --- a/python/treelite/sklearn/exporter.py +++ b/python/treelite/sklearn/exporter.py @@ -27,41 +27,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") @@ -124,6 +120,11 @@ def _export_tree( weighted_data_count[~weighted_data_count_mask] = np.nan nodes["weighted_n_node_samples"] = weighted_data_count 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 = (