Parse messages into Arrow Record Batches - #375
Draft
fpacifici wants to merge 12 commits into
Draft
Conversation
fpacifici
commented
Sep 13, 2026
fpacifici
commented
Collaborator
- docs(arrow): design plan for the Rust Arrow batch parser
- feat(arrow): add Arrow batch parser dependencies (phase 0)
- feat(arrow): expose RecordBatch to Python via PyCapsule (phase 1)
- refactor(arrow): generalize BatchStep over a flush producer (phase 2)
- feat(arrow): protobuf to Arrow extractors, TraceItem (phase 3)
- feat(arrow): the Arrow batch parser step (phase 4)
- feat(arrow): DSL and adapter wiring for ArrowBatchParser (phase 5)
- feat(arrow): example, end-to-end test, benchmark and docs (phase 6)
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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<Utf8, Utf8> 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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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_<n> 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<String>. 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) <noreply@anthropic.com>
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<RawMessage>; 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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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<RawMessage> 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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
…ding 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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
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) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.