polars_online.gbdt

XGBoost’s and LightGBM’s training inputs, built from Arrow files, Parquet files or a LazyFrame as fast as the libraries go.

Both libraries train in two stages. Binning places each feature’s bins and puts every value into one; boosting grows trees from per-bin sums and never reads a raw value again. This module does the first stage from data that streams, in Rust, and hands back the library’s own training object, so that xgb.train and lgb.train run unchanged, with their evaluation sets, early stopping and callbacks. Here the training rows are the example data’s two parquet files, read by a glob:

import xgboost as xgb

features = ["x0", "x1", "x2", "signal_a", "signal_b"]
train = po.gbdt.xgboost_matrix("ticks/*.parquet", label="y", features=features)
valid = po.gbdt.xgboost_matrix("today.parquet", label="y", features=features, ref=train)
booster = xgb.train({"tree_method": "hist", "max_bin": 256}, train, 500,
                    evals=[(valid, "valid")], early_stopping_rounds=20, verbose_eval=False)
  • xgboost_matrix() returns an xgboost.QuantileDMatrix. XGBoost reads its input four times, and the second pass sketches every feature’s quantiles over every row. This places the cuts from a sample of sample_rows rows instead, which removes that pass, about half the time before the first round. Without weights the sample is uniform. On 10 million unweighted rows of 50 features, at 50 rounds, XGBoost’s sketch scored a validation RMSE of 1.2379, and cuts from five seeds of the sample 1.2347 to 1.2411. LightGBM’s own sampling spans 1.2298 to 1.2443 over five seeds. With weight=, the sample is drawn in proportion to the weights and its quantiles are weighted, as XGBoost’s sketch is. On 50,000 normal values with weight 50 above 1.0 and max_bin=16, 14 of the 15 cuts fall above 1.0, as 14 of XGBoost’s own do, and the heaviest bin holds 6.26% of the weight, against 6.47% for XGBoost’s. sample_rows=None keeps XGBoost’s sketch and its own trees.

  • lightgbm_dataset() returns a constructed lightgbm.Dataset, binned on LightGBM’s own sample: a first pass gathers the rows LightGBM samples, a second pushes every row as float32; or, when the data fits, the whole stream goes in as Arrow.

The native engine reads Arrow IPC files itself, a record batch at a time, and any other source through Polars’ streaming engine (collect_batches). It hands each batch’s Arrow buffers to the library’s C API with no copy: nulls go by Arrow validity, not as a NaN copy. The library is loaded from the user’s own installed package, by the path that package loaded it from. Nothing is linked into this package, whose only dependency stays polars. engine="python" builds through the libraries’ Python APIs instead: numpy batches into an xgboost.DataIter with the same sample’s cuts, and a lightgbm.Sequence. engine="auto" falls back to it, with a warning, when the native engine cannot load the library. Both engines grow the same trees.

What the columns may hold. Every feature is numeric (any integer or float dtype), Boolean, or a pl.Enum, whose codes go in as a categorical feature. The Enum’s categories fix what each code means, so a model is read the same way at prediction (pass the codes, pl.col(c).to_physical()). A pl.Categorical is refused, since its codes depend on the process that made them; cast it to a pl.Enum. The label and the weight are numeric, with no nulls and finite values (a weight not negative).

How the figures in these docstrings were measured. On an Apple M4 Pro (14 cores, 48 GB), each route ran in a fresh process with 8 threads for the library, for Polars and for Rust. Other jobs shared the machine, so the routes ran interleaved, three times each, and a figure is the median, at a load average of 4.5 to 10. Memory is the peak physical footprint. The rows were seeded: 50 Float64 features of four kinds (normal, lognormal, uniform, and a normal rounded to 0.1), 5% nulls in every fifth column, one pl.Enum of 32 levels drawn unevenly, and a target built from them plus noise. Accuracy is the validation RMSE on separate rows from another seed.

Requires the library itself, which is not a dependency of this package (pip install xgboost or lightgbm), and numpy, which both libraries require.

Warning

This module is considered unstable: it is built on the libraries’ C APIs and on how their Python classes hold a native handle, so its keywords and what it supports may change in any release without that counting as a breaking change, and a new library version may need a new release here. Under POLARS_ONLINE_WARN_UNSTABLE=1 each call raises polars_online.UnstableWarning.

polars_online.gbdt.lightgbm_dataset(source: Source, *, label: str, weight: str | None = None, features: Sequence[str] | None = None, categorical: Sequence[str] | None = None, max_bin: int = 255, params: dict[str, Any] | None = None, reference: lightgbm.Dataset | None = None, in_memory: bool | Literal['auto'] = 'auto', threads: int | None = None, engine: Engine = 'auto') → lightgbm.Dataset[source]

A constructed lightgbm.Dataset of source, binned on LightGBM’s own sample.

LightGBM places each feature’s bins from a sample of bin_construct_sample_cnt rows (200,000 by default), drawn from the row count and data_random_seed, and then pushes every value into its bin. Here the source is counted (select(pl.len())), a first pass gathers the rows LightGBM samples, and a second pushes every row as float32, built in parallel: two passes, and the data never held. With in_memory, one pass holds the features as float32 Arrow and hands LightGBM the whole stream, which it pushes column by column. Both routes hand LightGBM the same float32 values and grow the same trees, which are LightGBM’s own from a float32 array of the data. The returned dataset is LightGBM’s own, and lgb.train takes it as it is. Here the training rows are the example data’s two parquet files, read by a glob:

import lightgbm as lgb

features = ["x0", "x1", "x2", "signal_a", "signal_b"]
train = po.gbdt.lightgbm_dataset("ticks/*.parquet", label="y", features=features,
                                 params={"verbose": -1})
valid = po.gbdt.lightgbm_dataset("today.parquet", label="y", features=features,
                                 reference=train)
booster = lgb.train({"objective": "regression", "verbose": -1}, train, 500,
                    valid_sets=[valid],
                    callbacks=[lgb.early_stopping(20, verbose=False)])

Weights do not move the bins, in LightGBM or here. LightGBM’s sample is uniform, and neither LGBM_DatasetCreateFromSampledColumn nor the BinMapper::FindBin it calls takes a weight, so its bins are those of the unweighted rows. xgboost_matrix() weights its cuts, because XGBoost’s own sketch does: each builds what its library builds. With weight 50 on the values above 1.0 of 50,000 normal ones and max_bin=16, LightGBM’s trees split at the same 14 bounds with the weights as without, 2 of them above 1.0, where 14 of XGBoost’s 15 cuts are.

The dataset is constructed and its raw data never kept, so training may not change a dataset parameter (max_bin, min_data_in_bin, data_random_seed, …): LightGBM refuses, as it does for any constructed dataset. Pass them here, in params.

Measured as polars_online.gbdt describes, on 10 million rows of 50 features from Arrow IPC files: the time before the first round, and the peak memory.

route

seconds

peak, GB

this, streamed in two passes

1.68

1.7

this, in_memory

1.45

2.9

LightGBM’s own route from a collected Polars frame

1.54

8.9

LightGBM’s Sequence

2.17

2.7

Parameters

source

Arrow IPC or Parquet paths (globs taken, one kind of file), a LazyFrame (re-run for each pass), or a DataFrame. The native engine reads Arrow IPC files itself, a record batch at a time.

label, weight

The label and the row weights: numeric columns with no nulls and finite values, a weight not negative.

features

The feature columns, in order; every other column by default.

categorical

The categorical features, pl.Enum or integer columns (a code below 0 is missing, as LightGBM reads it); every pl.Enum feature by default. A pl.Enum left out goes in as its codes, a number.

max_bin

Bins per feature, LightGBM’s max_bin.

params

Other dataset parameters, as lightgbm.Dataset(params=) takes them: verbose, min_data_in_bin, bin_construct_sample_cnt, data_random_seed, use_missing, zero_as_missing, … Not max_bin, categorical_feature or num_threads, which are this function’s own arguments.

reference

The training dataset, for a validation dataset binned on its bins; one pass. It must have the same features in order, when it records its feature names, and a reference built here must have each pl.Enum feature of the same categories in order, since a code is read as the training level at that position; a reference built elsewhere is checked by its feature names alone.

in_memory

Hand LightGBM the whole stream at once, held in memory: as Arrow under the native engine, as a float32 numpy array under the Python one. "auto" does when the features as float32 take at most a quarter of the memory available and, under the native engine, the LightGBM loaded takes an Arrow stream (4.7.0 on).

threads

Threads for LightGBM’s construction and the conversions; all by default.

engine

"native", "python", or "auto": native, falling back to Python with a warning when the native engine cannot load LightGBM.

Raises

ValueError, each named, for:

  • max_bin below 2, threads not a positive integer, an engine not one of the three, or an in_memory that is not True, False or "auto";

  • a parameter in params that is this function’s own argument;

  • a source that names no files, mixes Arrow IPC and Parquet, or names a file whose extension says neither;

  • a column missing, named twice, or of a type the libraries do not read (a pl.Categorical among them);

  • a categorical name that is not a feature, or is a float column;

  • a null or non-finite label or weight, or a negative weight;

  • an empty source, more rows than LightGBM’s 32-bit row count holds, or a source whose passes give different rows;

  • a reference that records other features, or built here on a pl.Enum of other categories;

  • (engine="native") a LightGBM the native engine does not support.

TypeError for a reference that is not a lightgbm.Dataset (one not yet constructed is constructed here), a source of another type, or features or categorical given as one name. FileNotFoundError for a glob of Arrow IPC files that matches nothing.

polars_online.gbdt.xgboost_matrix(source: Source, *, label: str, weight: str | None = None, features: Sequence[str] | None = None, max_bin: int = 256, sample_rows: int | None = 200000, ref: xgboost.DMatrix | None = None, threads: int | None = None, seed: int = 0, engine: Engine = 'auto') → xgboost.QuantileDMatrix[source]

An xgboost.QuantileDMatrix of source, binned in one stream.

XGBoost places each feature’s bins at its quantiles and reads the input in passes: it counts the values, sketches the quantiles over every row, indexes each value into its bin and builds the column matrix. The returned matrix is XGBoost’s own, and xgb.train takes it as it is. Here from Arrow IPC files, which the native engine reads itself:

import xgboost as xgb

df.head(300).write_ipc("train.arrow")              # the example data, as two Arrow files
df.tail(100).write_ipc("valid.arrow")
features = ["x0", "x1", "x2", "signal_a", "signal_b"]
train = po.gbdt.xgboost_matrix("train.arrow", label="y", features=features)
valid = po.gbdt.xgboost_matrix("valid.arrow", label="y", features=features, ref=train)
booster = xgb.train({"tree_method": "hist", "max_bin": 256}, train, 500,
                    evals=[(valid, "valid")], early_stopping_rounds=20, verbose_eval=False)

The cuts. With sample_rows, one pass first draws a sample of that many rows. Row i has a hash h, SplitMix64 of i and seed, and the priority w / u, where w is its weight (1 without weight) and u = (2h + 1) / 2**65. The sample_rows rows of largest priority are kept. Without weights these are the rows of smallest hash, a uniform sample. With weights it is priority sampling: a row heavier than tau, the largest priority not kept, is always kept and counts its own weight, and any other kept row counts tau. Summed over any set of rows, these counts estimate the set’s weight without bias, so the sample’s weighted quantiles estimate the source’s. A row of weight 0 counts for nothing and places no cut. The sample is drawn in proportion to the weights rather than uniformly because a uniform one loses heavy rows: under a weight of 1e9 on one row in a million, a uniform sample of 200,000 rows misses that row four times in five, and its cuts ignore it.

Each feature’s bin edges sit at the sample’s weighted quantiles. A feature with at most max_bin distinct sample values gets a bin per value. Any other gets max_bin bins of equal weight. Edge j is the first sorted sample value whose running weight exceeds j / max_bin of the total. Without weights that is the value of rank floor(j m / max_bin) for m values. At the quantile q, a cut’s share of the weight then errs by about sqrt(q (1 - q) / sample_rows), 0.11% at the median for 200,000 rows. Over ten seeds on a million rows, its standard deviation measured 0.29% to 0.46% with 20,000 rows, where the formula gives 0.35%, and 0.04% with 200,000 rows under weight 50 on a sixth of them. The cuts reach XGBoost as a ref, which removes its sketch pass. The sample is the same rows however the source is batched and whichever engine reads it. sample_rows=None lets XGBoost sketch every row, and grows XGBoost’s own trees.

Measured as polars_online.gbdt describes, on 10 million rows of 50 features: the time before the first round, and the peak memory. Polars’ read-ahead holds the LazyFrame’s 7.2 GB, and the README’s “Tuning memory with Polars’ own settings” cuts it.

route

seconds

peak, GB

this, from Arrow IPC files

4.0

3.6

this, from a LazyFrame over them

4.2

7.2

this, engine="python"

5.5

7.1

this, sample_rows=None

8.0

3.6

XGBoost’s own DataIter over Polars batches

8.7

4.8

collect() and numpy

7.0

8.7

Parameters

source

Arrow IPC or Parquet paths (globs taken, one kind of file), a LazyFrame, read in batches through collect_batches once per pass and so re-run each time, or a DataFrame. The native engine reads Arrow IPC files itself, a record batch at a time.

label, weight

The label and the row weights: numeric columns with no nulls and finite values, a weight not negative.

features

The feature columns, in order; every other column by default.

max_bin

Bins per feature, XGBoost’s max_bin: training must pass the same.

sample_rows

Rows of the sample the cuts come from; None for XGBoost’s own sketch. Ignored with ref.

ref

The training matrix, for a validation or test matrix binned on its cuts (xgb.train refuses a QuantileDMatrix evaluation set without it). It must have the same features in order, and a ref built here must have each pl.Enum feature of the same categories in order, since a code is read as the training level at that position; a ref built elsewhere is checked by its feature names alone. The returned matrix keeps a weak reference to it.

threads

Threads for XGBoost’s construction and the conversions; all by default.

seed

The sample’s seed.

engine

"native", "python", or "auto": native, falling back to Python with a warning when the native engine cannot load XGBoost.

Raises

ValueError, each named, for:

  • max_bin below 2, threads or sample_rows not a positive integer, or an engine not one of the three;

  • a source that names no files, mixes Arrow IPC and Parquet, or names a file whose extension says neither;

  • a column missing, named twice, or of a type the libraries do not read (a pl.Categorical among them);

  • a null or non-finite label or weight, or a negative weight;

  • an empty source, or one whose passes give different rows;

  • a ref binned on other features, or built here on a pl.Enum of other categories;

  • (engine="native") an XGBoost the native engine does not support.

TypeError for a ref that is not an xgboost.DMatrix, a source of another type, or features given as one name. FileNotFoundError for a glob of Arrow IPC files that matches nothing.