Fetching values from multiple columns in Pyspark

user_10215871 0
Viewed 68

Kindly have a look at input and expected output and suggest me efficient way of dealing it in Pyspark.

Input : A dataframe with

C1 x y  
11 a  1  
12 a  2  
13 b  5  
14 b  6  
16 c  &b  
17 c  7  
18 c  b

Wher C1 , x and y are columns

Expected output :

C1 x  y  z  
11 a  1  [1]  
12 a  2  [2]  
13 b  5  [5]  
14 b  6  [6]  
16 c  &b [5,6]  
17 c  7  [7]  
18 c  b  [5,6]

'z' is the expected column where it should fetch corresponding 'b' values for 'c'

It should fetch [5,6] for both &b and b in 5th row for 'c' .

Thanks in advance

pyspark python
1 Answers

Try this, self-join with collected list on rlike join condition is the way to go.

df.show() #sampledataframe

#+---+---+---+
#| C1|  x|  y|
#+---+---+---+
#| 11|  a|  1|
#| 12|  a|  2|
#| 15|  b|  5|
#| 16|  b|  6|
#| 17|  c| &b|
#| 17|  c|  7|
#| 18|  c|  b|
#+---+---+---+

df.join(df.groupBy("x").agg(F.collect_list("y").alias("y1"))\
          .withColumnRenamed("x","x1"),F.expr("""y rlike x1"""),'left')\
  .withColumn("z", F.when(F.col("y").cast("int").isNotNull(), F.array("y"))\
                    .otherwise(F.col("y1"))).drop("x1","y1").show()

#+---+---+---+------+
#| C1|  x|  y|     z|
#+---+---+---+------+
#| 11|  a|  1|   [1]|
#| 12|  a|  2|   [2]|
#| 15|  b|  5|   [5]|
#| 16|  b|  6|   [6]|
#| 17|  c| &b|[5, 6]|
#| 17|  c|  7|   [7]|
#| 18|  c|  b|[5, 6]|
#+---+---+---+------+
user_101073890
Related