topic Dynamic Python UDF in unity catalog in Data Engineering

Dynamic Python UDF in unity catalog

Dnirmania — Wed, 28 Aug 2024 17:47:40 GMT

Hi Team

I am trying to create a python UDF which I want to use for column masking. This function will take 2 input parameters(column name and groups name) and return the column value if user is part of group otherwise return masked value. I wrote following line of code but it seems we can't uses spark.sql in python UDF to execute sql query.

Can anyone help me to execute SQL query in python function?

CREATE OR REPLACE FUNCTION ryanlakehouse.default.column_masking(column_value STRING, groups_str String)

RETURNS STRING

LANGUAGE PYTHON

AS $$

def mask_column(column, groups_str):

groups = groups_str.split(',')

for group in groups:

if spark.sql(f"SELECT is_account_group_member('{group}')").collect()[0][0]:

return column

return '******'

column = "{column_value}"

groups_str = groups_str

return mask_column(column, groups_str)

$$;

Re: Dynamic Python UDF in unity catalog

menotron — Thu, 29 Aug 2024 01:06:47 GMT

Hi @Dnirmania, You could achieve something similar using this UDF:

%sql CREATE OR REPLACE FUNCTION ryanlakehouse.default.column_masking(column_value STRING, groups_str String) RETURNS STRING LANGUAGE SQL COMMENT 'Return the column value if user is part of group otherwise return masked value' RETURN CASE WHEN array_contains( transform( split(groups_str, ','), x -> is_account_group_member(trim(x)) ), True ) THEN column_value ELSE repeat("*", len(column_value)) END

Re: Dynamic Python UDF in unity catalog

Dnirmania — Thu, 29 Aug 2024 08:34:26 GMT

Thank you so much @menotron. Your suggestion worked perfectly.