如何在 Pyspark 中使用 @pandas_udf return 多个数据帧？

Question

我想为 Pyspark 创建 sklearn 的 train_test_split 函数。我正在使用 pandas udf 创建此函数

这就是我所做的。

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def load_dataset(dataset):
    
    feature_columns = cols
    label = 'y';
    X = dataset[feature_columns]
    Y = dataset[label]
 
    # splitting the dataset into train and test
    X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2)
    print(X_train)
 
    return X_train, X_test, y_train, y_test

我想要这些数据帧X_train、X_test、y_train、y_test。

我知道udf函数是这样调用的

df.groupby("key").apply(load_dataset).show()

但我不知道用什么代替 key 另外，这个 returns 单个数据框，我想要四个。

Answer 1

有什么问题：

df = inputDF.cache()
a,b = df.randomSplit([0.5, 0.5])

对于顺序重要的时间序列，使用：

df = df.withColumn("rank", percent_rank().over(Window.partitionBy().orderBy("departure_time")))

train_df = df.where("rank <= .8").drop("rank", "departure_time")

Answer 2

实际上我必须进行子采样。这就是为什么我必须从 train_test_split 函数中获取四个变量 returns 的原因。但是我连接了 X_test 和 y_test 并返回了一个数据帧。

@pandas_udf(schema, PandasUDFType.GROUPED_MAP)
def load_dataset(dataset):
    
    feature_columns = cols
    label = 'y';
    X = dataset[feature_columns]
    Y = dataset[label]
 
    # splitting the dataset into train and test
    X_train, X_test, y_train, y_test = train_test_split(X, Y, test_size=0.2)
    print(X_train)
    
    df_sample = pd.concat([X_test, y_test], axis=1)
 
    return df_sample

这段代码对我有用。

如何在 Pyspark 中使用 @pandas_udf return 多个数据帧？

How to return multiple dataframes using @pandas_udf in Pyspark?

python

user-defined-functions

pandas

scikit-learn

pyspark