在 PySpark Pandas UDF 中使用 scikit-learn train_test_split 函数时出现 ValueError

Question

我想为 Pyspark 创建一个 pandas udf 函数，我在其中使用 scikit-learn train_test_split 函数并返回一个数据帧。

我有一个这样的数据框：但是在我的数据框中，没有 id 列。所以我在 dataframe

中添加了 id 列

这就是我所做的。

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def load_dataset(df):
    X = df[X_columns]
    y = df[y_columns]
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size = 0.2, random_state = 42)
    
    df_sample_0 = pd.concat([y_test, X_test], axis=1)
 
    return df_sample_0

这就是我应用 groupby 的方式：

sample_df = final_df_spark.groupby("id").apply(load_dataset)

但是我收到这个错误：

ValueError: With n_samples=1, test_size=0.2 and train_size=None, the resulting train set will be empty. Adjust any of the aforementioned parameters.

我该如何解决这个错误？

Answer 1

我刚刚将 id 列替换为 Age 列。

sample_df = final_df_spark.groupby("Age").apply(load_dataset)

在 PySpark Pandas UDF 中使用 scikit-learn train_test_split 函数时出现 ValueError

ValueError when using scikit-learn train_test_split function in PySpark Pandas UDF

python

user-defined-functions

scikit-learn

pyspark