spark find which persons often go to the same counties - apache-spark

Question

The data source is:
val spark = SparkSession.builder().master("local[1,1]").config("spark.sql.shuffle.partitions", "1").config("spark.sql.crossJoin.enabled","true").getOrCreate()
spark.sparkContext.setLogLevel("error")
import spark.implicits._
val df=Seq(
("tom","America","2019"),
("jim","America","2019"),
("jack","America","2019"),
("tom","Russia","2019"),
("jim","Russia","2019"),
("jack","Russia","2019"),
("alex","Russia","2019"),
("tom","America","2018"),
("jim","America","2018"),
("tom","Germany","2018"),
("jim","England","2018")
).toDF("person","country","year")
I want to find which persons often go to the same countries for each year,and where they gone together,so what I expect is a json like this:
[{
"year": "2019",
"items": [{
"persons": ["tom", "jim", "jack"],
"common": ["America", "Russia"],
"times": 2
}, {
"persons": ["tom", "jack"],
"common": ["America", "Russia"],
"times": 2
}, {
"persons": ["tom", "jim"],
"common": ["America", "Russia"],
"times": 2
}, {
"persons": ["jack", "jim"],
"common": ["America", "Russia"],
"times": 2
}]
},
{
"year": "2018",
"items": [{
"persons": ["tom", "jim"],
"common": ["America"],
"times": 1
}]
}
]
Then I am not sure what model shall I use?
I tried Frequent Items Pattern:
val df1=df.where('year===2019)
val rdd1= df1.groupBy("country").agg(collect_set('person)).drop("country","year")
.as[Array[String]].rdd
val fpg = new FPGrowth()
.setMinSupport(0.3)
.setNumPartitions(10)
val schema = new StructType().add(new StructField("items", ArrayType(StringType))).add(new StructField("freq", LongType))
val model = fpg.run(rdd1);
val rdd2 = model.freqItemsets.map(itemset => Row(itemset.items, itemset.freq))
val df1 = spark.createDataFrame(rdd2, schema).where(size('items)>1)
.show()
loop for every year
val df2=df.where('year===2018)
val rdd2= df1.groupBy("country").agg(collect_set('person)).drop("country","year")
.as[Array[String]].rdd
....
val model = fpg.run(rdd12);
....
The result is :
for 2019
+----------------+----+
| items|freq|
+----------------+----+
| [jack, tom]| 2|
|[jack, tom, jim]| 2|
| [jack, jim]| 2|
| [tom, jim]| 2|
+----------------+----+
for 2018:
+----------+----+
| items|freq|
+----------+----+
|[tom, jim]| 1|
+----------+----+
But I can not get when and where they gone together,becuase the rdd I give to FPGRowth must be a RDD[Array[String]],no more columns allowed.
Is there any other better model?How can I achieve it?
I also want to know how many times each person group go together
Maybe what I should use collaborative filtering

Answer 1

Just self-join and aggregate
import org.apache.spark.sql.functions._
df.alias("left")
.join(df.alias("right"), Seq("country", "year"))
.where($"left.person" < $"right.person")
.groupBy(array($"left.person", $"right.person").alias("persons"))
.agg(collect_set(struct($"country", $"year")).alias("common"))

Answer 2

Try this:
val window = Window.partitionBy("country", "year")
df
.withColumn("persons", collect_set('person) over window)
.drop('person)
.distinct()
.groupBy('persons)
.agg(collect_set(struct('country, 'year)).alias("common"))
Output (tested):
+----------+----------------------------------+
|persons |common |
+----------+----------------------------------+
|[jim, tom]|[[America, 2019], [Russia, 2019]] |
|[tom] |[[Germany, 2018], [America, 2018]]|
|[jim] |[[Russia, 2018], [England, 2018]] |
+----------+----------------------------------+