feat: OpenTelemetry instrumentation for custom IO sources

[ptomecek]

What

OpenTelemetry instrumentation for custom Polars IO sources. A streaming engine's per-node metrics measure engine-side poll/CPU time and do not capture the time a Python IO source spends blocked fetching from a backend. This recovers that, per source execution.

Each source execution registered via register_io_source_with_is_pure emits one io_source.execute[<name>] span recording:

• polars_io_tools.next_elapsed_total_ms — the summed wall duration of every next() pull on the source iterator (caller-observed fetch latency; time the consumer spends processing between pulls is excluded);
• polars_io_tools.total_rows, polars_io_tools.n_columns, polars_io_tools.batch_count, and a polars_io_tools.outcome of exhausted / closed / error.

Attributes are namespaced under polars_io_tools.; the tracer is scoped to polars_io_tools with the package version.

Naming. The span's identity is the source's explain_name — a low-cardinality label that defaults to the source function's name (scan_db, scan_clickhouse, …; a source_generator closure inside scan_db resolves to scan_db). The free-form explain_detail is a per-instance description that distinguishes two instances of the same source (e.g. two SQL scans). Every public source constructor — scan_db, scan_clickhouse, scan_delta, scan_datadog, scan_narwhals, scan_synthetic_regression/scan_synthetic_panel, cache, cache_parquet, debug, concat_named, filtered_join/filtered_join_asof, execute_on_ray — takes a description argument for this. When the installed Polars build accepts explain_name/explain_detail, the same values are forwarded to the scan, so the scan's registered identity and description match the span's — feature-detected against register_io_source's signature, dormant until the corresponding upstream Polars change lands, activating automatically once it ships.

Always on, no-op without an SDK

There is no library-specific on/off switch. Following the OpenTelemetry convention for native library instrumentation, spans are created through the OpenTelemetry API only (no logfire/vendor dependency), which is a no-op unless the application has configured an SDK — attributes are set only on recording spans. This matches how e.g. elasticsearch-py ships native OTel support: on when you configure OTel, invisible otherwise.

OpenTelemetry is an optional, undeclared dependency: the code feature-detects the opentelemetry API and no-ops when it is absent. Any application that configures an OpenTelemetry SDK already provides the API transitively — which is the only case where spans do anything — so nothing new is imposed on users who don't use OTel. opentelemetry-sdk is added to the develop extra for the test suite.

Opt-out. Instrumentation is on by default; disable it with OTEL_PYTHON_INSTRUMENTATION_POLARS_IO_TOOLS_ENABLED=false — the OTEL_PYTHON_INSTRUMENTATION_<LIBRARY>_ENABLED convention used by e.g. elasticsearch-py (default true). This lets an application keep OpenTelemetry for the rest of its code while suppressing these spans. When disabled the wrapper is not installed at all, so there is zero overhead (not merely no-op emission).

How users use it

Configure any OTel SDK, then use the sources as normal. Each source is auto-named after its constructor; pass description= to tell instances apart:

from polars_io_tools.io_sources.lazy_sql_reader import scan_db

# Two SQL scans, told apart by their description.
lf  = scan_db(query_a, connection, description="trades")   # span io_source.execute[scan_db], explain_detail="trades"
lf2 = scan_db(query_b, connection, description="quotes")   # span io_source.execute[scan_db], explain_detail="quotes"
out = lf.join(lf2, on="id").collect(engine="streaming")

Third-party sources registered directly can set the identity explicitly:

from polars_io_tools.io_sources.util import register_io_source_with_is_pure

lf = register_io_source_with_is_pure(
    my_source, schema=..., explain_name="tickstore", explain_detail="collection=trades"
)

To nest source spans under a caller's span, publish a parent context (e.g. a query span) around the collect:

from opentelemetry import trace
from polars_io_tools.io_sources import profiling

profiling.set_source_span_parent(trace.set_span_in_context(my_span))
try:
    lf.collect(engine="streaming")
finally:
    profiling.set_source_span_parent(None)

Design / safety

• Standard emission: the span is created inline via the OpenTelemetry API; export is delegated to the application's span processor — use a batching processor (e.g. BatchSpanProcessor, which Logfire installs by default) to keep export off the Polars worker thread. It reuses OTel's own model (span == record, SpanProcessor/SpanExporter == sink), so there is no bespoke buffering or sink abstraction.
• No-throw: emission can never raise into the query, mask a source's exception, or change results.
• Close propagation: close() / GeneratorExit propagate to the inner iterator (releasing cursors/connections/prefetch threads), with a distinct closed outcome; a cleanup failure surfaces on a clean exit but never masks an active exception.
• Cheap when unused: timing is a couple of perf_counter_ns() reads per batch (not per row); the span is non-recording and carries no attributes when no SDK is configured.
• Timing covers every next() including the terminal StopIteration and any failing pull; the parent context is read at emit time.

Caveats (documented)

next_elapsed_total is caller-observed pull latency, not backend CPU: a driver that prefetches on a background thread (e.g. arrow-odbc's default fetch_concurrently=True) does work while the generator is suspended, which this does not attribute. total_rows counts rows yielded to the engine and can exceed a pushed n_rows due to batch overshoot. With is_pure=True, one span is emitted per physical execution (a self-join over a pure source runs once).

Tests

New tests/io_sources/test_profiling.py asserts on emitted spans via OpenTelemetry's InMemorySpanExporter: accumulation, terminal/pre-first-batch failures, consumer-pull exclusion, close() propagation to a custom iterator (including cleanup-error surfacing without masking an active error), exactly-once emission per outcome, emission failure never masking results, empty/deferred-schema cases, is_pure self-join, parent captured at emit time, source identity resolved from the original callable, distinct low-cardinality span identities, eager-factory-failure spans, registration-level wrapper composition (eager-work exclusion + close on exhaustion), the OTEL_PYTHON_INSTRUMENTATION_POLARS_IO_TOOLS_ENABLED=false opt-out, and feature-detected forwarding of explain_name/explain_detail. The single-thread Rayon deadlock regression runs both with and without an SDK configured. Full io_sources suite: 1316 passed.

API

Public source constructors (scan_db, scan_clickhouse, scan_delta, scan_datadog, scan_narwhals/from_narwhals, scan_synthetic_regression/scan_synthetic_panel, cache, cache_memory, cache_parquet, debug, concat_named, filtered_join/filtered_join_asof, execute_on_ray) gain a description argument. register_io_source_with_is_pure(..., explain_name=None, explain_detail=None); and in polars_io_tools.io_sources.profiling: profile_io_source_iterator, wrap_io_source_with_profiling, source_identity, set_source_span_parent, get_source_span_parent. Env var: OTEL_PYTHON_INSTRUMENTATION_POLARS_IO_TOOLS_ENABLED (default on; false disables with zero overhead).

[github-actions]

Test Results

1 454 tests  +29   1 406 ✅ +29   40s ⏱️ -7s
    2 suites ± 0      48 💤 ± 0 
    2 files   ± 0       0 ❌ ± 0 

Results for commit b162de5. ± Comparison against base commit 0ef48ea.

This pull request removes 1 and adds 30 tests. Note that renamed tests count towards both.

polars_io_tools.tests.io_sources.test_single_thread_deadlock ‑ test_iosource_collect_batches_does_not_deadlock_at_threads_one

polars_io_tools.tests.io_sources.test_profiling ‑ test_active_pull_error_not_masked_by_inner_close_error
polars_io_tools.tests.io_sources.test_profiling ‑ test_clean_exit_propagates_inner_close_error
polars_io_tools.tests.io_sources.test_profiling ‑ test_consumer_sleep_between_pulls_is_excluded
polars_io_tools.tests.io_sources.test_profiling ‑ test_default_source_identity_is_stable_for_callable_objects
polars_io_tools.tests.io_sources.test_profiling ‑ test_default_source_identity_trims_closure_to_enclosing_function
polars_io_tools.tests.io_sources.test_profiling ‑ test_deferred_schema_is_not_forced_for_telemetry
polars_io_tools.tests.io_sources.test_profiling ‑ test_distinct_source_names_produce_distinct_span_identities
polars_io_tools.tests.io_sources.test_profiling ‑ test_eager_callable_work_is_excluded_from_pull_time
polars_io_tools.tests.io_sources.test_profiling ‑ test_early_close_propagates_to_custom_iterator_and_emits_once[False]
polars_io_tools.tests.io_sources.test_profiling ‑ test_early_close_propagates_to_custom_iterator_and_emits_once[True]
…

♻️ This comment has been updated with latest results.