[Python] Use threading for UDF in databricks to use new api token

Discussão em 'Python' iniciado por Stack, Junho 11, 2025.

  1. Stack

    Stack Membro Participativo

    I have a dataframe that needs to be used row-wise for API calls. Therefore I implemented a UDF for it. It takes each row uses the data from that row to make an api call and stores an id from the response in the result dataframe. The call looks something like this:

    processDf = FilesDf.withColumn("processed", row_udf(col("input_id")).select("processed.*").cache()

    Now the issue is that the processing takes more than one hour and the headers I am passing inside the row_udf for the api call contain a token that is only valid for an hour. I set up the headers in a cell on top of the notebook.

    To overcome this I found threading in databricks.

    def get_new_token():
    app = TokenApplication(...)
    result = app.acquire_token_for_client(...)
    ...
    headers = {"Authorization": token}
    return headers

    headers = get_new_token()

    # Create a threading event to signal the thread to stop
    stop_event = threading.Event()

    def refresh_token_periodically():
    global headers
    while not stop_event.is_set():
    time.sleep(55 * 60) # Sleep for 55 minutes
    headers = get_new_token()

    # Start the background thread to refresh the token
    token_refresh_thread = threading.Thread(target=refresh_token_periodically)
    token_refresh_thread.daemon = True
    token_refresh_thread.start()


    I tested this with simple print/sleep statements in a new cell and saw that the token is refreshed. Now using it inside the udf does not work. I am wondering why.

    Can someone explain to me how the udf is triggering each row? Are all rows "prepared" at once and stacked in a waiting line so that the start of all rows is actually the same and they get the first token but the later rows are only really passed to the api later which then leads to a wrong token? Could this even work like I want it to?

    Continue reading...

Compartilhe esta Página