Spark DataFrame add column from subquery

user_1996296 0
Viewed 215

using SQL syntax I can add new column using subquery like that:

import spark.sqlContext.implicits._

     List(
      ("a", "1", "2"),
      ("b", "1", "3"),
      ("c", "1", "4"),
      ("d", "1", "5")
    ).toDF("name", "start", "end")
        .createOrReplaceTempView("base")

    List(
      ("a", "1", "2"),
      ("b", "2", "3"),
      ("c", "3", "4"),
      ("d", "4", "5"),
      ("f", "5", "6")
    ).toDF("name", "number", "_count")
      .createOrReplaceTempView("col")


   spark.sql(
     """
       |select a.name,
       |       (select Max(_count) from col b where b.number == a.end) - (select Max(_count) from col b where b.number == a.start) as result
       |from base a
       |""".stripMargin)
      .show(false)

How I can do that with DataFrame API?

apache-spark-sql apache-spark
2 Answers

I found the syntax:

import spark.sqlContext.implicits._

     val b = List(
      ("a", "1", "2"),
      ("b", "1", "3"),
      ("c", "1", "4"),
      ("d", "1", "5")
    ).toDF("name", "start", "end")

    List(
      ("a", "1", "2"),
      ("b", "2", "3"),
      ("c", "3", "4"),
      ("d", "4", "5"),
      ("f", "5", "6")
    ).toDF("name", "number", "_count")
      .createOrReplaceTempView("ref_table")


    b.withColumn("result", expr("((select max(_count) from ref_table r where r.number = end) - (select max(_count) from ref_table r where r.number = start)) as result")).show(false)
user_19962960

I think max is not required, we can follow below approach

    val base = List(
    ("a", "1", "2"),
    ("b", "1", "3"),
    ("c", "1", "4"),
    ("d", "1", "5")
    ).toDF("name", "start", "end")

    val col = List(
    ("a", "1", "2"),
    ("b", "2", "3"),
    ("c", "3", "4"),
    ("d", "4", "5"),
    ("f", "5", "6")
    ).toDF("name", "number", "_count")

    val df = base.join(col, col("number") === base("end")).select(base("name"), col("_count"))

    val df1 = base.join(col, col("number") === base("start")).select(base("name").alias("nameDf"), col("_count").alias("count"))

    df.join(df1, df("name") === df1("nameDf")).select($"name", ($"_count"- $"count").alias("result")).show(false)
user_101863230
Related