Make base_tree_regressor faster

[slievens]

Fixes #

Checklist

• [x ] My branch is up-to-date with main branch.
• [ x] Everything works and tested on latest stable Rust.
• [ ]x Coverage and Linting have been applied

Current behaviour

BaseTreeRegressor (used by RandomForestRegressor and ExtraTreesRegressor) could be faster during fit.

• Each tree sorts all feature columns again. The forest sorts them n_trees times.
• Each node has a samples: Vec<usize> with one entry per row of x. To split a node, the code makes a new vector of length n_rows. A node deep in the tree with few samples still costs O(n_rows).
• The split search reads order[j] for all rows of x and skips rows with a zero count.
• The code allocates a new list of features at each node, and uses a LinkedList as the queue of nodes.

New expected behaviour

The tree gives statistically equivalent results, but fitting is faster. The memory use for each tree is O(n_features × n_rows_in_bag).

• The forest sorts the columns one time and gives the order to each tree.
• At the start, each tree makes one sorted array for each feature. The arrays contain only the rows in the bootstrap sample. A node is a range start_idx..end_idx in these arrays. A split partitions the range of each column with a stable, branchless partition. Thus each column stays sorted in each child, and the work at each node is proportional to the number of samples in the node.
• The elements in the arrays have the smallest possible size for the data:
  • UnitElement (4 bytes): no weights, each row at most one time (for example, no bootstrap).
  • CountedElement (8 bytes): no weights, rows can repeat (bootstrap).
  • WeightedElement (16 bytes): with sample weights.
• The code does not partition a range when both children cannot split (because of max_depth or min_samples_split). It does not search for a split in a child that cannot split.
• To find the boundary between the two children, split reads only the sorted column of the split feature, and stops at the first false sample.

Behaviour that is visible to users:

• With a fixed seed, trees that use randomness are not equal to the trees from main. This includes mtry < n_features, Splitter::Random, and bootstrap. The nodes are now processed depth-first (stack) instead of breadth-first (queue), so the code uses the RNG in a different order. Also, children that cannot split do not use the RNG now. The statistical behaviour does not change.
• The split threshold of Splitter::Best is now (a.to_f64() + b.to_f64()) / 2 instead of (a + b).to_f64() / 2. This prevents overflow for integer TX and loss of precision for f32.

Change logs

No changes to the public API.

Changed

• BaseForestRegressor::fit sorts the feature columns one time for all trees.
• BaseTreeRegressor keeps one sorted array for each feature, and a node is a range in he vector of length n_rows for each node.
• The nodes are processed with a Vec as a stack. This replaces the LinkedList queue.
• The list of features to shuffle is allocated one time for each tree, not at each node.
• Splitter now derives Copy.

Added

• Test test_node_element_sizes, which makes sure that the element types stay small.

[codecov]

Codecov Report

❌ Patch coverage is 99.56332% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 85.31%. Comparing base (9eaae9e) to head (7d2b9ab).
⚠️ Report is 197 commits behind head on main.

Files with missing lines	Patch %	Lines
src/tree/base_tree_regressor.rs	99.55%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #481       +/-   ##
===========================================
+ Coverage   43.97%   85.31%   +41.33%     
===========================================
  Files          85       96       +11     
  Lines        7281     9368     +2087     
===========================================
+ Hits         3202     7992     +4790     
+ Misses       4079     1376     -2703     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Mec-iS]

Thanks @slievens, this is a well-motivated change. Sorting once per forest, using range-based nodes with stable partitioning, and switching to a DFS stack are all sound. I read the full diff (2 files, +420/-174). Findings below, roughly by priority.

Correctness / robustness

1. find_random_split lost its early break. The old loop stopped at the first sample with x > split_value. The new loop scans the whole start_idx..end_idx range and tests every element. The result is the same because the column is sorted, but it costs O(node) on every random split. Restore the break, or use partition_point/binary search on the sorted column.
2. Silent u32 truncation. row_idx as u32 and count as u32 (in WeightedElement, CountedElement, UnitElement and counts: Vec<u32>) wrap silently when n_rows > u32::MAX. For CountedElement, count << 32 also loses data if a count exceeds u32. Please add a check in fit that returns Failed, or at least a debug_assert!.
3. Implicit invariant on grow::<E>. CountedElement::mass() returns count and ignores weights, and UnitElement::mass() returns 1.0. This is only correct because fit_weak_learner dispatches Weighted whenever weights are Some. grow is generic, so add debug_assert!(sample_weights.is_none()) for the two non-weighted types, and a short doc comment on NodeElement::mass.
4. Midpoint threshold for large integers. The new (a.to_f64()+b.to_f64())/2 fixes overflow and f32 precision. For i64/u64 values above 2^53, to_f64() can collapse neighbouring values so that the threshold equals b. Then split (<= split_value) and find_best_split would disagree about the partition boundary. This is an edge case, but worth a test or a note. Also confirm that NaN in x is still not supported, since the partition logic assumes a total order.
5. sorted_by_feature[0] panics if there are zero features. Please verify that fit rejects num_attributes == 0, and the same for the forest path.

Performance / design

6. No benchmark numbers in the PR. Please add before/after timings for fit (e.g. 1k/10k/100k rows, 10/100 features, bootstrap on/off, Best vs Random, with/without weights) and peak memory. Ideally add a Criterion bench to benches/.
7. Cache locality in the split scan. x.get((row, j)) is still a random access into X for every element. Storing the feature value next to the row index would make scans sequential at the cost of larger elements. It may be worth a measured experiment, since the 4/8/16-byte sizing is a deliberate trade-off.
8. find_best_cutoff recomputes n by summing count() over the range. split already has true_count and false_count, so pass n in. That also removes one pass per node.
9. Per-tree allocations. Each tree allocates n_features × n_kept elements, plus partition_buffer and in_true_branch. This is fine, but it could be reused across trees in a forest, or built in parallel if rayon is used later. Please document the peak memory (~ n_features × n_kept × size_of::<E>() per tree in flight).
10. #[inline(never)] on stable_partition should carry a comment explaining why, for example a measurement, or the intent to keep it as a profiling boundary. The scratch[f] = e write still has bounds checks, so check the generated code before calling it branchless.

Tests / coverage

11. Codecov reports 69.95% patch coverage (64 lines missing), almost all in base_tree_regressor.rs. CountedElement, WeightedElement, Splitter::Random, and the early-leaf exits in split are likely under-tested.
12. test_node_element_sizes only checks sizes. Please add an equivalence test against a simple reference implementation (naive O(n²) best split, or the old code kept under #[cfg(test)]). It should compare with Splitter::Best, mtry == n_features, no bootstrap, and both weighted and unweighted data. The tree structure and predictions should be identical when there is no randomness. For randomised paths, assert statistical equivalence (e.g. OOB or test MSE within tolerance across seeds).
13. Add tests for max_depth and min_samples_split boundaries (the new "both children are leaves, so skip the partition" shortcut), min_samples_leaf rollback, duplicate feature values, constant features, and a column where all kept rows are tied.
14. size_of without an import relies on the prelude addition in Rust 1.80. Please check this against the crate's MSRV, or use std::mem::size_of.

Process / docs

15. The branch is 196 commits behind main and the PR is in a blocked state. Please rebase and rerun CI.
16. The checklist is malformed ([x ], [ x], [ ]x), and Fixes # is empty. Please fix both and add a CHANGELOG entry.
17. The fixed-seed behaviour change (DFS order, no RNG use for non-splittable children) is documented in the PR body. It should also go into the CHANGELOG, since users reproducing results across versions will see different trees.

Overall this looks promising. The main asks are the equivalence/coverage tests, benchmark evidence, the lost break, and the integer-cast guards. @Mec-iS, happy to defer to your judgement on MSRV and the checklist.

[Mec-iS]

Please use smartcore-benches to time the classes before and after the changes

[slievens]

I tried to address some of the concerns in commit slievens@d071328

[slievens]

I do not know how to show the results of the benchmark. Here is the output of running the benchmark on my machine, comparing against the version in the repo. The new version is always faster.

RandomForestRegressor::fit/n_samples: 100, n_features: 10, n_trees: 10
time: [28.880 µs 28.900 µs 28.922 µs]
change: [−70.208% −70.162% −70.117%] (p = 0.00 < 0.05)
Performance has improved.
Found 7 outliers among 100 measurements (7.00%)
4 (4.00%) high mild
3 (3.00%) high severe
RandomForestRegressor::fit/n_samples: 100, n_features: 10, n_trees: 50
time: [115.99 µs 116.12 µs 116.31 µs]
change: [−76.635% −76.584% −76.537%] (p = 0.00 < 0.05)
Performance has improved.
Found 2 outliers among 100 measurements (2.00%)
2 (2.00%) high severe
RandomForestRegressor::fit/n_samples: 100, n_features: 10, n_trees: 100
time: [226.76 µs 226.98 µs 227.22 µs]
change: [−77.144% −77.070% −76.987%] (p = 0.00 < 0.05)
Performance has improved.
Found 9 outliers among 100 measurements (9.00%)
2 (2.00%) high mild
7 (7.00%) high severe
RandomForestRegressor::fit/n_samples: 1000, n_features: 10, n_trees: 10
time: [238.56 µs 238.98 µs 239.54 µs]
change: [−74.612% −74.573% −74.529%] (p = 0.00 < 0.05)
Performance has improved.
Found 11 outliers among 100 measurements (11.00%)
4 (4.00%) high mild
7 (7.00%) high severe
RandomForestRegressor::fit/n_samples: 1000, n_features: 10, n_trees: 50
time: [875.86 µs 877.32 µs 879.11 µs]
change: [−81.383% −81.358% −81.328%] (p = 0.00 < 0.05)
Performance has improved.
Found 16 outliers among 100 measurements (16.00%)
4 (4.00%) high mild
12 (12.00%) high severe
RandomForestRegressor::fit/n_samples: 1000, n_features: 10, n_trees: 100
time: [1.6731 ms 1.6743 ms 1.6758 ms]
change: [−82.024% −82.001% −81.981%] (p = 0.00 < 0.05)
Performance has improved.
Found 7 outliers among 100 measurements (7.00%)
5 (5.00%) high mild
2 (2.00%) high severe
RandomForestRegressor::fit/n_samples: 1000, n_features: 50, n_trees: 10
time: [971.21 µs 972.45 µs 973.84 µs]
change: [−76.802% −76.757% −76.716%] (p = 0.00 < 0.05)
Performance has improved.
Found 9 outliers among 100 measurements (9.00%)
7 (7.00%) high mild
2 (2.00%) high severe
RandomForestRegressor::fit/n_samples: 1000, n_features: 50, n_trees: 50
time: [3.2175 ms 3.2204 ms 3.2235 ms]
change: [−84.541% −84.521% −84.498%] (p = 0.00 < 0.05)
Performance has improved.
Found 10 outliers among 100 measurements (10.00%)
7 (7.00%) high mild
3 (3.00%) high severe
RandomForestRegressor::fit/n_samples: 1000, n_features: 50, n_trees: 100
time: [6.0852 ms 6.0901 ms 6.0957 ms]
change: [−85.376% −85.335% −85.300%] (p = 0.00 < 0.05)
Performance has improved.
Found 11 outliers among 100 measurements (11.00%)
4 (4.00%) high mild
7 (7.00%) high severe

[Mec-iS]

thanks @slievens I have approved but:

• Boundary tests: none for max_depth, min_samples_split and min_samples_leaf, for duplicate or constant features, or for the weighted and counted element paths. Patch coverage is still about 70%. Codecov's patch check passes anyway.
• Invariant guards: there's no debug_assert!(sample_weights.is_none()) for CountedElement and UnitElement, whose mass() ignores weights.
• Large-integer thresholds: the midpoint threshold for i64/u64 values above 2^53, and NaN handling in x, are still unaddressed.
  Redundant sum: find_best_cutoff still recomputes n by summing counts, although split already has true_count and false_count.
• Benchmark scope: the results cover only RandomForestRegressor, at 1000 samples or fewer. They don't cover ExtraTrees, Splitter::Random, weighted fits, bootstrap off, larger data, or memory. Mec-iS asked for smartcore-benches, and I can't tell which harness was used.

You can go on and finish this or I will have a pass at those.

[slievens]

• "Boundary tests: none for max_depth, min_samples_split and min_samples_leaf, for duplicate or constant features, or for the weighted and counted element paths. Patch coverage is still about 70%. Codecov's patch check passes anyway."
  I will try to add additional tests for this.
• Invariant guards: there's no debug_assert!(sample_weights.is_none()) for CountedElement and UnitElement, whose mass() ignores weights.
  I will add debug_assert! statements which guard against wrong use.
• Large-integer thresholds: the midpoint threshold for i64/u64 values above 2^53, and NaN handling in x, are still unaddressed.
  I will look into this.
• Redundant sum: find_best_cutoff still recomputes n by summing counts, although split already has true_count and false_count.
  I will not address this now. My plan is to refactor the implementation to be generic over the criterion so that also the classifier can use is.
• Benchmark scope: the results cover only RandomForestRegressor, at 1000 samples or fewer. They don't cover ExtraTrees, Splitter::Random, weighted fits, bootstrap off, larger data, or memory. Mec-iS asked for smartcore-benches, and I can't tell which harness was used.
  I will add additional benchmarks which covers additional variations. I am not sure how large these can/should be.

[Mec-iS]

I will add additional benchmarks which covers additional variations. I am not sure how large these can/should be.

benchmarks go in smartcore-benches. Please follow the conventions there.

[slievens]

I added additional tests in 0b0ac49 which should hopefully improve code coverage. I think/hope this PR can now be merged.

[Mec-iS]

Thank you!