[jvm-packages] Add opt-in sparse data optimization to XGBoost4J-Spark
- Dominant language
- C++
- Stars
- 28.8k
- Forks
- 8.9k
- Avg merge
- 1d 12h
- Merged PRs (30d)
- 54
Description
## Summary
XGBoost4J-Spark currently densifies every Spark ML `SparseVector` in both
training and inference. This preserves Spark's implicit-zero semantics, but it
makes memory and conversion cost scale with the feature dimension instead of
the number of nonzero entries.
I would like to propose an opt-in sparse path for XGBoost4J-Spark, equivalent
to Python Spark's `enable_sparse_data_optim`:
- disabled by default;
- valid only when `missing == 0.0f`;
- preserve `SparseVector.indices` and `SparseVector.values` in both training
and inference;
- require no change to the XGBoost core or model format.
## Current behavior
The training path converts a Spark ML vector with:
```scala
features.toArray.map(_.toFloat)
```
The inference path uses the same dense representation through `asXGB`:
```scala
new XGBLabeledPoint(0.0f, v.size, null, v.toArray.map(_.toFloat))
```
Therefore a row with feature dimension `D` and `K` nonzero entries requires
`O(D)` conversion and intermediate storage even when `K << D`.
For example, a Spark `SparseVector` with 50 million dimensions and only a few
nonzero entries creates hundreds of megabytes of dense JVM arrays for one row
before the native `DMatrix` is constructed. Reducing `inferBatchSize` lowers
the batch-level peak but does not remove the per-row `O(D)` cost.
## Relationship to #12347
#12347 correctly fixed malformed `XGBLabeledPoint` construction during batch
prediction. Before the fix, a sparse row paired `indices` of length `K` with
dense `values` of length `D`, corrupting the CSR representation. The merged
fix uses a valid dense point (`indices = null`, `values.length = D`) and
restores correct predictions.
This proposal does not change or revert that correctness fix. It introduces a
separate, explicitly constrained sparse mode in which `indices` and `values`
are parallel arrays of length `K`.
## Sparse semantics
There is an important semantic difference between the two representations:
- an omitted entry in Spark's `SparseVector` represents `0.0`;
- an omitted entry in an XGBoost CSR matrix represents the configured missing
value.
Consequently, preserving Spark sparsity is not generally correct when
`missing != 0.0f`. With `missing == 0.0f`, however, a zero in dense input is
also treated as missing by XGBoost. Omitting zero entries therefore does not
change how XGBoost interprets the row.
Python Spark already applies this constraint when
`enable_sparse_data_optim=True` and rejects configurations where
`missing != 0.0`.
## Proposed API and behavior
Add a JVM-only boolean XGBoost4J-Spark parameter named
`enableSparseDataOptim`, with default `false`, and register it as a
non-XGBoost parameter so it is not forwarded to the native core.
When it is disabled, preserve all existing behavior.
When it is enabled:
1. Require `missing == 0.0f` and fail early with a clear error otherwise.
2. Require the feature column to be a Spark ML `Vector` column.
3. Convert `SparseVector` without densification:
```scala
new XGBLabeledPoint(
label,
vector.size,
vector.indices,
vector.values.map(_.toFloat),
weight,
group,
baseMargin
)
```
4. Apply the same representation in `XGBoostEstimator.toXGBLabeledPoint`,
`XGBoostModel.transform`, and single-instance prediction.
5. Continue to represent `DenseVector` as dense input.
6. Persist the parameter with the Spark estimator/model metadata in the same
way as existing Spark parameters.
The existing `DataBatch` sparse branch already accepts a `LabeledPoint` whose
`indices` and `values` are parallel `K`-element arrays, so the initial
implementation should remain within the JVM/Spark binding.
## Why opt-in instead of automatic
Automatically switching representations when `missing == 0.0f` could change
memory behavior and make the input contract less explicit. An opt-in parameter
also matches the Python Spark API and preserves backward compatibility.
Preserving sparsity when `missing != 0.0f` is not proposed because it would
silently change implicit zeros into missing values. Lowering `inferBatchSize`
is not sufficient because each row is still densified.
## Tests
The change should include JVM tests covering:
- parameter default and validation (`missing != 0.0f` is rejected);
- sparse conversion keeps `indices.length == values.length == nnz`;
- sparse and equivalent dense rows produce matching predictions with
`missing == 0.0f`;
- training and batch inference use the same representation semantics;
- mixed dense and sparse rows remain supported;
- estimator/model save and load preserve the parameter;
- converting a very high-dimensional, low-`nnz` vector does not create dense
JVM input arrays proportional to the feature dimension.
If useful, I can also include a small benchmark comparing allocation and
runtime against the current dense path as feature dimension increases while
`nnz` remains fixed.
## Related work
- [Discussion #11467](https://github.com/dmlc/xgboost/discussions/11467):
memory issues caused by sparse-to-dense conversion in XGBoost4J-Spark
- [Issue #12404](https://github.com/dmlc/xgboost/issues/12404): corrupted CSR
for `SparseVector` batch transform
- [PR #12347](https://github.com/dmlc/xgboost/pull/12347): correctness fix for
`SparseVector` batch prediction
- [Python Spark `enable_sparse_data_optim`](https://github.com/dmlc/xgboost/blob/master/python-package/xgboost/spark/params.py)
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with XGBoostEstimator.toXGBLabeledPoint, XGBoostModel.transform, and single-instance prediction, then inspect the existing DataBatch sparse branch. Add and validate the JVM-only parameter while preserving dense behavior and Spark metadata persistence. Done means training and inference retain parallel sparse indices and values, validation rejects nonzero missing values, and the listed JVM tests cover mixed vectors, predictions, save/load, and high-dimensional inputs.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- java, scala
- Domain
- data-engineering, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100