如何在使用make_csv_dataset创建的数据集上创建train_test_split

import tensorflow as tf
from sklearn.model_selection import train_test_split

url = "https://example.com/data.csv"  # 数据集的URL或本地文件路径
dataset = tf.data.experimental.make_csv_dataset(
    file_pattern=url,
    batch_size=32,
    label_name='label',
    num_epochs=1,
    ignore_errors=True
)

def extract_features_labels(element):
    features = tf.stack(list(element.values()), axis=1)
    labels = features[:, 0]  # 假设标签在第一列
    return features[:, 1:], labels

dataset = dataset.map(extract_features_labels)

features = []
labels = []

for feature, label in dataset:
    features.append(feature)
    labels.append(label)

features = tf.concat(features, axis=0)
labels = tf.concat(labels, axis=0)

train_features, test_features, train_labels, test_labels = train_test_split(
    features, labels, test_size=0.2, random_state=42
)