diff --git a/sentry_streams/Cargo.lock b/sentry_streams/Cargo.lock index e21a054a..1a1669ba 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,16 @@ dependencies = [ "metrics", "metrics-exporter-dogstatsd", "parking_lot", + "prost", + "prost-types", "pyo3", "rand 0.9.4", "rdkafka", "reqwest 0.12.15", "sentry", + "sentry-kafka-schemas", "sentry_arroyo", + "sentry_protos", "serde", "serde_json", "tokio", @@ -2388,6 +3049,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 +3126,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 +3192,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 +3229,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 +3446,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 +3514,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 +3555,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 +3603,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 +3745,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 +3763,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 +3780,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 +4256,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..51df17ff 100644 --- a/sentry_streams/Cargo.toml +++ b/sentry_streams/Cargo.toml @@ -23,6 +23,11 @@ metrics = "0.24.0" metrics-exporter-dogstatsd = "0.9.0" rand = "0.9" sentry = "0.48.1" +arrow = { version = "59", features = ["ffi"] } +prost = "0.14" +prost-types = "0.14" +sentry_protos = "0.70" +sentry-kafka-schemas = { version = "3", default-features = false } [lib] name = "rust_streams" 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..1b54c603 --- /dev/null +++ b/sentry_streams/docs/design/arrow-batch-parser.md @@ -0,0 +1,484 @@ +# Arrow Batch Parser — PoC Implementation Plan + +**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 + +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 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. | +| 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`; 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`. | +| 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` | no | 10 | +| `session_id` | `Utf8` | no | 11 | +| `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 | +| `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. + +> **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 +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 + +```toml +arrow = { version = "59", features = ["ffi"] } +prost = "0.14" +prost-types = "0.14" # prost_types::Timestamp, reached through TraceItem +sentry_protos = "0.70" +sentry-kafka-schemas = { version = "3", default-features = false } +``` + +`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. + +--- + +# 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"`. + +> 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)~~ — reverted + +Originally a `#[pyclass]` implementing `__arrow_c_array__`, `__arrow_c_schema__` and +`__arrow_c_stream__`, so Python received the `RecordBatch` itself with no copy. + +**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. + +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) + +**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 `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 }` → + `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` | 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 | +| 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.~~ **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`. + +**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` 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. + +## 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`.** 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` (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 **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 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. +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 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*. + +## 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 + +- **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 + 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. 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..4fa57995 --- /dev/null +++ b/sentry_streams/docs/source/arrow_batch_parser.rst @@ -0,0 +1,158 @@ +Arrow Batch Parser +================== + +``ArrowBatchParser`` batches raw Kafka payloads and decodes them into an Apache +Arrow ``RecordBatch`` entirely in Rust. The batch is handed to Python as an +`Arrow IPC stream `_ +-- an ordinary ``bytes`` payload, which every existing step already understands. + +.. code-block:: python + + import polars as pl + + from sentry_streams.pipeline import ArrowBatchParser, StreamSink, streaming_source + from sentry_streams.pipeline.pipeline import Map + + + def summarize(msg): + df = pl.read_ipc_stream(msg.payload) + return f"{df.height} rows".encode() + + + pipeline = ( + streaming_source(name="myinput", stream_name="snuba-items") + .apply( + ArrowBatchParser( + name="parse_arrow", schema_name="snuba-items", batch_size=1000 + ) + ) + .apply(Map(name="summarize", function=summarize)) + .sink(StreamSink[bytes](name="mysink", stream_name="transformed-events")) + ) + +It replaces ``Batch`` → ``Map(extract_bytes)`` → ``BatchParser``, in which every +message is copied into Python memory as ``bytes`` and decoded under the GIL. +Here the individual messages are never turned into Python objects; only the +assembled batch crosses over. + +.. note:: + + Serializing the batch and copying it into Python memory is a deliberate + simplification. Handing the ``RecordBatch`` across directly, through the Arrow + C data interface, removes both copies and is deferred rather than ruled out. + +Windowing is configured exactly like :class:`Batch`, by ``batch_size`` and/or +``batch_timedelta``, both overridable from ``steps_config``. + +``schema_name`` is the logical stream name whose ``sentry-kafka-schemas`` entry +names the message type, and so selects the extractor -- normally the source's +``stream_name``. It is declared on the step rather than inferred from the source, +so the step does not depend on where it sits and a deployment topic override +cannot change which extractor is used. + +The schema is hardcoded +----------------------- + +There is nothing to configure. Each supported message type has a hand-written +extractor in Rust that owns its Arrow schema, resolved from the source topic's +entry in ``sentry-kafka-schemas``. Several topics sharing a schema share one +extractor. + +**Adding a column, or a new message type, is a Rust change and a release.** This +is the deliberate trade of the current implementation: Rust has no equivalent of +Python's reflective ``ProtobufCodec``, and carrying a protobuf descriptor pool to +get one was judged not worth it yet. See +``docs/design/arrow-batch-parser.md``. + +Currently implemented: ``sentry_protos.snuba.v1.TraceItem`` (topic +``snuba-items``). + +How it can fail +--------------- + +Every failure is loud, and most happen at startup rather than in production: + +.. list-table:: + :header-rows: 1 + :widths: 30 70 + + * - Situation + - Result + * - Topic has no schema, or a JSON/msgpack schema + - Panic at startup, naming the topic and the schema type found + * - Message type has no extractor + - Panic at startup, listing the supported message types + * - Step placed after a ``Map`` or other converting step + - Type error under mypy; **panic** on the first batch if types were bypassed + * - Used with the pure-Python Arroyo adapter + - ``NotImplementedError`` + * - A payload fails to decode + - **Panic, failing the process** + +That last one deserves emphasis. Batching collapses offsets to the maximum per +partition, so there is no ``(partition, offset)`` for Arroyo's DLQ to reject a +single row with: one bad payload fails the whole batch, and the consumer will +re-read the same offsets on restart. This matches what the runtime already does +for an error on an ``AnyMessage``, but it means a malformed message is an outage, +not a dead letter. Per-row offset tracking is deferred work. + +Placement +--------- + +The step reads bytes off the wire, so it must come before any step that turns +messages into Python objects. Filters, broadcasts and routers are fine — they +forward messages untouched. A ``Map`` in between is a build-time error. + +Attribute columns +----------------- + +``TraceItem``'s ``map`` is split by value type into +``attr_str``, ``attr_int``, ``attr_double``, ``attr_bool`` and ``attr_bytes``. +Arrow has no usable union type here and Snuba EAP splits the same way. Two +consequences: + +* an attribute that changes type between messages lands in different columns; +* recursive values (``ArrayValue``, ``KeyValueList``) are **dropped**. Arrow has + no recursive type. ``AnyValue`` is a port of OpenTelemetry's type, so these + arms exist because OTel has them rather than because Sentry ingestion uses + them, and they do not appear in the canonical ``snuba-items`` example. If one + does arrive, that attribute is skipped -- the rest of its row is unaffected -- + and nothing is logged. If they turn out to occur in practice they should get + their own column rather than being flattened into ``attr_str``, where an + encoded array would be indistinguishable from a string that looks like one. + +Performance +----------- + +Measured with ``cargo test --release bench_arrow_vs_pylist -- --ignored +--nocapture``, 10 000 rows per window, ``TraceItem`` with one string attribute: + +.. list-table:: + :header-rows: 1 + + * - Path + - p50 + - p99 + * - ``ArrowBatchParser`` + - 3.6 ms + - 4.8 ms + * - ``Batch`` → ``BatchParser`` + - 4.2 ms + - 6.3 ms + +About 20% faster at the median and 35% at the tail — a real but modest win on +this workload, not the order of magnitude the "no Python round trip" framing +might suggest. Two things to keep in mind when reading it: + +* The comparison stops at decoded values. The Python path then has to *build* + something columnar (polars, parquet) from those objects, which this step has + already done. +* The source still hands Rust its payloads inside Python-owned ``RawMessage`` + objects, so decoding holds the GIL and pays a copy *per message*. Removing that + is work owned elsewhere; the numbers here understate what the step can do once + it lands, and the benchmark should be re-run then. + +At these figures a 1 000-row window decodes in well under a millisecond, so +running inline on the consumer thread is comfortably within +``max_poll_interval_ms``. Moving decoding to a threadpool is deferred until a +benchmark justifies it. 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/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..b21d9908 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, @@ -498,6 +499,23 @@ def reduce( step.override_config(loaded_config) step.validate() + if isinstance(step, ArrowBatchParser): + logger.info(f"Adding Arrow batch parser (native): {step.name} to pipeline") + self.__consumers[stream.source].add_step( + RuntimeOperator.ArrowBatchParser( + route=route, + step_name=step.name, + schema_name=step.schema_name, + max_batch_size=step.batch_size, + max_batch_time_ms=( + step.batch_timedelta.total_seconds() * 1000.0 + if step.batch_timedelta is not None + else None + ), + ) + ) + return stream + if isinstance(step, Batch): max_batch_time_ms: float | None if step.batch_timedelta is not None: 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..bc45a041 --- /dev/null +++ b/sentry_streams/sentry_streams/examples/arrow_trace_items.py @@ -0,0 +1,48 @@ +"""Decode a batch of TraceItem protobufs into an Apache Arrow RecordBatch in Rust. + +``ArrowBatchParser`` fuses batching and decoding: it reads the raw Kafka payloads +without turning the individual messages into Python objects, and hands the batch +over as an Arrow IPC stream -- ordinary ``bytes``, read here with polars. + +Compare with ``parquet_serializer.py``, which does the same job as +``Batch`` -> ``Map(extract_bytes)`` -> ``BatchParser`` and pays a Python round +trip per message. + +Run with:: + + python -m sentry_streams.runner \ + --name arrow-trace-items \ + --config deployment_config/arrow_trace_items.yaml \ + sentry_streams/examples/arrow_trace_items.py +""" + +import polars as pl + +from sentry_streams.pipeline import ArrowBatchParser, StreamSink, streaming_source +from sentry_streams.pipeline.message import Message +from sentry_streams.pipeline.pipeline import Map + +TOPIC = "snuba-items" + + +def summarize(msg: Message[bytes]) -> bytes: + """Read the Arrow batch through polars and emit a one-line summary. + + The payload is an Arrow IPC stream; polars reads all of it at once, with no + row-by-row conversion. + """ + df = pl.read_ipc_stream(msg.payload) + + by_type = df.group_by("item_type").agg(pl.len().alias("rows")).sort("rows", descending=True) + summary = ", ".join(f"{row[0]}={row[1]}" for row in by_type.iter_rows()) + return f"{df.height} trace items ({summary})".encode() + + +pipeline = ( + streaming_source(name="myinput", stream_name=TOPIC) + # Takes bytes, so it must come before any step that turns messages into + # Python objects. Placing it later is a type error. + .apply(ArrowBatchParser(name="parse_arrow", schema_name=TOPIC, batch_size=1000)) + .apply(Map(name="summarize", function=summarize)) + .sink(StreamSink[bytes](name="mysink", stream_name="transformed-events")) +) 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..42429895 100644 --- a/sentry_streams/sentry_streams/pipeline/pipeline.py +++ b/sentry_streams/sentry_streams/pipeline/pipeline.py @@ -681,6 +681,85 @@ def override_config(self, loaded_config: Mapping[str, Any]) -> None: self.batch_timedelta = timedelta(**loaded_kwargs) +@dataclass +class ArrowBatchParser( + Reduce[MeasurementUnit, bytes, bytes], + Generic[MeasurementUnit], +): + """ + Batches raw Kafka payloads and decodes them into an Apache Arrow + ``RecordBatch``, entirely in Rust. + + The emitted message payload is the batch serialized as an Arrow IPC stream, + so downstream steps receive ordinary ``bytes``:: + + import polars as pl + + def to_frame(msg: Message[bytes]) -> pl.DataFrame: + return pl.read_ipc_stream(msg.payload) + + This is the fused equivalent of ``Batch`` -> ``Map(extract_bytes)`` -> + ``BatchParser``: the individual messages are never turned into Python + objects. The batch itself is serialized and copied into Python memory, which + is a deliberate simplification -- handing the ``RecordBatch`` over directly + through the Arrow C data interface is deferred. + + ``schema_name`` is the logical stream name whose ``sentry-kafka-schemas`` + entry names the message type, and so selects the extractor -- usually the + source's ``stream_name``. It is given explicitly rather than inferred so the + step does not depend on where it sits in the pipeline. + + Limitations of the current implementation, all of which fail loudly: + + * **Protobuf topics only.** A JSON or msgpack topic raises at startup. + * **The Arrow schema is hardcoded in Rust**, per message type, so there is + nothing to configure here and adding a column needs a release. + * **Rust adapter only.** The pure-Python Arroyo adapter raises + ``NotImplementedError``. + * **It reads raw payloads**, so it takes ``bytes`` and must come before any + step that converts messages into Python objects. Placing it after one is a + type error, and a panic at runtime if the types were bypassed. + * A payload that fails to decode **fails the process**: batching collapses + offsets, so there is no single offset to dead-letter. + + Configured by batch size and/or batch_timedelta exactly like ``Batch``, and + both are overridable from the deployment config's ``steps_config``. + """ + + schema_name: str + batch_size: int | None = None + batch_timedelta: timedelta | None = timedelta(seconds=10) + step_type: StepType = StepType.REDUCE + + def validate(self) -> None: + """Validate that at least one of batch_size or batch_timedelta is set.""" + if self.batch_size is None and self.batch_timedelta is None: + raise ValueError("At least one of batch_size or batch_timedelta must be set.") + + @property + def group_by(self) -> Optional[GroupBy]: + return None + + @property + def windowing(self) -> Window[MeasurementUnit]: + return TumblingWindow(self.batch_size, self.batch_timedelta) + + @property + def aggregate_fn(self) -> Callable[[], Accumulator[Message[bytes], bytes]]: + raise NotImplementedError( + "ArrowBatchParser is implemented natively in Rust and has no Python accumulator." + ) + + def override_config(self, loaded_config: Mapping[str, Any]) -> None: + if loaded_config.get("batch_size") is not None: + self.batch_size = loaded_config.get("batch_size") + + if loaded_config.get("batch_timedelta") is not None: + loaded_kwargs = loaded_config.get("batch_timedelta") + assert isinstance(loaded_kwargs, Mapping) + self.batch_timedelta = timedelta(**loaded_kwargs) + + @dataclass class FlatMap(Transform[TIn, TOut], Generic[TIn, TOut]): """ diff --git a/sentry_streams/sentry_streams/rust_streams.pyi b/sentry_streams/sentry_streams/rust_streams.pyi index e4c359cf..fc0b8441 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/arrow_batch_parser.rs b/sentry_streams/src/arrow_batch_parser.rs new file mode 100644 index 00000000..8fa19301 --- /dev/null +++ b/sentry_streams/src/arrow_batch_parser.rs @@ -0,0 +1,644 @@ +//! The Arrow batch parser step: batches raw Kafka payloads and decodes them into +//! 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 +//! [`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_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}; +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. +/// +/// 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 +/// +/// 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], + 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) + }) +} + +/// 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, + 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 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, + RawMessage { + payload, + headers: vec![], + timestamp: ts, + // 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, + }, + ) + }) + .map_err(|e| StrategyError::Other(Box::new(e)))?; + + Ok(Message::new_any_message( + RoutedValue { + route: route.clone(), + payload: RoutedValuePayload::PyStreamingMessage(PyStreamingMessage::RawMessage { + 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::{StringArray, UInt64Array}; + use arrow::ipc::reader::StreamReader; + 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) + } + + 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::RawMessage { content }) => { + content + } + _ => panic!("expected a RawMessage carrying the serialized batch"), + }; + traced_with_gil!(|py| { + let borrowed = content.bind(py).borrow(); + (from_ipc_stream(&borrowed.payload), borrowed.schema.clone()) + }) + } + + #[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"); + + // `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 + /// 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 bytes: Vec = out[0].bind(py).extract().unwrap(); + assert_eq!(from_ipc_stream(&bytes).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); + } + + /// 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::RawMessage { + content, + }) => content, + _ => unreachable!(), + }; + 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("read_ipc_stream", (py_bytes,)).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" + ); + }); + } +} 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"); + }); + } + } } diff --git a/sentry_streams/src/extractors/mod.rs b/sentry_streams/src/extractors/mod.rs new file mode 100644 index 00000000..c2fd6248 --- /dev/null +++ b/sentry_streams/src/extractors/mod.rs @@ -0,0 +1,142 @@ +//! 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 +//! `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..d11bfa40 --- /dev/null +++ b/sentry_streams/src/extractors/trace_item.rs @@ -0,0 +1,752 @@ +//! 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. 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::{ + ArrayRef, BinaryBuilder, BooleanBuilder, Float64Builder, Int64Builder, MapBuilder, + MapFieldNames, RecordBatch, StringBuilder, TimestampMicrosecondBuilder, UInt32Builder, + UInt64Builder, +}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use prost::Message; +use prost_types::Timestamp; +use sentry_protos::snuba::v1::{any_value::Value, TraceItem, TraceItemType}; +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}"), + } +} + +/// 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); + } + // 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 => {} + } + } + + // 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())] + ); + } + + /// 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_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 { + 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, 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] + 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 b8819a2b..e15421b6 100644 --- a/sentry_streams/src/lib.rs +++ b/sentry_streams/src/lib.rs @@ -1,4 +1,5 @@ use pyo3::prelude::*; +mod arrow_batch_parser; mod batch_step; mod broadcaster; mod callers; @@ -6,6 +7,7 @@ mod commit_policy; mod committable; mod consumer; mod dev_null_sink; +mod extractors; mod filter_step; mod gcs_writer; mod header_filter_step; 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..577359d8 --- /dev/null +++ b/sentry_streams/tests/adapters/arroyo/test_arrow_batch_parser.py @@ -0,0 +1,188 @@ +"""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. +""" + +import subprocess +import sys +import tempfile +from pathlib import Path +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", schema_name=SOURCE_TOPIC, batch_size=100)) + .apply(Map(name="consume", function=consume_batch)) + .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_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() + + +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: + """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", schema_name=SOURCE_TOPIC, 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_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"} + + 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] = 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] = 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 + assert step.batch_timedelta is not None + assert step.batch_timedelta.total_seconds() == 3.0