Merge pull request #30 from Point72/feat/io-source-profiling

feat: OpenTelemetry instrumentation for custom IO sources

Author: ptomecek

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)

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)

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(

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)

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)

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)

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)

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,
     )

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)

polars_io_tools/io_sources/lazy_debug.py

@@ -14,13 +14,15 @@
 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.
 
     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)