feat(reporting): derive TimeWindowEvent event fact rows with mapInArrow

docs/impulse/docs/references/api/impulse_reporting/events/time_window_event.md

@@ -19,7 +19,12 @@ event instance per fixed-duration slice, tiling the container's ``start_ts`` / `
 span with windows of length ``window_length``.  The final slice is clamped to the
 container end.
 
-The event fact is computed from ``container_metrics`` alone (via
+The event fact is computed from ``container_metrics`` alone (in a ``mapInArrow``), so
+every filtered container gets windows regardless of its channel data.  Aggregations
+scoped to this event evaluate the :class:`TimeWindowExpression` in the solve.  Both use
+the same window function (``tile_windows``), so they produce identical windows, and the
+timestamp-based ``event_instance_id`` (like for other interval events) matches on both
+sides.
 
 
 #### \_\_init\_\_
@@ -143,7 +148,7 @@ containers' ``start_ts`` / ``stop_ts`` in the channel time frame
 gets windows.
 Each window becomes one event instance (``start_ts < end_ts``) whose
 ``event_instance_id`` hashes its boundaries. The solve uses the same window function
-for scoped aggregations (see :func:`window_intervals_udf`), so the ids match.
+for scoped aggregations (see :func:`tile_windows`), so the ids match.
 
 **Arguments**:
 

src/impulse_query_engine/analyze/query/events/time_window_expression.py

@@ -5,7 +5,6 @@
 
 import numpy as np
 import pandas as pd
-import pyspark.sql.functions as F
 
 from impulse_query_engine.analyze.metadata.tag_expression import TagExpression
 from impulse_query_engine.analyze.metadata.time_series_expression import (
@@ -73,13 +72,14 @@ def tile_windows(
 
     The one window implementation behind ``TimeWindowEvent``: the solve calls it through
     :meth:`TimeWindowExpression.build` (scoped aggregations), the event fact through
-    :func:`window_intervals_udf`.  ``event_instance_id`` hashes each window's boundaries, so
+    ``TimeWindowEvent.determine_events``.  ``event_instance_id`` hashes each window's boundaries, so
     both sides must produce identical windows, which a single function guarantees as long as
     both pass in the same values.  Both read the same Spark-computed bounds
-    (``solvers.utils.window_bounds.with_window_bounds``), but pandas hands them over as
-    ``int64`` or ``float64`` (nulls force ``float64``), or as ``None`` / ``NaN``.  The bounds are
-    therefore converted to ``float`` first: ``int64`` -> ``float64`` rounds to the nearest
-    double on either path, so the arithmetic below runs on identical doubles.
+    (``solvers.utils.window_bounds.with_window_bounds``), but receive them as Python or numpy
+    integers or floats (pandas turns long columns with nulls into ``float64``), or as
+    ``None`` / ``NaN``.  The bounds are therefore converted to ``float`` first: integer ->
+    double rounds to the nearest double on either path, so the arithmetic below runs on
+    identical doubles.
 
     Window ``i`` spans ``[start + i * W, min(start + (i + 1) * W, stop)]``, so the last one
     is clamped to *stop*; windows with ``start_i >= end_i`` (possible only through rounding)
@@ -128,48 +128,6 @@ def tile_windows(
     return starts[keep], ends[keep]
 
 
-def window_intervals_udf(window_length: float, max_windows: int = MAX_WINDOWS_PER_CONTAINER):
-    """Scalar pandas UDF giving each container's windows via :func:`tile_windows`.
-
-    Used by the reporting ``TimeWindowEvent`` for its event fact (one row per container), so
-    the event fact and the solve share one window implementation.
-
-    Parameters
-    ----------
-    window_length : float
-        Fixed window length, in the same unit as the bounds. Strictly positive.
-    max_windows : int, optional
-        Maximum number of windows per container (default
-        :data:`MAX_WINDOWS_PER_CONTAINER`). A container exceeding it fails the query with
-        an error naming the limit.
-
-    Returns
-    -------
-    callable
-        A pandas UDF ``(start, stop) -> struct<starts: array<double>, ends: array<double>>``
-        with the window starts and ends, in order; empty when a bound is null, NaN or
-        infinite, or the span is not strictly positive.
-    """
-    window_length = float(window_length)
-    max_windows = validate_max_windows(max_windows)
-
-    @F.pandas_udf("struct<starts: array<double>, ends: array<double>>")
-    def windows(start: pd.Series, stop: pd.Series) -> pd.DataFrame:
-        pairs = [
-            tile_windows(s, e, window_length, max_windows)
-            for s, e in zip(start, stop, strict=True)
-        ]
-        # object dtype keeps one array per row, also when all rows have equal window counts.
-        return pd.DataFrame(
-            {
-                "starts": pd.Series([p[0] for p in pairs], dtype=object),
-                "ends": pd.Series([p[1] for p in pairs], dtype=object),
-            }
-        )
-
-    return windows
-
-
 class TimeWindowExpression(TimeSeriesExpression):
     """Produce consecutive fixed-duration windows spanning a measurement container.
 
@@ -189,8 +147,8 @@ class TimeWindowExpression(TimeSeriesExpression):
 
     This is the query-engine counterpart of the reporting ``TimeWindowEvent``.  It evaluates
     to :class:`Intervals`, so it can scope a ``StatsAggregator`` (one statistic per window).
-    The windows come from :func:`tile_windows`, which the reporting event fact also uses (via
-    :func:`window_intervals_udf`), so both produce the same windows in the same order.
+    The windows come from :func:`tile_windows`, which the reporting event fact also uses (in
+    ``TimeWindowEvent.determine_events``), so both produce the same windows in the same order.
 
     Attributes
     ----------

src/impulse_reporting/events/time_window_event.py

@@ -4,7 +4,9 @@
 
 from collections.abc import Mapping
 
-import pyspark.sql.functions as f
+import numpy as np
+import pyarrow as pa
+import pyspark.sql.types as T
 from pyspark.sql import DataFrame, SparkSession
 
 from impulse_query_engine.analyze.metadata.time_series_expression import (
@@ -13,8 +15,8 @@
 from impulse_query_engine.analyze.query.events.time_window_expression import (
     MAX_WINDOWS_PER_CONTAINER,
     TimeWindowExpression,
+    tile_windows,
     validate_max_windows,
-    window_intervals_udf,
 )
 from impulse_query_engine.analyze.query.query_builder import QueryBuilder
 from impulse_query_engine.analyze.query.solvers.query_solver import QuerySolver
@@ -24,6 +26,15 @@
 from impulse_reporting.util.event_instance_util import generate_event_instance_id_column
 from impulse_reporting.util.report_entity_util import ReportEntityUtil
 
+# Columns of the window rows next to container_id, named as the event fact and its
+# event_instance_id hash expect them.
+_EVENT_NAME_COL = "event_name"
+_START_TS_COL = "start_ts"
+_END_TS_COL = "end_ts"
+
+# Windows after which _explode_windows emits an output batch.
+_BATCH_WINDOWS = 100_000
+
 
 class TimeWindowEvent(ContainerBoundaryEvent):
     """Event that divides each measurement container into consecutive fixed windows.
@@ -33,12 +44,12 @@ class TimeWindowEvent(ContainerBoundaryEvent):
     span with windows of length ``window_length``.  The final slice is clamped to the
     container end.
 
-    The event fact is computed from ``container_metrics`` alone (via
-    :func:`window_intervals_udf`), so every filtered container gets windows regardless of
-    its channel data.  Aggregations scoped to this event evaluate the
-    :class:`TimeWindowExpression` in the solve.  Both use the same window function
-    (``tile_windows``), so they produce identical windows, and the timestamp-based
-    ``event_instance_id`` (like for other interval events) matches on both sides.
+    The event fact is computed from ``container_metrics`` alone (in a ``mapInArrow``), so
+    every filtered container gets windows regardless of its channel data.  Aggregations
+    scoped to this event evaluate the :class:`TimeWindowExpression` in the solve.  Both use
+    the same window function (``tile_windows``), so they produce identical windows, and the
+    timestamp-based ``event_instance_id`` (like for other interval events) matches on both
+    sides.
     """
 
     def __init__(
@@ -184,7 +195,7 @@ def determine_events(
         gets windows.
         Each window becomes one event instance (``start_ts < end_ts``) whose
         ``event_instance_id`` hashes its boundaries. The solve uses the same window function
-        for scoped aggregations (see :func:`window_intervals_udf`), so the ids match.
+        for scoped aggregations (see :func:`tile_windows`), so the ids match.
 
         Parameters
         ----------
@@ -213,46 +224,127 @@ def determine_events(
         # uses for scoped aggregations (fails fast on the schema, e.g. when TIMESTAMP
         # boundaries lack solver_config.channel_time_unit).
         container_metrics_df = with_window_bounds(container_metrics_df, solver.config)
-        start_ts = f.col(solver.config.window_start_col)
-        stop_ts = f.col(solver.config.window_stop_col)
-
-        # One (event_name, windows) struct per event, exploded in a single pass over the
-        # containers.
-        per_event = f.array(
-            *[
-                f.struct(
-                    f.lit(event.get_name()).alias("event_name"),
-                    window_intervals_udf(
-                        event.window_length, max_windows=event.max_windows_per_container
-                    )(start_ts, stop_ts).alias("windows"),
-                )
+        windows_df = _explode_windows(
+            container_metrics_df,
+            id_col=solver.config.container_id_col,
+            start_col=solver.config.window_start_col,
+            stop_col=solver.config.window_stop_col,
+            windows=[
+                (event.get_name(), event.window_length, event.max_windows_per_container)
                 for event in events
-            ]
+            ],
         )
-
-        df = (
-            container_metrics_df.select(
-                f.col(solver.config.container_id_col).alias("container_id"),
-                f.explode(per_event).alias("event"),
-            )
-            .select(
-                "container_id",
-                f.col("event.event_name").alias("event_name"),
-                f.inline(
-                    f.arrays_zip(
-                        f.col("event.windows.starts").alias("start_ts"),
-                        f.col("event.windows.ends").alias("end_ts"),
-                    )
-                ),
-            )
-            .withColumn(
+        return (
+            windows_df.withColumn(
                 "event_instance_id",
-                generate_event_instance_id_column(event_type=TimeWindowEvent),
+                generate_event_instance_id_column(
+                    event_type=TimeWindowEvent,
+                    event_name_col=_EVENT_NAME_COL,
+                    start_ts_col=_START_TS_COL,
+                    end_ts_col=_END_TS_COL,
+                ),
             )
             .withColumn(
                 "event_id",
-                ReportEntityUtil.get_event_id_column(elements=events, element_name="event_name"),
+                ReportEntityUtil.get_event_id_column(
+                    elements=events, element_name=_EVENT_NAME_COL
+                ),
             )
             .select(EVENT_INSTANCE_FACT_SCHEMA.fieldNames())
         )
-        return df
+
+
+def _explode_windows(
+    df: DataFrame,
+    *,
+    id_col: str,
+    start_col: str,
+    stop_col: str,
+    windows: list[tuple[str, float, int]],
+) -> DataFrame:
+    """One row per window of each container, for all events in one pass.
+
+    Each container's windows come from :func:`tile_windows`, like those of the solve.  The
+    rows are built from numpy arrays in a ``mapInArrow``, in output batches of about 100,000
+    windows that never split a container's windows of one event, so the Python worker's
+    memory stays bounded by the batch size and ``max_windows``, independent of the number of
+    containers and events.
+
+    Parameters
+    ----------
+    df : pyspark.sql.DataFrame
+        One row per container, with *id_col* and the bounds in the channel time frame.
+    id_col : str
+        Container id column; kept with its name and type (e.g. long or string) in the output.
+    start_col, stop_col : str
+        Container bound columns (numeric).
+    windows : list of tuple
+        ``(name, window_length, max_windows)`` per event: the event name of its rows, the
+        window length (in the unit of the bounds) and the maximum number of windows per
+        container. A container exceeding it fails the query with an error naming the limit.
+
+    Returns
+    -------
+    pyspark.sql.DataFrame
+        Columns *id_col*, ``event_name`` (string), ``start_ts`` and ``end_ts`` (double); no
+        rows for a container whose bound is null, NaN or infinite, or whose span is not
+        strictly positive.
+    """
+    schema = T.StructType(
+        [
+            df.schema[id_col],
+            T.StructField(_EVENT_NAME_COL, T.StringType()),
+            T.StructField(_START_TS_COL, T.DoubleType()),
+            T.StructField(_END_TS_COL, T.DoubleType()),
+        ]
+    )
+    column_names = schema.names
+
+    def tile(batches):
+        yield from _window_batches(
+            batches, id_col, start_col, stop_col, windows, column_names, _BATCH_WINDOWS
+        )
+
+    return df.select(id_col, start_col, stop_col).mapInArrow(tile, schema)
+
+
+def _window_batches(batches, id_col, start_col, stop_col, windows, column_names, batch_windows):
+    """Yield :func:`_explode_windows` output batches named *column_names* for Arrow *batches*.
+
+    A batch is emitted once it holds at least *batch_windows* windows, after a container's
+    complete windows of one event, and at the end of each input batch.  A batch therefore
+    holds at most ``batch_windows - 1`` plus one event's ``max_windows`` windows.
+    """
+    event_names = pa.array([name for name, _, _ in windows], pa.string())
+    for batch in batches:
+        ids = batch.column(id_col)
+        bounds = zip(
+            batch.column(start_col).to_pylist(), batch.column(stop_col).to_pylist(), strict=True
+        )
+        parts, pending = [], 0
+        for row, (start, stop) in enumerate(bounds):
+            for event, (_, length, max_windows) in enumerate(windows):
+                starts, ends = tile_windows(start, stop, length, max_windows)
+                parts.append((row, event, starts, ends))
+                pending += len(starts)
+                if pending >= batch_windows:
+                    yield _record_batch(column_names, ids, event_names, parts)
+                    parts, pending = [], 0
+        if pending:
+            yield _record_batch(column_names, ids, event_names, parts)
+
+
+def _record_batch(column_names, ids, event_names, parts) -> pa.RecordBatch:
+    """Output batch for the ``(row, event, starts, ends)`` *parts* of one input batch."""
+    counts = [len(starts) for _, _, starts, _ in parts]
+    rows = np.repeat([row for row, _, _, _ in parts], counts)
+    events = np.repeat([event for _, event, _, _ in parts], counts)
+    return pa.RecordBatch.from_arrays(
+        [
+            ids.take(rows),
+            event_names.take(events),
+            pa.array(np.concatenate([starts for _, _, starts, _ in parts])),
+            pa.array(np.concatenate([ends for _, _, _, ends in parts])),
+        ],
+        names=column_names,
+    )