Is your feature request related to a problem? Please describe.
We have @with_columns for pyspark, but not for other dataframes.
Describe the solution you'd like
@with_columns works for any dataframe type.
Describe alternatives you've considered
N/A
Additional context
This would help with reusing transforms across dataframes without getting into naming issues.