dd training callbacks, feature validation, and update XGBoost API integration' or 'Implement callback mechanism and refactor gradient/hessian passing with array interfaces

[agene0001]

Fixed deprecated function warnings and added TODO features. Also bumped up version numbers and edition

Summary by CodeRabbit

• New Features

  • Per-round training callbacks with early-stop capability.
  • New distinct BinaryError metric for default-threshold binary classification.

• Added

  • Faster, modern CSR/CSC and binary data loading via a new array-interface path.
  • New benchmarks, verification guides, and example benchmark tooling; benchmark results included.
  • Package/edition version bump and updated examples.

• Fixed

  • Validation to detect mismatched feature names vs. column counts.

• Documentation

  • Updated changelog and detailed benchmark documentation.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds iteration-aware custom-objective training, per-round callbacks with early-stop, JSON array-interface marshaling for gradients/hessians, feature-name count validation, and tests updated to use the new BinaryError metric.

Changes

Booster & Training

Layer / File(s)	Summary
Array-interface helpers & JSON marshaling src/booster.rs	Adds NumPy-style JSON array-interface helper for f32 buffers; marshals gradients/hessians into JSON interfaces passed to XGBoosterTrainOneIter.
Iteration-aware update & train loop src/booster.rs	update_custom now accepts iteration; training loop threads current iteration into custom objectives and calls TrainOneIter with iteration.
Per-round callbacks & early stop src/booster.rs, src/parameters/mod.rs	Collects evaluation results into CallbackEnv each round and invokes registered TrainingCallbacks which can stop training by returning false.
Feature-name validation src/booster.rs	Adds private validate_features that errors when booster feature-name count and DMatrix column count both known and mismatch.
Tests src/booster.rs (tests)	Unit tests updated to use BinaryError instead of BinaryErrorRate(0.5).

DMatrix & IO

Layer / File(s)	Summary
(Not covered by these ranges) src/dmatrix.rs, benches/examples	CSR/CSC and URI/config loading changes, 64-bit index APIs, and benchmarks/examples are present in the PR but not in these booster-focused ranges.

Sequence Diagram(s)

sequenceDiagram
    participant Trainer as "Trainer"
    participant Booster as "Booster"
    participant Obj as "Custom Objective"
    participant Eval as "Evaluator"
    participant Callback as "Training Callback"

    Trainer->>Booster: train(dtrain, params with callbacks)
    loop per iteration
        Booster->>Obj: compute objective(dtrain, iteration)
        Obj-->>Booster: gradients, hessians
        Booster->>Booster: make_array_interface(grad, hess)
        Booster->>Eval: evaluate(dtrain)
        Eval-->>Booster: evaluation_results
        Booster->>Callback: invoke callback(&CallbackEnv)
        alt callback returns false
            Callback-->>Booster: false
            Booster-->>Trainer: stop training early
        else callback returns true
            Callback-->>Booster: true
            Booster->>Booster: XGBoosterTrainOneIter(..., iteration)
        end
    end

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

Poem

🐰 I hopped through arrays and rounds today,
JSON bridges carry grad and hess away,
Callbacks whisper "stop" or "go" each turn,
Error split lets metrics rightly learn,
A rabbit cheers the code that hums and plays!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Title check	❓ Inconclusive	The pull request title is unclear and appears to be incomplete or corrupted, containing fragmented phrases ('dd training callbacks', 'update XGBoost API integration') that don't form a coherent summary of the main changes.	Revise the title to be a single, clear sentence that accurately describes the primary change. For example: 'Implement training callbacks and refactor API integration with XGBoost 3.1.3' or similar.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

src/booster.rs (1)

233-243: Custom eval results are stored under the wrong key.

dmat_eval_results is documented as dataset → metric, but the custom eval path inserts eval_name as the outer key and dmat_name as the inner key, swapping axes. This produces incorrect output (e.g., custom-train) and makes callback data inconsistent.

🐛 Fix key placement

-                    let eval_results = dmat_eval_results
-                        .entry(eval_name.to_string())
-                        .or_insert_with(IndexMap::new);
-                    eval_results.insert(dmat_name.to_string(), eval_result);
+                    let eval_results = dmat_eval_results
+                        .entry(dmat_name.to_string())
+                        .or_insert_with(IndexMap::new);
+                    eval_results.insert(eval_name.to_string(), eval_result);

src/dmatrix.rs (1)

167-221: Design clarity improvement: Document that CSR/CSC inputs are copied by XGBoost.

The current implementation passes raw pointers to XGBoost's C API, which could appear unsafe without documentation. However, XGBoost copies the data during the FFI call (consistent with from_dense which uses the identical pattern), making the borrowed-slice API safe. Add a documentation comment explicitly stating that input buffers are copied and can be safely dropped after the call, or consider adding a safety note similar to new_with_cached_dmats() in Booster.

🤖 Fix all issues with AI agents

In `@src/dmatrix.rs`:
- Around line 258-265: The code in load_binary currently uses
path.as_ref().to_string_lossy(), which can silently alter non‑UTF8 paths; change
it to require valid UTF‑8 by using path.as_ref().to_str() and return an error if
it returns None (or map to a Result) before building the JSON config; then use
the obtained &str (escape backslashes/quotes as you already do) to create the
JSON config and CString, keeping the rest (config_cstr and
XGDMatrixCreateFromURI call) unchanged — reference symbols: load_binary, path,
escaped_path, config, config_cstr, and XGDMatrixCreateFromURI.
- Around line 12-46: The current make_array_interface_usize produces
platform-dependent typestr and can yield 4-byte types on 32-bit targets, which
breaks XGBoost's expectation of 64-bit CSR/CSC indices; change the API to use
fixed-width u64: replace or overload make_array_interface_usize to accept &[u64]
(or add new make_array_interface_u64) and ensure callers that pass indptr and
indices convert/collect their values to Vec<u64> before calling (update call
sites that currently pass indptr/indices slices), then emit typestr "<u8"
unconditionally so the array interface always encodes 8-byte unsigned indices.

♻️ Duplicate comments (1)

src/dmatrix.rs (1)

409-415: Array-interface lifetime/typing verification (same as above).

🧹 Nitpick comments (3)

Cargo.toml (1)

11-24: Confirm MSRV and prerelease dependency intent.

Edition 2024 (Line 11) raises MSRV and libc 1.0.0-alpha.2 is a pre-release. Please confirm the minimum supported Rust version is documented/CI-updated and that the alpha dependency is intentional; otherwise consider a stable libc version.

src/booster.rs (2)

13-28: Optional: centralize array-interface helper to reduce duplication.

make_array_interface duplicates the f32 helper in src/dmatrix.rs. Consider a shared utility to keep the spec in one place.

---

516-543: Consider validating features for the standard update path too.

validate_features is only invoked in boost (custom objective). If the goal is to guard all training flows, call it from update or at the start of the train loop.

[agene0001]

Hey @marcomq, I was wondering if you could take a look at this PR as I use it in one of my projects. I don't believe it breaks anything as the tests still pass. Thanks

[marcomq]

Thank you very much for the PR. It looks good and seems to be a good step forward.

However, it might introduce some breaking changes and might have performance impact as it internally uses now different functions that work with strings. These from_csc / from_dense functions are usually called in the inner loop and as we don't have performance benchmarks yet, I would prefer to not apply these changes silently.

Can you try to create new separate functions for the changed from_csc / from_dense, so it doesn't affect current functionality.

[marcomq]

This is a breaking change as it uses another C function that uses urls instead of files. For consistency, we should also use a new function here too. Let's keep the old function as it is and add load_binary_from_uri

[marcomq]

This is modern, but might be a breaking change for others. Can we keep it at 2021?

[marcomq]

I don't feel good about using an alpha version. Is this really necessary?

[marcomq]

xgboost is used as a library. Just keep "0.20" / "0.4". Rust will automatically try to select the latest most compatible version.
Only update it if necessary or breaking change.

In your later app, you can specify a more recent version, if you like to.

[marcomq]

Generally a good change.
But this looks like it may affect performance negatively as we are now allocating new strings for all the input. We must avoid undocumented perforformance regressions. Please create a new function, maybe from_dense_new or similar.

[marcomq]

Same here, this may have a negative performance impact.

[coderabbitai]

Actionable comments posted: 11

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

src/dmatrix.rs (1)

169-202: Breaking API change not addressed per maintainer feedback.

The maintainer previously requested keeping the old function and adding a new one (e.g., from_csr_new) to avoid undocumented performance regressions from the new string allocations. The signature change from &[usize] to &[u64] is also a breaking change for existing callers.

Consider either:

1. Restoring the old from_csr implementation and adding from_csr_v2 or similar
2. Adding a deprecation notice and migration guide if replacing is intentional

🤖 Fix all issues with AI agents

In `@benches/dmatrix_benchmark.rs`:
- Line 240: Update the printed threshold text so it matches the library logic:
change the println call that currently outputs "Threshold: 5000 non-zeros
(single-thread below, multi-thread above)" to state 30000 non-zeros instead;
locate and edit the println! invocation in benches/dmatrix_benchmark.rs (the
line that prints the threshold message) to reflect "Threshold: 30000 non-zeros
(single-thread below, multi-thread above)".
- Around line 122-130: The SINGLE_THREAD_THRESHOLD in from_csr_auto_tuned is set
to 5000 which diverges from the library implementation and bench_table.rs;
update the threshold constant used in from_csr_auto_tuned to 30000 so it matches
DMatrix::from_csr's behavior and the example benchmark; ensure the constant name
SINGLE_THREAD_THRESHOLD remains unchanged and that any logic using it continues
to select single-thread mode for sizes below 30000.
- Around line 306-310: The hardcoded 5000 threshold in the optimal selection
logic (the expression setting optimal using nnz, single_us, multi_us) must be
replaced with a named constant and updated to the correct threshold value;
introduce a const (e.g., OPTIMAL_NNZ_THRESHOLD) near the benchmarking constants
and use it in the condition: if (nnz < OPTIMAL_NNZ_THRESHOLD && single_us <=
multi_us) || (nnz >= OPTIMAL_NNZ_THRESHOLD && multi_us <= single_us) so the
threshold can be changed centrally and set to the updated value required by the
benchmark.
- Around line 329-333: There is a duplicate main when feature "print_table" is
enabled: the standalone fn main calling print_comparison_table conflicts with
the main generated by the Criterion macro; to fix it, gate one of them with cfg
so only one main is compiled — for example add #[cfg(not(feature =
"print_table"))] around the Criterion entry (criterion_main!(benches)) or
alternatively add #[cfg(feature = "print_table")] to the table-printer main;
reference the standalone main that calls print_comparison_table and the
criterion_main!(benches) macro when applying the cfg.

In `@BENCHMARK_VERIFICATION.md`:
- Around line 30-38: The fenced code block showing expected output (the block
starting with "Rows | NNZ | Deprecated (us)" in BENCHMARK_VERIFICATION.md) lacks
a language specifier; update that fenced block to include a language tag such as
"text" or "plaintext" (e.g., change ``` to ```text) so the static analysis
requirement for annotated code blocks is satisfied and the expected-output block
is treated as plain text.

In `@examples/basic/src/main.rs`:
- Around line 68-74: The printed "error=" value is actually computing accuracy
(num_correct as f32 / preds.len() as f32) which is misleading; update the
calculation in the println to use the true error rate by computing 1.0 -
(num_correct as f32 / preds.len() as f32) or (preds.len() - num_correct) as f32
/ preds.len() so the label "error=" matches the value (refer to the variables
num_correct and preds and the println! call).

In `@examples/bench_deprecated.rs`:
- Around line 12-13: The file is missing imports for ptr::null_mut and
Instant::now; add the appropriate imports (for example import std::ptr or
std::ptr::null_mut and std::time::Instant) so calls to ptr::null_mut() and
Instant::now() resolve. Update the top of the file near the existing extern
crate xgboost_sys declaration to include these imports so functions referencing
ptr::null_mut() and Instant::now() compile.
- Around line 54-56: The FFI call to xgboost_sys::XGDMatrixCreateFromCSREx is
incorrectly casting a Vec<u64> pointer to *const u32 (the indices variable)
causing misinterpreted memory; fix by converting or recreating indices as
Vec<u32> (or mapping indices.iter().map(|v| *v as u32).collect() immediately
before the call) and pass that u32 buffer's pointer to XGDMatrixCreateFromCSREx,
ensuring the indptr/indices types match the expected C signature.
- Around line 53-63: The call to XGDMatrixCreateFromCSREx is passing mismatched
pointer types and ignoring its return code: cast indptr.as_ptr() to *const
size_t, ensure indices is a &[u32] (or perform a checked conversion from &[u64]
to Vec<u32] to avoid truncation) and pass indices.as_ptr() as *const u32, keep
data.as_ptr(), indptr.len(), data.len(), num_cols and &mut handle as before,
then capture the returned int (e.g., let ret =
xgboost_sys::XGDMatrixCreateFromCSREx(...)) and check if ret != 0 and handle or
propagate the error (panic or map to Result) instead of ignoring it.

In `@examples/bench_table.rs`:
- Around line 75-86: The call to the FFI function XGDMatrixCreateFromCSR
currently ignores its return code and may leave handle null; update the code to
check the function's return value and verify handle is non-null before returning
or continuing. Specifically, capture the return code from
XGDMatrixCreateFromCSR, handle non-zero/error codes by returning an Err or
panicking with a clear message, and if the call appears to fail ensure you do
not call XGDMatrixFree on a null handle; alternatively return a Result from the
surrounding function so callers can handle failures instead of unconditionally
returning handle.

♻️ Duplicate comments (3)

Cargo.toml (2)

15-19: Prefer semver-compatible ranges over exact version pinning for library dependencies.

Exact version pinning (e.g., libc = "0.2.180") in a library crate can cause dependency resolution conflicts for downstream users. Consider using semver ranges like "0.2" or "0.4" to allow Cargo to select compatible versions, as mentioned in previous review feedback.

♻️ Suggested change

-libc = "0.2.180"
-derive_builder = "0.20.2"
-log = "0.4.29"
-tempfile = "3.24.0"
-indexmap = "2.13.0"
+libc = "0.2"
+derive_builder = "0.20"
+log = "0.4"
+tempfile = "3"
+indexmap = "2"

---

11-11: Rust edition 2024 may cause compatibility issues for downstream users.

Edition 2024 is very recent and may introduce breaking changes for users who haven't upgraded their toolchains. For a library crate, consider keeping edition = "2021" for broader compatibility, as mentioned in previous review feedback.

src/dmatrix.rs (1)

216-249: Same breaking change concern as from_csr.

This has the same issue - the signature change and new allocations weren't addressed per the maintainer's feedback requesting a new function to preserve backward compatibility.

🧹 Nitpick comments (4)

examples/bench_table.rs (1)

34-59: Code duplication with benches/dmatrix_benchmark.rs.

The generate_sparse_data function is duplicated verbatim across examples/bench_table.rs, benches/dmatrix_benchmark.rs, and examples/bench_deprecated.rs. Consider extracting this to a shared module (e.g., a bench_utils module) to improve maintainability.

examples/bench_deprecated.rs (1)

15-41: Duplicated generate_sparse_data function.

This is the third copy of this function (also in examples/bench_table.rs and benches/dmatrix_benchmark.rs). As noted earlier, consider extracting to a shared utility module.

benches/dmatrix_benchmark.rs (2)

18-36: Consider extracting shared array interface helpers.

These helper functions are duplicated from src/dmatrix.rs. Consider extracting them to a shared internal module to avoid drift and reduce maintenance burden.

---

171-177: Consider adjusting test cases for 30k threshold.

With the library's threshold of 30000, only "large_50k_nnz" and "xlarge_100k_nnz" would trigger multi-threaded mode in auto-tuned. Consider adding a test case near the 30k boundary (e.g., 25k and 35k nnz) to better validate the threshold behavior.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Add language specifier to fenced code blocks.

Per static analysis, fenced code blocks should have a language specified. For expected output, use text or plaintext.

📝 Suggested fix

 You should see output like:
-```
+```text
     Rows |      NNZ | Deprecated (us)
 ---------+----------+----------------
      100 |      979 |          56.00

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	You should see output like:
	```
	Rows | NNZ | Deprecated (us)
	---------+----------+----------------
	100 | 979 | 56.00
	1000 | 10037 | 487.00
	2000 | 19849 | 980.00
	...
	```
	You should see output like:

🧰 Tools

🪛 markdownlint-cli2 (0.18.1)

35-35: Fenced code blocks should have a language specified

(MD040, fenced-code-language)

🤖 Prompt for AI Agents

In `@BENCHMARK_VERIFICATION.md` around lines 30 - 38, The fenced code block
showing expected output (the block starting with "Rows | NNZ | Deprecated (us)"
in BENCHMARK_VERIFICATION.md) lacks a language specifier; update that fenced
block to include a language tag such as "text" or "plaintext" (e.g., change ```
to ```text) so the static analysis requirement for annotated code blocks is
satisfied and the expected-output block is treated as plain text.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Missing required imports.

The file uses ptr::null_mut() and Instant::now() but doesn't import std::ptr or std::time::Instant. This will cause compilation errors.

🐛 Proposed fix

 extern crate xgboost_sys;
+
+use std::ptr;
+use std::time::Instant;

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	extern crate xgboost_sys;
	extern crate xgboost_sys;
	use std::ptr;
	use std::time::Instant;

🤖 Prompt for AI Agents

In `@examples/bench_deprecated.rs` around lines 12 - 13, The file is missing
imports for ptr::null_mut and Instant::now; add the appropriate imports (for
example import std::ptr or std::ptr::null_mut and std::time::Instant) so calls
to ptr::null_mut() and Instant::now() resolve. Update the top of the file near
the existing extern crate xgboost_sys declaration to include these imports so
functions referencing ptr::null_mut() and Instant::now() compile.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

rg -n "XGDMatrixCreateFromCSREx" --type rust -A 10

Repository: marcomq/rust-xgboost

Length of output: 3036

---

🏁 Script executed:

fd -e rs -e toml | xargs rg -l "XGDMatrixCreateFromCSREx" | head -20

Repository: marcomq/rust-xgboost

Length of output: 92

---

🏁 Script executed:

# Search for xgboost-sys as a dependency
rg "xgboost.sys|xgboost_sys" --type toml -A 5

Repository: marcomq/rust-xgboost

Length of output: 600

---

🏁 Script executed:

# Search for function definition in bindings or lib files
rg -n "fn XGDMatrixCreateFromCSREx|extern.*XGDMatrixCreateFromCSREx" -A 5

Repository: marcomq/rust-xgboost

Length of output: 46

---

🏁 Script executed:

# Explore xgboost-sys directory structure
find xgboost-sys -name "*.rs" | head -20

Repository: marcomq/rust-xgboost

Length of output: 107

---

🏁 Script executed:

# Search specifically in xgboost-sys for the function
rg -n "XGDMatrixCreateFromCSREx" xgboost-sys/ -A 5

Repository: marcomq/rust-xgboost

Length of output: 46

---

🏁 Script executed:

# Read xgboost-sys lib.rs to understand bindings
cat xgboost-sys/src/lib.rs

Repository: marcomq/rust-xgboost

Length of output: 1057

---

🏁 Script executed:

# Check build.rs to see how bindings are generated
head -100 xgboost-sys/build.rs

Repository: marcomq/rust-xgboost

Length of output: 4529

---

🏁 Script executed:

# Find the C header file
find . -name "c_api.h" -o -name "c_api.h*"

Repository: marcomq/rust-xgboost

Length of output: 46

---

🏁 Script executed:

# Search for XGDMatrixCreateFromCSREx in header files
rg -n "XGDMatrixCreateFromCSREx" --type c --type h -B 2 -A 10

Repository: marcomq/rust-xgboost

Length of output: 212

---

🏁 Script executed:

# Search for any documentation or comments mentioning parameter types
rg -n "indptr|indices|XGDMatrix" examples/bench_deprecated.rs -B 5 -A 5

Repository: marcomq/rust-xgboost

Length of output: 4411

---

🏁 Script executed:

# Look for type definitions in the codebase
rg -n "size_t|bindex_t" --type rust

Repository: marcomq/rust-xgboost

Length of output: 46

---

🏁 Script executed:

# Check how other XGDMatrix functions handle return values
rg -n "XGDMatrix.*unsafe" -A 3 examples/

Repository: marcomq/rust-xgboost

Length of output: 46

---

🌐 Web query:

XGDMatrixCreateFromCSREx XGBoost C API signature parameters

💡 Result:

int XGDMatrixCreateFromCSREx(
    const size_t*   indptr,
    const unsigned* indices,
    const float*    data,
    size_t          nindptr,
    size_t          nelem,
    size_t          num_col,
    DMatrixHandle*  out
);

[1]

Parameters

• indptr: CSR row pointer array (length = nindptr). [2]
• indices: CSR column indices for each non-zero (length = nelem). [2]
• data: CSR non-zero values (length = nelem). [2]
• nindptr: number of rows + 1. [2]
• nelem: number of non-zero elements. [2]
• num_col: number of columns (if 0, XGBoost will try to infer it). [2]
• out: output DMatrixHandle (created DMatrix). [2]

Note: XGDMatrixCreateFromCSREx is deprecated since XGBoost 2.0.0; the recommended API is XGDMatrixCreateFromCSR() (array-interface based). [1]

---

🏁 Script executed:

# Search for any comments or documentation files about the deprecated API
rg -n "CSR|deprecated" --type md -A 5

Repository: marcomq/rust-xgboost

Length of output: 7608

---

🏁 Script executed:

# Look at the actual function signature from test cases if they exist
rg -n "XGDMatrixCreateFromCSREx|XGDMatrixCreateFromCSR" -B 10 -A 5

Repository: marcomq/rust-xgboost

Length of output: 17944

---

Fix type mismatches and add error handling for deprecated API call.

The XGDMatrixCreateFromCSREx function expects const size_t* indptr and const unsigned* indices, but the code passes:

• indptr.as_ptr() where indptr is &[u64] — size_t is platform-dependent (32-bit on some platforms, 64-bit on others)
• indices.as_ptr() as &[u64] cast to *const u32 — loses data when converting u64 to u32

Additionally, the function returns an error code (int) that is completely ignored. Cast both arrays to the correct pointer types and check the return value:

Example fix

unsafe {
    let ret = xgboost_sys::XGDMatrixCreateFromCSREx(
        indptr.as_ptr() as *const size_t,
        indices.as_ptr() as *const u32,  // or validate no data loss
        data.as_ptr(),
        indptr.len(),
        data.len(),
        num_cols,
        &mut handle,
    );
    if ret != 0 {
        panic!("XGDMatrixCreateFromCSREx failed with error code: {}", ret);
    }
}

🤖 Prompt for AI Agents

In `@examples/bench_deprecated.rs` around lines 53 - 63, The call to
XGDMatrixCreateFromCSREx is passing mismatched pointer types and ignoring its
return code: cast indptr.as_ptr() to *const size_t, ensure indices is a &[u32]
(or perform a checked conversion from &[u64] to Vec<u32] to avoid truncation)
and pass indices.as_ptr() as *const u32, keep data.as_ptr(), indptr.len(),
data.len(), num_cols and &mut handle as before, then capture the returned int
(e.g., let ret = xgboost_sys::XGDMatrixCreateFromCSREx(...)) and check if ret !=
0 and handle or propagate the error (panic or map to Result) instead of ignoring
it.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

Dangerous pointer cast from u64* to u32*.

The indices vector contains u64 values, but line 56 casts the pointer to *const u32. This reinterprets the memory layout incorrectly—each 64-bit index will be read as two 32-bit values, producing garbage data and potential memory access violations.

Either generate indices as Vec<u32> or convert the data before the FFI call.

🐛 Proposed fix - convert indices to u32

 #[cfg(feature = "deprecated_api")]
 fn from_csr_deprecated(
     indptr: &[u64],
     indices: &[u64],
     data: &[f32],
     num_cols: usize,
     num_rows: usize,
 ) -> xgboost_sys::DMatrixHandle {
     let mut handle = ptr::null_mut();
+    
+    // Convert indices to u32 for deprecated API
+    let indices_u32: Vec<u32> = indices.iter().map(|&x| x as u32).collect();

     unsafe {
         xgboost_sys::XGDMatrixCreateFromCSREx(
             indptr.as_ptr(),
-            indices.as_ptr() as *const u32,
+            indices_u32.as_ptr(),
             data.as_ptr(),
             indptr.len(),
             data.len(),
             num_cols,
             &mut handle,
         );
     }
     handle
 }

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	xgboost_sys::XGDMatrixCreateFromCSREx(
	indptr.as_ptr(),
	indices.as_ptr() as *const u32,
	#[cfg(feature = "deprecated_api")]
	fn from_csr_deprecated(
	indptr: &[u64],
	indices: &[u64],
	data: &[f32],
	num_cols: usize,
	num_rows: usize,
	) -> xgboost_sys::DMatrixHandle {
	let mut handle = ptr::null_mut();
	// Convert indices to u32 for deprecated API
	let indices_u32: Vec<u32> = indices.iter().map(|&x| x as u32).collect();
	unsafe {
	xgboost_sys::XGDMatrixCreateFromCSREx(
	indptr.as_ptr(),
	indices_u32.as_ptr(),
	data.as_ptr(),
	indptr.len(),
	data.len(),
	num_cols,
	&mut handle,
	);
	}
	handle
	}

🤖 Prompt for AI Agents

In `@examples/bench_deprecated.rs` around lines 54 - 56, The FFI call to
xgboost_sys::XGDMatrixCreateFromCSREx is incorrectly casting a Vec<u64> pointer
to *const u32 (the indices variable) causing misinterpreted memory; fix by
converting or recreating indices as Vec<u32> (or mapping indices.iter().map(|v|
*v as u32).collect() immediately before the call) and pass that u32 buffer's
pointer to XGDMatrixCreateFromCSREx, ensuring the indptr/indices types match the
expected C signature.

[coderabbitai]

⚠️ Potential issue | 🟡 Minor

Missing error handling for XGBoost FFI calls.

The return value of XGDMatrixCreateFromCSR is not checked. If the call fails, handle remains null, which could cause undefined behavior when passed to XGDMatrixFree. Consider checking the return code or at least asserting the handle is non-null after creation.

🛠️ Suggested fix

     unsafe {
-        xgboost_sys::XGDMatrixCreateFromCSR(
+        let ret = xgboost_sys::XGDMatrixCreateFromCSR(
             indptr_json.as_ptr(),
             indices_json.as_ptr(),
             data_json.as_ptr(),
             num_cols as xgboost_sys::bst_ulong,
             config.as_ptr(),
             &mut handle,
         );
+        assert!(ret == 0 && !handle.is_null(), "DMatrix creation failed");
     }

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	unsafe {
	xgboost_sys::XGDMatrixCreateFromCSR(
	indptr_json.as_ptr(),
	indices_json.as_ptr(),
	data_json.as_ptr(),
	num_cols as xgboost_sys::bst_ulong,
	config.as_ptr(),
	&mut handle,
	);
	}
	handle
	}
	unsafe {
	let ret = xgboost_sys::XGDMatrixCreateFromCSR(
	indptr_json.as_ptr(),
	indices_json.as_ptr(),
	data_json.as_ptr(),
	num_cols as xgboost_sys::bst_ulong,
	config.as_ptr(),
	&mut handle,
	);
	assert!(ret == 0 && !handle.is_null(), "DMatrix creation failed");
	}
	handle
	}

🤖 Prompt for AI Agents

In `@examples/bench_table.rs` around lines 75 - 86, The call to the FFI function
XGDMatrixCreateFromCSR currently ignores its return code and may leave handle
null; update the code to check the function's return value and verify handle is
non-null before returning or continuing. Specifically, capture the return code
from XGDMatrixCreateFromCSR, handle non-zero/error codes by returning an Err or
panicking with a clear message, and if the call appears to fail ensure you do
not call XGDMatrixFree on a null handle; alternatively return a Result from the
surrounding function so callers can handle failures instead of unconditionally
returning handle.

[coderabbitai]

⚠️ Potential issue | 🟠 Major

Function renamed per maintainer request not implemented.

The maintainer requested keeping the old load_binary function and adding a new load_binary_from_uri function since this uses a different C function (URI-based). The current implementation replaces the old behavior, which is a breaking change.

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

src/booster.rs (2)

244-253: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Store custom evals under the dataset key.

dmat_eval_results is dmat_name -> eval_name -> score everywhere else, but this block inserts custom metrics as custom -> dmat_name -> score. That makes callback data inconsistent and prints custom-train instead of train-custom.

💡 Proposed fix

                 if let Some(eval_fn) = params.custom_evaluation_fn {
                     let eval_name = "custom";
                     for (dmat, dmat_name) in eval_sets {
                         let margin = bst.predict_margin(dmat)?;
                         let eval_result = eval_fn(&margin, dmat);
                         let eval_results = dmat_eval_results
-                            .entry(eval_name.to_string())
+                            .entry(dmat_name.to_string())
                             .or_insert_with(IndexMap::new);
-                        eval_results.insert(dmat_name.to_string(), eval_result);
+                        eval_results.insert(eval_name.to_string(), eval_result);
                     }
                 }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/booster.rs` around lines 244 - 253, The custom evaluation block currently
inserts results under the evaluation-name key ("custom") causing the map shape
to be eval_name -> dmat_name -> score; change it to store under the dataset key
so the shape matches the rest (dmat_name -> eval_name -> score): inside the loop
over eval_sets (where params.custom_evaluation_fn, eval_sets,
bst.predict_margin, eval_result, dmat_name and dmat_eval_results are used)
replace the entry lookup that uses eval_name as the top-level key with one that
uses dmat_name.to_string() and then insert eval_name.to_string() -> eval_result
into that inner IndexMap so callbacks and prints become dataset-first (e.g.,
"train-custom" instead of "custom-train").

---

328-330: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Use raw margins for the custom objective callback.

XGBoost custom objectives are defined over raw predictions / leaf weights, not transformed outputs. The predict() method applies link function transformations, producing incorrect gradients and Hessians for objectives with a non-identity link. predict_margin() returns raw margins and is the correct source for objective_fn.

💡 Proposed fix

-        let pred = self.predict(dtrain)?;
-        let (gradient, hessian) = objective_fn(&pred.to_vec(), dtrain);
+        let pred = self.predict_margin(dtrain)?;
+        let (gradient, hessian) = objective_fn(&pred, dtrain);
         self.boost(dtrain, iteration, &gradient, &hessian)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/booster.rs` around lines 328 - 330, The code is using transformed
predictions from self.predict() to compute gradients/hessians for objective_fn,
but custom objectives must use raw margins; replace the call to
self.predict(...) with self.predict_margin(...) so objective_fn receives raw
predictions (margins) before boost is called (keep passing the resulting Vec to
objective_fn and then call self.boost(dtrain, iteration, &gradient, &hessian) as
before).

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@src/booster.rs`:
- Around line 244-253: The custom evaluation block currently inserts results
under the evaluation-name key ("custom") causing the map shape to be eval_name
-> dmat_name -> score; change it to store under the dataset key so the shape
matches the rest (dmat_name -> eval_name -> score): inside the loop over
eval_sets (where params.custom_evaluation_fn, eval_sets, bst.predict_margin,
eval_result, dmat_name and dmat_eval_results are used) replace the entry lookup
that uses eval_name as the top-level key with one that uses
dmat_name.to_string() and then insert eval_name.to_string() -> eval_result into
that inner IndexMap so callbacks and prints become dataset-first (e.g.,
"train-custom" instead of "custom-train").
- Around line 328-330: The code is using transformed predictions from
self.predict() to compute gradients/hessians for objective_fn, but custom
objectives must use raw margins; replace the call to self.predict(...) with
self.predict_margin(...) so objective_fn receives raw predictions (margins)
before boost is called (keep passing the resulting Vec to objective_fn and then
call self.boost(dtrain, iteration, &gradient, &hessian) as before).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 14d1c931-7390-4ba8-ab4a-274132a15916

📥 Commits

Reviewing files that changed from the base of the PR and between c732339 and 5f7d58e.

📒 Files selected for processing (1)

• src/booster.rs