dataframe选择某行某列的元素,如何从类型为Array的DataFrame列中提取出某个元素？...

问题：

'''

问题1：现有一数据表df如下，df有两列，其中values列的值为array类型，

现欲从中取出第二个元素，该如何操作？

'''

df = spark.createDataFrame([('a',[1,2,3]), ('b', [4,5,6])], ['key', 'values'])

df.show()

df.printSchema()

+---+---------+

|key| values|

+---+---------+

| a|[1, 2, 3]|

| b|[4, 5, 6]|

+---+---------+

root

|-- key: string (nullable = true)

|-- values: array (nullable = true)

| |-- element: long (containsNull = true)

'''

解决办法有三个，但是其中有两种都用到了select expression，可以归为一类，第三种方法用的是withColumn()函数。

'''

# 方法一、二

df.select(expr('key'), expr('values[1]')).show()

df.selectExpr('key', 'values[1]').show()

+---+---------+

|key|values[1]|

+---+---------+

| a| 2|

| b| 5|

+---+---------+

+---+---------+

|key|values[1]|

+---+---------+

| a| 2|

| b| 5|

+---+---------+

# 方法三

df.withColumn('c1', df['values'].getItem(1)).drop('values').show()

+---+---+

|key| c1|

+---+---+

| a| 2|

| b| 5|

+---+---+

上面提到的问题是从array列中提取子元素，如果把列的类型从array换成vector是否依然可行呢？答案是不行的。因此，我们需要使用别的方法：(1)现将column类型从vactor转化为array，这样便可以用上面介绍的方法解决了；(2)直接从vector中提取。下面我们分别使用这两种方法：

'''先将上面的array类型的DataFrame转换为vector类型.

参考 https://stackoverflow.com/questions/42138482/how-do-i-convert-an-array-i-e-list-column-to-vector

'''

from pyspark.ml.linalg import Vectors, VectorUDT

from pyspark.sql.functions import udf

list_to_vector_udf = udf(lambda l: Vectors.dense(l), VectorUDT())

df = df.select('key', list_to_vector_udf(df['values']).alias('values'))

df.show()

df.printSchema()

+---+-------------+

|key| values|

+---+-------------+

| a|[1.0,2.0,3.0]|

| b|[4.0,5.0,6.0]|

+---+-------------+

root

|-- key: string (nullable = true)

|-- values: vector (nullable = true)

# 方法一，直接从vector中提取元素, 直接使用udf

# 参考：https://stackoverflow.com/questions/44425159/access-element-of-a-vector-in-a-spark-dataframe-logistic-regression-probability

from pyspark.sql.functions import udf

from pyspark.sql.types import FloatType

firstelement=udf(lambda v:float(v[0]),FloatType())

df.select(firstelement('values').alias('val1')).show()

+----+

|val1|

+----+

| 1.0|

| 4.0|

+----+

# 方法二，先将vector转化为array，再从array中提取

# vector转array参考 https://stackoverflow.com/questions/48397543/how-to-transform-vector-into-array-for-frequent-pattern-analysis