Parse messages into Arrow Record Batches

sentry_streams/Cargo.toml

@@ -23,6 +23,11 @@ metrics = "0.24.0"
 metrics-exporter-dogstatsd = "0.9.0"
 rand = "0.9"
 sentry = "0.48.1"
+arrow = { version = "59", features = ["ffi"] }
+prost = "0.14"
+prost-types = "0.14"
+sentry_protos = "0.70"
+sentry-kafka-schemas = { version = "3", default-features = false }
 
 [lib]
 name = "rust_streams"

sentry_streams/docs/source/arrow_batch_parser.rst

@@ -0,0 +1,158 @@
+Arrow Batch Parser
+==================
+
+``ArrowBatchParser`` batches raw Kafka payloads and decodes them into an Apache
+Arrow ``RecordBatch`` entirely in Rust. The batch is handed to Python as an
+`Arrow IPC stream <https://arrow.apache.org/docs/format/Columnar.html#serialization-and-interprocess-communication-ipc>`_
+-- an ordinary ``bytes`` payload, which every existing step already understands.
+
+.. code-block:: python
+
+   import polars as pl
+
+   from sentry_streams.pipeline import ArrowBatchParser, StreamSink, streaming_source
+   from sentry_streams.pipeline.pipeline import Map
+
+
+   def summarize(msg):
+       df = pl.read_ipc_stream(msg.payload)
+       return f"{df.height} rows".encode()
+
+
+   pipeline = (
+       streaming_source(name="myinput", stream_name="snuba-items")
+       .apply(
+           ArrowBatchParser(
+               name="parse_arrow", schema_name="snuba-items", batch_size=1000
+           )
+       )
+       .apply(Map(name="summarize", function=summarize))
+       .sink(StreamSink[bytes](name="mysink", stream_name="transformed-events"))
+   )
+
+It replaces ``Batch`` → ``Map(extract_bytes)`` → ``BatchParser``, in which every
+message is copied into Python memory as ``bytes`` and decoded under the GIL.
+Here the individual messages are never turned into Python objects; only the
+assembled batch crosses over.
+
+.. note::
+
+   Serializing the batch and copying it into Python memory is a deliberate
+   simplification. Handing the ``RecordBatch`` across directly, through the Arrow
+   C data interface, removes both copies and is deferred rather than ruled out.
+
+Windowing is configured exactly like :class:`Batch`, by ``batch_size`` and/or
+``batch_timedelta``, both overridable from ``steps_config``.
+
+``schema_name`` is the logical stream name whose ``sentry-kafka-schemas`` entry
+names the message type, and so selects the extractor -- normally the source's
+``stream_name``. It is declared on the step rather than inferred from the source,
+so the step does not depend on where it sits and a deployment topic override
+cannot change which extractor is used.
+
+The schema is hardcoded
+-----------------------
+
+There is nothing to configure. Each supported message type has a hand-written
+extractor in Rust that owns its Arrow schema, resolved from the source topic's
+entry in ``sentry-kafka-schemas``. Several topics sharing a schema share one
+extractor.
+
+**Adding a column, or a new message type, is a Rust change and a release.** This
+is the deliberate trade of the current implementation: Rust has no equivalent of
+Python's reflective ``ProtobufCodec``, and carrying a protobuf descriptor pool to
+get one was judged not worth it yet. See
+``docs/design/arrow-batch-parser.md``.
+
+Currently implemented: ``sentry_protos.snuba.v1.TraceItem`` (topic
+``snuba-items``).
+
+How it can fail
+---------------
+
+Every failure is loud, and most happen at startup rather than in production:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Situation
+     - Result
+   * - Topic has no schema, or a JSON/msgpack schema
+     - Panic at startup, naming the topic and the schema type found
+   * - Message type has no extractor
+     - Panic at startup, listing the supported message types
+   * - Step placed after a ``Map`` or other converting step
+     - Type error under mypy; **panic** on the first batch if types were bypassed
+   * - Used with the pure-Python Arroyo adapter
+     - ``NotImplementedError``
+   * - A payload fails to decode
+     - **Panic, failing the process**
+
+That last one deserves emphasis. Batching collapses offsets to the maximum per
+partition, so there is no ``(partition, offset)`` for Arroyo's DLQ to reject a
+single row with: one bad payload fails the whole batch, and the consumer will
+re-read the same offsets on restart. This matches what the runtime already does
+for an error on an ``AnyMessage``, but it means a malformed message is an outage,
+not a dead letter. Per-row offset tracking is deferred work.
+
+Placement
+---------
+
+The step reads bytes off the wire, so it must come before any step that turns
+messages into Python objects. Filters, broadcasts and routers are fine — they
+forward messages untouched. A ``Map`` in between is a build-time error.
+
+Attribute columns
+-----------------
+
+``TraceItem``'s ``map<string, AnyValue>`` is split by value type into
+``attr_str``, ``attr_int``, ``attr_double``, ``attr_bool`` and ``attr_bytes``.
+Arrow has no usable union type here and Snuba EAP splits the same way. Two
+consequences:
+
+* an attribute that changes type between messages lands in different columns;
+* recursive values (``ArrayValue``, ``KeyValueList``) are **dropped**. Arrow has
+  no recursive type. ``AnyValue`` is a port of OpenTelemetry's type, so these
+  arms exist because OTel has them rather than because Sentry ingestion uses
+  them, and they do not appear in the canonical ``snuba-items`` example. If one
+  does arrive, that attribute is skipped -- the rest of its row is unaffected --
+  and nothing is logged. If they turn out to occur in practice they should get
+  their own column rather than being flattened into ``attr_str``, where an
+  encoded array would be indistinguishable from a string that looks like one.
+
+Performance
+-----------
+
+Measured with ``cargo test --release bench_arrow_vs_pylist -- --ignored
+--nocapture``, 10 000 rows per window, ``TraceItem`` with one string attribute:
+
+.. list-table::
+   :header-rows: 1
+
+   * - Path
+     - p50
+     - p99
+   * - ``ArrowBatchParser``
+     - 3.6 ms
+     - 4.8 ms
+   * - ``Batch`` → ``BatchParser``
+     - 4.2 ms
+     - 6.3 ms
+
+About 20% faster at the median and 35% at the tail — a real but modest win on
+this workload, not the order of magnitude the "no Python round trip" framing
+might suggest. Two things to keep in mind when reading it:
+
+* The comparison stops at decoded values. The Python path then has to *build*
+  something columnar (polars, parquet) from those objects, which this step has
+  already done.
+* The source still hands Rust its payloads inside Python-owned ``RawMessage``
+  objects, so decoding holds the GIL and pays a copy *per message*. Removing that
+  is work owned elsewhere; the numbers here understate what the step can do once
+  it lands, and the benchmark should be re-run then.
+
+At these figures a 1 000-row window decodes in well under a millisecond, so
+running inline on the consumer thread is comfortably within
+``max_poll_interval_ms``. Moving decoding to a threadpool is deferred until a
+benchmark justifies it.

sentry_streams/docs/source/index.rst

@@ -11,5 +11,6 @@
    build_pipeline
    configure_pipeline
    runtime/arroyo
+   arrow_batch_parser
    deployment
    rust

sentry_streams/sentry_streams/adapters/arroyo/adapter.py

@@ -44,6 +44,7 @@
     OutputType,
 )
 from sentry_streams.pipeline.pipeline import (
+    ArrowBatchParser,
     Broadcast,
     ComplexStep,
     Filter,
@@ -272,6 +273,13 @@ def reduce(
             stream.source in self.__consumers
         ), f"Stream starting at source {stream.source} not found when adding a reduce"
 
+        if isinstance(step, ArrowBatchParser):
+            raise NotImplementedError(
+                "ArrowBatchParser is only supported by the Rust Arroyo adapter (rust_arroyo), "
+                "not the pure Python Arroyo adapter. It decodes raw payloads into an Arrow "
+                "RecordBatch in Rust and has no Python implementation."
+            )
+
         self.__consumers[stream.source].add_step(ReduceStep(route=stream, pipeline_step=step))
         return stream
 

sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py

@@ -45,6 +45,7 @@
 )
 from sentry_streams.pipeline.message import Message
 from sentry_streams.pipeline.pipeline import (
+    ArrowBatchParser,
     Batch,
     Broadcast,
     ComplexStep,
@@ -498,6 +499,23 @@ def reduce(
         step.override_config(loaded_config)
         step.validate()
 
+        if isinstance(step, ArrowBatchParser):
+            logger.info(f"Adding Arrow batch parser (native): {step.name} to pipeline")
+            self.__consumers[stream.source].add_step(
+                RuntimeOperator.ArrowBatchParser(
+                    route=route,
+                    step_name=step.name,
+                    schema_name=step.schema_name,
+                    max_batch_size=step.batch_size,
+                    max_batch_time_ms=(
+                        step.batch_timedelta.total_seconds() * 1000.0
+                        if step.batch_timedelta is not None
+                        else None
+                    ),
+                )
+            )
+            return stream
+
         if isinstance(step, Batch):
             max_batch_time_ms: float | None
             if step.batch_timedelta is not None:

sentry_streams/sentry_streams/deployment_config/arrow_trace_items.yaml

@@ -0,0 +1,14 @@
+env: {}
+
+pipeline:
+  segments:
+    - steps_config:
+        myinput:
+          starts_segment: True
+          bootstrap_servers: ["127.0.0.1:9092"]
+        parse_arrow:
+          batch_size: 1000
+          batch_timedelta:
+            seconds: 5
+        mysink:
+          bootstrap_servers: ["127.0.0.1:9092"]

sentry_streams/sentry_streams/examples/arrow_trace_items.py

@@ -0,0 +1,48 @@
+"""Decode a batch of TraceItem protobufs into an Apache Arrow RecordBatch in Rust.
+
+``ArrowBatchParser`` fuses batching and decoding: it reads the raw Kafka payloads
+without turning the individual messages into Python objects, and hands the batch
+over as an Arrow IPC stream -- ordinary ``bytes``, read here with polars.
+
+Compare with ``parquet_serializer.py``, which does the same job as
+``Batch`` -> ``Map(extract_bytes)`` -> ``BatchParser`` and pays a Python round
+trip per message.
+
+Run with::
+
+    python -m sentry_streams.runner \
+        --name arrow-trace-items \
+        --config deployment_config/arrow_trace_items.yaml \
+        sentry_streams/examples/arrow_trace_items.py
+"""
+
+import polars as pl
+
+from sentry_streams.pipeline import ArrowBatchParser, StreamSink, streaming_source
+from sentry_streams.pipeline.message import Message
+from sentry_streams.pipeline.pipeline import Map
+
+TOPIC = "snuba-items"
+
+
+def summarize(msg: Message[bytes]) -> bytes:
+    """Read the Arrow batch through polars and emit a one-line summary.
+
+    The payload is an Arrow IPC stream; polars reads all of it at once, with no
+    row-by-row conversion.
+    """
+    df = pl.read_ipc_stream(msg.payload)
+
+    by_type = df.group_by("item_type").agg(pl.len().alias("rows")).sort("rows", descending=True)
+    summary = ", ".join(f"{row[0]}={row[1]}" for row in by_type.iter_rows())
+    return f"{df.height} trace items ({summary})".encode()
+
+
+pipeline = (
+    streaming_source(name="myinput", stream_name=TOPIC)
+    # Takes bytes, so it must come before any step that turns messages into
+    # Python objects. Placing it later is a type error.
+    .apply(ArrowBatchParser(name="parse_arrow", schema_name=TOPIC, batch_size=1000))
+    .apply(Map(name="summarize", function=summarize))
+    .sink(StreamSink[bytes](name="mysink", stream_name="transformed-events"))
+)

sentry_streams/sentry_streams/pipeline/__init__.py

@@ -1,4 +1,5 @@
 from sentry_streams.pipeline.pipeline import (
+    ArrowBatchParser,
     Batch,
     BatchParser,
     Filter,
@@ -16,6 +17,7 @@
 )
 
 __all__ = [
+    "ArrowBatchParser",
     "Batch",
     "BatchParser",
     "Filter",

sentry_streams/sentry_streams/pipeline/pipeline.py

@@ -681,6 +681,85 @@ def override_config(self, loaded_config: Mapping[str, Any]) -> None:
             self.batch_timedelta = timedelta(**loaded_kwargs)
 
 
+@dataclass
+class ArrowBatchParser(
+    Reduce[MeasurementUnit, bytes, bytes],
+    Generic[MeasurementUnit],
+):
+    """
+    Batches raw Kafka payloads and decodes them into an Apache Arrow
+    ``RecordBatch``, entirely in Rust.
+
+    The emitted message payload is the batch serialized as an Arrow IPC stream,
+    so downstream steps receive ordinary ``bytes``::
+
+        import polars as pl
+
+        def to_frame(msg: Message[bytes]) -> pl.DataFrame:
+            return pl.read_ipc_stream(msg.payload)
+
+    This is the fused equivalent of ``Batch`` -> ``Map(extract_bytes)`` ->
+    ``BatchParser``: the individual messages are never turned into Python
+    objects. The batch itself is serialized and copied into Python memory, which
+    is a deliberate simplification -- handing the ``RecordBatch`` over directly
+    through the Arrow C data interface is deferred.
+
+    ``schema_name`` is the logical stream name whose ``sentry-kafka-schemas``
+    entry names the message type, and so selects the extractor -- usually the
+    source's ``stream_name``. It is given explicitly rather than inferred so the
+    step does not depend on where it sits in the pipeline.
+
+    Limitations of the current implementation, all of which fail loudly:
+
+    * **Protobuf topics only.** A JSON or msgpack topic raises at startup.
+    * **The Arrow schema is hardcoded in Rust**, per message type, so there is
+      nothing to configure here and adding a column needs a release.
+    * **Rust adapter only.** The pure-Python Arroyo adapter raises
+      ``NotImplementedError``.
+    * **It reads raw payloads**, so it takes ``bytes`` and must come before any
+      step that converts messages into Python objects. Placing it after one is a
+      type error, and a panic at runtime if the types were bypassed.
+    * A payload that fails to decode **fails the process**: batching collapses
+      offsets, so there is no single offset to dead-letter.
+
+    Configured by batch size and/or batch_timedelta exactly like ``Batch``, and
+    both are overridable from the deployment config's ``steps_config``.
+    """
+
+    schema_name: str
+    batch_size: int | None = None
+    batch_timedelta: timedelta | None = timedelta(seconds=10)
+    step_type: StepType = StepType.REDUCE
+
+    def validate(self) -> None:
+        """Validate that at least one of batch_size or batch_timedelta is set."""
+        if self.batch_size is None and self.batch_timedelta is None:
+            raise ValueError("At least one of batch_size or batch_timedelta must be set.")
+
+    @property
+    def group_by(self) -> Optional[GroupBy]:
+        return None
+
+    @property
+    def windowing(self) -> Window[MeasurementUnit]:
+        return TumblingWindow(self.batch_size, self.batch_timedelta)
+
+    @property
+    def aggregate_fn(self) -> Callable[[], Accumulator[Message[bytes], bytes]]:
+        raise NotImplementedError(
+            "ArrowBatchParser is implemented natively in Rust and has no Python accumulator."
+        )
+
+    def override_config(self, loaded_config: Mapping[str, Any]) -> None:
+        if loaded_config.get("batch_size") is not None:
+            self.batch_size = loaded_config.get("batch_size")
+
+        if loaded_config.get("batch_timedelta") is not None:
+            loaded_kwargs = loaded_config.get("batch_timedelta")
+            assert isinstance(loaded_kwargs, Mapping)
+            self.batch_timedelta = timedelta(**loaded_kwargs)
+
+
 @dataclass
 class FlatMap(Transform[TIn, TOut], Generic[TIn, TOut]):
     """