PyData 生态系统的一大优势在于其出色的互操作性，能够支持数据在专注于探索性分析、模型训练与推理的各类库之间无缝流转。XGBoost 的最新版本引入了多项令人期待的新功能，例如类别变量的重新编码以及对 Polars DataFrame 的集成，进一步简化了数据处理流程。
本文将指导您如何将 Polars 的 GPU 引擎与 XGBoost 机器学习库相结合，重点介绍分类特征的无缝集成，以及 XGBoost 中新增的类别重新编码功能。
将 XGBoost 与 Polars GPU 引擎结合使用
Polars 是一个基于 Rust 构建的高性能 DataFrame 库，支持延迟计算和 GPU 加速，能够显著提升数据处理流程的效率。
在 GPU 加速的工作流中，将 Polars 与 XGBoost 结合使用的一个关键环节是理解延迟评估机制。Polars 的操作通常是延迟执行的，这意味着它们会先构建查询计划，但除非显式触发，否则不会立即执行。为了在 GPU 上自定义查询计划的执行过程，可以调用 ig1Q 的 ig0Q 方法，并指定 ig2Q 参数。
本教程使用 Microsoft 恶意软件预测数据集的一个小样本进行说明。该数据集可通过 Kaggle 获取，包含中等规模的数值型和分类型特征。在接下来的代码示例中，我们仅选取了部分列，以展示如何在 XGBoost 中处理分类型特征。
设置环境
在深入研究代码之前，请确保已安装以下库：xgboost、polars[gpu] 和 pyarrow。[gpu] 依赖项用于下载支持 GPU 的 Polars 版本。
pip install xgboost polars[gpu] pyarrow
使用 Polars 输入时，XGBoost 通过零拷贝方式处理 Polars DataFrame，从而提升效率。因此，Polars 与 XGBoost 之间的数据传递需借助 PyArrow 实现。如本示例所示，PyArrow 还可作为从 XGBoost 模型导出类别信息的数据交换格式。
数据准备和模型训练
首先，导入必要的库：
import polars as pl
import xgboost as xgb
数据集中包含三个特征，其中两个为分类特征，HasDetections 列是具有二进制值的预测目标。为充分发挥 Polars 执行引擎的性能优势，可基于 scan_csv 创建 LazyFrame 对象。
columns = [
"ProductName", # Categorical
"IsBeta", # Boolean
"Census_OSArchitecture", # Categorical
"HasDetections", # Binary target
]
# ignore_errors is set to True to let polars infer the schema
df_lazy = pl.scan_csv(
"./microsoft-malware-prediction/train.csv",
ignore_errors=True,
).select(columns)
# Cast the categorical features to the Polars Enum type
df_lazy = df_lazy.with_columns(
[
pl.col("ProductName").cast(
pl.Enum(
["fep", "mseprerelease", "win8defender", "scep", "mse", "windowsintune"]
)
),
pl.col("Census_OSArchitecture").cast(pl.Enum(["amd64", "x86", "arm64"])),
]
)
读取数据后，将 DataFrame 输入 XGBoost 模型，用于训练二分类任务。
X = df_lazy.drop("HasDetections")
y = df_lazy.select("HasDetections")
# Use GPU to fit the classification model. We let XGBoost handle the categorical features by setting the `enable_categorical` parameter.
clf = xgb.XGBClassifier(device="cuda", enable_categorical=True)
# Validation data is not created for simplicity
clf.fit(X, y)
该代码段仅使用 GPU 进行模型训练，而 DataFrame 的加载与处理仍在 CPU 上执行。调用 fit 方法时，XGBoost 会发出警告，建议将延迟帧转换为具体的数据帧以提升性能。为实现此目标并按数据帧自定义查询执行计划，可使用以下代码段：
# Convert the lazy frame into a concrete dataframe using a GPU
df = df_lazy.collect(engine="gpu")
X = df.drop("HasDetections")
y = df.select("HasDetections")
clf.fit(X, y)
或者，若要在模型训练之外为 Polars 启用全局 GPU 加速，可将引擎配置设置为 GPU。
import polars as pl
import xgboost as xgb
# set the engine before using polars
pl.Config.set_engine_affinity("gpu")
使用 XGBoost 对分类数据进行自动重新编码
最新发布的 XGBoost 版本引入了重新编码器，显著提升了对分类特征的处理能力。Polar 会依据输入值的顺序，将分类和枚举数据类型编码为整数。例如，对于三个类别 [“aa”, “bb”, “CC”]，Polar 可能会按如下方式对数据进行存储：
| 类别 | 值 | 类别 |
“cc” |
2 | “aa” |
“cc” |
1 | “bb” |
“bb” |
1 | “cc” |
“aa” |
0 |
该方案在 DataFrame 实现（包括 pandas 和 cuDF）之间通用。如需深入了解 Polars 中的分类类型与枚举类型，可参考 Polars 官方文档。
在早期的 XGBoost 版本中，用户需要在推理前手动对分类特征进行编码。这一过程不仅容易出错，而且操作复杂、效率较低。
例如，在训练数据集中包含三个类别的特征时，Polars 会通过映射将其编码为对应的数值。然而，在推理阶段，测试数据集可能仅包含这些类别的子集，该子集将通过相同的映射进行编码，可能导致编码结果在测试时无效。
借助更新的 XGBoost，增强型对象能够记忆训练数据集中的编码，并在 predict 方法中利用该编码自动重新编码类别。以下示例展示了如何将此功能与包含分类特征的合成数据集结合使用：
import numpy as np
import polars as pl
import xgboost as xgb
# Create a dataframe with a categorical feature (f1)
f0 = [1, 3, 2, 4, 4]
cats = ["aa", "cc", "bb", "ee", "ee"]
df = pl.DataFrame(
{"f0": f0, "f1": cats},
schema=[("f0", pl.Int64()), ("f1", pl.Categorical(ordering="lexical"))],
)
rng = np.random.default_rng(2025)
y = rng.normal(size=(df.shape[0]))
# Train a regression model
reg = xgb.XGBRegressor(enable_categorical=True, device="cuda")
reg.fit(df, y)
predt_0 = reg.predict(df)
# Use a subset of rows to create a different encoding, "aa" and "ee" are removed
df_new = pl.DataFrame(
{"f0": f0[1:3], "f1": cats[1:3]},
schema=[("f0", pl.Int64()), ("f1", pl.Categorical(ordering="lexical"))],
)
predt_1 = reg.predict(df_new)
# Check the resulting predictions are the same with the original encoding
np.testing.assert_allclose(predt_0[1:3], predt_1)
在此代码段中，通过训练 DataFrame 的类别子集构建测试 DataFrame，并验证在不同编码方案下输出预测结果是否保持一致。该方法无需依赖独立的转换流程，即可实现预测的稳定性验证。
此外，在处理大量特征时，XGBoost 中的重新编码比直接对 DataFrame 进行重新编码更加高效。其内部机制采用即时就地的方式执行重新编码，无需复制整个 DataFrame。同时，XGBoost 能够利用 GPU 并行处理所有分类变量列，进一步提升计算效率。
导出类别 (试验性)
如前所述，XGBoost 模型现已支持类别记忆功能。高级用户可通过访问高级模型中的底层增强对象，将保存的类别导出为 Arrow 数组列表。这一功能有助于验证模型是否使用了预期的类别进行训练。以下继续以示例说明：
# Get the underlying booster object
booster = reg.get_booster()
# Export the categories from the booster
categories = booster.get_categories(export_to_arrow=True)
# Export the categories into a list of arrow arrays
print(categories.to_arrow())
使用 PyArrow 进行数据交换时，需要启用 export_to_arrow 选项。
[('f0', None), ('f1', <pyarrow.lib.StringArray object at 0x735ba5407be0>
[
"aa",
"cc",
"bb",
"ee"
])]
f1 中类别的完整列表存储在增强程序中。自 XGBoost 3.1 版本起，用于导出类别的接口处于试验阶段。请注意，本文示例采用 scikit-learn 接口，因为它能自动处理大部分配置。若使用原生接口（即增强程序），操作可能更为复杂，特别是在继续训练的情况下。更多详细信息请参考 XGBoost 官方文档。
Polars 的 GPU 加速目前仅应用于其执行计划，生成的 DataFrame 仍存储在 CPU 内存中。因此，在推理过程中，XGBoost 需要将数据复制回 GPU，用户会收到关于性能影响的一次性警告。
开始使用 Polars 和 XGBoost
通过掌握延迟 Polars DataFrame 的实现方法，并充分利用 XGBoost 分类特征处理的新功能（包括特征的重新编码），您可以构建高效且稳定的 GPU 加速工作流。这种方法不仅简化了流程，还能显著提升机器学习模型的性能表现。
有关更多信息，请参阅 如何使用 Polars 与 NVIDIA RAPIDS 实现 GPU 加速。您还可以在 dmlc/xgboost 的 GitHub 仓库中提交关于 XGBoost 训练的反馈或提出问题。
