def replace_in_expression(
body: Expression, replace_lookup: Mapping[str,
Expression]) -> Expression:
ret = body.transform(lambda exp: replace_lookup[exp.column_name]
if isinstance(exp, Column) and exp.column_name in
replace_lookup else exp)
return ret
def transform_aggregate_functions_with_mismatched_nullable_parameters(
exp: Expression, ) -> Expression:
if (isinstance(exp, FunctionCall)
and exp.function_name in AGGREGATION_FUNCTIONS):
return exp.transform(cast_column_to_nullable)
return exp