spark create map of list columns

user_16405245 0
Viewed 496

I have a master table as

+-----------+----------+-------------+
|  Age      | Gender   | date        |
+-----------+----------+-------------+
|  [1,2]    |   M      |  2021-01-01 | 
|  [11,13]  |   M      |  2021-01-10 | 
|  [4,5]    |   M      |  2021-01-15 |
|  [3]      |   F      |  2021-01-30 |
|  [7]      |   F      |  2021-02-20 |
|  [2]      |   F      |  2021-02-30 |
|  [6]      |   M      |  2021-02-20 |

Output needed is

  1. map of date as key and values as list of age
  2. the values in map need to extracted for limited days, I have taken 15 days for a example. It could be anything like 60, 90, 30.
+-----------+----------+---------------------------------+
|  Age      | Gender   | date                            |
+-----------+----------+---------------------------------+
|  [1,2]    |   M      |  [2021-01-01-->[1,2,11,13,4,5]] | 
|  [11,13]  |   M      |  [2021-01-10-->[1,2,11,13,4,5]] | 
|  [4,5]    |   M      |  [2021-01-15-->[1,2,11,13,4,5]] |
|  [3]      |   F      |  [2021-01-30-->[3]]             |
|  [7]      |   F      |  [2021-02-20-->[7,2]]           |
|  [2]      |   F      |  [2021-02-30-->[7.2]]           |
|  [6]      |   M      |  [2021-02-20-->[6]]             |

I have tried spark.sql("""select Age,Gender,Map(date,age) as date from master""") but not able to find out how to control out based on number of days

pyspark apache-spark-sql apache-spark
1 Answers

You can solve this by windowing the dataframe based on the date limits in your example this is 15 days, after windowing group rows which match the criteria then group based on Gender and window group to collect ages.

Working Example

  1. Imports and Dataframe setup
from datetime import datetime
import pyspark.sql.functions as F
from pyspark.sql import Window

data = [([1, 2], "M", datetime.strptime("2021-01-01", "%Y-%m-%d")),
 ([11, 13], "M", datetime.strptime("2021-01-10", "%Y-%m-%d")),
 ([4, 5], "M", datetime.strptime("2021-01-15", "%Y-%m-%d")),
 ([3], "F", datetime.strptime("2021-01-30", "%Y-%m-%d")),
 ([7], "F", datetime.strptime("2021-02-20", "%Y-%m-%d")),
 ([2], "F", datetime.strptime("2021-02-28", "%Y-%m-%d")),
 ([6], "M", datetime.strptime("2021-02-20", "%Y-%m-%d")), ]


df = spark.createDataFrame(data, ("Age", "Gender", "date"))
  1. Window defintion and acceptable condition to determine which rows can belong in the same group (acceptable_window variable in the code below captures this).
window_spec = Window.partitionBy("Gender").orderBy("date")
acceptable_window = 15
df1 = df.withColumn("date_diff", F.datediff(F.col("date"), F.lag("date").over(window_spec))).withColumn("date_group", F.when(F.col("date_diff").isNull(), 1).when(F.col("date_diff") > acceptable_window, 1).otherwise(0))

Contents of df1

+--------+------+-------------------+---------+----------+
|     Age|Gender|               date|date_diff|date_group|
+--------+------+-------------------+---------+----------+
|     [3]|     F|2021-01-30 00:00:00|     null|         1|
|     [7]|     F|2021-02-20 00:00:00|       21|         1|
|     [2]|     F|2021-02-28 00:00:00|        8|         0|
|  [1, 2]|     M|2021-01-01 00:00:00|     null|         1|
|[11, 13]|     M|2021-01-10 00:00:00|        9|         0|
|  [4, 5]|     M|2021-01-15 00:00:00|        5|         0|
|     [6]|     M|2021-02-20 00:00:00|       36|         1|
+--------+------+-------------------+---------+----------+
  1. From df1 determine which rows will fall into the same group, for this we sum date_group over our window.
df2 = df1.withColumn("date_window", F.sum(F.col("date_group")).over(window_spec))

Contents of df2

+--------+------+-------------------+---------+----------+-----------+
|     Age|Gender|               date|date_diff|date_group|date_window|
+--------+------+-------------------+---------+----------+-----------+
|     [3]|     F|2021-01-30 00:00:00|     null|         1|          1|
|     [7]|     F|2021-02-20 00:00:00|       21|         1|          2|
|     [2]|     F|2021-02-28 00:00:00|        8|         0|          2|
|  [1, 2]|     M|2021-01-01 00:00:00|     null|         1|          1|
|[11, 13]|     M|2021-01-10 00:00:00|        9|         0|          1|
|  [4, 5]|     M|2021-01-15 00:00:00|        5|         0|          1|
|     [6]|     M|2021-02-20 00:00:00|       36|         1|          2|
+--------+------+-------------------+---------+----------+-----------+
  1. Use the date_window to flatten and collect Age that fall within this window and create a mapping between date and collected ages.
df2.withColumn("date", F.create_map(F.col("date"), F.flatten(F.collect_list("Age").over(Window.partitionBy("Gender", "date_window"))))).select("Age", "Gender", "date").show(200, False)

Output

+--------+------+---------------------------------------------+
|Age     |Gender|date                                         |
+--------+------+---------------------------------------------+
|[3]     |F     |{2021-01-30 00:00:00 -> [3]}                 |
|[7]     |F     |{2021-02-20 00:00:00 -> [7, 2]}              |
|[2]     |F     |{2021-02-28 00:00:00 -> [7, 2]}              |
|[1, 2]  |M     |{2021-01-01 00:00:00 -> [1, 2, 11, 13, 4, 5]}|
|[11, 13]|M     |{2021-01-10 00:00:00 -> [1, 2, 11, 13, 4, 5]}|
|[4, 5]  |M     |{2021-01-15 00:00:00 -> [1, 2, 11, 13, 4, 5]}|
|[6]     |M     |{2021-02-20 00:00:00 -> [6]}                 |
+--------+------+---------------------------------------------+
user_79895810
Related