Testare le pipeline Polars

Scalare e ottimizzare le pipeline di dati con Polars

Liam Brannigan

Data Scientist & Polars Contributor

Perché testare le pipeline?

Diagramma che mostra due DataFrame diversi in una cella.

Perché testare le pipeline?

Diagramma che mostra una trasformazione Polars testata contro un DataFrame atteso.

Richieste completate per reparto

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

Richieste completate per reparto

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




    )

Richieste completate per reparto

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


    )

Richieste completate per reparto

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")
    )

Input di test

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 |

Output atteso

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

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

Ottenere il risultato effettivo

actual = completed_by_department(

Ottenere il risultato effettivo

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

Ottenere il risultato effettivo

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

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

Confrontare con equals

actual.equals(expected)

False

Importare i test di Polars

from polars.testing import (
    assert_frame_equal,
    assert_schema_equal,
)

Testare lo schema

assert_schema_equal(


)

Testare lo schema

assert_schema_equal(
    actual.schema,
    expected.schema,
)

AssertionError: Schemi diversi (i tipi non corrispondono)
[left]: [String, Int32]
[right]: [String, Int64]

Correggere lo schema atteso

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   |

Testare lo schema

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

Schema test passed!

Testare il DataFrame

assert_frame_equal(


)

Testare il DataFrame

assert_frame_equal(
    actual,
    expected,
)

AssertionError: Altezza dei DataFrame diversa (conteggio righe non corrisponde)
[left]: 3
[right]: 2

Confronto tra effettivo e atteso

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

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

Correggere la query

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()

Asserzioni finali

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

All tests passed!

Ayo berlatih!

Scalare e ottimizzare le pipeline di dati con Polars