PySpark 2.2中数组列的每个元素的子字符串

Question

我想在PySpark 2.2中对数组列的每个元素进行子串。我的df看起来像下面的一个，它是 与this类似，尽管我的df中的每个元素在连字符分隔符之前具有相同的长度

+---------------------------------+----------------------+
|col1                             |new_column            |
+---------------------------------+----------------------+
|[hello-123, abcde-111]           |[hello, abcde]        |
|[hello-234, abcde-221, xyzhi-333]|[hello, abcde, xyzhi] |
|[hiiii-111, abbbb-333, xyzhu-222]|[hiiii, abbbb, xyzhu] |
+---------------------------------+----------------------+

我试图根据this答案调整前面问题中的udf，以获得上面new_column中的输出，但到目前为止运气不佳。有没有办法在PySpark 2.2中实现这一点

import pyspark.sql.functions as F
import pyspark.sql.types as T 

cust_udf = F.udf(lambda arr: [x[0:4] for x in arr], T.ArrayType(T.StringType()))
df1.withColumn('new_column', cust_udf(col("col1")))

Answer 1

你的udf方法适合我。此外，您可以将transform与substring一起使用：

import pyspark.sql.functions as f

df.withColumn('new_column', f.expr('transform(col1, x -> substring(x, 0, 5))')).show()

+          +          +
|                col1|          new_column|
+          +          +
|[hello-123, abcde...|      [hello, abcde]|
|[hello-234, abcde...|[hello, abcde, xy...|
|[hiiii-111, abbbb...|[hiiii, abbbb, xy...|
+          +          +

Answer 2

使用不同的方法解决了这个问题：分解数组，对元素进行子串，然后收集回数组

import pyspark.sql.functions as F
    
df1\
   .withColumn('idx', F.monotonically_increasing_id())\
   .withColumn('exploded_col', F.explode(col('col1')))\
   .withColumn('substr_col', F.substring(col('exploded_col'),1,5))\
   .groupBy(col('idx'))\
   .agg(F.collect_set('substr_col').alias('new_column'))