Partially replace dataframe with rows from another

Viewed 105

I have two similar dataframes, one has a single date and the other has multiple dates plus an additional column:

df:

| yyyy_mm_dd | id  | region | country  | product | count |
|------------|-----|--------|----------|---------|-------|
| 2021-06-14 | 111 | EMEA   | Spain    | P1      | 10    |
| 2021-06-14 | 111 | EMEA   | England  | P1      | 9     |
| 2021-06-14 | 111 | EMEA   | France   | P1      | 10    |
| 2021-06-14 | 111 | EMEA   | Spain    | P2      | 299   |
| 2021-06-14 | 111 | EMEA   | England  | P2      | 39    |
| 2021-06-14 | 111 | EMEA   | France   | P2      | 10    |
| 2021-06-14 | 112 | LATAM  | Brazil   | P1      | 64    |
| 2021-06-14 | 112 | LATAM  | Paraguay | P2      | 21    |
| 2021-06-14 | ... | ...    | ...      | ...     | ...   |

df1:

| yyyy_mm_dd | id  | region | country  | product | count | fullfilments |
|------------|-----|--------|----------|---------|-------|--------------|
| 2021-06-14 | 111 | EMEA   | Spain    | P1      | 1     | 1            |
| 2021-06-14 | 111 | EMEA   | England  | P1      | 1     | 3            |
| 2021-06-14 | 111 | EMEA   | France   | P1      | 2     | 4            |
| 2021-06-14 | 111 | EMEA   | Spain    | P2      | 1     | 1            |
| 2021-06-14 | 111 | EMEA   | England  | P2      | 2     | 1            |
| 2021-06-14 | 111 | EMEA   | France   | P2      | 1     | 5            |
| 2021-06-14 | 112 | LATAM  | Brazil   | P1      | 2     | 2            |
| 2021-06-14 | 112 | LATAM  | Paraguay | P2      | 21    | 1            |
| 2021-06-14 | ... | ...    | ...      | ...     | ...   | ...          |
| 2021-06-13 | 111 | EMEA   | Spain    | P1      | 0     | 1            |
| 2021-06-13 | 111 | EMEA   | England  | P2      | 0     | 2            |

Df1 has many dates of grouped data and df only has one date. I would like to replace the count column in df1 with the count in df for matching rows (yyyy_mm_dd, id, region, country, product) and retain fullfilments.

I could probably join both together and drop count in the first df, however I only want to replace where the date is matching and retain all other rows in df1.

2 Answers

You can simply join and use the coalesce function.

When you do the left join from the first dataframe to the second, the only matching records have the not null new_count value. Now, use the coalesce function that will return the first value when it is not null but the second value when the first is null.

coalesce(a   , b   ) => a
coalesce(a   , null) => a
coalesce(null, b   ) => b

From your dataframes,

from pyspark.sql import functions as f

df1 = spark.read.option("inferSchema","true").option("header","true").csv("test1.csv")

+----------+---+------+--------+-------+-----+
|yyyy_mm_dd|id |region|country |product|count|
+----------+---+------+--------+-------+-----+
|2021-06-14|111|EMEA  |Spain   |P1     |10   |
|2021-06-14|111|EMEA  |England |P1     |9    |
|2021-06-14|111|EMEA  |France  |P1     |10   |
|2021-06-14|111|EMEA  |Spain   |P2     |299  |
|2021-06-14|111|EMEA  |England |P2     |39   |
|2021-06-14|111|EMEA  |France  |P2     |10   |
|2021-06-14|112|LATAM |Brazil  |P1     |64   |
|2021-06-14|112|LATAM |Paraguay|P2     |21   |
+----------+---+------+--------+-------+-----+

df2 = spark.read.option("inferSchema","true").option("header","true").csv("test2.csv")

+----------+---+------+--------+-------+-----+------------+
|yyyy_mm_dd|id |region|country |product|count|fullfilments|
+----------+---+------+--------+-------+-----+------------+
|2021-06-14|111|EMEA  |Spain   |P1     |1    |1           |
|2021-06-14|111|EMEA  |England |P1     |1    |3           |
|2021-06-14|111|EMEA  |France  |P1     |2    |4           |
|2021-06-14|111|EMEA  |Spain   |P2     |1    |1           |
|2021-06-14|111|EMEA  |England |P2     |2    |1           |
|2021-06-14|111|EMEA  |France  |P2     |1    |5           |
|2021-06-14|112|LATAM |Brazil  |P1     |2    |2           |
|2021-06-14|112|LATAM |Paraguay|P2     |21   |1           |
|2021-06-13|111|EMEA  |Spain   |P1     |0    |1           |
|2021-06-13|111|EMEA  |England |P2     |0    |2           |
+----------+---+------+--------+-------+-----+------------+

the join of two dataframes are given by follows:

cols_to_join = ['yyyy_mm_dd', 'id', 'region', 'country', 'product']
df3 = df2.join(df1.withColumnRenamed('count', 'new_count'), cols_to_join, 'left') \
         .withColumn('count', f.coalesce('new_count', 'count')).drop('new_count')
df3.show(truncate=False)

+----------+---+------+--------+-------+-----+------------+
|yyyy_mm_dd|id |region|country |product|count|fullfilments|
+----------+---+------+--------+-------+-----+------------+
|2021-06-14|111|EMEA  |Spain   |P1     |10   |1           |
|2021-06-14|111|EMEA  |England |P1     |9    |3           |
|2021-06-14|111|EMEA  |France  |P1     |10   |4           |
|2021-06-14|111|EMEA  |Spain   |P2     |299  |1           |
|2021-06-14|111|EMEA  |England |P2     |39   |1           |
|2021-06-14|111|EMEA  |France  |P2     |10   |5           |
|2021-06-14|112|LATAM |Brazil  |P1     |64   |2           |
|2021-06-14|112|LATAM |Paraguay|P2     |21   |1           |
|2021-06-13|111|EMEA  |Spain   |P1     |0    |1           |
|2021-06-13|111|EMEA  |England |P2     |0    |2           |
+----------+---+------+--------+-------+-----+------------+

Everytime you need to retrieve a column from different dataframes you must join them:

import pyspark.sql.functions as f

df2 = df1.join(df.withColumnRenamed('count', 'new_count'),
               on=['yyyy_mm_dd', 'id', 'region', 'country', 'product'], how='left')

df2 = (df2
       .withColumn('count', f.coalesce('new_count', 'count'))
       .drop('new_count'))
df2.show(truncate=False)
Related