count number of rows given column names and values - pandas

Question

I am currently switching for R to Python so please be patient with me. Is the following a good way to count the number of rows given column names and values?

import pandas as pd
df = pd.DataFrame([["1", "2"], ["2", "4"], ["1", "4"]], columns=['A', 'B'])
cn1 = "A"
cn2 = "B"
cv1 = "1"
cv2 = "2"
no_rows = len(df[(df[cn1]==cv1) & (df[cn2]==cv2)].index)
print(no_rows)

Answer 1

First, it's a bad idea to input your numerics as strings in your dataframe. Use plain ints instead.

Your code currently forms a predicate, performs a slice on the frame and then finds the size of the frame. This is more work than necessary - the predicate itself is a series of booleans, and running a .sum() on it produces the number of matching values.

That, plus your current code is not general-purpose. A general-purpose implementation could look like

from typing import Dict, Any
import pandas as pd
def match_count(df: pd.DataFrame, **criteria: Any) -> int:
 pairs = iter(criteria.items())
 column, value = next(pairs)
 predicate = df[column] == value
 for column, value in pairs:
 predicate &= df[column] == value
 return predicate.sum()
def test() -> None:
 df = pd.DataFrame(
 [[1, 2],
 [2, 4],
 [1, 4]],
 columns=['A', 'B'],
 )
 print(match_count(df, A=1, B=2))
if __name__ == '__main__':
 test()

Answer 2

I usually use shape[0] because it's more readable, so in your case it would be:

no_rows = df[(df[cn1]==cv1) & (df[cn2]==cv2)].shape[0]

Answer 3

While this specific example can be completely refactored into Reinderien's top-notch functions, we don't always need something so elaborate (e.g., quick exploratory analysis).

Masking and counting come up very often in one form or another, so I think it's still worth reviewing how to do them idiomatically in pandas.

---

Revised code

Maintaining the spirit of the original code, I would use something like:

matches = df[cn1].eq(cv1) & df[cn2].eq(cv2)
len(df[matches]) # but remember that matches.sum() is faster

---

Comments on the original code

len(df[(df[cn1] == cv1) & (df[cn2] == cv2)].index)
^ ^ ^ ^
3 2 4 1

1. No need to use .index explicitly since DataFrame.__len__ does it automatically:

   class DataFrame(NDFrame, OpsMixin):
    ...
    def __len__(self) -> int:
    return len(self.index)
   

2. DataFrame.eq can sometimes be useful over ==:

   • supports axis / level broadcasting

   • arguably more readable when joining multiple tests

     df[cn1].eq(cv1) & df[cn2].eq(cv2) # (df[cn1] == cv1) & (df[cn2] == cv2)
     

   • arguably more readable when chaining methods (e.g., when comparing shifted columns)

     df[cn1].shift().eq(cv1).cumsum() # (df[cn1].shift() == cv1).cumsum()
     

3. If speed is important, len(df) is faster than df.shape[0] (h/t @root):

4. If you have a lot of conditions to join (e.g., generated via comprehension), consider np.logical_and.reduce:

   df[np.logical_and.reduce([
    df[cn1] == cv1,
    # ...
    df[cn2] == cv2,
   ])]