如何让 `train_test_split` 与数据框一起工作？

Question

我有一个包含 10,999 行的复杂数据框。

我正在尝试 运行 xgboost 用于机器学习。

我加载数据并尝试拆分它，正如我在教程中看到的以及 Whosebug 上发布的解决方案：How do I create test and train samples from one dataframe with pandas?

X_train, X_test = train_test_split(df, test_size=0.2)

但这失败了：

TypeError: Expected sequence or array-like, got <class 'pyspark.sql.dataframe.DataFrame'>

但这没有意义，我怎么可能将数据帧放入数组而不丢失大量有价值的信息？

所以有人建议我试试 pandas:

pandasDF = df.toPandas
X_train, X_test = train_test_split(pandasDF, test_size=0.2)

但这也失败了：

TypeError: Singleton array array(<bound method PandasConversionMixin.toPandas of DataFrame

如何将此数据框拆分为训练集和测试集？

Answer 1

使用这个选项：

pandasDF = df.toPandas()

如果需要时间，请在转换前使用此配置

spark.conf.set("spark.sql.execution.arrow.enabled", "true")