如何返回 PySpark 数据框中每行的数组平均值？

Question

假设我有如下数据：

from pyspark.sql import SparkSession
from pyspark.sql.types import ArrayType, DoubleType, StructField, StructType, LongType

spark = SparkSession.builder.appName("AveragingArray").getOrCreate()

# Define schema
schema = StructType([
    StructField("id", LongType(), True),
    StructField("arrays", ArrayType(DoubleType()), True)
])

# Sample data
data = [
    (1, [18.2, 50.9]),
    (2, [1.0, 3.4, 5.5]),
    (3, [10.0, 20.1, 30.0])
]

# Create DataFrame
df = spark.createDataFrame(data, schema=schema)

我想要一种简单的方法将

arrays

项转换为数组中值的平均值。例如，第一行是 34.55。

我尝试过这样的事情：

from pyspark.sql import functions as F

def average_arrays(df):
    # Explode the array to individual elements
    exploded_df = df.withColumn("arrays", F.explode("arrays"))
    
    # Group by the original identifier and calculate the average then return as int
    result_df = exploded_df.groupBy("id").agg(F.avg("arrays").alias("arrays"))
    df = df.withColumn('arrays', col('arrays')[0].cast('int')) 
    
    return result_df

但是，它仍然只是返回原始数组。我非常感谢这里的任何帮助，提前致谢。

Answer 1

我无法理解你的意思。我尝试运行你的代码。我发现它返回数组的平均值。我的结果如下。 +---+------------------+ |编号| 数组| +---+------------------+ | 1| 34.55| | 3|20.033333333333335| | 2|3.3000000000000003| +---+------------------+