[#137] Add an integration test for blocking or_groups

riley-harper · riley-harper · commit 1383518a01fd · 2024-06-06T19:17:04.000Z
diff --git a/hlink/tests/conftest.py b/hlink/tests/conftest.py
@@ -1160,6 +1160,63 @@ def blocking_explode_conf(spark, conf):
     return conf
 
 
+@pytest.fixture(scope="function")
+def blocking_or_groups_conf(spark, conf):
+    """
+    For testing the or_groups blocking functionality.
+    """
+    conf["column_mappings"] = [
+        {"column_name": "namefrst"},
+        {"column_name": "namelast"},
+        {"column_name": "birthyr"},
+        {"column_name": "sex"},
+        {"column_name": "bpl1"},
+        {"column_name": "bpl2"},
+        {"column_name": "bpl3"},
+    ]
+
+    conf["blocking"] = [
+        {
+            "column_name": "birthyr_3",
+            "dataset": "a",
+            "derived_from": "birthyr",
+            "expand_length": 3,
+            "explode": True,
+            "or_group": "birthyr",
+        },
+        {"column_name": "sex"},
+        {"column_name": "bpl1", "or_group": "bpl"},
+        {"column_name": "bpl2", "or_group": "bpl"},
+        {"column_name": "bpl3", "or_group": "bpl"},
+    ]
+    conf["comparison_features"] = [
+        {
+            "alias": "namefrst_jw",
+            "column_name": "namefrst",
+            "comparison_type": "jaro_winkler",
+        },
+        {
+            "alias": "namelast_jw",
+            "column_name": "namelast",
+            "comparison_type": "jaro_winkler",
+        },
+    ]
+    conf["comparisons"] = {
+        "comp_a": {
+            "feature_name": "namefrst_jw",
+            "threshold": 0.8,
+            "comparison_type": "threshold",
+        },
+        "comp_b": {
+            "feature_name": "namelast_jw",
+            "threshold": 0.8,
+            "comparison_type": "threshold",
+        },
+        "operator": "AND",
+    }
+    return conf
+
+
 @pytest.fixture(scope="function")
 def matching_household_conf(
     spark, conf, datasource_real_households, preprocessing, matching
diff --git a/hlink/tests/input_data/matching_or_group_test_a.csv b/hlink/tests/input_data/matching_or_group_test_a.csv
@@ -0,0 +1,59 @@
+id,namefrst,namelast,birthyr,sex,bpl1,bpl2,bpl3
+b5689d06-edd3-498e-8b5b-e04f2fa2f2a9,Catherine,Beebe,1866,2,,,
+a7118f06-949d-4d02-be0a-db33a6f8f3a8,Frances E,Bird,1870,2,,,
+85d089c0-b907-4d9c-95ab-c5fa4a3dd2bb,J S,Luff,1861,1,,,
+cddd9455-48e0-4b48-89a5-9ee315e00087,John,Smith,1884,1,,,
+8cb74256-6dfa-4d17-913a-59fa646c388a,Saml H,Russell,1833,1,,,
+1f8e1a74-d486-44ad-8d5c-51aedf86208e,Charles,Robertson,1884,1,,,
+61a1590f-1d3a-4666-8406-3d4aaf0770b4,John,Dickinson,1868,1,,,
+92277f0b-1476-41f5-9dc8-bf83672616d0,Joseph,Shissler,1874,1,,,
+322291a1-de91-439d-bba0-45fc2f47a2eb,David,Hall,1839,1,,,
+136f7105-ff59-4eac-9d95-44b002cbb448,John,Decame,1858,1,,,
+1138ab41-e234-4c72-b812-eaaf0fc5f76c,Nancy,Decame,1857,2,,,
+066ea4e1-f340-4231-b505-ec7bb9a07103,Peter N,Decame,1895,1,,,
+b7d96336-404e-490c-8c45-61f2287b52ff,Annam,Decame,1897,2,,,
+24bdff6a-5590-4494-8e8a-ac4a549c8890,Sarah,Decame,1900,2,,,
+c1fedaab-f026-4aa4-9320-e10f2432d539,James,Carney,1888,1,,,
+43a6ebe5-752b-4054-818d-6f6f75cc89e7,Alfred,Dell,1883,1,,,
+0d693015-2349-4363-9667-45036af7d0db,Chas,Syaex,1870,1,,,
+1d586e26-aac1-49df-a2ad-fe0a385a26bf,Sarah,Russell,1897,2,,,
+93b7ac89-f9db-49b2-a1f2-c189fecc14ae,Wm H,Hazard,1881,1,,,
+e51c36c9-570c-466d-aac1-bf380c9c20f1,Martha,Hazard,1880,2,,,
+9250341a-8336-494a-bc84-2b803efe64c6,Willie May,Hazard,1902,2,,,
+a70679f0-9313-4ef3-bf87-5dfe81beed5d,Samuel,Hazard,1906,2,,,
+4715bbf6-d3e2-4260-9ddd-6aece147e5c1,Samuel,Morgan,1878,1,,,
+77378570-5214-4ac5-8258-c5156e8b99b3,J Clauson,Mcfarland,1890,1,,,
+6542b541-6e10-411f-9b2a-7c0b93b0aa68,Eugene,Mcfarland,1892,1,,,
+396c4077-6a70-4a17-97fb-f8a0c06fdafe,Anna,Preston,1871,2,,,
+7e9dde5e-3fad-4b2e-b367-643c0dc8cabb,Rebecca N,Alexander,1861,2,,,
+f7d9e25f-c390-4222-ac24-4e93d72daa05,Martha,Ellis,1873,2,,,
+24b7afa1-8c49-4833-8292-c545c85d3b89,Otillia,Zeider,1876,2,,,
+4b416874-0c5c-4233-81ec-39223bc66f4f,Mary,Doyle,1846,2,,,
+a499b0dc-7ac0-4d61-b493-91a3036c712e ,ANNIE ,FAUBLE ,1884,2,1,,
+ae7261c3-7d71-4ea1-997f-5d1a68c18777 ,MARY ,REESE ,1875,2,,,
+ad6442b5-42bc-4c2e-a517-5a951d989a92 ,MARY ,REESE ,1899,2,1,2,3
+b0b6695f-dfa5-4e4d-bc75-798c27195fff ,SALLY ,REESE ,1901,2,,,
+9e807937-de09-414c-bfb2-ac821e112929 ,JOHN ,SHIELDS ,1880,1,1,,
+426f2cbe-32e1-45eb-9f86-89a2b9116b7e ,ANNE ,FAUBLE ,1884,2,,,
+a76697d9-b0c8-4774-bc3e-12a7e403c7e6 ,JOHN ,COLLINS ,1893,1,,,
+3575c9ba-1527-4ca2-aff0-d7c2d1efb421 ,MAGGIE ,COLLINS ,1894,2,,,
+49e53dbc-fe8e-4e55-8cb9-a1d93c284d98 ,MARY ,COLLINS ,1898,2,,,
+50b33ef6-259d-43af-8cdc-56a61f881169 ,WILLIAM H. ,SEWARD ,1856,1,,4,
+952754a5-48b4-462a-ac57-e4a059a9ef98 ,ESTHER ,BIERHAHN ,1870,2,,,
+ea6d77b3-2e2d-4c59-a0ac-6b297e8898e3 ,CHARLES ,CLEVELAND ,1865,1,,,
+60a5052e-6d67-455a-a3aa-bb79560c7d8d ,SUSAN ,WILSON ,1850,2,,,
+0d4472ec-6378-4aeb-b6c7-17e1c388bb94 ,ARCHER ,HARVEY ,1890,1,,,
+65ccbeb7-2c79-4fb0-b354-c67f150ad80c ,ELIZABETH ,MC LEAN ,1868,2,,,
+72cbe5fa-f558-4393-8423-1842fadf7f11 ,MARY A. ,FLEMMING ,1837,2,,,
+44693008-fd6f-48fe-9c52-e6c07baff361 ,BESSIE ,CHAMBERS ,1908,2,,,
+bcc0988e-2397-4f1b-8e76-4bfe1b05dbc6 ,THOMAS ,GRAHAM ,1846,1,,,
+a7b10530-b7c9-44d5-9125-c603f392d6d3 ,EDWARD ,DEKAY ,1875,1,,,
+1e635c1c-7faa-4270-acf3-a22635884b90 ,NATHEN ,THORPE ,1836,1,,,
+d3217545-3453-4d96-86c0-d6a3e60fb2f8 ,JOB ,FOSTER ,1884,1,,,
+2a35bae5-3120-4e2c-87da-694d4419c9ce ,JEZEBEL ,FOSTER ,1888,2,,,
+94460fc2-954b-469d-9726-f7126c30e5e2 ,ELIZA ,GOODWIN ,1871,2,,,
+620b6ebb-82e6-42db-8aae-300ca2be0c00 ,MARY ,GOODWIN ,1893,2,,,
+bfe1080e-2e67-4a8c-a6e1-ed94ea103712 ,JO ,GOODWIN ,1895,1,,6,7
+7fb55d25-2a7d-486d-9efa-27b9d7e60c24 ,PHINEAS ,TAYLOR ,1871,1,,5,
+a0f33b36-cef7-4949-a031-22b90f1055d4 ,MARY A. ,LORD ,1856,2,,,1
+1a76745c-acf8-48a0-9992-7fb10c11710b ,E.B. ,ALLEN ,1889,1,,,
diff --git a/hlink/tests/input_data/matching_or_group_test_b.csv b/hlink/tests/input_data/matching_or_group_test_b.csv
@@ -0,0 +1,27 @@
+id,namefrst,namelast,birthyr,sex,bpl1,bpl2,bpl3
+a499b0dc-7ac0-4d61-b493-91a3036c712e ,ANNIE ,FAUBLE ,1884,2,1,,
+ae7261c3-7d71-4ea1-997f-5d1a68c18777 ,MARY ,REESE ,1875,2,,,
+ad6442b5-42bc-4c2e-a517-5a951d989a92 ,MARY ,REESE ,1902,2,1,2,3
+9e807937-de09-414c-bfb2-ac821e112929 ,JOHN ,SHIELDS ,1889,1,1,,
+426f2cbe-32e1-45eb-9f86-89a2b9116b7e ,ANNE ,FAUBLE ,1884,2,,,
+a76697d9-b0c8-4774-bc3e-12a7e403c7e6 ,JOHN ,COLLINS ,1893,1,,,
+3575c9ba-1527-4ca2-aff0-d7c2d1efb421 ,MAGGIE ,COLLINS ,1894,2,,,
+49e53dbc-fe8e-4e55-8cb9-a1d93c284d98 ,MARY ,COLLINS ,1898,2,,,
+50b33ef6-259d-43af-8cdc-56a61f881169 ,WILLIAM H. ,SEWARD ,1866,1,,4,
+952754a5-48b4-462a-ac57-e4a059a9ef98 ,ESTHER ,BIERHAHN ,1870,2,,,
+ea6d77b3-2e2d-4c59-a0ac-6b297e8898e3 ,CHARLES ,CLEVELAND ,1865,1,,,
+60a5052e-6d67-455a-a3aa-bb79560c7d8d ,SUSAN ,WILSON ,1850,2,,,
+0d4472ec-6378-4aeb-b6c7-17e1c388bb94 ,ARCHER ,HARVEY ,1893,1,,,
+65ccbeb7-2c79-4fb0-b354-c67f150ad80c ,ELIZABETH ,MC LEAN ,1868,2,,,
+72cbe5fa-f558-4393-8423-1842fadf7f11 ,MARY A. ,FLEMMING ,1842,2,,,
+bcc0988e-2397-4f1b-8e76-4bfe1b05dbc6 ,THOMAS ,GRAHAM ,1846,1,,,
+a7b10530-b7c9-44d5-9125-c603f392d6d3 ,EDWARD ,DEKAY ,1875,1,,,
+1e635c1c-7faa-4270-acf3-a22635884b90 ,NATHEN ,THORPE ,1836,1,,,
+d3217545-3453-4d96-86c0-d6a3e60fb2f8 ,JOB ,FOSTER ,1884,1,,,
+2a35bae5-3120-4e2c-87da-694d4419c9ce ,JEZEBEL ,FOSTER ,1888,2,,,
+94460fc2-954b-469d-9726-f7126c30e5e2 ,ELIZA ,GOODWIN ,1871,2,,,
+620b6ebb-82e6-42db-8aae-300ca2be0c00 ,MARY ,GOODWIN ,1893,2,,,
+bfe1080e-2e67-4a8c-a6e1-ed94ea103712 ,JO ,GOODWIN ,1890,1,,6,7
+7fb55d25-2a7d-486d-9efa-27b9d7e60c24 ,PHINEAS ,TAYLOR ,1871,1,,5,
+a0f33b36-cef7-4949-a031-22b90f1055d4 ,MARY A. ,LORD ,1856,2,,,1
+1a76745c-acf8-48a0-9992-7fb10c11710b ,E.B. ,ALLEN ,1889,1,,,
diff --git a/hlink/tests/matching_blocking_explode_test.py b/hlink/tests/matching_blocking_explode_test.py
@@ -3,6 +3,7 @@
 # in this project's top-level directory, and also on-line at:
 #   https://github.com/ipums/hlink
 
+from pyspark.sql import Row
 import pytest
 import pandas as pd
 from hlink.linking.matching.link_step_match import extract_or_groups_from_blocking
@@ -123,6 +124,42 @@ def test_blocking_multi_layer_comparison(
         ) or (row["namelast_jw_x"] < 0.7)
 
 
+def test_blocking_or_groups(
+    spark, blocking_or_groups_conf, matching_or_groups_test_input, matching
+):
+    """Test the blocking or_group functionality. This feature supports
+    combining some or all blocking conditions with OR instead of AND."""
+    table_a, table_b = matching_or_groups_test_input
+    table_a.createOrReplaceTempView("prepped_df_a")
+    table_b.createOrReplaceTempView("prepped_df_b")
+
+    matching.run_step(0)
+    matching.run_step(1)
+
+    potential_matches = matching.spark.table("potential_matches")
+
+    results = potential_matches.select("id_a", "id_b").collect()
+
+    assert set(results) == {
+        Row(
+            id_a="ad6442b5-42bc-4c2e-a517-5a951d989a92 ",
+            id_b="ad6442b5-42bc-4c2e-a517-5a951d989a92 ",
+        ),
+        Row(
+            id_a="a499b0dc-7ac0-4d61-b493-91a3036c712e ",
+            id_b="a499b0dc-7ac0-4d61-b493-91a3036c712e ",
+        ),
+        Row(
+            id_a="7fb55d25-2a7d-486d-9efa-27b9d7e60c24 ",
+            id_b="7fb55d25-2a7d-486d-9efa-27b9d7e60c24 ",
+        ),
+        Row(
+            id_a="a0f33b36-cef7-4949-a031-22b90f1055d4 ",
+            id_b="a0f33b36-cef7-4949-a031-22b90f1055d4 ",
+        ),
+    }
+
+
 # TODO: test_step_2_length_b
 
 # TODO: test_step_2_has_matching_element
diff --git a/hlink/tests/plugins/datasources.py b/hlink/tests/plugins/datasources.py
@@ -879,6 +879,34 @@ def matching_test_input(spark, conf, tmpdir_factory):
     return pdfa, pdfb
 
 
+@pytest.fixture(scope="function")
+def matching_or_groups_test_input(spark):
+    prepped_a_data = "input_data/matching_or_group_test_a.csv"
+    prepped_b_data = "input_data/matching_or_group_test_b.csv"
+
+    package_path = os.path.dirname(hlink.tests.__file__)
+    pa_path = os.path.join(package_path, prepped_a_data)
+    pb_path = os.path.join(package_path, prepped_b_data)
+
+    schema = StructType(
+        [
+            StructField("id", StringType(), True),
+            StructField("namefrst", StringType(), True),
+            StructField("namelast", StringType(), True),
+            StructField("birthyr", LongType(), True),
+            StructField("sex", LongType(), True),
+            StructField("bpl1", LongType(), True),
+            StructField("bpl2", LongType(), True),
+            StructField("bpl3", LongType(), True),
+        ]
+    )
+
+    pdfa = spark.read.csv(pa_path, schema)
+    pdfb = spark.read.csv(pb_path, schema)
+
+    return pdfa, pdfb
+
+
 @pytest.fixture(scope="function")
 def datasource_mi_comparison(spark, conf):
     """Create the prepped_df_(a/b) dataframes and populate basic config values"""
