Add rule for domain union validation

sshenzha · sshenzha · commit 3bad7e566d02 · 2026-03-20T22:42:08.000Z
diff --git a/src/fastssv/rules/data_quality/__init__.py b/src/fastssv/rules/data_quality/__init__.py
@@ -8,10 +8,12 @@
 from .schema_validation import SchemaValidationRule
 from .column_type_validation import ColumnTypeValidationRule
 from .negative_concept_id_validation import NegativeConceptIdValidationRule
+from .union_concept_id_domain_indicator import UnionConceptIdDomainIndicatorRule
 
 __all__ = [
     "UnmappedConceptHandlingRule",
     "SchemaValidationRule",
     "ColumnTypeValidationRule",
     "NegativeConceptIdValidationRule",
+    "UnionConceptIdDomainIndicatorRule",
 ]
diff --git a/src/fastssv/rules/data_quality/union_concept_id_domain_indicator.py b/src/fastssv/rules/data_quality/union_concept_id_domain_indicator.py
@@ -0,0 +1,275 @@
+"""Union Concept ID Domain Indicator Rule.
+
+OMOP semantic rule OMOP_067:
+UNION queries should not combine concept_id columns from different domains without a domain
+indicator. Mixing condition_concept_id and drug_concept_id in a single column without context
+makes results uninterpretable.
+
+The Problem:
+    Each OMOP domain has its own concept_id column:
+    - condition_occurrence.condition_concept_id
+    - drug_exposure.drug_concept_id
+    - procedure_occurrence.procedure_concept_id
+    - measurement.measurement_concept_id
+    - observation.observation_concept_id
+
+    UNION queries that mix these without domain labels create ambiguous results.
+
+Example impact:
+    SELECT condition_concept_id AS concept_id
+    FROM condition_occurrence
+    UNION ALL
+    SELECT drug_concept_id AS concept_id
+    FROM drug_exposure
+    -- Returns: [201826, 1545958, 313217, ...]
+    -- Which are conditions? Which are drugs? UNKNOWN!
+    -- Results are uninterpretable without domain context
+
+Violation pattern:
+    SELECT condition_concept_id AS concept_id
+    FROM condition_occurrence
+    UNION ALL
+    SELECT drug_concept_id AS concept_id
+    FROM drug_exposure
+    -- No domain indicator!
+
+Correct patterns:
+    -- Option 1: Add literal domain column
+    SELECT 'Condition' AS domain, condition_concept_id AS concept_id
+    FROM condition_occurrence
+    UNION ALL
+    SELECT 'Drug' AS domain, drug_concept_id AS concept_id
+    FROM drug_exposure
+
+    -- Option 2: Use domain_id from concept table
+    SELECT c.domain_id, co.condition_concept_id AS concept_id
+    FROM condition_occurrence co
+    JOIN concept c ON co.condition_concept_id = c.concept_id
+    UNION ALL
+    SELECT c.domain_id, de.drug_concept_id AS concept_id
+    FROM drug_exposure de
+    JOIN concept c ON de.drug_concept_id = c.concept_id
+"""
+
+from typing import Dict, List, Optional, Set
+
+from sqlglot import exp
+
+from fastssv.core.base import Rule, RuleViolation, Severity
+from fastssv.core.helpers import (
+    extract_aliases,
+    normalize_name,
+    parse_sql,
+    resolve_table_col,
+)
+from fastssv.core.registry import register
+
+
+# --- Constants -------------------------------------------------------------
+
+TABLE_TO_DOMAIN = {
+    "condition_occurrence": "condition",
+    "drug_exposure": "drug",
+    "procedure_occurrence": "procedure",
+    "measurement": "measurement",
+    "observation": "observation",
+    "device_exposure": "device",
+}
+
+VALID_DOMAIN_LITERALS = set(TABLE_TO_DOMAIN.values())
+
+
+# --- Helpers ---------------------------------------------------------------
+
+def _norm(x: Optional[str]) -> Optional[str]:
+    return normalize_name(x) if x else None
+
+
+# 1. Flatten UNION chains properly
+def _collect_union_selects(node: exp.Expression) -> List[exp.Select]:
+    selects: List[exp.Select] = []
+
+    def _walk(n: exp.Expression):
+        if isinstance(n, exp.Union):
+            _walk(n.this)
+            _walk(n.expression)
+        elif isinstance(n, exp.Select):
+            selects.append(n)
+
+    _walk(node)
+    return selects
+
+
+# 2. Extract concept_id columns WITH table lineage
+def _get_concept_id_domains(
+    select: exp.Select,
+    aliases: Dict[str, str],
+) -> Set[str]:
+    domains: Set[str] = set()
+
+    for expr in select.expressions:
+        col = None
+
+        if isinstance(expr, exp.Alias):
+            if isinstance(expr.this, exp.Column):
+                col = expr.this
+        elif isinstance(expr, exp.Column):
+            col = expr
+
+        if not col:
+            continue
+
+        col_name = _norm(col.name)
+        if not col_name or not col_name.endswith("_concept_id"):
+            continue
+
+        table, _ = resolve_table_col(col, aliases)
+        table = _norm(table)
+
+        # If qualified, use the table name
+        if table and table in TABLE_TO_DOMAIN:
+            domains.add(TABLE_TO_DOMAIN[table])
+        # If unqualified, infer from tables in aliases
+        elif not table:
+            # Check tables referenced in this SELECT
+            for alias_table in aliases.values():
+                norm_table = _norm(alias_table)
+                if norm_table in TABLE_TO_DOMAIN:
+                    # Infer domain from column name matching table
+                    # e.g., condition_concept_id → condition_occurrence
+                    expected_prefix = col_name.replace("_concept_id", "")
+                    if expected_prefix in norm_table:
+                        domains.add(TABLE_TO_DOMAIN[norm_table])
+
+    return domains
+
+
+# 3. Strong domain indicator detection
+def _has_domain_indicator(select: exp.Select) -> bool:
+    for expr in select.expressions:
+        # Case 1: explicit domain column
+        if isinstance(expr, exp.Column):
+            if _norm(expr.name) in {"domain", "domain_id"}:
+                return True
+
+        # Case 2: aliased literal with domain meaning
+        if isinstance(expr, exp.Alias):
+            alias_name = _norm(expr.alias)
+
+            if isinstance(expr.this, exp.Literal):
+                value = _norm(str(expr.this.this))
+
+                # Must match known domain values
+                if value in VALID_DOMAIN_LITERALS:
+                    return True
+
+                # OR alias explicitly says domain
+                if alias_name in {"domain", "domain_id"}:
+                    return True
+
+    return False
+
+
+# 4. Core detection
+def _find_violations(tree: exp.Expression) -> List[str]:
+    issues: List[str] = []
+    seen: Set[str] = set()
+
+    for union in tree.find_all(exp.Union):
+        # Only process top-level UNIONs (skip nested ones to avoid duplicates)
+        parent = union.parent
+        is_nested = False
+        while parent:
+            if isinstance(parent, exp.Union):
+                is_nested = True
+                break
+            parent = parent.parent if hasattr(parent, 'parent') else None
+
+        if is_nested:
+            continue
+
+        selects = _collect_union_selects(union)
+
+        if len(selects) < 2:
+            continue
+
+        all_domains: Set[str] = set()
+
+        for select in selects:
+            aliases = extract_aliases(select)
+            domains = _get_concept_id_domains(select, aliases)
+            all_domains.update(domains)
+
+        # No concept_id usage → skip
+        if not all_domains:
+            continue
+
+        # Single domain → OK
+        if len(all_domains) <= 1:
+            continue
+
+        # Check for domain indicator
+        has_indicator = any(_has_domain_indicator(s) for s in selects)
+
+        if has_indicator:
+            continue
+
+        key = tuple(sorted(all_domains))
+        if key in seen:
+            continue
+        seen.add(key)
+
+        domains_str = ", ".join(sorted(d.capitalize() for d in all_domains))
+
+        issues.append(
+            f"UNION combines concept_id columns from multiple domains ({domains_str}) "
+            f"without a domain indicator column. This makes concept_id values ambiguous."
+        )
+
+    return issues
+
+
+# --- Rule ------------------------------------------------------------------
+
+@register
+class UnionConceptIdDomainIndicatorRule(Rule):
+    """Validates domain disambiguation in UNION queries with concept_id."""
+
+    rule_id = "semantic.union_concept_id_domain_indicator"
+    name = "Union Concept ID Domain Indicator"
+    description = (
+        "UNION queries combining concept_id values from multiple domains must include "
+        "a domain indicator column to avoid ambiguity."
+    )
+    severity = Severity.WARNING
+    suggested_fix = (
+        "Add a domain indicator column to each SELECT, e.g.: "
+        "SELECT 'Condition' AS domain, condition_concept_id FROM ... UNION ..."
+    )
+
+    def validate(self, sql: str, dialect: str = "postgres") -> List[RuleViolation]:
+        sql_lower = sql.lower()
+
+        # Fast pre-check
+        if "union" not in sql_lower or "_concept_id" not in sql_lower:
+            return []
+
+        trees, err = parse_sql(sql, dialect)
+        if err:
+            return []
+
+        violations: List[RuleViolation] = []
+
+        for tree in trees:
+            if not tree:
+                continue
+
+            issues = _find_violations(tree)
+
+            for msg in issues:
+                violations.append(self.create_violation(message=msg))
+
+        return violations
+
+
+__all__ = ["UnionConceptIdDomainIndicatorRule"]
diff --git a/tasks/IMPLEMENTATION_STATUS.md b/tasks/IMPLEMENTATION_STATUS.md
@@ -9,8 +9,8 @@ This checklist tracks which rules from `omop_rules.json` have been implemented i
 
 **Statistics:**
 - Total rules in JSON: 350+
-- Implemented: 62 rules (including covered rules)
-- Coverage: 17.7%
+- Implemented: 63 rules (including covered rules)
+- Coverage: 18.0%
 - Last updated: March 2026
 
 ---
@@ -153,12 +153,12 @@ This checklist tracks which rules from `omop_rules.json` have been implemented i
   - *Status: SKIPPED - This is an analytics best practice, not a semantic CDM violation. The rule would detect COUNT(*) vs COUNT(DISTINCT person_id) in prevalence queries, but this is more about query intent (counting patients vs events) which is highly context-dependent and subjective. Many legitimate use cases exist for counting total events. Better addressed through documentation, code review, and analytics training rather than static SQL validation. High false positive risk due to inability to determine query intent.*
 - [x] **OMOP_064**: drug_strength_valid_start_end_date_filter
   - *Implemented as: `domain_specific/drug/drug_strength_validity_filter.py`*
-- [ ] **OMOP_065**: observation_table_heterogeneous_domain
-  - *Suggested group: `domain_specific/observation/`*
+- [-] **OMOP_065**: observation_table_heterogeneous_domain
+  - *Status: SKIPPED - This is an analytics best practice, not a semantic violation. While the observation table is heterogeneous (surveys, social history, family history), requiring observation_concept_id filter is prescriptive about query patterns rather than catching semantic errors. Many legitimate use cases exist for querying observation without concept filters (metadata queries, exploratory analysis, data quality checks). Better addressed through analytics documentation and query guidelines rather than static SQL validation. High false positive risk for legitimate analytical queries.*
 - [x] **OMOP_066**: concept_domain_id_matches_target_table
   - *Implemented as: `concept_standardization/concept_domain_validation.py` (comprehensive merged rule covering all 35+ concept_id columns)*
-- [ ] **OMOP_067**: no_union_different_concept_id_types
-  - *Suggested group: `data_quality/`*
+- [x] **OMOP_067**: no_union_different_concept_id_types
+  - *Implemented as: `data_quality/union_concept_id_domain_indicator.py`*
 - [ ] **OMOP_068**: specimen_table_not_for_lab_results
   - *Suggested group: `domain_specific/specimen/`*
 - [ ] **OMOP_069**: cohort_definition_id_required_for_cohort_query
diff --git a/tests/test_semantic_validation.py b/tests/test_semantic_validation.py
