PySpark: Concat two dataframes with columns sums

Viewed 565

I have two PySpark dataframes which I would like to left join

Prev_table:
| user_id | earnings | start_date | end_date   |
|---------|--------|------------|------------|
| 1       | 10     | 2020-06-01 | 2020-06-10 |
| 2       | 20     | 2020-06-01 | 2020-06-10 |
| 3       | 30     | 2020-06-01 | 2020-06-10 |

New_table:
| user_id | profit |
|---------|--------|
| 1       | 100    |
| 2       | 200    |
| 5       | 500    |

The resultant table is a concatanation of the two dataframes grouped by the user_id. And the earnings column is Prev_table.earnings + New_table.profit.

| user_id | earnings | start_date | end_date   |
|---------|--------|------------|------------|
| 1       | 110    | 2020-06-01 | 2020-06-10 |
| 2       | 220    | 2020-06-01 | 2020-06-10 |
| 3       | 30     | 2020-06-01 | 2020-06-10 |
| 5       | 500    |            |            |

This is what Pandas concat would do, but in case of PySpark I believe I need to use PySparks union? Also I not sure how to do the sum of the two columns. I know I need to use something like combined_df.agg({"earnings": "sum"}).collect() but having trouble getting this to work.

Would appreciate any guidance on the workflow for this with PySpark. Thank you.

3 Answers

There can be a better way but one approach is , rename the profit to earnings, then fill missing columns in df2 , then union and groupby with agg:

Assuming Prev_table is df1 and New_table is df2

import pyspark.sql.functions as F

df3 = df2.select("user_id",F.col("profit").alias("earnings"))

(df1.union(df3.select("*",*[F.lit(None).alias(i) 
            for i in df1.columns if i not in df3.columns]))
.groupBy("user_id").agg(F.sum("earnings").alias("earnings")
 ,F.first("start_date",ignorenulls=True).alias("start_date")
 ,F.first("end_date",ignorenulls=True).alias("end_date")).orderBy("user_id")).show()

+-------+--------+----------+----------+
|user_id|earnings|start_date|  end_date|
+-------+--------+----------+----------+
|      1|     110|2020-06-01|2020-06-10|
|      2|     220|2020-06-01|2020-06-10|
|      3|      30|2020-06-01|2020-06-10|
|      5|     500|      null|      null|
+-------+--------+----------+----------+

Tried using join operations and then union all

import pyspark.sql.functions as f

import pyspark.sql.functions as f

df = df1.join(df2, df1.user_id==df2.user_id, how='left').select(df1.user_id, (df1.earnings.cast('int')+f.when(df2.profit.cast('int').isNull(),0).otherwise(df2.profit)).alias('earnings'),df1.start_date,df1.end_date)
df3 = df2.join(df, df.user_id==df2.user_id, how='leftanti').select(df2.user_id,df2.profit.alias('earnings'),f.lit(None).alias('start_date'),f.lit(None).alias('end_date'))
final_df = df.union(df3)
final_df.show()

+-------+--------+----------+----------+
|user_id|earnings|start_date|  end_date|
+-------+--------+----------+----------+
|      3|    30.0|2020-06-01|2020-06-10|
|      1|   110.0|2020-06-01|2020-06-10|
|      2|   220.0|2020-06-01|2020-06-10|
|      5|     500|      null|      null|
+-------+--------+----------+----------+

Alternative to other solutions(In scala, but should be used in pyspark with minimal change)-

Load the input provided

  val data1 =
      """
        |user_id | earnings | start_date | end_date
        |1       | 10     | 2020-06-01 | 2020-06-10
        |2       | 20     | 2020-06-01 | 2020-06-10
        |3       | 30     | 2020-06-01 | 2020-06-10
      """.stripMargin
    val stringDS1 = data1.split(System.lineSeparator())
      .map(_.split("\\|").map(_.replaceAll("""^[ \t]+|[ \t]+$""", "")).mkString(","))
      .toSeq.toDS()
    val df1 = spark.read
      .option("sep", ",")
      .option("inferSchema", "true")
      .option("header", "true")
      .option("nullValue", "null")
      .csv(stringDS1)
    df1.show(false)
    df1.printSchema()
    /**
      * +-------+--------+-------------------+-------------------+
      * |user_id|earnings|start_date         |end_date           |
      * +-------+--------+-------------------+-------------------+
      * |1      |10      |2020-06-01 00:00:00|2020-06-10 00:00:00|
      * |2      |20      |2020-06-01 00:00:00|2020-06-10 00:00:00|
      * |3      |30      |2020-06-01 00:00:00|2020-06-10 00:00:00|
      * +-------+--------+-------------------+-------------------+
      *
      * root
      * |-- user_id: integer (nullable = true)
      * |-- earnings: integer (nullable = true)
      * |-- start_date: timestamp (nullable = true)
      * |-- end_date: timestamp (nullable = true)
      */

    val data2 =
      """
        |user_id | profit
        |1       | 100
        |2       | 200
        |5       | 500
      """.stripMargin
    val stringDS2 = data2.split(System.lineSeparator())
      .map(_.split("\\|").map(_.replaceAll("""^[ \t]+|[ \t]+$""", "")).mkString(","))
      .toSeq.toDS()
    val df2 = spark.read
      .option("sep", ",")
      .option("inferSchema", "true")
      .option("header", "true")
      .option("nullValue", "null")
      .csv(stringDS2)
    df2.show(false)
    df2.printSchema()

    /**
      * +-------+------+
      * |user_id|profit|
      * +-------+------+
      * |1      |100   |
      * |2      |200   |
      * |5      |500   |
      * +-------+------+
      *
      * root
      * |-- user_id: integer (nullable = true)
      * |-- profit: integer (nullable = true)
      */

Join the table and get the required columns in mentioned specifics

  df1.createOrReplaceTempView("prev_table")
    df2.createOrReplaceTempView("new_table")

   val processedDF = spark.sql(
      """
        |SELECT coalesce(p.user_id, n.user_id) as user_id,
        |       (coalesce(earnings,0) + coalesce(profit, 0)) as earnings,
        |        start_date,
        |        end_date
        |FROM prev_table p FULL OUTER JOIN new_table n ON p.user_id=n.user_id
      """.stripMargin)

     processedDF.orderBy("user_id").show(false)

    /**
      * +-------+--------+-------------------+-------------------+
      * |user_id|earnings|start_date         |end_date           |
      * +-------+--------+-------------------+-------------------+
      * |1      |110     |2020-06-01 00:00:00|2020-06-10 00:00:00|
      * |2      |220     |2020-06-01 00:00:00|2020-06-10 00:00:00|
      * |3      |30      |2020-06-01 00:00:00|2020-06-10 00:00:00|
      * |5      |500     |null               |null               |
      * +-------+--------+-------------------+-------------------+
      */
Related