Merge pull request #213 from ipums/custom-col-mappings

Support custom column mapping transforms

Author: riley-harper

hlink/linking/core/column_mapping.py

@@ -2,7 +2,66 @@
 # For copyright and licensing information, see the NOTICE and LICENSE files
 # in this project's top-level directory, and also on-line at:
 #   https://github.com/ipums/hlink
-from typing import Any
+"""
+Column mappings for cleaning and preprocessing input data.
+
+This module provides functions for cleaning and preprocessing columns of Spark
+data frames. It depends on the idea of a "column mapping", which is a
+dictionary which specifies an input column, an optional output column alias,
+and a list of zero or more transforms to apply to the input column.
+
+```python
+# An example column mapping. The "column_name" attribute gives the name of the
+# input column, and "alias" gives the name of the output column. The alias is
+# optional and defaults to the input column name.
+{
+    "column_name": "namefrst",
+    "alias": "namefrst_std",
+    "transforms": [
+        {"type": "lowercase_strip"},
+        {"type": "rationalize_name_words"},
+        {"type": "remove_qmark_hyphen"},
+        {"type": "condense_strip_whitespace"},
+        {"type": "split"},
+        {"type": "array_index", "value": 0},
+    ]
+}
+```
+
+Hlink has many built-in column mapping transforms, computed by the
+`transform_*` functions in this module. Hlink also has support for custom
+column mapping transforms via the `custom_transforms` argument to
+`select_column_mapping`. This argument must be a mapping from strings to
+functions which compute the column mapping transforms. For example, say that
+you wanted to implement a custom column mapping transform named "reverse" which
+reverses a string. The first thing to do is to write a function which computes
+the transform and satisfies the column mapping transform interface (see the
+ColumnMappingTransform type alias below).
+
+```python
+from pyspark.sql import Column
+from pyspark.sql.functions import reverse
+
+# input_col is the input Column expression.
+# transform is the column mapping transform dictionary, like
+#   {"type": "reverse"}. This lets the transform accept arbitrary arguments from
+#   the configuration as needed.
+# context is a dictionary with additional context which may be helpful for some
+#   transforms. In particular, it always contains at least the key "dataset",
+#   which indicates whether the current dataset is dataset "a" or "b".
+def transform_reverse(input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]) -> Column:
+    return reverse(input_col)
+```
+
+Then, when you call `select_column_mapping`, you can pass
+`custom_transforms={"reverse": transform_reverse}`, and hlink will
+automatically use your custom transform when appropriate. Note that custom
+transforms which have the same name as a built-in transform override the
+built-in transform.
+"""
+
+from collections.abc import Mapping
+from typing import Any, Callable, TypeAlias
 
 from pyspark.sql import Column, DataFrame
 from pyspark.sql.functions import (
@@ -21,25 +80,40 @@
 from pyspark.sql.types import LongType
 
 
+ColumnMappingTransform: TypeAlias = Callable[
+    [Column, Mapping[str, Any], Mapping[str, Any]], Column
+]
+"""
+The form of column mapping transform functions. These take an input Column,
+the transform mapping from the configuration, and a mapping providing some
+additional context. They return a new output Column.
+"""
+
+
 def select_column_mapping(
-    column_mapping: dict[str, Any],
+    column_mapping: Mapping[str, Any],
     df_selected: DataFrame,
     is_a: bool,
     column_selects: list[str],
+    custom_transforms: Mapping[str, ColumnMappingTransform] | None = None,
 ) -> tuple[DataFrame, list[str]]:
     name = column_mapping["column_name"]
     if "override_column_a" in column_mapping and is_a:
         override_name = column_mapping["override_column_a"]
         column_select = col(override_name)
         if "override_transforms" in column_mapping:
             for transform in column_mapping["override_transforms"]:
-                column_select = apply_transform(column_select, transform, is_a)
+                column_select = apply_transform(
+                    column_select, transform, is_a, custom_transforms
+                )
     elif "override_column_b" in column_mapping and not is_a:
         override_name = column_mapping["override_column_b"]
         column_select = col(override_name)
         if "override_transforms" in column_mapping:
             for transform in column_mapping["override_transforms"]:
-                column_select = apply_transform(column_select, transform, is_a)
+                column_select = apply_transform(
+                    column_select, transform, is_a, custom_transforms
+                )
     elif "set_value_column_a" in column_mapping and is_a:
         value_to_set = column_mapping["set_value_column_a"]
         column_select = lit(value_to_set)
@@ -49,7 +123,9 @@ def select_column_mapping(
     elif "transforms" in column_mapping:
         column_select = col(name)
         for transform in column_mapping["transforms"]:
-            column_select = apply_transform(column_select, transform, is_a)
+            column_select = apply_transform(
+                column_select, transform, is_a, custom_transforms
+            )
     else:
         column_select = col(name)
 
@@ -59,7 +135,7 @@ def select_column_mapping(
     return df_selected.withColumn(alias, column_select), column_selects
 
 
-def _require_key(transform: dict[str, Any], key: str) -> Any:
+def _require_key(transform: Mapping[str, Any], key: str) -> Any:
     """
     Extract a key from a transform, or raise a helpful context-aware error if the
     key is not present.
@@ -78,7 +154,10 @@ def _require_key(transform: dict[str, Any], key: str) -> Any:
 
 #  These apply to the column mappings in the current config
 def apply_transform(
-    column_select: Column, transform: dict[str, Any], is_a: bool
+    column_select: Column,
+    transform: Mapping[str, Any],
+    is_a: bool,
+    custom_transforms: Mapping[str, ColumnMappingTransform] | None = None,
 ) -> Column:
     """Return a new column that is the result of applying the given transform
     to the given input column (column_select). The is_a parameter controls the
@@ -93,7 +172,7 @@ def apply_transform(
     dataset = "a" if is_a else "b"
     context = {"dataset": dataset}
     transform_type = transform["type"]
-    transforms = {
+    builtin_transforms = {
         "add_to_a": transform_add_to_a,
         "concat_to_a": transform_concat_to_a,
         "concat_to_b": transform_concat_to_b,
@@ -123,7 +202,9 @@ def apply_transform(
         "get_floor": transform_get_floor,
     }
 
-    transform_func = transforms.get(transform_type)
+    custom_transforms = custom_transforms or {}
+    builtin_func = builtin_transforms.get(transform_type)
+    transform_func = custom_transforms.get(transform_type, builtin_func)
 
     if transform_func is None:
         raise ValueError(f"Invalid transform type for {transform}")
@@ -132,7 +213,7 @@ def apply_transform(
 
 
 def transform_add_to_a(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     is_a = context["dataset"] == "a"
     if is_a:
@@ -142,7 +223,7 @@ def transform_add_to_a(
 
 
 def transform_concat_to_a(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     is_a = context["dataset"] == "a"
     if is_a:
@@ -153,7 +234,7 @@ def transform_concat_to_a(
 
 
 def transform_concat_to_b(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     is_a = context["dataset"] == "a"
     if is_a:
@@ -164,50 +245,50 @@ def transform_concat_to_b(
 
 
 def transform_concat_two_cols(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     column_to_append = _require_key(transform, "column_to_append")
     return concat(input_col, column_to_append)
 
 
 def transform_lowercase_strip(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return lower(trim(input_col))
 
 
 def transform_rationalize_name_words(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return regexp_replace(input_col, r"[^a-z?'\*\-]+", " ")
 
 
 def transform_remove_qmark_hyphen(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return regexp_replace(input_col, r"[?\*\-]+", "")
 
 
 def transform_remove_punctuation(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return regexp_replace(input_col, r"[?\-\\\/\"\':,.\[\]\{\}]+", "")
 
 
 def transform_replace_apostrophe(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return regexp_replace(input_col, r"'+", " ")
 
 
 def transform_remove_alternate_names(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return regexp_replace(input_col, r"(\w+)( or \w+)+", "$1")
 
 
 def transform_remove_suffixes(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     values = _require_key(transform, "values")
     suffixes = "|".join(values)
@@ -216,7 +297,7 @@ def transform_remove_suffixes(
 
 
 def transform_remove_stop_words(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     values = _require_key(transform, "values")
     words = "|".join(values)
@@ -225,7 +306,7 @@ def transform_remove_stop_words(
 
 
 def transform_remove_prefixes(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     values = _require_key(transform, "values")
     prefixes = "|".join(values)
@@ -234,7 +315,7 @@ def transform_remove_prefixes(
 
 
 def transform_condense_prefixes(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     values = _require_key(transform, "values")
     prefixes = "|".join(values)
@@ -243,38 +324,38 @@ def transform_condense_prefixes(
 
 
 def transform_condense_strip_whitespace(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return regexp_replace(trim(input_col), r"\s\s+", " ")
 
 
 def transform_remove_one_letter_names(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return regexp_replace(input_col, r"^((?:\w )+)(\w+)", r"$2")
 
 
 def transform_split(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return split(input_col, " ")
 
 
 def transform_length(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return length(input_col)
 
 
 def transform_array_index(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     value = _require_key(transform, "value")
     return input_col[value]
 
 
 def transform_mapping(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     mapped_column = input_col
     mappings = _require_key(transform, "mappings")
@@ -290,7 +371,7 @@ def transform_mapping(
 
 
 def transform_swap_words(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     mapped_column = input_col
     values = _require_key(transform, "values")
@@ -304,7 +385,7 @@ def transform_swap_words(
 
 
 def transform_substring(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     values = _require_key(transform, "values")
     if len(values) == 2:
@@ -318,27 +399,27 @@ def transform_substring(
 
 
 def transform_expand(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     expand_length = _require_key(transform, "value")
     return array([input_col + i for i in range(-expand_length, expand_length + 1)])
 
 
 def transform_cast_as_int(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return input_col.cast("int")
 
 
 def transform_divide_by_int(
-    input_col: Column, transform: dict[str, Any], context
+    input_col: Column, transform: Mapping[str, Any], context
 ) -> Column:
     divisor = _require_key(transform, "value")
     return input_col.cast("int") / divisor
 
 
 def transform_when_value(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     threshold = _require_key(transform, "value")
     if_value = _require_key(transform, "if_value")
@@ -347,6 +428,6 @@ def transform_when_value(
 
 
 def transform_get_floor(
-    input_col: Column, transform: dict[str, Any], context: dict[str, Any]
+    input_col: Column, transform: Mapping[str, Any], context: Mapping[str, Any]
 ) -> Column:
     return floor(input_col).cast("int")