Spark 변환의 행과 스키마 검사

이 테스트는 pytest-transforms의 실제 spark_session에서 CSV를 읽어 필터를 실행합니다. 출력에 쓸 DataFrame만 mock 핸들로 받아 확인하므로 실제 Foundry 데이터셋에는 기록하지 않습니다.

변환 모듈

src/myproject/datasets/find_aircraft_spark.py:

from pyspark.sql import functions as F
from transforms.api import Input, Output, transform


def compute_data(output, aircraft):
    rows = aircraft.dataframe()
    selected = rows.filter(
        (F.col("number_of_seats") > 300)
        & (F.col("operating_status") == "Yes")
    )
    output.write_dataframe(selected)


compute = transform.spark.using(
    output=Output("/Example Project/Outputs/aircraft_spark"),
    aircraft=Input("/Example Project/Inputs/aircraft"),
)(compute_data)

가상 자산 경로를 실제 값으로 바꾸고 해당 Spark API의 최소 버전 및 PySpark 테스트 환경을 준비합니다.

CSV 결과 비교

두 CSV를 src/test/resources/에 두고 src/test/test_find_aircraft_spark.py를 작성합니다.

from collections import Counter
from pathlib import Path
from unittest.mock import MagicMock

from myproject.datasets.find_aircraft_spark import compute_data

RESOURCES = Path(__file__).parent / "resources"


def test_find_aircraft(spark_session):
    source_frame = spark_session.read.csv(
        str(RESOURCES / "aircraft_mock.csv"), header=True, inferSchema=True
    )
    expected = spark_session.read.csv(
        str(RESOURCES / "expected_filtered_aircraft.csv"),
        header=True,
        inferSchema=True,
    )
    aircraft = MagicMock(spec_set=["dataframe"])
    aircraft.dataframe.return_value = source_frame
    output = MagicMock(spec_set=["write_dataframe"])

    compute_data(output=output, aircraft=aircraft)

    output.write_dataframe.assert_called_once()
    actual = output.write_dataframe.call_args.args[0]
    assert actual.columns == expected.columns
    assert actual.schema == expected.schema
    actual_rows = actual.collect()
    expected_rows = expected.collect()
    assert len(actual_rows) == len(expected_rows)
    assert Counter(tuple(row) for row in actual_rows) == Counter(
        tuple(row) for row in expected_rows
    )

예제 CSV는 작고 모든 열이 해시 가능한 스칼라이므로 결과를 한 번씩 collect하여 순서와 무관한 다중 집합으로 비교합니다. 중복 행의 개수도 유지합니다. 큰 데이터 전체를 테스트 목적으로 collect하는 패턴으로 확대하지 않습니다.

분산 DataFrame 상태로 비교해야 한다면 같은 열 순서와 호환 schema에서 양방향 exceptAll이 모두 비어 있는지 확인하는 방법도 사용할 수 있습니다. 한쪽 차집합만 확인하면 반대쪽의 추가 행을 놓칠 수 있습니다.

spark_session은 테스트 session에서 공유합니다. Time zone이나 Spark 설정을 바꾸는 테스트는 종료 전에 원래 값을 복원해야 하며, 다른 테스트의 실행 순서에 의존하지 않습니다.