def modify_column_by_udf()

in utilities/Hive_metastore_migration/src/hive_metastore_migration.py [0:0]

• 5 lines of code
• 4 McCabe index (conditional complexity)

    def modify_column_by_udf(df, udf, column_to_modify, new_column_name=None):
        """
        transform a column of the dataframe with the user-defined function, keeping all other columns unchanged.
        :param new_column_name: new column name. If None, old column name will be used
        :param df: dataframe
        :param udf: user-defined function
        :param column_to_modify: the name of the column to modify.
        :type column_to_modify: str
        :return: the dataframe with single column modified
        """
        if new_column_name is None:
            new_column_name = column_to_modify
        return df.select(
            *[udf(column).alias(new_column_name) if column == column_to_modify else column for column in df.columns])