PySpark：基于子串获取数组元素的索引

问题回顾 #

你有一个带数组列col1的DataFrame，想要找出数组里包含子串"58="的元素的索引。示例数据长这样：

+-----------------------------------------------------------+-----+
| col1                                                       |a_pos|
+-----------------------------------------------------------+-----+
|[8=FIX.4.4, 55=ITUBD264, 58=AID[43e39b2e-c6e2-4947]        | 0   |
+-----------------------------------------------------------+-----+

你试过用array_position(col1, "58=")，但它只支持精确匹配，没法处理子串的情况，不过已经在Pandas里用这段代码实现了需求：

df['idx'] = [max(range(len(l)), key=lambda x: '58=' in l[x]) for l in df['col1']]

针对SQL环境的解决方案 #

如果是在Spark SQL或者PostgreSQL这类支持数组操作的SQL环境中，你可以用下面的方法实现相同的逻辑：

1. Spark SQL 实现

方法一：用transform标记匹配项再找位置

这个方法先把数组里每个元素转换成“是否包含子串”的标记，再找第一个匹配项的索引（注意Spark数组索引是1-based）：

SELECT 
  col1,
  a_pos,
  array_position(
    transform(col1, x -> CASE WHEN x LIKE '%58=%' THEN 1 ELSE 0 END),
    1
  ) AS idx
FROM your_table_name

方法二：展开数组后筛选聚合（对应Pandas取最后一个匹配项的逻辑）

如果数组里可能有多个包含"58="的元素，想要取最后一个的索引，可以用posexplode展开数组后聚合：

WITH exploded_data AS (
  SELECT 
    col1,
    a_pos,
    pos,
    element
  FROM your_table_name
  LATERAL VIEW posexplode(col1) exploded AS pos, element
  WHERE element LIKE '%58=%'
)
SELECT 
  col1,
  a_pos,
  MAX(pos + 1) AS idx -- pos是0-based，加1转为1-based索引
FROM exploded_data
GROUP BY col1, a_pos

2. PostgreSQL 实现

PostgreSQL可以用unnest WITH ORDINALITY来展开数组并获取元素的索引，再筛选聚合：

SELECT 
  col1,
  a_pos,
  MAX(ordinal) AS idx
FROM your_table_name,
     unnest(col1) WITH ORDINALITY AS elem(element, ordinal)
WHERE element LIKE '%58=%'
GROUP BY col1, a_pos

内容的提问来源于stack exchange，提问作者user17271845