From b162de508e9014ef06e1468bd33d269823dda0c6 Mon Sep 17 00:00:00 2001 From: Pascal Tomecek <40371786+ptomecek@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:56:20 -0400 Subject: [PATCH] feat: OpenTelemetry instrumentation for custom IO sources Instrument custom Polars IO sources registered via `register_io_source_with_is_pure`. Each source execution emits one `io_source.execute[]` span recording `next_elapsed_total_ms` -- the summed duration of every `next()` pull on the source iterator, the caller-observed fetch latency that a streaming engine's own per-node metrics do not capture -- plus total rows, column count, batch count, an outcome (exhausted / closed / error), and `error.type` on failure. Instrumentation is on by default and follows the OpenTelemetry convention for native library instrumentation: spans are created through the OpenTelemetry API, which is a no-op unless the application has configured an SDK. OpenTelemetry is an optional, undeclared dependency -- when the API is not installed the instrumentation no-ops, and any application that configures an SDK already provides the API transitively. It can be disabled with `OTEL_PYTHON_INSTRUMENTATION_POLARS_IO_TOOLS_ENABLED=false` (the `OTEL_PYTHON_INSTRUMENTATION__ENABLED` opt-out convention, e.g. as elasticsearch-py uses), which skips installing the wrapper entirely for zero overhead -- letting an app keep OpenTelemetry for the rest of its code while suppressing these spans. When an SDK is configured, export is delegated to the application's span processor (use a batching processor, e.g. OpenTelemetry's BatchSpanProcessor which Logfire installs, to keep it off the Polars worker thread). Attributes are set only on recording spans and are namespaced under `polars_io_tools.`; the tracer is scoped to `polars_io_tools` with the package version. The span's identity is the source's `explain_name`, a low-cardinality label defaulting to the source function's name (a `source_generator` closure inside `scan_db` resolves to `scan_db`), or set via `explain_name`. Each 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`) also accepts a `description` argument -- a free-form per-instance label (the query, table, collection, ...) that becomes the span's `explain_detail` and lets two instances of the same source be told apart. 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. This is feature-detected against `register_io_source`'s signature (added upstream by pola-rs/polars#23978) -- a no-op on builds that lack the parameters, activating automatically once they ship. Profiling wraps the original source directly (innermost), with error-catching outermost, so timing excludes eager pre-iterator work and `close()` / `GeneratorExit` reach the real iterator (with a distinct `closed` outcome) on both early cancellation and exhaustion; a failure during the eager step still emits an `error` span. Emission is no-throw so it can never mask a source error or change results, and `set_source_span_parent()` lets a consumer nest source spans under a parent span. Adds tests (including registration-level composition tests) and extends the single-thread deadlock regression to run with and without an SDK configured. Signed-off-by: Pascal Tomecek <40371786+ptomecek@users.noreply.github.com> --- polars_io_tools/io_sources/concat_named.py | 5 +- polars_io_tools/io_sources/delta_io.py | 4 +- polars_io_tools/io_sources/join.py | 9 +- polars_io_tools/io_sources/lazy_cache.py | 4 +- .../io_sources/lazy_cache_memory.py | 4 +- .../io_sources/lazy_cache_parquet.py | 4 +- .../io_sources/lazy_clickhouse_reader.py | 4 +- .../io_sources/lazy_data_generator.py | 14 +- .../io_sources/lazy_datadog_reader.py | 4 +- polars_io_tools/io_sources/lazy_debug.py | 4 +- .../io_sources/lazy_narwhals_reader.py | 14 +- polars_io_tools/io_sources/lazy_ray.py | 4 +- polars_io_tools/io_sources/lazy_sql_reader.py | 7 +- polars_io_tools/io_sources/profiling.py | 321 ++++++++++ polars_io_tools/io_sources/util.py | 57 +- .../tests/io_sources/test_profiling.py | 606 ++++++++++++++++++ .../io_sources/test_single_thread_deadlock.py | 23 +- pyproject.toml | 1 + 18 files changed, 1064 insertions(+), 25 deletions(-) create mode 100644 polars_io_tools/io_sources/profiling.py create mode 100644 polars_io_tools/tests/io_sources/test_profiling.py diff --git a/polars_io_tools/io_sources/concat_named.py b/polars_io_tools/io_sources/concat_named.py index 3ac798b..dc6b733 100644 --- a/polars_io_tools/io_sources/concat_named.py +++ b/polars_io_tools/io_sources/concat_named.py @@ -18,6 +18,7 @@ def concat_named( identifier_cols: list[str | tuple[str, pl.DataType]], *, log_explain: bool = False, + description: str | None = None, **kwargs: Any, ) -> pl.LazyFrame: """ @@ -47,6 +48,8 @@ def concat_named( log_explain (bool, default False): If True, logs the LazyFrame execution plan for debugging purposes. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). + **kwargs (Any): Additional arguments passed to `pl.concat()` for concatenation. Returns: @@ -177,4 +180,4 @@ def source_gen( err_msg += f"\n\nWhile running the above, received error: {e.__class__.__name__}:{e}" raise RuntimeError(err_msg) from e - return register_io_source_with_is_pure(source_gen, schema=schema) + return register_io_source_with_is_pure(source_gen, schema=schema, explain_detail=description) diff --git a/polars_io_tools/io_sources/delta_io.py b/polars_io_tools/io_sources/delta_io.py index cea3a9b..4f3478a 100644 --- a/polars_io_tools/io_sources/delta_io.py +++ b/polars_io_tools/io_sources/delta_io.py @@ -578,6 +578,7 @@ def scan_delta( rechunk: bool | None = None, aws_profile: str | None = None, pushdown_predicate_deltalake: bool = True, + description: str | None = None, ) -> pl.LazyFrame: """ Lazily read from a Delta lake table with logical type translation. @@ -635,6 +636,7 @@ def scan_delta( reading file metadata. This can significantly reduce I/O on partitioned tables. When ``False``, falls back to standard ``pl.scan_delta`` behavior without partition pruning. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). Returns: LazyFrame @@ -771,4 +773,4 @@ def source_generator( yield from collect_lf_in_io_source(lf, batch_size) - return register_io_source_with_is_pure(source_generator, schema=exposed_schema) + return register_io_source_with_is_pure(source_generator, schema=exposed_schema, explain_detail=description) diff --git a/polars_io_tools/io_sources/join.py b/polars_io_tools/io_sources/join.py index 9a44198..7acf812 100644 --- a/polars_io_tools/io_sources/join.py +++ b/polars_io_tools/io_sources/join.py @@ -83,7 +83,7 @@ def _dummy_source( df = df.select(with_columns) yield df - lf = register_io_source_with_is_pure(_dummy_source, schema=schema) + lf = register_io_source_with_is_pure(_dummy_source, schema=schema, explain_name="filtered_join.predicate_rename") # We apply the rename here to have polars change the column names for us # NOTE: we are performing the rename right to left here. This might seem a bit # counterintuitive, but we are renaming the right columns to the left columns. @@ -106,6 +106,7 @@ def filtered_join( right_on: str | list[str] | None = None, nulls_equal: bool = False, log_explain: bool = False, + description: str | None = None, **join_kwargs, ) -> pl.LazyFrame: """ @@ -247,7 +248,7 @@ def source_generator( err_msg += f"\n\nError: {e.__class__.__name__}:{e}" raise RuntimeError(err_msg) from e - return register_io_source_with_is_pure(source_generator, schema=schema) + return register_io_source_with_is_pure(source_generator, schema=schema, explain_detail=description) def filtered_join_asof( @@ -263,6 +264,7 @@ def filtered_join_asof( strategy: Literal["backward", "forward", "nearest"] = "backward", tolerance: str | float | datetime.timedelta | None = None, # TODO: Only timedelta is supported for now log_explain: bool = True, + description: str | None = None, **join_kwargs, ) -> pl.LazyFrame: """ @@ -299,6 +301,7 @@ def filtered_join_asof( tolerance (timedelta, optional): Maximum time difference allowed for a match. Currently only timedelta is supported. When specified, enables temporal range expansion optimization. log_explain (bool, default True): Whether to log detailed execution plans for debugging + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). **join_kwargs: Additional keyword arguments passed to the underlying join_asof operation Returns: @@ -505,7 +508,7 @@ def source_generator( err_msg += f"\n\nError: {e.__class__.__name__}:{e}" raise RuntimeError(err_msg) from e - return register_io_source_with_is_pure(source_generator, schema=schema) + return register_io_source_with_is_pure(source_generator, schema=schema, explain_detail=description) def join_between( diff --git a/polars_io_tools/io_sources/lazy_cache.py b/polars_io_tools/io_sources/lazy_cache.py index 3fef677..da2b97b 100644 --- a/polars_io_tools/io_sources/lazy_cache.py +++ b/polars_io_tools/io_sources/lazy_cache.py @@ -139,6 +139,7 @@ def cache( cache_mode: Literal["cache", "ignore", "rebuild"] = "cache", validate: bool = True, log_explain: bool = False, + description: str | None = None, **kwargs, ) -> pl.LazyFrame: """ @@ -175,6 +176,7 @@ def cache( collected for the cache fill, so it adds no extra pass over the source. Set to False to skip it on hot paths where uniqueness is already guaranteed. log_explain: If True, logs the query plan when defining the function. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). **kwargs: Arguments to pass to the collect() method of the input data frame (i.e. to use a different engine) Notes: @@ -446,4 +448,4 @@ def source_generator( else: yield from df.iter_slices(n_rows=batch_size) - return register_io_source_with_is_pure(source_generator, schema=schema, validate_schema=True) + return register_io_source_with_is_pure(source_generator, schema=schema, validate_schema=True, explain_detail=description) diff --git a/polars_io_tools/io_sources/lazy_cache_memory.py b/polars_io_tools/io_sources/lazy_cache_memory.py index ff7037a..81a3582 100644 --- a/polars_io_tools/io_sources/lazy_cache_memory.py +++ b/polars_io_tools/io_sources/lazy_cache_memory.py @@ -72,6 +72,7 @@ def cache_memory( self_or_fn: pl.LazyFrame | Callable[[], pl.LazyFrame], *, schema: pl.Schema | Callable[[], pl.Schema], + description: str | None = None, ) -> pl.LazyFrame: """Collect a builder at most once into an in-memory buffer and replay it thereafter. @@ -97,6 +98,7 @@ def cache_memory( A callable defers resolution until Polars needs it, so ``collect_schema()`` on the result never forces the builder to run — useful when the schema is derived from the builder's own lazy plan. + description (str | None, default None): Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). Returns: pl.LazyFrame: A LazyFrame with ``schema``, backed by a generator over a one-time @@ -237,4 +239,4 @@ def source_generator( # register_io_source_with_is_pure (unlike the plain register_io_source, which raises ComputeError # for a callable schema) accepts a zero-arg callable schema and resolves it lazily. We hand it the # memoizing ``get_schema`` so Polars and the buffer reconciliation agree on one resolved schema. - return register_io_source_with_is_pure(io_source=source_generator, schema=get_schema) + return register_io_source_with_is_pure(io_source=source_generator, schema=get_schema, explain_detail=description) diff --git a/polars_io_tools/io_sources/lazy_cache_parquet.py b/polars_io_tools/io_sources/lazy_cache_parquet.py index 157d3f3..12d6f30 100644 --- a/polars_io_tools/io_sources/lazy_cache_parquet.py +++ b/polars_io_tools/io_sources/lazy_cache_parquet.py @@ -1429,6 +1429,7 @@ def cache_parquet( extra_partition_cols: str | list[str] | None = None, schema: pl.Schema | None = None, write_bounding_columns: list[str] | None = None, + description: str | None = None, ) -> pl.LazyFrame: """ Cache a LazyFrame to Parquet files with optional date-based partitioning. Supports daily, monthly, or yearly @@ -1489,6 +1490,7 @@ def cache_parquet( partition columns, not content), a bounded cache is typically regenerated with ``cache_mode=CacheMode.REBUILD`` so each run overwrites with the current predicate's rows; under ``CacheMode.CACHE`` a partition first written under one predicate is not re-written for a wider one. Default None leaves write behavior unchanged. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). Returns: pl.LazyFrame: If the cache has all data: a LazyFrame reading from the cache. @@ -1929,4 +1931,4 @@ def error_wrapper(e): end = time.time() log.debug("End: Loading data from cache at %s took %s seconds", time_unit_dir, end - start) - return register_io_source_with_is_pure(io_source=source_generator, schema=schema, validate_schema=False) + return register_io_source_with_is_pure(io_source=source_generator, schema=schema, validate_schema=False, explain_detail=description) diff --git a/polars_io_tools/io_sources/lazy_clickhouse_reader.py b/polars_io_tools/io_sources/lazy_clickhouse_reader.py index 89253f6..d31c630 100644 --- a/polars_io_tools/io_sources/lazy_clickhouse_reader.py +++ b/polars_io_tools/io_sources/lazy_clickhouse_reader.py @@ -25,7 +25,7 @@ def get_batch_reader_http(query: str, url: str, params: dict): return pa.ipc.open_stream(r.raw) -def scan_clickhouse(query: str, url: str, params: dict, fetch_size: int = 10000): +def scan_clickhouse(query: str, url: str, params: dict, fetch_size: int = 10000, description: str | None = None): # TODO: fetch_size param needs to be properly handled log.warning("fetch_size=%d is currently ignored and has no effect. Proper fetch_size support will be added in a future release.", fetch_size) dialect = "clickhouse" @@ -102,4 +102,4 @@ def select_cols(df) -> pl.DataFrame: err_msg += f"\n\nWhile running the above, received error: {e.__class__.__name__}:{e}" raise RuntimeError(err_msg) from e - return register_io_source_with_is_pure(source_generator, schema=schema) + return register_io_source_with_is_pure(source_generator, schema=schema, explain_detail=description) diff --git a/polars_io_tools/io_sources/lazy_data_generator.py b/polars_io_tools/io_sources/lazy_data_generator.py index 1549bca..d4f038a 100644 --- a/polars_io_tools/io_sources/lazy_data_generator.py +++ b/polars_io_tools/io_sources/lazy_data_generator.py @@ -46,6 +46,8 @@ def _register_source( mean_computer: MeanComputer, extras_schema: dict, chunk_sizes: np.ndarray | None = None, + explain_name: str, + explain_detail: str | None = None, ) -> pl.LazyFrame: feature_cols = [f"x{i}" for i in range(n_features)] response_cols = [f"y{i}" for i in range(n_responses)] @@ -120,7 +122,9 @@ def source_generator( log.debug("scan_synthetic: yielded %d rows; remaining_gen=%d, remaining_deliver=%d", df.height, remaining_gen, remaining_deliver) yield df - return register_io_source_with_is_pure(source_generator, schema=schema, is_pure=seed is not None) + return register_io_source_with_is_pure( + source_generator, schema=schema, is_pure=seed is not None, explain_name=explain_name, explain_detail=explain_detail + ) def _validate_common( @@ -160,6 +164,7 @@ def scan_synthetic_regression( n_chunks: int | None = None, seed: int | None = None, fetch_size: int = 10_000, + description: str | None = None, ) -> pl.LazyFrame: """ A lazy source of synthetic linear-regression data ``Y = X @ B + E`` with Gaussian noise. @@ -187,6 +192,7 @@ def scan_synthetic_regression( n_chunks: Number of contiguous chunks to split ``n_samples`` into. Required when ``chunk_key`` is set. Must satisfy ``1 <= n_chunks <= n_samples``. seed: Seed for ``np.random.default_rng``. If None, uses fresh entropy per call (and the source is registered with ``is_pure=False``). fetch_size: Default number of rows generated per batch when Polars does not provide a ``batch_size``. Must be >= 1. Defaults to 10_000. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). """ _validate_common( n_features=n_features, @@ -254,6 +260,8 @@ def mean_computer( mean_computer=mean_computer, extras_schema=extras_schema, chunk_sizes=chunk_sizes, + explain_name="scan_synthetic_regression", + explain_detail=description, ) @@ -275,6 +283,7 @@ def scan_synthetic_panel( epsilon_scale: float = 1.0, seed: int | None = None, fetch_size: int = 10_000, + description: str | None = None, ) -> pl.LazyFrame: """ A lazy source of synthetic panel data on a ``(date, symbol)`` grid. This generator draws independent per-row noise, uses ``x_i``/``y_i`` column names, and treats weights as a WLS variance model. Rows are yielded date-by-date so ``.set_sorted("date").group_by("date")`` streams cleanly under ``engine="streaming"``. @@ -302,6 +311,7 @@ def scan_synthetic_panel( epsilon_scale: When ``use_weights=False``, the noise stddev for every row. When ``use_weights=True``, the *reference* stddev at ``w=1``; actual per-row noise is ``N(loc, (epsilon_scale/√w)²)``. Must be >= 0. Defaults to 1.0. seed: Seed for ``np.random.default_rng``. If None, uses fresh entropy per call (and the source is registered with ``is_pure=False``). fetch_size: Default number of rows generated per batch when Polars does not provide a ``batch_size``. Must be >= 1. Defaults to 10_000. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). """ _validate_common( n_features=n_features, @@ -436,4 +446,6 @@ def mean_computer( mean_computer=mean_computer, extras_schema=extras_schema, chunk_sizes=chunk_sizes, + explain_name="scan_synthetic_panel", + explain_detail=description, ) diff --git a/polars_io_tools/io_sources/lazy_datadog_reader.py b/polars_io_tools/io_sources/lazy_datadog_reader.py index c0500ee..ce625b6 100644 --- a/polars_io_tools/io_sources/lazy_datadog_reader.py +++ b/polars_io_tools/io_sources/lazy_datadog_reader.py @@ -53,6 +53,7 @@ def scan_datadog( dd_interval: int | None = None, additional_schema: dict | None = None, overwrite_schema: bool = False, + description: str | None = None, ) -> pl.LazyFrame: """ Return a Polars `LazyFrame` that holds the result of a Datadog @@ -92,6 +93,7 @@ def scan_datadog( columns that you expect to be present in the response. overwrite_schema: If True, the `additional_schema` will overwrite the default schema (not add to it). You should probably avoid using this for safety reasons. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). Returns: pl.LazyFrame: A Polars LazyFrame @@ -277,4 +279,4 @@ def source_generator( yield df - return register_io_source_with_is_pure(source_generator, schema=schema) + return register_io_source_with_is_pure(source_generator, schema=schema, explain_detail=description) diff --git a/polars_io_tools/io_sources/lazy_debug.py b/polars_io_tools/io_sources/lazy_debug.py index 4ae88f0..73117b9 100644 --- a/polars_io_tools/io_sources/lazy_debug.py +++ b/polars_io_tools/io_sources/lazy_debug.py @@ -14,6 +14,7 @@ def debug( self: pl.LazyFrame, log_level: int | None = None, + description: str | None = None, ) -> pl.LazyFrame: """ A very simple pass-through lazy frame source to help with debugging experimentation of polars io sources and lazy frame behavior. @@ -21,6 +22,7 @@ def debug( Args: 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``). """ schema = self.collect_schema() @@ -53,4 +55,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) + return register_io_source_with_is_pure(source_generator, schema=schema, validate_schema=False, explain_detail=description) diff --git a/polars_io_tools/io_sources/lazy_narwhals_reader.py b/polars_io_tools/io_sources/lazy_narwhals_reader.py index 06c9b3a..bcc32fc 100644 --- a/polars_io_tools/io_sources/lazy_narwhals_reader.py +++ b/polars_io_tools/io_sources/lazy_narwhals_reader.py @@ -161,7 +161,7 @@ def polars_to_nw(pred: pl.Expr) -> nw.Expr | None: return builder.process_results() -def scan_narwhals(obj: Any, fetch_size: int) -> pl.LazyFrame: +def scan_narwhals(obj: Any, fetch_size: int, description: str | None = None) -> pl.LazyFrame: """ Turn an arbitrary Narwhals frame/series/lazyframe into a **Polars LazyFrame** so that the rest of the Polars optimisation pipeline can work unchanged. @@ -174,6 +174,7 @@ def scan_narwhals(obj: Any, fetch_size: int) -> pl.LazyFrame: does not pass a value for batch size; if it does, that will be used instead. There is no default value for this parameter, because scan_narwhals isn't supposed to be called directly. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). **kwargs: Additional arguments for the database connector Returns: pl.LazyFrame @@ -261,21 +262,21 @@ def source_generator( else: yield from pl_df.iter_slices(n_rows=bs) - return register_io_source_with_is_pure(io_source=source_generator, schema=schema) + return register_io_source_with_is_pure(io_source=source_generator, schema=schema, explain_detail=description) # from_narwhals takes either a NW lazyframe or a NW dataframe # if it's a dataframe, just call to_polars and call it a day. # if lazy, then use scan_narwhals to push things down @overload -def from_narwhals(obj: nw.DataFrame[Any], fetch_size: int = 10_000) -> pl.DataFrame: ... +def from_narwhals(obj: nw.DataFrame[Any], fetch_size: int = 10_000, description: str | None = None) -> pl.DataFrame: ... @overload -def from_narwhals(obj: nw.LazyFrame[Any], fetch_size: int = 10_000) -> pl.LazyFrame: ... +def from_narwhals(obj: nw.LazyFrame[Any], fetch_size: int = 10_000, description: str | None = None) -> pl.LazyFrame: ... -def from_narwhals(obj: FrameT, fetch_size: int = 10_000) -> pl.DataFrame | pl.LazyFrame: +def from_narwhals(obj: FrameT, fetch_size: int = 10_000, description: str | None = None) -> pl.DataFrame | pl.LazyFrame: """ Accept either a Narwhals `DataFrame` or `LazyFrame` and hands back the equivalent Polars object. @@ -283,6 +284,7 @@ def from_narwhals(obj: FrameT, fetch_size: int = 10_000) -> pl.DataFrame | pl.La Args: obj: Narwhals DataFrame **or** Narwhals LazyFrame. fetch_size (int, default 10_000): Passed through to `scan_narwhals` when `obj` is lazy. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). Only applies when `obj` is lazy and non-Polars (i.e. registered as an IO source). Returns: polars.DataFrame if `obj` is a Narwhals DataFrame @@ -297,7 +299,7 @@ def from_narwhals(obj: FrameT, fetch_size: int = 10_000) -> pl.DataFrame | pl.La return cast(pl.DataFrame | pl.LazyFrame, obj.to_native()) if callable(getattr(obj, "collect", None)): - return scan_narwhals(obj, fetch_size=fetch_size) + return scan_narwhals(obj, fetch_size=fetch_size, description=description) try: return obj.to_polars() # type: ignore[return-value] except AttributeError: diff --git a/polars_io_tools/io_sources/lazy_ray.py b/polars_io_tools/io_sources/lazy_ray.py index 4a879f7..93e6a2e 100644 --- a/polars_io_tools/io_sources/lazy_ray.py +++ b/polars_io_tools/io_sources/lazy_ray.py @@ -131,6 +131,7 @@ def execute_on_ray( return_as: Literal["arrow", "ipc", "parquet"] = "arrow", remote_options: dict | None = None, max_concurrency: int | None = 100, + description: str | None = None, ) -> pl.LazyFrame: """ Execute a Polars LazyFrame on an *already initialised* Ray cluster, @@ -157,6 +158,7 @@ def execute_on_ray( https://docs.ray.io/en/latest/ray-core/patterns/limit-pending-tasks.html. You may also wish to experiment with using resource hints to manage concurrency, although we do not recommend doing so; please see the following page for this alternative pattern: https://docs.ray.io/en/latest/ray-core/patterns/limit-running-tasks.html + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). Returns: pl.LazyFrame: A new LazyFrame whose execution plan includes information about @@ -308,4 +310,4 @@ def submit(idx: int, span: tuple[datetime, datetime]): pbar.close() - return register_io_source_with_is_pure(source_generator, schema=lambda: self.collect_schema()) + return register_io_source_with_is_pure(source_generator, schema=lambda: self.collect_schema(), explain_detail=description) diff --git a/polars_io_tools/io_sources/lazy_sql_reader.py b/polars_io_tools/io_sources/lazy_sql_reader.py index cad47bf..243284a 100644 --- a/polars_io_tools/io_sources/lazy_sql_reader.py +++ b/polars_io_tools/io_sources/lazy_sql_reader.py @@ -91,7 +91,9 @@ def get_schema_from_query_odbc( raise ValueError(f"Could not determine schema for query: {query}, with error: {e}") from e -def scan_db(query: str, connection: str, fetch_size: int = 10000, cast_map: dict[str, Any] | None = None, **kwargs) -> pl.LazyFrame: +def scan_db( + query: str, connection: str, fetch_size: int = 10000, cast_map: dict[str, Any] | None = None, description: str | None = None, **kwargs +) -> pl.LazyFrame: """ Create a LazyFrame from a SQL query with predicate pushdown support. @@ -121,6 +123,7 @@ def scan_db(query: str, connection: str, fetch_size: int = 10000, cast_map: dict optimize specific conversions (for example SQL Server seeks on ``CAST(datetime AS date)``) \ while others fall back to a scan. For a hot path on a large indexed table, prefer a \ filter expressed directly on the physical column instead of the cast one. + description: Optional free-form description of this source instance, attached to its OpenTelemetry span (``explain_detail``). **kwargs: Additional arguments for the database connector Returns: @@ -219,4 +222,4 @@ def select_cols(df) -> pl.DataFrame: err_msg += f"\n\nWhile running the above, received error: {e.__class__.__name__}:{e}" raise RuntimeError(err_msg) from e - return register_io_source_with_is_pure(source_generator, schema=schema) + return register_io_source_with_is_pure(source_generator, schema=schema, explain_detail=description) diff --git a/polars_io_tools/io_sources/profiling.py b/polars_io_tools/io_sources/profiling.py new file mode 100644 index 0000000..01cfe5b --- /dev/null +++ b/polars_io_tools/io_sources/profiling.py @@ -0,0 +1,321 @@ +"""Automatic OpenTelemetry profiling for registered IO sources. + +Each physical iterator execution is measured for its *pull latency* -- ``next_elapsed_total``, the summed wall duration of +every ``next()`` call on the source iterator. This is the fetch cost a streaming engine's own per-node metrics do not see. +It is caller-observed: it excludes work a driver performs on a background thread while the iterator is suspended (e.g. +arrow-odbc's ``fetch_concurrently``), and includes Arrow-to-Polars conversion and any pushed predicate/projection the +source applies. + +Sources registered through ``register_io_source_with_is_pure`` are instrumented automatically. Each execution is emitted +as one OpenTelemetry span through the OpenTelemetry API, which is a no-op unless the application has configured an SDK +(mirroring how native library instrumentation behaves); export is delegated to the application's span processor, so use a +batching processor to keep it off the Polars worker. Span attributes are namespaced under ``polars_io_tools.``. + +OpenTelemetry is an optional, undeclared dependency: when the ``opentelemetry`` API is not installed, instrumentation is +a no-op. Any application that has configured an OpenTelemetry SDK already provides the API transitively, so no extra +install is required to capture spans. + +Instrumentation is on by default and can be disabled -- e.g. to keep OpenTelemetry for the rest of an application while +suppressing these spans -- by setting ``OTEL_PYTHON_INSTRUMENTATION_POLARS_IO_TOOLS_ENABLED=false``, following the +OpenTelemetry Python ``OTEL_PYTHON_INSTRUMENTATION__ENABLED`` opt-out convention. When disabled the wrapper is +not installed at all, so there is zero overhead. +""" + +from __future__ import annotations + +import functools +import os +import sys +import time +from collections.abc import Callable, Generator, Iterable, Mapping +from typing import Any, Literal, TypeAlias + +try: + from opentelemetry.context import Context + + _OTEL_AVAILABLE = True +except ImportError: # pragma: no cover - exercised without OpenTelemetry installed + Context: TypeAlias = Any + _OTEL_AVAILABLE = False + +Outcome: TypeAlias = Literal["exhausted", "closed", "error"] + +_INSTRUMENTATION_NAME = "polars_io_tools" +_ATTR_PREFIX = f"{_INSTRUMENTATION_NAME}." + +_ENABLED_ENV_VAR = "OTEL_PYTHON_INSTRUMENTATION_POLARS_IO_TOOLS_ENABLED" + + +@functools.cache +def _instrumentation_enabled() -> bool: + """Whether IO-source instrumentation is active. + + On by default (the native-instrumentation convention); disable it by setting the environment variable + ``OTEL_PYTHON_INSTRUMENTATION_POLARS_IO_TOOLS_ENABLED=false``. This follows the OpenTelemetry Python + ``OTEL_PYTHON_INSTRUMENTATION__ENABLED`` opt-out convention, letting an application keep OpenTelemetry for + the rest of its code while suppressing these spans with zero overhead. + """ + return os.environ.get(_ENABLED_ENV_VAR, "").strip().lower() != "false" + + +@functools.cache +def _instrumentation_version() -> str: + """Return the installed ``polars_io_tools`` version for the tracer's instrumentation scope.""" + try: + from polars_io_tools import __version__ + except ImportError: # pragma: no cover - version metadata should always be importable + return "" + return __version__ + + +__all__ = ( + "get_source_span_parent", + "profile_io_source_iterator", + "set_source_span_parent", + "source_identity", + "wrap_io_source_with_profiling", +) + + +_span_parent: Context | None = None + + +def set_source_span_parent(ctx: Context | None) -> None: + """Publish a parent context so source spans nest under it. + + A single process-global, intended for sequential collects; concurrent collects make the parent association ambiguous. + """ + global _span_parent + _span_parent = ctx + + +def get_source_span_parent() -> Context | None: + """Return the currently published parent context, if any.""" + return _span_parent + + +def _emit_span( + *, + explain_name: str, + explain_detail: str | None, + elapsed_ns: int, + total_rows: int, + n_columns: int | None, + batch_count: int, + outcome: Outcome, + error_type: str | None, + start_ns: int, + end_ns: int, + parent: Context | None, +) -> None: + """Emit one ``io_source.execute[]`` span; a no-op without OpenTelemetry, and never raises. + + The span is created inline via the OpenTelemetry API, which is a no-op unless the application configured an SDK; + export is delegated to the application's span processor, so use a batching processor (e.g. ``BatchSpanProcessor``, + which Logfire installs by default) to keep it off the Polars worker thread. Telemetry must never affect the query, so + any error here is swallowed. + """ + if not _OTEL_AVAILABLE: + return + try: + from opentelemetry import trace + + # Name the span ``io_source.execute[]`` and attach ``explain_name`` so the span carries the + # identity the source is registered under. ``explain_name`` is expected to be low cardinality (the source kind, + # e.g. ``scan_db``); per-instance variation belongs in ``explain_detail``. + span = trace.get_tracer(_INSTRUMENTATION_NAME, _instrumentation_version()).start_span( + f"io_source.execute[{explain_name}]", context=parent, start_time=start_ns + ) + if span.is_recording(): + attributes: dict[str, Any] = { + f"{_ATTR_PREFIX}explain_name": explain_name, + f"{_ATTR_PREFIX}next_elapsed_total_ms": elapsed_ns / 1_000_000, + f"{_ATTR_PREFIX}total_rows": total_rows, + f"{_ATTR_PREFIX}batch_count": batch_count, + f"{_ATTR_PREFIX}outcome": outcome, + } + if explain_detail is not None: + attributes[f"{_ATTR_PREFIX}explain_detail"] = explain_detail + if n_columns is not None: + attributes[f"{_ATTR_PREFIX}n_columns"] = n_columns + if error_type is not None: + attributes["error.type"] = error_type # OTel semantic convention + span.set_attributes(attributes) + if outcome == "error": + from opentelemetry.trace import Status, StatusCode + + span.set_status(Status(StatusCode.ERROR)) + span.end(end_time=end_ns) + except BaseException: # noqa: BLE001, S110 - telemetry must never affect the query + pass + + +def _close_inner(inner: Any) -> None: + """Close the inner iterator, surfacing a cleanup failure only on a clean exit. + + When the wrapper is already unwinding an exception (a failing pull or a ``GeneratorExit`` from cancellation), a + ``close()`` error is swallowed so it cannot mask the active exception; on a clean exit it propagates, matching the + unwrapped iterator's behaviour. + """ + close = getattr(inner, "close", None) + if close is None: + return + if sys.exc_info()[0] is None: + close() + else: + try: + close() + except BaseException: # noqa: BLE001, S110 - do not mask the active exception + pass + + +def profile_io_source_iterator( + inner_iterable: Iterable[Any], + *, + explain_name: str, + explain_detail: str | None = None, + fallback_n_columns: int | None = None, +) -> Generator[Any, None, None]: + """Wrap one iterator execution in a self-measuring, close-propagating generator. + + The returned generator yields the inner batches unchanged and emits exactly one ``io_source.execute`` span when it + finishes (``exhausted`` / ``closed`` / ``error``). ``explain_name`` names the span and is attached as an attribute; + ``explain_detail`` is an optional free-form per-instance description. ``total_rows`` counts rows yielded and may + exceed a pushed ``n_rows`` when the final batch overshoots it. + """ + inner = iter(inner_iterable) + + def timed() -> Generator[Any, None, None]: + elapsed_ns = total_rows = batch_count = 0 + n_columns: int | None = None + outcome: Outcome = "exhausted" + error_type: str | None = None + wall_start_ns: int | None = None + try: + # Prime point: suspending inside the try means a close() before the + # first real pull still runs the finally (emit + propagate close). + yield + while True: + if wall_start_ns is None: + wall_start_ns = time.time_ns() + pull_start = time.perf_counter_ns() + try: + batch = next(inner) + except StopIteration: + elapsed_ns += time.perf_counter_ns() - pull_start + break + except BaseException: + elapsed_ns += time.perf_counter_ns() - pull_start + outcome = "error" + raise + elapsed_ns += time.perf_counter_ns() - pull_start + total_rows += batch.height + if n_columns is None: + n_columns = batch.width + batch_count += 1 + yield batch + except GeneratorExit: + outcome = "closed" + raise + except BaseException as exc: + outcome = "error" + error_type = type(exc).__name__ + raise + finally: + now_ns = time.time_ns() + # Emit before closing so telemetry is recorded regardless of whether the inner ``close()`` succeeds. The + # parent context is read here (emit time): a consumer publishes it after the iterator is built, so capturing + # earlier races it. + _emit_span( + explain_name=explain_name, + explain_detail=explain_detail, + elapsed_ns=elapsed_ns, + total_rows=total_rows, + n_columns=n_columns if n_columns is not None else fallback_n_columns, + batch_count=batch_count, + outcome=outcome, + error_type=error_type, + start_ns=wall_start_ns if wall_start_ns is not None else now_ns, + end_ns=now_ns, + parent=get_source_span_parent(), + ) + _close_inner(inner) + + generator = timed() + next(generator) # advance to the prime point + return generator + + +def source_identity(io_source: Callable[..., Any], explain_name: str | None = None) -> str: + """Resolve the ``explain_name`` for a source callable. + + Defaults to the callable's qualified name (or its class for callable instances), with the local-closure suffix + trimmed: a source defined as ``source_generator`` inside ``scan_db`` resolves to ``scan_db``, the enclosing function + that names the source kind, rather than ``scan_db..source_generator``. Resolve this from the *original* + callable, before any wrapping that would mask its identity. + """ + if explain_name is not None: + return explain_name + for attribute in ("__qualname__", "__name__"): + value = getattr(io_source, attribute, None) + if isinstance(value, str) and value: + # Drop the inner-closure boilerplate: the enclosing function is the meaningful, low-cardinality identity. + return value.split("..", 1)[0] + cls = type(io_source) + return f"{cls.__module__}.{cls.__qualname__}" + + +def _fallback_n_columns(args: tuple[Any, ...], kwargs: dict[str, Any], schema: Any) -> int | None: + with_columns = kwargs.get("with_columns", args[0] if args else None) + if with_columns is not None: + return len(with_columns) + if isinstance(schema, Mapping): + return len(schema) + return None + + +def wrap_io_source_with_profiling( + io_source: Callable[..., Iterable[Any]], + *, + schema: Any, + explain_name: str | None = None, + explain_detail: str | None = None, +) -> Callable[..., Generator[Any, None, None]]: + """Wrap an IO-source callable so each execution emits one ``io_source.execute`` span. + + Work performed eagerly by ``io_source`` before it returns its iterator is outside the pull-time measurement, but a + failure during that eager step still emits an ``error`` span; generator-body work is captured on the first pull. + ``explain_name`` names the span (defaults to the source function's name); ``explain_detail`` is an optional free-form + per-instance description. + """ + name = source_identity(io_source, explain_name) + + @functools.wraps(io_source) + def profiled(*args: Any, **kwargs: Any) -> Generator[Any, None, None]: + start_ns = time.time_ns() + try: + iterable = io_source(*args, **kwargs) + except BaseException as exc: + # A non-generator source can fail while eagerly building its iterator, before the timed generator exists. + _emit_span( + explain_name=name, + explain_detail=explain_detail, + elapsed_ns=0, + total_rows=0, + n_columns=_fallback_n_columns(args, kwargs, schema), + batch_count=0, + outcome="error", + error_type=type(exc).__name__, + start_ns=start_ns, + end_ns=time.time_ns(), + parent=get_source_span_parent(), + ) + raise + return profile_io_source_iterator( + iterable, + explain_name=name, + explain_detail=explain_detail, + fallback_n_columns=_fallback_n_columns(args, kwargs, schema), + ) + + return profiled diff --git a/polars_io_tools/io_sources/util.py b/polars_io_tools/io_sources/util.py index 86cecbb..0d65837 100644 --- a/polars_io_tools/io_sources/util.py +++ b/polars_io_tools/io_sources/util.py @@ -1,4 +1,7 @@ +from __future__ import annotations + import functools +import inspect import logging import os import socket @@ -331,7 +334,22 @@ def error_catching_io_source(*args, **source_kwargs): return error_catching_io_source -def register_io_source_with_is_pure(io_source, schema, wrap_with_error_catching: bool = True, **kwargs): +@functools.cache +def _register_io_source_supports_explain_labels() -> bool: + """Whether ``register_io_source`` accepts ``explain_name`` / ``explain_detail`` (pola-rs/polars#23978, unreleased).""" + from polars.io.plugins import register_io_source + + return {"explain_name", "explain_detail"} <= inspect.signature(register_io_source).parameters.keys() + + +def register_io_source_with_is_pure( + io_source, + schema, + wrap_with_error_catching: bool = True, + explain_name: str | None = None, + explain_detail: str | None = None, + **kwargs, +): """ Register an io source with is_pure=True if Polars version >= 1.33.1. @@ -339,6 +357,17 @@ def register_io_source_with_is_pure(io_source, schema, wrap_with_error_catching: Polars versions that support it (1.33.1+). Optionally wraps the source with `wrap_io_source_with_error_catching` for better diagnostics. + Each source execution is instrumented with OpenTelemetry: one + ``io_source.execute[]`` span recording its pull latency. This + is a no-op unless the application has configured an OpenTelemetry SDK, so it + is always installed and costs nothing when telemetry is not in use. It can + be disabled with ``OTEL_PYTHON_INSTRUMENTATION_POLARS_IO_TOOLS_ENABLED=false`` + (the wrapper is then not installed, so there is zero overhead). + ``explain_name`` names the span -- low cardinality, defaulting to the source + function's name; ``explain_detail`` is a free-form per-instance description. + On Polars builds that accept them, both are forwarded to the scan so its + registered identity and description match the span's. + Args: io_source (callable): The IO source function schema (dict, pl.Schema, or callable returning either): The schema for the IO source. May be passed as an eagerly-resolved @@ -347,6 +376,8 @@ def register_io_source_with_is_pure(io_source, schema, wrap_with_error_catching: Polars actually needs it (typically at collect time), which avoids forcing ``collect_schema()`` on input LazyFrames at construction. Callable schemas require Polars >= 1.22.0. + explain_name: Low-cardinality label (the source kind, e.g. ``scan_db``) naming the scan and the span; defaults to the source function's name. + explain_detail: Optional free-form per-instance description shown as the scan's ``explain_detail`` and attached to the span (e.g. the table / collection / query). **kwargs: Additional keyword arguments to pass to register_io_source Returns: @@ -354,10 +385,34 @@ def register_io_source_with_is_pure(io_source, schema, wrap_with_error_catching: """ from polars.io.plugins import register_io_source + from .profiling import _instrumentation_enabled, source_identity, wrap_io_source_with_profiling + + original_io_source = io_source + # Check if Polars version supports is_pure parameter (1.33.1+) if version.parse(pl.__version__) >= version.parse("1.33.1"): kwargs.setdefault("is_pure", True) + # Resolve the identity from the *original* callable -- error-catching wraps it + # in a generic closure that would otherwise mask distinct sources -- and use + # it consistently for the scan's ``explain_name`` and the profiling span. + name = source_identity(original_io_source, explain_name) + if _register_io_source_supports_explain_labels(): + kwargs.setdefault("explain_name", name) + if explain_detail is not None: + kwargs.setdefault("explain_detail", explain_detail) + + # Profile the original source directly (innermost) so timing excludes eager + # pre-iterator work and ``close()`` reaches the real iterator; error-catching + # then decorates the profiled source. Skipping the wrapper entirely when + # instrumentation is disabled keeps the disabled path zero-overhead. + if _instrumentation_enabled(): + io_source = wrap_io_source_with_profiling( + io_source, + schema=schema, + explain_name=name, + explain_detail=explain_detail, + ) if wrap_with_error_catching: io_source = wrap_io_source_with_error_catching(io_source) return register_io_source(io_source, schema=schema, **kwargs) diff --git a/polars_io_tools/tests/io_sources/test_profiling.py b/polars_io_tools/tests/io_sources/test_profiling.py new file mode 100644 index 0000000..fc66ec9 --- /dev/null +++ b/polars_io_tools/tests/io_sources/test_profiling.py @@ -0,0 +1,606 @@ +import time +from collections.abc import Iterator + +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_io_tools.io_sources.profiling import ( + get_source_span_parent, + profile_io_source_iterator, + set_source_span_parent, + wrap_io_source_with_profiling, +) +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() + set_source_span_parent(None) + yield + set_source_span_parent(None) + + +def _profile(iterable, *, fallback_n_columns=None): + return profile_io_source_iterator(iterable, explain_name="test.reader", fallback_n_columns=fallback_n_columns) + + +def _spans(): + return [s for s in _EXPORTER.get_finished_spans() if s.name.startswith("io_source.execute")] + + +def _span(): + spans = _spans() + assert len(spans) == 1 + return spans[0] + + +def _module_level_reader(): + # A module-level factory whose inner closure has qualname ``_module_level_reader..source_generator``, + # matching how the real readers define their ``source_generator``. + def source_generator(*_args, **_kwargs): + return iter([pl.DataFrame({"value": [1]})]) + + return source_generator + + +def test_multi_batch_accumulates_only_pull_time(): + sleep_s = 0.015 + + def source() -> Iterator[pl.DataFrame]: + for value in range(3): + time.sleep(sleep_s) + yield pl.DataFrame({"value": [value]}) + + assert len(list(_profile(source()))) == 3 + + span = _span() + assert span.attributes["polars_io_tools.next_elapsed_total_ms"] >= 3 * sleep_s * 0.8 * 1000 + assert span.attributes["polars_io_tools.total_rows"] == 3 + assert span.attributes["polars_io_tools.n_columns"] == 1 + assert span.attributes["polars_io_tools.batch_count"] == 3 + assert span.attributes["polars_io_tools.outcome"] == "exhausted" + + +def test_slow_terminal_stop_iteration_is_counted(): + terminal_sleep_s = 0.025 + + class SlowTerminalIterator: + def __init__(self) -> None: + self._yielded = False + + def __iter__(self): + return self + + def __next__(self): + if not self._yielded: + self._yielded = True + return pl.DataFrame({"value": [1]}) + time.sleep(terminal_sleep_s) + raise StopIteration + + assert len(list(_profile(SlowTerminalIterator()))) == 1 + span = _span() + assert span.attributes["polars_io_tools.next_elapsed_total_ms"] >= terminal_sleep_s * 0.8 * 1000 + assert span.attributes["polars_io_tools.outcome"] == "exhausted" + + +def test_failure_before_first_yield_is_counted_and_propagated_unchanged(): + error = ValueError("source failed") + failure_sleep_s = 0.025 + + class FailingIterator: + def __iter__(self): + return self + + def __next__(self): + time.sleep(failure_sleep_s) + raise error + + with pytest.raises(ValueError) as exc_info: + next(_profile(FailingIterator())) + + assert exc_info.value is error + span = _span() + assert span.attributes["polars_io_tools.next_elapsed_total_ms"] >= failure_sleep_s * 0.8 * 1000 + assert span.attributes["polars_io_tools.total_rows"] == 0 + assert span.attributes["polars_io_tools.batch_count"] == 0 + assert span.attributes["polars_io_tools.outcome"] == "error" + assert span.attributes["error.type"] == "ValueError" + + +def test_failure_reports_partial_rows(): + error = LookupError("failed after a batch") + + def source() -> Iterator[pl.DataFrame]: + yield pl.DataFrame({"value": [1, 2]}) + time.sleep(0.015) + raise error + + iterator = _profile(source()) + assert next(iterator).height == 2 + with pytest.raises(LookupError) as exc_info: + next(iterator) + + assert exc_info.value is error + span = _span() + assert span.attributes["polars_io_tools.total_rows"] == 2 + assert span.attributes["polars_io_tools.n_columns"] == 1 + assert span.attributes["polars_io_tools.batch_count"] == 1 + assert span.attributes["polars_io_tools.outcome"] == "error" + assert span.attributes["polars_io_tools.next_elapsed_total_ms"] >= 10 + + +def test_consumer_sleep_between_pulls_is_excluded(): + iterator = _profile(iter([pl.DataFrame({"value": [1]}), pl.DataFrame({"value": [2]})])) + + next(iterator) + time.sleep(0.1) + next(iterator) + with pytest.raises(StopIteration): + next(iterator) + + assert _span().attributes["polars_io_tools.next_elapsed_total_ms"] < 50 + + +@pytest.mark.parametrize("pull_first", [False, True]) +def test_early_close_propagates_to_custom_iterator_and_emits_once(pull_first): + class ClosableIterator: + def __init__(self) -> None: + self.closed = False + + def __iter__(self): + return self + + def __next__(self): + return pl.DataFrame({"value": [1]}) + + def close(self): + self.closed = True + + inner = ClosableIterator() + iterator = _profile(inner) + if pull_first: + next(iterator) + iterator.close() + iterator.close() + + assert inner.closed + span = _span() + assert span.attributes["polars_io_tools.outcome"] == "closed" + assert span.attributes["polars_io_tools.total_rows"] == (1 if pull_first else 0) + + +def test_exactly_once_emission_for_all_outcomes(): + assert list(_profile(iter(()))) == [] + assert _span().attributes["polars_io_tools.outcome"] == "exhausted" + _EXPORTER.clear() + + closed = _profile(iter([pl.DataFrame({"value": [1]})])) + next(closed) + closed.close() + assert _span().attributes["polars_io_tools.outcome"] == "closed" + _EXPORTER.clear() + + def failing() -> Iterator[pl.DataFrame]: + raise ZeroDivisionError("boom") + yield + + with pytest.raises(ZeroDivisionError): + next(_profile(failing())) + assert _span().attributes["polars_io_tools.outcome"] == "error" + + +def test_emission_failure_never_masks_completion_or_source_error(monkeypatch): + import opentelemetry.trace as ot + + def boom(*_args, **_kwargs): + raise RuntimeError("tracer boom") + + monkeypatch.setattr(ot, "get_tracer", boom) + + frames = list(_profile(iter([pl.DataFrame({"value": [1]})]))) + assert frames[0].height == 1 + + error = OSError("source error") + + def failing() -> Iterator[pl.DataFrame]: + raise error + yield + + with pytest.raises(OSError) as exc_info: + next(_profile(failing())) + assert exc_info.value is error + + +def test_no_yield_uses_known_projection_or_eager_schema_width(): + wrapped = wrap_io_source_with_profiling( + lambda *_args, **_kwargs: iter(()), + schema={"a": pl.Int64, "b": pl.Int64}, + explain_name="empty.reader", + ) + assert list(wrapped(["a"], None, None, None)) == [] + assert _span().attributes["polars_io_tools.n_columns"] == 1 + _EXPORTER.clear() + + wrapped = wrap_io_source_with_profiling( + lambda *_args, **_kwargs: iter(()), + schema={"a": pl.Int64, "b": pl.Int64}, + explain_name="empty.reader", + ) + assert list(wrapped(None, None, None, None)) == [] + assert _span().attributes["polars_io_tools.n_columns"] == 2 + + +def test_empty_first_batch_sets_width_and_zero_rows(): + frames = list(_profile(iter([pl.DataFrame(schema={"a": pl.Int64, "b": pl.String})]))) + assert frames[0].is_empty() + + span = _span() + assert span.attributes["polars_io_tools.total_rows"] == 0 + assert span.attributes["polars_io_tools.n_columns"] == 2 + assert span.attributes["polars_io_tools.batch_count"] == 1 + + +def test_deferred_schema_is_not_forced_for_telemetry(): + schema_calls = 0 + + def schema(): + nonlocal schema_calls + schema_calls += 1 + return {"a": pl.Int64} + + wrapped = wrap_io_source_with_profiling( + lambda *_args, **_kwargs: iter(()), + schema=schema, + explain_name="empty.reader", + ) + assert list(wrapped(None, None, None, None)) == [] + assert schema_calls == 0 + assert "polars_io_tools.n_columns" not in _span().attributes + + +def test_eager_callable_work_is_excluded_from_pull_time(): + def eager_source(*_args, **_kwargs): + time.sleep(0.05) + return iter([pl.DataFrame({"value": [1]})]) + + wrapped = wrap_io_source_with_profiling( + eager_source, + schema={"value": pl.Int64}, + explain_name="eager.reader", + ) + frames = list(wrapped(None, None, None, None)) + assert frames[0].height == 1 + assert _span().attributes["polars_io_tools.next_elapsed_total_ms"] < 30 + + +def test_source_parent_is_captured_at_emit_time(): + # The parent context is read when the span is emitted (after execution), NOT at construction -- a query observer + # publishes its context after the LazyFrame/iterator is built, so a parent set only after construction must apply. + parent = trace.get_tracer("test").start_span("parent") + iterator = _profile(iter([pl.DataFrame({"value": [1]})])) + set_source_span_parent(trace.set_span_in_context(parent)) # published after construction + try: + list(iterator) + finally: + set_source_span_parent(None) + parent.end() + + span = _span() + assert span.parent is not None + assert span.parent.span_id == parent.get_span_context().span_id + + +def test_source_parent_not_captured_at_construction(): + # A parent present at construction but cleared before execution completes is not captured. + parent = trace.get_tracer("test").start_span("parent") + set_source_span_parent(trace.set_span_in_context(parent)) + iterator = _profile(iter([pl.DataFrame({"value": [1]})])) + set_source_span_parent(None) # cleared before the source is drained + list(iterator) + parent.end() + + assert get_source_span_parent() is None + assert _span().parent is None + + +def test_default_source_identity_is_stable_for_callable_objects(): + class CallableSource: + def __call__(self, *_args, **_kwargs): + return iter([pl.DataFrame({"value": [1]})]) + + wrapped = wrap_io_source_with_profiling(CallableSource(), schema={"value": pl.Int64}) + list(wrapped(None, None, None, None)) + + name = _span().attributes["polars_io_tools.explain_name"] + assert name.endswith("CallableSource") + assert "0x" not in name + + +def test_default_source_identity_trims_closure_to_enclosing_function(): + # Mirrors the real readers: a ``source_generator`` closure defined inside ``scan_db`` should resolve to ``scan_db``, + # not ``scan_db..source_generator``. + wrapped = wrap_io_source_with_profiling(_module_level_reader(), schema={"value": pl.Int64}) + list(wrapped(None, None, None, None)) + + assert _span().attributes["polars_io_tools.explain_name"] == "_module_level_reader" + + +def test_pure_self_join_emits_one_physical_execution(): + calls = 0 + + def source(with_columns, predicate, n_rows, batch_size): + nonlocal calls + calls += 1 + frame = pl.DataFrame({"id": [1, 2], "value": [10, 20]}) + if with_columns is not None: + frame = frame.select(with_columns) + yield frame + + lf = register_io_source_with_is_pure( + source, + schema={"id": pl.Int64, "value": pl.Int64}, + explain_name="memory.reader", + ) + result = lf.join(lf, on="id").collect() + + assert result.height == 2 + assert calls == 1 + span = _span() + assert span.attributes["polars_io_tools.explain_name"] == "memory.reader" + assert span.attributes["polars_io_tools.outcome"] == "exhausted" + + +def _capture_registered_callback(source, monkeypatch, **kwargs): + # Assemble the callback exactly as register_io_source_with_is_pure hands it to polars, so tests exercise the real + # wrapper composition (profiling + error-catching) rather than profile_io_source_iterator in isolation. + captured = {} + + def fake_register(io_source, *, schema, **_kw): + captured["io_source"] = io_source + return object() + + monkeypatch.setattr("polars.io.plugins.register_io_source", fake_register) + register_io_source_with_is_pure(source, schema={"value": pl.Int64}, **kwargs) + return captured["io_source"] + + +def test_registration_excludes_eager_work_from_pull_time(monkeypatch): + # With error-catching enabled (the default), eager work the source does before returning its iterator must still be + # excluded from pull time -- i.e. profiling wraps the original source, not the error-catching generator. + def source(*_args, **_kwargs): + time.sleep(0.05) + return iter([pl.DataFrame({"value": [1]})]) + + callback = _capture_registered_callback(source, monkeypatch) + list(callback(None, None, None, None)) + + assert _span().attributes["polars_io_tools.next_elapsed_total_ms"] < 40 + + +def test_registration_closes_real_iterator_on_exhaustion(monkeypatch): + # Profiling wraps the original source, so its close() reaches the real iterator on exhaustion even under the + # error-catching wrapper (whose yield-from would not close a custom iterator on normal completion). + class RealIter: + def __init__(self): + self.closed = False + self.remaining = 2 + + def __iter__(self): + return self + + def __next__(self): + if self.remaining <= 0: + raise StopIteration + self.remaining -= 1 + return pl.DataFrame({"value": [1]}) + + def close(self): + self.closed = True + + iterator = RealIter() + callback = _capture_registered_callback(lambda *_a, **_k: iterator, monkeypatch) + list(callback(None, None, None, None)) + + assert iterator.closed + + +def test_registration_emits_error_span_for_eager_factory_failure(monkeypatch): + # A non-generator source can raise while building its iterator, before the timed generator exists; that failure must + # still produce an ``error`` span with ``error.type``. + def source(*_args, **_kwargs): + raise ConnectionError("connect failed") + + callback = _capture_registered_callback(source, monkeypatch, wrap_with_error_catching=False) + with pytest.raises(ConnectionError): + list(callback(None, None, None, None)) + + span = _span() + assert span.attributes["polars_io_tools.outcome"] == "error" + assert span.attributes["error.type"] == "ConnectionError" + + +def test_instrumentation_opt_out_disables_spans_with_zero_overhead(monkeypatch): + # The OTEL_PYTHON_INSTRUMENTATION__ENABLED opt-out: when disabled the profiling wrapper is not installed, so no + # span is emitted and the query still runs unchanged. + from polars_io_tools.io_sources import profiling + + monkeypatch.setenv("OTEL_PYTHON_INSTRUMENTATION_POLARS_IO_TOOLS_ENABLED", "false") + profiling._instrumentation_enabled.cache_clear() + + def source(with_columns, predicate, n_rows, batch_size): + yield pl.DataFrame({"id": [1, 2, 3]}) + + try: + result = register_io_source_with_is_pure(source, schema={"id": pl.Int64}, explain_name="scan_x").collect(engine="streaming") + finally: + profiling._instrumentation_enabled.cache_clear() + + assert result.height == 3 + assert _spans() == [] + + +def test_registration_wraps_source_and_forwards_kwargs(monkeypatch): + # Instrumentation is always installed, so the callback handed to register_io_source is the profiling wrapper, not the + # original; schema and pass-through kwargs are preserved. + captured = {} + sentinel = object() + + def fake_register(io_source, *, schema, **kwargs): + captured["io_source"] = io_source + captured["schema"] = schema + captured["kwargs"] = kwargs + return sentinel + + monkeypatch.setattr("polars.io.plugins.register_io_source", fake_register) + + def source(with_columns, predicate, n_rows, batch_size): + return iter([pl.DataFrame({"value": [1]})]) + + result = register_io_source_with_is_pure( + source, + schema={"value": pl.Int64}, + wrap_with_error_catching=False, + explain_name="ignored.reader", + validate_schema=True, + ) + + assert result is sentinel + assert captured["io_source"] is not source + assert captured["schema"] == {"value": pl.Int64} + assert captured["kwargs"]["validate_schema"] is True + + +def test_clean_exit_propagates_inner_close_error(): + class ClosingRaises: + def __iter__(self): + return self + + def __next__(self): + raise StopIteration + + def close(self): + raise RuntimeError("close failed") + + with pytest.raises(RuntimeError, match="close failed"): + list(_profile(ClosingRaises())) + + # The span is still emitted (before closing), with a clean outcome. + assert _span().attributes["polars_io_tools.outcome"] == "exhausted" + + +def test_active_pull_error_not_masked_by_inner_close_error(): + class BothRaise: + def __iter__(self): + return self + + def __next__(self): + raise ValueError("pull failed") + + def close(self): + raise RuntimeError("close failed") + + with pytest.raises(ValueError, match="pull failed"): + next(_profile(BothRaise())) + + assert _span().attributes["polars_io_tools.outcome"] == "error" + + +def test_registration_resolves_identity_from_original_callable(): + # Error-catching wraps the source in a generic closure; the profiling identity must come from the original callable. + + class AlphaSource: + def __call__(self, with_columns, predicate, n_rows, batch_size): + frame = pl.DataFrame({"id": [1, 2]}) + yield frame.select(with_columns) if with_columns else frame + + class BetaSource: + def __call__(self, with_columns, predicate, n_rows, batch_size): + frame = pl.DataFrame({"id": [3, 4]}) + yield frame.select(with_columns) if with_columns else frame + + for source in (AlphaSource(), BetaSource()): + register_io_source_with_is_pure(source, schema={"id": pl.Int64}).collect(engine="streaming") + + names = {span.attributes["polars_io_tools.explain_name"] for span in _spans()} + assert any(name.endswith("AlphaSource") for name in names) + assert any(name.endswith("BetaSource") for name in names) + assert not any("error_catching" in name for name in names) + + +def test_distinct_source_names_produce_distinct_span_identities(): + # Two registrations of the same reader with different identities must emit distinctly named spans, carrying the + # ``explain_name`` / ``explain_detail`` they are registered under. + + def make_reader(): + def _reader(with_columns, predicate, n_rows, batch_size): + frame = pl.DataFrame({"id": [1, 2]}) + yield frame.select(with_columns) if with_columns else frame + + return _reader + + for collection in ("trades", "quotes"): + register_io_source_with_is_pure( + make_reader(), + schema={"id": pl.Int64}, + explain_name=f"tickstore.{collection}", + explain_detail=f"collection={collection}", + ).collect(engine="streaming") + + spans = {span.name: span for span in _spans()} + assert set(spans) == {"io_source.execute[tickstore.trades]", "io_source.execute[tickstore.quotes]"} + for collection in ("trades", "quotes"): + span = spans[f"io_source.execute[tickstore.{collection}]"] + assert span.attributes["polars_io_tools.explain_name"] == f"tickstore.{collection}" + assert span.attributes["polars_io_tools.explain_detail"] == f"collection={collection}" + + +def test_explain_labels_forwarded_to_register_io_source_when_supported(monkeypatch): + # When the Polars build exposes explain_name/explain_detail, the resolved identity and detail are forwarded to the + # scan so its node label matches the profiling span. Feature-detected against register_io_source's signature. + from polars_io_tools.io_sources import util + + captured = {} + + def fake_register(io_source, *, schema, explain_name=None, explain_detail=None, **kwargs): + captured["explain_name"] = explain_name + captured["explain_detail"] = explain_detail + return object() + + monkeypatch.setattr("polars.io.plugins.register_io_source", fake_register) + util._register_io_source_supports_explain_labels.cache_clear() + + def source(with_columns, predicate, n_rows, batch_size): + yield pl.DataFrame({"value": [1]}) + + try: + register_io_source_with_is_pure( + source, + schema={"value": pl.Int64}, + wrap_with_error_catching=False, + explain_name="tickstore.trades", + explain_detail="collection=trades", + ) + finally: + util._register_io_source_supports_explain_labels.cache_clear() + + assert captured["explain_name"] == "tickstore.trades" + assert captured["explain_detail"] == "collection=trades" diff --git a/polars_io_tools/tests/io_sources/test_single_thread_deadlock.py b/polars_io_tools/tests/io_sources/test_single_thread_deadlock.py index 04f3918..038d763 100644 --- a/polars_io_tools/tests/io_sources/test_single_thread_deadlock.py +++ b/polars_io_tools/tests/io_sources/test_single_thread_deadlock.py @@ -39,7 +39,24 @@ def _run_in_subprocess(script: str) -> subprocess.CompletedProcess: ) -def test_iosource_collect_batches_does_not_deadlock_at_threads_one(): +@pytest.mark.parametrize( + "otel_setup", + [ + "", + """ + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + provider = TracerProvider() + provider.add_span_processor(BatchSpanProcessor(InMemorySpanExporter())) + trace.set_tracer_provider(provider) + """, + ], + ids=["otel-no-op", "otel-batch-processor"], +) +def test_iosource_collect_batches_does_not_deadlock_at_threads_one(otel_setup): """Reproduces the polars-io-tools / polars deadlock and asserts the fix holds. Without the helper's thread-pool gate this script hangs forever at @@ -58,6 +75,8 @@ def test_iosource_collect_batches_does_not_deadlock_at_threads_one(): assert pl.thread_pool_size() == 1, pl.thread_pool_size() + __OTEL_SETUP__ + # This is the minimal pattern that deadlocks polars 1.39.3+1.40.1 # without the fix in collect_lf_in_io_source. Removing any one of: # the join, the pushed-down predicate, the inner filter+collect_batches, @@ -90,7 +109,7 @@ def reader(with_columns, predicate, n_rows, batch_size): assert out.height == 100_000, out.height print("OK", out.height) """ - ) + ).replace("__OTEL_SETUP__", textwrap.dedent(otel_setup)) try: result = _run_in_subprocess(script) diff --git a/pyproject.toml b/pyproject.toml index 3f9d3ce..eaacf19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,6 +63,7 @@ develop = [ "hatchling", "mdformat", "mdformat-tables>=1", + "opentelemetry-sdk", "pytest", "pytest-cov", "ruff",