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

from pyspark.sql.functions import pandas_udf
from pyspark.sql.types import StructType, StructField, IntegerType, StringType

@pandas_udf(returnType=StructType([
    StructField("df1_col1", IntegerType()),
    StructField("df1_col2", StringType()),
    StructField("df2_col1", IntegerType()),
    StructField("df2_col2", StringType())
]))
def process_data(input_df):
    # 使用pandas进行数据处理
    df1 = input_df[['col1', 'col2']]
    df2 = input_df[['col3', 'col4']]
    
    # 返回多个数据帧
    return df1, df2

# 加载数据到Spark DataFrame
data = [(1, 'A', 10, 'X'), (2, 'B', 20, 'Y'), (3, 'C', 30, 'Z')]
df = spark.createDataFrame(data, ['col1', 'col2', 'col3', 'col4'])

# 使用pandas_udf函数处理数据
result_df = df.withColumn("result", process_data(df))

result_df.select("result.df1_col1", "result.df1_col2").show()
result_df.select("result.df2_col1", "result.df2_col2").show()