percentage count per group and pivot with pyspark

Question

I have dataframe with columns from and to.Both are country codes and they show starting country and destination country.

+----+---+
|from| to|
+----+---+
|  TR| tr|
|  TR| tr|
|  TR| tr|
|  TR| gr|
|  ES| tr|
|  GR| tr|
|  CZ| it|
|  LU| it|
|  AR| it|
|  DE| it|
|  IT| it|
|  IT| it|
|  US| it|
|  GR| fr|

Is there a way to get a dataframe that shows the percentage of each destination country per country of origin, with column all the destination country code?

the percentage must be out of the total destinations by the same country of origin(row).

e.g.

+----+---+----+---+----+
|from| tr|  it| fr|  gr|
+----+---+----+---+----+
|  TR|0.6|0.12|0.2|0.09|
|  IT|0.3| 0.3|0.3| 0.8|
|  US|0.1|0.34|0.3| 0.2|

Answer 1

You can pivot with count and adjust the result. First some imports:

from pyspark.sql.functions import col, lit, coalesce
from itertools import chain

Find levels:

levels = [x for x in chain(*df.select("to").distinct().collect())]

pivot:

pivoted = df.groupBy("from").pivot("to", levels).count()

compute row count expression:

row_count = sum(coalesce(col(x), lit(0)) for x in levels)

create a list of adjusted columns:

adjusted = [(col(c) / row_count).alias(c) for c in levels]

and select:

pivoted.select(col("from"), *adjusted)