Use date field in traffic aggregation

thisisartium · Sep 20, 2022 · 8f8fb12 · 8f8fb12
1 parent a0f8ce6
commit 8f8fb12
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/sample_pipelines/traffic/analytics.py b/sample_pipelines/traffic/analytics.py
@@ -16,7 +16,7 @@ def aggregate_incidents(paths, secrets, spark, glueContext):
             "geojson",
             h3_pyspark.h3_to_geo_boundary(F.col("h3_index"), F.lit(True)),
         )
-        .withColumn("month", F.date_format("published_date", "yyyy-MM"))
+        .withColumn("month", F.trunc("published_date", "MM"))
         .groupBy("month", "h3_index", "issue_reported")
         .agg(F.count("*").alias("count"), F.first("geojson").alias("geojson"))
     )

diff --git a/tests/sample_pipelines/traffic/analytics_test.py b/tests/sample_pipelines/traffic/analytics_test.py
@@ -1,5 +1,6 @@
 from chispa.dataframe_comparer import assert_df_equality
 from dateutil.parser import isoparse
+from datetime import date
 from pytest import fixture
 
 from sample_pipelines.traffic.analytics import aggregate_incidents
@@ -63,9 +64,9 @@ def test_aggregation(paths, secrets, spark, glueContext, input_df):
 
     expected = spark.createDataFrame(
         [
-            ("2018-06", "89489e3464fffff", "Incident 1", 2),
-            ("2018-06", "89489e3464fffff", "Incident 2", 1),
-            ("2018-07", "89489e3464fffff", "Incident 1", 1),
+            (date.fromisoformat("2018-06-01"), "89489e3464fffff", "Incident 1", 2),
+            (date.fromisoformat("2018-06-01"), "89489e3464fffff", "Incident 2", 1),
+            (date.fromisoformat("2018-07-01"), "89489e3464fffff", "Incident 1", 1),
         ],
         ["month", "h3_index", "issue_reported", "count"],
     )