获取数组中项目的索引,该数组是 Spark 数据帧中的一列

问题描述 投票:0回答:3

我可以通过执行以下操作,根据数组字段中是否存在特定值来过滤 Spark 数据帧(在 PySpark 中):

from pyspark.sql.functions import array_contains
spark_df.filter(array_contains(spark_df.array_column_name, "value that I want")).show() 

有没有办法获取数组中找到该项目的索引? 看起来应该存在,但我没有找到它。 谢谢。

apache-spark pyspark
3个回答
14
投票

在 Spark 2.4+ 中,有

array_position
功能:

df = spark.createDataFrame([(["c", "b", "a"],), ([],)], ['data'])
df.show()
#+---------+
#|     data|
#+---------+
#|[c, b, a]|
#|       []|
#+---------+

from pyspark.sql.functions import array_position
df.select(df.data, array_position(df.data, "a").alias('a_pos')).show()
#+---------+-----+
#|     data|a_pos|
#+---------+-----+
#|[c, b, a]|    3|
#|       []|    0|
#+---------+-----+

文档注释:

  1. 在给定数组中定位仅第一次出现给定值的位置;

  2. 该位置不是从零开始的索引,而是从1开始的索引。如果在数组中找不到给定值,则返回 0。


3
投票

我使用的是spark 2.3版本,所以我尝试使用udf。

df = spark.createDataFrame([(["c", "b", "a","e","f"],)], ['arraydata'])
+---------------+
|      arraydata|
+---------------+
|[c, b, a, e, f]|
+---------------+

user_func = udf (lambda x,y: [i for i, e in enumerate(x) if e==y ])

检查项目“b”的索引位置:

newdf = df.withColumn('item_position',user_func(df.arraydata,lit('b')))

>>> newdf.show();
+---------------+-------------+
|      arraydata|item_position|
+---------------+-------------+
|[c, b, a, e, f]|          [1]|
+---------------+-------------+

检查项目“e”的索引位置:

newdf = df.withColumn('item_position',user_func(df.arraydata,lit('e')))

>>> newdf.show();
+---------------+-------------+
|      arraydata|item_position|
+---------------+-------------+
|[c, b, a, e, f]|          [3]|
+---------------+-------------+

0
投票

这可用于查找多个位置:

F.array_compact(F.transform('data', lambda x, i: F.when(x == 'b', i)))

结果是零基础的。


完整示例(Spark 3.4+):

from pyspark.sql import functions as F
df = spark.createDataFrame([(['a', 'b', 'b'],), (['c'],)], ['data'])
df.show()
# +---------+
# |     data|
# +---------+
# |[a, b, b]|
# |      [c]|
# +---------+

df = df.withColumn(
    'positions_b',
    F.array_compact(F.transform('data', lambda x, i: F.when(x == 'b', i)))
)
df.show()
# +---------+-----------+
# |     data|positions_b|
# +---------+-----------+
# |[a, b, b]|     [1, 2]|
# |      [c]|         []|
# +---------+-----------+
© www.soinside.com 2019 - 2024. All rights reserved.