在Scala中将Spark数据框列与其行连接起来

Question

I am trying build a string by concatenating values from a data frame. For example:

val df = Seq(
  ("20181001","10"),     
  ("20181002","40"),
  ("20181003","50")).toDF("Date","Key")
df.show

Output of the DF is like below.

Here I want to build the condition based on the values of the data frame, such as: (Date=20181001 and key=10) or (Date=20181002 and key=40) or (Date=20181003 and key=50) and so on..The generated condition will serve as input for another process. Here the columns in the data frame can be dynamic.

The below snippet I am trying and it is forming the string as needed but its a static one.Also not very sure how it will perform when i have to generate the condition for more than 10 columns. Any suggestion is highly appreciated.

val df = Seq(
  ("20181001","10"),     
  ("20181002","40"),
  ("20181003","50")).toDF("Date","Key")

val colList = df.columns
var cond1 = ""
var finalCond =""
for (row <- df.rdd.collect)
 {
    cond1 = "("
    var pk = row.mkString(",").split(",")(0)
    cond1 = cond1+colList(0)+"="+pk
    var ak = row.mkString(",").split(",")(1)
    cond1 = cond1 +" and " + colList(1)+ "=" +ak +")"
    finalCond = finalCond + cond1 + " or " 
    cond1= ""    
 }
 print("Condition:" +finalCond.dropRight(3))

Answer 1

Check this DF solution.

scala> val df = Seq(
       |   ("20181001","10"),
       |   ("20181002","40"),
       |   ("20181003","50")).toDF("Date","Key")
df: org.apache.spark.sql.DataFrame = [Date: string, Key: string]

scala> val df2 = df.withColumn("gencond",concat(lit("(Date="), 'Date, lit(" and Key=") ,'Key,lit(")")))
df2: org.apache.spark.sql.DataFrame = [Date: string, Key: string ... 1 more field]


scala> df2.agg(collect_list('gencond)).show(false)
+------------------------------------------------------------------------------------+
|collect_list(gencond)                                                               |
+------------------------------------------------------------------------------------+
|[(Date=20181001 and Key=10), (Date=20181002 and Key=40), (Date=20181003 and Key=50)]|
+------------------------------------------------------------------------------------+

EDIT1

You can read them from parquet files and just change the names as in this solution. In the final step, again replace the names from the parquet header. Check this.

scala> val df = Seq(("101","Jack"),("103","wright")).toDF("id","name")  // Original names from parquet
df: org.apache.spark.sql.DataFrame = [id: string, name: string]

scala> val df2= df.select("*").toDF("Date","Key")  // replace it with Date/Key as we used in this question
df2: org.apache.spark.sql.DataFrame = [Date: string, Key: string]

scala> val df3 = df2.withColumn("gencond",concat(lit("(Date="), 'Date, lit(" and Key=") ,'Key,lit(")")))
df3: org.apache.spark.sql.DataFrame = [Date: string, Key: string ... 1 more field]

scala> val df4=df3.agg(collect_list('gencond).as("list"))
df4: org.apache.spark.sql.DataFrame = [list: array<string>]

scala> df4.select(concat_ws(" or ",'list)).show(false)
+----------------------------------------------------+
|concat_ws( or , list)                               |
+----------------------------------------------------+
|(Date=101 and Key=Jack) or (Date=103 and Key=wright)|
+----------------------------------------------------+

scala> val a = df.columns(0)
a: String = id

scala> val b = df.columns(1)
b: String = name

scala>  df4.select(concat_ws(" or ",'list).as("new1")).select(regexp_replace('new1,"Date",a).as("colx")).select(regexp_replace('colx,"Key",b).as("colxy")).show(false)
+--------------------------------------------------+
|colxy                                             |
+--------------------------------------------------+
|(id=101 and name=Jack) or (id=103 and name=wright)|
+--------------------------------------------------+


scala>

Answer 2

Calling collect pulls the result back to the driver program so if you have a huge DataFrame you may well run out of memory.

If you are sure you are only dealing with a smallish number of rows that isn't a problem.

You could do something like:

df.map(row => s"($Date={row.getString(0)} and Key=${row.getString(1)})").collect.mkString("Condition: ", " or ", "")

Output:

res2: String = Condition: (Date=20181001 and Key=10) or (Date=20181002 and Key=40) or (Date=20181003 and Key=50)

Answer 3

Using udf you can do for variable number of columns like below

val list=List("Date","Key")

def getCondString(row:Row):String={
    "("+list.map(cl=>cl+"="+row.getAs[String](cl)).mkString(" and ")+")"
  }

val getCondStringUDF=udf(getCondString _)
df.withColumn("row", getCondStringUDF(struct(df.columns.map(df.col(_)):_*))).select("row").rdd.map(_(0).toString()).collect().mkString(" or ")