Trim String Characters in Pyspark dataframe

user_13495460 0
Viewed 503

Suppose if I have dataframe in which I have the values in a column like :

ABC00909083888
ABC93890380380
XYZ7394949
XYZ3898302
PQR3799_ABZ
MGE8983_ABZ

I want to trim these values like, remove first 3 characters and remove last 3 characters if it ends with ABZ.

00909083888
93890380380
7394949
3898302
3799
8983

Tried some methods but did not work.

from pyspark.sql import functions as f

new_df = df.withColumn("new_column", f.when((condition on some column),
                                               f.substring('Existing_COL', 4, f.length(f.col("Existing_COL"))), ))

Can anyone please tell me which function I can use in pyspark. Trim only removes white space or tab something characters.

pyspark python
2 Answers

Based upon your input and expected output. See below logic -

from pyspark.sql.functions import *

df = spark.createDataFrame(data = [("ABC00909083888",) ,("ABC93890380380",) ,("XYZ7394949",) ,("XYZ3898302",) ,("PQR3799_ABZ",) ,("MGE8983_ABZ",)], schema = ["values",])

(df.withColumn("new_vals", when(col('values').rlike("(_ABZ$)"), regexp_replace(col('values'),r'(_ABZ$)', '')).otherwise(col('values')))
  .withColumn("final_vals", expr(("substring(new_vals, 4 ,length(new_vals))")))
).show()

Output

+--------------+--------------+-----------+
|        values|      new_vals| final_vals|
+--------------+--------------+-----------+
|ABC00909083888|ABC00909083888|00909083888|
|ABC93890380380|ABC93890380380|93890380380|
|    XYZ7394949|    XYZ7394949|    7394949|
|    XYZ3898302|    XYZ3898302|    3898302|
|   PQR3799_ABZ|       PQR3799|       3799|
|   MGE8983_ABZ|       MGE8983|       8983|
+--------------+--------------+-----------+
user_151125630

If I get you correctly and if you don't insist on using pyspark substring or trim functions, you can easily define a function to do what you want and then make use of that with udfs in spark:

from pyspark.sql.types import StringType
from pyspark.sql.functions import udf


def mysub(word):
  if word.endswith('_ABZ'):
    word = word[:-4]
  return word[3:]


udf1 = udf(lambda x: mysub(x), StringType())
df.withColumn('new_label',udf1('label')).show()

The output will be like:

+---+--------------+-----------+
| id|         label|  new_label|
+---+--------------+-----------+
|  1|ABC00909083888|00909083888|
|  2|ABC93890380380|93890380380|
|  3|    XYZ7394949|    7394949|
|  4|    XYZ3898302|    3898302|
|  5|   PQR3799_ABZ|       3799|
|  6|   MGE8983_ABZ|       8983|
+---+--------------+-----------+

Please let me know if I got you wrong in some cases.

user_89611700
Related