Dask - map_partition

Question

我有一个带有纬度和经度集（约 32m 条记录）的 Dask DataFrame。我正在尝试使用如下函数计算纬度/经度之间的距离：

import numpy as np
from geopy import distance

def calc_distance(df, lat_col_name_1, lon_col_name_1, lat_col_name_2, lon_col_name_2):
if df[lat_col_name_1] != np.nan and df[lon_col_name_1] != np.nan and df[lat_col_name_2] != np.nan and df[lon_col_name_2] != np.nan:
    return distance.distance((df[lat_col_name_1], df[lon_col_name_1]), (df[lat_col_name_2], df[lon_col_name_2])).miles
else:
    return np.nan 

我尝试使用 map_partitions 调用此函数（以创建索引和距离的 DataFrame 以及使用 assign 调用 map_paritions。我想使用 assign 以避免将 DataFrame 重新连接在一起（似乎很昂贵）。它确实不像 np.nan 检查。我得到了一个

ValueError：Series 的真值不明确。使用 a.empty， a.bool()、a.item()、a.any() 或 a.all()。

我有零纬度/经度的记录，所以我需要在计算距离时考虑到这一点。

使用 map_partitions

distance = big_df.map_partitions(calc_distance, 
                                    lat_col_name_1='latitude_1', 
                                    lon_col_name_1='longitude_1', 
                                    lat_col_name_2='latitude_2', 
                                    lon_col_name_2='longitude_2', 
                                    meta={'distance': np.float64})

使用 map_partitions 和分配

def calc_distance_miles(lat1, lon1, lat2, lon2):
    if lat1 != np.nan and lon1 != np.nan and lat2 != np.nan and lon2 != np.nan:
        return distance.distance((lat1, lon1), (lat2, lon2)).miles
    else:
        return np.nan
    

big_df = big_df.map_partitions(lambda df: df.assign(
    distance=calc_distance_miles(df['latitude_1'], df['longitude_1'], df['latitude_2'], df['longitude_2'])
), meta={'distance': np.float64}
)

Answer 1

map_partitions 不像 df.apply - 函数 calc_distance 是通过 dask.dataframe 的分区调用的，该分区的类型为 pd.DataFrame。

因此，df[lat_col_name_1] 是一个系列，df[lat_col_name_1] != np.nan 是一个布尔系列（它将始终返回此错误 - 参见例如 Truth value of a Series is ambiguous. Use a.empty, a.bool(), a.item(), a.any() or a.all()）。

有比按元素计算距离更快的数组方法，但与您尝试做的类似的 dask.dataframe 是使用 map_partitions 然后应用：

def calc_distance(series, lat_col_name_1, lon_col_name_1, lat_col_name_2, lon_col_name_2):
    if series[
        [lat_col_name_1, lon_col_name_1, lat_col_name_2, lon_col_name_2]
    ].notnull().all():

        return distance.distance(
            (series[lat_col_name_1], series[lon_col_name_1]),
            (series[lat_col_name_2], series[lon_col_name_2]),
        ).miles

    else:
        return np.nan 

def calc_distance_df(df, **kwargs):
    return df.apply(calc_distance, axis=1, **kwargs)

distances = big_df.map_partitions(
    calc_distance_df,
    meta=np.float64,
    lat_col_name_1=lat_col_name_1,
    lon_col_name_1=lon_col_name_1,
    lat_col_name_2=lat_col_name_2,
    lon_col_name_2=lon_col_name_2,
)