How to extract characters from a left of a substring and right of the same substring in PySpark column?

user_8522409 0
Viewed 528

My Pyspark dataframe is something like this:

|ID|A|
+--+-------+
|1|7800028|
|2|700024|
|3|720004|
|4|70004|
|5|700004|

I want to remove the 3 zeros occurring together and get numbers to the left and right of the three zeroes in separate columns. Something like this:

|ID|B|C
+--+-------+
|1|78|28|
|2|7|24|
|3|72|4|
|4|7|4|
|5|70|4

The problem is col A can be of varied length with values in B ranging from 0-99 and values in C ranging from 0-99. Therefore I can't seem to use substring to get B. C is still doable through substring function.

pyspark substring python
2 Answers

Use the PySpark split() function to split the values in column "A". Reference: https://spark.apache.org/docs/latest/api/python/pyspark.sql.html?highlight=split#pyspark.sql.functions.split

data = [[1, 7800028], [2, 700024], [3, 720004], [4, 70004]]
data_df = spark.createDataFrame(data, ["ID", "A"])

data_df.show()

+---+-------+-------+
| ID|      A|    new|
+---+-------+-------+
|  1|7800028|7800028|
|  2| 700024| 700024|
|  3| 720004| 720004|
|  4|  70004|  70004|
+---+-------+-------+

from pyspark.sql import functions as F

new_df = (data_df
           .withColumn("B", F.split(data_df["A"].cast("string"), "000")[0])
           .withColumn("C", F.split(data_df["A"].cast("string"), "000")[1])
          )

new_df.show()

+---+-------+---+---+
| ID|      A|  B|  C|
+---+-------+---+---+
|  1|7800028| 78| 28|
|  2| 700024|  7| 24|
|  3| 720004| 72|  4|
|  4|  70004|  7|  4|
+---+-------+---+---+
user_125773940

You can use the pattern 000(?=[1-9]|0$) to split the string, where (?=[1-9]|0$) is an anchor to make sure that the last 0 must either be followed by a non-zero digit or zero as the last digit of the whole number, example:

spark.sql("""
    with t as ( 
        select *, split(A,'000(?=[1-9]|0$)') as arr
        from values (1,7800028),(2,700024),(3,720004),(4,70004),(5,700004),(6,120000) as (id,A)
    ) 
    select id, A, arr[0] as B, arr[1] as C 
    from t
""").show()
+---+-------+---+---+
| id|      A|  B|  C|
+---+-------+---+---+
|  1|7800028| 78| 28|
|  2| 700024|  7| 24|
|  3| 720004| 72|  4|
|  4|  70004|  7|  4|
|  5| 700004| 70|  4|
|  6| 120000| 12|  0|
+---+-------+---+---+
user_95107290
Related