Search code examples
pandaspandas-groupbyhaving

Filtering/Querying Pandas DataFrame after multiple grouping/agg


I have a dataframe that I first group, Counting QuoteLine Items grouped by stock(1-true, 0-false) and mfg type (K-Kit, M-manufactured, P-Purchased). Ultimately, I am interested in quotes that ALL items are either NonStock/Kit and/or Stock/['M','P'] :

grouped = df.groupby(['QuoteNum', 'typecode', 'stock']).agg({"QuoteLine": "count"})

and I get this:

                                 QuoteLine-count
QuoteNum    typecode    stock   
10001          K          0         1
10003          M          0         1
10005          M          0         3
                          1         1
10006          M          1         1
...           ...        ...       ...
26961          P          1         1
26962          P          1         1
26963          P          1         2
26964          K          0         1   
               M          1         2

If I unstack it twice:

grouped = df.groupby(['QuoteNum', 'typecode', 'stock']).agg({"QuoteLine": "count"}).unstack().unstack()

# I get
    QuoteLine-count
stock           0                       1
typecode    K       M       P       K       M       P
QuoteNum                        
10001       1.0     NaN     NaN     NaN     NaN     NaN
10003       NaN     1.0     NaN     NaN     NaN     NaN
10005       NaN     3.0     NaN     NaN     1.0     NaN
10006       NaN     NaN     NaN     NaN     1.0     NaN
10007       2.0     NaN     NaN     NaN     NaN     NaN
...         ...     ...     ...     ...     ...     ...
26959       NaN     NaN     NaN     NaN     NaN     1.0
26961       NaN     1.0     NaN     NaN     NaN     1.0
26962       NaN     NaN     NaN     NaN     NaN     1.0
26963       NaN     NaN     NaN     NaN     NaN     2.0
26964      1.0      NaN     NaN     NaN     2.0     NaN

Now I need to filter out all records where, this is where I need help

    # pseudo-code
    (stock == 0 and typecode in ['M','P']) -> values are NOT NaN (don't want those)
    and 
    (stock == 1 and typecode='K') -> values are NOT NaN (don't want those either)

so I'm left with these records:
Basically: Columns "0/M, 0/P, 1/K" must be all NaNs and other columns have at least one non NaN value
    QuoteLine-count
stock           0                       1
typecode    K       M       P       K       M       P
QuoteNum                        
10001       1.0     NaN     NaN     NaN     NaN     NaN
10006       NaN     NaN     NaN     NaN     1.0     NaN
10007       2.0     NaN     NaN     NaN     NaN     NaN
...         ...     ...     ...     ...     ...     ...
26959       NaN     NaN     NaN     NaN     NaN     1.0
26962       NaN     NaN     NaN     NaN     NaN     1.0
26963       NaN     NaN     NaN     NaN     NaN     2.0
26964      1.0      NaN     NaN     NaN     2.0     NaN

Solution

  • IIUC, use boolean mask to set rows that match your conditions to NaN then unstack desired levels:

    # Shortcut (for readability)
    lvl_vals = grouped.index.get_level_values
    
    m1 = (lvl_vals('typecode') == 'K') & (lvl_vals('stock') == 0)
    m2 = (lvl_vals('typecode').isin(['M', 'P'])) & (lvl_vals('stock') == 1)
    grouped[m1|m2] = np.nan
    out = grouped.unstack(level=['stock', 'typecode']) \
                 .loc[lambda x: x.isna().all(axis=1)]
    

    Output result:

    >>> out
             QuoteLine-count            
    stock                  0       1    
    typecode               K   M   M   P
    QuoteNum                            
    10001                NaN NaN NaN NaN
    10006                NaN NaN NaN NaN
    26961                NaN NaN NaN NaN
    26962                NaN NaN NaN NaN
    26963                NaN NaN NaN NaN
    26964                NaN NaN NaN NaN