Join pandas DataFrames on available non-null keys

user_9148880 0
Viewed 238

I have a "base" DataFrame a that contains an identifier seq and join keys, as well as a "value" DataFrame b which will be merged:

import numpy as np
import pandas as pd

a = pd.DataFrame({
    'id':range(11),
    'seq':[
        1, 1, 1, 2, 2, 3, 3, 3, 4, 5, 5],
    'class':[
        'alpha', 'beta', 'gaga', np.nan, np.nan, np.nan, np.nan, np.nan,
        np.nan, 'beta', 'beta'],
    'style':[
        'x', 'x', 'x', 'y', 'y', np.nan, np.nan, np.nan, np.nan, np.nan,
        np.nan],
    'drama':[
        'no', 'no', 'no', 'yes', 'oh yes', 'no', 'yes', 'oh yes', np.nan,
        'yes', 'oh yes']})

b = pd.DataFrame({
    'class':[
        'gaga', 'alpha', 'alpha', 'alpha', 'alpha', 'alpha', 'beta', 'beta',
        'alpha', 'gaga', 'beta', 'beta', 'alpha', 'beta', 'gaga', 'alpha',
        'beta', 'alpha', 'beta', 'gaga', 'gaga', 'beta', 'beta', 'beta',
        'gaga'],
    'style':[
        'y', 'y', 'y', 'x', 'x', 'x', 'y', 'x', 'x', 'x', 'y', 'y', 'y', 'y',
        'x', 'y', 'y', 'y', 'y', 'x', 'x', 'x', 'y', 'x', 'y'],
    'drama':[
        'yes', 'no', 'no', 'no', 'no', 'oh yes', 'oh yes', 'oh yes', 'oh yes',
        'no', 'yes', 'oh yes', 'no', 'no', 'yes', 'yes', 'no', 'yes', 'oh yes',
        'oh yes', 'oh yes', 'no', 'oh yes', 'yes', 'yes'],
    'start':[
        838, 727, 700, 840, 530, 507, 871, 585, 120, 164, 562, 750, 953, 733,
        337, 307, 277, 972, 3, 805, 539, 600, 8, 382, 147],
    'end':[
        198, 328, 591, 427, 151, 126, 132, 149, 856, 725, 608, 726, 178, 521,
        316, 154, 633, 4, 113, 881, 258, 32, 354, 259, 958]})

What I would like to do is create a function that allows for a sort of "recursive" join against b, which will join only on available non-null columns for each row. In a's case:

  • Where class, style and drama are non-null, will join on three keys;
  • Where class is null, will join on style and drama only;
  • Where both class and style are null, will join only on drama;
  • Where all three columns are null, will "join" on the entire value DataFrame b;
  • The null columns are not necessarily the same: for instance, if style is null, it will join on class and drama.

Result-wise, in this example the output would be the same as manually doing the not-so-intelligent:

ll = []
x, y, z, w, v = a.loc[:2], a.loc[3:4], a.loc[5:7], a.loc[8:8], a.loc[9:]

ll.append(x.merge(b, on=['class', 'style', 'drama']))
ll.append(y.drop(columns='class').merge(b, on=['style', 'drama']))
ll.append(z.drop(columns=['class', 'style']).merge(b, on='drama'))
ll.append(v.drop(columns='style').merge(b, on=['class', 'drama']))

# For all null values, get the entire DataFrame
w['placeholder'] = 1
b['placeholder'] = 1
ll.append(w
    .drop(columns=['class', 'style', 'drama'])
    .merge(b, on='placeholder')
    .drop(columns='placeholder'))

result = pd.concat(ll)

However, doing manually in this case was only possible because I already knew beforehand how to segregate the "groups" (x, y, z, w and v) and which columns I would use for the merge operation in each one.

I have half-implemented it with very limited usability and using in my opinion a subpar way to deal with the columns:

def recjoin(base: pd.DataFrame, other: pd.DataFrame, keys: list) -> pd.DataFrame:
    missing_cols = base.columns[base.isnull().any()]

    if len(missing_cols) == 0:
        result = base.merge(other, on=idx)
    else:
        nonmissing = base.dropna(subset=keys)
        result_nonmissing = nonmissing.merge(other, on=keys)

        id_missing = base.index.difference(nonmissing.index)
        missing = base.loc[id_missing].drop(columns=missing_cols)

        if isinstance(keys, str):
            keys = [keys]

        alt_keys = list(pd.Index(keys).difference(missing_cols))
        result_missing = missing.merge(other, on=alt_keys)

        result = pd.concat([nonmissing, missing])

    return result

This way it will work if a.loc[:4] was passed to base, but not if it were a.loc[:7], because there is a variable number of NaN columns in the second case:

In [1]: a.loc[:4]
Out[1]:
   id  seq  class style   drama
0   0    1  alpha     x      no
1   1    1   beta     x      no
2   2    1   gaga     x      no
3   3    2    NaN     y     yes
4   4    2    NaN     y  oh yes

In [2]: a.loc[:7]
Out[2]:
   id  seq  class style   drama
0   0    1  alpha     x      no
1   1    1   beta     x      no
2   2    1   gaga     x      no
3   3    2    NaN     y     yes
4   4    2    NaN     y  oh yes
5   5    3    NaN   NaN      no
6   6    3    NaN   NaN     yes
7   7    3    NaN   NaN  oh yes

What is the best approach in this case so we don't fall on a iterrows solution?

pandas join python
2 Answers

Here is an idea using itertools.combinations to create all the possible combinations of different lengths in the columns to use for the merge. Then for each loop, you need to select the rows from a that are nan for the columns to drop and not nan for the one to merge (what you do by hand in the example above). Finally, concat all the dataframes

from itertools import combinations

# define the colms to 
cols = set(['class', 'style', 'drama'])
l = []
for i in range(0, len(cols)+1):
    for comb in combinations(cols, i):
        cols_drop = list(comb)
        cols_merge = list(cols-set(comb))
        # get all the rows with nan for all columns to drop 
        # and notna for the columns to merge
        m = a[cols_merge].notna().all(1) & a[cols_drop].isna().all(1)
#         print (cols_drop, cols_merge) # if you want to understand what is happening
#         print(a[m])                   #  on row selections
        l.append(a[m].drop(cols_drop, axis=1).assign(placeholder=1)
                     .merge(b.assign(placeholder=1), on=cols_merge+['placeholder']))

res = (
    pd.concat(l, ignore_index=True)
      .drop('placeholder', axis=1)
)

and you get

print(res)
    id  seq  class style   drama  start  end
0    0    1  alpha     x      no    840  427
1    0    1  alpha     x      no    530  151
2    1    1   beta     x      no    600   32
3    2    1   gaga     x      no    164  725
4    3    2   gaga     y     yes    838  198
5    3    2   beta     y     yes    562  608
6    3    2  alpha     y     yes    307  154
7    3    2  alpha     y     yes    972    4
8    3    2   gaga     y     yes    147  958
9    4    2   beta     y  oh yes    871  132
10   4    2   beta     y  oh yes    750  726
11   4    2   beta     y  oh yes      3  113
12   4    2   beta     y  oh yes      8  354
13   9    5   beta     y     yes    562  608
14   9    5   beta     x     yes    382  259
...
user_92747320

The solution is based on the evaluation of a query string so we need to change the column name class:

a = a.rename(columns={"class": "klass"})
b = b.rename(columns={"class": "klass"})

ACOLS = ["klass", "style", "drama"]
BCOLS = ["start", "end"]

To match your rules, we need to select all values of a null column: so we replace all nan values by .* then create query string qs for each row of a.

Query the dataframe b to get a list of lists of all rows that match the filter and build a new dataframe before join to the dataframe a:

# "klass.str.contains('^{}$') & style.str.contains('^{}$') & drama.str.contains('^{}$')"
QUERY = " & ".join(f"{c}.str.contains('^{{}}$')" for c in ACOLS)

qs = a[ACOLS].fillna(".*") \
             .apply(lambda c: QUERY.format(*c.tolist()), axis="columns")

data = qs.apply(lambda q: b.query(q, engine="python")[BCOLS].values).explode()
data = pd.DataFrame(data.tolist(), index=data.index, columns=BCOLS)

out = a.join(data).reset_index(drop=True)

>>> out
    id  seq  klass style   drama  start  end
0    0    1  alpha     x      no    840  427
1    0    1  alpha     x      no    530  151
2    1    1   beta     x      no    600   32
3    2    1   gaga     x      no    164  725
4    3    2    NaN     y     yes    838  198
..  ..  ...    ...   ...     ...    ...  ...
65  10    5   beta   NaN  oh yes    871  132
66  10    5   beta   NaN  oh yes    585  149
67  10    5   beta   NaN  oh yes    750  726
68  10    5   beta   NaN  oh yes      3  113
69  10    5   beta   NaN  oh yes      8  354

[70 rows x 7 columns]

Example of a query string:

>>> a.loc[5, ACOLS]
klass    NaN
style    NaN
drama     no
Name: 5, dtype: object

>>> qs.loc[5]
"klass.str.contains('^.*$') and style.str.contains('^.*$') and drama.str.contains('^no$')"
user_152399510
Related