The very first "when" function results in the posted error message (see image). The print statement of the count of df_td_amm works. A printSchema of the "df_td_amm" data frame confirms that "AGE" is a column. A select statement is also successful, so wondering why the column isn't accessible in the "when" function. I've tried referencing the column df_td_amm.AGE and get the same error.
I've tried this on two other small dataframes not created using Pandas and encounter the same error. For two other multi-column dataframes, the "when" function works fine referencing the columns as expected.
import pandas as pd
import pyspark.sql
from pyspark.sql.functions import *
from pyspark.sql.types import *
months = list(range(26))
pdf = pd.DataFrame(months, columns = ['AGE'])
df_td_amm = spark.createDataFrame(pdf).cache()
print(df_td_amm.count())
df_age_bucket = df_td_amm \
.withColumn("COMPOUNDING_FREQ", \
when(col("AGE") <= 3, "00-03 Months")
.when(col("AGE") <= 6, "04-06 Months")
.when(col("AGE") <= 9, "07-09 Months")
.other("Other"))
display(df_age_bucket)
