PySpark：从列中提取/收集第一个数组元素

Question

我有一个看起来像的数据框

 |-- alleleFrequencies: array (nullable = true)
 |    |-- element: double (containsNull = true)

element是一个双精度数组

我希望将这些数据放入numpy数组中，因此我很天真地做到了：

allele_freq1 = np.array(df1.select("alleleFrequencies").collect())

但这给出了

[[list([0.5, 0.5])]
 [list([0.5, 0.5])]
 [list([1.0])]...

这不是我想要的简单一维数组

我也尝试过

allele_freq1 = np.array(df1.select("alleleFrequencies")[0].collect())

但这给了

TypeError: 'Column' object is not callable

我也尝试过

allele_freq1 = np.array(df1.select("alleleFrequencies[0]").collect())

但这给了

org.apache.spark.sql.AnalysisException: cannot resolve '`alleleFrequencies[0]`' given input columns...

如何将alleleFrequencies列中的第一项放入numpy数组中？

我检查了How to extract an element from a array in pyspark，但看不到那里的解决方案如何适用于我的情况

Answer 1

allele_freq1 = np.array(df1.select(col("alleleFrequencies").getItem(0)).collect())
print(allele_freq1)
print(type(allele_freq1))