pandas: Why query() instead of bracket operator?

Viewed 1718

What is the difference between using the .query method in a pandas.DataFrame versus the bracket operator? To me, it seems they do the same thing, with just different syntax. Am I missing something? Why would I choose one way over the other?

1 Answers

I'd say both are almost same but query() came into existence to provide few benefits over the bracket operator.
There will always be trade offs depending on the size of the data, required performance and ease of maintaining code.

As per pandas doc:

Ease of Coding and Code Maintainability

You can get the value of the frame where column b has values between the values of columns a and c. For example:

df = DataFrame({
    'a': np.random.choice(10, 5),
    'b': np.random.choice(10, 5),
    'c': np.random.choice(10, 5),
})

# pure python
df[(df['a'] < df['b']) & (df['b'] < df['c'])]

# query
df.query('(a<b) & (b<c)')

Performance of query()

DataFrame.query() using numexpr is slightly faster than Python for large frames.
Note: You will only see the performance benefits of using the numexpr engine with DataFrame.query() if your frame has more than approximately 200,000 rows.

Syntax Comparison

query() and bracket Syntax Comparison

# Full numpy-like syntax:
df = pd.DataFrame(np.random.randint(n, size=(n, 3)), columns=list('abc'))
df.query('(a < b) & (b < c)')
df[(df['a'] < df['b']) & (df['b'] < df['c'])]

# Slightly nicer by removing the parentheses (by binding making comparison operators bind tighter than & and |).
df.query('a < b & b < c')
# Use English instead of symbols:
df.query('a < b and b < c')
# Pretty close to how you might write it on paper:
df.query('a < b < c')

query() Use Cases

A use case for query() is when you have a collection of DataFrame objects that have a subset of column names (or index levels/names) in common. You can pass the same query to both frames without having to specify which frame you’re interested in querying

df1 = DataFrame({
    'a':np.random.choice(10, 5),
    'b':np.random.choice(10, 5),
    'c':np.random.choice(10, 5),
})

df2 = DataFrame({
    'a':np.random.choice(np.arange(5,15,3), 8),
    'b':np.random.choice(np.arange(5,13,2), 8),
    'c':np.random.choice(np.arange(3,11,3), 8),
})
expr = '4<a<c<9'
map(lambda frame:frame.query(expr), [df1, df2])

Few other benefits which can be considered

The in and not in operators

query() also supports special use of Python’s in and not in comparison operators, providing a succinct syntax for calling the isin method of a Series or DataFrame.

# get all rows where columns "a" and "b" have overlapping values
df = pd.DataFrame({'y': list('aabbccddeeff'), 'z': list('aaaabbbbcccc'),
                   'c': np.random.randint(5, size=12),
                   'd': np.random.randint(9, size=12)})
df.query('y in z')
# How you’d do it in pure Python
df[df['y'].isin(df['z'])]

df.query('y not in z')
# pure Python
df[~df['y'].isin(df['z'])]

You can combine this with other expressions for very succinct queries:

# rows where cols a and b have overlapping values and col c's values are less than col d's
df.query('y in z and c > d')
# pure python
df[(df.y.isin(df.z)) & (df.c > df.d)]

Variable reference using '@'

You can refer to variables in the environment by prefixing them with an ‘@’ character like @z

variable_z = list('aaaabbbbcccc')
df.query('y in @variable_z')

Note: Note that in and not in are evaluated in Python, since numexpr has no equivalent of this operation. However,only the in/not in expression itself is evaluated in vanilla Python.

# For example, in the expression
    df.query('a in b + c + d')
    # (b + c + d) is evaluated by numexpr and then the in operation is evaluated in plain Python. In general, any
    # operations that can be evaluated using numexpr will be.

Special use of the == operator with list objects

Comparing a list of values to a column using ==/!= works similarly to in/not in.

df.query('z == ["y", "z", "c"]')
# it is equivalent to df.query('z in ["y", "z", "c"]')

# which one is faster in or ==?
# both are almost same

df.query('z == ["y", "z", "c"]')
# pure python
df[df['z'].isin(['y','z','c'])]

df.query('c == [1, 2]')
df.query('c != [1, 2]')

df.query('[1, 2] in c')
df.query('[1, 2] not in c')

# pure Python
df[df['c'].isin([1, 2])]

Boolean operators

You can negate boolean expressions with the word not or the ~ operator.

df = pd.DataFrame(np.random.rand(n, 3), columns=list('abc'))
df['bools'] = np.random.rand(len(df)) > 0.5
df.query('~bools')
df.query('not bools')
df.query('not bools') == df[~df['bools']]

Of course, expressions can be arbitrarily complex too:

# short query syntax
shorter = df.query('a < b < c and (not bools) or bools > 2')

# equivalent in pure Python
longer = df[(df['a'] < df['b']) & (df['b'] < df['c']) & (~df['bools']) | (df['bools'] > 2)]
Related