PySpark: извлечь / собрать первый элемент массива из столбца

Question

У меня есть фрейм данных, который выглядит как

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

element - это массив значений типа double

Я хочу получить эти данные в массив numpy, который у меня наивносделано так:

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

но это дает

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

, который не является простым 1D массивом, как то, что я хочу

Я также пробовал

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?

Я проверил Как извлечь элемент из массива в pyspark , но я не вижу, как там применяется решениев моей ситуации

Answer 1

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