Scala dataframe calling function inside withColumn? - scala

Question

Here is what I am trying to do, I have two tables that have exactly the same column names.
Table look somewhat like this:
-----------
A B C D
-----------
1 2 3 4
5 6 3 4
7 8 3 4
The logic of the problem I need to have is, compare A B C D columns in Table1 with Table2. If A,B match each other, return a new column with value 0, else return 0. If C from table A is 3, return 0, else return 1. Only one value should be returned for each row, priority: C>D>A=B.
I joined two tables(dataFrames), result in a combinedDf. This is how I join them: Table1.join(Table2,table1($"A")=table2($"A"))
So here is what I did:
def func(A:mutable.WrappedArray[String],B:mutable.WrappedArray[String],C:String,D:String) =
{if(C=="3") "0";
else if(D=="4")"1";
else if ((0 to A.length-1).exists(i => A(i) == B(i)))"0" else "1"}
For this function I want to put A,B columns from Table1 in to one array and A,B column from Table2 into another array and running a for loop to check the equality. (I need the array because for real case, I have a random number of columns I need to compare).
And here is how I call the function.
combinedDf.withColumn("returnVal",func(array(col("table1.A"),col("table1.B")),
array(col("table2.A"),col("table2.B")),col("table1.C"),col("table1.D")))
But it's just doesn't work, even though I put the columns inside array using the array function its' still telling me type mismatch.
Error Message: <console:67>: error: type mismatch; found:org.apache.spark.Column required: String
Thanks in advance!

Answer 1

You can try this, however help me understand one thing why do you need to combine the dataframes, and what you mean by if A and B matches(my assumption is per row, am i right?), If your A,B,C,D columns are string then change integer to String.
def func(A:Integer,B:Integer,C:Integer,D:Integer) =
{
if(C == 3) "0"
else if(D == 4) "1"
else if (A == B) "0"
else "1"
}
val udfFunc = udf(func _)
combinedDf.withColumn("returnVal",
udfFunc(col("table1.A"), col("table1.B"),
col("table1.C"),col("table1.D")
)
)