predict.py 文件源码

def pivot_stations(df, engine):
    query = """
    SELECT
      d.*,
      s.name AS arrivalname
    FROM distance d INNER JOIN station s ON d.stationto = s.id
    """
    distances = pd.read_sql_query(query, con=engine)
    stations = distances['arrivalname'].unique().tolist()
    dist_pivot = pd.pivot_table(distances, values='distance', index=['stationfrom', 'date', 'trip'],
                                columns=['arrivalname'], aggfunc=np.mean)
    dist_pivot = dist_pivot.reindex(df.index.rename(['stationfrom', 'date', 'vehicle']))
    df = df.join(dist_pivot, how='outer')
    return df, stations