Problems with extending DataFrame-object with custom method in DBR 14.0

[hanspetter]

I noticed that extending the DataFrame object with a custom method does not work in DBR 14.0 This works fine in when using DBR 13.3. The below code throws a "TypeError: 'Column' object is not callable".

from pyspark.sql.functions import col,lit
from pyspark.sql.dataframe import DataFrame

def custom_method(df):
  new_df = df.withColumn("new_column", lit(1))
  return new_df

setattr(DataFrame, "custom_method", custom_method)
#Also tried: DataFrame.custom_method = custom_method

d = [{'empid': 1, 'name': 'Test'}]

my_df = spark.createDataFrame(d)

my_df = my_df.custom_method()

 Is this expected behaviour?