sparklightgbm 0.9.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,9 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ .pytest_cache/
5
+ .ruff_cache/
6
+ .venv/
7
+ .release-venv/
8
+ build/
9
+ dist/
@@ -0,0 +1,16 @@
1
+ # Changelog
2
+
3
+ ## 0.9.0 - 2026-09-22
4
+
5
+ - Prepare metadata, documentation, source distribution, wheel, and CI for a public PyPI release.
6
+ - Define the supported Python 3.9-3.12, PySpark 3.4-3.5 compatibility policy.
7
+ - Make `num_workers=None` the documented default and conservatively cap automatic cluster training at four workers while reserving one reported Spark slot.
8
+ - Keep validation shards distributed during native multi-worker training.
9
+ - Synchronize decomposable validation metrics across barrier workers for consistent distributed early stopping.
10
+ - Batch inference by Spark partition and reuse native boosters in Python workers.
11
+ - Make distributed objective inference globally consistent and reduce model artifacts returned to the driver.
12
+ - Add library-specific configuration, network, and worker errors plus a reproducible benchmark harness.
13
+
14
+ ## 0.1.0b1 - 2026-09-20
15
+
16
+ Beta release for objective runtime testing. Includes Spark DataFrame estimators for classification, regression, quantile regression, and LambdaRank; native predictions, SHAP contributions, feature importance, early stopping in single-worker mode, categorical values, weights, missing values, seeds, native model persistence, and experimental native distributed training through Spark barriers.
@@ -0,0 +1,3 @@
1
+ # Contributing
2
+
3
+ Create a focused branch, install `.[dev]`, and run `pytest` before opening a pull request. Public API changes should include documentation and a Spark-local integration test. Keep LightGBM calls in the execution layer and do not introduce SynapseML or copied implementation code. Contributions must be compatible with the MIT license.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Celso Renato Quiroz Valle
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,293 @@
1
+ Metadata-Version: 2.5
2
+ Name: sparklightgbm
3
+ Version: 0.9.0
4
+ Summary: A lightweight bridge between Apache Spark DataFrames and official LightGBM.
5
+ Project-URL: Homepage, https://github.com/Celso-RQ-Valle/SparkLightGBM
6
+ Project-URL: Repository, https://github.com/Celso-RQ-Valle/SparkLightGBM
7
+ Project-URL: Documentation, https://github.com/Celso-RQ-Valle/SparkLightGBM#readme
8
+ Project-URL: Issues, https://github.com/Celso-RQ-Valle/SparkLightGBM/issues
9
+ Project-URL: Changelog, https://github.com/Celso-RQ-Valle/SparkLightGBM/blob/main/CHANGELOG.md
10
+ Author: SparkLightGBM contributors
11
+ Maintainer: SparkLightGBM maintainers
12
+ License-Expression: MIT
13
+ License-File: LICENSE
14
+ Keywords: apache-spark,distributed-training,lightgbm,machine-learning,pyspark
15
+ Classifier: Intended Audience :: Developers
16
+ Classifier: Intended Audience :: Science/Research
17
+ Classifier: License :: OSI Approved :: MIT License
18
+ Classifier: Operating System :: OS Independent
19
+ Classifier: Programming Language :: Python :: 3
20
+ Classifier: Programming Language :: Python :: 3 :: Only
21
+ Classifier: Programming Language :: Python :: 3.9
22
+ Classifier: Programming Language :: Python :: 3.10
23
+ Classifier: Programming Language :: Python :: 3.11
24
+ Classifier: Programming Language :: Python :: 3.12
25
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
26
+ Requires-Python: <3.13,>=3.9
27
+ Requires-Dist: lightgbm>=4.0
28
+ Requires-Dist: numpy>=1.21
29
+ Provides-Extra: dev
30
+ Requires-Dist: build>=1; extra == 'dev'
31
+ Requires-Dist: pyspark<4,>=3.4; extra == 'dev'
32
+ Requires-Dist: pytest>=7; extra == 'dev'
33
+ Requires-Dist: ruff>=0.6; extra == 'dev'
34
+ Requires-Dist: twine>=5; extra == 'dev'
35
+ Provides-Extra: spark
36
+ Requires-Dist: pyspark<4,>=3.4; extra == 'spark'
37
+ Provides-Extra: test
38
+ Requires-Dist: pyspark<4,>=3.4; extra == 'test'
39
+ Requires-Dist: pytest>=7; extra == 'test'
40
+ Description-Content-Type: text/markdown
41
+
42
+ # SparkLightGBM
43
+
44
+ SparkLightGBM is a lightweight bridge between Apache Spark DataFrames and the official [LightGBM](https://github.com/lightgbm-org/LightGBM) Python package. It supports native LightGBM training across Spark executors without requiring SynapseML or a Scala/JVM extension.
45
+
46
+ The project aims for efficient training and inference, numerical correctness, distributed scalability, and production reliability while keeping installation simple, dependencies minimal, and behavior portable across Spark environments. It does not aim to reproduce every feature of larger Spark integrations. Simplicity, performance, compatibility, and predictable behavior take priority over feature count.
47
+
48
+ ## Status
49
+
50
+ `0.9.x` is the pre-1.0 real-world validation series. It is intended for practical evaluation, but the public API may still change based on results across different environments, Spark configurations, datasets, and workloads. Version `1.0.0` is reserved for the stable release after that validation.
51
+
52
+ ## Installation
53
+
54
+ ```bash
55
+ pip install "sparklightgbm[spark]"
56
+ ```
57
+
58
+ Use `pip install sparklightgbm` when PySpark is already supplied by a managed Spark environment. The supported compatibility policy for `0.9.x` is Python 3.9-3.12, PySpark 3.4-3.5, LightGBM 4.x or newer, and NumPy 1.21 or newer. PySpark, LightGBM, NumPy, and this package must be available on the driver and every executor. Linux and Windows are exercised in CI; distributed multi-node training is primarily expected on Linux clusters.
59
+
60
+ The optional development installation is:
61
+
62
+ ```bash
63
+ pip install "sparklightgbm[dev]"
64
+ ```
65
+
66
+ SparkLightGBM never calls `toPandas()` and does not reimplement LightGBM. CPU is the default. There is no required SynapseML dependency.
67
+
68
+ ## Quick start
69
+
70
+ The `features_col` value must identify a Spark column containing a numeric feature vector or a numeric array/list. Spark ML `Vector` values and ordinary Spark arrays are accepted. The label column must contain numeric labels or targets.
71
+
72
+ ```python
73
+ from sparklightgbm import LightGBMClassifier
74
+
75
+ estimator = LightGBMClassifier(
76
+ features_col="features",
77
+ label_col="label",
78
+ num_workers=None,
79
+ num_iterations=300,
80
+ learning_rate=0.05,
81
+ num_leaves=31,
82
+ max_depth=-1,
83
+ min_data_in_leaf=20,
84
+ feature_fraction=0.9,
85
+ bagging_fraction=0.8,
86
+ bagging_freq=1,
87
+ lambda_l1=0.0,
88
+ lambda_l2=1.0,
89
+ min_gain_to_split=0.0,
90
+ max_bin=255,
91
+ seed=7,
92
+ feature_fraction_seed=7,
93
+ bagging_seed=7,
94
+ data_random_seed=7,
95
+ )
96
+
97
+ model = estimator.fit(train_df)
98
+ predictions = model.transform(test_df)
99
+ predictions.select("prediction", "probability").show()
100
+ ```
101
+
102
+ ## Estimators and input parameters
103
+
104
+ All three estimators accept the shared parameters below. Additional keyword arguments are forwarded to native LightGBM, allowing supported LightGBM parameters to be used without waiting for a wrapper release.
105
+
106
+ | Parameter | Accepted type | Default | Description |
107
+ | --- | --- | --- | --- |
108
+ | `features_col` | `str` | `"features"` | Spark column containing a numeric vector or array. |
109
+ | `label_col` | `str` | `"label"` | Spark column containing the target. |
110
+ | `prediction_col` | `str` | `"prediction"` | Output column for the predicted value or class. |
111
+ | `raw_prediction_col` | `str \| None` | `"rawPrediction"` | Output column for raw LightGBM scores; `None` omits it. |
112
+ | `probability_col` | `str \| None` | `"probability"` | Classifier probability output column; `None` omits it. |
113
+ | `leaf_prediction_col` | `str \| None` | `"leafPrediction"` | Output column containing integer leaf indices; `None` omits it. |
114
+ | `weight_col` | `str \| None` | `None` | Optional Spark column containing per-row weights. |
115
+ | `group_col` | `str \| None` | `None` | Optional Spark column containing ranking group identifiers; required by `LightGBMRanker`. |
116
+ | `categorical_feature` | `list[int] \| str \| None` | `None` | LightGBM categorical feature indices or native categorical setting. Feature values must be numeric category codes. |
117
+ | `validation_data` | Spark `DataFrame` \| `None` | `None` | Validation DataFrame with matching feature, label, weight, and group names. Can also be passed to `fit()`. |
118
+ | `early_stopping_rounds` | `int \| None` | `None` | Rounds without validation improvement; requires validation data. Distributed mode supports exactly aggregatable metrics. |
119
+ | `seed` | `int` | `0` | Seed passed to LightGBM. |
120
+ | `num_workers` | `int \| None` | `None` | Automatically selects a conservative worker count. Local Spark and ranking use driver training. Clustered classification and regression use the smallest of four workers, the input partition count, and available Spark slots minus one. Set a positive integer to override the strategy. |
121
+ | `local_listen_port` | `int` | `12400` | Distributed listener base port; worker `n` uses `local_listen_port + n`. |
122
+ | `prediction_batch_size` | `int` | `1024` | Rows scored per native LightGBM prediction batch in each Spark partition. |
123
+ | `objective` | `str \| None` | `None` | Native objective: defaults to `binary`/`multiclass`, `regression`, or `lambdarank`; `quantile` is supported for regression. |
124
+ | `**params` | Native value | - | Parameters forwarded to `lightgbm.Dataset` or `lightgbm.train`, as applicable. |
125
+
126
+ Estimator-specific parameters:
127
+
128
+ | Estimator | Parameter | Accepted type / default | Description |
129
+ | --- | --- | --- | --- |
130
+ | `LightGBMClassifier` | `num_class` | `int \| None`, default `None` | Number of classes; if omitted, inferred from labels. |
131
+ | `LightGBMRegressor` | - | - | Supports native regression parameters and `objective="quantile"`; set native `alpha` for the target quantile. |
132
+ | `LightGBMRanker` | - | - | Uses LambdaRank by default. `group_col` identifies groups; rows are sorted by group before native training. |
133
+
134
+ ## Common native LightGBM parameters
135
+
136
+ The following are commonly used native parameters, not a separate SparkLightGBM parameter system. Their appropriate values depend on the data, objective, validation strategy, and resource constraints.
137
+
138
+ | Parameter | What it controls |
139
+ | --- | --- |
140
+ | `num_iterations` | Maximum number of boosting rounds. `n_estimators` and `num_boost_round` are also accepted by SparkLightGBM. |
141
+ | `learning_rate` | Contribution of each new tree; interacts with the number of iterations. |
142
+ | `num_leaves` | Maximum leaves per tree and therefore much of the model's capacity. |
143
+ | `max_depth` | Optional tree-depth limit; negative values leave depth unconstrained. |
144
+ | `min_data_in_leaf` | Minimum observations allowed in a leaf, controlling leaf granularity and regularization. |
145
+ | `feature_fraction` | Fraction of features considered for each tree. |
146
+ | `bagging_fraction` | Fraction of rows used when bagging is active. |
147
+ | `bagging_freq` | Frequency of bagging; `0` disables it. |
148
+ | `lambda_l1` | L1 regularization applied to leaf weights. |
149
+ | `lambda_l2` | L2 regularization applied to leaf weights. |
150
+ | `min_gain_to_split` | Minimum gain required to create a split. |
151
+ | `max_bin` | Maximum histogram bins used for numeric features; affects accuracy, memory, and speed. |
152
+ | `is_unbalance` | Enables automatic binary-class imbalance handling. Do not combine it with `scale_pos_weight`. |
153
+ | `scale_pos_weight` | Explicit positive-class weight for binary classification. Do not combine it with `is_unbalance`. |
154
+ | `seed` | Top-level seed used by SparkLightGBM and passed to LightGBM. |
155
+ | `data_random_seed` | Seed used while constructing histogram bins. |
156
+ | `feature_fraction_seed` | Seed used for feature subsampling. |
157
+ | `bagging_seed` | Seed used for row subsampling. |
158
+ | `drop_seed` | Seed used by DART boosting. |
159
+ | `deterministic` | Requests stable CPU results; LightGBM may require related parameters for fully reproducible runs. |
160
+
161
+ `fit(params={...})` can supply or override native parameters for a training call:
162
+
163
+ ```python
164
+ from sparklightgbm import LightGBMRegressor
165
+
166
+ model = LightGBMRegressor(features_col="features", label_col="target").fit(
167
+ train_df,
168
+ params={"objective": "quantile", "alpha": 0.9, "n_estimators": 200, "learning_rate": 0.03},
169
+ )
170
+ ```
171
+
172
+ SparkLightGBM documents only common parameters. Availability, aliases, interactions, and exact semantics are determined by the installed LightGBM version; see the [official LightGBM parameter documentation](https://lightgbm.readthedocs.io/en/latest/Parameters.html) for the complete reference. SparkLightGBM uses `lightgbm.Dataset` and `lightgbm.train`, not the scikit-learn wrapper.
173
+
174
+ The public training signature is `fit(dataset, params=None, validation_data=None)`: `dataset` and `validation_data` are Spark `DataFrame` objects, and `params` is a `dict[str, object]` of native LightGBM parameter overrides. A `validation_data` argument passed to `fit()` takes precedence over the constructor value.
175
+
176
+ ## Spark input and output
177
+
178
+ Training reads rows from Spark partitions. Supported training columns are:
179
+
180
+ - `features_col`: numeric Spark ML vector, numeric array, or numeric Python list.
181
+ - `label_col`: numeric binary/multiclass label, regression target, or ranking relevance score.
182
+ - `weight_col`: optional numeric row weight.
183
+ - `group_col`: optional numeric/integer ranking group identifier.
184
+
185
+ Missing feature values should use LightGBM-compatible values such as `NaN`. A null feature vector is skipped; a null label or malformed feature value raises an error during collection or native training.
186
+
187
+ `model.transform(df)` preserves the input DataFrame and adds configured output columns. It creates `prediction` for every estimator, raw scores when enabled, leaf indices when enabled, and probabilities for classifiers when enabled. Output names can be overridden per call:
188
+
189
+ The transform signature is `transform(dataset, prediction_col=None, raw_prediction_col=None, probability_col=None, leaf_prediction_col=None)`. Every column-name override is a `str` or `None`; `probability_col` is used only for classifier models.
190
+
191
+ ```python
192
+ predictions = model.transform(
193
+ test_df,
194
+ prediction_col="score",
195
+ raw_prediction_col="raw_score",
196
+ probability_col="class_probability",
197
+ leaf_prediction_col="leaf_index",
198
+ )
199
+ ```
200
+
201
+ ## Validation and early stopping
202
+
203
+ Use a Spark DataFrame with matching column names:
204
+
205
+ ```python
206
+ model = LightGBMRegressor(n_estimators=500, early_stopping_rounds=30).fit(
207
+ train_df,
208
+ validation_data=validation_df,
209
+ )
210
+ ```
211
+
212
+ Validation data remains partitioned in distributed mode. Each worker builds local training and validation datasets, while Spark barrier synchronization aggregates decomposable native metrics before the early-stopping callback. Metrics whose exact global value cannot be reconstructed from shard-level scalar results (currently `auc`, `average_precision`, `map`, and `ndcg`) are rejected for distributed early stopping; use `num_workers=1` for those metrics.
213
+
214
+ ## Local and distributed execution
215
+
216
+ The worker setting changes where training data is materialized:
217
+
218
+ | Setting | Execution and memory behavior | Intended use |
219
+ | --- | --- | --- |
220
+ | `num_workers=1` | Spark partitions are read and the complete training dataset—and validation dataset, when present—is collected into NumPy arrays on the driver. Native LightGBM trains in the driver process. | Local development, compatibility fallback, ranking, and datasets that safely fit in driver memory. |
221
+ | `num_workers>1` | Spark repartitions the data into a barrier stage. Each executor-side worker converts only its shard to contiguous NumPy buffers and participates in LightGBM's native `data_parallel` network. Training and validation rows are not collected to the driver; the driver receives bounded metadata and the trained model artifact. | Cluster datasets that should remain distributed. Each worker shard must fit in that executor's memory. |
222
+ | `num_workers=None` | Selects `1` for local Spark and ranking. On a cluster, classification and regression reserve one slot and select the smallest of `4`, the remaining Spark parallelism, and the input partition count. The result is always at least `1`. | Portable, conservative default that avoids claiming every reported slot. Set an explicit value when cluster capacity or scheduling calls for another count. |
223
+
224
+ `SparkContext.defaultParallelism` is an estimate, not a dynamic cluster-capacity reservation. On shared or autoscaling clusters, set `num_workers` explicitly when the scheduler policy requires a specific limit. An explicit value always takes precedence, including in local mode; the requested barrier tasks must be schedulable concurrently.
225
+
226
+ Distributed workers default to one native LightGBM thread per Spark task to avoid CPU oversubscription. Pass `num_threads` explicitly when the Spark resource configuration provides additional CPU capacity per task.
227
+
228
+ Distributed execution requires:
229
+
230
+ - SparkLightGBM, NumPy, and the same compatible LightGBM build on every executor.
231
+ - Spark barrier execution support.
232
+ - Executor hostnames or IP addresses resolvable and reachable from every other executor.
233
+ - Listener ports available between executors; set `local_listen_port` if `12400 + worker_id` is unavailable.
234
+ - Correct executor networking. `SPARK_LOCAL_IP` can provide the advertised worker address when hostname resolution is unsuitable.
235
+
236
+ Inference uses partition-level NumPy batches instead of row-wise Python UDFs. A native booster is cached per reused Python worker, and classification prediction, probability, and raw prediction are derived from the same raw-score batch.
237
+
238
+ ## Current limitations
239
+
240
+ - The `0.9.x` API is in pre-1.0 validation and may change before `1.0.0`.
241
+ - Distributed ranking is not yet supported because query groups must remain complete and worker-local. Ranking therefore uses the single-worker path.
242
+ - Distributed early stopping supports metrics that can be exactly aggregated from worker-level results. Non-decomposable metrics currently rejected for this mode include `auc`, `average_precision`, `map`, and `ndcg`; use `num_workers=1` when early stopping depends on them.
243
+ - Distributed jobs require Spark barrier scheduling plus stable, mutually reachable executor addresses and ports. Executor loss aborts the coordinated native training job.
244
+ - Models use native LightGBM persistence rather than Spark ML `MLWriter`/`MLReader`, and the estimators are not yet Spark ML `Estimator`/`Model` stages for `Pipeline` or `CrossValidator`.
245
+ - CPU is the supported default execution path; GPU execution is not currently documented or tested by this project.
246
+
247
+ ## Predictions and explainability
248
+
249
+ Fitted models expose native LightGBM outputs:
250
+
251
+ ```python
252
+ raw = model.predict_raw([[1.0, 2.0]])
253
+ probability = model.predict_probability([[1.0, 2.0]]) # classifiers
254
+ leaves = model.predict_leaf([[1.0, 2.0]])
255
+ shap = model.predict_shap([[1.0, 2.0]])
256
+ split_importance = model.feature_importance("split")
257
+ gain_importance = model.feature_importance("gain")
258
+ ```
259
+
260
+ `predict_shap()` requests native LightGBM contribution values. The output includes one value per feature plus the expected-value contribution. `feature_importance()` accepts `"split"` or `"gain"`.
261
+
262
+ ## Native model persistence
263
+
264
+ Models are saved in LightGBM's native format. A JSON sidecar containing bridge metadata is written next to the model file:
265
+
266
+ ```python
267
+ model.save_native_model("artifacts/model.txt")
268
+
269
+ from sparklightgbm import LightGBMClassificationModel
270
+ restored = LightGBMClassificationModel.load_native_model(
271
+ "artifacts/model.txt",
272
+ features_col="features",
273
+ )
274
+ ```
275
+
276
+ The native model file is portable across Spark jobs with a compatible LightGBM installation. It is not a Spark ML `PipelineModel` format.
277
+
278
+ `save_native_model(path)` accepts a filesystem path (`str` or `pathlib.Path`). `LightGBMClassificationModel.load_native_model(path, **kwargs)`, `LightGBMRegressionModel.load_native_model(path, **kwargs)`, and `LightGBMRankingModel.load_native_model(path, **kwargs)` accept the native model path plus optional `features_col`, `prediction_col`, `raw_prediction_col`, `probability_col`, `leaf_prediction_col`, and `kind` strings used to reconstruct Spark output behavior.
279
+
280
+ ## Compatibility and errors
281
+
282
+ SparkLightGBM validates required columns and checks that supported PySpark and LightGBM installations are available on the driver. The `0.9.x` compatibility policy is Python 3.9-3.12 and PySpark 3.4-3.5; newer combinations are not claimed until validated. Runtime imports also need to be available on executors. Missing dependencies, missing columns, missing ranking groups, unsupported distributed validation settings, and invalid native parameters fail with errors from the bridge or LightGBM.
283
+
284
+ ## Development
285
+
286
+ ```bash
287
+ pip install -e ".[dev]"
288
+ python -m pytest -q
289
+ python -m build
290
+ python -m twine check dist/*
291
+ ```
292
+
293
+ See [CHANGELOG.md](CHANGELOG.md) and [CONTRIBUTING.md](CONTRIBUTING.md).
@@ -0,0 +1,252 @@
1
+ # SparkLightGBM
2
+
3
+ SparkLightGBM is a lightweight bridge between Apache Spark DataFrames and the official [LightGBM](https://github.com/lightgbm-org/LightGBM) Python package. It supports native LightGBM training across Spark executors without requiring SynapseML or a Scala/JVM extension.
4
+
5
+ The project aims for efficient training and inference, numerical correctness, distributed scalability, and production reliability while keeping installation simple, dependencies minimal, and behavior portable across Spark environments. It does not aim to reproduce every feature of larger Spark integrations. Simplicity, performance, compatibility, and predictable behavior take priority over feature count.
6
+
7
+ ## Status
8
+
9
+ `0.9.x` is the pre-1.0 real-world validation series. It is intended for practical evaluation, but the public API may still change based on results across different environments, Spark configurations, datasets, and workloads. Version `1.0.0` is reserved for the stable release after that validation.
10
+
11
+ ## Installation
12
+
13
+ ```bash
14
+ pip install "sparklightgbm[spark]"
15
+ ```
16
+
17
+ Use `pip install sparklightgbm` when PySpark is already supplied by a managed Spark environment. The supported compatibility policy for `0.9.x` is Python 3.9-3.12, PySpark 3.4-3.5, LightGBM 4.x or newer, and NumPy 1.21 or newer. PySpark, LightGBM, NumPy, and this package must be available on the driver and every executor. Linux and Windows are exercised in CI; distributed multi-node training is primarily expected on Linux clusters.
18
+
19
+ The optional development installation is:
20
+
21
+ ```bash
22
+ pip install "sparklightgbm[dev]"
23
+ ```
24
+
25
+ SparkLightGBM never calls `toPandas()` and does not reimplement LightGBM. CPU is the default. There is no required SynapseML dependency.
26
+
27
+ ## Quick start
28
+
29
+ The `features_col` value must identify a Spark column containing a numeric feature vector or a numeric array/list. Spark ML `Vector` values and ordinary Spark arrays are accepted. The label column must contain numeric labels or targets.
30
+
31
+ ```python
32
+ from sparklightgbm import LightGBMClassifier
33
+
34
+ estimator = LightGBMClassifier(
35
+ features_col="features",
36
+ label_col="label",
37
+ num_workers=None,
38
+ num_iterations=300,
39
+ learning_rate=0.05,
40
+ num_leaves=31,
41
+ max_depth=-1,
42
+ min_data_in_leaf=20,
43
+ feature_fraction=0.9,
44
+ bagging_fraction=0.8,
45
+ bagging_freq=1,
46
+ lambda_l1=0.0,
47
+ lambda_l2=1.0,
48
+ min_gain_to_split=0.0,
49
+ max_bin=255,
50
+ seed=7,
51
+ feature_fraction_seed=7,
52
+ bagging_seed=7,
53
+ data_random_seed=7,
54
+ )
55
+
56
+ model = estimator.fit(train_df)
57
+ predictions = model.transform(test_df)
58
+ predictions.select("prediction", "probability").show()
59
+ ```
60
+
61
+ ## Estimators and input parameters
62
+
63
+ All three estimators accept the shared parameters below. Additional keyword arguments are forwarded to native LightGBM, allowing supported LightGBM parameters to be used without waiting for a wrapper release.
64
+
65
+ | Parameter | Accepted type | Default | Description |
66
+ | --- | --- | --- | --- |
67
+ | `features_col` | `str` | `"features"` | Spark column containing a numeric vector or array. |
68
+ | `label_col` | `str` | `"label"` | Spark column containing the target. |
69
+ | `prediction_col` | `str` | `"prediction"` | Output column for the predicted value or class. |
70
+ | `raw_prediction_col` | `str \| None` | `"rawPrediction"` | Output column for raw LightGBM scores; `None` omits it. |
71
+ | `probability_col` | `str \| None` | `"probability"` | Classifier probability output column; `None` omits it. |
72
+ | `leaf_prediction_col` | `str \| None` | `"leafPrediction"` | Output column containing integer leaf indices; `None` omits it. |
73
+ | `weight_col` | `str \| None` | `None` | Optional Spark column containing per-row weights. |
74
+ | `group_col` | `str \| None` | `None` | Optional Spark column containing ranking group identifiers; required by `LightGBMRanker`. |
75
+ | `categorical_feature` | `list[int] \| str \| None` | `None` | LightGBM categorical feature indices or native categorical setting. Feature values must be numeric category codes. |
76
+ | `validation_data` | Spark `DataFrame` \| `None` | `None` | Validation DataFrame with matching feature, label, weight, and group names. Can also be passed to `fit()`. |
77
+ | `early_stopping_rounds` | `int \| None` | `None` | Rounds without validation improvement; requires validation data. Distributed mode supports exactly aggregatable metrics. |
78
+ | `seed` | `int` | `0` | Seed passed to LightGBM. |
79
+ | `num_workers` | `int \| None` | `None` | Automatically selects a conservative worker count. Local Spark and ranking use driver training. Clustered classification and regression use the smallest of four workers, the input partition count, and available Spark slots minus one. Set a positive integer to override the strategy. |
80
+ | `local_listen_port` | `int` | `12400` | Distributed listener base port; worker `n` uses `local_listen_port + n`. |
81
+ | `prediction_batch_size` | `int` | `1024` | Rows scored per native LightGBM prediction batch in each Spark partition. |
82
+ | `objective` | `str \| None` | `None` | Native objective: defaults to `binary`/`multiclass`, `regression`, or `lambdarank`; `quantile` is supported for regression. |
83
+ | `**params` | Native value | - | Parameters forwarded to `lightgbm.Dataset` or `lightgbm.train`, as applicable. |
84
+
85
+ Estimator-specific parameters:
86
+
87
+ | Estimator | Parameter | Accepted type / default | Description |
88
+ | --- | --- | --- | --- |
89
+ | `LightGBMClassifier` | `num_class` | `int \| None`, default `None` | Number of classes; if omitted, inferred from labels. |
90
+ | `LightGBMRegressor` | - | - | Supports native regression parameters and `objective="quantile"`; set native `alpha` for the target quantile. |
91
+ | `LightGBMRanker` | - | - | Uses LambdaRank by default. `group_col` identifies groups; rows are sorted by group before native training. |
92
+
93
+ ## Common native LightGBM parameters
94
+
95
+ The following are commonly used native parameters, not a separate SparkLightGBM parameter system. Their appropriate values depend on the data, objective, validation strategy, and resource constraints.
96
+
97
+ | Parameter | What it controls |
98
+ | --- | --- |
99
+ | `num_iterations` | Maximum number of boosting rounds. `n_estimators` and `num_boost_round` are also accepted by SparkLightGBM. |
100
+ | `learning_rate` | Contribution of each new tree; interacts with the number of iterations. |
101
+ | `num_leaves` | Maximum leaves per tree and therefore much of the model's capacity. |
102
+ | `max_depth` | Optional tree-depth limit; negative values leave depth unconstrained. |
103
+ | `min_data_in_leaf` | Minimum observations allowed in a leaf, controlling leaf granularity and regularization. |
104
+ | `feature_fraction` | Fraction of features considered for each tree. |
105
+ | `bagging_fraction` | Fraction of rows used when bagging is active. |
106
+ | `bagging_freq` | Frequency of bagging; `0` disables it. |
107
+ | `lambda_l1` | L1 regularization applied to leaf weights. |
108
+ | `lambda_l2` | L2 regularization applied to leaf weights. |
109
+ | `min_gain_to_split` | Minimum gain required to create a split. |
110
+ | `max_bin` | Maximum histogram bins used for numeric features; affects accuracy, memory, and speed. |
111
+ | `is_unbalance` | Enables automatic binary-class imbalance handling. Do not combine it with `scale_pos_weight`. |
112
+ | `scale_pos_weight` | Explicit positive-class weight for binary classification. Do not combine it with `is_unbalance`. |
113
+ | `seed` | Top-level seed used by SparkLightGBM and passed to LightGBM. |
114
+ | `data_random_seed` | Seed used while constructing histogram bins. |
115
+ | `feature_fraction_seed` | Seed used for feature subsampling. |
116
+ | `bagging_seed` | Seed used for row subsampling. |
117
+ | `drop_seed` | Seed used by DART boosting. |
118
+ | `deterministic` | Requests stable CPU results; LightGBM may require related parameters for fully reproducible runs. |
119
+
120
+ `fit(params={...})` can supply or override native parameters for a training call:
121
+
122
+ ```python
123
+ from sparklightgbm import LightGBMRegressor
124
+
125
+ model = LightGBMRegressor(features_col="features", label_col="target").fit(
126
+ train_df,
127
+ params={"objective": "quantile", "alpha": 0.9, "n_estimators": 200, "learning_rate": 0.03},
128
+ )
129
+ ```
130
+
131
+ SparkLightGBM documents only common parameters. Availability, aliases, interactions, and exact semantics are determined by the installed LightGBM version; see the [official LightGBM parameter documentation](https://lightgbm.readthedocs.io/en/latest/Parameters.html) for the complete reference. SparkLightGBM uses `lightgbm.Dataset` and `lightgbm.train`, not the scikit-learn wrapper.
132
+
133
+ The public training signature is `fit(dataset, params=None, validation_data=None)`: `dataset` and `validation_data` are Spark `DataFrame` objects, and `params` is a `dict[str, object]` of native LightGBM parameter overrides. A `validation_data` argument passed to `fit()` takes precedence over the constructor value.
134
+
135
+ ## Spark input and output
136
+
137
+ Training reads rows from Spark partitions. Supported training columns are:
138
+
139
+ - `features_col`: numeric Spark ML vector, numeric array, or numeric Python list.
140
+ - `label_col`: numeric binary/multiclass label, regression target, or ranking relevance score.
141
+ - `weight_col`: optional numeric row weight.
142
+ - `group_col`: optional numeric/integer ranking group identifier.
143
+
144
+ Missing feature values should use LightGBM-compatible values such as `NaN`. A null feature vector is skipped; a null label or malformed feature value raises an error during collection or native training.
145
+
146
+ `model.transform(df)` preserves the input DataFrame and adds configured output columns. It creates `prediction` for every estimator, raw scores when enabled, leaf indices when enabled, and probabilities for classifiers when enabled. Output names can be overridden per call:
147
+
148
+ The transform signature is `transform(dataset, prediction_col=None, raw_prediction_col=None, probability_col=None, leaf_prediction_col=None)`. Every column-name override is a `str` or `None`; `probability_col` is used only for classifier models.
149
+
150
+ ```python
151
+ predictions = model.transform(
152
+ test_df,
153
+ prediction_col="score",
154
+ raw_prediction_col="raw_score",
155
+ probability_col="class_probability",
156
+ leaf_prediction_col="leaf_index",
157
+ )
158
+ ```
159
+
160
+ ## Validation and early stopping
161
+
162
+ Use a Spark DataFrame with matching column names:
163
+
164
+ ```python
165
+ model = LightGBMRegressor(n_estimators=500, early_stopping_rounds=30).fit(
166
+ train_df,
167
+ validation_data=validation_df,
168
+ )
169
+ ```
170
+
171
+ Validation data remains partitioned in distributed mode. Each worker builds local training and validation datasets, while Spark barrier synchronization aggregates decomposable native metrics before the early-stopping callback. Metrics whose exact global value cannot be reconstructed from shard-level scalar results (currently `auc`, `average_precision`, `map`, and `ndcg`) are rejected for distributed early stopping; use `num_workers=1` for those metrics.
172
+
173
+ ## Local and distributed execution
174
+
175
+ The worker setting changes where training data is materialized:
176
+
177
+ | Setting | Execution and memory behavior | Intended use |
178
+ | --- | --- | --- |
179
+ | `num_workers=1` | Spark partitions are read and the complete training dataset—and validation dataset, when present—is collected into NumPy arrays on the driver. Native LightGBM trains in the driver process. | Local development, compatibility fallback, ranking, and datasets that safely fit in driver memory. |
180
+ | `num_workers>1` | Spark repartitions the data into a barrier stage. Each executor-side worker converts only its shard to contiguous NumPy buffers and participates in LightGBM's native `data_parallel` network. Training and validation rows are not collected to the driver; the driver receives bounded metadata and the trained model artifact. | Cluster datasets that should remain distributed. Each worker shard must fit in that executor's memory. |
181
+ | `num_workers=None` | Selects `1` for local Spark and ranking. On a cluster, classification and regression reserve one slot and select the smallest of `4`, the remaining Spark parallelism, and the input partition count. The result is always at least `1`. | Portable, conservative default that avoids claiming every reported slot. Set an explicit value when cluster capacity or scheduling calls for another count. |
182
+
183
+ `SparkContext.defaultParallelism` is an estimate, not a dynamic cluster-capacity reservation. On shared or autoscaling clusters, set `num_workers` explicitly when the scheduler policy requires a specific limit. An explicit value always takes precedence, including in local mode; the requested barrier tasks must be schedulable concurrently.
184
+
185
+ Distributed workers default to one native LightGBM thread per Spark task to avoid CPU oversubscription. Pass `num_threads` explicitly when the Spark resource configuration provides additional CPU capacity per task.
186
+
187
+ Distributed execution requires:
188
+
189
+ - SparkLightGBM, NumPy, and the same compatible LightGBM build on every executor.
190
+ - Spark barrier execution support.
191
+ - Executor hostnames or IP addresses resolvable and reachable from every other executor.
192
+ - Listener ports available between executors; set `local_listen_port` if `12400 + worker_id` is unavailable.
193
+ - Correct executor networking. `SPARK_LOCAL_IP` can provide the advertised worker address when hostname resolution is unsuitable.
194
+
195
+ Inference uses partition-level NumPy batches instead of row-wise Python UDFs. A native booster is cached per reused Python worker, and classification prediction, probability, and raw prediction are derived from the same raw-score batch.
196
+
197
+ ## Current limitations
198
+
199
+ - The `0.9.x` API is in pre-1.0 validation and may change before `1.0.0`.
200
+ - Distributed ranking is not yet supported because query groups must remain complete and worker-local. Ranking therefore uses the single-worker path.
201
+ - Distributed early stopping supports metrics that can be exactly aggregated from worker-level results. Non-decomposable metrics currently rejected for this mode include `auc`, `average_precision`, `map`, and `ndcg`; use `num_workers=1` when early stopping depends on them.
202
+ - Distributed jobs require Spark barrier scheduling plus stable, mutually reachable executor addresses and ports. Executor loss aborts the coordinated native training job.
203
+ - Models use native LightGBM persistence rather than Spark ML `MLWriter`/`MLReader`, and the estimators are not yet Spark ML `Estimator`/`Model` stages for `Pipeline` or `CrossValidator`.
204
+ - CPU is the supported default execution path; GPU execution is not currently documented or tested by this project.
205
+
206
+ ## Predictions and explainability
207
+
208
+ Fitted models expose native LightGBM outputs:
209
+
210
+ ```python
211
+ raw = model.predict_raw([[1.0, 2.0]])
212
+ probability = model.predict_probability([[1.0, 2.0]]) # classifiers
213
+ leaves = model.predict_leaf([[1.0, 2.0]])
214
+ shap = model.predict_shap([[1.0, 2.0]])
215
+ split_importance = model.feature_importance("split")
216
+ gain_importance = model.feature_importance("gain")
217
+ ```
218
+
219
+ `predict_shap()` requests native LightGBM contribution values. The output includes one value per feature plus the expected-value contribution. `feature_importance()` accepts `"split"` or `"gain"`.
220
+
221
+ ## Native model persistence
222
+
223
+ Models are saved in LightGBM's native format. A JSON sidecar containing bridge metadata is written next to the model file:
224
+
225
+ ```python
226
+ model.save_native_model("artifacts/model.txt")
227
+
228
+ from sparklightgbm import LightGBMClassificationModel
229
+ restored = LightGBMClassificationModel.load_native_model(
230
+ "artifacts/model.txt",
231
+ features_col="features",
232
+ )
233
+ ```
234
+
235
+ The native model file is portable across Spark jobs with a compatible LightGBM installation. It is not a Spark ML `PipelineModel` format.
236
+
237
+ `save_native_model(path)` accepts a filesystem path (`str` or `pathlib.Path`). `LightGBMClassificationModel.load_native_model(path, **kwargs)`, `LightGBMRegressionModel.load_native_model(path, **kwargs)`, and `LightGBMRankingModel.load_native_model(path, **kwargs)` accept the native model path plus optional `features_col`, `prediction_col`, `raw_prediction_col`, `probability_col`, `leaf_prediction_col`, and `kind` strings used to reconstruct Spark output behavior.
238
+
239
+ ## Compatibility and errors
240
+
241
+ SparkLightGBM validates required columns and checks that supported PySpark and LightGBM installations are available on the driver. The `0.9.x` compatibility policy is Python 3.9-3.12 and PySpark 3.4-3.5; newer combinations are not claimed until validated. Runtime imports also need to be available on executors. Missing dependencies, missing columns, missing ranking groups, unsupported distributed validation settings, and invalid native parameters fail with errors from the bridge or LightGBM.
242
+
243
+ ## Development
244
+
245
+ ```bash
246
+ pip install -e ".[dev]"
247
+ python -m pytest -q
248
+ python -m build
249
+ python -m twine check dist/*
250
+ ```
251
+
252
+ See [CHANGELOG.md](CHANGELOG.md) and [CONTRIBUTING.md](CONTRIBUTING.md).
@@ -0,0 +1,39 @@
1
+ # Architecture audit
2
+
3
+ ## Execution paths
4
+
5
+ Single-worker training intentionally provides a compatibility fallback:
6
+
7
+ ```text
8
+ Spark DataFrame -> partition rows -> driver NumPy arrays -> LightGBM Dataset -> Booster
9
+ ```
10
+
11
+ This path collects training and validation rows and is suitable only when they fit in driver memory.
12
+
13
+ Distributed classifier and regressor training follows:
14
+
15
+ ```text
16
+ Spark DataFrame -> repartitioned barrier stage -> executor-local contiguous NumPy arrays
17
+ -> executor-local LightGBM Datasets -> native LightGBM network
18
+ -> distributed tree construction -> model string metadata to driver
19
+ ```
20
+
21
+ Training rows are never collected by the driver. The driver collects only bounded metadata (class labels), one serialized model per worker, and barrier results. Validation rows are tagged, co-partitioned with training shards, and remain executor-local. Native shard metrics that are mathematically decomposable are combined through the Spark barrier before early stopping, ensuring every worker stops on the same iteration.
22
+
23
+ Distributed ranking remains disabled because preserving complete query groups per worker needs a dedicated partitioning contract.
24
+
25
+ ## Data movement and memory
26
+
27
+ The distributed path does not use pandas. Spark rows are decoded once into feature, label, and optional-weight buffers, then converted into contiguous NumPy matrices. One training and, when configured, one validation matrix are resident per worker. The current Python-row boundary remains a measurable cost and is a benchmark target; replacing it requires evidence that Arrow or another optional path improves end-to-end behavior without reducing portability.
28
+
29
+ Inference uses `RDD.mapPartitions`, constructs or reuses one Booster per Python worker, batches feature vectors into contiguous matrices, and calls native prediction once per requested mode and batch. Null feature vectors produce null outputs consistently.
30
+
31
+ ## Known constraints
32
+
33
+ - Barrier jobs require all worker slots concurrently and stable executor networking.
34
+ - Exact distributed early stopping is unavailable for non-decomposable shard metrics such as AUC and average precision.
35
+ - Worker loss aborts native distributed training; transparent recovery would require a full coordinated restart.
36
+ - Spark ML `Estimator`, `Model`, and persistence interfaces are not implemented yet.
37
+ - Model artifacts are native LightGBM files plus lightweight metadata, not `PipelineModel` artifacts.
38
+
39
+ These constraints are kept explicit instead of adding silent fallbacks that collect distributed data to the driver.