Posso usare regexp_replace o qualcosa di equivalente per sostituire più valori in una colonna dataframe pyspark con una riga di codice?

Question

Posso usare regexp_replace o qualcosa di equivalente per sostituire più valori in una colonna dataframe pyspark con una riga di codice?

Ecco il codice per creare il mio dataframe:

from pyspark import SparkContext, SparkConf, SQLContext
from datetime import datetime

sc = SparkContext().getOrCreate()
sqlContext = SQLContext(sc)

data1 = [
  ('George', datetime(2010, 3, 24, 3, 19, 58), 13),
  ('George', datetime(2020, 9, 24, 3, 19, 6), 8),
  ('George', datetime(2009, 12, 12, 17, 21, 30), 5),
  ('Micheal', datetime(2010, 11, 22, 13, 29, 40), 12),
  ('Maggie', datetime(2010, 2, 8, 3, 31, 23), 8),
  ('Ravi', datetime(2009, 1, 1, 4, 19, 47), 2),
  ('Xien', datetime(2010, 3, 2, 4, 33, 51), 3),
]
 
df1 = sqlContext.createDataFrame(data1, ['name', 'trial_start_time', 'purchase_time'])
df1.show(truncate=False)

Ecco il dataframe:

+-------+-------------------+-------------+
|name   |trial_start_time   |purchase_time|
+-------+-------------------+-------------+
|George |2010-03-24 07:19:58|13           |
|George |2020-09-24 07:19:06|8            |
|George |2009-12-12 22:21:30|5            |
|Micheal|2010-11-22 18:29:40|12           |
|Maggie |2010-02-08 08:31:23|8            |
|Ravi   |2009-01-01 09:19:47|2            |
|Xien   |2010-03-02 09:33:51|3            |
+-------+-------------------+-------------+

Ecco un esempio funzionante per sostituire una stringa:

from pyspark.sql.functions import regexp_replace, regexp_extract, col
df1.withColumn("name", regexp_replace('name', "Ravi", "Ravi_renamed")).show()

Ecco l'output:

+------------+-------------------+-------------+
|        name|   trial_start_time|purchase_time|
+------------+-------------------+-------------+
|      George|2010-03-24 07:19:58|           13|
|      George|2020-09-24 07:19:06|            8|
|      George|2009-12-12 22:21:30|            5|
|     Micheal|2010-11-22 18:29:40|           12|
|      Maggie|2010-02-08 08:31:23|            8|
|Ravi_renamed|2009-01-01 09:19:47|            2|
|        Xien|2010-03-02 09:33:51|            3|
+------------+-------------------+-------------+

In Panda potrei sostituire più stringhe in una riga di codice con un'espressione lambda:

df1[name].apply(lambda x: x.replace('George','George_renamed1').replace('Ravi', 'Ravi_renamed2')

Non sono sicuro che ciò possa essere fatto in pyspark con regexp_replace. Forse un'altra alternativa? Quando ho letto sull'uso delle espressioni lambda in pyspark sembra che devo creare funzioni udf (che sembrano diventare un po 'lunghe). Ma sono curioso di poter semplicemente eseguire qualche tipo di espressione regex su più stringhe come sopra in una riga di codice.

Answer 1

Questo è quello che stai cercando:

Utilizzo when()(più leggibile)

df1.withColumn('name', 
               when(col('name') == 'George', 'George_renamed1')
               .when(col('name') == 'Ravi', 'Ravi_renamed2')
               .otherwise(col('name'))
              )

Con mapping expr (meno esplicito ma utile se ci sono molti valori da sostituire)

df1 = df1.withColumn('name', F.expr("coalesce(map('George', 'George_renamed1', 'Ravi', 'Ravi_renamed2')[name], name)"))

o se hai già un elenco da utilizzare, ad es name_changes = ['George', 'George_renamed1', 'Ravi', 'Ravi_renamed2']

# str()[1:-1] to convert list to string and remove [ ]
df1 = df1.withColumn('name', expr(f'coalesce(map({str(name_changes)[1:-1]})[name], name)'))

quanto sopra ma solo usando le funzioni importate da pyspark

mapping_expr = create_map([lit(x) for x in name_changes])

df1 = df1.withColumn('name', coalesce(mapping_expr[df1['name']], 'name'))

Risultato

df1.withColumn('name', F.expr("coalesce(map('George', 'George_renamed1', 'Ravi', 'Ravi_renamed2')[name],name)")).show()
+---------------+-------------------+-------------+
|           name|   trial_start_time|purchase_time|
+---------------+-------------------+-------------+
|George_renamed1|2010-03-24 03:19:58|           13|
|George_renamed1|2020-09-24 03:19:06|            8|
|George_renamed1|2009-12-12 17:21:30|            5|
|        Micheal|2010-11-22 13:29:40|           12|
|         Maggie|2010-02-08 03:31:23|            8|
|  Ravi_renamed2|2009-01-01 04:19:47|            2|
|           Xien|2010-03-02 04:33:51|            3|
+---------------+-------------------+-------------+