Merge pull request #34 from Point72/feat/probe-source

feat: add `.piot.probe()` pass-through telemetry source

Author: ptomecek

polars_io_tools/io_sources/__init__.py

@@ -18,6 +18,7 @@
 from .lazy_debug import debug, debug as _lazy_debug
 from .lazy_iter_rows import *
 from .lazy_narwhals_reader import *
+from .lazy_probe import probe, probe as _lazy_probe
 from .lazy_sql_reader import *
 from .pushdown_combine import *
 from .pushdown_pivot import *
@@ -69,6 +70,10 @@ def __init__(self, lf: pl.LazyFrame) -> None:
     def debug(self, *args, **kwargs) -> pl.LazyFrame:
         return _lazy_debug(self._lf, *args, **kwargs)
 
+    @functools.wraps(_lazy_probe)
+    def probe(self, *args, **kwargs) -> pl.LazyFrame:
+        return _lazy_probe(self._lf, *args, **kwargs)
+
     @functools.wraps(_lazy_cache)
     def cache(self, *args, **kwargs) -> pl.LazyFrame:
         if self._DISABLE_OPTIMIZATIONS:

polars_io_tools/io_sources/lazy_debug.py

@@ -15,6 +15,8 @@ def debug(
     self: pl.LazyFrame,
     log_level: int | None = None,
     description: str | None = None,
+    *,
+    is_pure: bool = False,
 ) -> pl.LazyFrame:
     """
     A very simple pass-through lazy frame source to help with debugging experimentation of polars io sources and lazy frame behavior.
@@ -23,6 +25,7 @@ def debug(
         self: The input data frame to cache columns of.
         log_level: If provided, will log at the given level. If None, will print. Defaults to None.
         description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``).
+        is_pure: Whether the wrapped input is pure (deterministic and side-effect-free). Defaults to ``False`` so repeated uses are not deduplicated -- ``debug`` logs on every execution, so collapsing executions would drop output you asked for. Set ``True`` only for a known-pure input where deduplication is wanted.
     """
     schema = self.collect_schema()
 
@@ -55,4 +58,4 @@ def source_generator(
             raise RuntimeError(err_msg) from e
 
     # TODO: Turn on validate_schema when this is solved: https://github.com/pola-rs/polars/issues/22110
-    return register_io_source_with_is_pure(source_generator, schema=schema, validate_schema=False, explain_detail=description)
+    return register_io_source_with_is_pure(source_generator, schema=schema, is_pure=is_pure, validate_schema=False, explain_detail=description)

polars_io_tools/io_sources/lazy_probe.py

@@ -0,0 +1,61 @@
+from collections.abc import Iterator
+
+import polars as pl
+
+from .util import collect_lf_in_io_source, register_io_source_with_is_pure
+
+__all__ = ("probe",)
+
+
+def probe(self: pl.LazyFrame, description: str | None = None, *, is_pure: bool = False) -> pl.LazyFrame:
+    """A result-preserving pass-through source that emits one OpenTelemetry span per execution.
+
+    Inserting ``.piot.probe()`` at a point in a lazy pipeline turns that point into a measured
+    ``io_source.execute[probe]`` span: ``next_elapsed_total_ms`` is the time spent pulling the sub-plan
+    below the probe and ``total_rows`` is the number of rows that flow through it. Predicate and
+    projection pushdown are forwarded to the input, so the probe is not an optimization barrier for
+    them. It does add a ``PythonScan`` node that re-enters Polars via ``collect_lf_in_io_source``, so it
+    is a measurement point rather than a zero-cost tap.
+
+    Unlike ``debug``, it performs no logging and does not compute ``explain()``, so it is cheap enough to
+    leave in a pipeline. The span is a no-op unless an OpenTelemetry SDK is configured (see
+    ``polars_io_tools.io_sources.profiling``).
+
+    Overhead is on the order of milliseconds and grows with the number of rows passing through the probe
+    (each batch makes a Rust-Python round trip); it is negligible relative to any non-trivial read but can
+    dominate a small, fully-cached one. The overhead is the io-source boundary itself, not the telemetry.
+    Under a single-worker Polars pool the input is materialized in full before the first batch (see
+    ``collect_lf_in_io_source``), so memory and latency can be higher than the streaming path.
+    ``next_elapsed_total_ms`` is caller-observed pull latency, so it can undercount a reader that decodes
+    on background threads.
+
+    Args:
+        self: The input LazyFrame to pass through unchanged.
+        description: Optional free-form description of this probe instance, attached to its OpenTelemetry span (``explain_detail``) -- e.g. the pipeline stage being measured.
+        is_pure: Whether the wrapped input is pure (deterministic and side-effect-free). Defaults to ``False``: a pass-through cannot infer its input's purity, and ``False`` preserves results for an impure input and measures every occurrence. Set ``True`` when the input is known pure to let Polars deduplicate repeated uses of the same probe (one execution, one span) -- i.e. to propagate a pure input's dedup-ability through the probe.
+
+    Returns:
+        pl.LazyFrame: A LazyFrame equivalent to ``self`` whose execution emits one telemetry span.
+    """
+
+    def source_generator(
+        with_columns: list[str] | None,
+        predicate: pl.Expr | None,
+        n_rows: int | None,
+        batch_size: int | None,
+    ) -> Iterator[pl.DataFrame]:
+        df = self
+        if predicate is not None:
+            df = df.filter(predicate)
+        if with_columns is not None:
+            df = df.select(with_columns)
+        if n_rows is not None:
+            df = df.head(n_rows)
+        yield from collect_lf_in_io_source(df, batch_size)
+
+    # The probe's purity is its input's purity, which it cannot infer, so ``is_pure`` defaults to ``False``:
+    # claiming purity for an impure input would let Polars deduplicate repeated uses and change results. The
+    # schema is passed as a callable so it is resolved lazily rather than forced at construction.
+    return register_io_source_with_is_pure(
+        source_generator, schema=self.collect_schema, is_pure=is_pure, validate_schema=False, explain_detail=description
+    )

polars_io_tools/tests/io_sources/test_lazy_debug.py

@@ -43,3 +43,12 @@ def test_debug_filter(caplog):
     assert "debug called with" in caplog.text
     assert "`with_columns=None`" in caplog.text
     assert '`predicate=[(col("a")) > (1)]`' in caplog.text
+
+
+def test_debug_logs_per_use_by_default(caplog):
+    # is_pure defaults to False, so reusing the same debug is not deduplicated: its log-side-effect runs per occurrence.
+    df = pl.DataFrame({"a": [1]}).lazy()
+    debug = df.piot.debug(log_level=logging.INFO)
+    caplog.set_level(logging.INFO)
+    pl.concat([debug, debug]).collect(engine="streaming")
+    assert caplog.text.count("debug called with") == 2

polars_io_tools/tests/io_sources/test_lazy_probe.py

@@ -0,0 +1,97 @@
+import polars as pl
+import pytest
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from polars.testing import assert_frame_equal
+
+import polars_io_tools  # noqa: F401  -- registers the `.piot` namespace
+from polars_io_tools.io_sources.util import register_io_source_with_is_pure
+
+_EXPORTER = InMemorySpanExporter()
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _otel_provider():
+    provider = trace.get_tracer_provider()
+    if not hasattr(provider, "add_span_processor"):
+        trace.set_tracer_provider(TracerProvider())
+        provider = trace.get_tracer_provider()
+    if hasattr(provider, "add_span_processor"):
+        provider.add_span_processor(SimpleSpanProcessor(_EXPORTER))
+    yield
+
+
+@pytest.fixture(autouse=True)
+def _reset():
+    _EXPORTER.clear()
+    yield
+
+
+def _probe_spans():
+    return [s for s in _EXPORTER.get_finished_spans() if s.name.startswith("io_source.execute[probe]")]
+
+
+def test_probe_is_result_preserving_passthrough():
+    df = pl.DataFrame({"id": [1, 2, 3], "v": [10, 20, 30]})
+    out = df.lazy().piot.probe().collect(engine="streaming")
+    assert_frame_equal(out, df)
+
+
+def test_probe_forwards_n_rows_pushdown():
+    out = pl.LazyFrame({"id": list(range(100))}).piot.probe().head(5).collect(engine="streaming")
+    assert out["id"].to_list() == [0, 1, 2, 3, 4]
+    assert _probe_spans()[-1].attributes["polars_io_tools.total_rows"] == 5
+
+
+def test_probe_does_not_dedupe_impure_input():
+    # A pass-through cannot claim purity: reusing the same probe over an impure input must not collapse executions.
+    state = {"n": 0}
+
+    def impure(with_columns, predicate, n_rows, batch_size):
+        state["n"] += 1
+        yield pl.DataFrame({"run": [state["n"]]})
+
+    probe = register_io_source_with_is_pure(impure, schema=pl.Schema({"run": pl.Int64}), is_pure=False).piot.probe()
+    out = pl.concat([probe, probe]).collect(engine="streaming")
+    assert sorted(out["run"].to_list()) == [1, 2]
+
+
+def test_probe_emits_one_span_with_description():
+    df = pl.DataFrame({"id": [1, 2, 3]})
+    df.lazy().piot.probe(description="stage1").collect(engine="streaming")
+
+    spans = _probe_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.attributes["polars_io_tools.explain_name"] == "probe"
+    assert span.attributes["polars_io_tools.explain_detail"] == "stage1"
+    assert span.attributes["polars_io_tools.total_rows"] == 3
+
+
+def test_probe_forwards_predicate_pushdown():
+    # A pushed predicate reaches the probe, so only matching rows flow through it (and are counted).
+    df = pl.DataFrame({"id": [1, 2, 3, 4], "grp": [0, 1, 0, 1]})
+    out = df.lazy().piot.probe().filter(pl.col("grp") == 1).collect(engine="streaming")
+
+    assert out.sort("id")["id"].to_list() == [2, 4]
+    assert _probe_spans()[-1].attributes["polars_io_tools.total_rows"] == 2
+
+
+def test_probe_forwards_is_pure(monkeypatch):
+    # is_pure defaults to False (safe for an input of unknown purity) and can be opted into for a known-pure input.
+    captured = {}
+
+    def fake_register(io_source, *, schema, **kwargs):
+        captured.update(kwargs)
+        return object()
+
+    monkeypatch.setattr("polars.io.plugins.register_io_source", fake_register)
+
+    pl.LazyFrame({"a": [1]}).piot.probe()
+    assert captured["is_pure"] is False
+
+    captured.clear()
+    pl.LazyFrame({"a": [1]}).piot.probe(is_pure=True)
+    assert captured["is_pure"] is True