Как удалить подстроку символов из столбца PySpark Dataframe StringType (), условно основываясь на длине строк в столбцах?

Question

У меня есть PySpark Dataframe со столбцом StringType (), который состоит в основном из 15 символов. Тем не менее, некоторые строки имеют 11 символов. Пример:

df = 
+--------------+--------+
|             code|state|
+--------------+--------+
|'334445532234553'|wa   |
|'332452132234553'|mn   |
|'45532234553'    |fl   |
|'679645532234553'|mo   |
|'918535532234553'|ar   |
|'174925532234553'|wi   |
|'45532234553'    |al   |
|'928405532234553'|ca   |
+--------------+--------+

Мне нужно, чтобы во всех строках было по 11 символов, а последние 4 символа были удалены из любой строки, содержащей 15 символов. Так что это мой желаемый вывод:

df.show(8) = 
+-------------+-----+
|         code|state|
+-------------+-----+
|'33444553223'|wa   |
|'33245213223'|mn   |
|'45532234553'|fl   |
|'67964553223'|mo   |
|'91853553223'|ar   |
|'17492553223'|wi   |
|'45532234553'|al   |
|'92840553223'|ca   |
+-------------+-----+

Пока у меня есть это преобразование, которое удаляет последние 4 символа из ВСЕХ строк в моем столбце под названием 'code':

from pyspark.sql.functions import substring, length, col, expr

df = df.withColumn("code",expr("substring(code, 1, length(code)-4)"))

Так что мне нужно что-то сделать, чтобы сделать это условным для длины строки в строке.

РЕДАКТИРОВАТЬ С помощью @gmds я нашел это решение:

df.withColumn("code",expr("substring(code, 1, 11)"))

Answer 1

Вы можете использовать встроенную функцию длины вместе с подстрокой:

from pyspark.sql.functions import substring, length

df = df.withColumn("code", when(length(df.code) > 11, substring(df.code, 1, length(df.code) - 4)).otherwise(df.code))

Answer 2

Как насчет этого:

df.withColumn('code', df['code'].substr(1, 11))

Ваша идея была правильной; просто вы указали изменяющееся значение для длины подстроки, когда вам действительно нужна константа.

Answer 3

Я нашел ответ на свой вопрос, основываясь на информации, полученной от @gmds.Вот оно:

df.withColumn("code",expr("substring(code, 1, 11)"))

Этот код устанавливает 11 как константу, что означает, что независимо от того, содержит ли столбец значение длиной 11 или 15 символов, после преобразования все они будут иметь константу 11 символов.