[breaking][MT] Support trees to dataframe.#12293
Conversation
There was a problem hiding this comment.
Pull request overview
This PR modernizes Booster.trees_to_dataframe() by switching from parsing the text dump to reading the JSON model, while extending the returned schema to better support multi-output / vector-leaf trees (including a new Target column) and standardizing missing values via pandas nullable dtypes.
Changes:
- Reimplemented
Booster.trees_to_dataframe()to parsesave_raw(raw_format="json"), addTarget, support vector-leaf trees, and emit pandas nullable dtypes. - Expanded parse-tree test helpers and added new assertions/tests for ID integrity, leaf/split invariants, and mixed scalar + vector-leaf boosters.
- Updated indicator-feature test to rely on
feature_types=["i", ...]instead of a feature map file; deprecatedfmapusage.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| tests/python/test_parse_tree.py | Adds stronger structural assertions for trees_to_dataframe and a new vector-leaf mixed test; updates indicator test to use feature_types. |
| python-package/xgboost/testing/parse_tree.py | Adds categorical + vector-leaf dataframe validation helpers (including JSON-based expected leaf vectors). |
| python-package/xgboost/core.py | Rewrites trees_to_dataframe() to use JSON model parsing, adds Target, supports vector-leaf expansion, and standardizes dtypes/missing values. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
RAMitchell
left a comment
There was a problem hiding this comment.
If we wanted this dataframe to be reliably correct we could do a prediction from it and compare against xgboosts prediction.
e.g. right now the following could be wrong and not caught by tests:
- swap yes and no for numerical splits
- wrong missing direction
| # A vector-leaf split is shared by all targets, so it has no single | ||
| # target; scalar-tree splits belong to the tree's target. | ||
| target_ids.append(None if is_vector else tree_target) | ||
| node_ids.append(nid) |
There was a problem hiding this comment.
I will assume all of this is correct
There was a problem hiding this comment.
yea, we will be implementing and maintaining a Python inference library for this test.
|
The dataframe schema doesn't contain missing branches nor the intercept. For dense data and a constant intercept, we can make comparisons. But I would like to keep this out of XGBoost. Throwaway code: def predict_from_dataframe(
df: pd.DataFrame, X: np.ndarray, feature_names: List[str], n_targets: int
) -> np.ndarray:
"""Sum leaf outputs per target by walking the exported tree dataframe.
Parameters
----------
df :
Output of ``Booster.trees_to_dataframe()``.
"""
name2idx = {name: i for i, name in enumerate(feature_names)}
trees = _build_trees(df)
X = np.asarray(X, dtype=np.float32)
n = X.shape[0]
out = np.zeros((n, n_targets), dtype=np.float64)
for nodes in trees.values():
for i in range(n):
nid = 0
while not nodes[nid]["leaf"]:
node = nodes[nid]
nid = _next_node(node, float(X[i, name2idx[node["feature"]]]))
for target, value in nodes[nid]["values"].items():
out[i, target] += value
return out
def _child_node(ref: str) -> int:
return int(str(ref).rsplit("-", 1)[1])
def _build_trees(df: pd.DataFrame) -> Dict[int, Dict[int, dict]]:
"""Parse the dataframe into ``{tree_id: {node_id: node_info}}``."""
trees: Dict[int, Dict[int, dict]] = {}
for tid, group in df.groupby("Tree"):
nodes: Dict[int, dict] = {}
for row in group.itertuples(index=False):
nid = int(row.Node)
if row.Feature == "Leaf":
node = nodes.setdefault(nid, {"leaf": True, "values": {}})
node["values"][int(row.Target)] = float(row.Gain)
else:
cats: Optional[set] = None
if isinstance(row.Category, list):
cats = {int(c) for c in row.Category}
split = None if pd.isna(row.Split) else float(row.Split)
nodes[nid] = {
"leaf": False,
"feature": str(row.Feature),
"split": split,
"cats": cats,
"yes": _child_node(row.Yes),
"no": _child_node(row.No),
"missing": _child_node(row.Missing),
}
trees[int(tid)] = nodes
return trees
def _next_node(node: dict, value: float) -> int:
if np.isnan(value):
return node["missing"]
if node["cats"] is not None:
return node["yes"] if int(value) in node["cats"] else node["no"]
if node["split"] is not None:
return node["yes"] if np.float32(value) < np.float32(node["split"]) else node["no"]
# Indicator (boolean): threshold is always 1.0 and the Yes/No labels depend on the
# (unexported) default direction, so identify the physical children by node-id order
# (left child id < right child id) and apply ``value < 1.0 -> left``.
left, right = sorted((node["yes"], node["no"]))
return left if value < 1.0 else right |
fmap.Targetcolumn has been added to the schema. It works with both scalar trees and vector trees.pandas.NAuniformly to denote that something is not available, instead of mixingNoneandnan.The
trees_to_dataframewas ported from the R implementation. This PR adds a newTargetcolumn, but the rest should remain semantically consistent with the original implementation.ref: #9043