Polars 파이프라인 테스트

Polars로 데이터 파이프라인 확장 및 최적화하기

Liam Brannigan

Data Scientist & Polars Contributor

파이프라인을 테스트하는 이유

한 셀이 다른 두 DataFrame을 보여주는 다이어그램.

파이프라인을 테스트하는 이유

Polars 변환을 기대 DataFrame과 비교하는 다이어그램.

부서별 완료된 요청

def completed_by_department(requests: pl.LazyFrame) -> pl.LazyFrame:
    return

부서별 완료된 요청

def completed_by_department(requests: pl.LazyFrame) -> pl.LazyFrame:
    return (
        requests




    )

부서별 완료된 요청

def completed_by_department(requests: pl.LazyFrame) -> pl.LazyFrame:
    return (
        requests
        .group_by("DEPARTMENT")
        .len()


    )

부서별 완료된 요청

def completed_by_department(requests: pl.LazyFrame) -> pl.LazyFrame:
    return (
        requests
        .group_by("DEPARTMENT")
        .len()
        .with_columns(pl.col("len").cast(pl.Int32))
        .sort("DEPARTMENT")
    )

테스트 입력

sample = pl.DataFrame({
    "SR_NUMBER": ["SR1", "SR2", "SR3", "SR4"],
    "DEPARTMENT": ["Sanitation", "Water", "Sanitation", "Aviation"],
    "STATUS": ["Completed", "Open", "Completed", "Completed"],
})

shape: (4, 3)
| SR_NUMBER | DEPARTMENT | STATUS    |
| ---       | ---        | ---       |
| str       | str        | str       |
|-----------|------------|-----------|
| SR1       | Sanitation | Completed |
| SR2       | Water      | Open      |
| SR3       | Sanitation | Completed |
| SR4       | Aviation   | Completed |

기대 출력

expected = pl.DataFrame(
    {
        "DEPARTMENT": ["Aviation", "Sanitation"],
        "len": [1, 2],
    }
)

shape: (2, 2)
| DEPARTMENT | len |
| ---        | --- |
| str        | i64 |
|------------|-----|
| Aviation   | 1   |
| Sanitation | 2   |

실제 결과 가져오기

actual = completed_by_department(

실제 결과 가져오기

actual = completed_by_department(
    sample.lazy()
).collect()

실제 결과 가져오기

actual = completed_by_department(
    sample.lazy()
).collect()

shape: (3, 2)
| DEPARTMENT | len |
| ---        | --- |
| str        | i32 |
|------------|-----|
| Aviation   | 1   |
| Sanitation | 2   |
| Water      | 1   |

equals로 비교하기

actual.equals(expected)

False

Polars 테스트 가져오기

from polars.testing import (
    assert_frame_equal,
    assert_schema_equal,
)

스키마 테스트

assert_schema_equal(


)

스키마 테스트

assert_schema_equal(
    actual.schema,
    expected.schema,
)

AssertionError: Schemas are different (dtypes do not match)
[left]: [String, Int32]
[right]: [String, Int64]

기대 스키마 수정

expected = pl.DataFrame(
    {
        "DEPARTMENT": ["Aviation", "Sanitation"], "len": [1, 2],
    },
    schema={"DEPARTMENT": pl.String, "len": pl.Int32},
)

shape: (2, 2)
| DEPARTMENT | len |
| ---        | --- |
| str        | i32 |
|------------|-----|
| Aviation   | 1   |
| Sanitation | 2   |

스키마 테스트

assert_schema_equal(
    actual.schema,
    expected.schema,
)
print("Schema test passed!")

Schema test passed!

프레임 테스트

assert_frame_equal(


)

프레임 테스트

assert_frame_equal(
    actual,
    expected,
)

AssertionError: DataFrames are different height (row count) mismatch
[left]: 3
[right]: 2

실제 결과와 기대 결과 비교

shape: (3, 2)
| DEPARTMENT | len |
| ---        | --- |
| str        | i32 |
|------------|-----|
| Aviation   | 1   |
| Sanitation | 2   |
| Water      | 1   |

shape: (2, 2)
| DEPARTMENT | len |
| ---        | --- |
| str        | i32 |
|------------|-----|
| Aviation   | 1   |
| Sanitation | 2   |

쿼리 수정

def completed_by_department(requests: pl.LazyFrame) -> pl.LazyFrame:
    return (
        requests
        .filter(pl.col("STATUS") == "Completed")
        .group_by("DEPARTMENT")
        .len()
        .with_columns(pl.col("len").cast(pl.Int32))
        .sort("DEPARTMENT")
    )

actual = completed_by_department(
    sample.lazy()
).collect()

최종 검증

assert_schema_equal(actual.schema, expected.schema)
assert_frame_equal(actual, expected)
print("All tests passed!")

All tests passed!

연습해 봅시다!

Polars로 데이터 파이프라인 확장 및 최적화하기