I have two DataFrame
in my spark (v1.5.0) code:
aDF = [user_id : Int, user_purchases: array<int> ]
bDF = [user_id : Int, user_purchases: array<int> ]
What I want to do is to join these two dataframes, but I only need the lines where the intersection between aDF.user_purchases
and bDF.user_purchases
has more than 2 elements (intersection > 2).
Do I have to use RDD API or is it possible to use some function from org.apache.sql.functions ?
I don't see any function built-in, but you can use UDF:
import scala.collection.mutable.WrappedArray;
val intersect = udf ((a : WrappedArray[Int], b : WrappedArray[Int]) => {
var count = 0;
a.foreach (x => {
if (b.contains(x)) count = count + 1;
});
count;
});
// test data sets
val one = sc.parallelize(List(
(1, Array(1, 2, 3)),
(2, Array(1,2 ,3, 4)),
(3, Array(1, 2,3)),
(4, Array(1,2))
)).toDF("user", "arr");
val two = sc.parallelize(List(
(1, Array(1, 2, 3)),
(2, Array(1,2 ,3, 4)),
(3, Array(1, 2, 3)),
(4, Array(1))
)).toDF("user", "arr");
// usage
one.join(two, one("user") === two("user"))
.select (one("user"), intersect(one("arr"), two("arr")).as("intersect"))
.where(col("intersect") > 2).show
// version from comment
one.join(two)
.select (one("user"), two("user"), intersect(one("arr"), two("arr")).as("intersect")).
where('intersect > 2).show
One possible solution is to find interesting pairs and augment these with arrays. First let's import some functions:
import org.apache.spark.sql.functions.explode
and rename columns:
val aDF_ = aDF.toDF("a_user_id", "a_user_purchases")
val bDF_ = bDF.toDF("b_user_id", "b_user_purchases")
Pairs matching the predicate can be identified as:
val filtered = aDF_.withColumn("purchase", explode($"a_user_purchases"))
.join(bDF_.withColumn("purchase", explode($"b_user_purchases")), Seq("purchase"))
.groupBy("a_user_id", "b_user_id")
.count()
.where($"count" > 2)
Finally filtered data can joined with the input datasets to obtain full result:
filtered.join(aDF_, Seq("a_user_id")).join(bDF_, Seq("b_user_id")).drop("count")
In Spark 2.4 or later you can also use built-in functions:
import org.apache.spark.sql.functions.{size, array_intersect}
aDF_
.crossJoin(bDF_)
.where(size(
array_intersect($"a_user_purchases", $"b_user_purchases"
)) > 2)
although this might be still slower than a more targeted hash join.
The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.