Python Pandas query() 함수로 열 값 기준 행 필터링하는 방법

Pandas에서 열(Column) 값을 기준으로 DataFrame의 행을 필터링하려면 query() 함수를 사용할 수 있습니다. 이 함수 안에 원하는 조건식을 문자열 형태로 전달하면, 해당 조건을 만족하는 행만 손쉽게 추출할 수 있습니다.

1. 라이브러리 임포트

먼저 필요한 라이브러리를 임포트합니다.

import pandas as pd

2. 예제 데이터 준비

팀 순위 데이터를 리스트 형태로 준비합니다.

Team = [['India', 1, 100], ['Australia', 2, 85], ['England', 3, 75],
        ['New Zealand', 4, 65], ['South Africa', 5, 50], ['Bangladesh', 6, 40]]

이 데이터를 바탕으로 DataFrame을 생성하고, 각 열의 이름을 지정합니다.

dataFrame = pd.DataFrame(Team, columns=['Country', 'Rank', 'Points'])

3. query() 함수로 조건 필터링

query() 함수를 사용해 "Rank"가 5인 팀만 조회해 보겠습니다.

dataFrame.query("Rank == 5")

예제 1: 단일 조건으로 필터링

다음은 전체 코드입니다.

import pandas as pd

# 팀 순위 데이터 (리스트 형태)
Team = [['India', 1, 100], ['Australia', 2, 85], ['England', 3, 75],
        ['New Zealand', 4, 65], ['South Africa', 5, 50], ['Bangladesh', 6, 40]]

# DataFrame 생성 및 열 이름 지정
dataFrame = pd.DataFrame(Team, columns=['Country', 'Rank', 'Points'])

print("DataFrame...\n", dataFrame)

# query()를 사용해 Rank가 5인 행 필터링
print("\nRank가 5인 팀 조회..\n", dataFrame.query("Rank == 5"))

출력 결과

위 코드를 실행하면 다음과 같은 결과가 출력됩니다.

DataFrame...
       Country  Rank  Points
0        India     1     100
1    Australia     2      85
2      England     3      75
3  New Zealand     4      65
4 South Africa     5      50
5   Bangladesh     6      40

Rank가 5인 팀 조회..
       Country  Rank  Points
4 South Africa     5      50

예제 2: 복수 조건(AND 연산)으로 필터링

query() 함수는 여러 조건을 조합할 수도 있습니다. 아래 예제는 Points가 70보다 크고 Rank가 3 미만인 행을 필터링합니다.

import pandas as pd

# 팀 순위 데이터 (리스트 형태)
Team = [['India', 1, 100], ['Australia', 2, 85], ['England', 3, 75],
        ['New Zealand', 4, 65], ['South Africa', 5, 50], ['Bangladesh', 6, 40]]

# DataFrame 생성 및 열 이름 지정
dataFrame = pd.DataFrame(Team, columns=['Country', 'Rank', 'Points'])

print("DataFrame...\n", dataFrame)

# Points > 70 이고 Rank < 3 인 행 필터링
print("\nPoints가 70 초과이고 Rank가 3 미만인 팀 조회..\n")
print(dataFrame.query("Points > 70 and Rank < 3"))

출력 결과

실행 결과는 다음과 같습니다.

DataFrame...
       Country  Rank  Points
0        India     1     100
1    Australia     2      85
2      England     3      75
3  New Zealand     4      65
4 South Africa     5      50
5   Bangladesh     6      40

Points가 70 초과이고 Rank가 3 미만인 팀 조회..

    Country  Rank  Points
0     India     1     100
1 Australia     2      85

정리

query() 함수를 활용하면 SQL 스타일의 직관적인 조건식 문자열로 DataFrame의 행을 필터링할 수 있습니다. 주요 특징은 다음과 같습니다.

- 비교 연산자(==, >, <, >=, <=, !=)를 자유롭게 사용할 수 있습니다.
- and, or, not 등의 논리 연산자로 복수 조건을 조합할 수 있습니다.
- 조건식을 문자열로 작성하므로 코드가 간결하고 가독성이 뛰어납니다.
- 대규모 데이터에서는 불리언 마스크 방식(df[df['Rank'] == 5])보다 성능상 이점을 가질 수도 있습니다.