Merge pull request #143 from ipums/multiple_exploded_columns

Support blocking sections with multiple exploded columns

Author: riley-harper

hlink/linking/matching/link_step_explode.py

@@ -124,49 +124,37 @@ def _explode(
             if exploding_column.get("expand_length", False):
                 expand_length = exploding_column["expand_length"]
                 derived_from_column = exploding_column["derived_from"]
-                explode_selects = [
-                    (
-                        explode(self._expand(derived_from_column, expand_length)).alias(
-                            exploding_column_name
-                        )
-                        if exploding_column_name == column
-                        else column
-                    )
-                    for column in all_column_names
-                ]
+
+                explode_col_expr = explode(
+                    self._expand(derived_from_column, expand_length)
+                )
             else:
-                explode_selects = [
-                    (
-                        explode(col(exploding_column_name)).alias(exploding_column_name)
-                        if exploding_column_name == c
-                        else c
-                    )
-                    for c in all_column_names
-                ]
+                explode_col_expr = explode(col(exploding_column_name))
+
             if "dataset" in exploding_column:
                 derived_from_column = exploding_column["derived_from"]
-                explode_selects_with_derived_column = [
-                    (
-                        col(derived_from_column).alias(exploding_column_name)
-                        if exploding_column_name == column
-                        else column
-                    )
-                    for column in all_column_names
-                ]
+                no_explode_col_expr = col(derived_from_column)
+
                 if exploding_column["dataset"] == "a":
-                    exploded_df = (
-                        exploded_df.select(explode_selects)
-                        if is_a
-                        else exploded_df.select(explode_selects_with_derived_column)
-                    )
+                    expr = explode_col_expr if is_a else no_explode_col_expr
+                    exploded_df = exploded_df.withColumn(exploding_column_name, expr)
                 elif exploding_column["dataset"] == "b":
-                    exploded_df = (
-                        exploded_df.select(explode_selects)
-                        if not (is_a)
-                        else exploded_df.select(explode_selects_with_derived_column)
-                    )
+                    expr = explode_col_expr if not is_a else no_explode_col_expr
+                    exploded_df = exploded_df.withColumn(exploding_column_name, expr)
             else:
-                exploded_df = exploded_df.select(explode_selects)
+                exploded_df = exploded_df.withColumn(
+                    exploding_column_name, explode_col_expr
+                )
+
+        # If there are exploding columns, then select out "all_column_names".
+        # Otherwise, just let all of the columns through without selecting
+        # specific ones. I believe this is an artifact of a previous
+        # implementation, but the tests currently enforce it. It may or may not
+        # be a breaking change to remove this. We'd have to look into the
+        # ramifications.
+        if len(all_exploding_columns) > 0:
+            exploded_df = exploded_df.select(sorted(all_column_names))
+
         return exploded_df
 
     def _expand(self, column_name: str, expand_length: int) -> Column:

hlink/tests/matching_blocking_explode_test.py

@@ -124,6 +124,59 @@ def test_blocking_multi_layer_comparison(
         ) or (row["namelast_jw_x"] < 0.7)
 
 
+def test_blocking_multiple_exploded_columns(
+    spark, blocking_explode_conf, matching_test_input, matching
+):
+    """
+    Matching supports multiple exploded blocking columns. Each column is
+    exploded independently. See GitHub issue #142.
+    """
+    table_a, table_b = matching_test_input
+    table_a.createOrReplaceTempView("prepped_df_a")
+    table_b.createOrReplaceTempView("prepped_df_b")
+
+    blocking_explode_conf["blocking"] = [
+        {
+            "column_name": "birthyr_3",
+            "dataset": "a",
+            "derived_from": "birthyr",
+            "expand_length": 3,
+            "explode": True,
+        },
+        {
+            "column_name": "birthyr_4",
+            "dataset": "a",
+            "derived_from": "birthyr",
+            "expand_length": 4,
+            "explode": True,
+        },
+        {"column_name": "sex"},
+    ]
+
+    matching.run_step(0)
+
+    exploded_a = spark.table("exploded_df_a").toPandas()
+    exploded_b = spark.table("exploded_df_b").toPandas()
+
+    input_size_a = spark.table("prepped_df_a").count()
+    input_size_b = spark.table("prepped_df_b").count()
+    output_size_a = len(exploded_a)
+    output_size_b = len(exploded_b)
+
+    assert "sex" in exploded_a.columns
+    assert "birthyr_3" in exploded_a.columns
+    assert "birthyr_4" in exploded_a.columns
+    assert "sex" in exploded_b.columns
+    assert "birthyr_3" in exploded_b.columns
+    assert "birthyr_4" in exploded_b.columns
+
+    # birthyr_3 multiplies the number of columns by 2 * 3 + 1 = 7
+    # birthyr_4 multiplies the number of columns by 2 * 4 + 1 = 9
+    assert input_size_a * 63 == output_size_a
+    # Both columns are only exploded in dataset A
+    assert input_size_b == output_size_b
+
+
 def test_blocking_or_groups(
     spark, blocking_or_groups_conf, matching_or_groups_test_input, matching
 ):