@@ -1347,6 +1347,9 @@ def generate_batches(self, table_name: str) -> Any:
13471347 # Apply formulas
13481348 df_batch = self ._apply_formula_columns (df_batch , table_name )
13491349
1350+ # Apply conditional nulls (null_if)
1351+ df_batch = self ._apply_null_if (df_batch , table_name )
1352+
13501353 # Post-process
13511354 df_batch = self ._fix_correlated_columns (df_batch , table_name )
13521355
@@ -1683,6 +1686,27 @@ def _apply_single_constraint(self, df: pd.DataFrame, constraint: Any) -> pd.Data
16831686
16841687 return df
16851688
1689+ def _apply_null_if (self , df : pd .DataFrame , table_name : str ) -> pd .DataFrame :
1690+ """Set column to NaN/NaT where a sibling column matches a trigger value.
1691+
1692+ Reads ``null_if`` from each column's distribution_params:
1693+ null_if: {"column": "status", "values": ["cancelled", "refunded"]}
1694+ or the shorthand single-value form:
1695+ null_if: {"column": "status", "value": "cancelled"}
1696+ """
1697+ columns = self .config .columns .get (table_name , [])
1698+ for col in columns :
1699+ spec = col .distribution_params .get ("null_if" )
1700+ if not spec :
1701+ continue
1702+ ref_col = spec .get ("column" )
1703+ trigger_values = spec .get ("values" ) or ([spec ["value" ]] if "value" in spec else [])
1704+ if not ref_col or not trigger_values or ref_col not in df .columns or col .name not in df .columns :
1705+ continue
1706+ mask = df [ref_col ].isin (trigger_values )
1707+ df .loc [mask , col .name ] = np .nan
1708+ return df
1709+
16861710 def _apply_formula_columns (self , df : pd .DataFrame , table_name : str ) -> pd .DataFrame :
16871711 """Apply formula-based derived columns using context for lookups."""
16881712 try :
0 commit comments