PySpark structured streaming with applyInPandasWithState worked example

				
					def update_state(key:Tuple[int], pdfs:Iterator[pd.DataFrame], state:GroupState)->Iterator[pd.DataFrame]:

    # Initialise variables 
    str_state_status = "state_status: success" # Status is returned in output
    (int_id,) = key # Used to get the id from the key
    list_IoT_readings = [] # Used to contain all readings for a particular group
    tuple_state = () # Used to store individual IoT readings in a state
    df_output = pd.DataFrame() # initialise dataframe to return to output

    try:
        if state.hasTimedOut:
            ####################################################
            # State has timed out so return the final state and clear it from memory 

            # Get final tuple stored in the state
            (list_prev_IoT_readings,) = state.get

            ####################################################
            # Construct final id, IoT readings, IoT state, and status into a dataframe to return to output

            # Get final state from state tuple
            (left_engine_vibration_pass, right_engine_vibration_pass, has_landed) = agg_analyse_data(list_prev_IoT_readings, threshold=threshold)

            # Set status to include any error information in group state
            str_state_status = "state_status: state group timed out"

            # Create dataframe to yeild with flight id, flight readings and flight
            dict_data = {'flight_id': [id], 'IoT_readings': [list_prev_IoT_readings], 'left_engine_vibration_pass': [left_engine_vibration_pass], 'right_engine_vibration_pass': [right_engine_vibration_pass], 'state_status': [str_state_status]}
            
            df_output = pd.DataFrame(dict_data)

            # Remove state 
            state.remove() 

        else:
            ####################################################
            # State has not timed out so proccess new data

            #####################################
            # Get the new data

            # Loop through each pandas dataframe in the iterator. Each dataframe contains data for a group
            for df_for_group in pdfs:

                # Get a list of all data for a particular group
                list_IoT_readings = convert_pd_df_to_list(df_for_group)

            if state.exists:
                ####################################################
                # State already exists so append the new data
                # to the existing data

                # Get the list of data stored in the current tuple state
                (list_prev_IoT_readings,) = state.get

                # Append the new data list to the existing list of data points
                combined_IoT_readings = list_prev_IoT_readings + list_IoT_readings

                # place IoT readings list into a tuple
                tuple_state = (combined_IoT_readings,)

                # Update the current state
                state.update(tuple_state)

                ####################################################
                # Construct IoT readings state and dataframe to return to output

                (left_engine_vibration_pass, right_engine_vibration_pass, has_landed) = agg_analyse_data(combined_IoT_readings, threshold=95)

                str_state_status = "state_status: state group updated"

                # Create dataframe to yeild with flight id, flight readings and flight
                dict_data = {'flight_id': [int_id], 'IoT_readings': [combined_IoT_readings], 'left_engine_vibration_pass': [left_engine_vibration_pass], 'right_engine_vibration_pass': [right_engine_vibration_pass], 'state_status': [str_state_status]}
                
                df_output = pd.DataFrame(dict_data)

                # Stop holding group data in state once the flight has landed
                if(has_landed == True):
                    # Remove state 
                    state.remove() 

            else:
                ####################################################
                # State does not exist yet so initialise it. 
                
                # Place IoT readings list into a tuple
                tuple_state = (list_IoT_readings,)

                # Update the state with the initialised tuple
                state.update(tuple_state)

                ####################################################
                # Construct IoT readings state and dataframe to return to output

                (left_engine_vibration_pass, right_engine_vibration_pass, has_landed) = agg_analyse_data(list_IoT_readings, threshold=threshold)

                str_state_status = "state_status: state group initiated"

                # Create dataframe to yeild with flight id, flight readings and flight
                dict_data = {'flight_id': [int_id], 'IoT_readings': [list_IoT_readings], 'left_engine_vibration_pass': [left_engine_vibration_pass], 'right_engine_vibration_pass': [right_engine_vibration_pass], 'state_status': [str_state_status]}
                
                df_output = pd.DataFrame(dict_data)

                # Stop holding group data in state once the flight has landed
                if(has_landed == True):
                    # Remove state 
                    state.remove() 


            ####################################################
            # Reset timeout period for group state 

            # Assume maximum duration of flight is 3 hours so timeout after this
            #timeout_duration_ms = 1000 * 60 * 60 * 0.005
            #state.setTimeoutDuration(timeout_duration_ms)

    except Exception as e:
        str_state_status = f"state_status: Error occured - {e}"
    
        (list_prev_IoT_readings,) = state.get

        ############################################
        # return error data
        dict_data = {'flight_id': [int_id], 'IoT_readings': [list_prev_IoT_readings], 'left_engine_vibration_pass': [None], 'right_engine_vibration_pass': [None], 'state_status': [str_state_status]}

        df_output = pd.DataFrame(dict_data)

    # Yield dataframe back to calling function
    yield df_output