我可以通过执行以下操作,根据数组字段中是否存在特定值来过滤 Spark 数据帧(在 PySpark 中):
from pyspark.sql.functions import array_contains
spark_df.filter(array_contains(spark_df.array_column_name, "value that I want")).show()
有没有办法获取数组中找到该项目的索引? 看起来应该存在,但我没有找到它。 谢谢。
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开始的索引。如果在数组中找不到给定值,则返回 0。
我使用的是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 ])
newdf = df.withColumn('item_position',user_func(df.arraydata,lit('b')))
>>> newdf.show();
+---------------+-------------+
| arraydata|item_position|
+---------------+-------------+
|[c, b, a, e, f]| [1]|
+---------------+-------------+
newdf = df.withColumn('item_position',user_func(df.arraydata,lit('e')))
>>> newdf.show();
+---------------+-------------+
| arraydata|item_position|
+---------------+-------------+
|[c, b, a, e, f]| [3]|
+---------------+-------------+
这可用于查找多个位置:
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]| []|
# +---------+-----------+