diff --git a/src/polars_plugin/mod.rs b/src/polars_plugin/mod.rs index ae459e4..fd2be33 100644 --- a/src/polars_plugin/mod.rs +++ b/src/polars_plugin/mod.rs @@ -3,6 +3,7 @@ mod utils; mod validation_error; use polars::prelude::*; +use polars_arrow::bitmap::Bitmap; use polars_core::POOL; use pyo3_polars::derive::polars_expr; use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; @@ -86,13 +87,12 @@ pub fn all_rules_required( let failures = compute_rule_failures(rule_inputs, kwargs.null_is_valid)?; // If there's any failure, we know that validation failed and use the failure object for an - // informative error message. If no failure exists, we simply return a series with a single - // boolean value to filter by. When filtering on a series with a single value of `true`, polars - // neither actually runs the filter logic, nor does it copy any data. It's essentially a no-op - // that is not optimized away in a lazy frame. + // informative error message. If no failure exists, we return a bitmap with the "unset bits" + // flags set to 0 to trigger a fast-path in the filter executed by the caller. if failures.is_empty() { - let column = Column::new_scalar(PlSmallStr::EMPTY, Scalar::from(true), 1); - return Ok(column.take_materialized_series()); + let bitmap = Bitmap::new_with_value(true, inputs[0].len()); + let ca = BooleanChunked::from_bitmap(PlSmallStr::EMPTY, bitmap); + return Ok(ca.into_series()); } // Aggregate failures into a validation error diff --git a/tests/benches/test_schema.py b/tests/benches/test_schema.py index ab57a78..02df090 100644 --- a/tests/benches/test_schema.py +++ b/tests/benches/test_schema.py @@ -33,8 +33,14 @@ class RowWiseValidationSchema(dy.Schema): @pytest.mark.benchmark(group="schema-row-wise") -def test_row_wise_validate(benchmark: BenchmarkFixture, dataset: pl.DataFrame) -> None: - benchmark(RowWiseValidationSchema.validate, dataset) +@pytest.mark.parametrize("eager", [True, False]) +def test_row_wise_validate( + benchmark: BenchmarkFixture, dataset: pl.DataFrame, eager: bool +) -> None: + def benchmark_fn() -> None: + RowWiseValidationSchema.validate(dataset, eager=eager).lazy().collect() + + benchmark(benchmark_fn) @pytest.mark.benchmark(group="schema-row-wise") diff --git a/tests/schema/test_validate.py b/tests/schema/test_validate.py index a2efbe5..908d4e9 100644 --- a/tests/schema/test_validate.py +++ b/tests/schema/test_validate.py @@ -1,6 +1,8 @@ # Copyright (c) QuantCo 2025-2026 # SPDX-License-Identifier: BSD-3-Clause +from pathlib import Path + import polars as pl import polars.exceptions as plexc import pytest @@ -319,6 +321,25 @@ def test_multiple_unique_columns_both_invalid( assert not MultiUniqueSchema.is_valid(df) +def test_lazy_validation_scan_parquet_length(tmp_path: Path) -> None: + # Arrange + schema = create_schema("test", {"a": dy.Int64(), "b": dy.String()}) + df = schema.sample(10) + df.write_parquet(tmp_path / "test.parquet") + + # Act + height = ( + pl.scan_parquet(tmp_path / "test.parquet") + .pipe(schema.validate, eager=False) + .select(pl.len()) + .collect() + .item() + ) + + # Assert + assert height == 10 + + # ----------------------------------- PERFORMANCE ------------------------------------ #