如何加快涉及具有复杂条件的 2 个不同数据帧的 pandas 数据帧迭代？答案

【问题标题】：How to speed up pandas dataframe iteration involving 2 different dataframes with a complex condition?如何加快涉及具有复杂条件的 2 个不同数据帧的 pandas 数据帧迭代？
【发布时间】：2021-11-29 23:00:32
【问题描述】：

我有一个大约 300000 行的 pandas 数据框 A。每行都有一个纬度和经度值。

我还有一个大约 10000 行的第二个 pandas 数据帧 B，它有一个 ID 号、最大和最小纬度以及最大和最小经度。

对于A中的每一行，我需要B中对应行的ID，这样A中的行的经纬度就包含在B中的行所代表的边界框内。

到目前为止，我有以下内容：

ID_list = []

for index, row in A.iterrows():
    filtered_B = B.apply(lambda x : x['ID'] if row['latitude'] >= x['min_latitude']
                                            and row['latitude'] < x['max_latitude'] \
                                            and row['longitude'] >= x['min_longitude'] \
                                            and row['longitude'] < x['max_longitude'] \
                                            else None, axis = 1)
    ID_list.append(B.loc[filtered_B == True]['ID']

创建 ID_list 变量的目的是将其作为 ID 列添加到 A。包括大于或等于和小于条件，以便 A 中的每一行只有一个来自 B 的 ID。

上面的代码在技术上是可行的，但它每分钟完成大约 1000 行，这对于这么大的数据集是不可行的。

任何提示将不胜感激，谢谢。

编辑：示例数据框：

答：

location	latitude	longitude
1	-33.81263	151.23691
2	-33.994823	151.161274
3	-33.320154	151.662009
4	-33.99019	151.1567332

乙：

ID	min_latitude	max_latitude	min_longitude	max_longitude
9ae8704	-33.815	-33.810	151.234	151.237
2ju1423	-33.555	-33.543	151.948	151.957
3ef4522	-33.321	-33.320	151.655	151.668
0uh0478	-33.996	-33.990	151.152	151.182

预期输出：

ID_list = [9ae8704, 0uh0478, 3ef4522, 0uh0478]

【问题讨论】：

请分享示例数据帧，以及预期的输出。
您没有添加预期的输出数据帧
A 中的值一定会在 B 中找到吗？如果不是，那么 A 中该行的输出应该是什么？会有多个匹配项吗？
如果没有找到，填上nan。而且只有一场比赛
尝试 np.vectorize 或此处提到的其他替代方法，而不是 apply，看看它是否仍然很慢。 stackoverflow.com/questions/52673285/…

标签： python pandas

【解决方案1】：

我认为 Geopandas 将是确保涵盖所有边缘情况的最佳解决方案，例如子午线/赤道交叉点等 - 空间查询正是 Geopandas 的设计目标。不过安装起来可能会很痛苦。

numpy 中的一种天真的方法（假设大多数点不会在任何地方改变符号）将计算每个子句作为差异的符号，然后仅保留所有符号都符合您的标准的匹配项。

对于像这样的大量重复计算，通常最好将其从 pandas 中转储到 numpy 中，进行处理，然后再将其放回 pandas。

a_lats = A.latitude.values.reshape(-1,1)
b_min_lats = B.min_latitude.values.reshape(1,-1)
b_max_lats = B.max_latitude.values.reshape(1,-1)

a_lons = A.longitude.values.reshape(-1,1)
b_min_lons =B.min_longitude.values.reshape(1,-1)
b_max_lons = B.max_longitude.values.reshape(1,-1)

north_of_min_lat = np.sign(a_lats - b_min_lats)
south_of_max_lat = np.sign(b_max_lats - a_lats)

west_of_min_lon = np.sign(a_lons - b_min_lons)
east_of_max_lon = np.sign(b_max_lons - a_lons)

margin_matches = (north_of_min_lat + south_of_max_lat + west_of_min_lon + east_of_max_lon)

match_indexes = (margin_matches == 4).nonzero()

matches = [(A.location[i], B.ID[j]) for i, j in zip(match_indexes[0], match_indexes[1])]
print(matches)

PS - 如果您使用 CuPy 并将所有对 numpy 的引用替换为 cupy，您可以轻松地在 GPU 上运行它。

【讨论】：

【解决方案2】：

我会使用 geopandas 来做到这一点，它利用了 rtree 索引。

import geopandas as gpd
from shapely.geometry import box

a_gdf = gpd.GeoDataFrame(a[['location']], geometry=gpd.points_from_xy(a.longitude,
                                                                      a.latitude))
b_gdf = gpd.GeoDataFrame(
    b[['ID']], 
    geometry=[box(*bounds) for _, bounds in b.loc[:, ['min_longitude',
                                                      'min_latitude', 
                                                      'max_longitude', 
                                                      'max_latitude']].iterrows()])

gpd.sjoin(a_gdf, b_gdf)

输出：

	location	geometry	index_right	ID
0	1	POINT (151.23691 -33.81263)	0	9ae8704
1	2	POINT (151.161274 -33.994823)	3	0uh0478
3	4	POINT (151.1567332 -33.99019000000001)	3	0uh0478
2	3	POINT (151.662009 -33.320154)	2	3ef4522

【讨论】：

【解决方案3】：

你的代码消耗

还有这段代码

A = pd.DataFrame.from_dict(
    {'location': {0: 1, 1: 2, 2: 3, 3: 4},
 'latitude': {0: -33.81263, 1: -33.994823, 2: -33.320154, 3: -33.99019},
 'longitude': {0: 151.23691, 1: 151.161274, 2: 151.662009, 3: 151.1567332}}
)
B = pd.DataFrame.from_dict(
    {'ID': {0: '9ae8704', 1: '2ju1423', 2: '3ef4522', 3: '0uh0478'},
 'min_latitude': {0: -33.815, 1: -33.555, 2: -33.321, 3: -33.996},
 'max_latitude': {0: -33.81, 1: -33.543, 2: -33.32, 3: -33.99},
 'min_longitude': {0: 151.234, 1: 151.948, 2: 151.655, 3: 151.152},
 'max_longitude': {0: 151.237, 1: 151.957, 2: 151.668, 3: 151.182}}
)


def func(latitude, longitude):
    for y, x in B.iterrows():
        if (latitude >= x.min_latitude and latitude < x.max_latitude and longitude >= x.min_longitude and longitude < x.max_longitude):
            return x['ID']
np.vectorize(func)
A.apply(lambda x: func(x.latitude, x.longitude), axis=1).to_list()

消耗

所以新解决方案的速度提高了 2.33 倍

【讨论】：

一个更好的测试可能是在一个大数组上

【解决方案4】：

一种相当快速的方法可能是按纬度和经度对点进行预排序，然后使用 @987654321 分别通过纬度 (lat_min np.intersect1d 相交。

在我的测试中，对于 300K 点和 10K 非重叠框，运行时间不到 10 秒。

这是一个示例实现：

# create `ids` series to be populated with box IDs for each point
ids = pd.Series(np.nan, index=a.index)

# create series with points sorted by lats and lons
lats = a['latitude'].sort_values()
lons = a['longitude'].sort_values()

# iterate over boxes
for bi, r in b.set_index('ID').iterrows():
    # find points inside the box by latitude:
    i1, i2 = np.searchsorted(lats, [r['min_latitude'], r['max_latitude']])
    ix_lat = lats.index[i1:i2]
    
    # find points inside the box by longitude:
    j1, j2 = np.searchsorted(lons, [r['min_longitude'], r['max_longitude']])
    ix_lon = lons.index[j1:j2]
    
    # find points inside the box as intersection and set values in ids:
    ix = np.intersect1d(ix_lat, ix_lon)
    ids.loc[ix] = bi

ids.tolist()

输出（根据提供的样本数据）：

['9ae8704', '0uh0478', '3ef4522', '0uh0478']

【讨论】：

【解决方案5】：

我们可以在b 上创建一个多区间索引，然后将常规的loc 索引与a 行中的元组一起使用。当我们有一个包含低值和高值的表来存储变量时，区间索引在这种情况下很有用。

from io import StringIO

import pandas as pd

a = pd.read_table(StringIO("""
    location    latitude    longitude
    1   -33.81263   151.23691
    2   -33.994823  151.161274
    3   -33.320154  151.662009
    4   -33.99019   151.1567332
"""), sep='\s+')

b = pd.read_table(StringIO("""
    ID  min_latitude    max_latitude    min_longitude   max_longitude
    9ae8704     -33.815     -33.810     151.234     151.237
    2ju1423     -33.555     -33.543     151.948     151.957
    3ef4522     -33.321     -33.320     151.655     151.668
    0uh0478     -33.996     -33.990     151.152     151.182
"""), sep='\s+')

lat_index = pd.IntervalIndex.from_arrays(b['min_latitude'], b['max_latitude'], closed='left')
lon_index = pd.IntervalIndex.from_arrays(b['min_longitude'], b['max_longitude'], closed='left')
index = pd.MultiIndex.from_tuples(list(zip(lat_index, lon_index)), names=['lat_range', 'lon_range'])
b = b.set_index(index)

print(b.loc[list(zip(a.latitude, a.longitude)), 'ID'].tolist())

上面甚至可以处理a 中没有对应行的b 行，方法是用nan 优雅地填充这些值。

【讨论】：

区间索引是快速且正确的方法；但是，如果间隔重叠，您的解决方案将不够
重叠间隔确实会导致意外行为，但问题的 cmets 澄清间隔不会重叠。

【解决方案6】：

一个不错的选择可能是执行交叉产品合并并删除不需要的列。例如，您可能会这样做：

AB_cross = A.merge(
    B
    how = "cross"
)

现在我们有一个巨大的数据框，其中包含所有可能的匹配项，其中 B 中的 ID（或者可能没有，我们还不知道）可能具有符合 A 中点的边界。这速度很快，但会在内存中产生一个大型数据集，因为我们现在有一个 30000x10000 行长的数据集。

现在，我们需要通过相应地过滤数据集来应用我们的逻辑。这是一个 numpy 过程（据我所知），所以它是矢量化的并且非常快！我还要说，使用between 可能更容易让您的代码更加语义化。

请注意，下面我使用.between(inclusive = 'left') 表示您想要查看long/lat 是否为min_long

ID_list = AB_cross['ID'].loc[
    AB_cross['longitude'].between(AB_cross['min_longitude'], AB_cross['max_longitude'], inclusive = 'left') &
    AB_cross['latitude'].between(AB_cross['min_latitude'], AB_cross['max_latitude'], inclusive = 'left')
]

【讨论】：

我想我需要购买更多的内存才能做到这一点：D
我可以拆分我的数据框并使用此方法，但感谢您的建议