Pruebas de pipelines en Polars
Escala y optimiza canalizaciones de datos con Polars
Liam Brannigan
Data Scientist & Polars Contributor
¿Por qué probar pipelines?
¿Por qué probar pipelines?
Solicitudes completadas por departamento
def completed_by_department(requests: pl.LazyFrame) -> pl.LazyFrame:
return
Solicitudes completadas por departamento
def completed_by_department(requests: pl.LazyFrame) -> pl.LazyFrame:
return (
requests
)
Solicitudes completadas por departamento
def completed_by_department(requests: pl.LazyFrame) -> pl.LazyFrame:
return (
requests
.group_by("DEPARTMENT")
.len()
)
Solicitudes completadas por departamento
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")
)
Entrada de prueba
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 |
Salida esperada
expected = pl.DataFrame(
{
"DEPARTMENT": ["Aviation", "Sanitation"],
"len": [1, 2],
}
)
shape: (2, 2)
| DEPARTMENT | len |
| --- | --- |
| str | i64 |
|------------|-----|
| Aviation | 1 |
| Sanitation | 2 |
Obtener el resultado real
actual = completed_by_department(
Obtener el resultado real
actual = completed_by_department(
sample.lazy()
).collect()
Obtener el resultado real
actual = completed_by_department(
sample.lazy()
).collect()
shape: (3, 2)
| DEPARTMENT | len |
| --- | --- |
| str | i32 |
|------------|-----|
| Aviation | 1 |
| Sanitation | 2 |
| Water | 1 |
Comparar con equals
actual.equals(expected)
False
Importar pruebas de Polars
from polars.testing import (
assert_frame_equal,
assert_schema_equal,
)
Probar el esquema
assert_schema_equal(
)
Probar el esquema
assert_schema_equal(
actual.schema,
expected.schema,
)
AssertionError: Schemas are different (dtypes do not match)
[left]: [String, Int32]
[right]: [String, Int64]
Corregir el esquema esperado
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 |
Probar el esquema
assert_schema_equal(
actual.schema,
expected.schema,
)
print("Schema test passed!")
Schema test passed!
Probar el DataFrame
assert_frame_equal(
)
Probar el DataFrame
assert_frame_equal(
actual,
expected,
)
AssertionError: DataFrames are different height (row count) mismatch
[left]: 3
[right]: 2
Comparar real y esperado
shape: (3, 2)
| DEPARTMENT | len |
| --- | --- |
| str | i32 |
|------------|-----|
| Aviation | 1 |
| Sanitation | 2 |
| Water | 1 |
shape: (2, 2)
| DEPARTMENT | len |
| --- | --- |
| str | i32 |
|------------|-----|
| Aviation | 1 |
| Sanitation | 2 |
Corregir la consulta
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()
Aserciones finales
assert_schema_equal(actual.schema, expected.schema)
assert_frame_equal(actual, expected)
print("All tests passed!")
All tests passed!
¡Vamos a practicar!
Escala y optimiza canalizaciones de datos con Polars
