From c1dac1bf867cf94fcc15d5d0d5ab8427b9504509 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 1 Sep 2026 00:35:44 -0700 Subject: [PATCH 1/2] Update export_model(): Add split_kind field --- python/treelite/sklearn/exporter.py | 47 ++++++----------------------- 1 file changed, 9 insertions(+), 38 deletions(-) diff --git a/python/treelite/sklearn/exporter.py b/python/treelite/sklearn/exporter.py index ee0678a5..782488ef 100644 --- a/python/treelite/sklearn/exporter.py +++ b/python/treelite/sklearn/exporter.py @@ -27,42 +27,6 @@ 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( - { - "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 +90,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 = ( From 93074692a63b3f0b91c0ac0e2229eb87f45e5857 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Tue, 1 Sep 2026 16:26:28 -0700 Subject: [PATCH 2/2] Address reviewer feedback Co-authored-by: Tim Head --- python/treelite/sklearn/exporter.py | 30 +++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/python/treelite/sklearn/exporter.py b/python/treelite/sklearn/exporter.py index 782488ef..83b8a570 100644 --- a/python/treelite/sklearn/exporter.py +++ b/python/treelite/sklearn/exporter.py @@ -27,6 +27,28 @@ def _ensure_numpy(x: Any) -> np.ndarray: raise ValueError(f"x is not a valid NumPy array. {x.type=}") +# 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( + { + "left_child", + "right_child", + "feature", + "threshold", + "left_cat_bitset", + "impurity", + "n_node_samples", + "weighted_n_node_samples", + "missing_go_to_left", + "split_kind", + } +) + + class _TaskType(IntEnum): # pylint: disable=invalid-name kBinaryClf = 0 @@ -54,6 +76,14 @@ def _export_tree( raise NotImplementedError( "Trees with categorical splits cannot yet be exported as scikit-learn" ) + unknown_fields = sorted(set(NODE_DTYPE.names) - _KNOWN_NODE_FIELDS) + if unknown_fields: + raise NotImplementedError( + f"scikit-learn {sklearn_version} stores tree node fields that this " + f"version of Treelite does not know how to populate: {unknown_fields}. " + "Exporting would produce a model that predicts incorrectly. " + "Please upgrade Treelite." + ) n_nodes = tree_accessor.get_field("num_nodes").tolist()[0] if parse_version(sklearn_version) >= parse_version("1.10.0.dev0"):