Skip to content

[breaking][MT] Support trees to dataframe.#12293

Open
trivialfis wants to merge 9 commits into
dmlc:masterfrom
trivialfis:tree-df-json
Open

[breaking][MT] Support trees to dataframe.#12293
trivialfis wants to merge 9 commits into
dmlc:masterfrom
trivialfis:tree-df-json

Conversation

@trivialfis

@trivialfis trivialfis commented Jul 9, 2026

Copy link
Copy Markdown
Member
  • Use the JSON model instead of parsing the text dump.
  • Deprecate the use of the fmap.
  • A new Target column has been added to the schema. It works with both scalar trees and vector trees.
  • Use pandas.NA uniformly to denote that something is not available, instead of mixing None and nan.
  • Support vector leaf.

The trees_to_dataframe was ported from the R implementation. This PR adds a new Target column, but the rest should remain semantically consistent with the original implementation.

ref: #9043

- Deprecated the `fmap`.
- Use the pandas nullable dtype uniformly to avoid mixing nan with None.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 parse save_raw(raw_format="json"), add Target, 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; deprecated fmap usage.

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.

Comment thread tests/python/test_parse_tree.py
Comment thread python-package/xgboost/testing/parse_tree.py
Comment thread python-package/xgboost/testing/parse_tree.py
Comment thread python-package/xgboost/core.py Outdated
Comment thread python-package/xgboost/core.py

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

Comment thread python-package/xgboost/core.py
Comment thread python-package/xgboost/core.py
Comment thread python-package/xgboost/core.py

@RAMitchell RAMitchell left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will assume all of this is correct

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yea, we will be implementing and maintaining a Python inference library for this test.

@trivialfis

trivialfis commented Jul 9, 2026

Copy link
Copy Markdown
Member Author

@RAMitchell

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants