如何在不创建结果 dataframe 的情况下使用 pandas 从内部连接中获取索引对？

Question

I would like to merge two dataframes, but instead of returning the full join just return two columns showing the paired rows.

I've written an example function below, but it creates the inner join of the two dataframes, including all columns that are not needed. If there are many columns and rows, this can use a lot of memory.

example

def get_index_pairs(df1, df2, on):
     return pd.merge(df1.reset_index(), df2.reset_index(), on=on)[['index_x', 'index_y']]
df1 = pd.DataFrame( dict( key = ["a","b","c","d"], v1=[1,2,3,4]))
df2 = pd.DataFrame( dict( key = ["b","d","f","g"], v2=[10,20,30,40]))

pairs = get_index_pairs(df1, df2, on='key')

print(pairs)

output

   index_x  index_y
0        1        0
1        3        1

I'm looking for a more memory efficient version of get_index_pairs .

Answer 1

Solution for Duplicate Keys (edited)

Data

The data is slightly adjusted for demo purpose. Specifically, key="d" will have a 2*2 Cartesian join.

import pandas as pd
import numpy as np
df1 = pd.DataFrame( dict( key = ["a","b","d","d"], v1=[1,2,3,4]))
df2 = pd.DataFrame( dict( key = ["b","d","d","g"], v2=[10,20,30,40]))

Code

Use np.argwhere() to return all matched indices.

ls = []
for i1, k1 in enumerate(df1["key"]):  # the only slow step (explicit for loop)
    for i2 in np.argwhere((df2["key"] == k1).values):
        ls.append([i1, i2[0]])

df_ans = pd.DataFrame(ls, columns=["index_x", "index_y"])

Result

print(df_ans)

   index_x  index_y
0        1        0
1        2        1   <-  Cartesian join on "d" like what
2        2        2   <-  would be produced by an
3        3        1   <-  inner join
4        3        2   <-

Note

The OP asked for reduced memory usage , not faster execution. Otherwise pd.merge() would be preferred in terms of speed.

Old Solution (no dup keys)

Construct dict s containing {key: index} pairs for your search. Only a constant multiple of the size of dfN["key"] would be consumed during the search process.

Code

# {key: index} mappings
dic1 = dict(zip(df1["key"].values, range(len(df1))))
dic2 = dict(zip(df2["key"].values, range(len(df2))))

# collect matched results
ls = []
for k1, v1 in dic1.items():  # the only slow step (explicit for loop)
    if k1 in dic2:  # fast (hashed search)
        ls.append([v1, dic2[k1]])

df_ans = pd.DataFrame(ls, columns=["index_x", "index_y"])

Result

print(df_ans)

   index_x  index_y
0        1        0
1        3        1

Answer 2

Another option is to use pandas DataFrame align , to align the axes bases on the key :

def get_index_pairs(df1, df2, on):
    left, right = pd.DataFrame.align(
        df1.set_index(on), df2.set_index(on), axis="index", join="inner"
    )
    left_index = df1.index[df1.loc[:, on].isin(left.index)]
    right_index = df2.index[df2.loc[:, on].isin(right.index)]
    return pd.DataFrame({"index_x": left_index, "index_y": right_index})


df1 = pd.DataFrame(dict(key=["a", "b", "c", "d"], v1=[1, 2, 3, 4]))
df2 = pd.DataFrame(dict(key=["b", "d", "f", "g"], v2=[10, 20, 30, 40]))

pairs = get_index_pairs(df1, df2, on="key")

print(pairs)

     index_x  index_y
0        1        0
1        3        1