[Refactor] faster get_leafs method for xbooster/lgb_constructor.py

[RektPunk]

Description

This PR optimizes the get_leafs method (specifically for output_type="margin") by replacing iterative DataFrame column assignment with the Collect-then-Construct pattern using np.column_stack.

Key Changes

• Switched from iterative column assignment (df[col] = value) to Collect-then-Construct pattern using np.column_stack.

Performance Benchmark

Measured using %%timeit on the example ipynb:

Method	Execution Time (Mean ± Std. Dev.)	Speedup
Original (Iterative)	6.24 ms ± 139 μs	1.0x
Improved (Stacked)	4.54 ms ± 44.3 μ	~1.4x faster

Summary by Sourcery

Enhancements:

• Replace iterative per-column DataFrame assignment with collection of per-tree prediction arrays and a single np.column_stack-based DataFrame construction to speed up margin leaf retrieval.

[sourcery-ai]

Reviewer's guide (collapsed on small PRs)

Reviewer's Guide

Refactors the margin-output branch of get_leafs in lgb_constructor to batch LightGBM per-tree predictions into a NumPy column stack before constructing the DataFrame, eliminating per-column assignment for better performance.

Class diagram for updated get_leafs method in lgb_constructor

classDiagram
    class LGBConstructor {
        +model
        +get_leafs(X, output_type, n_trees)
    }

    class LightGBMModel {
        +predict(X, raw_score, start_iteration, num_iteration)
    }

    LGBConstructor --> LightGBMModel : uses

Loading

Flow diagram for optimized margin branch in get_leafs

flowchart LR
    A[start] --> B["Call get_leafs with X, output_type, n_trees"]
    B --> C{output_type == margin}
    C -->|no| D["Execute non margin logic (unchanged)"]
    D --> Z[return df_leafs]
    C -->|yes| E["Initialize tree_results as empty list"]
    E --> F["Set i = 0"]
    F --> G{ i < n_trees }
    G -->|no| H["Stack tree_results with np.column_stack"]
    H --> I["Construct df_leafs = DataFrame(stacked, index=X.index, columns=_colnames)"]
    I --> Z[return df_leafs]
    G -->|yes| J["res = model.predict(X, raw_score=True, start_iteration=i, num_iteration=1)"]
    J --> K["Append res to tree_results"]
    K --> L["i = i + 1"]
    L --> G

Loading

File-Level Changes

Change	Details	Files
Optimize margin-output leaf prediction construction by switching from iterative DataFrame column assignment to a collect-then-construct pattern using NumPy.	Replace creation of an empty DataFrame and per-iteration df_leafs[f"tree_{i}"] assignment with accumulation of per-tree prediction arrays in a Python list Use np.column_stack on the collected per-tree results to build a 2D array of predictions in one shot Construct the output DataFrame once from the stacked NumPy array, setting the index to X.index and columns to the precomputed column names	xbooster/lgb_constructor.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey - I've left some high level feedback:

• The new DataFrame construction relies on _colnames, which doesn’t appear in this snippet; double-check that _colnames is defined in this branch and that its length always matches n_trees to avoid runtime errors or misaligned columns.
• Using index=X.index changes behavior compared to the previous implementation (which used the default RangeIndex); ensure X is guaranteed to be a pandas object with an index attribute, or guard against numpy arrays to avoid attribute errors.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- The new `DataFrame` construction relies on `_colnames`, which doesn’t appear in this snippet; double-check that `_colnames` is defined in this branch and that its length always matches `n_trees` to avoid runtime errors or misaligned columns.
- Using `index=X.index` changes behavior compared to the previous implementation (which used the default RangeIndex); ensure `X` is guaranteed to be a pandas object with an `index` attribute, or guard against numpy arrays to avoid attribute errors.

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[xRiskLab]

Great work on these performance optimizations!

All your changes from PRs #10, #11, #13, and #14 have been merged into the release/v0.2.8rc1 release candidate branch. Your optimizations are now part of the upcoming v0.2.8 release.

What's included:

• PR [Refactor] make faster construct_scorecard method for xbooster/lgb_constructor.py #10: Faster construct_scorecard() for LightGBM
• PR [Refactor] make faster _convert_tree_to_points use np.vectorize instead of merge in xbooster/lgb_constructor.py #11: Faster _convert_tree_to_points() using vectorized lookup
• PR [Refactor] faster get_leafs method for xbooster/lgb_constructor.py #13: Faster get_leafs() method for LightGBM
• PR [Refactor] make xgb_construct fast  #14: Faster construct_scorecard() for XGBoost

Impact:

• Replaced loop-based pd.concat() operations with vectorized melt() + groupby()
• Replaced pd.merge() in loops with dictionary mapping using .map()
• Significant performance improvements for models with many trees
• All changes maintain backward compatibility and numerical equivalence

The RC branch (release/v0.2.8rc1) combines your performance improvements with the new SHAP integration features. Once testing is complete, this will be released as v0.2.8.

Thank you for these excellent contributions!