@daniel_sahal thanks for your reply! 

There is no timeout set on the service side actually. The same MyOcr function, when invoked outside a UDF has no issues processing the records. It only fails when the function is wrapped inside a UDF. So, I am guessing it has something to do with the Python UDF Timeout or some other issue. We did change the configuration parameter `spark.databricks.sql.execution.pythonUDFTimeout` to 200 seconds to allow the UDF to wait 200 seconds for each row before timing out. But that hasn't helped us either. 

Maybe there's a different interpretation to the spark.databricks.sql.execution.pythonUDFTimeout? Our understanding was that this refers to timeout at the row level, i.e, wait 200 seconds before a record goes to fail state. However, we could not find details on this anywhere in the databricks documentation except on one of the other topics in this community

Again, we are not sure if this is because of the timeout either. It could be something entirely different too.