"""Create unique RAD_ID for each record in a dataset.
The dataset should contain:
1. place name (nullable)
2. address
3. city
4. location_id (nullable)
5. state
6. zip code
7. latitude (nullable)
8. longitude (nullable)
The `CreateRadIdNativeOperator` is an Airflow operator designed to identify and resolve duplicate records in a dataset by assigning unique RAD_IDs (Record Aggregation and Deduplication IDs).
## What it does
The operator:
1. Identifies potential duplicate records in a dataset using fuzzy matching on address and entity name
2. Creates clusters of duplicate records using a union-find algorithm
3. Assigns a common RAD_ID to records in the same cluster (using the minimum ID in the cluster)
4. Returns the full dataset with new RAD_ID values and a flag indicating duplicate records
## How the SQL works
The operator uses two SQL queries (`SQL_WITH_LOC` and `SQL_NO_LOC`):
1. Both queries:
- Create a CTE (Common Table Expression) to prepare data for comparison
- Join the table to itself to find potential duplicate pairs
- Compare addresses using SOUNDEX (phonetic algorithm) and LEVENSHTEIN_DISTANCE (edit distance)
- Ensure pairs are unique by using `a.entity_id_old < b.entity_id_old`
2. Key matching criteria:
- `SOUNDEX(a.address) = SOUNDEX(b.address)` - addresses sound similar
- `LEVENSHTEIN_DISTANCE(LOWER(a.address), LOWER(b.address)) < 5` - addresses are at most 4 edits apart
- `LEVENSHTEIN_DISTANCE(LOWER(a.entity_name), LOWER(b.entity_name)) < 3` - entity names are at most 2 edits apart
3. Additional criteria when `location_id` is available:
- `a.location_id = b.location_id` - exact location ID match
- Street number matching
## How to use it
```python
create_rad_id_task = CreateRadIdNativeOperator(
task_id='create_rad_id',
table_name='your_athena_table',
id_col_old='original_id_column',
id_col_new='rad_id', # optional, defaults to "rad_id"
city_col='city_column', # optional, defaults to "city"
state_col='state_column', # optional, defaults to "state"
address_col='address_column', # optional, defaults to "address_primary"
place_name_col='place_name_column', # optional, defaults to "entity_name"
location_id_col='location_id_column', # optional, if available
dag=dag
)
```
## When to use it
This operator is ideal for:
- Datasets containing place/business information with potential duplicates
- Records that have address, city, and name information
- Cases where you need to identify and consolidate duplicate entities
- Data cleansing and preprocessing tasks
## When not to use it
Avoid using this operator when:
1. Your dataset doesn't contain address and name information
2. You need perfect matching rather than fuzzy matching
3. The dataset is extremely large (the SQL performs self-joins which can be expensive)
4. You need customized deduplication logic beyond what the preset thresholds allow
## Important notes
1. As mentioned in the log message: "the thresholds are preset. You may not find all dups with current thresholds."
2. The operator returns a dictionary of records but doesn't write back to Athena (there's a "TO FIX by YUAN" comment regarding this)
3. The deduplication is based primarily on address and entity name similarity
4. The operator adds an `is_dup` column to flag records that were identified as duplicates
The operator uses AWS Wrangler to interact with Athena and pandas for data manipulation.
"""
from config import Config
import pandas as pd
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
import awswrangler as wr
SOURCE_DB = Config.DATA_LAKE_DB
SQL_WITH_LOC = """
WITH raw AS (
SELECT
{id_col_old} AS entity_id_old,
CONCAT(regexp_replace({address_col},'[^\\x20-\\x7E]', ''), ', ', regexp_replace({city_col},'[^\\x20-\\x7E]', '')) AS address,
regexp_replace({place_name_col},'[^\\x20-\\x7E]', '') AS entity_name,
COALESCE({location_id_col}, 'NULL') AS location_id,
-- Get street number fom the address_col
regexp_replace(regexp_replace({address_col},'[^0-9]', ' '), ' +', ' ') AS street_number
FROM {table_name}
)
SELECT
a.entity_id_old AS id1,
b.entity_id_old AS id2
FROM raw AS a
JOIN raw AS b
ON a.entity_id_old < b.entity_id_old
AND SOUNDEX(a.address) = SOUNDEX(b.address)
AND LEVENSHTEIN_DISTANCE(LOWER(a.address), LOWER(b.address)) < 5
AND LEVENSHTEIN_DISTANCE(LOWER(a.entity_name), LOWER(b.entity_name)) < 3
AND a.location_id = b.location_id
AND COALESCE(CAST(a.street_number AS VARCHAR), '') = COALESCE(CAST(b.street_number AS VARCHAR), '')
"""
SQL_NO_LOC = """
WITH raw AS (
SELECT
{id_col_old} AS entity_id_old,
CONCAT({address_col}, ', ', {city_col}) AS address,
{place_name_col} AS entity_name
FROM {table_name}
)
SELECT
a.entity_id_old AS id1,
b.entity_id_old AS id2
FROM raw AS a
JOIN raw AS b
ON a.entity_id_old < b.entity_id_old
AND SOUNDEX(a.address) = SOUNDEX(b.address)
AND LEVENSHTEIN_DISTANCE(LOWER(a.address), LOWER(b.address)) < 5
AND LEVENSHTEIN_DISTANCE(LOWER(a.entity_name), LOWER(b.entity_name)) < 3
"""
[docs]
class CreateRadIdNativeOperator(BaseOperator):
"""Create unique RAD_ID for each record in a dataset.
The dataset should contain:
1. place name (nullable)
2. address
3. city
4. location_id (nullable)
5. state
6. zip code
7. latitude (nullable)
8. longitude (nullable)
:param table_name: The name of the table.
:param id_col_old: The name of the original id column.
:param id_col_new: The name of the new rad_id column.
:param city_col: The name of the city column.
:param state_col: The name of the state column.
:param address_col: The name of the address column.
:param place_name_col: The name of the place name column.
:param location_id_col: The name of the location id column. If None, the operator will not consider location_id in the duplicate detection.
:return: A dictionary of full records with new rad_id, and is_dup column
"""
template_fields = (
"table_name",
"id_col_old",
"id_col_new",
"location_id_col",
"city_col",
"state_col",
"address_col",
"place_name_col",
)
@apply_defaults
def __init__(
self,
table_name: str = None,
id_col_old: str = None,
id_col_new: str = "rad_id",
city_col: str = "city",
state_col: str = "state",
address_col: str = "address_primary",
place_name_col: str = "entity_name",
location_id_col: str = None,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.table_name = table_name
self.id_col_old = id_col_old
self.id_col_new = id_col_new
self.city_col = city_col
self.state_col = state_col
self.address_col = address_col
self.place_name_col = place_name_col
self.location_id_col = location_id_col
[docs]
def execute(self, context):
if not self.location_id_col:
sql = SQL_NO_LOC.format(
id_col_old=self.id_col_old,
address_col=self.address_col,
place_name_col=self.place_name_col,
table_name=self.table_name,
city_col=self.city_col,
)
else:
sql = SQL_WITH_LOC.format(
id_col_old=self.id_col_old,
address_col=self.address_col,
place_name_col=self.place_name_col,
table_name=self.table_name,
city_col=self.city_col,
location_id_col=self.location_id_col,
)
self.log.info("Running SQL query to find duplicate pairs")
self.log.info(sql)
dup_pair_df = wr.athena.read_sql_query(sql, database=SOURCE_DB)
self.log.info("Found %d duplicate pairs", len(dup_pair_df))
self.log.info(
"WARNING: the thresholds are preset. You may not find all dups with current thresholds. Check the Operator for details."
)
if len(dup_pair_df) == 0:
print("No duplicate pairs found.")
return []
# 2. Build a mapping from original id to new rad_id via union-find.
mapping = self._create_rad_id_mapping(dup_pair_df)
# 3. Read the full dataset from Athena again.
self.log.info("Reading full dataset from Athena")
full_df = wr.athena.read_sql_query(
f"SELECT * FROM {self.table_name}",
database=SOURCE_DB,
)
self.log.info("Total records: %d", len(full_df))
# 4. Assign new rad_id: if record id exists in mapping, use it; else, use original id.
self.log.info("Assigning new rad_id based on duplicate clusters")
full_df[self.id_col_new] = (
full_df[self.id_col_old].astype(str).apply(lambda x: mapping.get(x, x))
)
# 5. Add a column to indicate if the record is a duplicate for ease of futher analysis
full_df["is_dup"] = (
full_df[self.id_col_old].astype(str).apply(lambda x: x in mapping)
)
self.log.info("Sample output with new rad_id:")
self.log.info(full_df[[self.id_col_old, self.id_col_new]].head().to_string())
# 6. TO FIX by YUAN. Directly upload a new table to Athena datalake.
return full_df.to_dict(orient="records")
def _create_rad_id_mapping(self, dup_pair_df: pd.DataFrame) -> dict:
"""
Given a DataFrame of duplicate pairs with columns 'id1' and 'id2', create a mapping
from each record (id) to a new rad_id representing its connected component.
"""
# Ensure ids are strings.
dup_pair_df["id1"] = dup_pair_df["id1"].astype(str)
dup_pair_df["id2"] = dup_pair_df["id2"].astype(str)
# Build list of duplicate pairs.
pairs = list(zip(dup_pair_df["id1"], dup_pair_df["id2"]))
# Build set of all IDs in duplicate pairs.
all_ids = set()
for a, b in pairs:
all_ids.add(a)
all_ids.add(b)
all_ids = list(all_ids)
# Initialize union-find structures.
parent = {i: i for i in all_ids}
rank = {i: 0 for i in all_ids}
def find(x):
if parent[x] != x:
parent[x] = find(parent[x])
return parent[x]
def union(x, y):
root_x = find(x)
root_y = find(y)
if root_x != root_y:
if rank[root_x] < rank[root_y]:
parent[root_x] = root_y
elif rank[root_x] > rank[root_y]:
parent[root_y] = root_x
else:
parent[root_y] = root_x
rank[root_x] += 1
# Process each duplicate pair.
for a, b in pairs:
union(a, b)
# Group ids by their root.
clusters = {}
for i in all_ids:
root = find(i)
clusters.setdefault(root, []).append(i)
# For each cluster, assign the smallest id as the new rad_id.
mapping = {}
for cluster_ids in clusters.values():
new_id = min(cluster_ids)
for record_id in cluster_ids:
mapping[record_id] = new_id
return mapping