Skip to content

[backport] Update export_model(): Add split_kind field (#698) #700

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
chyunsu3 merged 2 commits into from
Sep 2, 2026
Merged

[backport] Update export_model(): Add split_kind field (#698) #700

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7d94e21
Update export_model(): Add split_kind field (#698)
chyunsu3 Sep 1, 2026
d7c10a2
Merge branch 'release_4.7' into backport_698
chyunsu3 Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 34 additions & 33 deletions python/treelite/sklearn/exporter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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": ["<i8", "<i8", "<i8", "<f8", "<f8", "<i8", "<f8", "u1"],
"offsets": [0, 8, 16, 24, 32, 40, 48, 56],
"itemsize": 64,
"left_child",
"right_child",
"feature",
"threshold",
"left_cat_bitset",
"impurity",
"n_node_samples",
"weighted_n_node_samples",
"missing_go_to_left",
"split_kind",
}
)

_node_dtype_sklearn1_10 = np.dtype(
[
("left_child", np.intp), # 8 bytes (offset 0)
("right_child", np.intp), # 8 bytes (offset 8)
("feature", np.intp), # 8 bytes (offset 16)
("threshold", np.float64), # 8 bytes (offset 24)
("left_cat_bitset", (np.uint32, BITSET_LENGTH)), # 32 bytes (offset 32)
("impurity", np.float64), # 8 bytes (offset 64)
("n_node_samples", np.intp), # 8 bytes (offset 72)
("weighted_n_node_samples", np.float64), # 8 bytes (offset 80)
("missing_go_to_left", np.uint8), # 1 byte (offset 88)
],
align=True,
)


class _TaskType(IntEnum):
# pylint: disable=invalid-name
Expand All @@ -78,6 +64,7 @@ def _export_tree(
try:
from sklearn import __version__ as sklearn_version
from sklearn.tree import DecisionTreeClassifier
from sklearn.tree._tree import NODE_DTYPE, TREE_LEAF
from sklearn.tree._tree import Tree as SKLearnTree
except ImportError as e:
raise TreeliteError("This function requires scikit-learn package") from e
Expand All @@ -88,15 +75,24 @@ 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"):
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")
Expand All @@ -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 = (
Expand Down
Loading