Custom functions can't be used in temporary views with CTEs

The following query works in Spark 3.0 not Spark 3.1.
 
Start with:
spark.udf.registerJavaFunction("custom_func", "com.stuff.path.custom_func", LongType())
 
Works: * select custom_func()

• create temporary view blaah as select custom_func()
• with step_1 as ( select custom_func() ) select * from step_1

Broken:
create temporary view blaah as with step_1 as ( select custom_func() ) select * from step_1
 
followed by:
select * from blaah
 
Error:
Error in SQL statement: AnalysisException: No handler for UDF/UDAF/UDTF 'com.stuff.path.custom_func';