如何在pyspark上创建分层的拆分训练、验证和测试集？

from pyspark.ml.tuning import TrainValidationSplit
from pyspark.ml.evaluation import MulticlassClassificationEvaluator

# 假设你已经加载了你的数据集并进行了预处理
data = ...

# 定义特征列和标签列
feature_cols = ...
label_col = ...

# 使用StratifiedSampling进行分层抽样
stratified_data = data.stat.sampleBy(label_col, fractions={0: 0.6, 1: 0.2, 2: 0.2}, seed=42)

# 将数据集分为训练集和验证集
train_data, validation_data = stratified_data.randomSplit([0.8, 0.2], seed=42)

# 创建模型并进行训练
model = ...
model.fit(train_data)

# 在验证集上进行评估
evaluator = MulticlassClassificationEvaluator(labelCol=label_col, metricName='accuracy')
validation_accuracy = evaluator.evaluate(model.transform(validation_data))

# 创建测试集
test_data = data.subtract(stratified_data)

# 在测试集上进行评估
test_accuracy = evaluator.evaluate(model.transform(test_data))

# 打印结果
print("Validation Accuracy: ", validation_accuracy)
print("Test Accuracy: ", test_accuracy)