How to filter data using pandas.DataFrame.query in Python

When you are dealing with data in Python, especially with libraries like pandas, understanding how to query DataFrames is essential. The query method is a powerful tool that lets you filter your data using a syntax that is almost like SQL. This can simplify your code and make it much more readable.

First, let’s take a look at how to create a simple DataFrame. You can use the DataFrame constructor from pandas to initialize your dataset.

import pandas as pd

data = {
    'name': ['Alice', 'Bob', 'Charlie', 'David'],
    'age': [24, 27, 22, 32],
    'city': ['New York', 'Los Angeles', 'Chicago', 'New York']
}

df = pd.DataFrame(data)

Now that we have our DataFrame df, we can start querying it. The basic structure of a query is simpler. You can filter rows based on column values using boolean expressions. For example, if you want to find all entries where the age is greater than 25, you can use the following query:

result = df.query('age > 25')
print(result)

This will return a DataFrame containing only the rows where the age is greater than 25. The syntax is clean and easy to understand, which is one of the biggest benefits of using the query method.

You can also query based on string values. If you want to find all people living in New York, you can do it like this:

result = df.query('city == "New York"')
print(result)

It’s also possible to combine multiple conditions. For instance, if you want to find all people older than 24 and living in New York, you can use the and operator:

result = df.query('age > 24 and city == "New York"')
print(result)

Using parentheses can help clarify complex queries, especially when you mix different conditions. For example, if you want to find people who either live in New York or are younger than 25:

result = df.query('(city == "New York") or (age < 25)')
print(result)

This flexibility allows you to perform complex filtering efficiently within a single line of code. To further enhance your queries, you can use functions like startswith or in for string matching.

result = df.query('name.str.startswith("A")')
print(result)

While most of the time, query will take care of your needs, it is good to remember that it operates on the assumption that your column names are valid Python variable names. If your column names contain spaces or special characters, you'll need to use backticks:

df = pd.DataFrame({'first name': ['Alice', 'Bob'], 'age': [24, 27]})
result = df.query('first name == "Alice"')
print(result)

Understanding these basics will set a solid foundation for you to master more complex filters later on. As you get comfortable with the syntax, you'll find that querying DataFrames becomes an intuitive part of your data manipulation toolkit, so that you can focus more on deriving insights rather than getting bogged down in syntax.

Mastering complex filters with pandas DataFrame query

When your filtering needs grow beyond simple comparisons, query can handle more sophisticated expressions involving multiple columns, functions, and even variables defined outside the DataFrame.

Suppose you want to filter rows where the age is between 23 and 30, and the city is either "New York" or "Chicago". You can combine and, or, and parentheses to express this clearly:

result = df.query('(age >= 23 and age <= 30) and (city == "New York" or city == "Chicago")')
print(result)

This approach keeps your code readable and concise, avoiding the clutter of chained boolean masks.

Sometimes, you want to use a list of values for filtering, similar to SQL's IN clause. While query doesn't support in directly in the expression, you can pass external variables to the query environment using the @ symbol.

cities = ["New York", "Chicago"]
result = df.query('city in @cities')
print(result)

This syntax lets you leverage Python variables inside your query strings, making your code modular and easier to maintain.

Complex string operations are also possible. For example, to find rows where the name contains the letter 'a' (case insensitive), you can use the pandas string methods inside query like this:

result = df.query('name.str.contains("a", case=False)')
print(result)

Keep in mind that these string methods are vectorized and efficient, so they scale well even with large datasets.

When working with numerical ranges and multiple filters, you might want to combine query with other pandas methods for even more power. For instance, chaining query with sort_values to get the top results after filtering:

result = df.query('age > 22').sort_values('age', ascending=False)
print(result)

This pattern is common in data analysis workflows: filter first, then sort or aggregate.

Finally, if you need to reference columns with names that clash with Python keywords or contain special characters, always use backticks around the column name. For example, if your DataFrame has a column named class:

df = pd.DataFrame({
    'name': ['Alice', 'Bob', 'Charlie'],
    'class': ['A', 'B', 'A'],
    'score': [85, 92, 78]
})

result = df.query('class == "A" and score > 80')
print(result)

This ensures pandas interprets the column name correctly without syntax errors.