Data cleaning region tags (GoogleCloudPlatform#4785)

* add region tags

* remove whitespace

* add mising region tag, fix whitespcae

* remove gender

* remove collect

* update gender column removal, switch args position

* fix lint

* add region tag, address brad comments, fix test

* remove timeout

* address brad comment

Author: leahecole

data-science-onramp/data-cleaning/clean.py

@@ -10,6 +10,7 @@
 from pyspark.sql.types import FloatType, IntegerType, StringType
 
 
+# [START datascienceonramp_tripdurationudf]
 def trip_duration_udf(duration):
     """Convert trip duration to seconds. Return None if negative."""
     if not duration:
@@ -33,11 +34,17 @@ def trip_duration_udf(duration):
     return int(time)
 
 
+# [END datascienceonramp_tripdurationudf]
+
+# [START datascienceonramp_stationnameudf]
 def station_name_udf(name):
     """Replaces '/' with '&'."""
     return name.replace("/", "&") if name else None
 
 
+# [END datascienceonramp_stationnameudf]
+
+# [START datascienceonramp_usertypeudf]
 def user_type_udf(user):
     """Converts user type to 'Subscriber' or 'Customer'."""
     if not user:
@@ -49,17 +56,10 @@ def user_type_udf(user):
         return "Customer"
 
 
-def gender_udf(gender):
-    """Converts gender to 'Male' or 'Female'."""
-    if not gender:
-        return None
-
-    if gender.lower().startswith("m"):
-        return "Male"
-    elif gender.lower().startswith("f"):
-        return "Female"
+# [END datascienceonramp_usertypeudf]
 
 
+# [START datascienceonramp_stationlocationudf]
 def angle_udf(angle):
     """Converts DMS notation to degrees. Return None if not in DMS or degrees notation."""
     if not angle:
@@ -74,6 +74,9 @@ def angle_udf(angle):
         return float(degrees[0])
 
 
+# [END datascienceonramp_stationlocationudf]
+
+# [START datascienceonramp_timeconvertudf]
 def compute_time(duration, start, end):
     """Calculates duration, start time, and end time from each other if one value is null."""
     time_format = "%Y-%m-%dT%H:%M:%S"
@@ -94,15 +97,17 @@ def compute_time(duration, start, end):
     if duration:
         # Convert to timedelta
         duration = datetime.timedelta(seconds=duration)
-
+    # [END datascienceonramp_timeconvertudf]
+    # [START datascienceonramp_timemissingvalueudf]
     # Calculate missing value
     if start and end and not duration:
         duration = end - start
     elif duration and end and not start:
         start = end - duration
     elif duration and start and not end:
         end = start + duration
-
+    # [END datascienceonramp_timemissingvalueudf]
+    # [START datascienceonramp_timereturnudf]
     # Transform to primitive types
     if duration:
         duration = int(duration.total_seconds())
@@ -114,6 +119,9 @@ def compute_time(duration, start, end):
     return (duration, start, end)
 
 
+# [END datascienceonramp_timereturnudf]
+
+# [START datascienceonramp_timehelperudf]
 def compute_duration_udf(duration, start, end):
     """Calculates duration from start and end time if null."""
     return compute_time(duration, start, end)[0]
@@ -129,9 +137,12 @@ def compute_end_udf(duration, start, end):
     return compute_time(duration, start, end)[2]
 
 
+# [END datascienceonramp_timehelperudf]
+
+# [START datascienceonramp_sparksession]
 if __name__ == "__main__":
-    TABLE = sys.argv[1]
-    BUCKET_NAME = sys.argv[2]
+    BUCKET_NAME = sys.argv[1]
+    TABLE = sys.argv[2]
 
     # Create a SparkSession, viewable via the Spark UI
     spark = SparkSession.builder.appName("data_cleaning").getOrCreate()
@@ -142,13 +153,20 @@ def compute_end_udf(duration, start, end):
     except Py4JJavaError as e:
         raise Exception(f"Error reading {TABLE}") from e
 
+# [END datascienceonramp_sparksession]
+
+# [START datascienceonramp_removecolumn]
+    # remove unused column
+    df = df.drop("gender")
+# [END datascienceonramp_removecolumn]
+
+# [START datascienceonramp_sparksingleudfs]
     # Single-parameter udfs
     udfs = {
         "start_station_name": UserDefinedFunction(station_name_udf, StringType()),
         "end_station_name": UserDefinedFunction(station_name_udf, StringType()),
         "tripduration": UserDefinedFunction(trip_duration_udf, IntegerType()),
         "usertype": UserDefinedFunction(user_type_udf, StringType()),
-        "gender": UserDefinedFunction(gender_udf, StringType()),
         "start_station_latitude": UserDefinedFunction(angle_udf, FloatType()),
         "start_station_longitude": UserDefinedFunction(angle_udf, FloatType()),
         "end_station_latitude": UserDefinedFunction(angle_udf, FloatType()),
@@ -157,7 +175,8 @@ def compute_end_udf(duration, start, end):
 
     for name, udf in udfs.items():
         df = df.withColumn(name, udf(name))
-
+    # [END datascienceonramp_sparksingleudfs]
+    # [START datascienceonramp_sparkmultiudfs]
     # Multi-parameter udfs
     multi_udfs = {
         "tripduration": {
@@ -176,10 +195,12 @@ def compute_end_udf(duration, start, end):
 
     for name, obj in multi_udfs.items():
         df = df.withColumn(name, obj["udf"](*obj["params"]))
-
+    # [END datascienceonramp_sparkmultiudfs]
+    # [START datascienceonramp_displaysamplerows]
     # Display sample of rows
     df.show(n=20)
-
+    # [END datascienceonramp_displaysamplerows]
+    # [START datascienceonramp_writetogcs]
     # Write results to GCS
     if "--dry-run" in sys.argv:
         print("Data will not be uploaded to GCS")
@@ -222,3 +243,4 @@ def compute_end_udf(duration, start, end):
         print(
             "Data successfully uploaded to " + "gs://" + BUCKET_NAME + "/" + final_path
         )
+# [END datascienceonramp_writetogcs]

data-science-onramp/data-cleaning/clean_test.py

@@ -49,11 +49,11 @@
         },
     },
 }
-DATAPROC_JOB = {    # Dataproc job configuration
+DATAPROC_JOB = {  # Dataproc job configuration
     "placement": {"cluster_name": DATAPROC_CLUSTER},
     "pyspark_job": {
         "main_python_file_uri": f"gs://{BUCKET_NAME}/{BUCKET_BLOB}",
-        "args": [BQ_TABLE, BUCKET_NAME, "--dry-run"],
+        "args": [BUCKET_NAME, BQ_TABLE, "--dry-run"],
         "jar_file_uris": ["gs://spark-lib/bigquery/spark-bigquery-latest_2.12.jar"],
     },
 }
@@ -70,12 +70,9 @@ def setup_and_teardown_table():
     # Load table from dataframe
     df = pd.read_csv(CSV_FILE)
     job_config = bigquery.LoadJobConfig(
-        autodetect=True,
-        write_disposition="WRITE_TRUNCATE"
-    )
-    operation = bq_client.load_table_from_dataframe(
-        df, BQ_TABLE, job_config=job_config
+        autodetect=True, write_disposition="WRITE_TRUNCATE"
     )
+    operation = bq_client.load_table_from_dataframe(df, BQ_TABLE, job_config=job_config)
 
     # Wait for job to complete
     operation.result()
@@ -108,8 +105,7 @@ def setup_and_teardown_cluster():
     operation = cluster_client.delete_cluster(
         project_id=PROJECT_ID,
         region=CLUSTER_REGION,
-        cluster_name=DATAPROC_CLUSTER,
-        timeout=300
+        cluster_name=DATAPROC_CLUSTER
     )
     operation.result()
 
@@ -172,20 +168,7 @@ def test_clean():
 
     # gender
     assert not is_in_table("M", out)
-    assert not is_in_table("m", out)
-    assert not is_in_table("male", out)
-    assert not is_in_table("MALE", out)
     assert not is_in_table("F", out)
-    assert not is_in_table("f", out)
-    assert not is_in_table("female", out)
-    assert not is_in_table("FEMALE", out)
-    assert not is_in_table("U", out)
-    assert not is_in_table("u", out)
-    assert not is_in_table("unknown", out)
-    assert not is_in_table("UNKNOWN", out)
-
-    assert is_in_table("Male", out)
-    assert is_in_table("Female", out)
 
     # customer plan
     assert not is_in_table("subscriber", out)