XGBoost4J - Scala 数据帧到稀疏 dmatrix

Question

将 scala 数据帧转换为 XGBoost4J 的稀疏 dmatrix 的最有效和可扩展的方法是什么？

假设我有一个数据框 train，其中包含 row_index、column_index 和 value 列，它会是这样的

new DMatrix(train.select("row_index"), train.select("column_index"), train.select("Value"), DMatrix.SparseType.CSR, n_col)

但是，上面的代码会导致类型不匹配，因为 DMatrix 需要 Array[Long]。

train.select(F.collect_list("row_index")).first().getList[Long](0) 似乎是一个可能的选择，但它似乎不是内存友好和可扩展的。

我在 Databricks 上执行此操作，因此欢迎使用其他受支持语言（python、SQL、scala）的解决方案。

Answer 1

答案是按行使用稀疏向量，而不是尝试创建稀疏矩阵或 dmatrix。

train.rdd.map(r => (r.getInt(0), (r.getInt(1), r.getInt(2).toDouble))).groupByKey().map(r => (r._1, Vectors.sparse(n_col, r._2.toSeq))).toDF

我测试了使用Matrix::sparseMatrix 和xgboost::dmatrix 对R 中的数据样本进行评分，结果匹配。