From e8e085e98cfdee76512a7abb6c05d34884109a33 Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 15:33:07 -0700 Subject: [PATCH 01/13] docs(arrow): design plan for the Rust Arrow batch parser Records the agreed design for a fused batch+decode pipeline primitive that turns raw Kafka payloads into an Apache Arrow RecordBatch entirely in Rust, handed to Python over the Arrow PyCapsule interface. PoC scope: protobuf only, hand-written extractor per message type with a hardcoded Arrow schema, TraceItem only. Captures the decisions, the runtime constraints that drove them, the accepted consequences, and a six-phase implementation plan. Co-Authored-By: Claude Opus 5 (1M context) --- .../docs/design/arrow-batch-parser.md | 459 ++++++++++++++++++ 1 file changed, 459 insertions(+) create mode 100644 sentry_streams/docs/design/arrow-batch-parser.md diff --git a/sentry_streams/docs/design/arrow-batch-parser.md b/sentry_streams/docs/design/arrow-batch-parser.md new file mode 100644 index 00000000..f858163c --- /dev/null +++ b/sentry_streams/docs/design/arrow-batch-parser.md @@ -0,0 +1,459 @@ +# Arrow Batch Parser — PoC Implementation Plan + +**Status:** design agreed, not yet implemented +**Scope:** proof of concept — protobuf only, `sentry_protos.snuba.v1.TraceItem` only + +## Goal + +A pipeline primitive that batches raw Kafka payloads and decodes them into an Apache +Arrow `RecordBatch` entirely in Rust, with no per-message round trip through Python. +The result is handed to Python as a `RoutedValuePayload::PyStreamingMessage`. + +It replaces the Python path `Batch` → `Map(extract_bytes)` → `BatchParser`, in which +every message is copied into Python memory as `bytes` and decoded by +`sentry_kafka_schemas` under the GIL. + +### PoC scope + +Protobuf only, hand-written extractor per message type, Arrow schema hardcoded in the +extractor. Only `TraceItem` implemented. The Python DSL declares no schema. + +### Non-goals + +JSON and msgpack topics (a JSON topic is a startup panic); descriptor-driven generic +extraction; per-row dead-lettering; `TraceItem.outcomes`; replacing the Python +`BatchParser`, which keeps working on all adapters. + +## Decisions + +| # | Decision | +|---|---| +| 1 | **Fused step** — batching and decoding in one primitive, not a parser after `Batch`. | +| 2 | Output is a `#[pyclass]` implementing the **Arrow PyCapsule interface**. No `pyarrow` dependency. | +| 4 | New primitive; Python `BatchParser` untouched. | +| 5 | Decoding runs **inline** on the consumer thread. | +| 9 | Offsets collapse to `max` per partition, as `batch_step.rs` does today. | +| 10/17 | Input contract **`RawMessage` only** — build-time check in the adapter, runtime backstop. | +| 11 | **Generalize `BatchStep`** over a flush-producer trait rather than forking it. | +| 12 | `Reduce` subclass, `StepType.REDUCE`, `isinstance` branch in `reduce()`. | +| 14 | Failure is `panic!`, matching `transformer.rs:45-46`. | +| 22 | `map` → **type-split maps** `attr_str/int/double/bool/bytes`. | +| 23 | **Protobuf only.** | +| 24 | `sentry-protos` for types and `prost` decode; `sentry-kafka-schemas` (`default-features = false`) for topic → schema. | +| 25 | Extractors indexed by the **raw resource string**; several topics sharing a schema share one extractor. | +| 26 | Resolution and validation at **step construction**; failures panic at startup. | +| 27 | `TraceItem` schema: all fields except `outcomes`. | + +**Why no descriptor pool.** An earlier iteration used `prost-reflect` + `DescriptorPool` +for `get_field_by_name` — Rust's only equivalent of what Python's `ProtobufCodec` gets +free from protobuf's reflective runtime. Dropped for the PoC. The trade is explicit: +**adding a column requires a Rust change and a release here.** See *Deferred*. + +**Why `default-features = false`.** `validate_protobuf` is the only item behind +`type_generation`, and we decode with `TraceItem::decode` directly. `get_schema`, +`schema_type` and `raw_schema` sit outside it. Disabling drops `typify`, `syn`, +`prettyplease`, `schemars` and the crate's own `prost`/`sentry_protos` pins. + +## Integration points in existing code + +Constraints discovered by reading the runtime. These drive several choices below. + +| Fact | Location | Consequence | +|---|---|---| +| `operators::build()` receives no topic or schema. `build_chain` has `schema` but does not pass it down. | `src/operators.rs:125`, `src/consumer.rs` | The schema name must be **carried in the `RuntimeOperator` variant**, supplied by the Python adapter. Do not widen `build()`. | +| The adapter captures `schema_name = step.stream_name` *before* `override_config` can change the topic. | `rust_arroyo.py:303` | Schema lookup survives deployment topic overrides. The adapter must stash it per source so `reduce()` can read it. | +| Source wraps every payload in a `Py` immediately. | `consumer.rs::to_routed_value` | Payload bytes live in **Python-owned memory**; reading them needs the GIL. **Temporary** — see *Assumed future work*. | +| `Batch` stores `committable` (`offset+1`), collapsed to `max` per partition. | `batch_step.rs:104` | No per-row offsets → no DLQ. Consequence 1. | +| A non-`InvalidMessageError` Python exception, or any error on an `AnyMessage`, panics. | `transformer.rs:45-46` | Panicking is consistent with the runtime's existing behaviour, not a new failure mode. | +| `StrategyError::InvalidMessage` makes arroyo **continue**; `StrategyError::Other` stops the consumer. | `processing/mod.rs:321-349` | Neither is usable for a batch with collapsed offsets; hence `panic!`. | + +## `TraceItem` Arrow schema + +| Column | Arrow type | Null | Source | +|---|---|---|---| +| `organization_id` | `UInt64` | no | 1 | +| `project_id` | `UInt64` | no | 2 | +| `trace_id` | `Utf8` | no | 3 | +| `item_id` | `Binary` | no | 4 (bytes, little endian) | +| `item_type` | `Utf8` | no | 5, enum **name** via `as_str_name()` | +| `timestamp` | `Timestamp(us, "UTC")` | **yes** | 6 | +| `client_sample_rate` | `Float64` | no | 8 | +| `server_sample_rate` | `Float64` | no | 9 | +| `conversation_id` | `Utf8` | yes | 10, proto3 `optional` | +| `session_id` | `Utf8` | yes | 11, proto3 `optional` | +| `retention_days` | `UInt32` | no | 100 | +| `received` | `Timestamp(us, "UTC")` | **yes** | 101 | +| `downsampled_retention_days` | `UInt32` | no | 102 | +| `attr_str` | `Map` | no | 7, `AnyValue` arm 1 (+ 5, 6 JSON-encoded) | +| `attr_int` | `Map` | no | 7, arm 3 | +| `attr_double` | `Map` | no | 7, arm 4 | +| `attr_bool` | `Map` | no | 7, arm 2 | +| `attr_bytes` | `Map` | no | 7, arm 7 | + +**Nullability follows protobuf presence, not the `optional` keyword.** Message-typed +fields always have explicit presence in proto3, so prost yields `Option` +for *both* `timestamp` and `received` — both nullable. Implicit-presence scalars cannot +distinguish unset from zero, so they are non-nullable columns carrying the default. + +`ArrayValue` (arm 5) and `KeyValueList` (arm 6) are recursive; Arrow has no recursive +types, so they are JSON-encoded into `attr_str`. + +## Dependencies + +```toml +arrow = { version = "59", features = ["ffi"] } +prost = "0.14" +sentry_protos = "0.70" +sentry-kafka-schemas = { version = "3", default-features = false } +``` + +No new Python dependencies; `pyarrow` is deliberately not added. + +--- + +# Phases + +Phases 1 and 2 are independent and may run in parallel. 3 depends on 0; 4 on 1+2+3; +5 on 4; 6 on 5. + +## Phase 0 — Dependencies + +1. Add the four crates above. +2. **Verify `default-features = false` still exposes** `get_schema`, `Schema::schema_type` + and `Schema::raw_schema`. If `raw_schema` turns out to be gated, fall back to default + features and record the build-time cost in this document. +3. `cargo build`, `cargo test` green with no code changes. + +**Acceptance:** clean build; a throwaway test asserting +`get_schema("snuba-items", None).unwrap().raw_schema()` equals +`"sentry_protos.snuba.v1.trace_item_pb2.TraceItem"`. + +## Phase 1 — `PyRecordBatch` (Arrow → Python) + +**File:** `src/py_record_batch.rs` *(new)*; register in `src/lib.rs`; stubs in +`sentry_streams/rust_streams.pyi`. + +```rust +#[pyclass(name = "ArrowRecordBatch", module = "sentry_streams.rust_streams")] +pub struct PyRecordBatch { pub(crate) batch: RecordBatch } + +#[pymethods] +impl PyRecordBatch { + #[getter] fn num_rows(&self) -> usize; + #[getter] fn num_columns(&self) -> usize; + fn __repr__(&self) -> String; + + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_array__<'py>(&self, py: Python<'py>, requested_schema: Option>) + -> PyResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)>; + + fn __arrow_c_schema__<'py>(&self, py: Python<'py>) -> PyResult>; + + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_stream__<'py>(&self, py: Python<'py>, requested_schema: Option>) + -> PyResult>; +} +``` + +**Implement all three, not just `__arrow_c_array__`.** Table-level consumers — +`pl.DataFrame(obj)`, `pa.table(obj)` — look for `__arrow_c_stream__`; array-level +consumers use `__arrow_c_array__`. Implementing only one makes the object work in some +call sites and not others. + +Mechanics: + +- Array: `StructArray::from(batch.clone())` → `arrow::ffi::to_ffi(&struct_array.to_data())` + → two capsules. +- Stream: `FFI_ArrowArrayStream::new(Box::new(RecordBatchIterator::new(...)))`, one batch. +- **Capsule names must be exactly** `arrow_schema`, `arrow_array`, `arrow_array_stream`, + as NUL-terminated `CString`. A wrong name fails at the consumer with an opaque error. +- `PyCapsule::new` takes ownership; `FFI_ArrowSchema`/`FFI_ArrowArray`'s `Drop` invokes + the C release callback, so no manual destructor is needed. +- `requested_schema` is accepted and **ignored** — the protocol permits returning the + native schema when a cast is unsupported. Document it in the docstring. + +**Tests** + +| Test | Assertion | +|---|---| +| `polars.DataFrame(rb)` | values, column names, dtypes match | +| `pyarrow.record_batch(rb)` *(dev-dep only)* | round-trips `__arrow_c_array__` | +| `pyarrow.table(rb)` | round-trips `__arrow_c_stream__` | +| consume twice | second call still yields a valid batch (no double-release) | +| nested `Map` column | survives the FFI boundary | + +**Acceptance:** a Rust-built `RecordBatch` reaches polars with correct values and +schema. pyarrow may be a dev-only dependency; it must not enter runtime deps. + +## Phase 2 — Generalize `BatchStep` (pure refactor) + +**File:** `src/batch_step.rs`. + +```rust +pub(crate) trait BatchFlushProducer: Send + Sync { + fn produce( + &self, + route: &Route, + elements: &[PyStreamingMessage], + committable: BTreeMap, + ) -> Result, StrategyError>; +} + +/// Existing behaviour, lifted verbatim out of `Batch::flush`. +pub(crate) struct PyListFlushProducer; +``` + +1. Move the body of `Batch::flush` (`batch_step.rs:166-199`) into + `PyListFlushProducer::produce`. Do not change it. +2. `Batch` gains `producer: Arc`; `flush()` delegates. +3. `BatchStep::new` takes the producer; `build_batch_step` keeps its signature and + passes `PyListFlushProducer`. +4. Leave untouched: watermark buffering, `pending_batch`, `drain_outbound`, + `record_rejected_submit`, `oldest_batch_row_timestamp`, `join`. + +**Acceptance:** **no behaviour change.** The existing `batch_step.rs` test module +passes unmodified — not adapted, unmodified. Lands as its own commit so any regression +is attributable. + +## Phase 3 — Extractor module + +**Files:** `src/extractors/mod.rs`, `src/extractors/trace_item.rs` *(new)*. + +```rust +#[derive(Debug)] +pub enum ExtractorError { + Decode { index: usize, source: prost::DecodeError }, + Build(arrow::error::ArrowError), + Field { field: &'static str, detail: String }, +} + +pub trait Extractor: Send + Sync { + fn resource(&self) -> &'static str; + fn schema(&self) -> SchemaRef; + fn extract(&self, payloads: &[&[u8]]) -> Result; +} + +pub fn get_extractor(resource: &str) -> Option<&'static dyn Extractor>; +pub fn registered_resources() -> Vec<&'static str>; +``` + +`extract` is **batch-wise**: one Arrow builder per column, fed across all rows, finished +once. Not row-at-a-time — that is the shape Arrow builders want and it keeps a second +extractor down to one file plus one registry line. + +### `TraceItem` extractor — implementation notes + +- **`MapBuilder` field names.** arrow-rs defaults to `entries`/`keys`/`values`; the Arrow + spec is `entries`/`key`/`value`. Set `MapFieldNames` explicitly to the spec names or + pyarrow/polars interop breaks in confusing ways. +- **`append(true)` per row on every map builder**, including rows with no attributes of + that type. Skipping it silently misaligns every subsequent row. +- **Sort attribute keys per row.** prost decodes `map` into a + `HashMap` with nondeterministic iteration order; unsorted output makes batches + irreproducible and tests flaky. +- **Unknown enum values must not panic.** `TraceItemType::try_from(i32)` fails on a value + from a newer producer; render `TYPE_UNKNOWN_`. Enum additions are routine + forward-compatible producer changes and crashing on them would be a self-inflicted + outage. *(Flagged — override if you want strictness.)* +- **Timestamps:** `prost_types::Timestamp { seconds, nanos }` → + `seconds * 1_000_000 + nanos / 1_000`, with `checked_mul`/`checked_add` and an + `ExtractorError::Field` on overflow rather than a silent wrap. +- Builders: `UInt64Builder`, `StringBuilder`, `BinaryBuilder`, + `TimestampMicrosecondBuilder::new().with_timezone("UTC")`, `Float64Builder`, + `UInt32Builder`, and five `MapBuilder`. + +**Tests** (`src/extractors/trace_item.rs`, `mod tests`) — build `TraceItem` values +in-process with prost, encode, extract, assert: + +| Case | Expectation | +|---|---| +| all scalar fields populated | every column matches | +| each `AnyValue` arm (string, bool, int, double, bytes) | lands in its own `attr_*` map | +| `ArrayValue` / `KeyValueList` | JSON-encoded into `attr_str` | +| absent `timestamp` / `received` | null, not epoch zero | +| absent `conversation_id` / `session_id` | null | +| unset implicit-presence scalars | zero/empty, non-null | +| unknown `item_type` number | `TYPE_UNKNOWN_`, no panic | +| same attributes, different insertion order | byte-identical `RecordBatch` | +| rows with 0, 1, many attributes mixed in one batch | map offsets aligned | +| truncated payload | `ExtractorError::Decode { index }` naming the row | +| `Timestamp` at i64 boundary | `ExtractorError::Field`, no wrap | + +**Acceptance:** the table above passes; `extract` on an empty slice yields a +zero-row batch with the correct schema. + +## Phase 4 — The parser step + +**File:** `src/arrow_batch_parser.rs` *(new)*. + +```rust +pub(crate) struct ArrowFlushProducer { + extractor: &'static dyn Extractor, + step_name: String, +} +impl BatchFlushProducer for ArrowFlushProducer { /* ... */ } + +pub fn build_arrow_batch_parser_step( + route: &Route, + schema_name: &str, + step_name: String, + max_batch_size: Option, + max_batch_time: Option, + next: Box>, +) -> Box>; +``` + +Construction-time resolution, each failure a `panic!` naming step, topic and cause: + +1. `get_schema(schema_name, None)` — panic if the topic is unknown. +2. `schema.schema_type == SchemaType::Protobuf` — panic otherwise, naming the actual type. +3. `schema.raw_schema()` → resource string. +4. `get_extractor(resource)` — panic if absent, listing `registered_resources()`. + +`get_schema` runs **exactly once here.** It `Box::leak`s on protobuf topics and must +never be called per message. + +`produce` then: + +1. Calls `with_payloads` (below) to obtain `&[&[u8]]`, rejecting any `PyAnyMessage` + element with a panic naming the step (the decision-10 runtime backstop). +2. Calls `extractor.extract` inside that scope. +3. Wraps the `RecordBatch` in `PyRecordBatch`, then a `PyAnyMessage` with + `schema = Some(schema_name)` and the flush timestamp, exactly as + `PyListFlushProducer` does. +5. Emits `RoutedValuePayload::PyStreamingMessage`. + +### The payload seam — get this shape right the first time + +`consumer.rs::to_routed_value` boxes every payload into a `Py` at the +source, so payload bytes currently live in Python-owned memory and reading them needs +the GIL. Extraction therefore holds the GIL for the duration of a batch decode. That is +**temporary** (see *Assumed future work*), and the code must be written so removing it +is a deletion rather than a refactor. + +```rust +/// The only place that knows where payload bytes live. +fn with_payloads(elements: &[BatchElement], f: impl FnOnce(&[&[u8]]) -> R) -> R; +``` + +Today: acquire the GIL, collect `PyRef` guards, borrow `&[u8]` from each, +call `f`. Once the source emits native messages: collect the slices and call `f`. The +call site does not change. It is a scope rather than a plain function because the +`PyRef` guards must outlive the slices. + +**Do not copy payloads out of Python memory to avoid holding the GIL.** Collecting +`Vec>` and releasing the GIL would work today and would become permanent dead +weight the moment the source goes native — a per-message copy in the one step whose +purpose is to eliminate per-message copies. + +`Extractor::extract` takes `&[&[u8]]` precisely so that it, and every test in phase 3, +is indifferent to where the bytes live. The **output** side is unaffected either way: +the result must remain a `PyStreamingMessage`, since Python consumes the `RecordBatch`. + +**Tests:** flush producing a correct batch; mixed `PyAnyMessage` in the window panics; +an empty window produces nothing; committable and the synthetic watermark match +`PyListFlushProducer`'s behaviour for the same input. + +## Phase 5 — DSL and adapter wiring + +**`src/operators.rs`** — new variant. Note `schema_name`, per *Integration points*: + +```rust +#[pyo3(name = "ArrowBatchParser")] +ArrowBatchParser { + route: Route, + step_name: String, + schema_name: String, + max_batch_size: Option, + max_batch_time_ms: Option, +}, +``` + +with a `build()` arm mirroring `RuntimeOperator::Batch` (`operators.rs:232-240`), +converting ms → `Duration` the same way. + +**`sentry_streams/pipeline/pipeline.py`** — `ArrowBatchParser`, a `Reduce` subclass +carrying `batch_size` and `batch_timedelta` only. `override_config` and `validate` +mirror `Batch` (`pipeline.py:674-681`). No schema, no format, no type name. + +**`adapters/arroyo/rust_arroyo.py`** + +1. In `source()`, stash the pre-override schema name: + `self.__source_schemas[source_name] = schema_name`. +2. In `reduce()`, add an `isinstance(step, ArrowBatchParser)` branch before the `Batch` + branch, emitting `RuntimeOperator.ArrowBatchParser(..., schema_name=self.__source_schemas[stream.source], ...)`. +3. Build-time input check: walk the pipeline's incoming edges from this step; if any + predecessor is not a `RawMessage`-preserving step (source, `HeadersFilter`), raise + naming the step and the offending predecessor. + +**`adapters/arroyo/adapter.py`** — `NotImplementedError` pointing at the Rust adapter, +documented Rust-only in the style of `HeadersFilter`. + +**Also:** export from `pipeline/__init__.py` (both the import and `__all__`); add +`RuntimeOperator.ArrowBatchParser` and `ArrowRecordBatch` to `rust_streams.pyi`. + +**Tests:** placing the step after a Python `Map` fails at build time naming both steps; +a JSON topic panics at startup with the actual `schema_type`; the pure-Python adapter +raises `NotImplementedError`; `make typecheck` clean. + +## Phase 6 — Example, end-to-end, benchmark + +1. `sentry_streams/examples/arrow_trace_items.py` — `snuba-items` → `ArrowBatchParser` + → a `Map` consuming the batch via polars. +2. End-to-end test through the full step with real `TraceItem` payloads. +3. **Benchmark against `Batch` + `BatchParser`.** This validates the premise of the + exercise and tells us whether inline decoding (decision 5) holds against + `max_poll_interval_ms=60000`. Record throughput and p99 batch decode time here. +4. Docs page: the hardcoded-schema contract, Rust-adapter-only, failure behaviour. + +--- + +## Accepted consequences + +1. **A malformed message panics the process.** Decisions 9 and 14: one bad payload fails + the whole batch, cannot be dead-lettered, and crash-loops on restart since the + consumer re-reads the same offsets. Arroyo's DLQ needs an exact `(partition, offset)`; + batching collapses offsets to `max`. +2. **The Arrow schema lives in Rust.** Adding a column is a code change and a release — + the explicit PoC trade for dropping the descriptor pool. +3. **An attribute changing type between messages lands in different columns** across + batches. Inherent to decision 22; Snuba EAP has the same property. +4. **Recursive attribute values become JSON strings**, not structured data. +5. **Protobuf only.** A JSON or msgpack topic panics at startup. +6. **Rust adapter only**, unlike the Python `BatchParser`. +7. **Decoding holds the GIL** for the batch and stalls the consumer loop. Temporary, + and confined to `with_payloads`; see phase 4 and *Assumed future work*. + +## Assumed future work (owned elsewhere) + +**The source will stop boxing payloads into Python memory.** `RoutedValuePayload` will +carry a Rust-native message — `messages.rs` already has an unused `StreamingMessage` +enum reserved for it — instead of `Py` built by `into_pyraw` per message. + +This is outside the scope of this work, but the plan assumes it lands. When it does: + +- `with_payloads` loses its `traced_with_gil!` block and its `PyRef` guards. Nothing + else in this step changes. +- Consequence 7 disappears; decoding becomes genuinely GIL-free. +- Phase 6's benchmark should be re-run — the numbers taken before the change understate + what the step is capable of. +- The `BatchElement` type alias used by `BatchFlushProducer` and `with_payloads` is + where the element type changes; it exists so that swap is one line. + +Decision 1 (fused rather than a parser after `Batch`) is unaffected: `Batch`'s flush +still builds a Python list of `bytes`, so a separate parser would still pay a round trip +even once the source is native. + +## Deferred + +- **Descriptor-driven extraction** — `prost-reflect` + `DescriptorPool` restores + `get_field_by_name`, making a new column configuration rather than a release. + Descriptors from vendored `.proto` compiled by `protox`, or better from an upstream PR + adding `.file_descriptor_set_path(...)` to `sentry-protos`' generator so the crate + ships a `FILE_DESCRIPTOR_SET` and no vendoring is needed. +- The JSON path, and whether its schema is derived from the topic's JSON Schema (needing + a `$ref` resolver, the crate's being private) or declared. +- Per-row offset tracking, turning consequence 1 into a dead-lettered message and a + surviving batch. +- Threadpool decoding, if the benchmark justifies it. +- `TraceItem.outcomes`; msgpack topics. From a0b5f541796e7ea9d6c6a183d2b484fccbfdda03 Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 15:43:01 -0700 Subject: [PATCH 02/13] feat(arrow): add Arrow batch parser dependencies (phase 0) Adds the four crates the Arrow batch parser needs, plus tests pinning the two assumptions the design rests on. sentry-kafka-schemas is taken with default-features = false: validate_protobuf is the only item behind the type_generation feature, and we decode with prost directly, so get_schema/schema_type/raw_schema remain available while typify, syn, prettyplease and schemars drop out of the build. The tests assert that snuba-items resolves to the TraceItem resource string, and that a sentry_protos type round-trips through the prost::Message trait we import -- i.e. that both crates agree on a prost version. Resolved: arrow 59.3.0, prost 0.14.4, sentry_protos 0.70.0, sentry-kafka-schemas 3.0.1. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- sentry_streams/Cargo.lock | 797 +++++++++++++++++++++++++++++++++++++- sentry_streams/Cargo.toml | 4 + sentry_streams/src/lib.rs | 42 ++ 3 files changed, 833 insertions(+), 10 deletions(-) diff --git a/sentry_streams/Cargo.lock b/sentry_streams/Cargo.lock index e21a054a..626113e5 100644 --- a/sentry_streams/Cargo.lock +++ b/sentry_streams/Cargo.lock @@ -172,7 +172,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" dependencies = [ "cfg-if", + "const-random", + "getrandom 0.3.2", "once_cell", + "serde", "version_check", "zerocopy", ] @@ -257,6 +260,226 @@ version = "1.0.98" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" +[[package]] +name = "arrow" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c14b3d39f306bc28fd639d59f06e17a0f377d0021e1b7e9054e4d6fedc98774" +dependencies = [ + "arrow-arith", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-csv", + "arrow-data", + "arrow-ipc", + "arrow-json", + "arrow-ord", + "arrow-row", + "arrow-schema", + "arrow-select", + "arrow-string", +] + +[[package]] +name = "arrow-arith" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce2961626677665b2195eb59242af4c7befe7b8737ca2050295389362380104e" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "num-traits", +] + +[[package]] +name = "arrow-array" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e5f6adeffdf587d7a31db5d2266189624b526730cd3627f9ff9fedae97ad584" +dependencies = [ + "ahash", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "chrono", + "half", + "hashbrown 0.17.1", + "libc", + "num-complex", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-buffer" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "097d193003ce7995d5d087089069ec2a6e0187faf5a6f8c9f38af2645d987182" +dependencies = [ + "bytes", + "half", + "num-bigint 0.5.1", + "num-traits", +] + +[[package]] +name = "arrow-cast" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "635c9c635668ad26adf76cce8fb276c4be7cf06e63bd516de7da514f9680ee53" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-ord", + "arrow-schema", + "arrow-select", + "atoi", + "base64 0.23.1", + "chrono", + "half", + "lexical-core", + "num-traits", + "ryu", +] + +[[package]] +name = "arrow-csv" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c2ebf8d631e79b02c16cf5ae860561272c26024ec88fce389a56aaddd558e86" +dependencies = [ + "arrow-array", + "arrow-cast", + "arrow-schema", + "chrono", + "csv", + "csv-core", + "regex", +] + +[[package]] +name = "arrow-data" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ba2f832eaeca24b8f26143dba750e42ee4ab51cf7d65e701ca9607cfda9f358" +dependencies = [ + "arrow-buffer", + "arrow-schema", + "half", + "num-integer", + "num-traits", +] + +[[package]] +name = "arrow-ipc" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcc41681ea80f521df14c36725b74d4c60702c47f0793af2be469c04527e2599" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "flatbuffers", +] + +[[package]] +name = "arrow-json" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2f57d7a81969f24ccf80809587b76c09897e6f829d2d65a5976bfb3218851f1" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-ord", + "arrow-schema", + "arrow-select", + "chrono", + "half", + "indexmap", + "itoa", + "lexical-core", + "memchr", + "num-traits", + "ryu", + "serde_core", + "serde_json", + "simdutf8", +] + +[[package]] +name = "arrow-ord" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c900759f3bd8354fd4196bc4403eee846894dc2adf66b4225472006a0bf18c5" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", +] + +[[package]] +name = "arrow-row" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4c6425032e28266e3fc4ff680805e57e670d6ea92473043f3e65b7ed6ac79f2" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "half", +] + +[[package]] +name = "arrow-schema" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10fab8d4563491417ba801fab29d205104d20d4bdf37bda6cd1cf425cff598cd" +dependencies = [ + "bitflags", +] + +[[package]] +name = "arrow-select" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc58569193c2525915f3cc6310edba3792f1200f65d6e9ed330aa33e691493b8" +dependencies = [ + "ahash", + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "num-traits", +] + +[[package]] +name = "arrow-string" +version = "59.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e0813f3c35c1cfea65e14c20a953440f7783c088b7ad2d0db162ccdeefcec14" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-data", + "arrow-schema", + "arrow-select", + "memchr", + "num-traits", + "regex", + "regex-syntax", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -268,6 +491,15 @@ dependencies = [ "syn", ] +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -280,6 +512,49 @@ version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26" +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "futures-util", + "http 1.3.1", + "http-body", + "http-body-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "sync_wrapper", + "tower", + "tower-layer", + "tower-service", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http 1.3.1", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", +] + [[package]] name = "backtrace" version = "0.3.74" @@ -301,12 +576,33 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64ct" version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + [[package]] name = "bitflags" version = "2.11.1" @@ -328,6 +624,12 @@ version = "3.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1628fb46dfa0b37568d12e5edd512553eccf6a22a78e8bde00bb4aed84d5bdbf" +[[package]] +name = "bytecount" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "175812e0be2bccb6abe50bb8d566126198344f707e304f45c648fd8f2cc0365e" + [[package]] name = "bytes" version = "1.11.1" @@ -445,6 +747,26 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "const-random" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87e00182fe74b066627d63b85fd550ac2998d4b0bd86bfed477a0ae4c7c71359" +dependencies = [ + "const-random-macro", +] + +[[package]] +name = "const-random-macro" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9d839f2a20b0aee515dc581a6172f2321f96cab76c1a38a4c584a194955390e" +dependencies = [ + "getrandom 0.2.15", + "once_cell", + "tiny-keccak", +] + [[package]] name = "convert_case" version = "0.10.0" @@ -495,6 +817,33 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + [[package]] name = "ctrlc" version = "3.4.6" @@ -578,6 +927,12 @@ dependencies = [ "syn", ] +[[package]] +name = "either" +version = "1.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" + [[package]] name = "encoding_rs" version = "0.8.35" @@ -609,6 +964,17 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "fancy-regex" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + [[package]] name = "fastrand" version = "2.3.0" @@ -627,6 +993,16 @@ dependencies = [ "winapi", ] +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" +dependencies = [ + "bitflags", + "rustc_version", +] + [[package]] name = "fnv" version = "1.0.7" @@ -669,6 +1045,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fraction" +version = "0.15.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e076045bb43dac435333ed5f04caf35c7463631d0dae2deb2638d94dd0a5b872" +dependencies = [ + "lazy_static", + "num", +] + [[package]] name = "futures-channel" version = "0.3.31" @@ -726,7 +1112,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dbf67f30198e045a039264c01fb44659ce82402d7771c50938beb41a5ac87733" dependencies = [ "async-trait", - "base64", + "base64 0.22.1", "bytes", "chrono", "home", @@ -753,8 +1139,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4567c8db10ae91089c99af84c68c38da3ec2f087c3f82960bcdbf3656b6f4d7" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi 0.11.0+wasi-snapshot-preview1", + "wasm-bindgen", ] [[package]] @@ -794,6 +1182,18 @@ dependencies = [ "tracing", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "num-traits", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.15.2" @@ -809,6 +1209,12 @@ dependencies = [ "foldhash 0.2.0", ] +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + [[package]] name = "heck" version = "0.5.0" @@ -912,6 +1318,7 @@ dependencies = [ "http 1.3.1", "http-body", "httparse", + "httpdate", "itoa", "pin-project-lite", "smallvec", @@ -937,6 +1344,19 @@ dependencies = [ "tower-service", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-tls" version = "0.6.0" @@ -959,7 +1379,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -1167,6 +1587,24 @@ version = "1.70.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" +[[package]] +name = "iso8601" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ffd3254cf2b0fc53e38414bdba99719f3e269db8a6519731b68a3a90040c41b" +dependencies = [ + "nom", +] + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.15" @@ -1183,6 +1621,34 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonschema" +version = "0.18.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa0f4bea31643be4c6a678e9aa4ae44f0db9e5609d5ca9dc9083d06eb3e9a27a" +dependencies = [ + "ahash", + "anyhow", + "base64 0.22.1", + "bytecount", + "fancy-regex", + "fraction", + "getrandom 0.2.15", + "iso8601", + "itoa", + "memchr", + "num-cmp", + "once_cell", + "parking_lot", + "percent-encoding", + "regex", + "serde", + "serde_json", + "time", + "url", + "uuid", +] + [[package]] name = "language-tags" version = "0.3.2" @@ -1195,12 +1661,75 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" +[[package]] +name = "lexical-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" +dependencies = [ + "lexical-parse-float", + "lexical-parse-integer", + "lexical-util", + "lexical-write-float", + "lexical-write-integer", +] + +[[package]] +name = "lexical-parse-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" +dependencies = [ + "lexical-parse-integer", + "lexical-util", +] + +[[package]] +name = "lexical-parse-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" +dependencies = [ + "lexical-util", +] + +[[package]] +name = "lexical-util" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" + +[[package]] +name = "lexical-write-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" +dependencies = [ + "lexical-util", + "lexical-write-integer", +] + +[[package]] +name = "lexical-write-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" +dependencies = [ + "lexical-util", +] + [[package]] name = "libc" version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + [[package]] name = "libz-sys" version = "1.1.22" @@ -1247,6 +1776,12 @@ version = "0.4.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + [[package]] name = "memchr" version = "2.7.4" @@ -1375,6 +1910,15 @@ dependencies = [ "libc", ] +[[package]] +name = "nom" +version = "8.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df9761775871bdef83bee530e60050f7e54b1105350d6884eb0fb4f46c2f9405" +dependencies = [ + "memchr", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -1384,12 +1928,91 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "num" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" +dependencies = [ + "num-bigint 0.4.8", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-cmp" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63335b2e2c34fae2fb0aa2cecfd9f0832a1e24b3b32ecec612c3426d46dc8aaa" + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf97ec579c3c42f953ef76dbf8d55ac91fb219dde70e49aa4a6b7d74e9919050" +[[package]] +name = "num-integer" +version = "0.1.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-iter" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c92800bd69a1eac91786bcfe9da64a897eb72911b8dc3095decbd07429e8048b" +dependencies = [ + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" +dependencies = [ + "num-bigint 0.4.8", + "num-integer", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.19" @@ -1397,6 +2020,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" dependencies = [ "autocfg", + "libm", ] [[package]] @@ -1783,6 +2407,38 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "prost" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" +dependencies = [ + "anyhow", + "itertools", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "prost-types" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" +dependencies = [ + "prost", +] + [[package]] name = "pyo3" version = "0.29.0" @@ -2040,7 +2696,7 @@ version = "0.12.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d19c46a6fdd48bc4dab94b6103fccc55d34c67cc0ad04653aad4ea2a07cd7bbb" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -2084,7 +2740,7 @@ version = "0.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "62e0021ea2c22aed41653bc7e1419abb2c97e038ff2c33d0e1309e49a97deec0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-core", @@ -2136,6 +2792,7 @@ name = "rust_streams" version = "0.1.0" dependencies = [ "anyhow", + "arrow", "chrono", "clap", "ctrlc", @@ -2144,12 +2801,15 @@ dependencies = [ "metrics", "metrics-exporter-dogstatsd", "parking_lot", + "prost", "pyo3", "rand 0.9.4", "rdkafka", "reqwest 0.12.15", "sentry", + "sentry-kafka-schemas", "sentry_arroyo", + "sentry_protos", "serde", "serde_json", "tokio", @@ -2388,6 +3048,22 @@ dependencies = [ "sentry-core", ] +[[package]] +name = "sentry-kafka-schemas" +version = "3.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7845abe17525edcb1a6ddf137d89cad154aa6295a5b0eda8aacb956a7835f3a3" +dependencies = [ + "jsonschema", + "prost", + "sentry_protos", + "serde", + "serde_json", + "serde_yaml", + "thiserror 2.0.17", + "url", +] + [[package]] name = "sentry-panic" version = "0.48.1" @@ -2449,6 +3125,18 @@ dependencies = [ "uuid", ] +[[package]] +name = "sentry_protos" +version = "0.70.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9032bdd97373cb7dede51d6367c63dc715b4db51fb1a3124604e242a8a40c131" +dependencies = [ + "prost", + "prost-types", + "tonic", + "tonic-prost", +] + [[package]] name = "serde" version = "1.0.228" @@ -2503,6 +3191,19 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_yaml" +version = "0.9.34+deprecated" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" +dependencies = [ + "indexmap", + "itoa", + "ryu", + "serde", + "unsafe-libyaml", +] + [[package]] name = "sharded-slab" version = "0.1.7" @@ -2527,6 +3228,12 @@ dependencies = [ "libc", ] +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + [[package]] name = "sketches-ddsketch" version = "0.3.0" @@ -2738,6 +3445,15 @@ dependencies = [ "time-core", ] +[[package]] +name = "tiny-keccak" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237" +dependencies = [ + "crunchy", +] + [[package]] name = "tinystr" version = "0.7.6" @@ -2797,6 +3513,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.15" @@ -2827,6 +3554,46 @@ dependencies = [ "winnow", ] +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "async-trait", + "axum", + "base64 0.22.1", + "bytes", + "h2", + "http 1.3.1", + "http-body", + "http-body-util", + "hyper", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "socket2 0.6.3", + "sync_wrapper", + "tokio", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost", + "tonic", +] + [[package]] name = "tower" version = "0.5.2" @@ -2835,11 +3602,15 @@ checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9" dependencies = [ "futures-core", "futures-util", + "indexmap", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", + "tracing", ] [[package]] @@ -2973,6 +3744,12 @@ version = "0.2.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" +[[package]] +name = "unsafe-libyaml" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" + [[package]] name = "untrusted" version = "0.9.0" @@ -2985,7 +3762,7 @@ version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dea7109cdcd5864d4eeb1b58a1648dc9bf520360d7af16ec26d0a9354bafcfc0" dependencies = [ - "base64", + "base64 0.22.1", "der", "log", "native-tls", @@ -3002,7 +3779,7 @@ version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" dependencies = [ - "base64", + "base64 0.22.1", "http 1.3.1", "httparse", "log", @@ -3478,18 +4255,18 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.24" +version = "0.8.57" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2586fea28e186957ef732a5f8b3be2da217d65c5969d4b1e17f973ebbe876879" +checksum = "d35102a9f36d089ccae9e4c6802bc118be4487b80aaffc0ab4e0cf5ce92d2873" dependencies = [ "zerocopy-derive", ] [[package]] name = "zerocopy-derive" -version = "0.8.24" +version = "0.8.57" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a996a8f63c5c4448cd959ac1bab0aaa3306ccfd060472f85943ee0750f0169be" +checksum = "146c01f5ab44258da43cf276c74a2763db2ff3969c9c652c3f2de07041d0b2bc" dependencies = [ "proc-macro2", "quote", diff --git a/sentry_streams/Cargo.toml b/sentry_streams/Cargo.toml index 07847758..f64cff97 100644 --- a/sentry_streams/Cargo.toml +++ b/sentry_streams/Cargo.toml @@ -23,6 +23,10 @@ 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" +sentry_protos = "0.70" +sentry-kafka-schemas = { version = "3", default-features = false } [lib] name = "rust_streams" diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index b8819a2b..96ebae9a 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -52,3 +52,45 @@ fn rust_streams(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; Ok(()) } + +/// Phase 0 of the Arrow batch parser plan: prove the new dependencies are wired +/// correctly before any code depends on them. +/// +/// In particular this pins the two assumptions the plan rests on: +/// * `sentry-kafka-schemas` with `default-features = false` still exposes the +/// topic -> schema lookup (only `validate_protobuf` sits behind +/// `type_generation`), and +/// * `sentry_protos` and `prost` agree on a `prost` version, so a generated +/// type can actually be decoded through the `prost::Message` trait we import. +/// +/// See `docs/design/arrow-batch-parser.md`. +#[cfg(test)] +mod dependency_wiring_tests { + use prost::Message; + use sentry_kafka_schemas::{get_schema, SchemaType}; + use sentry_protos::snuba::v1::TraceItem; + + #[test] + fn snuba_items_resolves_to_the_trace_item_resource() { + let schema = get_schema("snuba-items", None).expect("snuba-items must have a schema"); + + assert_eq!(schema.schema_type, SchemaType::Protobuf); + assert_eq!( + schema.raw_schema(), + "sentry_protos.snuba.v1.trace_item_pb2.TraceItem" + ); + } + + #[test] + fn sentry_protos_types_decode_through_our_prost() { + let item = TraceItem { + organization_id: 7, + ..Default::default() + }; + + let encoded = item.encode_to_vec(); + let decoded = TraceItem::decode(encoded.as_slice()).expect("round trip"); + + assert_eq!(decoded.organization_id, 7); + } +} From 5839f5a1a9f6c6678d774c2a5513e4a9b25f579e Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 15:49:14 -0700 Subject: [PATCH 03/13] feat(arrow): expose RecordBatch to Python via PyCapsule (phase 1) Adds ArrowRecordBatch, a pyclass wrapping an Arrow RecordBatch and exporting it over the Arrow PyCapsule interface. No pyarrow dependency: any consumer speaking the protocol reads the batch without a copy. All three dunders are implemented, not just __arrow_c_array__. Table-level consumers (pl.DataFrame, pa.table) look for __arrow_c_stream__ while array-level consumers use __arrow_c_array__; implementing only one makes the object work at some call sites and not others. requested_schema is accepted and ignored -- the protocol permits returning the native schema when a cast is unsupported, and we never cast. Tested by round-tripping each capsule back through arrow's FFI importer, which also proves exporting twice does not double-release, and by reading the batch from polars as an independent Arrow implementation. The fixture carries a Map column so the nested case is covered everywhere rather than only in a dedicated test. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- .../sentry_streams/rust_streams.pyi | 16 + sentry_streams/src/lib.rs | 2 + sentry_streams/src/py_record_batch.rs | 384 ++++++++++++++++++ 3 files changed, 402 insertions(+) create mode 100644 sentry_streams/src/py_record_batch.rs diff --git a/sentry_streams/sentry_streams/rust_streams.pyi b/sentry_streams/sentry_streams/rust_streams.pyi index e4c359cf..606cde96 100644 --- a/sentry_streams/sentry_streams/rust_streams.pyi +++ b/sentry_streams/sentry_streams/rust_streams.pyi @@ -211,3 +211,19 @@ class PyWatermark: def timestamp(self) -> int: ... @property def last_message_time(self) -> float | None: ... + +class ArrowRecordBatch: + """An Apache Arrow RecordBatch produced by the Rust runtime. + + Readable by any consumer implementing the Arrow PyCapsule interface, for + example ``polars.DataFrame(batch)`` or ``pyarrow.table(batch)``. There is no + Python constructor: instances come out of the Arrow batch parser step. + """ + + @property + def num_rows(self) -> int: ... + @property + def num_columns(self) -> int: ... + def __arrow_c_array__(self, requested_schema: object | None = None) -> Tuple[Any, Any]: ... + def __arrow_c_schema__(self) -> Any: ... + def __arrow_c_stream__(self, requested_schema: object | None = None) -> Any: ... diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index 96ebae9a..d9def81a 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -16,6 +16,7 @@ mod metrics_config; mod mocks; mod operators; mod pipeline_stats; +mod py_record_batch; mod python_operator; mod routers; mod routes; @@ -50,6 +51,7 @@ fn rust_streams(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; Ok(()) } diff --git a/sentry_streams/src/py_record_batch.rs b/sentry_streams/src/py_record_batch.rs new file mode 100644 index 00000000..7d4ae5bf --- /dev/null +++ b/sentry_streams/src/py_record_batch.rs @@ -0,0 +1,384 @@ +//! Hands an Arrow `RecordBatch` to Python over the Arrow PyCapsule interface. +//! +//! Deliberately no `pyarrow` dependency: any consumer that speaks the PyCapsule +//! protocol (polars, pyarrow, duckdb, ...) can read the batch without a copy. +//! +//! See `docs/design/arrow-batch-parser.md`, phase 1. + +use arrow::array::{Array, RecordBatch, RecordBatchIterator, StructArray}; +use arrow::error::ArrowError; +use arrow::ffi::{to_ffi, FFI_ArrowSchema}; +use arrow::ffi_stream::FFI_ArrowArrayStream; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyCapsule; +use std::ffi::CStr; + +/// Capsule names mandated by the Arrow PyCapsule interface. Getting one wrong is +/// not a soft failure: consumers reject the capsule with an opaque error, so they +/// are named constants and asserted in the tests. +const SCHEMA_CAPSULE_NAME: &CStr = c"arrow_schema"; +const ARRAY_CAPSULE_NAME: &CStr = c"arrow_array"; +const STREAM_CAPSULE_NAME: &CStr = c"arrow_array_stream"; + +/// An Arrow `RecordBatch` produced by the Rust runtime, readable from Python by +/// anything that speaks the Arrow PyCapsule interface: +/// +/// ```python +/// import polars as pl +/// df = pl.DataFrame(batch) +/// ``` +#[pyclass( + name = "ArrowRecordBatch", + module = "sentry_streams.rust_streams", + frozen +)] +pub struct PyRecordBatch { + pub(crate) batch: RecordBatch, +} + +impl PyRecordBatch { + // Constructed by the Arrow batch parser step (phase 4); until then only the + // tests build one. + #[allow(dead_code)] + pub(crate) fn new(batch: RecordBatch) -> Self { + Self { batch } + } +} + +fn arrow_err(e: ArrowError) -> PyErr { + PyRuntimeError::new_err(format!("Arrow C data interface export failed: {e}")) +} + +#[pymethods] +impl PyRecordBatch { + #[getter] + fn num_rows(&self) -> usize { + self.batch.num_rows() + } + + #[getter] + fn num_columns(&self) -> usize { + self.batch.num_columns() + } + + fn __repr__(&self) -> String { + format!( + "ArrowRecordBatch(num_rows={}, num_columns={})", + self.batch.num_rows(), + self.batch.num_columns() + ) + } + + /// Export as a single Arrow array (a struct array, one field per column). + /// + /// `requested_schema` is accepted and ignored: the PyCapsule interface allows + /// a producer to return its native schema when it cannot perform the + /// requested cast, and we never cast. + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_array__<'py>( + &self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)> { + let _ = requested_schema; + + let struct_array = StructArray::from(self.batch.clone()); + let (ffi_array, ffi_schema) = to_ffi(&struct_array.to_data()).map_err(arrow_err)?; + + // The capsule takes ownership; FFI_ArrowSchema/FFI_ArrowArray's Drop + // invokes the C release callback, so no manual destructor is needed. + let schema_capsule = PyCapsule::new_with_value(py, ffi_schema, SCHEMA_CAPSULE_NAME)?; + let array_capsule = PyCapsule::new_with_value(py, ffi_array, ARRAY_CAPSULE_NAME)?; + Ok((schema_capsule, array_capsule)) + } + + /// Export just the schema, for consumers that inspect before reading. + fn __arrow_c_schema__<'py>(&self, py: Python<'py>) -> PyResult> { + let ffi_schema = + FFI_ArrowSchema::try_from(self.batch.schema().as_ref()).map_err(arrow_err)?; + PyCapsule::new_with_value(py, ffi_schema, SCHEMA_CAPSULE_NAME) + } + + /// Export as a stream of exactly one batch. + /// + /// Table-level consumers (`pl.DataFrame(obj)`, `pa.table(obj)`) look for this + /// rather than `__arrow_c_array__`, which is why both exist. + #[pyo3(signature = (requested_schema=None))] + fn __arrow_c_stream__<'py>( + &self, + py: Python<'py>, + requested_schema: Option>, + ) -> PyResult> { + let _ = requested_schema; + + let batch = self.batch.clone(); + let schema = batch.schema(); + let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); + let stream = FFI_ArrowArrayStream::new(Box::new(reader)); + + PyCapsule::new_with_value(py, stream, STREAM_CAPSULE_NAME) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + ArrayRef, Int64Array, MapBuilder, MapFieldNames, StringArray, StringBuilder, + }; + use arrow::datatypes::Schema; + use arrow::ffi::from_ffi; + use arrow::ffi_stream::ArrowArrayStreamReader; + use arrow::record_batch::RecordBatchReader; + use pyo3::types::PyCapsuleMethods; + use std::sync::Arc; + + fn capsule_name(capsule: &Bound<'_, PyCapsule>) -> &'static CStr { + // SAFETY: the name is a `&'static CStr` we set ourselves and no Python + // code has had the chance to rename the capsule. + unsafe { capsule.name().unwrap().unwrap().as_cstr() } + } + + /// Consume the capsules the way a real consumer does: move the array out and + /// leave the exported struct released, so double-release bugs would show up. + fn import_array( + schema_capsule: &Bound<'_, PyCapsule>, + array_capsule: &Bound<'_, PyCapsule>, + ) -> RecordBatch { + let data = unsafe { + let schema_ptr = schema_capsule + .pointer_checked(Some(SCHEMA_CAPSULE_NAME)) + .unwrap() + .as_ptr() as *const FFI_ArrowSchema; + let array_ptr = array_capsule + .pointer_checked(Some(ARRAY_CAPSULE_NAME)) + .unwrap() + .as_ptr() as *mut arrow::ffi::FFI_ArrowArray; + let array = std::ptr::replace(array_ptr, arrow::ffi::FFI_ArrowArray::empty()); + from_ffi(array, &*schema_ptr).unwrap() + }; + RecordBatch::from(StructArray::from(data)) + } + + fn import_schema(capsule: &Bound<'_, PyCapsule>) -> Schema { + unsafe { + let ptr = capsule + .pointer_checked(Some(SCHEMA_CAPSULE_NAME)) + .unwrap() + .as_ptr() as *const FFI_ArrowSchema; + Schema::try_from(&*ptr).unwrap() + } + } + + /// Two scalar columns plus a `Map`, so the nested case is + /// exercised by every test rather than only by a dedicated one. + fn sample_batch() -> RecordBatch { + let ids: ArrayRef = Arc::new(Int64Array::from(vec![1_i64, 2, 3])); + let names: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])); + + let mut attrs = MapBuilder::new( + Some(MapFieldNames { + entry: "entries".to_string(), + key: "key".to_string(), + value: "value".to_string(), + }), + StringBuilder::new(), + StringBuilder::new(), + ); + // row 0: two entries, row 1: none, row 2: one entry + attrs.keys().append_value("k1"); + attrs.values().append_value("v1"); + attrs.keys().append_value("k2"); + attrs.values().append_value("v2"); + attrs.append(true).unwrap(); + attrs.append(true).unwrap(); + attrs.keys().append_value("k3"); + attrs.values().append_value("v3"); + attrs.append(true).unwrap(); + let attrs: ArrayRef = Arc::new(attrs.finish()); + + RecordBatch::try_from_iter(vec![("id", ids), ("name", names), ("attrs", attrs)]).unwrap() + } + + fn py_batch(py: Python<'_>) -> Py { + Py::new(py, PyRecordBatch::new(sample_batch())).unwrap() + } + + #[test] + fn exposes_shape_to_python() { + crate::testutils::initialize_python(); + Python::attach(|py| { + let rb = py_batch(py); + let b = rb.bind(py); + assert_eq!( + b.getattr("num_rows").unwrap().extract::().unwrap(), + 3 + ); + assert_eq!( + b.getattr("num_columns") + .unwrap() + .extract::() + .unwrap(), + 3 + ); + let repr = b.repr().unwrap().extract::().unwrap(); + assert!(repr.contains("ArrowRecordBatch"), "got {repr}"); + assert!(repr.contains('3'), "repr should mention the shape: {repr}"); + }); + } + + #[test] + fn arrow_c_array_round_trips_through_ffi() { + crate::testutils::initialize_python(); + Python::attach(|py| { + let rb = py_batch(py); + let (schema_capsule, array_capsule) = rb.get().__arrow_c_array__(py, None).unwrap(); + + assert_eq!(capsule_name(&schema_capsule), SCHEMA_CAPSULE_NAME); + assert_eq!(capsule_name(&array_capsule), ARRAY_CAPSULE_NAME); + + // Import the capsules back and compare against the original batch. + let round_tripped = import_array(&schema_capsule, &array_capsule); + assert_eq!(round_tripped, sample_batch()); + }); + } + + #[test] + fn arrow_c_schema_describes_the_batch() { + crate::testutils::initialize_python(); + Python::attach(|py| { + let rb = py_batch(py); + let capsule = rb.get().__arrow_c_schema__(py).unwrap(); + assert_eq!(capsule_name(&capsule), SCHEMA_CAPSULE_NAME); + + assert_eq!(&import_schema(&capsule), sample_batch().schema().as_ref()); + }); + } + + #[test] + fn arrow_c_stream_yields_the_batch() { + crate::testutils::initialize_python(); + Python::attach(|py| { + let rb = py_batch(py); + let capsule = rb.get().__arrow_c_stream__(py, None).unwrap(); + assert_eq!(capsule_name(&capsule), STREAM_CAPSULE_NAME); + + let ptr = capsule + .pointer_checked(Some(STREAM_CAPSULE_NAME)) + .unwrap() + .as_ptr() as *mut FFI_ArrowArrayStream; + let mut reader = unsafe { ArrowArrayStreamReader::from_raw(ptr) }.unwrap(); + assert_eq!(reader.schema().as_ref(), sample_batch().schema().as_ref()); + let batch = reader.next().unwrap().unwrap(); + assert_eq!(batch, sample_batch()); + assert!( + reader.next().is_none(), + "stream must hold exactly one batch" + ); + }); + } + + /// Exporting must not consume the batch: the object stays usable, which is + /// what a Python caller passing it to two consumers would expect. + #[test] + fn can_be_exported_twice() { + crate::testutils::initialize_python(); + Python::attach(|py| { + let rb = py_batch(py); + let _first = rb.get().__arrow_c_array__(py, None).unwrap(); + let (schema_capsule, array_capsule) = rb.get().__arrow_c_array__(py, None).unwrap(); + + assert_eq!( + import_array(&schema_capsule, &array_capsule), + sample_batch() + ); + }); + } + + /// `requested_schema` is accepted and ignored; the protocol allows returning + /// the native schema when the requested cast is unsupported. + #[test] + fn requested_schema_is_ignored_not_rejected() { + crate::testutils::initialize_python(); + Python::attach(|py| { + let rb = py_batch(py); + let requested = rb.get().__arrow_c_schema__(py).unwrap().into_any(); + assert!(rb + .get() + .__arrow_c_array__(py, Some(requested.clone())) + .is_ok()); + assert!(rb.get().__arrow_c_stream__(py, Some(requested)).is_ok()); + }); + } + + /// The real acceptance criterion: an independent Arrow implementation reads + /// our batch with the right values, names and dtypes. + #[test] + fn polars_reads_the_batch() { + crate::testutils::initialize_python(); + Python::attach(|py| { + // polars is a declared runtime dependency of this package, so a + // missing import is a broken environment, not a reason to skip. + let pl = py.import("polars").expect("polars must be importable"); + let rb = py_batch(py); + let df = pl.call_method1("DataFrame", (rb,)).unwrap(); + + let columns: Vec = df.getattr("columns").unwrap().extract().unwrap(); + assert_eq!(columns, vec!["id", "name", "attrs"]); + assert_eq!( + df.call_method0("__len__") + .unwrap() + .extract::() + .unwrap(), + 3 + ); + + let ids: Vec = df + .get_item("id") + .unwrap() + .call_method0("to_list") + .unwrap() + .extract() + .unwrap(); + assert_eq!(ids, vec![1, 2, 3]); + + let names: Vec> = df + .get_item("name") + .unwrap() + .call_method0("to_list") + .unwrap() + .extract() + .unwrap(); + assert_eq!( + names, + vec![Some("a".to_string()), None, Some("c".to_string())] + ); + + let dtype = df + .get_item("id") + .unwrap() + .getattr("dtype") + .unwrap() + .str() + .unwrap() + .extract::() + .unwrap(); + assert_eq!(dtype, "Int64"); + }); + } + + #[test] + fn empty_batch_keeps_its_schema() { + crate::testutils::initialize_python(); + Python::attach(|py| { + let schema = sample_batch().schema(); + let empty = RecordBatch::new_empty(schema.clone()); + let rb = Py::new(py, PyRecordBatch::new(empty)).unwrap(); + assert_eq!(rb.get().num_rows(), 0); + + let capsule = rb.get().__arrow_c_schema__(py).unwrap(); + assert_eq!(&import_schema(&capsule), schema.as_ref()); + }); + } +} From 4f08736c446686b24d81839a25859fe2b67448a0 Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 15:52:54 -0700 Subject: [PATCH 04/13] refactor(arrow): generalize BatchStep over a flush producer (phase 2) Lifts the body of Batch::flush into a BatchFlushProducer trait, implemented by PyListFlushProducer with the existing behaviour moved verbatim. BatchStep owns windowing, watermark ordering and backpressure; a producer owns only the shape of the emitted message. The Arrow batch parser will plug in a producer that decodes the same payloads into a RecordBatch instead of a Python list. No behaviour change: the existing test module is byte-identical and passes unmodified. BatchStep::new and Batch::from_initial keep their signatures and default to PyListFlushProducer; the generalized forms are with_producer and from_initial_with_producer. from_initial is now reached only from tests, so it is marked #[cfg(test)] rather than left as dead weight in the build. Adds a BatchElement alias for the window's element type. It is a PyStreamingMessage today because the source hands Rust payloads already boxed in Python memory; it exists so that swapping to a Rust-native message is a one-line change. The new producer_seam test asserts a custom producer receives the whole window with offsets collapsed to max per partition, and that its message -- not a Python list -- reaches downstream. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- sentry_streams/src/batch_step.rs | 275 +++++++++++++++++++++++++++---- 1 file changed, 240 insertions(+), 35 deletions(-) diff --git a/sentry_streams/src/batch_step.rs b/sentry_streams/src/batch_step.rs index f3e8cdbd..b75fadf0 100644 --- a/sentry_streams/src/batch_step.rs +++ b/sentry_streams/src/batch_step.rs @@ -20,6 +20,7 @@ use sentry_arroyo::processing::strategies::{ use sentry_arroyo::types::{Message, Partition}; use sentry_arroyo::utils::timing::Deadline; use std::collections::{BTreeMap, VecDeque}; +use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; const METRIC_BATCH_SIZE: &str = "streams.pipeline.batch.size"; @@ -58,6 +59,74 @@ fn list_item_for_streaming_message( } } +/// One element of an open batch window. +/// +/// Today the source hands Rust every payload already boxed in Python memory +/// (`consumer.rs::to_routed_value`), so an element is a [`PyStreamingMessage`]. +/// When the source starts emitting Rust-native messages this alias is the single +/// line that changes. See `docs/design/arrow-batch-parser.md`. +pub(crate) type BatchElement = PyStreamingMessage; + +/// Turns a flushed batch window into the message sent downstream. +/// +/// [`BatchStep`] owns windowing, watermark ordering and backpressure; a producer +/// owns only the shape of the emitted payload. [`PyListFlushProducer`] builds the +/// Python list the `Batch` step has always built; the Arrow batch parser swaps in +/// a producer that decodes the same payloads into an Arrow `RecordBatch` without +/// ever materialising them as Python objects. +pub(crate) trait BatchFlushProducer: Send + Sync { + fn produce( + &self, + route: &Route, + elements: &[BatchElement], + committable: BTreeMap, + ) -> Result, StrategyError>; +} + +/// The historical `Batch` behaviour: one `PyAnyMessage` whose payload is a Python +/// `list`, one item per element, schema taken from the first element. +pub(crate) struct PyListFlushProducer; + +impl BatchFlushProducer for PyListFlushProducer { + fn produce( + &self, + route: &Route, + elements: &[BatchElement], + committable: BTreeMap, + ) -> Result, StrategyError> { + let route = route.clone(); + let ts = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs_f64()) + .unwrap_or(0.0); + + let content = traced_with_gil!(|py| -> PyResult> { + let first_schema = first_element_schema(py, &elements[0]); + let py_items: Result>, _> = elements + .iter() + .map(|el| list_item_for_streaming_message(py, el)) + .collect(); + let py_items = py_items.map_err(|e: PyErr| e)?; + let list = PyList::new(py, &py_items)?.unbind(); + let inner = PyAnyMessage { + payload: list.into_any(), + headers: vec![], + timestamp: ts, + schema: first_schema, + }; + into_pyany(py, inner) + }) + .map_err(|e| StrategyError::Other(Box::new(e)))?; + + let py_streaming = PyStreamingMessage::PyAnyMessage { content }; + let rv = RoutedValue { + route, + payload: RoutedValuePayload::PyStreamingMessage(py_streaming), + }; + Ok(Message::new_any_message(rv, committable)) + } +} + /// Count- and/or time-based window of streaming elements for one route. On flush, output is /// always a batched `PyAnyMessage` with a list payload. pub(crate) struct Batch { @@ -67,13 +136,17 @@ pub(crate) struct Batch { batch_deadline: Option, /// Wall time when the first element opened this batch window. created_at: Instant, - elements: Vec, + elements: Vec, batch_offsets: BTreeMap, + producer: Arc, } impl Batch { /// First element in a window. `committable` and `first` are from the same [`RoutedValue`] /// (see [`BatchStep::submit`]). Later elements may use either `PyAnyMessage` or `RawMessage`. + /// Convenience form defaulting to [`PyListFlushProducer`]. Only the tests + /// use it now: [`BatchStep`] always supplies its own producer. + #[cfg(test)] pub fn from_initial( route: Route, max_batch_size: Option, @@ -82,6 +155,26 @@ impl Batch { // we will return when the batch is flushed. committable: BTreeMap, first: PyStreamingMessage, + ) -> Self { + Self::from_initial_with_producer( + route, + max_batch_size, + max_batch_time, + committable, + first, + Arc::new(PyListFlushProducer), + ) + } + + /// As [`Self::from_initial`], but emitting through `producer` instead of + /// building a Python list. + pub fn from_initial_with_producer( + route: Route, + max_batch_size: Option, + max_batch_time: Option, + committable: BTreeMap, + first: BatchElement, + producer: Arc, ) -> Self { let mut batch_offsets: BTreeMap = BTreeMap::new(); for (p, o) in committable { @@ -98,10 +191,11 @@ impl Batch { created_at: Instant::now(), elements: vec![first], batch_offsets, + producer, } } - pub fn append(&mut self, committable: BTreeMap, pysm: PyStreamingMessage) { + pub fn append(&mut self, committable: BTreeMap, pysm: BatchElement) { for (p, o) in committable { self.batch_offsets .entry(p) @@ -163,39 +257,11 @@ impl Batch { }) } + /// Delegates to the step's [`BatchFlushProducer`]; the window itself has no + /// opinion on the shape of the emitted message. pub fn flush(&self) -> Result, StrategyError> { - let route = self.route.clone(); - let committable = self.batch_offsets.clone(); - let ts = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|d| d.as_secs_f64()) - .unwrap_or(0.0); - - let content = traced_with_gil!(|py| -> PyResult> { - let first_schema = first_element_schema(py, &self.elements[0]); - let py_items: Result>, _> = self - .elements - .iter() - .map(|el| list_item_for_streaming_message(py, el)) - .collect(); - let py_items = py_items.map_err(|e: PyErr| e)?; - let list = PyList::new(py, &py_items)?.unbind(); - let inner = PyAnyMessage { - payload: list.into_any(), - headers: vec![], - timestamp: ts, - schema: first_schema, - }; - into_pyany(py, inner) - }) - .map_err(|e| StrategyError::Other(Box::new(e)))?; - - let py_streaming = PyStreamingMessage::PyAnyMessage { content }; - let rv = RoutedValue { - route, - payload: RoutedValuePayload::PyStreamingMessage(py_streaming), - }; - Ok(Message::new_any_message(rv, committable)) + self.producer + .produce(&self.route, &self.elements, self.batch_offsets.clone()) } } @@ -208,6 +274,8 @@ pub struct BatchStep { max_batch_time: Option, /// `None` until the first streaming message in a window. batch: Option, + /// Shapes the message emitted on flush. Shared by every window this step opens. + producer: Arc, /// Watermarks received while the current batch window is open; on successful batch send they /// are appended to [`Self::outbound`]. /// @@ -241,6 +309,27 @@ impl BatchStep { max_batch_time: Option, step_name: String, next_step: Box>, + ) -> Self { + Self::with_producer( + route, + max_batch_size, + max_batch_time, + step_name, + next_step, + Arc::new(PyListFlushProducer), + ) + } + + /// As [`Self::new`], but emitting flushed windows through `producer`. + /// Everything else -- windowing, watermark ordering, backpressure -- is + /// identical, which is the point of the seam. + pub fn with_producer( + route: Route, + max_batch_size: Option, + max_batch_time: Option, + step_name: String, + next_step: Box>, + producer: Arc, ) -> Self { let step_labels = vec![("step".to_string(), step_name.clone())]; Self { @@ -250,6 +339,7 @@ impl BatchStep { max_batch_size, max_batch_time, batch: None, + producer, watermark_buffer: Vec::new(), outbound: VecDeque::new(), pending_batch: false, @@ -496,12 +586,13 @@ impl ProcessingStrategy for BatchStep { } RoutedValuePayload::PyStreamingMessage(pysm) => { if self.batch.is_none() { - self.batch = Some(Batch::from_initial( + self.batch = Some(Batch::from_initial_with_producer( self.route.clone(), self.max_batch_size, self.max_batch_time, committable, pysm, + Arc::clone(&self.producer), )); } else { self.batch @@ -1052,4 +1143,118 @@ mod tests { ); } } + + mod producer_seam { + //! The generalisation added for the Arrow batch parser: [`BatchStep`] + //! delegates the *shape* of the flushed message to a + //! [`BatchFlushProducer`], and nothing else about the step changes. + + use crate::batch_step::{BatchElement, BatchFlushProducer, BatchStep}; + use crate::fake_strategy::FakeStrategy; + use crate::messages::{PyAnyMessage, PyStreamingMessage, RoutedValuePayload}; + use crate::routes::{Route, RoutedValue}; + use crate::testutils::build_routed_value; + use crate::utils::traced_with_gil; + use pyo3::prelude::*; + use pyo3::IntoPyObject; + use sentry_arroyo::processing::strategies::{ProcessingStrategy, StrategyError}; + use sentry_arroyo::types::{Message, Partition, Topic}; + use std::collections::BTreeMap; + use std::sync::{Arc, Mutex}; + + /// Emits a fixed marker payload and records what it was handed, so the + /// test can assert on the elements and committable the step passes in. + struct RecordingProducer { + seen: Arc)>>>, + } + + impl BatchFlushProducer for RecordingProducer { + fn produce( + &self, + route: &Route, + elements: &[BatchElement], + committable: BTreeMap, + ) -> Result, StrategyError> { + self.seen + .lock() + .unwrap() + .push((elements.len(), committable.clone())); + + let content = traced_with_gil!(|py| { + crate::messages::into_pyany( + py, + PyAnyMessage { + payload: "produced-by-seam" + .into_pyobject(py) + .unwrap() + .into_any() + .unbind(), + headers: vec![], + timestamp: 0.0, + schema: None, + }, + ) + }) + .unwrap(); + + Ok(Message::new_any_message( + RoutedValue { + route: route.clone(), + payload: RoutedValuePayload::PyStreamingMessage( + PyStreamingMessage::PyAnyMessage { content }, + ), + }, + committable, + )) + } + } + + #[test] + fn custom_producer_shapes_the_flushed_message() { + crate::testutils::initialize_python(); + let route = Route::new("s".into(), vec!["w".into()]); + let seen = Arc::new(Mutex::new(Vec::new())); + let sub = Arc::new(Mutex::new(Vec::new())); + let wms = Arc::new(Mutex::new(Vec::new())); + + let mut step = BatchStep::with_producer( + route, + Some(2), + None, + "test_seam".to_string(), + Box::new(FakeStrategy::new(sub.clone(), wms, false)), + Arc::new(RecordingProducer { seen: seen.clone() }), + ); + + let partition = Partition::new(Topic::new("t"), 0); + traced_with_gil!(|py| { + for (i, offset) in [7_u64, 9].into_iter().enumerate() { + let payload = (i as i32).into_pyobject(py).unwrap().into_any().unbind(); + let msg = Message::new_any_message( + build_routed_value(py, payload, "s", vec!["w".into()]), + BTreeMap::from([(partition, offset)]), + ); + step.submit(msg).unwrap(); + } + step.poll().unwrap(); + }); + + // The producer saw the whole window, with offsets collapsed to max. + let seen = seen.lock().unwrap(); + assert_eq!(seen.len(), 1, "exactly one flush"); + assert_eq!( + seen[0].0, 2, + "producer receives every element in the window" + ); + assert_eq!(seen[0].1, BTreeMap::from([(partition, 9)])); + + // ... and its message, not a Python list, reached downstream. + let out = sub.lock().unwrap(); + assert_eq!(out.len(), 1); + traced_with_gil!(|py| { + let payload: String = out[0].bind(py).extract().unwrap(); + assert_eq!(payload, "produced-by-seam"); + }); + } + } } From 3b67c5e86e35126eabab2f02e7f86b4422b45b7d Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 16:00:13 -0700 Subject: [PATCH 05/13] feat(arrow): protobuf to Arrow extractors, TraceItem (phase 3) Adds the extractor module: a registry keyed by the sentry-kafka-schemas resource string, and a hand-written TraceItem extractor with a hardcoded Arrow schema. Several topics sharing a schema share one extractor; adding a message type is one new file plus one line in the registry. Extraction is batch-wise -- one Arrow builder per column fed across all rows and finished once -- which is the shape Arrow builders want and keeps a second extractor small. Details that are easy to get wrong, each covered by a test: * map entry fields are named entries/key/value per the Arrow spec, not arrow-rs's entries/keys/values default, or pyarrow and polars interop breaks confusingly; * every map builder is closed on every row, including rows with no attributes of that type, or all later rows silently misalign; * attribute keys are sorted, because prost decodes the map into a HashMap whose iteration order varies run to run; * an unknown item_type enum value renders as TRACE_ITEM_TYPE_UNKNOWN_ rather than panicking -- a newer producer adding a member is routine, and crashing on it would make someone else's deploy our outage; * timestamps use checked arithmetic and report a field error on overflow. Two corrections to the design doc, applied here: * conversation_id and session_id are plain proto3 strings in sentry_protos 0.70.0, not Option. They have implicit presence, so by the plan's own rule they are non-nullable columns carrying "" -- the table had them nullable on the assumption they were declared optional. * bytes nested inside recursive attribute values are base64-encoded, following proto3 canonical JSON. Top-level bytes attributes still keep their raw bytes in attr_bytes, which is the case the design rejected JSON for. Adds prost-types and base64 as direct dependencies; both were already in the lock file transitively. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- sentry_streams/Cargo.lock | 2 + sentry_streams/Cargo.toml | 2 + .../docs/design/arrow-batch-parser.md | 25 +- sentry_streams/src/extractors/mod.rs | 138 ++++ sentry_streams/src/extractors/trace_item.rs | 766 ++++++++++++++++++ sentry_streams/src/lib.rs | 4 + 6 files changed, 933 insertions(+), 4 deletions(-) create mode 100644 sentry_streams/src/extractors/mod.rs create mode 100644 sentry_streams/src/extractors/trace_item.rs diff --git a/sentry_streams/Cargo.lock b/sentry_streams/Cargo.lock index 626113e5..1d387e2a 100644 --- a/sentry_streams/Cargo.lock +++ b/sentry_streams/Cargo.lock @@ -2793,6 +2793,7 @@ version = "0.1.0" dependencies = [ "anyhow", "arrow", + "base64 0.22.1", "chrono", "clap", "ctrlc", @@ -2802,6 +2803,7 @@ dependencies = [ "metrics-exporter-dogstatsd", "parking_lot", "prost", + "prost-types", "pyo3", "rand 0.9.4", "rdkafka", diff --git a/sentry_streams/Cargo.toml b/sentry_streams/Cargo.toml index f64cff97..23fa2383 100644 --- a/sentry_streams/Cargo.toml +++ b/sentry_streams/Cargo.toml @@ -25,6 +25,8 @@ rand = "0.9" sentry = "0.48.1" arrow = { version = "59", features = ["ffi"] } prost = "0.14" +prost-types = "0.14" +base64 = "0.22" sentry_protos = "0.70" sentry-kafka-schemas = { version = "3", default-features = false } diff --git a/sentry_streams/docs/design/arrow-batch-parser.md b/sentry_streams/docs/design/arrow-batch-parser.md index f858163c..a11cfdcf 100644 --- a/sentry_streams/docs/design/arrow-batch-parser.md +++ b/sentry_streams/docs/design/arrow-batch-parser.md @@ -79,8 +79,8 @@ Constraints discovered by reading the runtime. These drive several choices below | `timestamp` | `Timestamp(us, "UTC")` | **yes** | 6 | | `client_sample_rate` | `Float64` | no | 8 | | `server_sample_rate` | `Float64` | no | 9 | -| `conversation_id` | `Utf8` | yes | 10, proto3 `optional` | -| `session_id` | `Utf8` | yes | 11, proto3 `optional` | +| `conversation_id` | `Utf8` | no | 10 | +| `session_id` | `Utf8` | no | 11 | | `retention_days` | `UInt32` | no | 100 | | `received` | `Timestamp(us, "UTC")` | **yes** | 101 | | `downsampled_retention_days` | `UInt32` | no | 102 | @@ -95,18 +95,34 @@ fields always have explicit presence in proto3, so prost yields `Option **Corrected during phase 3.** This table originally marked `conversation_id` and +> `session_id` nullable, on the assumption they were declared `optional`. In +> `sentry_protos` 0.70.0 both are plain `String`, not `Option` — implicit +> presence, despite the "if any" comments in the proto. Applying the rule above, +> they are non-nullable columns carrying `""` when unset. Should they ever gain +> `optional` upstream, prost will change their type and the extractor will fail to +> compile, which is the right way to find out. + `ArrayValue` (arm 5) and `KeyValueList` (arm 6) are recursive; Arrow has no recursive -types, so they are JSON-encoded into `attr_str`. +types, so they are JSON-encoded into `attr_str`. Bytes *nested inside* such a value +are base64-encoded, following proto3's canonical JSON mapping — there is no way to put +raw bytes in a JSON string. This is not the case the plan rejected earlier: top-level +`bytes` attributes never pass through JSON, they keep their raw bytes in `attr_bytes`. ## Dependencies ```toml arrow = { version = "59", features = ["ffi"] } prost = "0.14" +prost-types = "0.14" # prost_types::Timestamp, reached through TraceItem +base64 = "0.22" # bytes nested in recursive attribute values sentry_protos = "0.70" sentry-kafka-schemas = { version = "3", default-features = false } ``` +`prost-types` and `base64` were added in phase 3; both were already in the lock file +transitively, so neither costs build time. + No new Python dependencies; `pyarrow` is deliberately not added. --- @@ -252,7 +268,8 @@ extractor down to one file plus one registry line. `HashMap` with nondeterministic iteration order; unsorted output makes batches irreproducible and tests flaky. - **Unknown enum values must not panic.** `TraceItemType::try_from(i32)` fails on a value - from a newer producer; render `TYPE_UNKNOWN_`. Enum additions are routine + from a newer producer; render `TRACE_ITEM_TYPE_UNKNOWN_` (the enum's own + prefix, so string comparisons downstream stay uniform). Enum additions are routine forward-compatible producer changes and crashing on them would be a self-inflicted outage. *(Flagged — override if you want strictness.)* - **Timestamps:** `prost_types::Timestamp { seconds, nanos }` → diff --git a/sentry_streams/src/extractors/mod.rs b/sentry_streams/src/extractors/mod.rs new file mode 100644 index 00000000..6d68b4e3 --- /dev/null +++ b/sentry_streams/src/extractors/mod.rs @@ -0,0 +1,138 @@ +//! Hand-written protobuf -> Arrow extractors, one per message type. +//! +//! Each extractor owns a hardcoded Arrow schema and knows how to turn a batch of +//! encoded payloads into a `RecordBatch`. This is the PoC trade recorded in +//! `docs/design/arrow-batch-parser.md`: Rust has no equivalent of Python's +//! reflective `ProtobufCodec`, so rather than carry a descriptor pool we write +//! the mapping by hand and accept that adding a column is a code change. +//! +//! Adding a message type is one new file plus one line in [`REGISTRY`]. +//! +//! Extractors are indexed by the raw resource string from `sentry-kafka-schemas` +//! (for example `sentry_protos.snuba.v1.trace_item_pb2.TraceItem`), so several +//! topics sharing a schema share one extractor. + +pub mod trace_item; + +use arrow::array::RecordBatch; +use arrow::datatypes::SchemaRef; +use std::collections::BTreeMap; +use std::fmt; +use std::sync::LazyLock; + +/// Why a batch could not be turned into a `RecordBatch`. +/// +/// Every variant is fatal for the whole batch: offsets are collapsed to max per +/// partition on flush, so there is no `(partition, offset)` to dead-letter a +/// single row with. The step turns these into a panic; see decision 14. +#[derive(Debug)] +pub enum ExtractorError { + /// A payload was not valid protobuf. `index` is the row's position in the + /// batch, which is the most specific thing we can say about it. + Decode { + index: usize, + source: prost::DecodeError, + }, + /// Arrow rejected the assembled columns. + Build(arrow::error::ArrowError), + /// A field held a value we cannot represent, for example a timestamp that + /// overflows microseconds. + Field { field: &'static str, detail: String }, +} + +impl fmt::Display for ExtractorError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ExtractorError::Decode { index, source } => { + write!(f, "row {index} is not valid protobuf: {source}") + } + ExtractorError::Build(e) => write!(f, "could not build the Arrow record batch: {e}"), + ExtractorError::Field { field, detail } => { + write!(f, "field `{field}` could not be represented: {detail}") + } + } + } +} + +impl std::error::Error for ExtractorError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + ExtractorError::Decode { source, .. } => Some(source), + ExtractorError::Build(e) => Some(e), + ExtractorError::Field { .. } => None, + } + } +} + +impl From for ExtractorError { + fn from(e: arrow::error::ArrowError) -> Self { + ExtractorError::Build(e) + } +} + +/// Decodes a batch of payloads of one message type into an Arrow `RecordBatch`. +pub trait Extractor: Send + Sync { + /// The `sentry-kafka-schemas` resource string this extractor handles. + fn resource(&self) -> &'static str; + + /// The Arrow schema of every batch this extractor produces, including for an + /// empty batch. + fn schema(&self) -> SchemaRef; + + /// Decode every payload and assemble one `RecordBatch`. + /// + /// Batch-wise by design: one Arrow builder per column fed across all rows and + /// finished once, rather than a batch per row. That is the shape Arrow + /// builders want, and it keeps a second extractor down to one file. + fn extract(&self, payloads: &[&[u8]]) -> Result; +} + +static TRACE_ITEM: trace_item::TraceItemExtractor = trace_item::TraceItemExtractor; + +/// Every extractor the runtime knows about, by resource string. +static REGISTRY: LazyLock> = LazyLock::new(|| { + let extractors: [&'static dyn Extractor; 1] = [&TRACE_ITEM]; + extractors.into_iter().map(|e| (e.resource(), e)).collect() +}); + +/// The extractor for `resource`, or `None` if the message type is not supported. +pub fn get_extractor(resource: &str) -> Option<&'static dyn Extractor> { + REGISTRY.get(resource).copied() +} + +/// Every supported resource string, for error messages that tell the operator +/// what they could have used instead. +pub fn registered_resources() -> Vec<&'static str> { + REGISTRY.keys().copied().collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn trace_item_is_registered_under_its_resource_string() { + let resource = "sentry_protos.snuba.v1.trace_item_pb2.TraceItem"; + let extractor = get_extractor(resource).expect("TraceItem must be registered"); + assert_eq!(extractor.resource(), resource); + } + + /// The registry key must match what `sentry-kafka-schemas` actually reports + /// for the topic, otherwise lookup silently fails at startup. + #[test] + fn registry_key_matches_the_schema_registry() { + let schema = sentry_kafka_schemas::get_schema("snuba-items", None).unwrap(); + assert!( + get_extractor(schema.raw_schema()).is_some(), + "snuba-items reports resource {:?}, which is not registered; known: {:?}", + schema.raw_schema(), + registered_resources() + ); + } + + #[test] + fn unknown_resource_has_no_extractor() { + assert!(get_extractor("nope.NotAThing").is_none()); + assert!(!registered_resources().is_empty()); + } +} diff --git a/sentry_streams/src/extractors/trace_item.rs b/sentry_streams/src/extractors/trace_item.rs new file mode 100644 index 00000000..f943c8b5 --- /dev/null +++ b/sentry_streams/src/extractors/trace_item.rs @@ -0,0 +1,766 @@ +//! Extractor for `sentry_protos.snuba.v1.TraceItem`. +//! +//! The Arrow schema is hardcoded here; see `docs/design/arrow-batch-parser.md` +//! for the column-by-column mapping and the reasoning behind it. +//! +//! Two things are worth knowing before changing this file: +//! +//! * **Nullability follows protobuf presence, not the `optional` keyword.** +//! Message-typed fields (`timestamp`, `received`) always have explicit presence +//! in proto3, so they are nullable columns. Implicit-presence scalars cannot +//! distinguish "unset" from "zero", so they are non-nullable columns carrying +//! the default -- `conversation_id` and `session_id` included, despite their +//! "if any" comments in the proto. +//! * **`map` 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. One consequence: an +//! attribute that changes type between messages lands in different columns. + +use crate::extractors::{Extractor, ExtractorError}; +use arrow::array::{ + ArrayRef, BinaryBuilder, BooleanBuilder, Float64Builder, Int64Builder, MapBuilder, + MapFieldNames, RecordBatch, StringBuilder, TimestampMicrosecondBuilder, UInt32Builder, + UInt64Builder, +}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use base64::engine::general_purpose::STANDARD as BASE64; +use base64::Engine as _; +use prost::Message; +use prost_types::Timestamp; +use sentry_protos::snuba::v1::{any_value::Value, AnyValue, TraceItem, TraceItemType}; +use serde_json::{Map as JsonMap, Value as Json}; +use std::sync::{Arc, LazyLock}; + +pub struct TraceItemExtractor; + +pub const RESOURCE: &str = "sentry_protos.snuba.v1.trace_item_pb2.TraceItem"; + +const TIMESTAMP_TZ: &str = "UTC"; + +/// Arrow spec names for map entries. arrow-rs defaults to `entries`/`keys`/`values`; +/// the spec (and therefore pyarrow and polars) wants `entries`/`key`/`value`. +/// Getting this wrong breaks interop in confusing ways rather than loudly. +fn map_field_names() -> MapFieldNames { + MapFieldNames { + entry: "entries".to_string(), + key: "key".to_string(), + value: "value".to_string(), + } +} + +/// The exact `DataType` a [`MapBuilder`] with [`map_field_names`] produces. +/// Declared explicitly so the schema and the built columns cannot drift apart. +fn map_type(value_type: DataType) -> DataType { + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", value_type, true), + ] + .into(), + ), + false, + )), + false, + ) +} + +fn timestamp_type() -> DataType { + DataType::Timestamp(TimeUnit::Microsecond, Some(TIMESTAMP_TZ.into())) +} + +static SCHEMA: LazyLock = LazyLock::new(|| { + Arc::new(Schema::new(vec![ + Field::new("organization_id", DataType::UInt64, false), + Field::new("project_id", DataType::UInt64, false), + Field::new("trace_id", DataType::Utf8, false), + Field::new("item_id", DataType::Binary, false), + Field::new("item_type", DataType::Utf8, false), + Field::new("timestamp", timestamp_type(), true), + Field::new("client_sample_rate", DataType::Float64, false), + Field::new("server_sample_rate", DataType::Float64, false), + Field::new("conversation_id", DataType::Utf8, false), + Field::new("session_id", DataType::Utf8, false), + Field::new("retention_days", DataType::UInt32, false), + Field::new("received", timestamp_type(), true), + Field::new("downsampled_retention_days", DataType::UInt32, false), + Field::new("attr_str", map_type(DataType::Utf8), false), + Field::new("attr_int", map_type(DataType::Int64), false), + Field::new("attr_double", map_type(DataType::Float64), false), + Field::new("attr_bool", map_type(DataType::Boolean), false), + Field::new("attr_bytes", map_type(DataType::Binary), false), + ])) +}); + +/// Microseconds since the epoch. +/// +/// `checked_*` rather than a silent wrap: a bogus timestamp should be a loud +/// error, not a row that quietly claims to be from the year 1754. +fn timestamp_micros(ts: &Timestamp, field: &'static str) -> Result { + ts.seconds + .checked_mul(1_000_000) + .and_then(|s| s.checked_add(i64::from(ts.nanos / 1_000))) + .ok_or_else(|| ExtractorError::Field { + field, + detail: format!( + "{}s + {}ns overflows microseconds since the epoch", + ts.seconds, ts.nanos + ), + }) +} + +/// The enum's protobuf name, or a rendered placeholder for a value we do not +/// know about. +/// +/// A newer producer adding an enum member is a routine forward-compatible +/// change; panicking on it would turn someone else's deploy into our outage. +fn item_type_name(value: i32) -> String { + match TraceItemType::try_from(value) { + Ok(t) => t.as_str_name().to_string(), + Err(_) => format!("TRACE_ITEM_TYPE_UNKNOWN_{value}"), + } +} + +/// `ArrayValue` and `KeyValueList` are recursive and Arrow has no recursive type, +/// so they are flattened to JSON and stored in `attr_str`. +/// +/// Bytes nested inside such a value are base64-encoded, following the proto3 +/// canonical JSON mapping. (Top-level `bytes` attributes do *not* go through +/// here: they keep their raw bytes in `attr_bytes`.) +fn any_value_to_json(value: &AnyValue) -> Json { + match &value.value { + None => Json::Null, + Some(Value::StringValue(s)) => Json::String(s.clone()), + Some(Value::BoolValue(b)) => Json::Bool(*b), + Some(Value::IntValue(i)) => Json::Number((*i).into()), + Some(Value::DoubleValue(d)) => serde_json::Number::from_f64(*d) + .map(Json::Number) + .unwrap_or(Json::Null), + Some(Value::BytesValue(b)) => Json::String(BASE64.encode(b)), + Some(Value::ArrayValue(a)) => Json::Array(a.values.iter().map(any_value_to_json).collect()), + Some(Value::KvlistValue(kv)) => { + let mut out = JsonMap::with_capacity(kv.values.len()); + for entry in &kv.values { + let v = entry + .value + .as_ref() + .map(any_value_to_json) + .unwrap_or(Json::Null); + out.insert(entry.key.clone(), v); + } + Json::Object(out) + } + } +} + +/// The five type-split attribute map builders. +struct AttributeBuilders { + str_: MapBuilder, + int: MapBuilder, + double: MapBuilder, + bool_: MapBuilder, + bytes: MapBuilder, +} + +impl AttributeBuilders { + fn new() -> Self { + Self { + str_: MapBuilder::new( + Some(map_field_names()), + StringBuilder::new(), + StringBuilder::new(), + ), + int: MapBuilder::new( + Some(map_field_names()), + StringBuilder::new(), + Int64Builder::new(), + ), + double: MapBuilder::new( + Some(map_field_names()), + StringBuilder::new(), + Float64Builder::new(), + ), + bool_: MapBuilder::new( + Some(map_field_names()), + StringBuilder::new(), + BooleanBuilder::new(), + ), + bytes: MapBuilder::new( + Some(map_field_names()), + StringBuilder::new(), + BinaryBuilder::new(), + ), + } + } + + /// Append one row's attributes. + /// + /// Keys are sorted first: prost decodes `map` into a + /// `HashMap`, whose iteration order varies run to run, and an unsorted batch + /// would not be reproducible. + fn append_row(&mut self, item: &TraceItem) -> Result<(), ExtractorError> { + let mut keys: Vec<&String> = item.attributes.keys().collect(); + keys.sort_unstable(); + + for key in keys { + let value = &item.attributes[key]; + match &value.value { + Some(Value::StringValue(s)) => { + self.str_.keys().append_value(key); + self.str_.values().append_value(s); + } + Some(Value::IntValue(i)) => { + self.int.keys().append_value(key); + self.int.values().append_value(*i); + } + Some(Value::DoubleValue(d)) => { + self.double.keys().append_value(key); + self.double.values().append_value(*d); + } + Some(Value::BoolValue(b)) => { + self.bool_.keys().append_value(key); + self.bool_.values().append_value(*b); + } + Some(Value::BytesValue(b)) => { + self.bytes.keys().append_value(key); + self.bytes.values().append_value(b); + } + Some(Value::ArrayValue(_)) | Some(Value::KvlistValue(_)) => { + self.str_.keys().append_value(key); + self.str_ + .values() + .append_value(any_value_to_json(value).to_string()); + } + // An attribute whose oneof is unset carries no information. + None => {} + } + } + + // Every builder must be closed on every row, including the ones that got + // no entries; skipping one silently misaligns all later rows. + self.str_.append(true)?; + self.int.append(true)?; + self.double.append(true)?; + self.bool_.append(true)?; + self.bytes.append(true)?; + Ok(()) + } + + fn finish(mut self) -> [ArrayRef; 5] { + [ + Arc::new(self.str_.finish()), + Arc::new(self.int.finish()), + Arc::new(self.double.finish()), + Arc::new(self.bool_.finish()), + Arc::new(self.bytes.finish()), + ] + } +} + +impl Extractor for TraceItemExtractor { + fn resource(&self) -> &'static str { + RESOURCE + } + + fn schema(&self) -> SchemaRef { + SCHEMA.clone() + } + + fn extract(&self, payloads: &[&[u8]]) -> Result { + let n = payloads.len(); + + let mut organization_id = UInt64Builder::with_capacity(n); + let mut project_id = UInt64Builder::with_capacity(n); + let mut trace_id = StringBuilder::with_capacity(n, n * 32); + let mut item_id = BinaryBuilder::with_capacity(n, n * 16); + let mut item_type = StringBuilder::with_capacity(n, n * 24); + let mut timestamp = TimestampMicrosecondBuilder::with_capacity(n); + let mut client_sample_rate = Float64Builder::with_capacity(n); + let mut server_sample_rate = Float64Builder::with_capacity(n); + let mut conversation_id = StringBuilder::with_capacity(n, n * 16); + let mut session_id = StringBuilder::with_capacity(n, n * 16); + let mut retention_days = UInt32Builder::with_capacity(n); + let mut received = TimestampMicrosecondBuilder::with_capacity(n); + let mut downsampled_retention_days = UInt32Builder::with_capacity(n); + let mut attributes = AttributeBuilders::new(); + + for (index, payload) in payloads.iter().enumerate() { + let item = TraceItem::decode(*payload) + .map_err(|source| ExtractorError::Decode { index, source })?; + + organization_id.append_value(item.organization_id); + project_id.append_value(item.project_id); + trace_id.append_value(&item.trace_id); + item_id.append_value(&item.item_id); + item_type.append_value(item_type_name(item.item_type)); + timestamp.append_option( + item.timestamp + .as_ref() + .map(|t| timestamp_micros(t, "timestamp")) + .transpose()?, + ); + client_sample_rate.append_value(item.client_sample_rate); + server_sample_rate.append_value(item.server_sample_rate); + conversation_id.append_value(&item.conversation_id); + session_id.append_value(&item.session_id); + retention_days.append_value(item.retention_days); + received.append_option( + item.received + .as_ref() + .map(|t| timestamp_micros(t, "received")) + .transpose()?, + ); + downsampled_retention_days.append_value(item.downsampled_retention_days); + attributes.append_row(&item)?; + } + + let [attr_str, attr_int, attr_double, attr_bool, attr_bytes] = attributes.finish(); + + let columns: Vec = vec![ + Arc::new(organization_id.finish()), + Arc::new(project_id.finish()), + Arc::new(trace_id.finish()), + Arc::new(item_id.finish()), + Arc::new(item_type.finish()), + Arc::new(timestamp.finish().with_timezone(TIMESTAMP_TZ)), + Arc::new(client_sample_rate.finish()), + Arc::new(server_sample_rate.finish()), + Arc::new(conversation_id.finish()), + Arc::new(session_id.finish()), + Arc::new(retention_days.finish()), + Arc::new(received.finish().with_timezone(TIMESTAMP_TZ)), + Arc::new(downsampled_retention_days.finish()), + attr_str, + attr_int, + attr_double, + attr_bool, + attr_bytes, + ]; + + Ok(RecordBatch::try_new(SCHEMA.clone(), columns)?) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + Array, BinaryArray, BooleanArray, Float64Array, Int64Array, MapArray, StringArray, + TimestampMicrosecondArray, UInt32Array, UInt64Array, + }; + use prost::Message; + use prost_types::Timestamp; + use sentry_protos::snuba::v1::{ + any_value::Value, AnyValue, ArrayValue, KeyValue, KeyValueList, TraceItem, + }; + use std::collections::HashMap; + + fn extract(items: &[TraceItem]) -> RecordBatch { + let encoded: Vec> = items.iter().map(|i| i.encode_to_vec()).collect(); + let refs: Vec<&[u8]> = encoded.iter().map(|v| v.as_slice()).collect(); + TraceItemExtractor.extract(&refs).expect("extract") + } + + fn attr(kind: Value) -> AnyValue { + AnyValue { value: Some(kind) } + } + + /// Reads one row of a `Map` column as (key, value) pairs. + fn map_row(batch: &RecordBatch, column: &str, row: usize) -> Vec<(String, String)> { + let col = batch + .column_by_name(column) + .unwrap_or_else(|| panic!("no column {column}")) + .as_any() + .downcast_ref::() + .expect("map column"); + assert!(col.is_valid(row), "map rows are never null"); + let entries = col.value(row); + let keys = entries + .column(0) + .as_any() + .downcast_ref::() + .expect("map keys are Utf8"); + let values = entries.column(1); + + (0..entries.len()) + .map(|i| (keys.value(i).to_string(), scalar_to_string(values, i))) + .collect() + } + + fn scalar_to_string(values: &arrow::array::ArrayRef, i: usize) -> String { + use arrow::datatypes::DataType; + match values.data_type() { + DataType::Utf8 => values + .as_any() + .downcast_ref::() + .unwrap() + .value(i) + .to_string(), + DataType::Int64 => values + .as_any() + .downcast_ref::() + .unwrap() + .value(i) + .to_string(), + DataType::Float64 => values + .as_any() + .downcast_ref::() + .unwrap() + .value(i) + .to_string(), + DataType::Boolean => values + .as_any() + .downcast_ref::() + .unwrap() + .value(i) + .to_string(), + DataType::Binary => { + let b = values.as_any().downcast_ref::().unwrap(); + String::from_utf8_lossy(b.value(i)).to_string() + } + other => panic!("unexpected map value type {other:?}"), + } + } + + fn str_col<'a>(batch: &'a RecordBatch, name: &str) -> &'a StringArray { + batch + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + } + + fn ts_col<'a>(batch: &'a RecordBatch, name: &str) -> &'a TimestampMicrosecondArray { + batch + .column_by_name(name) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + } + + #[test] + fn every_scalar_field_lands_in_its_column() { + let item = TraceItem { + organization_id: 42, + project_id: 7, + trace_id: "abc123".into(), + item_id: vec![1, 2, 3, 4], + item_type: 1, // TRACE_ITEM_TYPE_SPAN + timestamp: Some(Timestamp { + seconds: 1_700_000_000, + nanos: 123_456_000, + }), + client_sample_rate: 0.5, + server_sample_rate: 0.25, + conversation_id: "conv".into(), + session_id: "sess".into(), + retention_days: 90, + received: Some(Timestamp { + seconds: 1_700_000_001, + nanos: 0, + }), + downsampled_retention_days: 30, + ..Default::default() + }; + + let batch = extract(&[item]); + assert_eq!(batch.num_rows(), 1); + + let u64_col = |n: &str| { + batch + .column_by_name(n) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + }; + let u32_col = |n: &str| { + batch + .column_by_name(n) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + }; + let f64_col = |n: &str| { + batch + .column_by_name(n) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0) + }; + + assert_eq!(u64_col("organization_id"), 42); + assert_eq!(u64_col("project_id"), 7); + assert_eq!(str_col(&batch, "trace_id").value(0), "abc123"); + assert_eq!( + batch + .column_by_name("item_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + &[1, 2, 3, 4] + ); + assert_eq!( + str_col(&batch, "item_type").value(0), + "TRACE_ITEM_TYPE_SPAN" + ); + assert_eq!( + ts_col(&batch, "timestamp").value(0), + 1_700_000_000_123_456_i64 + ); + assert_eq!(f64_col("client_sample_rate"), 0.5); + assert_eq!(f64_col("server_sample_rate"), 0.25); + assert_eq!(str_col(&batch, "conversation_id").value(0), "conv"); + assert_eq!(str_col(&batch, "session_id").value(0), "sess"); + assert_eq!(u32_col("retention_days"), 90); + assert_eq!( + ts_col(&batch, "received").value(0), + 1_700_000_001_000_000_i64 + ); + assert_eq!(u32_col("downsampled_retention_days"), 30); + } + + #[test] + fn each_any_value_arm_lands_in_its_own_map() { + let item = TraceItem { + attributes: HashMap::from([ + ("s".into(), attr(Value::StringValue("hello".into()))), + ("b".into(), attr(Value::BoolValue(true))), + ("i".into(), attr(Value::IntValue(-9))), + ("d".into(), attr(Value::DoubleValue(1.5))), + ("y".into(), attr(Value::BytesValue(b"raw".to_vec()))), + ]), + ..Default::default() + }; + + let batch = extract(&[item]); + assert_eq!( + map_row(&batch, "attr_str", 0), + vec![("s".into(), "hello".into())] + ); + assert_eq!( + map_row(&batch, "attr_bool", 0), + vec![("b".into(), "true".into())] + ); + assert_eq!( + map_row(&batch, "attr_int", 0), + vec![("i".into(), "-9".into())] + ); + assert_eq!( + map_row(&batch, "attr_double", 0), + vec![("d".into(), "1.5".into())] + ); + assert_eq!( + map_row(&batch, "attr_bytes", 0), + vec![("y".into(), "raw".into())] + ); + } + + #[test] + fn recursive_values_are_json_encoded_into_attr_str() { + let item = TraceItem { + attributes: HashMap::from([ + ( + "arr".into(), + attr(Value::ArrayValue(ArrayValue { + values: vec![ + attr(Value::IntValue(1)), + attr(Value::StringValue("two".into())), + ], + })), + ), + ( + "kv".into(), + attr(Value::KvlistValue(KeyValueList { + values: vec![KeyValue { + key: "inner".into(), + value: Some(attr(Value::BoolValue(false))), + }], + })), + ), + ]), + ..Default::default() + }; + + let batch = extract(&[item]); + let row: HashMap = map_row(&batch, "attr_str", 0).into_iter().collect(); + assert_eq!(row["arr"], r#"[1,"two"]"#); + assert_eq!(row["kv"], r#"{"inner":false}"#); + } + + #[test] + fn absent_message_fields_are_null_not_epoch_zero() { + let batch = extract(&[TraceItem::default()]); + assert!(ts_col(&batch, "timestamp").is_null(0)); + assert!(ts_col(&batch, "received").is_null(0)); + } + + /// `conversation_id` and `session_id` are plain proto3 strings in + /// sentry_protos 0.70: they have implicit presence, so "unset" and "empty" + /// are indistinguishable and the column is non-nullable. + #[test] + fn implicit_presence_scalars_carry_defaults_and_are_never_null() { + let batch = extract(&[TraceItem::default()]); + for name in [ + "organization_id", + "project_id", + "trace_id", + "item_id", + "item_type", + "client_sample_rate", + "server_sample_rate", + "conversation_id", + "session_id", + "retention_days", + "downsampled_retention_days", + ] { + let col = batch.column_by_name(name).unwrap(); + assert!(!col.is_null(0), "{name} must not be null"); + assert!( + !batch.schema().field_with_name(name).unwrap().is_nullable(), + "{name} must be a non-nullable column" + ); + } + assert_eq!(str_col(&batch, "conversation_id").value(0), ""); + assert_eq!( + str_col(&batch, "item_type").value(0), + "TRACE_ITEM_TYPE_UNSPECIFIED" + ); + } + + /// A newer producer sending an enum value we do not know about is a routine + /// forward-compatible change. Crashing on it would be a self-inflicted outage. + #[test] + fn unknown_item_type_renders_rather_than_panicking() { + let item = TraceItem { + item_type: 31337, + ..Default::default() + }; + let batch = extract(&[item]); + assert_eq!( + str_col(&batch, "item_type").value(0), + "TRACE_ITEM_TYPE_UNKNOWN_31337" + ); + } + + /// prost decodes `map` into a HashMap with nondeterministic + /// iteration order. Without sorting, batches would not be reproducible. + #[test] + fn attribute_order_is_deterministic() { + let forward = TraceItem { + attributes: HashMap::from([ + ("a".into(), attr(Value::IntValue(1))), + ("b".into(), attr(Value::IntValue(2))), + ("c".into(), attr(Value::IntValue(3))), + ]), + ..Default::default() + }; + let first = extract(&[forward.clone()]); + let second = extract(&[forward]); + assert_eq!(first, second); + assert_eq!( + map_row(&first, "attr_int", 0), + vec![ + ("a".into(), "1".into()), + ("b".into(), "2".into()), + ("c".into(), "3".into()) + ] + ); + } + + /// Every map builder must be appended to on every row, including rows with no + /// attributes of that type; skipping one silently misaligns all later rows. + #[test] + fn map_offsets_stay_aligned_across_mixed_rows() { + let none = TraceItem::default(); + let one = TraceItem { + attributes: HashMap::from([("only".into(), attr(Value::IntValue(1)))]), + ..Default::default() + }; + let many = TraceItem { + attributes: HashMap::from([ + ("x".into(), attr(Value::IntValue(10))), + ("y".into(), attr(Value::IntValue(20))), + ("z".into(), attr(Value::StringValue("s".into()))), + ]), + ..Default::default() + }; + + let batch = extract(&[none, one, many]); + assert_eq!(batch.num_rows(), 3); + assert_eq!(map_row(&batch, "attr_int", 0), vec![]); + assert_eq!( + map_row(&batch, "attr_int", 1), + vec![("only".into(), "1".into())] + ); + assert_eq!( + map_row(&batch, "attr_int", 2), + vec![("x".into(), "10".into()), ("y".into(), "20".into())] + ); + assert_eq!( + map_row(&batch, "attr_str", 2), + vec![("z".into(), "s".into())] + ); + assert_eq!(map_row(&batch, "attr_str", 0), vec![]); + } + + #[test] + fn truncated_payload_names_the_row() { + let good = TraceItem { + organization_id: 1, + ..Default::default() + } + .encode_to_vec(); + let bad = vec![0xffu8, 0xff, 0xff]; + let payloads: Vec<&[u8]> = vec![good.as_slice(), bad.as_slice()]; + + match TraceItemExtractor.extract(&payloads) { + Err(ExtractorError::Decode { index, .. }) => assert_eq!(index, 1), + other => panic!("expected a decode error naming row 1, got {other:?}"), + } + } + + #[test] + fn timestamp_overflow_is_an_error_not_a_silent_wrap() { + let item = TraceItem { + timestamp: Some(Timestamp { + seconds: i64::MAX, + nanos: 0, + }), + ..Default::default() + }; + let encoded = item.encode_to_vec(); + match TraceItemExtractor.extract(&[encoded.as_slice()]) { + Err(ExtractorError::Field { field, .. }) => assert_eq!(field, "timestamp"), + other => panic!("expected a field error, got {other:?}"), + } + } + + #[test] + fn empty_batch_still_carries_the_schema() { + let batch = TraceItemExtractor.extract(&[]).expect("empty batch"); + assert_eq!(batch.num_rows(), 0); + assert_eq!(batch.schema(), TraceItemExtractor.schema()); + assert_eq!( + batch.num_columns(), + TraceItemExtractor.schema().fields().len() + ); + } + + #[test] + fn produced_batch_matches_the_declared_schema() { + let batch = extract(&[TraceItem::default()]); + assert_eq!(batch.schema(), TraceItemExtractor.schema()); + } +} diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index d9def81a..64e1853d 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -6,6 +6,10 @@ mod commit_policy; mod committable; mod consumer; mod dev_null_sink; +// Consumed by the Arrow batch parser step (phase 4). Until that lands nothing +// outside this module's own tests calls into it. +#[allow(dead_code)] +mod extractors; mod filter_step; mod gcs_writer; mod header_filter_step; From fd79874bcf9a5535e6304860c4d8aa6fb3740e83 Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 16:03:29 -0700 Subject: [PATCH 06/13] feat(arrow): the Arrow batch parser step (phase 4) Adds ArrowFlushProducer, which decodes a flushed window into a RecordBatch and hands it to Python as an ArrowRecordBatch inside a PyAnyMessage, and build_arrow_batch_parser_step, which pairs it with a BatchStep. Windowing, watermark ordering and backpressure are the Batch step's, untouched. Topic -> schema -> extractor is resolved once, at construction. Each failure -- unknown topic, non-protobuf schema, no registered extractor -- panics at startup naming the step, the topic and what was actually found, because none of them will fix themselves at runtime. It also matters that get_schema runs exactly once: it leaks on protobuf topics. with_payloads is the only place that knows where payload bytes live. Today it takes the GIL and borrows from Py; when the source starts emitting Rust-native messages it loses the GIL block and the PyRef guards and nothing else changes. It is a scope rather than an accessor because the guards must outlive the slices, and it deliberately does not copy the payloads to release the GIL early -- that would work today and become a per-message copy in the one step built to eliminate per-message copies. A malformed payload panics rather than dead-lettering: offsets collapse to max per partition on flush, so there is no (partition, offset) for arroyo's DLQ to reject a row with. A PyAnyMessage in the window panics too -- it means a Python step ran in between and the bytes are gone. The adapter will reject that at build time in phase 5; this is the backstop. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- sentry_streams/src/arrow_batch_parser.rs | 388 +++++++++++++++++++++++ sentry_streams/src/lib.rs | 8 +- sentry_streams/src/py_record_batch.rs | 3 - 3 files changed, 394 insertions(+), 5 deletions(-) create mode 100644 sentry_streams/src/arrow_batch_parser.rs diff --git a/sentry_streams/src/arrow_batch_parser.rs b/sentry_streams/src/arrow_batch_parser.rs new file mode 100644 index 00000000..2128b287 --- /dev/null +++ b/sentry_streams/src/arrow_batch_parser.rs @@ -0,0 +1,388 @@ +//! The Arrow batch parser step: batches raw Kafka payloads and decodes them into +//! an Apache Arrow `RecordBatch` without ever materialising them as Python +//! objects. +//! +//! It reuses [`BatchStep`] wholesale -- windowing, watermark ordering and +//! backpressure are identical to the `Batch` step -- and supplies its own +//! [`BatchFlushProducer`]. See `docs/design/arrow-batch-parser.md`, phase 4. + +use crate::batch_step::{BatchElement, BatchFlushProducer, BatchStep}; +use crate::extractors::{get_extractor, registered_resources, Extractor}; +use crate::messages::{into_pyany, PyAnyMessage, PyStreamingMessage, RoutedValuePayload}; +use crate::py_record_batch::PyRecordBatch; +use crate::routes::{Route, RoutedValue}; +use crate::utils::traced_with_gil; +use pyo3::prelude::*; +use sentry_arroyo::processing::strategies::{ProcessingStrategy, StrategyError}; +use sentry_arroyo::types::{Message, Partition}; +use sentry_kafka_schemas::{get_schema, SchemaType}; +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +/// Runs `f` over the batch's payload bytes. +/// +/// **This is the only place that knows where payload bytes live.** Today the +/// source boxes every payload into a `Py` (`consumer.rs`), so the +/// bytes are Python-owned and reading them needs the GIL; the borrow is held for +/// the duration of the decode. When the source starts emitting Rust-native +/// messages this function loses its GIL block and its `PyRef` guards, and +/// nothing else in the step changes. +/// +/// It is a scope rather than a plain accessor because the `PyRef` guards must +/// outlive the slices handed to `f`. +/// +/// Do **not** copy the payloads out to release the GIL sooner. It would work +/// today and would become permanent dead weight the moment the source goes +/// native -- a per-message copy in the one step whose whole purpose is to remove +/// per-message copies. +fn with_payloads( + step_name: &str, + elements: &[BatchElement], + f: impl FnOnce(&[&[u8]]) -> R, +) -> R { + traced_with_gil!(|py| { + let guards: Vec> = elements + .iter() + .map(|element| match element { + PyStreamingMessage::RawMessage { content } => content.bind(py).borrow(), + // Decision 10: this step reads bytes off the wire. A PyAnyMessage + // means a Python step ran in between and the bytes are gone. + // The adapter rejects this at build time; this is the backstop. + PyStreamingMessage::PyAnyMessage { .. } => panic!( + "step '{step_name}': the Arrow batch parser only accepts raw messages, \ + but the window contains a message already converted to a Python object. \ + Place this step directly after the source." + ), + }) + .collect(); + + let payloads: Vec<&[u8]> = guards.iter().map(|g| g.payload.as_slice()).collect(); + f(&payloads) + }) +} + +/// Decodes a flushed window into a `RecordBatch` and hands it to Python. +pub(crate) struct ArrowFlushProducer { + extractor: &'static dyn Extractor, + step_name: String, + schema_name: String, +} + +impl ArrowFlushProducer { + /// Resolve topic -> schema -> extractor once, at step construction. + /// + /// Every failure here is a configuration error that will never fix itself at + /// runtime, so each one panics at startup rather than at the first message. + /// `get_schema` leaks on protobuf topics and must never be called per + /// message, which is the other reason this happens exactly once. + fn resolve(step_name: String, schema_name: &str) -> Self { + let schema = get_schema(schema_name, None).unwrap_or_else(|e| { + panic!( + "step '{step_name}': no schema registered for topic '{schema_name}': {e}. \ + The Arrow batch parser resolves its extractor from the source topic's schema." + ) + }); + + if schema.schema_type != SchemaType::Protobuf { + panic!( + "step '{step_name}': topic '{schema_name}' has schema type {:?}, but the Arrow \ + batch parser supports protobuf only.", + schema.schema_type + ); + } + + let resource = schema.raw_schema(); + let extractor = get_extractor(resource).unwrap_or_else(|| { + panic!( + "step '{step_name}': topic '{schema_name}' carries '{resource}', which has no \ + Arrow extractor. Known message types: {:?}", + registered_resources() + ) + }); + + Self { + extractor, + step_name, + schema_name: schema_name.to_string(), + } + } +} + +impl BatchFlushProducer for ArrowFlushProducer { + fn produce( + &self, + route: &Route, + elements: &[BatchElement], + committable: BTreeMap, + ) -> Result, StrategyError> { + let ts = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_secs_f64()) + .unwrap_or(0.0); + + let batch = with_payloads(&self.step_name, elements, |payloads| { + self.extractor.extract(payloads) + }) + .unwrap_or_else(|e| { + // Offsets are collapsed to max per partition on flush, so there is no + // (partition, offset) for arroyo's DLQ to reject a single row with, + // and no way to fail the batch without failing the window. Panicking + // matches what the runtime already does for an error on an + // AnyMessage; see transformer.rs. + panic!( + "step '{}': could not decode a batch of {} message(s) from topic '{}': {e}", + self.step_name, + elements.len(), + self.schema_name, + ) + }); + + let content = traced_with_gil!(|py| -> PyResult> { + let py_batch = Py::new(py, PyRecordBatch::new(batch))?; + into_pyany( + py, + PyAnyMessage { + payload: py_batch.into_any(), + headers: vec![], + timestamp: ts, + schema: Some(self.schema_name.clone()), + }, + ) + }) + .map_err(|e| StrategyError::Other(Box::new(e)))?; + + Ok(Message::new_any_message( + RoutedValue { + route: route.clone(), + payload: RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::PyAnyMessage { + content, + }), + }, + committable, + )) + } +} + +pub fn build_arrow_batch_parser_step( + route: &Route, + schema_name: &str, + step_name: String, + max_batch_size: Option, + max_batch_time: Option, + next: Box>, +) -> Box> { + let producer = ArrowFlushProducer::resolve(step_name.clone(), schema_name); + Box::new(BatchStep::with_producer( + route.clone(), + max_batch_size, + max_batch_time, + step_name, + next, + Arc::new(producer), + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::fake_strategy::FakeStrategy; + use crate::testutils::{build_raw_routed_value, build_routed_value}; + use arrow::array::{RecordBatch, StringArray, UInt64Array}; + use prost::Message as _; + use pyo3::types::PyAnyMethods; + use pyo3::IntoPyObject; + use sentry_arroyo::types::{Partition, Topic}; + use sentry_protos::snuba::v1::TraceItem; + use std::sync::Mutex; + + const SNUBA_ITEMS: &str = "snuba-items"; + + fn route() -> Route { + Route::new("s".into(), vec!["w".into()]) + } + + fn trace_item(org: u64) -> Vec { + TraceItem { + organization_id: org, + trace_id: format!("trace-{org}"), + ..Default::default() + } + .encode_to_vec() + } + + fn producer() -> ArrowFlushProducer { + ArrowFlushProducer::resolve("test_arrow".to_string(), SNUBA_ITEMS) + } + + /// Pull the `RecordBatch` back out of the emitted message, the way a Python + /// consumer would see it. + fn batch_of(message: Message) -> (RecordBatch, Option) { + let payload = message.into_payload(); + let content = match payload.payload { + RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::PyAnyMessage { + content, + }) => content, + _ => panic!("expected a PyAnyMessage carrying the record batch"), + }; + traced_with_gil!(|py| { + let borrowed = content.bind(py).borrow(); + let schema = borrowed.schema.clone(); + let rb: PyRef = borrowed.payload.bind(py).extract().unwrap(); + (rb.batch.clone(), schema) + }) + } + + #[test] + fn decodes_a_window_of_raw_messages_into_one_record_batch() { + crate::testutils::initialize_python(); + let partition = Partition::new(Topic::new("t"), 0); + let committable = BTreeMap::from([(partition, 11_u64)]); + + let (batch, schema) = traced_with_gil!(|py| { + let elements: Vec = [1_u64, 2, 3] + .into_iter() + .map(|org| { + match build_raw_routed_value(py, trace_item(org), "s", vec!["w".into()]).payload + { + RoutedValuePayload::PyStreamingMessage(m) => m, + _ => unreachable!(), + } + }) + .collect(); + + let message = producer() + .produce(&route(), &elements, committable.clone()) + .expect("produce"); + assert_eq!( + message.committable().collect::>(), + committable + ); + batch_of(message) + }); + + assert_eq!(batch.num_rows(), 3); + let orgs = batch + .column_by_name("organization_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(orgs.values(), &[1, 2, 3]); + let traces = batch + .column_by_name("trace_id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(traces.value(0), "trace-1"); + + // Downstream Python needs to know which stream the batch came from. + assert_eq!(schema.as_deref(), Some(SNUBA_ITEMS)); + } + + /// Decision 10: this step only accepts RawMessage. A PyAnyMessage means some + /// Python step ran in between and the payload bytes are gone. + #[test] + #[should_panic(expected = "test_arrow")] + fn a_python_payload_in_the_window_panics() { + crate::testutils::initialize_python(); + traced_with_gil!(|py| { + let payload = 1i32.into_pyobject(py).unwrap().into_any().unbind(); + let element = match build_routed_value(py, payload, "s", vec!["w".into()]).payload { + RoutedValuePayload::PyStreamingMessage(m) => m, + _ => unreachable!(), + }; + let _ = producer().produce(&route(), &[element], BTreeMap::new()); + }); + } + + /// A malformed payload cannot be dead-lettered -- offsets are collapsed to + /// max per partition -- so it fails the process. Decisions 9 and 14. + #[test] + #[should_panic(expected = "row 1")] + fn a_malformed_payload_panics_naming_the_row() { + crate::testutils::initialize_python(); + traced_with_gil!(|py| { + let payloads = vec![trace_item(1), vec![0xff, 0xff, 0xff]]; + let elements: Vec = payloads + .into_iter() + .map( + |p| match build_raw_routed_value(py, p, "s", vec!["w".into()]).payload { + RoutedValuePayload::PyStreamingMessage(m) => m, + _ => unreachable!(), + }, + ) + .collect(); + let _ = producer().produce(&route(), &elements, BTreeMap::new()); + }); + } + + #[test] + #[should_panic(expected = "not-a-real-topic")] + fn an_unknown_topic_panics_at_construction() { + ArrowFlushProducer::resolve("test_arrow".to_string(), "not-a-real-topic"); + } + + /// JSON topics are out of scope for the PoC and must fail loudly at startup + /// rather than at the first message. + #[test] + #[should_panic(expected = "Json")] + fn a_json_topic_panics_at_construction() { + let topic = "events"; + let schema = sentry_kafka_schemas::get_schema(topic, None).unwrap(); + assert_eq!( + schema.schema_type, + SchemaType::Json, + "{topic} is expected to be a JSON topic" + ); + ArrowFlushProducer::resolve("test_arrow".to_string(), topic); + } + + /// The step is a BatchStep underneath: same windowing, same watermarks. + #[test] + fn behaves_as_a_batch_step_end_to_end() { + crate::testutils::initialize_python(); + let sub = Arc::new(Mutex::new(Vec::new())); + let wms = Arc::new(Mutex::new(Vec::new())); + let mut step = build_arrow_batch_parser_step( + &route(), + SNUBA_ITEMS, + "test_arrow".to_string(), + Some(2), + None, + Box::new(FakeStrategy::new(sub.clone(), wms, false)), + ); + + traced_with_gil!(|py| { + for org in [1_u64, 2] { + let msg = Message::new_any_message( + build_raw_routed_value(py, trace_item(org), "s", vec!["w".into()]), + BTreeMap::new(), + ); + step.submit(msg).unwrap(); + } + step.poll().unwrap(); + }); + + let out = sub.lock().unwrap(); + assert_eq!(out.len(), 1, "one batch downstream, not two rows"); + traced_with_gil!(|py| { + let rb: PyRef = out[0].bind(py).extract().unwrap(); + assert_eq!(rb.batch.num_rows(), 2); + }); + } + + /// An empty window never reaches a producer, but the extractor's empty batch + /// must still be schema-correct if it ever does. + #[test] + fn an_empty_window_still_produces_a_typed_batch() { + crate::testutils::initialize_python(); + let message = producer() + .produce(&route(), &[], BTreeMap::new()) + .expect("produce"); + let (batch, _) = batch_of(message); + assert_eq!(batch.num_rows(), 0); + assert!(batch.num_columns() > 0); + } +} diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index 64e1853d..60a438b1 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -1,4 +1,8 @@ use pyo3::prelude::*; +// Reachable once `operators::build()` gains its ArrowBatchParser arm (phase 5); +// until then the whole chain is dead in the non-test build. +#[allow(dead_code)] +mod arrow_batch_parser; mod batch_step; mod broadcaster; mod callers; @@ -6,8 +10,8 @@ mod commit_policy; mod committable; mod consumer; mod dev_null_sink; -// Consumed by the Arrow batch parser step (phase 4). Until that lands nothing -// outside this module's own tests calls into it. +// Reachable once `operators::build()` gains its ArrowBatchParser arm (phase 5); +// until then the whole chain is dead in the non-test build. #[allow(dead_code)] mod extractors; mod filter_step; diff --git a/sentry_streams/src/py_record_batch.rs b/sentry_streams/src/py_record_batch.rs index 7d4ae5bf..36af8f22 100644 --- a/sentry_streams/src/py_record_batch.rs +++ b/sentry_streams/src/py_record_batch.rs @@ -38,9 +38,6 @@ pub struct PyRecordBatch { } impl PyRecordBatch { - // Constructed by the Arrow batch parser step (phase 4); until then only the - // tests build one. - #[allow(dead_code)] pub(crate) fn new(batch: RecordBatch) -> Self { Self { batch } } From 3856ccaa8a4e7c05788389336fb463527242d7ba Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 16:13:44 -0700 Subject: [PATCH 07/13] feat(arrow): DSL and adapter wiring for ArrowBatchParser (phase 5) Adds the RuntimeOperator.ArrowBatchParser variant with its build() arm, the ArrowBatchParser DSL primitive (a Reduce subclass mirroring Batch), the rust_arroyo wiring, a NotImplementedError in the pure-Python adapter, and the type stubs. schema_name travels in the operator variant because operators::build() is handed no topic. The adapter captures it from step.stream_name before override_config runs, so overriding the physical topic in a deployment config does not change which schema the extractor is resolved from. The build-time placement check is implemented differently from the design sketch, which assumed the adapter could walk the pipeline's incoming edges. reduce() only ever receives the step and its Route, so instead the adapter tracks which routes still carry raw payloads: the source marks its route raw, map and every other reduce clear it, and filters, broadcast and router propagate it, since all three forward messages untouched. Placing the step anywhere else is a ValueError naming the step, rather than a panic in production. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- .../docs/design/arrow-batch-parser.md | 13 +- .../sentry_streams/adapters/arroyo/adapter.py | 8 ++ .../adapters/arroyo/rust_arroyo.py | 79 +++++++++- .../sentry_streams/pipeline/__init__.py | 2 + .../sentry_streams/pipeline/pipeline.py | 69 +++++++++ .../sentry_streams/rust_streams.pyi | 9 ++ sentry_streams/src/lib.rs | 6 - sentry_streams/src/operators.rs | 35 +++++ .../arroyo/test_arrow_batch_parser.py | 136 ++++++++++++++++++ 9 files changed, 344 insertions(+), 13 deletions(-) create mode 100644 sentry_streams/tests/adapters/arroyo/test_arrow_batch_parser.py diff --git a/sentry_streams/docs/design/arrow-batch-parser.md b/sentry_streams/docs/design/arrow-batch-parser.md index a11cfdcf..1df2c8b5 100644 --- a/sentry_streams/docs/design/arrow-batch-parser.md +++ b/sentry_streams/docs/design/arrow-batch-parser.md @@ -399,9 +399,12 @@ mirror `Batch` (`pipeline.py:674-681`). No schema, no format, no type name. `self.__source_schemas[source_name] = schema_name`. 2. In `reduce()`, add an `isinstance(step, ArrowBatchParser)` branch before the `Batch` branch, emitting `RuntimeOperator.ArrowBatchParser(..., schema_name=self.__source_schemas[stream.source], ...)`. -3. Build-time input check: walk the pipeline's incoming edges from this step; if any - predecessor is not a `RawMessage`-preserving step (source, `HeadersFilter`), raise - naming the step and the offending predecessor. +3. Build-time input check. *(Implemented differently from the sketch: `reduce()` is + handed only the step and the `Route`, never the pipeline graph, so the walk is not + available.)* The adapter instead tracks which routes still carry raw payloads — the + source marks its route raw, `map`/`flat_map` and every other `reduce` clear it, and + filters, `broadcast` and `router` propagate it, since they forward messages + untouched. `ArrowBatchParser` raises if its route is not raw. **`adapters/arroyo/adapter.py`** — `NotImplementedError` pointing at the Rust adapter, documented Rust-only in the style of `HeadersFilter`. @@ -409,7 +412,9 @@ documented Rust-only in the style of `HeadersFilter`. **Also:** export from `pipeline/__init__.py` (both the import and `__all__`); add `RuntimeOperator.ArrowBatchParser` and `ArrowRecordBatch` to `rust_streams.pyi`. -**Tests:** placing the step after a Python `Map` fails at build time naming both steps; +**Tests:** placing the step after a Python `Map` fails at build time naming the step; +a filter between source and parser is accepted; a deployment topic override does not +change the resolved schema; a JSON topic panics at startup with the actual `schema_type`; the pure-Python adapter raises `NotImplementedError`; `make typecheck` clean. diff --git a/sentry_streams/sentry_streams/adapters/arroyo/adapter.py b/sentry_streams/sentry_streams/adapters/arroyo/adapter.py index e7e31fe5..5a11a85e 100644 --- a/sentry_streams/sentry_streams/adapters/arroyo/adapter.py +++ b/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 diff --git a/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py b/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py index 94d4e67b..a7b34b3c 100644 --- a/sentry_streams/sentry_streams/adapters/arroyo/rust_arroyo.py +++ b/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, @@ -256,6 +257,12 @@ def __init__( self.__consumers: MutableMapping[str, ArroyoConsumer] = {} self.__chains = TransformChains() self.__sentry_dsn = sentry_sdk_config.get("dsn") if sentry_sdk_config else None + # Logical stream name per source, captured before any topic override, so + # steps that resolve a schema are unaffected by deployment overrides. + self.__source_schemas: MutableMapping[str, str] = {} + # Routes whose messages still carry raw Kafka payloads. ArrowBatchParser + # reads bytes off the wire, so it can only be placed on one of these. + self.__raw_routes: set[tuple[str, tuple[str, ...]]] = set() @classmethod def build( # type: ignore[override] @@ -277,6 +284,32 @@ def __close_chain(self, stream: Route) -> None: finalize_chain(self.__chains, stream, self.__metrics_config) ) + @staticmethod + def __route_key(stream: Route) -> tuple[str, tuple[str, ...]]: + return (stream.source, tuple(stream.waypoints)) + + def __mark_route_raw(self, stream: Route) -> None: + self.__raw_routes.add(self.__route_key(stream)) + + def __mark_route_converted(self, stream: Route) -> None: + """Record that messages on this route are now Python objects. + + Steps that only forward messages -- filters, broadcast, router -- leave + the payload alone and so do not call this. + """ + self.__raw_routes.discard(self.__route_key(stream)) + + def __route_is_raw(self, stream: Route) -> bool: + return self.__route_key(stream) in self.__raw_routes + + def __propagate_raw(self, stream: Route, branches: Mapping[str, Route]) -> Mapping[str, Route]: + """Broadcast and router forward messages untouched, so each branch keeps + whatever the incoming route had.""" + if self.__route_is_raw(stream): + for branch in branches.values(): + self.__mark_route_raw(branch) + return branches + def get_consumer(self, source: str) -> ArroyoConsumer: return self.__consumers[source] @@ -321,7 +354,10 @@ def source(self, step: Source[Any]) -> Route: dlq_config=dlq_config, sentry_dsn=self.__sentry_dsn, ) - return Route(source_name, []) + self.__source_schemas[source_name] = schema_name + route = Route(source_name, []) + self.__mark_route_raw(route) + return route def sink(self, step: Sink[Any], stream: Route) -> Route: """ @@ -434,6 +470,7 @@ def map(self, step: Map[Any, Any], stream: Route) -> Route: if self.__chains.exists(stream): self.__chains.add_map(stream, step) + self.__mark_route_converted(stream) return stream def flat_map(self, step: FlatMap[Any, Any], stream: Route) -> Route: @@ -498,6 +535,40 @@ def reduce( step.override_config(loaded_config) step.validate() + if isinstance(step, ArrowBatchParser): + if not self.__route_is_raw(stream): + raise ValueError( + f"Step '{step.name}' is an ArrowBatchParser, which decodes raw Kafka " + f"payloads, but the messages reaching it on route {stream} have already " + "been converted to Python objects by an earlier step. Place it directly " + "after the source, with only filters, broadcasts or routers in between." + ) + + schema_name = self.__source_schemas.get(stream.source) + if schema_name is None: + raise ValueError( + f"Step '{step.name}': no schema recorded for source '{stream.source}'. " + "ArrowBatchParser resolves its Arrow schema from the source topic." + ) + + 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=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 + ), + ) + ) + # The batch is an ArrowRecordBatch from here on, not raw payloads. + self.__mark_route_converted(stream) + return stream + if isinstance(step, Batch): max_batch_time_ms: float | None if step.batch_timedelta is not None: @@ -513,6 +584,7 @@ def reduce( max_batch_time_ms=max_batch_time_ms, ) ) + self.__mark_route_converted(stream) return stream step = MetricsReportingReduce(step, name) @@ -522,6 +594,7 @@ def reduce( self.__consumers[stream.source].add_step( RuntimeOperator.PythonAdapter(route, ReduceDelegateFactory(step)) ) + self.__mark_route_converted(stream) return stream def broadcast( @@ -546,7 +619,7 @@ def broadcast( route, downstream_routes=[branch.root.name for branch in step.routes] ) ) - return build_branches(stream, step.routes) + return self.__propagate_raw(stream, build_branches(stream, step.routes)) def router( self, @@ -585,7 +658,7 @@ def routing_function(msg: Message[Any]) -> str: route, routing_function, cast(Sequence[str], step.routing_table.values()) ) ) - return build_branches(stream, step.routing_table.values()) + return self.__propagate_raw(stream, build_branches(stream, step.routing_table.values())) def run(self) -> None: """ diff --git a/sentry_streams/sentry_streams/pipeline/__init__.py b/sentry_streams/sentry_streams/pipeline/__init__.py index fe28502e..ee102a41 100644 --- a/sentry_streams/sentry_streams/pipeline/__init__.py +++ b/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", diff --git a/sentry_streams/sentry_streams/pipeline/pipeline.py b/sentry_streams/sentry_streams/pipeline/pipeline.py index cc303841..1f7c5f1c 100644 --- a/sentry_streams/sentry_streams/pipeline/pipeline.py +++ b/sentry_streams/sentry_streams/pipeline/pipeline.py @@ -681,6 +681,75 @@ def override_config(self, loaded_config: Mapping[str, Any]) -> None: self.batch_timedelta = timedelta(**loaded_kwargs) +@dataclass +class ArrowBatchParser( + Reduce[MeasurementUnit, InputType, Any], + Generic[MeasurementUnit, InputType], +): + """ + Batches raw Kafka payloads and decodes them into an Apache Arrow + ``RecordBatch``, entirely in Rust. + + The emitted message payload is a ``rust_streams.ArrowRecordBatch``, readable + by anything implementing the Arrow PyCapsule interface:: + + import polars as pl + + def to_frame(msg): + return pl.DataFrame(msg.payload) + + This is the fused equivalent of ``Batch`` -> ``Map(extract_bytes)`` -> + ``BatchParser``, without copying every message into Python memory on the way. + + 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 must read raw payloads**, so it has to come before any step that + converts messages into Python objects. + * 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``. + """ + + 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[InputType], Any]]: + 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]): """ diff --git a/sentry_streams/sentry_streams/rust_streams.pyi b/sentry_streams/sentry_streams/rust_streams.pyi index 606cde96..4ac8700a 100644 --- a/sentry_streams/sentry_streams/rust_streams.pyi +++ b/sentry_streams/sentry_streams/rust_streams.pyi @@ -143,6 +143,15 @@ class RuntimeOperator: max_batch_time_ms: float | None = None, ) -> Self: ... @classmethod + def ArrowBatchParser( + cls, + route: Route, + step_name: str, + schema_name: str, + max_batch_size: int | None = None, + max_batch_time_ms: float | None = None, + ) -> Self: ... + @classmethod def PythonAdapter(cls, route: Route, delegate_Factory: RustOperatorFactory) -> Self: ... class ArroyoConsumer: diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index 60a438b1..e8ea824c 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -1,7 +1,4 @@ use pyo3::prelude::*; -// Reachable once `operators::build()` gains its ArrowBatchParser arm (phase 5); -// until then the whole chain is dead in the non-test build. -#[allow(dead_code)] mod arrow_batch_parser; mod batch_step; mod broadcaster; @@ -10,9 +7,6 @@ mod commit_policy; mod committable; mod consumer; mod dev_null_sink; -// Reachable once `operators::build()` gains its ArrowBatchParser arm (phase 5); -// until then the whole chain is dead in the non-test build. -#[allow(dead_code)] mod extractors; mod filter_step; mod gcs_writer; diff --git a/sentry_streams/src/operators.rs b/sentry_streams/src/operators.rs index 523742ac..26effeb8 100644 --- a/sentry_streams/src/operators.rs +++ b/sentry_streams/src/operators.rs @@ -1,3 +1,4 @@ +use crate::arrow_batch_parser::build_arrow_batch_parser_step; use crate::batch_step::build_batch_step; use crate::broadcaster::Broadcaster; use crate::header_filter_step::build_header_int_filter; @@ -112,6 +113,23 @@ pub enum RuntimeOperator { /// Wall-clock duration in milliseconds; `None` means no time limit (size-only batch). max_batch_time_ms: Option, }, + /// Batches raw Kafka payloads and decodes them into an Apache Arrow `RecordBatch` + /// entirely in Rust, emitting one `PyAnyMessage` whose payload is an + /// `ArrowRecordBatch`. + /// + /// `schema_name` is the source topic's *logical* name, captured by the adapter + /// before any deployment topic override, and is what resolves the extractor. + /// It has to travel in the variant because [`build`] is handed no topic. + #[pyo3(name = "ArrowBatchParser")] + ArrowBatchParser { + route: Route, + step_name: String, + schema_name: String, + /// `None` means no size limit (time-only window). + max_batch_size: Option, + /// Wall-clock duration in milliseconds; `None` means no time limit. + max_batch_time_ms: Option, + }, /// Delegates messages processing to a Python operator that provides /// the same kind of interface as an Arroyo strategy. This is meant /// to simplify the porting of python strategies to Rust. @@ -238,6 +256,23 @@ pub fn build( let max_t = max_batch_time_ms.map(|ms| Duration::from_secs_f64((ms / 1000.0).max(0.0))); build_batch_step(route, *max_batch_size, max_t, step_name.clone(), next) } + RuntimeOperator::ArrowBatchParser { + route, + step_name, + schema_name, + max_batch_size, + max_batch_time_ms, + } => { + let max_t = max_batch_time_ms.map(|ms| Duration::from_secs_f64((ms / 1000.0).max(0.0))); + build_arrow_batch_parser_step( + route, + schema_name, + step_name.clone(), + *max_batch_size, + max_t, + next, + ) + } RuntimeOperator::PythonAdapter { route, delegate_factory, diff --git a/sentry_streams/tests/adapters/arroyo/test_arrow_batch_parser.py b/sentry_streams/tests/adapters/arroyo/test_arrow_batch_parser.py new file mode 100644 index 00000000..9d0c6525 --- /dev/null +++ b/sentry_streams/tests/adapters/arroyo/test_arrow_batch_parser.py @@ -0,0 +1,136 @@ +"""Adapter wiring for :class:`ArrowBatchParser`. + +The decoding itself is tested in Rust; what is worth testing here is the +placement contract, which is the part a pipeline author can get wrong. +""" + +from typing import Any, Mapping + +import pytest + +from sentry_streams.adapters.arroyo.adapter import ArroyoAdapter +from sentry_streams.adapters.arroyo.rust_arroyo import RustArroyoAdapter +from sentry_streams.adapters.stream_adapter import RuntimeTranslator +from sentry_streams.pipeline.message import Message +from sentry_streams.pipeline.pipeline import ( + ArrowBatchParser, + Map, + Pipeline, + StreamSink, + streaming_source, +) +from sentry_streams.runner import iterate_edges + +SOURCE_TOPIC = "snuba-items" + +STEPS_CONFIG: Mapping[str, Any] = { + "myinput": { + "bootstrap_servers": ["localhost:9092"], + "auto_offset_reset": "earliest", + "consumer_group": "test_group", + "override_params": {}, + }, + "kafkasink": {"bootstrap_servers": ["localhost:9092"], "override_params": {}}, +} + + +def consume_batch(msg: Message[Any]) -> Any: + """Module level, because the adapter pickle-checks transform chains.""" + return msg.payload + + +def keep_everything(msg: Message[Any]) -> bool: + return True + + +def build_adapter(steps_config: Mapping[str, Any] = STEPS_CONFIG) -> RustArroyoAdapter: + return RustArroyoAdapter.build( + {"steps_config": steps_config}, + {"type": "dummy"}, + ) + + +def parser_directly_after_source() -> Pipeline[Any]: + return ( + streaming_source(name="myinput", stream_name=SOURCE_TOPIC) + .apply(ArrowBatchParser(name="parse", batch_size=100)) + .apply(Map(name="consume", function=consume_batch)) + .sink(StreamSink(name="kafkasink", stream_name="transformed-events")) + ) + + +def parser_after_a_map() -> Pipeline[Any]: + return ( + streaming_source(name="myinput", stream_name=SOURCE_TOPIC) + .apply(Map(name="decode", function=consume_batch)) + .apply(ArrowBatchParser(name="parse", batch_size=100)) + .sink(StreamSink(name="kafkasink", stream_name="transformed-events")) + ) + + +def test_parser_directly_after_source_builds() -> None: + adapter = build_adapter() + iterate_edges(parser_directly_after_source(), RuntimeTranslator(adapter)) + assert adapter.get_consumer("myinput") is not None + + +def test_parser_after_a_python_step_is_rejected_at_build_time() -> None: + """The step reads bytes off the wire, so a preceding Map has already thrown + them away. Better a build error than a panic in production.""" + adapter = build_adapter() + with pytest.raises(ValueError) as excinfo: + iterate_edges(parser_after_a_map(), RuntimeTranslator(adapter)) + + message = str(excinfo.value) + assert "parse" in message, message + assert "already" in message and "Python objects" in message, message + + +def test_a_filter_between_source_and_parser_is_allowed() -> None: + """Filters forward messages untouched, so the raw payload survives.""" + from sentry_streams.pipeline.pipeline import PredicateFilter + + pipeline = ( + streaming_source(name="myinput", stream_name=SOURCE_TOPIC) + .apply(PredicateFilter(name="keep", function=keep_everything)) + .apply(ArrowBatchParser(name="parse", batch_size=100)) + .sink(StreamSink(name="kafkasink", stream_name="transformed-events")) + ) + adapter = build_adapter() + iterate_edges(pipeline, RuntimeTranslator(adapter)) + assert adapter.get_consumer("myinput") is not None + + +def test_schema_survives_a_deployment_topic_override() -> None: + """The extractor is resolved from the logical stream name, so overriding the + physical topic in the deployment config must not change which schema is used. + A wrong name here would be a startup panic in Rust.""" + steps_config = dict(STEPS_CONFIG) + steps_config["myinput"] = {**STEPS_CONFIG["myinput"], "topic": "snuba-items-rerouted"} + + adapter = build_adapter(steps_config) + iterate_edges(parser_directly_after_source(), RuntimeTranslator(adapter)) + assert adapter.get_consumer("myinput") is not None + + +def test_pure_python_adapter_refuses_the_step() -> None: + adapter = ArroyoAdapter.build({"steps_config": STEPS_CONFIG}) + with pytest.raises(NotImplementedError, match="rust_arroyo"): + iterate_edges(parser_directly_after_source(), RuntimeTranslator(adapter)) + + +def test_validate_requires_a_size_or_a_time_bound() -> None: + step: ArrowBatchParser[Any, Any] = ArrowBatchParser( + name="parse", batch_size=None, batch_timedelta=None + ) + with pytest.raises(ValueError, match="batch_size or batch_timedelta"): + step.validate() + + +def test_override_config_applies_deployment_settings() -> None: + step: ArrowBatchParser[Any, Any] = ArrowBatchParser(name="parse", batch_size=10) + step.override_config({"batch_size": 500, "batch_timedelta": {"seconds": 3}}) + step.validate() + assert step.batch_size == 500 + assert step.batch_timedelta is not None + assert step.batch_timedelta.total_seconds() == 3.0 From b99fcf08407f01d10b3fb6703d8cc59a4a1d7169 Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 16:27:13 -0700 Subject: [PATCH 08/13] feat(arrow): example, end-to-end test, benchmark and docs (phase 6) Adds the arrow_trace_items example with its deployment config, a test that carries real TraceItem protobufs through the step and reads the result from polars, a benchmark, and a docs page. The benchmark is an #[ignore]d test rather than a criterion suite, so it costs no dependency: cargo test --release bench_arrow_vs_pylist -- --ignored --nocapture 10k rows per window, TraceItem with one string attribute: ArrowBatchParser p50 3.5ms p99 4.1ms Batch -> BatchParser p50 4.3ms p99 6.4ms About 20% faster at the median and 35% at the tail. That is a real but modest win, well short of what "no Python round trip" suggests, and the docs say so rather than quoting the flattering comparison against Batch's flush alone -- which is 20x faster than either but does not decode anything, so it is not a complete path. Two caveats point the same way: the comparison stops at decoded values, where the Python path still has to build something columnar from those objects, and the Arrow path is still paying the Py copy and holding the GIL. The benchmark should be re-run once the source emits Rust-native messages. Decision 5 holds either way: a 1000-row window decodes in well under a millisecond, nowhere near max_poll_interval_ms, so inline decoding on the consumer thread is fine and threadpool decoding stays deferred. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- .../docs/design/arrow-batch-parser.md | 26 ++- .../docs/source/arrow_batch_parser.rst | 135 +++++++++++ sentry_streams/docs/source/index.rst | 1 + .../deployment_config/arrow_trace_items.yaml | 14 ++ .../examples/arrow_trace_items.py | 47 ++++ sentry_streams/src/arrow_batch_parser.rs | 215 ++++++++++++++++++ 6 files changed, 434 insertions(+), 4 deletions(-) create mode 100644 sentry_streams/docs/source/arrow_batch_parser.rst create mode 100644 sentry_streams/sentry_streams/deployment_config/arrow_trace_items.yaml create mode 100644 sentry_streams/sentry_streams/examples/arrow_trace_items.py diff --git a/sentry_streams/docs/design/arrow-batch-parser.md b/sentry_streams/docs/design/arrow-batch-parser.md index 1df2c8b5..0f7d181d 100644 --- a/sentry_streams/docs/design/arrow-batch-parser.md +++ b/sentry_streams/docs/design/arrow-batch-parser.md @@ -1,6 +1,6 @@ # Arrow Batch Parser — PoC Implementation Plan -**Status:** design agreed, not yet implemented +**Status:** implemented, phases 0-6 (see git history on `fpacifici/arrow_batches`) **Scope:** proof of concept — protobuf only, `sentry_protos.snuba.v1.TraceItem` only ## Goal @@ -423,9 +423,27 @@ raises `NotImplementedError`; `make typecheck` clean. 1. `sentry_streams/examples/arrow_trace_items.py` — `snuba-items` → `ArrowBatchParser` → a `Map` consuming the batch via polars. 2. End-to-end test through the full step with real `TraceItem` payloads. -3. **Benchmark against `Batch` + `BatchParser`.** This validates the premise of the - exercise and tells us whether inline decoding (decision 5) holds against - `max_poll_interval_ms=60000`. Record throughput and p99 batch decode time here. +3. **Benchmark against `Batch` + `BatchParser`.** Implemented as an `#[ignore]`d test + rather than a criterion suite, so it adds no dependency: + `cargo test --release bench_arrow_vs_pylist -- --ignored --nocapture`. + + **Results** (10 000 rows/window, `TraceItem` with one string attribute, 20 windows): + + | Path | p50 | p99 | rows/s (p50) | + |---|---|---|---| + | `ArrowBatchParser` | 3.5 ms | 4.1 ms | 2.86 M | + | `Batch` → `BatchParser` | 4.3 ms | 6.4 ms | 2.33 M | + | `Batch` flush alone (not a complete path) | 0.15 ms | 0.20 ms | 65.8 M | + + About **20% faster at p50 and 35% at p99** — real, but well short of what the "no + Python round trip" framing suggests, and worth being straight about. Two caveats + both point the same way: the comparison stops at *decoded values*, where the Python + path still has to build something columnar from those objects, and the Arrow path is + still paying the `Py` copy and holding the GIL (see *Assumed future + work*). Re-run once the source goes native. + + Decision 5 holds comfortably: a 1000-row window decodes in well under a millisecond, + nowhere near `max_poll_interval_ms`. Threadpool decoding stays deferred. 4. Docs page: the hardcoded-schema contract, Rust-adapter-only, failure behaviour. --- diff --git a/sentry_streams/docs/source/arrow_batch_parser.rst b/sentry_streams/docs/source/arrow_batch_parser.rst new file mode 100644 index 00000000..4e056d1c --- /dev/null +++ b/sentry_streams/docs/source/arrow_batch_parser.rst @@ -0,0 +1,135 @@ +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 +``ArrowRecordBatch``, readable by anything that implements the `Arrow PyCapsule +interface `_. + +.. 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.DataFrame(msg.payload) # zero-copy, via __arrow_c_stream__ + return f"{df.height} rows".encode() + + + pipeline = ( + streaming_source(name="myinput", stream_name="snuba-items") + .apply(ArrowBatchParser(name="parse_arrow", 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. + +Windowing is configured exactly like :class:`Batch`, by ``batch_size`` and/or +``batch_timedelta``, both overridable from ``steps_config``. + +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 + - ``ValueError`` at build time + * - 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`` 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 JSON-encoded into + ``attr_str``, with any nested bytes base64-encoded per proto3's canonical JSON + mapping. + +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.5 ms + - 4.1 ms + * - ``Batch`` → ``BatchParser`` + - 4.3 ms + - 6.4 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. diff --git a/sentry_streams/docs/source/index.rst b/sentry_streams/docs/source/index.rst index 08a0f19c..09386da6 100644 --- a/sentry_streams/docs/source/index.rst +++ b/sentry_streams/docs/source/index.rst @@ -11,5 +11,6 @@ build_pipeline configure_pipeline runtime/arroyo + arrow_batch_parser deployment rust diff --git a/sentry_streams/sentry_streams/deployment_config/arrow_trace_items.yaml b/sentry_streams/sentry_streams/deployment_config/arrow_trace_items.yaml new file mode 100644 index 00000000..d7f25e6f --- /dev/null +++ b/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"] diff --git a/sentry_streams/sentry_streams/examples/arrow_trace_items.py b/sentry_streams/sentry_streams/examples/arrow_trace_items.py new file mode 100644 index 00000000..63ca4bb5 --- /dev/null +++ b/sentry_streams/sentry_streams/examples/arrow_trace_items.py @@ -0,0 +1,47 @@ +"""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 copying them into Python memory, and hands the result over as an +``ArrowRecordBatch``. Any consumer implementing the Arrow PyCapsule interface can +read it -- here, 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 + + +def summarize(msg: Message[object]) -> bytes: + """Read the Arrow batch through polars and emit a one-line summary. + + ``pl.DataFrame(batch)`` goes through ``__arrow_c_stream__``; nothing is + converted row by row. + """ + df = pl.DataFrame(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="snuba-items") + # Must come before any step that turns messages into Python objects: it reads + # the raw payloads. Placing it later is a build-time error. + .apply(ArrowBatchParser(name="parse_arrow", batch_size=1000)) + .apply(Map(name="summarize", function=summarize)) + .sink(StreamSink[bytes](name="mysink", stream_name="transformed-events")) +) diff --git a/sentry_streams/src/arrow_batch_parser.rs b/sentry_streams/src/arrow_batch_parser.rs index 2128b287..e464e551 100644 --- a/sentry_streams/src/arrow_batch_parser.rs +++ b/sentry_streams/src/arrow_batch_parser.rs @@ -385,4 +385,219 @@ mod tests { assert_eq!(batch.num_rows(), 0); assert!(batch.num_columns() > 0); } + + /// The whole point of the exercise, verified end to end: raw protobuf off the + /// wire reaches polars as a typed frame without a Python object per message. + #[test] + fn a_python_consumer_reads_the_batch_with_polars() { + crate::testutils::initialize_python(); + traced_with_gil!(|py| { + let items: Vec = (1..=4) + .map(|org| TraceItem { + organization_id: org, + trace_id: format!("trace-{org}"), + item_type: 1, + attributes: std::collections::HashMap::from([( + "service".to_string(), + sentry_protos::snuba::v1::AnyValue { + value: Some(sentry_protos::snuba::v1::any_value::Value::StringValue( + format!("svc-{org}"), + )), + }, + )]), + ..Default::default() + }) + .collect(); + + let elements: Vec = items + .iter() + .map(|item| { + match build_raw_routed_value(py, item.encode_to_vec(), "s", vec!["w".into()]) + .payload + { + RoutedValuePayload::PyStreamingMessage(m) => m, + _ => unreachable!(), + } + }) + .collect(); + + let message = producer() + .produce(&route(), &elements, BTreeMap::new()) + .expect("produce"); + + let payload = message.into_payload(); + let content = match payload.payload { + RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::PyAnyMessage { + content, + }) => content, + _ => unreachable!(), + }; + let py_batch = content.bind(py).borrow().payload.clone_ref(py); + + let pl = py.import("polars").expect("polars must be importable"); + let df = pl.call_method1("DataFrame", (py_batch,)).unwrap(); + + assert_eq!( + df.call_method0("__len__") + .unwrap() + .extract::() + .unwrap(), + 4 + ); + let trace_ids: Vec = df + .get_item("trace_id") + .unwrap() + .call_method0("to_list") + .unwrap() + .extract() + .unwrap(); + assert_eq!(trace_ids, vec!["trace-1", "trace-2", "trace-3", "trace-4"]); + + let item_types: Vec = df + .get_item("item_type") + .unwrap() + .call_method0("to_list") + .unwrap() + .extract() + .unwrap(); + assert!(item_types.iter().all(|t| t == "TRACE_ITEM_TYPE_SPAN")); + + // The type-split attribute map survives the FFI boundary as a nested + // column rather than being flattened or dropped. + let columns: Vec = df.getattr("columns").unwrap().extract().unwrap(); + assert!(columns.contains(&"attr_str".to_string()), "{columns:?}"); + let dtype = df + .get_item("attr_str") + .unwrap() + .getattr("dtype") + .unwrap() + .str() + .unwrap() + .extract::() + .unwrap(); + assert!( + dtype.contains("List") || dtype.contains("Struct"), + "attr_str should arrive as a nested column, got {dtype}" + ); + }); + } + + /// Phase 6's benchmark. Not run by default -- it is a measurement, not an + /// assertion: + /// + /// cargo test --release bench_arrow_vs_pylist -- --ignored --nocapture + /// + /// It compares the Arrow producer against the Python-list producer the + /// existing `Batch` step uses, on the same window, and so answers whether + /// inline decoding on the consumer thread is affordable. + #[test] + #[ignore = "benchmark: run explicitly with --ignored --nocapture"] + fn bench_arrow_vs_pylist() { + use crate::batch_step::PyListFlushProducer; + use std::time::Instant; + + crate::testutils::initialize_python(); + const ROWS: usize = 10_000; + const REPEATS: usize = 20; + + traced_with_gil!(|py| { + let elements: Vec = (0..ROWS) + .map(|i| { + let item = TraceItem { + organization_id: i as u64, + trace_id: format!("trace-{i}"), + item_type: 1, + attributes: std::collections::HashMap::from([( + "service".to_string(), + sentry_protos::snuba::v1::AnyValue { + value: Some( + sentry_protos::snuba::v1::any_value::Value::StringValue( + "checkout".to_string(), + ), + ), + }, + )]), + ..Default::default() + }; + match build_raw_routed_value(py, item.encode_to_vec(), "s", vec!["w".into()]) + .payload + { + RoutedValuePayload::PyStreamingMessage(m) => m, + _ => unreachable!(), + } + }) + .collect(); + + let arrow = producer(); + let pylist = PyListFlushProducer; + + // The path this step replaces: Batch -> Map(extract_bytes) -> BatchParser. + // Decoding in Python is what makes it a fair comparison; the list build + // alone is not the competitor. + let decode_in_python = py + .eval( + c"lambda batch, codec: [codec.decode(p, validate=False) for p in batch]", + None, + None, + ) + .unwrap(); + let codec = py + .import("sentry_kafka_schemas") + .unwrap() + .call_method1("get_codec", ("snuba-items",)) + .unwrap(); + + let mut arrow_times = Vec::with_capacity(REPEATS); + let mut pylist_times = Vec::with_capacity(REPEATS); + let mut python_times = Vec::with_capacity(REPEATS); + for _ in 0..REPEATS { + let t = Instant::now(); + arrow + .produce(&route(), &elements, BTreeMap::new()) + .expect("arrow produce"); + arrow_times.push(t.elapsed().as_secs_f64()); + + let t = Instant::now(); + let listed = pylist + .produce(&route(), &elements, BTreeMap::new()) + .expect("pylist produce"); + pylist_times.push(t.elapsed().as_secs_f64()); + + let batch_list = match listed.into_payload().payload { + RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::PyAnyMessage { + content, + }) => content.bind(py).borrow().payload.clone_ref(py), + _ => unreachable!(), + }; + let t = Instant::now(); + decode_in_python + .call1((batch_list, &codec)) + .expect("python decode"); + python_times.push(t.elapsed().as_secs_f64() + pylist_times[pylist_times.len() - 1]); + } + + let report = |label: &str, mut times: Vec| { + times.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let p50 = times[times.len() / 2]; + let p99 = times[(times.len() as f64 * 0.99) as usize % times.len()]; + println!( + "{label:>10}: p50 {:>8.2}ms p99 {:>8.2}ms {:>10.0} rows/s", + p50 * 1000.0, + p99 * 1000.0, + ROWS as f64 / p50 + ); + }; + + println!("\n{ROWS} rows/window, {REPEATS} windows"); + report("arrow", arrow_times); + report("list only", pylist_times); + report("list+parse", python_times); + println!( + " arrow = ArrowBatchParser: decode to a RecordBatch in Rust\n \ + list only = Batch's flush alone, a Python list of `bytes` (not a \ + complete path)\n list+parse = Batch -> BatchParser, the path this \ + step replaces\n" + ); + }); + } } From f2df92be93c85b45621d8e715d8cf660dd5bb79a Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 16:52:09 -0700 Subject: [PATCH 09/13] refactor(arrow): drop the adapter's placement and schema tracking The adapter carried ~55 lines of state to pre-empt two failures the step already catches. For a PoC that is not worth it. schema_name is now declared on ArrowBatchParser rather than stashed per source and looked up at wiring time. It is usually the source's stream_name, but declaring it means the step no longer depends on where it sits, and a deployment topic override still cannot change which extractor is resolved -- now because the value never came from the topic in the first place. The placement check is gone with the route-rawness tracking that backed it (__raw_routes, __mark_route_raw/_converted, __route_is_raw, __propagate_raw). The contract moves into the type signature instead: ArrowBatchParser is a Reduce[MeasurementUnit, bytes, Any], so placing it after a step that produces anything else is a mypy error on apply(). If the types are bypassed, the with_payloads panic is still there to catch it on the first batch. Tests follow: the build-time rejection test becomes a pair that runs mypy out of process, one asserting misplacement is rejected with the specific arg-type error and one asserting correct placement is clean. Running mypy is the only way to test this, since the whole point is that it is not detectable at runtime. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- .../docs/design/arrow-batch-parser.md | 25 +++-- .../docs/source/arrow_batch_parser.rst | 14 ++- .../adapters/arroyo/rust_arroyo.py | 63 +---------- .../examples/arrow_trace_items.py | 10 +- .../sentry_streams/pipeline/pipeline.py | 17 ++- .../arroyo/test_arrow_batch_parser.py | 106 +++++++++++++----- 6 files changed, 127 insertions(+), 108 deletions(-) diff --git a/sentry_streams/docs/design/arrow-batch-parser.md b/sentry_streams/docs/design/arrow-batch-parser.md index 0f7d181d..98b4ade8 100644 --- a/sentry_streams/docs/design/arrow-batch-parser.md +++ b/sentry_streams/docs/design/arrow-batch-parser.md @@ -33,14 +33,14 @@ extraction; per-row dead-lettering; `TraceItem.outcomes`; replacing the Python | 4 | New primitive; Python `BatchParser` untouched. | | 5 | Decoding runs **inline** on the consumer thread. | | 9 | Offsets collapse to `max` per partition, as `batch_step.rs` does today. | -| 10/17 | Input contract **`RawMessage` only** — build-time check in the adapter, runtime backstop. | +| 10/17 | Input contract **`RawMessage` only** — expressed as `bytes` in the DSL signature (mypy), with the runtime panic as the backstop. | | 11 | **Generalize `BatchStep`** over a flush-producer trait rather than forking it. | | 12 | `Reduce` subclass, `StepType.REDUCE`, `isinstance` branch in `reduce()`. | | 14 | Failure is `panic!`, matching `transformer.rs:45-46`. | | 22 | `map` → **type-split maps** `attr_str/int/double/bool/bytes`. | | 23 | **Protobuf only.** | | 24 | `sentry-protos` for types and `prost` decode; `sentry-kafka-schemas` (`default-features = false`) for topic → schema. | -| 25 | Extractors indexed by the **raw resource string**; several topics sharing a schema share one extractor. | +| 25 | Extractors indexed by the **raw resource string**; several topics sharing a schema share one extractor. The topic is declared on the step as `schema_name`. | | 26 | Resolution and validation at **step construction**; failures panic at startup. | | 27 | `TraceItem` schema: all fields except `outcomes`. | @@ -399,12 +399,14 @@ mirror `Batch` (`pipeline.py:674-681`). No schema, no format, no type name. `self.__source_schemas[source_name] = schema_name`. 2. In `reduce()`, add an `isinstance(step, ArrowBatchParser)` branch before the `Batch` branch, emitting `RuntimeOperator.ArrowBatchParser(..., schema_name=self.__source_schemas[stream.source], ...)`. -3. Build-time input check. *(Implemented differently from the sketch: `reduce()` is - handed only the step and the `Route`, never the pipeline graph, so the walk is not - available.)* The adapter instead tracks which routes still carry raw payloads — the - source marks its route raw, `map`/`flat_map` and every other `reduce` clear it, and - filters, `broadcast` and `router` propagate it, since they forward messages - untouched. `ArrowBatchParser` raises if its route is not raw. +3. ~~Build-time input check.~~ **Dropped as PoC scope.** It was first implemented as + route-rawness tracking in the adapter (the sketch's edge walk is not available — + `reduce()` never sees the pipeline graph), but that is ~55 lines of adapter state to + pre-empt a failure the step already catches. The contract now lives in the DSL + signature: `ArrowBatchParser` is a `Reduce[MeasurementUnit, bytes, Any]`, so + misplacing it is a mypy error on `apply()`, and the `with_payloads` panic remains the + runtime backstop. Likewise `schema_name` is declared on the step rather than + stashed per source, which removes the other piece of adapter state. **`adapters/arroyo/adapter.py`** — `NotImplementedError` pointing at the Rust adapter, documented Rust-only in the style of `HeadersFilter`. @@ -412,9 +414,10 @@ documented Rust-only in the style of `HeadersFilter`. **Also:** export from `pipeline/__init__.py` (both the import and `__all__`); add `RuntimeOperator.ArrowBatchParser` and `ArrowRecordBatch` to `rust_streams.pyi`. -**Tests:** placing the step after a Python `Map` fails at build time naming the step; -a filter between source and parser is accepted; a deployment topic override does not -change the resolved schema; +**Tests:** placing the step after a Python `Map` is a mypy error and correct placement +is not (both run mypy out of process, since the mistake is by construction not +detectable at runtime); a filter between source and parser is accepted; a deployment +topic override does not change the resolved schema; a JSON topic panics at startup with the actual `schema_type`; the pure-Python adapter raises `NotImplementedError`; `make typecheck` clean. diff --git a/sentry_streams/docs/source/arrow_batch_parser.rst b/sentry_streams/docs/source/arrow_batch_parser.rst index 4e056d1c..b7e52b40 100644 --- a/sentry_streams/docs/source/arrow_batch_parser.rst +++ b/sentry_streams/docs/source/arrow_batch_parser.rst @@ -21,7 +21,11 @@ interface None: finalize_chain(self.__chains, stream, self.__metrics_config) ) - @staticmethod - def __route_key(stream: Route) -> tuple[str, tuple[str, ...]]: - return (stream.source, tuple(stream.waypoints)) - - def __mark_route_raw(self, stream: Route) -> None: - self.__raw_routes.add(self.__route_key(stream)) - - def __mark_route_converted(self, stream: Route) -> None: - """Record that messages on this route are now Python objects. - - Steps that only forward messages -- filters, broadcast, router -- leave - the payload alone and so do not call this. - """ - self.__raw_routes.discard(self.__route_key(stream)) - - def __route_is_raw(self, stream: Route) -> bool: - return self.__route_key(stream) in self.__raw_routes - - def __propagate_raw(self, stream: Route, branches: Mapping[str, Route]) -> Mapping[str, Route]: - """Broadcast and router forward messages untouched, so each branch keeps - whatever the incoming route had.""" - if self.__route_is_raw(stream): - for branch in branches.values(): - self.__mark_route_raw(branch) - return branches - def get_consumer(self, source: str) -> ArroyoConsumer: return self.__consumers[source] @@ -354,10 +322,7 @@ def source(self, step: Source[Any]) -> Route: dlq_config=dlq_config, sentry_dsn=self.__sentry_dsn, ) - self.__source_schemas[source_name] = schema_name - route = Route(source_name, []) - self.__mark_route_raw(route) - return route + return Route(source_name, []) def sink(self, step: Sink[Any], stream: Route) -> Route: """ @@ -470,7 +435,6 @@ def map(self, step: Map[Any, Any], stream: Route) -> Route: if self.__chains.exists(stream): self.__chains.add_map(stream, step) - self.__mark_route_converted(stream) return stream def flat_map(self, step: FlatMap[Any, Any], stream: Route) -> Route: @@ -536,27 +500,12 @@ def reduce( step.validate() if isinstance(step, ArrowBatchParser): - if not self.__route_is_raw(stream): - raise ValueError( - f"Step '{step.name}' is an ArrowBatchParser, which decodes raw Kafka " - f"payloads, but the messages reaching it on route {stream} have already " - "been converted to Python objects by an earlier step. Place it directly " - "after the source, with only filters, broadcasts or routers in between." - ) - - schema_name = self.__source_schemas.get(stream.source) - if schema_name is None: - raise ValueError( - f"Step '{step.name}': no schema recorded for source '{stream.source}'. " - "ArrowBatchParser resolves its Arrow schema from the source topic." - ) - 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=schema_name, + schema_name=step.schema_name, max_batch_size=step.batch_size, max_batch_time_ms=( step.batch_timedelta.total_seconds() * 1000.0 @@ -565,8 +514,6 @@ def reduce( ), ) ) - # The batch is an ArrowRecordBatch from here on, not raw payloads. - self.__mark_route_converted(stream) return stream if isinstance(step, Batch): @@ -584,7 +531,6 @@ def reduce( max_batch_time_ms=max_batch_time_ms, ) ) - self.__mark_route_converted(stream) return stream step = MetricsReportingReduce(step, name) @@ -594,7 +540,6 @@ def reduce( self.__consumers[stream.source].add_step( RuntimeOperator.PythonAdapter(route, ReduceDelegateFactory(step)) ) - self.__mark_route_converted(stream) return stream def broadcast( @@ -619,7 +564,7 @@ def broadcast( route, downstream_routes=[branch.root.name for branch in step.routes] ) ) - return self.__propagate_raw(stream, build_branches(stream, step.routes)) + return build_branches(stream, step.routes) def router( self, @@ -658,7 +603,7 @@ def routing_function(msg: Message[Any]) -> str: route, routing_function, cast(Sequence[str], step.routing_table.values()) ) ) - return self.__propagate_raw(stream, build_branches(stream, step.routing_table.values())) + return build_branches(stream, step.routing_table.values()) def run(self) -> None: """ diff --git a/sentry_streams/sentry_streams/examples/arrow_trace_items.py b/sentry_streams/sentry_streams/examples/arrow_trace_items.py index 63ca4bb5..04d8a9f8 100644 --- a/sentry_streams/sentry_streams/examples/arrow_trace_items.py +++ b/sentry_streams/sentry_streams/examples/arrow_trace_items.py @@ -23,6 +23,8 @@ from sentry_streams.pipeline.message import Message from sentry_streams.pipeline.pipeline import Map +TOPIC = "snuba-items" + def summarize(msg: Message[object]) -> bytes: """Read the Arrow batch through polars and emit a one-line summary. @@ -38,10 +40,10 @@ def summarize(msg: Message[object]) -> bytes: pipeline = ( - streaming_source(name="myinput", stream_name="snuba-items") - # Must come before any step that turns messages into Python objects: it reads - # the raw payloads. Placing it later is a build-time error. - .apply(ArrowBatchParser(name="parse_arrow", batch_size=1000)) + 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")) ) diff --git a/sentry_streams/sentry_streams/pipeline/pipeline.py b/sentry_streams/sentry_streams/pipeline/pipeline.py index 1f7c5f1c..bafc22ee 100644 --- a/sentry_streams/sentry_streams/pipeline/pipeline.py +++ b/sentry_streams/sentry_streams/pipeline/pipeline.py @@ -683,8 +683,8 @@ def override_config(self, loaded_config: Mapping[str, Any]) -> None: @dataclass class ArrowBatchParser( - Reduce[MeasurementUnit, InputType, Any], - Generic[MeasurementUnit, InputType], + Reduce[MeasurementUnit, bytes, Any], + Generic[MeasurementUnit], ): """ Batches raw Kafka payloads and decodes them into an Apache Arrow @@ -701,6 +701,11 @@ def to_frame(msg): This is the fused equivalent of ``Batch`` -> ``Map(extract_bytes)`` -> ``BatchParser``, without copying every message into Python memory on the way. + ``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. @@ -708,8 +713,9 @@ def to_frame(msg): nothing to configure here and adding a column needs a release. * **Rust adapter only.** The pure-Python Arroyo adapter raises ``NotImplementedError``. - * **It must read raw payloads**, so it has to come before any step that - converts messages into Python objects. + * **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. @@ -717,6 +723,7 @@ def to_frame(msg): 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 @@ -735,7 +742,7 @@ def windowing(self) -> Window[MeasurementUnit]: return TumblingWindow(self.batch_size, self.batch_timedelta) @property - def aggregate_fn(self) -> Callable[[], Accumulator[Message[InputType], Any]]: + def aggregate_fn(self) -> Callable[[], Accumulator[Message[bytes], Any]]: raise NotImplementedError( "ArrowBatchParser is implemented natively in Rust and has no Python accumulator." ) diff --git a/sentry_streams/tests/adapters/arroyo/test_arrow_batch_parser.py b/sentry_streams/tests/adapters/arroyo/test_arrow_batch_parser.py index 9d0c6525..577359d8 100644 --- a/sentry_streams/tests/adapters/arroyo/test_arrow_batch_parser.py +++ b/sentry_streams/tests/adapters/arroyo/test_arrow_batch_parser.py @@ -4,6 +4,10 @@ placement contract, which is the part a pipeline author can get wrong. """ +import subprocess +import sys +import tempfile +from pathlib import Path from typing import Any, Mapping import pytest @@ -53,37 +57,84 @@ def build_adapter(steps_config: Mapping[str, Any] = STEPS_CONFIG) -> RustArroyoA def parser_directly_after_source() -> Pipeline[Any]: return ( streaming_source(name="myinput", stream_name=SOURCE_TOPIC) - .apply(ArrowBatchParser(name="parse", batch_size=100)) + .apply(ArrowBatchParser(name="parse", schema_name=SOURCE_TOPIC, batch_size=100)) .apply(Map(name="consume", function=consume_batch)) .sink(StreamSink(name="kafkasink", stream_name="transformed-events")) ) -def parser_after_a_map() -> Pipeline[Any]: - return ( - streaming_source(name="myinput", stream_name=SOURCE_TOPIC) - .apply(Map(name="decode", function=consume_batch)) - .apply(ArrowBatchParser(name="parse", batch_size=100)) - .sink(StreamSink(name="kafkasink", stream_name="transformed-events")) - ) - - def test_parser_directly_after_source_builds() -> None: adapter = build_adapter() iterate_edges(parser_directly_after_source(), RuntimeTranslator(adapter)) assert adapter.get_consumer("myinput") is not None -def test_parser_after_a_python_step_is_rejected_at_build_time() -> None: - """The step reads bytes off the wire, so a preceding Map has already thrown - them away. Better a build error than a panic in production.""" - adapter = build_adapter() - with pytest.raises(ValueError) as excinfo: - iterate_edges(parser_after_a_map(), RuntimeTranslator(adapter)) +def test_parser_after_a_python_step_is_a_type_error() -> None: + """The step takes bytes, so placing it after a step that produces something + else is a type error. There is no build-time check: mypy catches it, and if + the types are bypassed the Rust step panics on the first batch. + + Run mypy out of process, because the mistake is by construction not + detectable at runtime.""" + code = """ +from sentry_streams.pipeline.pipeline import ( + ArrowBatchParser, + Map, + StreamSink, + streaming_source, +) +from sentry_streams.pipeline.message import Message + + +def to_text(msg: Message[bytes]) -> str: + return msg.payload.decode() + - message = str(excinfo.value) - assert "parse" in message, message - assert "already" in message and "Python objects" in message, message +pipeline = ( + streaming_source("myinput", "snuba-items") + .apply(Map("decode", function=to_text)) # bytes -> str + .apply(ArrowBatchParser("parse", schema_name="snuba-items")) # wants bytes! + .sink(StreamSink("mysink", "transformed-events")) +) +""" + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "misplaced.py" + path.write_text(code) + result = subprocess.run( + [sys.executable, "-m", "mypy", str(path), "--show-error-codes"], + capture_output=True, + text=True, + ) + + assert result.returncode > 0, result.stdout + # Specifically the placement error, not merely "mypy was unhappy". + assert ( + 'Argument 1 to "apply" of "Pipeline" has incompatible type' in result.stdout + ), result.stdout + assert "ArrowBatchParser" in result.stdout, result.stdout + + +def test_parser_directly_after_source_type_checks() -> None: + """The mirror of the test above: correct placement is accepted.""" + code = """ +from sentry_streams.pipeline.pipeline import ArrowBatchParser, StreamSink, streaming_source + +pipeline = ( + streaming_source("myinput", "snuba-items") + .apply(ArrowBatchParser("parse", schema_name="snuba-items")) + .sink(StreamSink("mysink", "transformed-events")) +) +""" + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "placed.py" + path.write_text(code) + result = subprocess.run( + [sys.executable, "-m", "mypy", str(path), "--show-error-codes"], + capture_output=True, + text=True, + ) + + assert result.returncode == 0, result.stdout def test_a_filter_between_source_and_parser_is_allowed() -> None: @@ -93,7 +144,7 @@ def test_a_filter_between_source_and_parser_is_allowed() -> None: pipeline = ( streaming_source(name="myinput", stream_name=SOURCE_TOPIC) .apply(PredicateFilter(name="keep", function=keep_everything)) - .apply(ArrowBatchParser(name="parse", batch_size=100)) + .apply(ArrowBatchParser(name="parse", schema_name=SOURCE_TOPIC, batch_size=100)) .sink(StreamSink(name="kafkasink", stream_name="transformed-events")) ) adapter = build_adapter() @@ -101,10 +152,9 @@ def test_a_filter_between_source_and_parser_is_allowed() -> None: assert adapter.get_consumer("myinput") is not None -def test_schema_survives_a_deployment_topic_override() -> None: - """The extractor is resolved from the logical stream name, so overriding the - physical topic in the deployment config must not change which schema is used. - A wrong name here would be a startup panic in Rust.""" +def test_schema_is_independent_of_the_deployment_topic() -> None: + """schema_name is declared on the step, so overriding the physical topic in + the deployment config cannot change which extractor is resolved.""" steps_config = dict(STEPS_CONFIG) steps_config["myinput"] = {**STEPS_CONFIG["myinput"], "topic": "snuba-items-rerouted"} @@ -120,15 +170,17 @@ def test_pure_python_adapter_refuses_the_step() -> None: def test_validate_requires_a_size_or_a_time_bound() -> None: - step: ArrowBatchParser[Any, Any] = ArrowBatchParser( - name="parse", batch_size=None, batch_timedelta=None + step: ArrowBatchParser[Any] = ArrowBatchParser( + name="parse", schema_name=SOURCE_TOPIC, batch_size=None, batch_timedelta=None ) with pytest.raises(ValueError, match="batch_size or batch_timedelta"): step.validate() def test_override_config_applies_deployment_settings() -> None: - step: ArrowBatchParser[Any, Any] = ArrowBatchParser(name="parse", batch_size=10) + step: ArrowBatchParser[Any] = ArrowBatchParser( + name="parse", schema_name=SOURCE_TOPIC, batch_size=10 + ) step.override_config({"batch_size": 500, "batch_timedelta": {"seconds": 3}}) step.validate() assert step.batch_size == 500 From 4a972437ee1403f8e0d75b64e8c38878d007cfd1 Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 17:06:33 -0700 Subject: [PATCH 10/13] refactor(arrow): drop recursive attribute values instead of JSON-encoding them AnyValue's array_value and kvlist_value arms were flattened to JSON and stored in attr_str. Two problems with that. It invented an ambiguity the input did not have: a JSON-encoded array in attr_str is indistinguishable from a string attribute whose value happens to look like JSON. And it was not needed. AnyValue is a port of OpenTelemetry's type, so the recursive arms exist because OTel has them, not because Sentry ingestion populates them -- the canonical snuba-items example in sentry-kafka-schemas uses only string, int, double and bool. Those attributes are now skipped, like an attribute whose oneof is unset. The rest of the row is unaffected, which the test pins along with the following row's map offsets. The accepted cost is silent loss if a producer ever does send one. Failing the batch was considered and rejected: such a message is valid protobuf we merely choose not to represent, and panicking on it would be the same self-inflicted outage as panicking on an unknown enum value. Recorded in both the design doc and the user-facing docs, with the note that if these show up in practice they should get their own attr_json column rather than going back into attr_str. Removes ~25 lines and the base64 dependency. serde_json stays; it was already a dependency of the crate for other reasons. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- sentry_streams/Cargo.lock | 1 - sentry_streams/Cargo.toml | 1 - .../docs/design/arrow-batch-parser.md | 33 +++++--- .../docs/source/arrow_batch_parser.rst | 11 ++- sentry_streams/src/extractors/trace_item.rs | 84 ++++++++----------- 5 files changed, 65 insertions(+), 65 deletions(-) diff --git a/sentry_streams/Cargo.lock b/sentry_streams/Cargo.lock index 1d387e2a..1a1669ba 100644 --- a/sentry_streams/Cargo.lock +++ b/sentry_streams/Cargo.lock @@ -2793,7 +2793,6 @@ version = "0.1.0" dependencies = [ "anyhow", "arrow", - "base64 0.22.1", "chrono", "clap", "ctrlc", diff --git a/sentry_streams/Cargo.toml b/sentry_streams/Cargo.toml index 23fa2383..51df17ff 100644 --- a/sentry_streams/Cargo.toml +++ b/sentry_streams/Cargo.toml @@ -26,7 +26,6 @@ sentry = "0.48.1" arrow = { version = "59", features = ["ffi"] } prost = "0.14" prost-types = "0.14" -base64 = "0.22" sentry_protos = "0.70" sentry-kafka-schemas = { version = "3", default-features = false } diff --git a/sentry_streams/docs/design/arrow-batch-parser.md b/sentry_streams/docs/design/arrow-batch-parser.md index 98b4ade8..17bb9a2d 100644 --- a/sentry_streams/docs/design/arrow-batch-parser.md +++ b/sentry_streams/docs/design/arrow-batch-parser.md @@ -37,7 +37,7 @@ extraction; per-row dead-lettering; `TraceItem.outcomes`; replacing the Python | 11 | **Generalize `BatchStep`** over a flush-producer trait rather than forking it. | | 12 | `Reduce` subclass, `StepType.REDUCE`, `isinstance` branch in `reduce()`. | | 14 | Failure is `panic!`, matching `transformer.rs:45-46`. | -| 22 | `map` → **type-split maps** `attr_str/int/double/bool/bytes`. | +| 22 | `map` → **type-split maps** `attr_str/int/double/bool/bytes`; the two recursive arms are dropped. | | 23 | **Protobuf only.** | | 24 | `sentry-protos` for types and `prost` decode; `sentry-kafka-schemas` (`default-features = false`) for topic → schema. | | 25 | Extractors indexed by the **raw resource string**; several topics sharing a schema share one extractor. The topic is declared on the step as `schema_name`. | @@ -84,7 +84,7 @@ Constraints discovered by reading the runtime. These drive several choices below | `retention_days` | `UInt32` | no | 100 | | `received` | `Timestamp(us, "UTC")` | **yes** | 101 | | `downsampled_retention_days` | `UInt32` | no | 102 | -| `attr_str` | `Map` | no | 7, `AnyValue` arm 1 (+ 5, 6 JSON-encoded) | +| `attr_str` | `Map` | no | 7, `AnyValue` arm 1 | | `attr_int` | `Map` | no | 7, arm 3 | | `attr_double` | `Map` | no | 7, arm 4 | | `attr_bool` | `Map` | no | 7, arm 2 | @@ -104,10 +104,22 @@ distinguish unset from zero, so they are non-nullable columns carrying the defau > compile, which is the right way to find out. `ArrayValue` (arm 5) and `KeyValueList` (arm 6) are recursive; Arrow has no recursive -types, so they are JSON-encoded into `attr_str`. Bytes *nested inside* such a value -are base64-encoded, following proto3's canonical JSON mapping — there is no way to put -raw bytes in a JSON string. This is not the case the plan rejected earlier: top-level -`bytes` attributes never pass through JSON, they keep their raw bytes in `attr_bytes`. +type, so **those attributes are dropped**. + +> **Revised after phase 3.** They were originally JSON-encoded into `attr_str`, with +> nested bytes base64-encoded per proto3 canonical JSON. Two things were wrong with +> that. It creates an ambiguity the input did not have — a JSON-encoded array in +> `attr_str` is indistinguishable from a string attribute whose value happens to look +> like JSON. And it is unnecessary: `AnyValue` is a port of OpenTelemetry's type, so the +> recursive arms exist because OTel has them, and the canonical `snuba-items` example +> uses only `string`/`int`/`double`/`bool`. Dropping them removed ~25 lines, the +> `base64` dependency, and the ambiguity. +> +> The accepted cost is **silent loss** if a producer ever does send one. Failing the +> batch was considered and rejected: such a message is valid protobuf we merely choose +> not to represent, and panicking on it is the same self-inflicted outage as panicking +> on an unknown enum value. If these ever show up in practice, give them their own +> `attr_json` column rather than folding them back into `attr_str`. ## Dependencies @@ -115,13 +127,12 @@ raw bytes in a JSON string. This is not the case the plan rejected earlier: top- arrow = { version = "59", features = ["ffi"] } prost = "0.14" prost-types = "0.14" # prost_types::Timestamp, reached through TraceItem -base64 = "0.22" # bytes nested in recursive attribute values sentry_protos = "0.70" sentry-kafka-schemas = { version = "3", default-features = false } ``` -`prost-types` and `base64` were added in phase 3; both were already in the lock file -transitively, so neither costs build time. +`prost-types` was added in phase 3, and was already in the lock file transitively, so +it costs no build time. No new Python dependencies; `pyarrow` is deliberately not added. @@ -286,7 +297,7 @@ in-process with prost, encode, extract, assert: |---|---| | all scalar fields populated | every column matches | | each `AnyValue` arm (string, bool, int, double, bytes) | lands in its own `attr_*` map | -| `ArrayValue` / `KeyValueList` | JSON-encoded into `attr_str` | +| `ArrayValue` / `KeyValueList` | dropped, leaving the row's other attributes intact | | absent `timestamp` / `received` | null, not epoch zero | | absent `conversation_id` / `session_id` | null | | unset implicit-presence scalars | zero/empty, non-null | @@ -461,7 +472,7 @@ raises `NotImplementedError`; `make typecheck` clean. the explicit PoC trade for dropping the descriptor pool. 3. **An attribute changing type between messages lands in different columns** across batches. Inherent to decision 22; Snuba EAP has the same property. -4. **Recursive attribute values become JSON strings**, not structured data. +4. **Recursive attribute values are dropped**, silently. See the schema section. 5. **Protobuf only.** A JSON or msgpack topic panics at startup. 6. **Rust adapter only**, unlike the Python `BatchParser`. 7. **Decoding holds the GIL** for the batch and stalls the consumer loop. Temporary, diff --git a/sentry_streams/docs/source/arrow_batch_parser.rst b/sentry_streams/docs/source/arrow_batch_parser.rst index b7e52b40..1d51e40d 100644 --- a/sentry_streams/docs/source/arrow_batch_parser.rst +++ b/sentry_streams/docs/source/arrow_batch_parser.rst @@ -104,9 +104,14 @@ 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 JSON-encoded into - ``attr_str``, with any nested bytes base64-encoded per proto3's canonical JSON - mapping. +* 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 ----------- diff --git a/sentry_streams/src/extractors/trace_item.rs b/sentry_streams/src/extractors/trace_item.rs index f943c8b5..d11bfa40 100644 --- a/sentry_streams/src/extractors/trace_item.rs +++ b/sentry_streams/src/extractors/trace_item.rs @@ -13,8 +13,9 @@ //! "if any" comments in the proto. //! * **`map` 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. One consequence: an -//! attribute that changes type between messages lands in different columns. +//! union type here, and Snuba EAP splits the same way. Two consequences: an +//! attribute that changes type between messages lands in different columns, and +//! `AnyValue`'s two recursive arms are dropped -- see `AttributeBuilders`. use crate::extractors::{Extractor, ExtractorError}; use arrow::array::{ @@ -23,12 +24,9 @@ use arrow::array::{ UInt64Builder, }; use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; -use base64::engine::general_purpose::STANDARD as BASE64; -use base64::Engine as _; use prost::Message; use prost_types::Timestamp; -use sentry_protos::snuba::v1::{any_value::Value, AnyValue, TraceItem, TraceItemType}; -use serde_json::{Map as JsonMap, Value as Json}; +use sentry_protos::snuba::v1::{any_value::Value, TraceItem, TraceItemType}; use std::sync::{Arc, LazyLock}; pub struct TraceItemExtractor; @@ -123,38 +121,6 @@ fn item_type_name(value: i32) -> String { } } -/// `ArrayValue` and `KeyValueList` are recursive and Arrow has no recursive type, -/// so they are flattened to JSON and stored in `attr_str`. -/// -/// Bytes nested inside such a value are base64-encoded, following the proto3 -/// canonical JSON mapping. (Top-level `bytes` attributes do *not* go through -/// here: they keep their raw bytes in `attr_bytes`.) -fn any_value_to_json(value: &AnyValue) -> Json { - match &value.value { - None => Json::Null, - Some(Value::StringValue(s)) => Json::String(s.clone()), - Some(Value::BoolValue(b)) => Json::Bool(*b), - Some(Value::IntValue(i)) => Json::Number((*i).into()), - Some(Value::DoubleValue(d)) => serde_json::Number::from_f64(*d) - .map(Json::Number) - .unwrap_or(Json::Null), - Some(Value::BytesValue(b)) => Json::String(BASE64.encode(b)), - Some(Value::ArrayValue(a)) => Json::Array(a.values.iter().map(any_value_to_json).collect()), - Some(Value::KvlistValue(kv)) => { - let mut out = JsonMap::with_capacity(kv.values.len()); - for entry in &kv.values { - let v = entry - .value - .as_ref() - .map(any_value_to_json) - .unwrap_or(Json::Null); - out.insert(entry.key.clone(), v); - } - Json::Object(out) - } - } -} - /// The five type-split attribute map builders. struct AttributeBuilders { str_: MapBuilder, @@ -227,12 +193,12 @@ impl AttributeBuilders { self.bytes.keys().append_value(key); self.bytes.values().append_value(b); } - Some(Value::ArrayValue(_)) | Some(Value::KvlistValue(_)) => { - self.str_.keys().append_value(key); - self.str_ - .values() - .append_value(any_value_to_json(value).to_string()); - } + // Recursive values have no Arrow representation, and are not + // produced in practice: `AnyValue` mirrors OpenTelemetry's type, + // so these arms exist because OTel has them. Dropped rather than + // flattened into `attr_str`, where a JSON-encoded array would be + // indistinguishable from a string attribute that looks like one. + Some(Value::ArrayValue(_)) | Some(Value::KvlistValue(_)) => {} // An attribute whose oneof is unset carries no information. None => {} } @@ -567,10 +533,14 @@ mod tests { ); } + /// Recursive attribute values are dropped, not flattened. Folding them into + /// `attr_str` as JSON would make an array attribute indistinguishable from a + /// string attribute whose value happens to look like JSON. #[test] - fn recursive_values_are_json_encoded_into_attr_str() { + fn recursive_values_are_dropped_and_do_not_disturb_their_row() { let item = TraceItem { attributes: HashMap::from([ + ("kept".into(), attr(Value::StringValue("here".into()))), ( "arr".into(), attr(Value::ArrayValue(ArrayValue { @@ -593,10 +563,26 @@ mod tests { ..Default::default() }; - let batch = extract(&[item]); - let row: HashMap = map_row(&batch, "attr_str", 0).into_iter().collect(); - assert_eq!(row["arr"], r#"[1,"two"]"#); - assert_eq!(row["kv"], r#"{"inner":false}"#); + let batch = extract(&[item, TraceItem::default()]); + + // The scalar attribute alongside them survives, and only it. + assert_eq!( + map_row(&batch, "attr_str", 0), + vec![("kept".to_string(), "here".to_string())] + ); + for column in ["attr_int", "attr_double", "attr_bool", "attr_bytes"] { + assert_eq!(map_row(&batch, column, 0), vec![], "{column}"); + } + // ... and the dropped entries did not shift the following row's offsets. + for column in [ + "attr_str", + "attr_int", + "attr_double", + "attr_bool", + "attr_bytes", + ] { + assert_eq!(map_row(&batch, column, 1), vec![], "{column}"); + } } #[test] From 66a986ff1933dbbd3c75929ab414a24587c50f25 Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 17:12:59 -0700 Subject: [PATCH 11/13] test(arrow): drop the phase 0 dependency wiring scaffolding Those two tests existed to prove the new crates were wired up before any code depended on them. Code now depends on them, so they are redundant: * get_schema, raw_schema and therefore the default-features = false choice are pinned by extractors::tests::registry_key_matches_the_schema_registry; * schema_type and the protobuf gate by ArrowFlushProducer::resolve, exercised by every parser test and asserted directly by a_json_topic_panics_at_construction; * the sentry_protos/prost version agreement by every extractor test that round-trips a TraceItem through encode_to_vec. lib.rs goes back to module declarations and the pymodule. The design doc notes where the coverage moved, so phase 0 does not read as missing its acceptance test. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- .../docs/design/arrow-batch-parser.md | 6 +++ sentry_streams/src/extractors/mod.rs | 4 ++ sentry_streams/src/lib.rs | 42 ------------------- 3 files changed, 10 insertions(+), 42 deletions(-) diff --git a/sentry_streams/docs/design/arrow-batch-parser.md b/sentry_streams/docs/design/arrow-batch-parser.md index 17bb9a2d..28b840f6 100644 --- a/sentry_streams/docs/design/arrow-batch-parser.md +++ b/sentry_streams/docs/design/arrow-batch-parser.md @@ -155,6 +155,12 @@ Phases 1 and 2 are independent and may run in parallel. 3 depends on 0; 4 on 1+2 `get_schema("snuba-items", None).unwrap().raw_schema()` equals `"sentry_protos.snuba.v1.trace_item_pb2.TraceItem"`. +> Throwaway as intended: those tests were **deleted** once later phases covered the same +> ground. `get_schema`/`raw_schema` (and so `default-features = false`) are pinned by +> `extractors::tests::registry_key_matches_the_schema_registry`, `schema_type` by +> `ArrowFlushProducer::resolve` and its tests, and the `prost` version agreement by every +> extractor test that round-trips a `TraceItem`. + ## Phase 1 — `PyRecordBatch` (Arrow → Python) **File:** `src/py_record_batch.rs` *(new)*; register in `src/lib.rs`; stubs in diff --git a/sentry_streams/src/extractors/mod.rs b/sentry_streams/src/extractors/mod.rs index 6d68b4e3..c2fd6248 100644 --- a/sentry_streams/src/extractors/mod.rs +++ b/sentry_streams/src/extractors/mod.rs @@ -1,4 +1,8 @@ //! Hand-written protobuf -> Arrow extractors, one per message type. +//! This is meant to exist for the PoC, if this works we will need a better +//! way for the rust code to generate the Arrow schema from the protobuf message +//! definitions. This may be done by exposing the protobuf definitions from +//! sentry proto. //! //! Each extractor owns a hardcoded Arrow schema and knows how to turn a batch of //! encoded payloads into a `RecordBatch`. This is the PoC trade recorded in diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index e8ea824c..c200ffe8 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -56,45 +56,3 @@ fn rust_streams(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; Ok(()) } - -/// Phase 0 of the Arrow batch parser plan: prove the new dependencies are wired -/// correctly before any code depends on them. -/// -/// In particular this pins the two assumptions the plan rests on: -/// * `sentry-kafka-schemas` with `default-features = false` still exposes the -/// topic -> schema lookup (only `validate_protobuf` sits behind -/// `type_generation`), and -/// * `sentry_protos` and `prost` agree on a `prost` version, so a generated -/// type can actually be decoded through the `prost::Message` trait we import. -/// -/// See `docs/design/arrow-batch-parser.md`. -#[cfg(test)] -mod dependency_wiring_tests { - use prost::Message; - use sentry_kafka_schemas::{get_schema, SchemaType}; - use sentry_protos::snuba::v1::TraceItem; - - #[test] - fn snuba_items_resolves_to_the_trace_item_resource() { - let schema = get_schema("snuba-items", None).expect("snuba-items must have a schema"); - - assert_eq!(schema.schema_type, SchemaType::Protobuf); - assert_eq!( - schema.raw_schema(), - "sentry_protos.snuba.v1.trace_item_pb2.TraceItem" - ); - } - - #[test] - fn sentry_protos_types_decode_through_our_prost() { - let item = TraceItem { - organization_id: 7, - ..Default::default() - }; - - let encoded = item.encode_to_vec(); - let decoded = TraceItem::decode(encoded.as_slice()).expect("round trip"); - - assert_eq!(decoded.organization_id, 7); - } -} From 59d737b52d8180b678ecb8de6387c7bd1614db50 Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 17:31:50 -0700 Subject: [PATCH 12/13] refactor(arrow): hand the batch to Python as Arrow IPC bytes Drops the PyCapsule surface entirely. ArrowFlushProducer now serializes the RecordBatch to an Arrow IPC stream and emits it as a RawMessage via into_pyraw, so Python receives ordinary bytes and reads them with polars.read_ipc_stream. This costs a serialization and a copy into Python memory. In exchange it deletes py_record_batch.rs whole -- capsule naming, release-callback ownership, the double-consume hazard, and three dunders that had to agree with each other -- along with its pyclass registration, its type stubs and its eight tests. The individual messages are still never turned into Python objects, which was the point; only the assembled batch crosses over. Re-running the benchmark says the copies are not what limits the step: including serialization it is 3.6ms p50 against 3.5ms before, about 0.1ms per 10k rows, or 3%. The per-message copy at the source is the one that matters. Restoring a zero-copy handoff later is confined to ArrowFlushProducer::produce: the extractor still produces a plain RecordBatch and knows nothing about how it is delivered. Two details worth noting: * schema is set to None on the emitted RawMessage, not to schema_name. In this runtime schema means "the codec this payload decodes with" (see msg_codecs._get_codec_from_msg), and an Arrow IPC stream is not a snuba-items message. Claiming otherwise would invite a Parser step to try. * ArrowBatchParser is now Reduce[MeasurementUnit, bytes, bytes], so the output type is honest and a downstream step that expects bytes type checks. Refs docs/design/arrow-batch-parser.md Co-Authored-By: Claude Opus 5 (1M context) --- .../docs/design/arrow-batch-parser.md | 85 ++-- .../docs/source/arrow_batch_parser.rst | 24 +- .../examples/arrow_trace_items.py | 13 +- .../sentry_streams/pipeline/pipeline.py | 17 +- .../sentry_streams/rust_streams.pyi | 16 - sentry_streams/src/arrow_batch_parser.rs | 98 +++-- sentry_streams/src/lib.rs | 2 - sentry_streams/src/py_record_batch.rs | 381 ------------------ 8 files changed, 127 insertions(+), 509 deletions(-) delete mode 100644 sentry_streams/src/py_record_batch.rs diff --git a/sentry_streams/docs/design/arrow-batch-parser.md b/sentry_streams/docs/design/arrow-batch-parser.md index 28b840f6..1b54c603 100644 --- a/sentry_streams/docs/design/arrow-batch-parser.md +++ b/sentry_streams/docs/design/arrow-batch-parser.md @@ -29,7 +29,7 @@ extraction; per-row dead-lettering; `TraceItem.outcomes`; replacing the Python | # | Decision | |---|---| | 1 | **Fused step** — batching and decoding in one primitive, not a parser after `Batch`. | -| 2 | Output is a `#[pyclass]` implementing the **Arrow PyCapsule interface**. No `pyarrow` dependency. | +| 2 | Output is the batch **serialized as an Arrow IPC stream** in a `RawMessage`. Python reads it with `polars.read_ipc_stream`. No `pyarrow` dependency. | | 4 | New primitive; Python `BatchParser` untouched. | | 5 | Decoding runs **inline** on the consumer thread. | | 9 | Offsets collapse to `max` per partition, as `batch_step.rs` does today. | @@ -161,62 +161,20 @@ Phases 1 and 2 are independent and may run in parallel. 3 depends on 0; 4 on 1+2 > `ArrowFlushProducer::resolve` and its tests, and the `prost` version agreement by every > extractor test that round-trips a `TraceItem`. -## Phase 1 — `PyRecordBatch` (Arrow → Python) +## Phase 1 — ~~`PyRecordBatch` (Arrow → Python)~~ — reverted -**File:** `src/py_record_batch.rs` *(new)*; register in `src/lib.rs`; stubs in -`sentry_streams/rust_streams.pyi`. +Originally a `#[pyclass]` implementing `__arrow_c_array__`, `__arrow_c_schema__` and +`__arrow_c_stream__`, so Python received the `RecordBatch` itself with no copy. -```rust -#[pyclass(name = "ArrowRecordBatch", module = "sentry_streams.rust_streams")] -pub struct PyRecordBatch { pub(crate) batch: RecordBatch } - -#[pymethods] -impl PyRecordBatch { - #[getter] fn num_rows(&self) -> usize; - #[getter] fn num_columns(&self) -> usize; - fn __repr__(&self) -> String; - - #[pyo3(signature = (requested_schema=None))] - fn __arrow_c_array__<'py>(&self, py: Python<'py>, requested_schema: Option>) - -> PyResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)>; - - fn __arrow_c_schema__<'py>(&self, py: Python<'py>) -> PyResult>; - - #[pyo3(signature = (requested_schema=None))] - fn __arrow_c_stream__<'py>(&self, py: Python<'py>, requested_schema: Option>) - -> PyResult>; -} -``` +**Removed as PoC scope.** The step now serializes the batch to an Arrow IPC stream and +emits it as a `RawMessage` via `into_pyraw`, so Python gets `bytes` and calls +`polars.read_ipc_stream`. That costs a serialization plus a copy into Python memory, +and in exchange deletes the entire FFI surface: capsule naming, release-callback +ownership, the double-consume hazard, and the three dunders that had to agree. -**Implement all three, not just `__arrow_c_array__`.** Table-level consumers — -`pl.DataFrame(obj)`, `pa.table(obj)` — look for `__arrow_c_stream__`; array-level -consumers use `__arrow_c_array__`. Implementing only one makes the object work in some -call sites and not others. - -Mechanics: - -- Array: `StructArray::from(batch.clone())` → `arrow::ffi::to_ffi(&struct_array.to_data())` - → two capsules. -- Stream: `FFI_ArrowArrayStream::new(Box::new(RecordBatchIterator::new(...)))`, one batch. -- **Capsule names must be exactly** `arrow_schema`, `arrow_array`, `arrow_array_stream`, - as NUL-terminated `CString`. A wrong name fails at the consumer with an opaque error. -- `PyCapsule::new` takes ownership; `FFI_ArrowSchema`/`FFI_ArrowArray`'s `Drop` invokes - the C release callback, so no manual destructor is needed. -- `requested_schema` is accepted and **ignored** — the protocol permits returning the - native schema when a cast is unsupported. Document it in the docstring. - -**Tests** - -| Test | Assertion | -|---|---| -| `polars.DataFrame(rb)` | values, column names, dtypes match | -| `pyarrow.record_batch(rb)` *(dev-dep only)* | round-trips `__arrow_c_array__` | -| `pyarrow.table(rb)` | round-trips `__arrow_c_stream__` | -| consume twice | second call still yields a valid batch (no double-release) | -| nested `Map` column | survives the FFI boundary | - -**Acceptance:** a Rust-built `RecordBatch` reaches polars with correct values and -schema. pyarrow may be a dev-only dependency; it must not enter runtime deps. +The copies are the thing to remove later, and they are removable independently — the +extractor still produces a plain `RecordBatch`, so restoring a zero-copy handoff is a +change to `ArrowFlushProducer::produce` and nothing else. ## Phase 2 — Generalize `BatchStep` (pure refactor) @@ -451,16 +409,20 @@ raises `NotImplementedError`; `make typecheck` clean. | Path | p50 | p99 | rows/s (p50) | |---|---|---|---| - | `ArrowBatchParser` | 3.5 ms | 4.1 ms | 2.86 M | - | `Batch` → `BatchParser` | 4.3 ms | 6.4 ms | 2.33 M | - | `Batch` flush alone (not a complete path) | 0.15 ms | 0.20 ms | 65.8 M | + | `ArrowBatchParser` (incl. IPC serialization + copy) | 3.6 ms | 4.8 ms | 2.77 M | + | `Batch` → `BatchParser` | 4.2 ms | 6.3 ms | 2.39 M | + | `Batch` flush alone (not a complete path) | 0.16 ms | 0.25 ms | 60.7 M | + + Re-measured after phase 1 was reverted, so these *include* serializing the batch and + copying it into Python memory — which turned out to cost about 0.1 ms per 10 000 + rows, roughly 3% of the step. The copies the PoC accepts are not what limits it. - About **20% faster at p50 and 35% at p99** — real, but well short of what the "no + About **15% faster at p50 and 25% at p99** — real, but well short of what the "no Python round trip" framing suggests, and worth being straight about. Two caveats both point the same way: the comparison stops at *decoded values*, where the Python path still has to build something columnar from those objects, and the Arrow path is - still paying the `Py` copy and holding the GIL (see *Assumed future - work*). Re-run once the source goes native. + still paying the per-message `Py` copy and holding the GIL (see *Assumed + future work*). Re-run once the source goes native. Decision 5 holds comfortably: a 1000-row window decodes in well under a millisecond, nowhere near `max_poll_interval_ms`. Threadpool decoding stays deferred. @@ -481,6 +443,7 @@ raises `NotImplementedError`; `make typecheck` clean. 4. **Recursive attribute values are dropped**, silently. See the schema section. 5. **Protobuf only.** A JSON or msgpack topic panics at startup. 6. **Rust adapter only**, unlike the Python `BatchParser`. +8. **The batch is serialized and copied into Python memory.** Deliberate, see phase 1. 7. **Decoding holds the GIL** for the batch and stalls the consumer loop. Temporary, and confined to `with_payloads`; see phase 4 and *Assumed future work*. @@ -506,6 +469,8 @@ even once the source is native. ## Deferred +- **Zero-copy handoff to Python**, via the Arrow C data interface, replacing the IPC + serialization and the copy it implies. See phase 1. - **Descriptor-driven extraction** — `prost-reflect` + `DescriptorPool` restores `get_field_by_name`, making a new column configuration rather than a release. Descriptors from vendored `.proto` compiled by `protox`, or better from an upstream PR diff --git a/sentry_streams/docs/source/arrow_batch_parser.rst b/sentry_streams/docs/source/arrow_batch_parser.rst index 1d51e40d..4fa57995 100644 --- a/sentry_streams/docs/source/arrow_batch_parser.rst +++ b/sentry_streams/docs/source/arrow_batch_parser.rst @@ -3,8 +3,8 @@ 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 -``ArrowRecordBatch``, readable by anything that implements the `Arrow PyCapsule -interface `_. +`Arrow IPC stream `_ +-- an ordinary ``bytes`` payload, which every existing step already understands. .. code-block:: python @@ -15,7 +15,7 @@ interface ``Map(extract_bytes)`` -> ``BatchParser`` and pays a Python round @@ -26,13 +25,13 @@ TOPIC = "snuba-items" -def summarize(msg: Message[object]) -> bytes: +def summarize(msg: Message[bytes]) -> bytes: """Read the Arrow batch through polars and emit a one-line summary. - ``pl.DataFrame(batch)`` goes through ``__arrow_c_stream__``; nothing is - converted row by row. + The payload is an Arrow IPC stream; polars reads all of it at once, with no + row-by-row conversion. """ - df = pl.DataFrame(msg.payload) + 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()) diff --git a/sentry_streams/sentry_streams/pipeline/pipeline.py b/sentry_streams/sentry_streams/pipeline/pipeline.py index bafc22ee..42429895 100644 --- a/sentry_streams/sentry_streams/pipeline/pipeline.py +++ b/sentry_streams/sentry_streams/pipeline/pipeline.py @@ -683,23 +683,26 @@ def override_config(self, loaded_config: Mapping[str, Any]) -> None: @dataclass class ArrowBatchParser( - Reduce[MeasurementUnit, bytes, Any], + 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 a ``rust_streams.ArrowRecordBatch``, readable - by anything implementing the Arrow PyCapsule interface:: + 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): - return pl.DataFrame(msg.payload) + 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``, without copying every message into Python memory on the way. + ``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 @@ -742,7 +745,7 @@ def windowing(self) -> Window[MeasurementUnit]: return TumblingWindow(self.batch_size, self.batch_timedelta) @property - def aggregate_fn(self) -> Callable[[], Accumulator[Message[bytes], Any]]: + def aggregate_fn(self) -> Callable[[], Accumulator[Message[bytes], bytes]]: raise NotImplementedError( "ArrowBatchParser is implemented natively in Rust and has no Python accumulator." ) diff --git a/sentry_streams/sentry_streams/rust_streams.pyi b/sentry_streams/sentry_streams/rust_streams.pyi index 4ac8700a..fc0b8441 100644 --- a/sentry_streams/sentry_streams/rust_streams.pyi +++ b/sentry_streams/sentry_streams/rust_streams.pyi @@ -220,19 +220,3 @@ class PyWatermark: def timestamp(self) -> int: ... @property def last_message_time(self) -> float | None: ... - -class ArrowRecordBatch: - """An Apache Arrow RecordBatch produced by the Rust runtime. - - Readable by any consumer implementing the Arrow PyCapsule interface, for - example ``polars.DataFrame(batch)`` or ``pyarrow.table(batch)``. There is no - Python constructor: instances come out of the Arrow batch parser step. - """ - - @property - def num_rows(self) -> int: ... - @property - def num_columns(self) -> int: ... - def __arrow_c_array__(self, requested_schema: object | None = None) -> Tuple[Any, Any]: ... - def __arrow_c_schema__(self) -> Any: ... - def __arrow_c_stream__(self, requested_schema: object | None = None) -> Any: ... diff --git a/sentry_streams/src/arrow_batch_parser.rs b/sentry_streams/src/arrow_batch_parser.rs index e464e551..91131ad8 100644 --- a/sentry_streams/src/arrow_batch_parser.rs +++ b/sentry_streams/src/arrow_batch_parser.rs @@ -1,6 +1,11 @@ //! The Arrow batch parser step: batches raw Kafka payloads and decodes them into -//! an Apache Arrow `RecordBatch` without ever materialising them as Python -//! objects. +//! an Apache Arrow `RecordBatch` without ever materialising the individual +//! messages as Python objects. +//! +//! The batch leaves as an Arrow IPC stream in a `RawMessage`, so Python receives +//! ordinary `bytes` and reads them with `polars.read_ipc_stream`. That costs a +//! serialization and a copy into Python memory; handing the `RecordBatch` over +//! directly through the Arrow C data interface is deferred. //! //! It reuses [`BatchStep`] wholesale -- windowing, watermark ordering and //! backpressure are identical to the `Batch` step -- and supplies its own @@ -8,10 +13,12 @@ use crate::batch_step::{BatchElement, BatchFlushProducer, BatchStep}; use crate::extractors::{get_extractor, registered_resources, Extractor}; -use crate::messages::{into_pyany, PyAnyMessage, PyStreamingMessage, RoutedValuePayload}; -use crate::py_record_batch::PyRecordBatch; +use crate::messages::{into_pyraw, PyStreamingMessage, RawMessage, RoutedValuePayload}; use crate::routes::{Route, RoutedValue}; use crate::utils::traced_with_gil; +use arrow::array::RecordBatch; +use arrow::error::ArrowError; +use arrow::ipc::writer::StreamWriter; use pyo3::prelude::*; use sentry_arroyo::processing::strategies::{ProcessingStrategy, StrategyError}; use sentry_arroyo::types::{Message, Partition}; @@ -62,6 +69,21 @@ fn with_payloads( }) } +/// Serialize a batch as an Arrow IPC stream. +/// +/// This costs a copy into a `Vec` and another into Python memory. That is +/// deliberate for now: it keeps the handoff an ordinary `bytes` payload, which +/// every existing step already understands, rather than an FFI object. Python +/// reads it back with `polars.read_ipc_stream`. +fn to_ipc_stream(batch: &RecordBatch) -> Result, ArrowError> { + let mut buffer = Vec::new(); + let mut writer = StreamWriter::try_new(&mut buffer, batch.schema().as_ref())?; + writer.write(batch)?; + writer.finish()?; + drop(writer); + Ok(buffer) +} + /// Decodes a flushed window into a `RecordBatch` and hands it to Python. pub(crate) struct ArrowFlushProducer { extractor: &'static dyn Extractor, @@ -138,15 +160,27 @@ impl BatchFlushProducer for ArrowFlushProducer { ) }); - let content = traced_with_gil!(|py| -> PyResult> { - let py_batch = Py::new(py, PyRecordBatch::new(batch))?; - into_pyany( + let payload = to_ipc_stream(&batch).unwrap_or_else(|e| { + panic!( + "step '{}': could not serialize a {}-row Arrow batch from topic '{}': {e}", + self.step_name, + batch.num_rows(), + self.schema_name, + ) + }); + + let content = traced_with_gil!(|py| { + into_pyraw( py, - PyAnyMessage { - payload: py_batch.into_any(), + RawMessage { + payload, headers: vec![], timestamp: ts, - schema: Some(self.schema_name.clone()), + // Deliberately not `schema_name`. In this runtime `schema` means + // "the schema this payload can be decoded with" (see + // `msg_codecs._get_codec_from_msg`), and these bytes are an Arrow + // IPC stream, not a message of the source's schema. + schema: None, }, ) }) @@ -155,7 +189,7 @@ impl BatchFlushProducer for ArrowFlushProducer { Ok(Message::new_any_message( RoutedValue { route: route.clone(), - payload: RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::PyAnyMessage { + payload: RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::RawMessage { content, }), }, @@ -188,7 +222,8 @@ mod tests { use super::*; use crate::fake_strategy::FakeStrategy; use crate::testutils::{build_raw_routed_value, build_routed_value}; - use arrow::array::{RecordBatch, StringArray, UInt64Array}; + use arrow::array::{StringArray, UInt64Array}; + use arrow::ipc::reader::StreamReader; use prost::Message as _; use pyo3::types::PyAnyMethods; use pyo3::IntoPyObject; @@ -215,21 +250,26 @@ mod tests { ArrowFlushProducer::resolve("test_arrow".to_string(), SNUBA_ITEMS) } - /// Pull the `RecordBatch` back out of the emitted message, the way a Python - /// consumer would see it. + fn from_ipc_stream(bytes: &[u8]) -> RecordBatch { + let mut reader = StreamReader::try_new(bytes, None).expect("valid Arrow IPC stream"); + let batch = reader.next().expect("one batch").expect("readable batch"); + assert!(reader.next().is_none(), "stream carries exactly one batch"); + batch + } + + /// Pull the emitted payload apart the way a Python consumer would: raw bytes + /// out of a `RawMessage`, decoded as an Arrow IPC stream. fn batch_of(message: Message) -> (RecordBatch, Option) { let payload = message.into_payload(); let content = match payload.payload { - RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::PyAnyMessage { - content, - }) => content, - _ => panic!("expected a PyAnyMessage carrying the record batch"), + RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::RawMessage { content }) => { + content + } + _ => panic!("expected a RawMessage carrying the serialized batch"), }; traced_with_gil!(|py| { let borrowed = content.bind(py).borrow(); - let schema = borrowed.schema.clone(); - let rb: PyRef = borrowed.payload.bind(py).extract().unwrap(); - (rb.batch.clone(), schema) + (from_ipc_stream(&borrowed.payload), borrowed.schema.clone()) }) } @@ -277,8 +317,9 @@ mod tests { .unwrap(); assert_eq!(traces.value(0), "trace-1"); - // Downstream Python needs to know which stream the batch came from. - assert_eq!(schema.as_deref(), Some(SNUBA_ITEMS)); + // `schema` means "decodable with this codec" in this runtime, and an Arrow + // IPC stream is not a snuba-items message, so it is deliberately unset. + assert_eq!(schema, None); } /// Decision 10: this step only accepts RawMessage. A PyAnyMessage means some @@ -368,8 +409,8 @@ mod tests { let out = sub.lock().unwrap(); assert_eq!(out.len(), 1, "one batch downstream, not two rows"); traced_with_gil!(|py| { - let rb: PyRef = out[0].bind(py).extract().unwrap(); - assert_eq!(rb.batch.num_rows(), 2); + let bytes: Vec = out[0].bind(py).extract().unwrap(); + assert_eq!(from_ipc_stream(&bytes).num_rows(), 2); }); } @@ -427,15 +468,16 @@ mod tests { let payload = message.into_payload(); let content = match payload.payload { - RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::PyAnyMessage { + RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::RawMessage { content, }) => content, _ => unreachable!(), }; - let py_batch = content.bind(py).borrow().payload.clone_ref(py); + let py_bytes = content.bind(py).getattr("payload").unwrap(); + // Exactly what a downstream Map would do with the payload. let pl = py.import("polars").expect("polars must be importable"); - let df = pl.call_method1("DataFrame", (py_batch,)).unwrap(); + let df = pl.call_method1("read_ipc_stream", (py_bytes,)).unwrap(); assert_eq!( df.call_method0("__len__") diff --git a/sentry_streams/src/lib.rs b/sentry_streams/src/lib.rs index c200ffe8..e15421b6 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -18,7 +18,6 @@ mod metrics_config; mod mocks; mod operators; mod pipeline_stats; -mod py_record_batch; mod python_operator; mod routers; mod routes; @@ -53,6 +52,5 @@ fn rust_streams(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; - m.add_class::()?; Ok(()) } diff --git a/sentry_streams/src/py_record_batch.rs b/sentry_streams/src/py_record_batch.rs deleted file mode 100644 index 36af8f22..00000000 --- a/sentry_streams/src/py_record_batch.rs +++ /dev/null @@ -1,381 +0,0 @@ -//! Hands an Arrow `RecordBatch` to Python over the Arrow PyCapsule interface. -//! -//! Deliberately no `pyarrow` dependency: any consumer that speaks the PyCapsule -//! protocol (polars, pyarrow, duckdb, ...) can read the batch without a copy. -//! -//! See `docs/design/arrow-batch-parser.md`, phase 1. - -use arrow::array::{Array, RecordBatch, RecordBatchIterator, StructArray}; -use arrow::error::ArrowError; -use arrow::ffi::{to_ffi, FFI_ArrowSchema}; -use arrow::ffi_stream::FFI_ArrowArrayStream; -use pyo3::exceptions::PyRuntimeError; -use pyo3::prelude::*; -use pyo3::types::PyCapsule; -use std::ffi::CStr; - -/// Capsule names mandated by the Arrow PyCapsule interface. Getting one wrong is -/// not a soft failure: consumers reject the capsule with an opaque error, so they -/// are named constants and asserted in the tests. -const SCHEMA_CAPSULE_NAME: &CStr = c"arrow_schema"; -const ARRAY_CAPSULE_NAME: &CStr = c"arrow_array"; -const STREAM_CAPSULE_NAME: &CStr = c"arrow_array_stream"; - -/// An Arrow `RecordBatch` produced by the Rust runtime, readable from Python by -/// anything that speaks the Arrow PyCapsule interface: -/// -/// ```python -/// import polars as pl -/// df = pl.DataFrame(batch) -/// ``` -#[pyclass( - name = "ArrowRecordBatch", - module = "sentry_streams.rust_streams", - frozen -)] -pub struct PyRecordBatch { - pub(crate) batch: RecordBatch, -} - -impl PyRecordBatch { - pub(crate) fn new(batch: RecordBatch) -> Self { - Self { batch } - } -} - -fn arrow_err(e: ArrowError) -> PyErr { - PyRuntimeError::new_err(format!("Arrow C data interface export failed: {e}")) -} - -#[pymethods] -impl PyRecordBatch { - #[getter] - fn num_rows(&self) -> usize { - self.batch.num_rows() - } - - #[getter] - fn num_columns(&self) -> usize { - self.batch.num_columns() - } - - fn __repr__(&self) -> String { - format!( - "ArrowRecordBatch(num_rows={}, num_columns={})", - self.batch.num_rows(), - self.batch.num_columns() - ) - } - - /// Export as a single Arrow array (a struct array, one field per column). - /// - /// `requested_schema` is accepted and ignored: the PyCapsule interface allows - /// a producer to return its native schema when it cannot perform the - /// requested cast, and we never cast. - #[pyo3(signature = (requested_schema=None))] - fn __arrow_c_array__<'py>( - &self, - py: Python<'py>, - requested_schema: Option>, - ) -> PyResult<(Bound<'py, PyCapsule>, Bound<'py, PyCapsule>)> { - let _ = requested_schema; - - let struct_array = StructArray::from(self.batch.clone()); - let (ffi_array, ffi_schema) = to_ffi(&struct_array.to_data()).map_err(arrow_err)?; - - // The capsule takes ownership; FFI_ArrowSchema/FFI_ArrowArray's Drop - // invokes the C release callback, so no manual destructor is needed. - let schema_capsule = PyCapsule::new_with_value(py, ffi_schema, SCHEMA_CAPSULE_NAME)?; - let array_capsule = PyCapsule::new_with_value(py, ffi_array, ARRAY_CAPSULE_NAME)?; - Ok((schema_capsule, array_capsule)) - } - - /// Export just the schema, for consumers that inspect before reading. - fn __arrow_c_schema__<'py>(&self, py: Python<'py>) -> PyResult> { - let ffi_schema = - FFI_ArrowSchema::try_from(self.batch.schema().as_ref()).map_err(arrow_err)?; - PyCapsule::new_with_value(py, ffi_schema, SCHEMA_CAPSULE_NAME) - } - - /// Export as a stream of exactly one batch. - /// - /// Table-level consumers (`pl.DataFrame(obj)`, `pa.table(obj)`) look for this - /// rather than `__arrow_c_array__`, which is why both exist. - #[pyo3(signature = (requested_schema=None))] - fn __arrow_c_stream__<'py>( - &self, - py: Python<'py>, - requested_schema: Option>, - ) -> PyResult> { - let _ = requested_schema; - - let batch = self.batch.clone(); - let schema = batch.schema(); - let reader = RecordBatchIterator::new(vec![Ok(batch)], schema); - let stream = FFI_ArrowArrayStream::new(Box::new(reader)); - - PyCapsule::new_with_value(py, stream, STREAM_CAPSULE_NAME) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{ - ArrayRef, Int64Array, MapBuilder, MapFieldNames, StringArray, StringBuilder, - }; - use arrow::datatypes::Schema; - use arrow::ffi::from_ffi; - use arrow::ffi_stream::ArrowArrayStreamReader; - use arrow::record_batch::RecordBatchReader; - use pyo3::types::PyCapsuleMethods; - use std::sync::Arc; - - fn capsule_name(capsule: &Bound<'_, PyCapsule>) -> &'static CStr { - // SAFETY: the name is a `&'static CStr` we set ourselves and no Python - // code has had the chance to rename the capsule. - unsafe { capsule.name().unwrap().unwrap().as_cstr() } - } - - /// Consume the capsules the way a real consumer does: move the array out and - /// leave the exported struct released, so double-release bugs would show up. - fn import_array( - schema_capsule: &Bound<'_, PyCapsule>, - array_capsule: &Bound<'_, PyCapsule>, - ) -> RecordBatch { - let data = unsafe { - let schema_ptr = schema_capsule - .pointer_checked(Some(SCHEMA_CAPSULE_NAME)) - .unwrap() - .as_ptr() as *const FFI_ArrowSchema; - let array_ptr = array_capsule - .pointer_checked(Some(ARRAY_CAPSULE_NAME)) - .unwrap() - .as_ptr() as *mut arrow::ffi::FFI_ArrowArray; - let array = std::ptr::replace(array_ptr, arrow::ffi::FFI_ArrowArray::empty()); - from_ffi(array, &*schema_ptr).unwrap() - }; - RecordBatch::from(StructArray::from(data)) - } - - fn import_schema(capsule: &Bound<'_, PyCapsule>) -> Schema { - unsafe { - let ptr = capsule - .pointer_checked(Some(SCHEMA_CAPSULE_NAME)) - .unwrap() - .as_ptr() as *const FFI_ArrowSchema; - Schema::try_from(&*ptr).unwrap() - } - } - - /// Two scalar columns plus a `Map`, so the nested case is - /// exercised by every test rather than only by a dedicated one. - fn sample_batch() -> RecordBatch { - let ids: ArrayRef = Arc::new(Int64Array::from(vec![1_i64, 2, 3])); - let names: ArrayRef = Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])); - - let mut attrs = MapBuilder::new( - Some(MapFieldNames { - entry: "entries".to_string(), - key: "key".to_string(), - value: "value".to_string(), - }), - StringBuilder::new(), - StringBuilder::new(), - ); - // row 0: two entries, row 1: none, row 2: one entry - attrs.keys().append_value("k1"); - attrs.values().append_value("v1"); - attrs.keys().append_value("k2"); - attrs.values().append_value("v2"); - attrs.append(true).unwrap(); - attrs.append(true).unwrap(); - attrs.keys().append_value("k3"); - attrs.values().append_value("v3"); - attrs.append(true).unwrap(); - let attrs: ArrayRef = Arc::new(attrs.finish()); - - RecordBatch::try_from_iter(vec![("id", ids), ("name", names), ("attrs", attrs)]).unwrap() - } - - fn py_batch(py: Python<'_>) -> Py { - Py::new(py, PyRecordBatch::new(sample_batch())).unwrap() - } - - #[test] - fn exposes_shape_to_python() { - crate::testutils::initialize_python(); - Python::attach(|py| { - let rb = py_batch(py); - let b = rb.bind(py); - assert_eq!( - b.getattr("num_rows").unwrap().extract::().unwrap(), - 3 - ); - assert_eq!( - b.getattr("num_columns") - .unwrap() - .extract::() - .unwrap(), - 3 - ); - let repr = b.repr().unwrap().extract::().unwrap(); - assert!(repr.contains("ArrowRecordBatch"), "got {repr}"); - assert!(repr.contains('3'), "repr should mention the shape: {repr}"); - }); - } - - #[test] - fn arrow_c_array_round_trips_through_ffi() { - crate::testutils::initialize_python(); - Python::attach(|py| { - let rb = py_batch(py); - let (schema_capsule, array_capsule) = rb.get().__arrow_c_array__(py, None).unwrap(); - - assert_eq!(capsule_name(&schema_capsule), SCHEMA_CAPSULE_NAME); - assert_eq!(capsule_name(&array_capsule), ARRAY_CAPSULE_NAME); - - // Import the capsules back and compare against the original batch. - let round_tripped = import_array(&schema_capsule, &array_capsule); - assert_eq!(round_tripped, sample_batch()); - }); - } - - #[test] - fn arrow_c_schema_describes_the_batch() { - crate::testutils::initialize_python(); - Python::attach(|py| { - let rb = py_batch(py); - let capsule = rb.get().__arrow_c_schema__(py).unwrap(); - assert_eq!(capsule_name(&capsule), SCHEMA_CAPSULE_NAME); - - assert_eq!(&import_schema(&capsule), sample_batch().schema().as_ref()); - }); - } - - #[test] - fn arrow_c_stream_yields_the_batch() { - crate::testutils::initialize_python(); - Python::attach(|py| { - let rb = py_batch(py); - let capsule = rb.get().__arrow_c_stream__(py, None).unwrap(); - assert_eq!(capsule_name(&capsule), STREAM_CAPSULE_NAME); - - let ptr = capsule - .pointer_checked(Some(STREAM_CAPSULE_NAME)) - .unwrap() - .as_ptr() as *mut FFI_ArrowArrayStream; - let mut reader = unsafe { ArrowArrayStreamReader::from_raw(ptr) }.unwrap(); - assert_eq!(reader.schema().as_ref(), sample_batch().schema().as_ref()); - let batch = reader.next().unwrap().unwrap(); - assert_eq!(batch, sample_batch()); - assert!( - reader.next().is_none(), - "stream must hold exactly one batch" - ); - }); - } - - /// Exporting must not consume the batch: the object stays usable, which is - /// what a Python caller passing it to two consumers would expect. - #[test] - fn can_be_exported_twice() { - crate::testutils::initialize_python(); - Python::attach(|py| { - let rb = py_batch(py); - let _first = rb.get().__arrow_c_array__(py, None).unwrap(); - let (schema_capsule, array_capsule) = rb.get().__arrow_c_array__(py, None).unwrap(); - - assert_eq!( - import_array(&schema_capsule, &array_capsule), - sample_batch() - ); - }); - } - - /// `requested_schema` is accepted and ignored; the protocol allows returning - /// the native schema when the requested cast is unsupported. - #[test] - fn requested_schema_is_ignored_not_rejected() { - crate::testutils::initialize_python(); - Python::attach(|py| { - let rb = py_batch(py); - let requested = rb.get().__arrow_c_schema__(py).unwrap().into_any(); - assert!(rb - .get() - .__arrow_c_array__(py, Some(requested.clone())) - .is_ok()); - assert!(rb.get().__arrow_c_stream__(py, Some(requested)).is_ok()); - }); - } - - /// The real acceptance criterion: an independent Arrow implementation reads - /// our batch with the right values, names and dtypes. - #[test] - fn polars_reads_the_batch() { - crate::testutils::initialize_python(); - Python::attach(|py| { - // polars is a declared runtime dependency of this package, so a - // missing import is a broken environment, not a reason to skip. - let pl = py.import("polars").expect("polars must be importable"); - let rb = py_batch(py); - let df = pl.call_method1("DataFrame", (rb,)).unwrap(); - - let columns: Vec = df.getattr("columns").unwrap().extract().unwrap(); - assert_eq!(columns, vec!["id", "name", "attrs"]); - assert_eq!( - df.call_method0("__len__") - .unwrap() - .extract::() - .unwrap(), - 3 - ); - - let ids: Vec = df - .get_item("id") - .unwrap() - .call_method0("to_list") - .unwrap() - .extract() - .unwrap(); - assert_eq!(ids, vec![1, 2, 3]); - - let names: Vec> = df - .get_item("name") - .unwrap() - .call_method0("to_list") - .unwrap() - .extract() - .unwrap(); - assert_eq!( - names, - vec![Some("a".to_string()), None, Some("c".to_string())] - ); - - let dtype = df - .get_item("id") - .unwrap() - .getattr("dtype") - .unwrap() - .str() - .unwrap() - .extract::() - .unwrap(); - assert_eq!(dtype, "Int64"); - }); - } - - #[test] - fn empty_batch_keeps_its_schema() { - crate::testutils::initialize_python(); - Python::attach(|py| { - let schema = sample_batch().schema(); - let empty = RecordBatch::new_empty(schema.clone()); - let rb = Py::new(py, PyRecordBatch::new(empty)).unwrap(); - assert_eq!(rb.get().num_rows(), 0); - - let capsule = rb.get().__arrow_c_schema__(py).unwrap(); - assert_eq!(&import_schema(&capsule), schema.as_ref()); - }); - } -} From e08e7468d25498344836c318801b855b97cfbe26 Mon Sep 17 00:00:00 2001 From: Filippo Pacifici Date: Sun, 13 Sep 2026 17:53:32 -0700 Subject: [PATCH 13/13] Fix commit --- sentry_streams/src/arrow_batch_parser.rs | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/sentry_streams/src/arrow_batch_parser.rs b/sentry_streams/src/arrow_batch_parser.rs index 91131ad8..8fa19301 100644 --- a/sentry_streams/src/arrow_batch_parser.rs +++ b/sentry_streams/src/arrow_batch_parser.rs @@ -36,13 +36,12 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; /// messages this function loses its GIL block and its `PyRef` guards, and /// nothing else in the step changes. /// -/// It is a scope rather than a plain accessor because the `PyRef` guards must -/// outlive the slices handed to `f`. +/// This is needed because the consumer immediately moves the raw message into +/// Python memory. This behavior is removed in: +/// https://github.com/getsentry/streams/pull/376 /// -/// Do **not** copy the payloads out to release the GIL sooner. It would work -/// today and would become permanent dead weight the moment the source goes -/// native -- a per-message copy in the one step whose whole purpose is to remove -/// per-message copies. +/// When The PR above will be merged, we will simplify this and avoid taking the +/// GIL at ewvery message. fn with_payloads( step_name: &str, elements: &[BatchElement],