diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 3f1a5c9..d9b6612 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -27,7 +27,7 @@ jobs: # Build the Rust project # This command compiles the project and checks for errors - name: Build project - run: cargo build --verbose + run: cargo build --verbose --release # Check the formatting - name: Formatting check diff --git a/Cargo.lock b/Cargo.lock index fa33d00..0dccb05 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8,6 +8,7 @@ version = "0.1.0" dependencies = [ "anyhow", "arrow", + "async-stream", "async-trait", "axum", "chrono", @@ -20,11 +21,12 @@ dependencies = [ "lazy_static", "once_cell", "parquet", + "polars", "prometheus", "rand 0.8.5", "rayon", "regex", - "reqwest", + "reqwest 0.11.27", "rnltk", "serde", "serde_json", @@ -36,6 +38,7 @@ dependencies = [ "tokio", "tokio-executor-trait", "tokio-reactor-trait", + "tokio-stream", "tracing", "tracing-appender", "tracing-subscriber", @@ -240,6 +243,33 @@ dependencies = [ "num-traits", ] +[[package]] +name = "argminmax" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70f13d10a41ac8d2ec79ee34178d61e6f47a29c2edfe7ef1721c7383b0359e65" +dependencies = [ + "num-traits", +] + +[[package]] +name = "array-init-cursor" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed51fe0f224d1d4ea768be38c51f9f831dee9d05c163c11fba0b8c44387b1fc3" + +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" + [[package]] name = "arrow" version = "54.3.1" @@ -609,6 +639,28 @@ dependencies = [ "reactor-trait", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "async-task" version = "4.7.1" @@ -635,6 +687,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "atoi_simd" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a49e05797ca52e312a0c658938b7d00693ef037799ef7187678f212d7684cf" +dependencies = [ + "debug_unsafe", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -741,6 +802,15 @@ version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "55248b47b0caf0546f7988906588779981c43bb1bc9d0c44087278f80cdb44ba" +[[package]] +name = "bincode" +version = "1.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +dependencies = [ + "serde", +] + [[package]] name = "bitflags" version = "1.3.2" @@ -752,6 +822,22 @@ name = "bitflags" version = "2.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967" +dependencies = [ + "serde", +] + +[[package]] +name = "blake3" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3888aaa89e4b2a40fca9848e400f6a658a5a3978de7be858e209cafa8be9a4a0" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", +] [[package]] name = "block-buffer" @@ -834,6 +920,12 @@ dependencies = [ "serde_with", ] +[[package]] +name = "boxcar" +version = "0.2.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26c4925bc979b677330a8c7fe7a8c94af2dbb4a2d37b4a20a80d884400f46baa" + [[package]] name = "brotli" version = "7.0.0" @@ -866,6 +958,20 @@ name = "bytemuck" version = "1.23.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c76a5792e44e4abe34d3abf15636779261d45a7450612059293d1d2cfc63422" +dependencies = [ + "bytemuck_derive", +] + +[[package]] +name = "bytemuck_derive" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ecc273b49b3205b83d648f0690daa588925572cc5063745bfe547fe7ec8e1a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] [[package]] name = "byteorder" @@ -878,6 +984,9 @@ name = "bytes" version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" +dependencies = [ + "serde", +] [[package]] name = "calendrical_calculations" @@ -889,6 +998,15 @@ dependencies = [ "displaydoc", ] +[[package]] +name = "castaway" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0abae9be0aaf9ea96a3b1b8b1b55c602ca751eba1b1500220cea4ecbafe7c0d5" +dependencies = [ + "rustversion", +] + [[package]] name = "cbc" version = "0.1.2" @@ -915,6 +1033,12 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9555578bc9e57714c812a1f84e4fc5b4d21fcb063490c624de019f7464c91268" +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + [[package]] name = "chrono" version = "0.4.41" @@ -930,6 +1054,27 @@ dependencies = [ "windows-link", ] +[[package]] +name = "chrono-tz" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "efdce149c370f133a071ca8ef6ea340b7b88748ab0810097a9e2976eaa34b4f3" +dependencies = [ + "chrono", + "chrono-tz-build", + "phf", +] + +[[package]] +name = "chrono-tz-build" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f10f8c9340e31fc120ff885fcdb54a0b48e474bbd77cab557f0c30a3e569402" +dependencies = [ + "parse-zoneinfo", + "phf_codegen", +] + [[package]] name = "cipher" version = "0.4.4" @@ -998,6 +1143,32 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" +[[package]] +name = "comfy-table" +version = "7.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a65ebfec4fb190b6f90e944a817d60499ee0744e582530e2c9900a22e591d9a" +dependencies = [ + "crossterm", + "unicode-segmentation", + "unicode-width", +] + +[[package]] +name = "compact_str" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b79c4069c6cad78e2e0cdfcbd26275770669fb39fd308a752dc110e83b9af32" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -1046,6 +1217,12 @@ dependencies = [ "tiny-keccak", ] +[[package]] +name = "constant_time_eq" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c74b8349d32d297c9134b8c88677813a227df8f779daa29bfc29c183fe3dca6" + [[package]] name = "cookie-factory" version = "0.3.3" @@ -1133,12 +1310,43 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "crossbeam-queue" +version = "0.3.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f58bbc28f91df819d0aa2a2c00cd19754769c2fad90579b3592b1c9ba7a3115" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-utils" version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crossterm" +version = "0.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "829d955a0bb380ef178a640b91779e3987da38c9aea133b20614cfed8cdea9c6" +dependencies = [ + "bitflags 2.9.1", + "crossterm_winapi", + "parking_lot", + "rustix 0.38.44", + "winapi", +] + +[[package]] +name = "crossterm_winapi" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acdd7c62a3665c7f6830a51635d9ac9b23ed385797f70a83bb8bafe9c572ab2b" +dependencies = [ + "winapi", +] + [[package]] name = "crunchy" version = "0.2.3" @@ -1217,6 +1425,12 @@ version = "2.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a2330da5de22e8a3cb63252ce2abb30116bf5265e89c0e01bc17015ce30a476" +[[package]] +name = "debug_unsafe" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85d3cef41d236720ed453e102153a53e4cc3d2fde848c0078a50cf249e8e3e5b" + [[package]] name = "der" version = "0.7.10" @@ -1428,6 +1642,12 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "ethnum" +version = "1.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca81e6b4777c89fd810c25a4be2b1bd93ea034fbe58e6a75216a34c6b82c539b" + [[package]] name = "event-listener" version = "2.5.3" @@ -1464,6 +1684,18 @@ dependencies = [ "async-trait", ] +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + +[[package]] +name = "fast-float2" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8eb564c5c7423d25c886fb561d1e4ee69f72354d16918afa32c08811f6b6a55" + [[package]] name = "fastrand" version = "1.9.0" @@ -1525,6 +1757,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a3d7db9596fecd151c5f638c0ee5d5bd487b6e0ea232e5dc96d5250f6f94b1d" dependencies = [ "crc32fast", + "libz-rs-sys", "miniz_oxide", ] @@ -1545,6 +1778,12 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + [[package]] name = "foreign-types" version = "0.3.2" @@ -1569,6 +1808,16 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "fs4" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8640e34b88f7652208ce9e88b1a37a2ae95227d84abec377ccd3c5cfeb141ed4" +dependencies = [ + "rustix 1.0.7", + "windows-sys 0.59.0", +] + [[package]] name = "futures" version = "0.3.31" @@ -1703,8 +1952,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" dependencies = [ "cfg-if", + "js-sys", "libc", "wasi 0.11.1+wasi-snapshot-preview1", + "wasm-bindgen", ] [[package]] @@ -1714,9 +1965,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" dependencies = [ "cfg-if", + "js-sys", "libc", "r-efi", "wasi 0.14.2+wasi-0.2.4", + "wasm-bindgen", ] [[package]] @@ -1725,6 +1978,12 @@ version = "0.31.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07e28edb80900c19c28f1072f2e8aeca7fa06b23cd4169cefe1af5aa3260783f" +[[package]] +name = "glob" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2" + [[package]] name = "h2" version = "0.3.26" @@ -1744,6 +2003,25 @@ dependencies = [ "tracing", ] +[[package]] +name = "h2" +version = "0.4.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9421a676d1b147b16b82c9225157dc629087ef8ec4d5e2960f9437a90dac0a5" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http 1.3.1", + "indexmap 2.9.0", + "slab", + "tokio", + "tokio-util", + "tracing", +] + [[package]] name = "half" version = "2.6.0" @@ -1769,6 +2047,8 @@ checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" dependencies = [ "ahash", "allocator-api2", + "rayon", + "serde", ] [[package]] @@ -1776,6 +2056,13 @@ name = "hashbrown" version = "0.15.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5971ac85611da7067dbfcabef3c70ebb5606018acd9e2a3903a0da507521e0d5" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", + "rayon", + "serde", +] [[package]] name = "heck" @@ -1906,6 +2193,12 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "humantime" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b112acc8b3adf4b107a8ec20977da0273a8c386765a3ec0229bd500a1443f9f" + [[package]] name = "hyper" version = "0.14.32" @@ -1916,7 +2209,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "h2", + "h2 0.3.26", "http 0.2.12", "http-body 0.4.6", "httparse", @@ -1939,6 +2232,7 @@ dependencies = [ "bytes", "futures-channel", "futures-util", + "h2 0.4.10", "http 1.3.1", "http-body 1.0.1", "httparse", @@ -1989,6 +2283,7 @@ dependencies = [ "hyper 1.6.0", "hyper-util", "rustls 0.23.28", + "rustls-native-certs 0.8.1", "rustls-pki-types", "tokio", "tokio-rustls 0.26.2", @@ -2001,6 +2296,7 @@ version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc2fdfdbff08affe55bb779f33b053aa1fe5dd5b54c257343c17edfa55711bdb" dependencies = [ + "base64 0.22.1", "bytes", "futures-channel", "futures-core", @@ -2008,7 +2304,9 @@ dependencies = [ "http 1.3.1", "http-body 1.0.1", "hyper 1.6.0", + "ipnet", "libc", + "percent-encoding", "pin-project-lite", "socket2 0.5.10", "tokio", @@ -2639,6 +2937,16 @@ version = "2.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "469fb0b9cefa57e3ef31275ee7cacb78f2fdca44e4765491884a2b119d4eb130" +[[package]] +name = "iri-string" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbc5ebe9c3a1a7a5127f920a418f7585e9e758e911d0466ed004f393b0e380b2" +dependencies = [ + "memchr", + "serde", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.1" @@ -2813,6 +3121,15 @@ dependencies = [ "redox_syscall 0.5.13", ] +[[package]] +name = "libz-rs-sys" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "172a788537a2221661b480fee8dc5f96c580eb34fa88764d3205dc356c7e4221" +dependencies = [ + "zlib-rs", +] + [[package]] name = "linux-raw-sys" version = "0.3.8" @@ -2859,6 +3176,31 @@ version = "0.4.27" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "lz4" +version = "1.28.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a20b523e860d03443e98350ceaac5e71c6ba89aea7d960769ec3ce37f4de5af4" +dependencies = [ + "lz4-sys", +] + +[[package]] +name = "lz4-sys" +version = "1.11.1+lz4-1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bd8c0d6c6ed0cd30b3652886bb8711dc4bb01d637a68105a3d5158039b418e6" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "lz4_flex" version = "0.11.4" @@ -2915,6 +3257,15 @@ version = "2.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" +[[package]] +name = "memmap2" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd3f7eed9d3848f8b98834af67102b720745c4ec028fcd0aa0239277e7de374f" +dependencies = [ + "libc", +] + [[package]] name = "mime" version = "0.3.17" @@ -3022,6 +3373,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "now" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d89e9874397a1f0a52fc1f197a8effd9735223cb2390e9dcc83ac6cd02923d0" +dependencies = [ + "chrono", +] + [[package]] name = "nu-ansi-term" version = "0.46.0" @@ -3127,6 +3487,41 @@ dependencies = [ "memchr", ] +[[package]] +name = "object_store" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781f96d79ed0f961a7021424ab01840efbda64ae7a505aaea195efc91eaaec4" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "chrono", + "form_urlencoded", + "futures", + "http 1.3.1", + "http-body-util", + "humantime", + "hyper 1.6.0", + "itertools 0.14.0", + "parking_lot", + "percent-encoding", + "quick-xml", + "rand 0.9.1", + "reqwest 0.12.20", + "ring", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror 2.0.12", + "tokio", + "tracing", + "url", + "walkdir", + "wasm-bindgen-futures", + "web-time", +] + [[package]] name = "oid-registry" version = "0.8.1" @@ -3344,6 +3739,15 @@ dependencies = [ "syn", ] +[[package]] +name = "parse-zoneinfo" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f2a05b18d44e2957b88f96ba460715e295bc1d7510468a2f3d3b44535d26c24" +dependencies = [ + "regex", +] + [[package]] name = "paste" version = "1.0.15" @@ -3375,6 +3779,44 @@ version = "2.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e" +[[package]] +name = "phf" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_codegen" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a" +dependencies = [ + "phf_generator", + "phf_shared", +] + +[[package]] +name = "phf_generator" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d" +dependencies = [ + "phf_shared", + "rand 0.8.5", +] + +[[package]] +name = "phf_shared" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project-lite" version = "0.2.16" @@ -3446,6 +3888,506 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" +[[package]] +name = "planus" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3daf8e3d4b712abe1d690838f6e29fb76b76ea19589c4afa39ec30e12f62af71" +dependencies = [ + "array-init-cursor", + "hashbrown 0.15.4", +] + +[[package]] +name = "polars" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "443824f43bca39b178353d6c09e4b44e115b21f107a5654d5f980d20b432a303" +dependencies = [ + "getrandom 0.2.16", + "polars-arrow", + "polars-core", + "polars-error", + "polars-io", + "polars-lazy", + "polars-ops", + "polars-parquet", + "polars-sql", + "polars-time", + "polars-utils", + "version_check", +] + +[[package]] +name = "polars-arrow" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "809c5340e9e6c16eee5a07585161bae99f903f53af7402075efec23ee75fce5b" +dependencies = [ + "atoi_simd", + "bitflags 2.9.1", + "bytemuck", + "chrono", + "chrono-tz", + "dyn-clone", + "either", + "ethnum", + "getrandom 0.2.16", + "hashbrown 0.15.4", + "itoa", + "lz4", + "num-traits", + "polars-arrow-format", + "polars-error", + "polars-schema", + "polars-utils", + "serde", + "simdutf8", + "streaming-iterator", + "strum_macros", + "version_check", + "zstd", +] + +[[package]] +name = "polars-arrow-format" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "863c04c514be005eced7db7053e20d49f7e7a58048a282fa52dfea1fd5434e78" +dependencies = [ + "planus", + "serde", +] + +[[package]] +name = "polars-compute" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b8802ff2cccea01a845ea8267a7600e495747ed109035bb5020c33eb8717ff4" +dependencies = [ + "atoi_simd", + "bytemuck", + "chrono", + "either", + "fast-float2", + "hashbrown 0.15.4", + "itoa", + "num-traits", + "polars-arrow", + "polars-error", + "polars-utils", + "rand 0.8.5", + "ryu", + "serde", + "skiplist", + "strength_reduce", + "strum_macros", + "version_check", +] + +[[package]] +name = "polars-core" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fc3c99d7000be1be11665e1e260b93dc3b927342b9da3b53d9a1ac264e4343d" +dependencies = [ + "bitflags 2.9.1", + "boxcar", + "bytemuck", + "chrono", + "chrono-tz", + "comfy-table", + "either", + "hashbrown 0.14.5", + "hashbrown 0.15.4", + "indexmap 2.9.0", + "itoa", + "num-traits", + "polars-arrow", + "polars-compute", + "polars-error", + "polars-row", + "polars-schema", + "polars-utils", + "rand 0.8.5", + "rand_distr", + "rayon", + "regex", + "serde", + "serde_json", + "strum_macros", + "uuid", + "version_check", + "xxhash-rust", +] + +[[package]] +name = "polars-error" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1397c17712e61a55fdd45c033a69f0451fde2973ff2609c22e363e21d68f11ef" +dependencies = [ + "object_store", + "parking_lot", + "polars-arrow-format", + "regex", + "signal-hook", + "simdutf8", +] + +[[package]] +name = "polars-expr" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33d3aa6722c9a3e0b721ec2bcdc4affd9e50e4cb606cd81bb94535a9a5a6ade9" +dependencies = [ + "bitflags 2.9.1", + "hashbrown 0.15.4", + "num-traits", + "polars-arrow", + "polars-compute", + "polars-core", + "polars-io", + "polars-ops", + "polars-plan", + "polars-row", + "polars-time", + "polars-utils", + "rand 0.8.5", + "rayon", + "recursive", +] + +[[package]] +name = "polars-io" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a632d442a99821250a8fa66f7d488bf5ee98e5f515e65256b12956cb81fc110" +dependencies = [ + "async-trait", + "atoi_simd", + "blake3", + "bytes", + "chrono", + "fast-float2", + "fs4", + "futures", + "glob", + "hashbrown 0.15.4", + "home", + "itoa", + "memchr", + "memmap2", + "num-traits", + "object_store", + "percent-encoding", + "polars-arrow", + "polars-core", + "polars-error", + "polars-parquet", + "polars-schema", + "polars-time", + "polars-utils", + "rayon", + "regex", + "reqwest 0.12.20", + "ryu", + "serde", + "serde_json", + "simdutf8", + "tokio", + "tokio-util", + "url", +] + +[[package]] +name = "polars-lazy" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4ed0c87bdc8820447a38ae8efdb5a51a5a93e8bd528cffb05d05cf1145e4161" +dependencies = [ + "bitflags 2.9.1", + "chrono", + "either", + "memchr", + "polars-arrow", + "polars-compute", + "polars-core", + "polars-expr", + "polars-io", + "polars-mem-engine", + "polars-ops", + "polars-plan", + "polars-stream", + "polars-time", + "polars-utils", + "rayon", + "version_check", +] + +[[package]] +name = "polars-mem-engine" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "675294ddf9174029e48caa4e59b0665ea64bfb784a366b197690895a6ed65c68" +dependencies = [ + "futures", + "memmap2", + "polars-arrow", + "polars-core", + "polars-error", + "polars-expr", + "polars-io", + "polars-ops", + "polars-plan", + "polars-time", + "polars-utils", + "rayon", + "recursive", + "tokio", +] + +[[package]] +name = "polars-ops" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1eb4db68956f857c52eeda072d87644a7b42eac41d55073af94dfac8441af6cf" +dependencies = [ + "argminmax", + "base64 0.22.1", + "bytemuck", + "chrono", + "chrono-tz", + "either", + "hashbrown 0.15.4", + "hex", + "indexmap 2.9.0", + "libm", + "memchr", + "num-traits", + "polars-arrow", + "polars-compute", + "polars-core", + "polars-error", + "polars-schema", + "polars-utils", + "rayon", + "regex", + "regex-syntax 0.8.5", + "strum_macros", + "unicode-normalization", + "unicode-reverse", + "version_check", +] + +[[package]] +name = "polars-parquet" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c849c10edd9511ccd4ec4130e283ee3a8b3bb48a7d74ac6354c1c20add81065" +dependencies = [ + "async-stream", + "base64 0.22.1", + "brotli", + "bytemuck", + "ethnum", + "flate2", + "futures", + "hashbrown 0.15.4", + "lz4", + "num-traits", + "polars-arrow", + "polars-compute", + "polars-error", + "polars-parquet-format", + "polars-utils", + "serde", + "simdutf8", + "snap", + "streaming-decompression", + "zstd", +] + +[[package]] +name = "polars-parquet-format" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c025243dcfe8dbc57e94d9f82eb3bef10b565ab180d5b99bed87fd8aea319ce1" +dependencies = [ + "async-trait", + "futures", +] + +[[package]] +name = "polars-plan" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71fb4412c42bf637c2c02a617381c682ed425d9c8e4bd1fcb85cf352ed2a67c6" +dependencies = [ + "bitflags 2.9.1", + "bytemuck", + "bytes", + "chrono", + "chrono-tz", + "either", + "futures", + "hashbrown 0.15.4", + "memmap2", + "num-traits", + "percent-encoding", + "polars-arrow", + "polars-compute", + "polars-core", + "polars-io", + "polars-ops", + "polars-parquet", + "polars-time", + "polars-utils", + "rayon", + "recursive", + "regex", + "strum_macros", + "version_check", +] + +[[package]] +name = "polars-row" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08fb77ac1d37340d9cfe57cf58000cf3d9cce429e10d25066952c6145c684cc0" +dependencies = [ + "bitflags 2.9.1", + "bytemuck", + "polars-arrow", + "polars-compute", + "polars-error", + "polars-utils", +] + +[[package]] +name = "polars-schema" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ada7c7e2fbbeffbdd67628cd8a89f02b0a8d21c71d34e297e2463a7c17575203" +dependencies = [ + "indexmap 2.9.0", + "polars-error", + "polars-utils", + "serde", + "version_check", +] + +[[package]] +name = "polars-sql" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a8e512b1f05ffda9963fe8f6a7c62dcba86be85218bc033ecdad2802cc1b1a0" +dependencies = [ + "bitflags 2.9.1", + "hex", + "polars-core", + "polars-error", + "polars-lazy", + "polars-ops", + "polars-plan", + "polars-time", + "polars-utils", + "rand 0.8.5", + "regex", + "serde", + "sqlparser", +] + +[[package]] +name = "polars-stream" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b0a02d8050acd9b64ed7e36c5bc96f6d4f46a940220f9c0e34c96b51f830f8c" +dependencies = [ + "async-channel", + "async-trait", + "atomic-waker", + "bitflags 2.9.1", + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-queue", + "crossbeam-utils", + "futures", + "memmap2", + "parking_lot", + "percent-encoding", + "pin-project-lite", + "polars-arrow", + "polars-core", + "polars-error", + "polars-expr", + "polars-io", + "polars-mem-engine", + "polars-ops", + "polars-parquet", + "polars-plan", + "polars-utils", + "rand 0.8.5", + "rayon", + "recursive", + "slotmap", + "tokio", + "version_check", +] + +[[package]] +name = "polars-time" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72e84a30110880ffede8d93c085fc429ab1b8bf1acf3d6d489143dd34be374c4" +dependencies = [ + "atoi_simd", + "bytemuck", + "chrono", + "chrono-tz", + "now", + "num-traits", + "polars-arrow", + "polars-compute", + "polars-core", + "polars-error", + "polars-ops", + "polars-utils", + "rayon", + "regex", + "strum_macros", +] + +[[package]] +name = "polars-utils" +version = "0.49.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a05e033960552c47fc35afe14d5af5b29696acc97ae5d3c585ebc33c246cc15f" +dependencies = [ + "bincode", + "bytemuck", + "bytes", + "compact_str", + "flate2", + "foldhash", + "hashbrown 0.15.4", + "indexmap 2.9.0", + "libc", + "memmap2", + "num-traits", + "polars-error", + "rand 0.8.5", + "raw-cpuid", + "rayon", + "regex", + "rmp-serde", + "serde", + "serde_json", + "slotmap", + "stacker", + "version_check", +] + [[package]] name = "polling" version = "2.8.0" @@ -3559,20 +4501,94 @@ dependencies = [ name = "protobuf" version = "3.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d65a1d4ddae7d8b5de68153b48f6aa3bba8cb002b243dbdbc55a5afbc98f99f4" +checksum = "d65a1d4ddae7d8b5de68153b48f6aa3bba8cb002b243dbdbc55a5afbc98f99f4" +dependencies = [ + "once_cell", + "protobuf-support", + "thiserror 1.0.69", +] + +[[package]] +name = "protobuf-support" +version = "3.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e36c2f31e0a47f9280fb347ef5e461ffcd2c52dd520d8e216b52f93b0b0d7d6" +dependencies = [ + "thiserror 1.0.69", +] + +[[package]] +name = "psm" +version = "0.1.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e944464ec8536cd1beb0bbfd96987eb5e3b72f2ecdafdc5c769a37f1fa2ae1f" +dependencies = [ + "cc", +] + +[[package]] +name = "quick-xml" +version = "0.37.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "331e97a1af0bf59823e6eadffe373d7b27f485be8748f71471c662c1f269b7fb" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "quinn" +version = "0.11.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "626214629cda6781b6dc1d316ba307189c85ba657213ce642d9c77670f8202c8" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash", + "rustls 0.23.28", + "socket2 0.5.10", + "thiserror 2.0.12", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49df843a9161c85bb8aae55f101bc0bac8bcafd637a620d9122fd7e0b2f7422e" dependencies = [ - "once_cell", - "protobuf-support", - "thiserror 1.0.69", + "bytes", + "getrandom 0.3.3", + "lru-slab", + "rand 0.9.1", + "ring", + "rustc-hash", + "rustls 0.23.28", + "rustls-pki-types", + "slab", + "thiserror 2.0.12", + "tinyvec", + "tracing", + "web-time", ] [[package]] -name = "protobuf-support" -version = "3.7.2" +name = "quinn-udp" +version = "0.5.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3e36c2f31e0a47f9280fb347ef5e461ffcd2c52dd520d8e216b52f93b0b0d7d6" +checksum = "fcebb1209ee276352ef14ff8732e24cc2b02bbac986cd74a4c81bcb2f9881970" dependencies = [ - "thiserror 1.0.69", + "cfg_aliases", + "libc", + "once_cell", + "socket2 0.5.10", + "tracing", + "windows-sys 0.59.0", ] [[package]] @@ -3649,6 +4665,25 @@ dependencies = [ "getrandom 0.3.3", ] +[[package]] +name = "rand_distr" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32cb0b9bc82b0a0876c2dd994a7e7a2683d3e7390ca40e6886785ef0c7e3ee31" +dependencies = [ + "num-traits", + "rand 0.8.5", +] + +[[package]] +name = "raw-cpuid" +version = "11.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6df7ab838ed27997ba19a4664507e6f82b41fe6e20be42929332156e5e85146" +dependencies = [ + "bitflags 2.9.1", +] + [[package]] name = "rawpointer" version = "0.2.1" @@ -3706,6 +4741,26 @@ dependencies = [ "futures-io", ] +[[package]] +name = "recursive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0786a43debb760f491b1bc0269fe5e84155353c67482b9e60d0cfb596054b43e" +dependencies = [ + "recursive-proc-macro-impl", + "stacker", +] + +[[package]] +name = "recursive-proc-macro-impl" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" +dependencies = [ + "quote", + "syn", +] + [[package]] name = "redox_syscall" version = "0.3.5" @@ -3819,7 +4874,7 @@ dependencies = [ "encoding_rs", "futures-core", "futures-util", - "h2", + "h2 0.3.26", "http 0.2.12", "http-body 0.4.6", "hyper 0.14.32", @@ -3849,6 +4904,48 @@ dependencies = [ "winreg", ] +[[package]] +name = "reqwest" +version = "0.12.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eabf4c97d9130e2bf606614eb937e86edac8292eaa6f422f995d7e8de1eb1813" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-util", + "h2 0.4.10", + "http 1.3.1", + "http-body 1.0.1", + "http-body-util", + "hyper 1.6.0", + "hyper-rustls 0.27.7", + "hyper-util", + "js-sys", + "log", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls 0.23.28", + "rustls-native-certs 0.8.1", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper 1.0.2", + "tokio", + "tokio-rustls 0.26.2", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", +] + [[package]] name = "ring" version = "0.17.14" @@ -3863,6 +4960,28 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rmp" +version = "0.8.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "228ed7c16fa39782c3b3468e974aec2795e9089153cd08ee2e9aefb3613334c4" +dependencies = [ + "byteorder", + "num-traits", + "paste", +] + +[[package]] +name = "rmp-serde" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52e599a477cf9840e92f2cde9a7189e67b42c57532749bf90aea6ec10facd4db" +dependencies = [ + "byteorder", + "rmp", + "serde", +] + [[package]] name = "rnltk" version = "0.4.0" @@ -3883,6 +5002,12 @@ version = "0.1.25" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "989e6739f80c4ad5b13e0fd7fe89531180375b18520cc8c82080e4dc4035b84f" +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + [[package]] name = "rustc_version" version = "0.4.1" @@ -4030,6 +5155,7 @@ version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "229a4a4c221013e7e1f1a043678c5cc39fe5171437c88fb47151a21e6f5b5c79" dependencies = [ + "web-time", "zeroize", ] @@ -4084,6 +5210,15 @@ dependencies = [ "cipher", ] +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "schannel" version = "0.1.27" @@ -4326,6 +5461,16 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "signal-hook" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d881a16cf4426aa584979d30bd82cb33429027e42122b169753d6ef1085ed6e2" +dependencies = [ + "libc", + "signal-hook-registry", +] + [[package]] name = "signal-hook-registry" version = "1.4.5" @@ -4354,12 +5499,36 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" +[[package]] +name = "siphasher" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56199f7ddabf13fe5074ce809e7d3f42b42ae711800501b5b16ea82ad029c39d" + +[[package]] +name = "skiplist" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eec25f46463fcdc5e02f388c2780b1b58e01be81a8378e62ec60931beccc3f6" +dependencies = [ + "rand 0.8.5", +] + [[package]] name = "slab" version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04dc19736151f35336d325007ac991178d504a119863a2fcb3758cdb5e52c50d" +[[package]] +name = "slotmap" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbff4acf519f630b3a3ddcfaea6c06b42174d9a44bc70c620e9ed1649d58b82a" +dependencies = [ + "version_check", +] + [[package]] name = "smallvec" version = "1.15.1" @@ -4434,18 +5603,61 @@ dependencies = [ "unicode-segmentation", ] +[[package]] +name = "sqlparser" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05a528114c392209b3264855ad491fcce534b94a38771b0a0b97a79379275ce8" +dependencies = [ + "log", +] + [[package]] name = "stable_deref_trait" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" +[[package]] +name = "stacker" +version = "0.1.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cddb07e32ddb770749da91081d8d0ac3a16f1a569a18b20348cd371f5dead06b" +dependencies = [ + "cc", + "cfg-if", + "libc", + "psm", + "windows-sys 0.59.0", +] + [[package]] name = "static_assertions" version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" +[[package]] +name = "streaming-decompression" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf6cc3b19bfb128a8ad11026086e31d3ce9ad23f8ea37354b31383a187c44cf3" +dependencies = [ + "fallible-streaming-iterator", +] + +[[package]] +name = "streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b2231b7c3057d5e4ad0156fb3dc807d900806020c5ffa3ee6ff2c8c76fb8520" + +[[package]] +name = "strength_reduce" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe895eb47f22e2ddd4dabc02bce419d2e643c8e3b585c78158b349195bc24d82" + [[package]] name = "strsim" version = "0.11.1" @@ -4475,6 +5687,19 @@ dependencies = [ "syn", ] +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "rustversion", + "syn", +] + [[package]] name = "subtle" version = "2.6.1" @@ -4503,6 +5728,9 @@ name = "sync_wrapper" version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] [[package]] name = "synstructure" @@ -4711,6 +5939,21 @@ dependencies = [ "zerovec 0.11.2", ] +[[package]] +name = "tinyvec" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09b3661f17e86524eccd4371ab0429194e0d7c008abb45f7a7495b1719463c71" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + [[package]] name = "tokenizers" version = "0.21.1" @@ -4873,6 +6116,24 @@ dependencies = [ "tracing", ] +[[package]] +name = "tower-http" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adc82fd73de2a9722ac5da747f12383d2bfdb93591ee6c58486e0097890f05f2" +dependencies = [ + "bitflags 2.9.1", + "bytes", + "futures-util", + "http 1.3.1", + "http-body 1.0.1", + "iri-string", + "pin-project-lite", + "tower", + "tower-layer", + "tower-service", +] + [[package]] name = "tower-layer" version = "0.3.3" @@ -5006,6 +6267,15 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" +[[package]] +name = "unicode-normalization" +version = "0.1.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5033c97c4262335cded6d6fc3e5c18ab755e1a3dc96376350f3d8e9f009ad956" +dependencies = [ + "tinyvec", +] + [[package]] name = "unicode-normalization-alignments" version = "0.1.12" @@ -5015,6 +6285,15 @@ dependencies = [ "smallvec", ] +[[package]] +name = "unicode-reverse" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b6f4888ebc23094adfb574fdca9fdc891826287a6397d2cd28802ffd6f20c76" +dependencies = [ + "unicode-segmentation", +] + [[package]] name = "unicode-segmentation" version = "1.12.0" @@ -5094,6 +6373,17 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "uuid" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cf4199d1e5d15ddd86a694e4d0dffa9c323ce759fea589f00fef9d81cc1931d" +dependencies = [ + "getrandom 0.3.3", + "js-sys", + "wasm-bindgen", +] + [[package]] name = "valuable" version = "0.1.1" @@ -5118,6 +6408,16 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "317211a0dc0ceedd78fb2ca9a44aed3d7b9b26f81870d485c07122b4350673b7" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -5213,6 +6513,19 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "web-sys" version = "0.3.77" @@ -5293,6 +6606,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" +[[package]] +name = "winapi-util" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" +dependencies = [ + "windows-sys 0.59.0", +] + [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" @@ -5657,6 +6979,12 @@ dependencies = [ "rustix 1.0.7", ] +[[package]] +name = "xxhash-rust" +version = "0.8.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" + [[package]] name = "yoke" version = "0.7.5" @@ -5818,6 +7146,12 @@ dependencies = [ "syn", ] +[[package]] +name = "zlib-rs" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "626bd9fa9734751fc50d6060752170984d7053f5a39061f524cda68023d4db8a" + [[package]] name = "zstd" version = "0.13.3" diff --git a/Cargo.toml b/Cargo.toml index f908287..6e1a286 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -60,5 +60,8 @@ testcontainers = "0.24.0" rand = "0.8" tokio = { version = "1", features = ["macros", "rt-multi-thread"] } serde_json = "1.0" +async-stream = "0.3.6" +tokio-stream = "0.1.17" +polars = { version = "0.49.1", features = ["lazy", "parquet"] } # {{ No need for the old default main.rs binary now }} diff --git a/src/bin/producer.rs b/src/bin/producer.rs index 6189719..68c6d89 100644 --- a/src/bin/producer.rs +++ b/src/bin/producer.rs @@ -1,70 +1,29 @@ // src/bin/producer.rs use clap::Parser; -use futures::StreamExt; // For consuming the results stream -use indicatif::{HumanDuration, ProgressBar, ProgressStyle}; // Added indicatif -use lapin::{ - options::{BasicAckOptions, BasicConsumeOptions, BasicPublishOptions, QueueDeclareOptions}, - protocol::basic::AMQPProperties, - types::FieldTable, -}; -use std::time::Instant; // Added Instant -use tracing::{error, info, info_span, warn}; // Added tracing & info_span -use tracing_appender::{non_blocking, rolling}; // Added for file logging -use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer}; // Added tracing_subscriber components -use TextBlaster::config::ParquetInputConfig; -use TextBlaster::data_model::{ProcessingOutcome, TextDocument}; // Import both TextDocument and ProcessingOutcome -use TextBlaster::error::{PipelineError, Result}; // Use the library's Result type -use TextBlaster::pipeline::readers::ParquetReader; -use TextBlaster::pipeline::writers::parquet_writer::ParquetWriter; // For retry delay -use TextBlaster::utils::common::connect_rabbitmq; // Updated for shared functions +use indicatif::{ProgressBar, ProgressStyle}; +use TextBlaster::config::producer::Args; // Removed HumanDuration + // lapin::options are used by producer_logic, not directly here anymore for these specific ones + // use lapin::{ + // options::{BasicAckOptions, BasicConsumeOptions, QueueDeclareOptions}, + // types::FieldTable, + // }; + // std::time::Instant is used by producer_logic +use tracing::{error, info, warn}; +use tracing_appender::{non_blocking, rolling}; +use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer}; +// TextBlaster::data_model imports are used by producer_logic +// TextBlaster::error::{PipelineError, Result} is used +use TextBlaster::error::{PipelineError, Result}; +// TextBlaster::pipeline::writers::parquet_writer::ParquetWriter is used by producer_logic +use TextBlaster::utils::common::connect_rabbitmq; use TextBlaster::utils::prometheus_metrics::setup_prometheus_metrics; +// chrono::Utc is used by producer_logic +// TextBlaster::utils::prometheus_metrics::* is used by producer_logic -use chrono::Utc; // For consumer tag in aggregate_results -use TextBlaster::utils::prometheus_metrics::*; // Import shared metrics +// const PARQUET_WRITE_BATCH_SIZE: usize = 500; // This const is now in producer_logic.rs -const PARQUET_WRITE_BATCH_SIZE: usize = 500; // Configurable batch size for writing - -// Define command-line arguments -#[derive(Parser, Debug)] -#[command(author, version, about, long_about = None)] -struct Args { - /// Path to the input Parquet file - #[arg(short, long)] - input_file: String, - - /// Text column name in the Parquet file - #[arg(long, default_value = "text")] - text_column: String, - - /// Optional ID column name in the Parquet file - #[arg(long)] - id_column: Option, - - /// RabbitMQ connection string (e.g., amqp://guest:guest@localhost:5672/%2f) - #[arg(short, long, default_value = "amqp://guest:guest@localhost:5672/%2f")] - amqp_addr: String, - - /// Name of the queue to publish tasks to - #[arg(short = 'q', long, default_value = "task_queue")] - task_queue: String, - - /// Name of the queue to consume results/outcomes from - #[arg(short = 'r', long, default_value = "results_queue")] - results_queue: String, - - /// Path to the output Parquet file - #[arg(short = 'o', long, default_value = "output_processed.parquet")] - output_file: String, - - /// Path to the excluded output Parquet file - #[arg(short = 'e', long, default_value = "excluded.parquet")] - excluded_file: String, - - /// Optional: Port for the Prometheus metrics HTTP endpoint - #[arg(long)] - metrics_port: Option, -} +// Args struct is now imported from TextBlaster::config::Args // --- Prometheus Metrics (now imported from TextBlaster::utils::prometheus_metrics) --- @@ -86,292 +45,9 @@ fn create_progress_bar(total_items: u64, message: &str, template: &str) -> Progr pb } -// Function to publish tasks to RabbitMQ -async fn publish_tasks( - args: &Args, - conn: &lapin::Connection, - publishing_pb: &ProgressBar, -) -> Result { - let publish_channel = conn.create_channel().await?; - - // Declare the task queue (durable) - let _task_queue_info = publish_channel - .queue_declare( - &args.task_queue, - QueueDeclareOptions { - durable: true, // Ensure tasks survive broker restart - ..Default::default() - }, - Default::default(), - ) - .await?; - info!("Declared durable task queue '{}'", args.task_queue); - - // Configure and create the Parquet Reader - let parquet_config = ParquetInputConfig { - path: args.input_file.clone(), - text_column: args.text_column.clone(), - id_column: args.id_column.clone(), - batch_size: Some(1024), // Example batch size for reading - }; - let reader = ParquetReader::new(parquet_config); - - info!("Reading documents and publishing tasks..."); - let mut published_count = 0u64; - let mut read_errors = 0u64; // This counts errors before attempting to publish - let doc_iterator = reader.read_documents()?; - let publish_start_time = Instant::now(); - - for doc_result in doc_iterator { - publishing_pb.tick(); - match doc_result { - Ok(doc) => { - let _doc_span = info_span!("publishing_doc", doc_id = %doc.id).entered(); - match serde_json::to_vec(&doc) { - Ok(payload) => { - info!("Serialized document for publishing."); - let task_publish_timer = TASK_PUBLISHING_DURATION_SECONDS.start_timer(); - let publish_confirm = publish_channel - .basic_publish( - "", // Default exchange - &args.task_queue, - BasicPublishOptions::default(), - &payload, - AMQPProperties::default().with_delivery_mode(2), // Persistent - ) - .await? - .await; // Wait for broker ack/nack - task_publish_timer.observe_duration(); - - match publish_confirm { - Ok(_) => { - published_count += 1; - TASKS_PUBLISHED_TOTAL.inc(); - ACTIVE_TASKS_IN_FLIGHT.inc(); - publishing_pb.inc(1); - info!("Successfully published task and received confirmation."); - } - Err(e) => { - TASK_PUBLISH_ERRORS_TOTAL.inc(); - error!( - // doc_id is inherited from span - error = %e, - "FATAL: Failed broker confirmation for task. Stopping." - ); - // The main function will call publishing_pb.finish_with_message on error. - return Err(PipelineError::QueueError(format!( - "Publish confirmation failed: {}", - e - ))); - } - } - } - Err(e) => { - TASK_PUBLISH_ERRORS_TOTAL.inc(); - warn!(/* doc_id inherited from span */ error = %e, "Failed to serialize task. Skipping."); - read_errors += 1; - } - } - } - Err(e) => { - // doc_id is not available here as reading the document itself failed. - warn!(error = %e, "Error reading document for task. Skipping."); - read_errors += 1; - } - } - } - let publishing_duration = publish_start_time.elapsed(); - publishing_pb.finish_with_message(format!( - "Finished publishing {} tasks in {}. Read/Serialization Errors: {}", - published_count, - HumanDuration(publishing_duration), - read_errors - )); - Ok(published_count) -} - -// Function to aggregate results from RabbitMQ -async fn aggregate_results( - args: &Args, - conn: &lapin::Connection, - published_count: u64, - aggregation_pb: &ProgressBar, -) -> Result<(u64, u64, u64)> { - info!("\nStarting results aggregation phase..."); - - let consume_channel = conn.create_channel().await?; - let _results_queue_info = consume_channel - .queue_declare( - &args.results_queue, - QueueDeclareOptions { - durable: true, - ..Default::default() - }, - Default::default(), - ) - .await?; - info!("Declared durable results queue '{}'", args.results_queue); - - if let Some(parent_dir) = std::path::Path::new(&args.output_file).parent() { - tokio::fs::create_dir_all(parent_dir).await?; - } - let mut parquet_writer_output = ParquetWriter::new(&args.output_file)?; - info!("Initialized Parquet writer for: {}", args.output_file); - - let mut parquet_writer_excluded = ParquetWriter::new(&args.excluded_file)?; - info!("Initialized Parquet writer for: {}", args.excluded_file); - - let mut results_batch: Vec = Vec::with_capacity(PARQUET_WRITE_BATCH_SIZE); - let mut excluded_batch: Vec = Vec::with_capacity(PARQUET_WRITE_BATCH_SIZE); - let mut outcomes_received_count = 0u64; - let mut success_count = 0u64; - let mut filtered_count = 0u64; - let mut outcome_deserialization_errors = 0u64; - - let consumer_tag = format!( - "producer-aggregator-{}-{}", - std::process::id(), - Utc::now().timestamp() - ); - let mut consumer = consume_channel - .basic_consume( - &args.results_queue, - &consumer_tag, - BasicConsumeOptions::default(), - FieldTable::default(), - ) - .await?; - - info!( - "Waiting for outcomes from queue '{}'. Expecting {} outcomes.", - args.results_queue, published_count - ); - let aggregation_start_time = Instant::now(); - - while outcomes_received_count < published_count { - aggregation_pb.tick(); - match consumer.next().await { - Some(Ok(delivery)) => { - match serde_json::from_slice::(&delivery.data) { - Ok(outcome) => { - outcomes_received_count += 1; - RESULTS_RECEIVED_TOTAL.inc(); - aggregation_pb.inc(1); - ACTIVE_TASKS_IN_FLIGHT.dec(); - - match outcome { - ProcessingOutcome::Success(doc) => { - success_count += 1; - RESULTS_SUCCESS_TOTAL.inc(); - info!(doc_id = %doc.id, "Received successful processing outcome."); - results_batch.push(doc); - if results_batch.len() >= PARQUET_WRITE_BATCH_SIZE { - parquet_writer_output.write_batch(&results_batch)?; - info!( - batch_size = results_batch.len(), - total_received = outcomes_received_count, - total_expected = published_count, - success_count = success_count, - filtered_count = filtered_count, - "Written Parquet batch of processed documents." - ); - results_batch.clear(); - } - } - ProcessingOutcome::Filtered { document, reason } => { - filtered_count += 1; - RESULTS_FILTERED_TOTAL.inc(); - info!(doc_id = %document.id, %reason, "Received filtered processing outcome."); - excluded_batch.push(document); - if excluded_batch.len() >= PARQUET_WRITE_BATCH_SIZE { - parquet_writer_excluded.write_batch(&excluded_batch)?; - info!( - batch_size = excluded_batch.len(), - total_received = outcomes_received_count, - total_expected = published_count, - success_count = success_count, - filtered_count = filtered_count, - "Written Parquet batch of filtered documents." - ); - excluded_batch.clear(); - } - } - ProcessingOutcome::Error { - document, - error_message, - worker_id, - } => { - error!(doc_id = %document.id, worker_id = %worker_id, error = %error_message, "Task processing failed"); - RESULTS_ERROR_TOTAL.inc(); - } - } - } - Err(e) => { - outcome_deserialization_errors += 1; - RESULT_DESERIALIZATION_ERRORS_TOTAL.inc(); - ACTIVE_TASKS_IN_FLIGHT.dec(); - warn!( - delivery_tag = %delivery.delivery_tag, - error = %e, - payload = %String::from_utf8_lossy(&delivery.data), - "Failed to deserialize outcome." - ); - } - } - if let Err(ack_err) = delivery.ack(BasicAckOptions::default()).await { - error!( - delivery_tag = %delivery.delivery_tag, - error = %ack_err, - "Failed to ack outcome. Might lead to duplicate counts." - ); - } - } - Some(Err(e)) => { - error!(error = %e, "Error receiving outcome. Will attempt to finalize current results."); - break; // Exit loop on consumer error, then try to write remaining batches. - } - None => { - warn!("Consumer stream closed unexpectedly. Will attempt to finalize current results."); - break; // Exit loop if stream closes, then try to write remaining batches. - } - } - } - let aggregation_duration = aggregation_start_time.elapsed(); - aggregation_pb.finish_with_message(format!( - "Finished consuming (Received {}/{}, Desaerial. Errors: {}) in {}.", - outcomes_received_count, - published_count, - outcome_deserialization_errors, - HumanDuration(aggregation_duration) - )); - - if !results_batch.is_empty() { - info!( - "Writing final batch of successfully processed documents ({} docs)...", - results_batch.len() - ); - parquet_writer_output.write_batch(&results_batch)?; - } - if !excluded_batch.is_empty() { - info!( - "Writing final batch of excluded documents ({} docs)...", - excluded_batch.len() - ); - parquet_writer_excluded.write_batch(&excluded_batch)?; - } +// publish_tasks and aggregate_results functions are now in TextBlaster::producer_logic - info!("Closing Parquet writer for output_processed.parquet..."); - parquet_writer_output.close()?; - info!("Parquet writer (output_processed.parquet) closed successfully."); - - info!("Closing Parquet writer for excluded.parquet..."); - parquet_writer_excluded.close()?; - info!("Parquet writer (excluded.parquet) closed successfully."); - - Ok((outcomes_received_count, success_count, filtered_count)) -} - -// setup_prometheus_metrics and metrics_handler removed, now imported from utils +// setup_prometheus_metrics and metrics_handler removed, now imported from utils (already done) #[tokio::main] async fn main() -> Result<()> { @@ -420,15 +96,29 @@ async fn main() -> Result<()> { info!("Output File: {}", args.output_file); // 1. Connect to RabbitMQ - let conn = connect_rabbitmq(&args.amqp_addr).await?; + let conn = connect_rabbitmq(&args.amqp_addr).await?; // This returns a lapin::Connection + + // Create channels for publishing and consuming results + let task_publish_channel = conn.create_channel().await.map_err(PipelineError::from)?; + let results_consume_channel = conn.create_channel().await.map_err(PipelineError::from)?; + + // Optionally, set task_publish_channel to confirm mode if desired for all publishes + // task_publish_channel.confirm_select(lapin::options::ConfirmSelectOptions::default()).await + // .map_err(|e| PipelineError::QueueError(format!("Failed to set task channel to confirm mode: {}", e)))?; // --- Progress Bars --- let publishing_pb_template = "{spinner:.green} [{elapsed_precise}] {msg} Tasks published: {pos} ({per_sec}, ETA: {eta})"; let publishing_pb = create_progress_bar(0, "Publishing tasks", publishing_pb_template); - // 2. Publish Tasks - let published_count = match publish_tasks(&args, &conn, &publishing_pb).await { + // 2. Publish Tasks - now passing the channel directly + let published_count = match TextBlaster::producer_logic::publish_tasks( + &args, + &task_publish_channel, + &publishing_pb, + ) + .await + { Ok(count) => count, Err(e) => { error!("Failed during task publishing: {}", e); @@ -453,8 +143,16 @@ async fn main() -> Result<()> { ); // 3. Aggregate Results + // aggregate_results now takes the results_consume_channel directly let (outcomes_received_count, success_count, filtered_count) = - match aggregate_results(&args, &conn, published_count, &aggregation_pb).await { + match TextBlaster::producer_logic::aggregate_results( + &args, + &results_consume_channel, + published_count, + &aggregation_pb, + ) + .await + { Ok(counts) => counts, Err(e) => { error!("Failed during result aggregation: {}", e); diff --git a/src/bin/worker.rs b/src/bin/worker.rs index 8283a05..6f7a5a3 100644 --- a/src/bin/worker.rs +++ b/src/bin/worker.rs @@ -5,7 +5,7 @@ use futures::StreamExt; // For processing the consumer stream use indicatif::{ProgressBar, ProgressStyle}; use std::time::{Duration, Instant}; // Added for progress bar // Added for progress bar speed calculation // {{ Use the new load_pipeline_config function }} -use TextBlaster::config::{load_pipeline_config, PipelineConfig, StepConfig}; // Added config imports and load_pipeline_config +use TextBlaster::config::pipeline::{load_pipeline_config, PipelineConfig, StepConfig}; // Added config imports and load_pipeline_config use TextBlaster::data_model::{ProcessingOutcome, TextDocument}; // Updated import use TextBlaster::error::{PipelineError, Result}; // Use the library's Result type use TextBlaster::executor::{PipelineExecutor, ProcessingStep}; diff --git a/src/config/mod.rs b/src/config/mod.rs new file mode 100644 index 0000000..d62ca31 --- /dev/null +++ b/src/config/mod.rs @@ -0,0 +1,5 @@ +// src/config/mod.rs + +pub mod parquet; +pub mod pipeline; +pub mod producer; diff --git a/src/config/parquet.rs b/src/config/parquet.rs new file mode 100644 index 0000000..f14a47a --- /dev/null +++ b/src/config/parquet.rs @@ -0,0 +1,11 @@ +use serde::Deserialize; + +// Keep existing config +#[derive(Deserialize, Debug, Clone)] +pub struct ParquetInputConfig { + pub path: String, // Path to the Parquet file or directory + pub text_column: String, // Name of the column containing the main text + pub id_column: Option, // Optional: Name of a column to use as document ID + // Add other column mappings as needed (e.g., for metadata) + pub batch_size: Option, // Optional: Arrow batch size for reading +} diff --git a/src/config/pipeline.rs b/src/config/pipeline.rs new file mode 100644 index 0000000..44ce82c --- /dev/null +++ b/src/config/pipeline.rs @@ -0,0 +1,393 @@ +use crate::error::{PipelineError, Result}; +use serde::Deserialize; +use std::collections::HashSet; +use std::fs; // For reading the file +use std::path::Path; // For path handling // Assuming these are your error types + +// --- Pipeline Configuration --- + +/// Represents the overall pipeline configuration read from YAML. +#[derive(Deserialize, Debug, Clone)] +pub struct PipelineConfig { + pub pipeline: Vec, +} + +impl PipelineConfig { + pub fn validate(&self) -> Result<()> { + for step_config in &self.pipeline { + step_config.validate()?; + } + Ok(()) + } +} + +/// Represents a single step in the processing pipeline. +/// Uses Serde's externally tagged enum representation. +#[derive(Deserialize, Debug, Clone)] +#[serde(tag = "type")] // The 'type' field in YAML determines which variant +pub enum StepConfig { + C4QualityFilter(C4QualityParams), + GopherRepetitionFilter(GopherRepetitionParams), + GopherQualityFilter(GopherQualityParams), + C4BadWordsFilter(C4BadWordsParams), // New + LanguageDetectionFilter(LanguageDetectionParams), + FineWebQualityFilter(FineWebQualityFilterParams), // Renamed and new params struct + TokenCounter(TokenCounterParams), + // Add other filter/step types here as needed +} + +impl StepConfig { + /// Returns a string slice representing the name of the step type. + pub fn name(&self) -> &'static str { + match self { + StepConfig::C4QualityFilter(_) => "C4QualityFilter", + StepConfig::GopherRepetitionFilter(_) => "GopherRepetitionFilter", + StepConfig::GopherQualityFilter(_) => "GopherQualityFilter", + StepConfig::C4BadWordsFilter(_) => "C4BadWordsFilter", // New + StepConfig::LanguageDetectionFilter(_) => "LanguageDetectionFilter", + StepConfig::FineWebQualityFilter(_) => "FineWebQualityFilter", // Renamed + StepConfig::TokenCounter(_) => "TokenCounter", // Add cases for other StepConfig variants here + } + } + + pub fn validate(&self) -> Result<()> { + match self { + StepConfig::C4QualityFilter(params) => params.validate(), + StepConfig::GopherRepetitionFilter(params) => params.validate(), + StepConfig::GopherQualityFilter(params) => params.validate(), + StepConfig::C4BadWordsFilter(params) => params.validate(), + StepConfig::LanguageDetectionFilter(params) => params.validate(), + StepConfig::FineWebQualityFilter(params) => params.validate(), + StepConfig::TokenCounter(params) => params.validate(), + } + } +} + +/// Parameters for the C4QualityFilter. +#[derive(Deserialize, Debug, Clone)] +pub struct C4QualityParams { + pub split_paragraph: bool, + pub remove_citations: bool, + pub filter_no_terminal_punct: bool, + pub min_num_sentences: usize, + pub min_words_per_line: usize, + pub max_word_length: usize, + pub filter_lorem_ipsum: bool, + pub filter_javascript: bool, + pub filter_curly_bracket: bool, + pub filter_policy: bool, +} + +impl C4QualityParams { + pub fn validate(&self) -> Result<()> { + if self.min_num_sentences == 0 { + return Err(PipelineError::ConfigValidationError( + "C4QualityParams: min_num_sentences must be greater than 0".to_string(), + )); + } + if self.min_words_per_line == 0 { + return Err(PipelineError::ConfigValidationError( + "C4QualityParams: min_words_per_line must be greater than 0".to_string(), + )); + } + if self.max_word_length == 0 { + return Err(PipelineError::ConfigValidationError( + "C4QualityParams: max_word_length must be greater than 0".to_string(), + )); + } + Ok(()) + } +} + +/// Parameters for the GopherRepetitionFilter. +#[derive(Deserialize, Debug, Clone)] +pub struct GopherRepetitionParams { + // Use Option for optional fields + pub dup_line_frac: Option, + pub dup_para_frac: Option, + pub dup_line_char_frac: Option, + pub dup_para_char_frac: Option, + // Vec<(usize, f64)> represents lists like [[2, 0.2], [3, 0.18]] in YAML + #[serde(default)] // Use default (empty vec) if not specified + pub top_n_grams: Vec<(usize, f64)>, + #[serde(default)] + pub dup_n_grams: Vec<(usize, f64)>, +} + +impl GopherRepetitionParams { + pub fn validate(&self) -> Result<()> { + let fractions = [ + ("dup_line_frac", self.dup_line_frac), + ("dup_para_frac", self.dup_para_frac), + ("dup_line_char_frac", self.dup_line_char_frac), + ("dup_para_char_frac", self.dup_para_char_frac), + ]; + for (name, val) in fractions.iter() { + if let Some(v) = val { + if !(0.0..=1.0).contains(v) { + return Err(PipelineError::ConfigValidationError(format!( + "GopherRepetitionParams: {} must be between 0.0 and 1.0, got {}", + name, v + ))); + } + } + } + + for (name, n_grams) in [ + ("top_n_grams", &self.top_n_grams), + ("dup_n_grams", &self.dup_n_grams), + ] + .iter() + { + for (idx, (size, fraction)) in n_grams.iter().enumerate() { + if *size == 0 { + return Err(PipelineError::ConfigValidationError(format!( + "GopherRepetitionParams: n-gram size in {} at index {} must be greater than 0", + name, idx + ))); + } + if !(0.0..=1.0).contains(fraction) { + return Err(PipelineError::ConfigValidationError(format!( + "GopherRepetitionParams: n-gram fraction in {} at index {} must be between 0.0 and 1.0, got {}", + name, idx, fraction + ))); + } + } + } + Ok(()) + } +} + +/// Parameters for the GopherQualityFilter. +#[derive(Deserialize, Debug, Clone)] +pub struct GopherQualityParams { + // Use Option for optional fields + pub min_doc_words: Option, + pub max_doc_words: Option, + pub min_avg_word_length: Option, + pub max_avg_word_length: Option, + pub max_symbol_word_ratio: Option, + pub max_bullet_lines_ratio: Option, + pub max_ellipsis_lines_ratio: Option, + pub max_non_alpha_words_ratio: Option, + pub min_stop_words: Option, + // Optional list of stop words; if None, the filter's default will be used. + pub stop_words: Option>, +} + +impl GopherQualityParams { + pub fn validate(&self) -> Result<()> { + if let Some(min_doc_words) = self.min_doc_words { + if min_doc_words == 0 { + return Err(PipelineError::ConfigValidationError( + "GopherQualityParams: min_doc_words must be greater than 0".to_string(), + )); + } + } + if let Some(max_doc_words) = self.max_doc_words { + if max_doc_words == 0 { + return Err(PipelineError::ConfigValidationError( + "GopherQualityParams: max_doc_words must be greater than 0".to_string(), + )); + } + } + if let (Some(min_val), Some(max_val)) = (self.min_doc_words, self.max_doc_words) { + if min_val > max_val { + return Err(PipelineError::ConfigValidationError(format!( + "GopherQualityParams: min_doc_words ({}) cannot be greater than max_doc_words ({})", + min_val, max_val + ))); + } + } + + if let Some(min_avg_word_length) = self.min_avg_word_length { + if min_avg_word_length <= 0.0 { + return Err(PipelineError::ConfigValidationError( + "GopherQualityParams: min_avg_word_length must be greater than 0.0".to_string(), + )); + } + } + if let Some(max_avg_word_length) = self.max_avg_word_length { + if max_avg_word_length <= 0.0 { + return Err(PipelineError::ConfigValidationError( + "GopherQualityParams: max_avg_word_length must be greater than 0.0".to_string(), + )); + } + } + if let (Some(min_val), Some(max_val)) = (self.min_avg_word_length, self.max_avg_word_length) + { + if min_val > max_val { + return Err(PipelineError::ConfigValidationError(format!( + "GopherQualityParams: min_avg_word_length ({}) cannot be greater than max_avg_word_length ({})", + min_val, max_val + ))); + } + } + + let ratio_params = [ + ("max_symbol_word_ratio", self.max_symbol_word_ratio), + ("max_bullet_lines_ratio", self.max_bullet_lines_ratio), + ("max_ellipsis_lines_ratio", self.max_ellipsis_lines_ratio), + ("max_non_alpha_words_ratio", self.max_non_alpha_words_ratio), + ]; + for (name, val) in ratio_params.iter() { + if let Some(v) = val { + if *v < 0.0 { + return Err(PipelineError::ConfigValidationError(format!( + "GopherQualityParams: {} must be non-negative, got {}", + name, v + ))); + } + } + } + + // This test is pointless... + // if let Some(min_stop_words) = self.min_stop_words { + // // min_stop_words can be 0, so no check for > 0 needed here. + // // This is valid as per the original description "greater than or equal to 0". + // if min_stop_words < 0 { + // return Err(PipelineError::ConfigValidationError(format!( + // "GopherQualityParams: min_stop_words must be non-negative, got {}", + // min_stop_words + // ))); + // } + // } + + Ok(()) + } +} + +#[derive(Deserialize, Debug, Clone)] +pub struct C4BadWordsParams { + pub keep_fraction: f32, + pub fail_on_missing_language: bool, + pub seed: Option, + pub default_language: String, + #[serde(skip)] // This field will not be deserialized from YAML + pub cache_base_path: Option, +} + +impl C4BadWordsParams { + pub fn validate(&self) -> Result<()> { + if !(0.0..=1.0).contains(&self.keep_fraction) { + return Err(PipelineError::ConfigValidationError(format!( + "C4BadWordsParams: keep_fraction must be between 0.0 and 1.0, got {}", + self.keep_fraction + ))); + } + if self.default_language.is_empty() { + return Err(PipelineError::ConfigValidationError( + "C4BadWordsParams: default_language cannot be empty".to_string(), + )); + } + Ok(()) + } +} + +// Parameters for the LangaugeDetectionFilter +#[derive(Deserialize, Debug, Clone)] +pub struct LanguageDetectionParams { + pub min_confidence: f64, + pub allowed_languages: Vec, +} + +impl LanguageDetectionParams { + pub fn validate(&self) -> Result<()> { + if !(0.0..=1.0).contains(&self.min_confidence) { + return Err(PipelineError::ConfigValidationError(format!( + "LanguageDetectionParams: min_confidence must be between 0.0 and 1.0, got {}", + self.min_confidence + ))); + } + if self.allowed_languages.is_empty() { + return Err(PipelineError::ConfigValidationError( + "LanguageDetectionParams: allowed_languages cannot be empty".to_string(), + )); + } + Ok(()) + } +} + +// Parameters for the FineWebQualityFilter (new filter based on Python logic). +#[derive(Deserialize, Debug, Clone, Default)] // Added Default for easier construction in worker +pub struct FineWebQualityFilterParams { + pub line_punct_thr: f64, + pub line_punct_exclude_zero: bool, + pub stop_chars: Option>, // Will be converted to HashSet in setup + pub short_line_thr: f64, + pub short_line_length: usize, // serde will handle u64 -> usize if value fits + pub char_duplicates_ratio: f64, + pub new_line_ratio: f64, + // pub language: String, +} + +impl FineWebQualityFilterParams { + pub fn validate(&self) -> Result<()> { + let params_to_check = [ + ("line_punct_thr", self.line_punct_thr), + ("short_line_thr", self.short_line_thr), + ("char_duplicates_ratio", self.char_duplicates_ratio), + ("new_line_ratio", self.new_line_ratio), + ]; + + for (name, value) in params_to_check.iter() { + if !(0.0..=1.0).contains(value) { + return Err(PipelineError::ConfigValidationError(format!( + "FineWebQualityFilterParams: {} must be between 0.0 and 1.0, got {}", + name, value + ))); + } + } + + if self.short_line_length == 0 { + return Err(PipelineError::ConfigValidationError( + "FineWebQualityFilterParams: short_line_length must be greater than 0".to_string(), + )); + } + Ok(()) + } +} + +// Parameters for the TokenCounter +#[derive(Deserialize, Debug, Clone)] +pub struct TokenCounterParams { + pub tokenizer_name: String, +} + +impl TokenCounterParams { + pub fn validate(&self) -> Result<()> { + // Add specific validation logic for TokenCounterParams if needed in the future + // For example, check if tokenizer_name is not empty or refers to a known tokenizer + if self.tokenizer_name.is_empty() { + return Err(PipelineError::ConfigValidationError( + "TokenCounterParams: tokenizer_name cannot be empty".to_string(), + )); + } + Ok(()) + } +} + +// {{ Add the new function to load pipeline configuration }} +/// Loads and parses the pipeline configuration YAML file. +pub fn load_pipeline_config>(config_path: P) -> Result { + let path_ref = config_path.as_ref(); + let config_content = fs::read_to_string(path_ref).map_err(|e| { + PipelineError::ConfigError(format!( + "Failed to read pipeline config file '{}': {}", + path_ref.display(), + e + )) + })?; + + let config: PipelineConfig = serde_yaml::from_str(&config_content).map_err(|e| { + PipelineError::ConfigError(format!( + "Failed to parse pipeline config YAML from '{}': {}", + path_ref.display(), + e + )) + })?; + + config.validate()?; // Validate the loaded configuration + + Ok(config) +} diff --git a/src/config/producer.rs b/src/config/producer.rs new file mode 100644 index 0000000..f4b69b7 --- /dev/null +++ b/src/config/producer.rs @@ -0,0 +1,43 @@ +// --- Command-Line Arguments Struct --- +// Moved from src/bin/producer.rs to make it accessible by library tests and other binaries if needed. +use clap::Parser; // Add clap::Parser import + +#[derive(Parser, Debug, Clone)] // Added Clone for testability if needed +#[command(author, version, about, long_about = None)] +pub struct Args { + /// Path to the input Parquet file + #[arg(short, long)] + pub input_file: String, + + /// Text column name in the Parquet file + #[arg(long, default_value = "text")] + pub text_column: String, + + /// Optional ID column name in the Parquet file + #[arg(long)] + pub id_column: Option, + + /// RabbitMQ connection string (e.g., amqp://guest:guest@localhost:5672/%2f) + #[arg(short, long, default_value = "amqp://guest:guest@localhost:5672/%2f")] + pub amqp_addr: String, + + /// Name of the queue to publish tasks to + #[arg(short = 'q', long, default_value = "task_queue")] + pub task_queue: String, + + /// Name of the queue to consume results/outcomes from + #[arg(short = 'r', long, default_value = "results_queue")] + pub results_queue: String, + + /// Path to the output Parquet file + #[arg(short = 'o', long, default_value = "output_processed.parquet")] + pub output_file: String, + + /// Path to the excluded output Parquet file + #[arg(short = 'e', long, default_value = "excluded.parquet")] + pub excluded_file: String, + + /// Optional: Port for the Prometheus metrics HTTP endpoint + #[arg(long)] + pub metrics_port: Option, +} diff --git a/src/lib.rs b/src/lib.rs index 4a54396..e0c553a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,3 +14,9 @@ pub mod utils; // pub use error::{Result, PipelineError}; // pub use data_model::TextDocument; // pub use executor::{PipelineExecutor, ProcessingStep}; + +pub mod producer_logic; // Declare the new module + +// The AmqpConnectionManager trait and its implementation have been removed from here +// as per the new strategy focusing on TaskPublisherChannel defined in producer_logic.rs. +// This should resolve the persistent "expected trait, found struct lapin::Channel" error in this file. diff --git a/src/pipeline/filters/c4_filters.rs b/src/pipeline/filters/c4_filters.rs index e1ee629..69385e2 100644 --- a/src/pipeline/filters/c4_filters.rs +++ b/src/pipeline/filters/c4_filters.rs @@ -1,4 +1,4 @@ -use crate::config::C4BadWordsParams; +use crate::config::pipeline::C4BadWordsParams; use crate::data_model::TextDocument; use crate::error::{PipelineError, Result}; use crate::executor::ProcessingStep; @@ -533,7 +533,7 @@ impl ProcessingStep for C4BadWordsFilter { #[cfg(test)] mod tests { use super::*; - use crate::config::C4BadWordsParams; + use crate::config::pipeline::C4BadWordsParams; use crate::data_model::TextDocument; use crate::error::PipelineError; use std::collections::HashMap; diff --git a/src/pipeline/readers/parquet_reader.rs b/src/pipeline/readers/parquet_reader.rs index 7303ccf..cd760d7 100644 --- a/src/pipeline/readers/parquet_reader.rs +++ b/src/pipeline/readers/parquet_reader.rs @@ -1,6 +1,6 @@ // src/readers/parquet_reader.rs -use crate::config::ParquetInputConfig; +use crate::config::parquet::ParquetInputConfig; use crate::data_model::TextDocument; use crate::error::{PipelineError, Result}; // Use the crate's error types diff --git a/src/producer_logic.rs b/src/producer_logic.rs new file mode 100644 index 0000000..d804068 --- /dev/null +++ b/src/producer_logic.rs @@ -0,0 +1,360 @@ +// src/producer_logic.rs +use crate::config::parquet::ParquetInputConfig; +use crate::config::producer::Args; +use crate::data_model::ProcessingOutcome; // Still needed for aggregate_results +use crate::data_model::TextDocument; // TextDocument is needed for aggregate_results +use crate::error::{PipelineError, Result as AppResult}; +use crate::pipeline::readers::ParquetReader; +use crate::pipeline::writers::parquet_writer::ParquetWriter; +use crate::utils::prometheus_metrics::*; +use async_trait::async_trait; +use chrono::Utc; +use futures::{pin_mut, Stream, StreamExt}; +use indicatif::ProgressBar; +use lapin::{ + options::{ + BasicAckOptions, BasicConsumeOptions, BasicPublishOptions, ConfirmSelectOptions, + QueueDeclareOptions, + }, + protocol::basic::AMQPProperties, + publisher_confirm::Confirmation, + types::FieldTable, + Channel as LapinChannel, // Alias lapin::Channel to avoid confusion + Consumer, + Result as LapinResult, +}; +use serde_json; +use std::time::Instant; +use tracing::{error, info, info_span, warn}; // For aggregate_results consumer tag + +pub const PARQUET_WRITE_BATCH_SIZE: usize = 500; + +#[async_trait] +pub trait TaskPublisherChannel: Send + Sync { + async fn queue_declare( + &self, + name: &str, + options: QueueDeclareOptions, + arguments: FieldTable, + ) -> LapinResult<()>; + async fn basic_publish( + &self, + exchange: &str, + routing_key: &str, + options: BasicPublishOptions, + payload: &[u8], + properties: AMQPProperties, + ) -> LapinResult; + async fn confirm_select(&self, options: ConfirmSelectOptions) -> LapinResult<()>; +} + +#[async_trait] +impl TaskPublisherChannel for LapinChannel { + async fn queue_declare( + &self, + name: &str, + options: QueueDeclareOptions, + arguments: FieldTable, + ) -> LapinResult<()> { + LapinChannel::queue_declare(self, name, options, arguments).await?; + Ok(()) + } + async fn basic_publish( + &self, + exchange: &str, + routing_key: &str, + options: BasicPublishOptions, + payload: &[u8], + properties: AMQPProperties, + ) -> LapinResult { + let publisher_confirmation = + LapinChannel::basic_publish(self, exchange, routing_key, options, payload, properties) + .await?; + publisher_confirmation.await + } + async fn confirm_select(&self, options: ConfirmSelectOptions) -> LapinResult<()> { + LapinChannel::confirm_select(self, options).await + } +} + +pub async fn publish_tasks( + args: &Args, + publish_channel: &CH, + publishing_pb: &ProgressBar, +) -> AppResult { + publish_channel + .queue_declare( + &args.task_queue, + QueueDeclareOptions { + durable: true, + ..Default::default() + }, + FieldTable::default(), + ) + .await?; + info!("Declared durable task queue '{}'", args.task_queue); + + let parquet_config = ParquetInputConfig { + path: args.input_file.clone(), + text_column: args.text_column.clone(), + id_column: args.id_column.clone(), + batch_size: Some(1024), + }; + let reader = ParquetReader::new(parquet_config); + + info!("Reading documents and publishing tasks..."); + let mut published_count = 0u64; + let mut read_errors = 0u64; + let doc_iterator = reader.read_documents()?; + + let publish_start_time = Instant::now(); + + for doc_result in doc_iterator { + publishing_pb.tick(); + match doc_result { + Ok(doc) => { + let _doc_span = info_span!("publishing_doc", doc_id = %doc.id).entered(); + match serde_json::to_vec(&doc) { + Ok(payload) => { + info!("Serialized document for publishing."); + let task_publish_timer = TASK_PUBLISHING_DURATION_SECONDS.start_timer(); + let confirmation = publish_channel + .basic_publish( + "", + &args.task_queue, + BasicPublishOptions::default(), + &payload, + AMQPProperties::default().with_delivery_mode(2), + ) + .await?; + task_publish_timer.observe_duration(); + + match confirmation { + Confirmation::Ack(_) | Confirmation::NotRequested => { + published_count += 1; + TASKS_PUBLISHED_TOTAL.inc(); + ACTIVE_TASKS_IN_FLIGHT.inc(); + publishing_pb.inc(1); + if matches!(confirmation, Confirmation::Ack(_)) { + info!("Successfully published task and received ACK."); + } else { + info!("Successfully published task (no confirmation requested/received)."); + } + } + Confirmation::Nack(_) => { + TASK_PUBLISH_ERRORS_TOTAL.inc(); + error!(doc_id = %doc.id, "FATAL: Broker NACKed task. Stopping."); + return Err(PipelineError::QueueError(format!( + "Publish confirmation failed (NACK) for doc {}", + doc.id + ))); + } + } + } + Err(e) => { + TASK_PUBLISH_ERRORS_TOTAL.inc(); + warn!(doc_id = %doc.id, error = %e, "Failed to serialize task. Skipping."); + read_errors += 1; + } + } + } + Err(e) => { + warn!(error = %e, "Error reading document for task. Skipping."); + read_errors += 1; + } + } + } + let publishing_duration = publish_start_time.elapsed(); + publishing_pb.finish_with_message(format!( + "Finished publishing {} tasks in {}. Read/Serialization Errors: {}", + published_count, + indicatif::HumanDuration(publishing_duration), + read_errors + )); + Ok(published_count) +} + +#[async_trait] +pub trait ResultConsumerChannel: Send + Sync { + async fn queue_declare( + &self, + name: &str, + options: QueueDeclareOptions, + arguments: FieldTable, + ) -> LapinResult<()>; + async fn basic_consume( + &self, + queue: &str, + consumer_tag: &str, + options: BasicConsumeOptions, + arguments: FieldTable, + ) -> LapinResult; +} + +#[async_trait] +impl ResultConsumerChannel for LapinChannel { + async fn queue_declare( + &self, + name: &str, + options: QueueDeclareOptions, + arguments: FieldTable, + ) -> LapinResult<()> { + let _ = LapinChannel::queue_declare(self, name, options, arguments).await; + Ok(()) + } + async fn basic_consume( + &self, + queue: &str, + consumer_tag: &str, + options: BasicConsumeOptions, + arguments: FieldTable, + ) -> LapinResult { + LapinChannel::basic_consume(self, queue, consumer_tag, options, arguments).await + } +} + +pub async fn aggregate_results_from_stream( + args: &Args, + stream: S, + published_count: u64, + aggregation_pb: &ProgressBar, +) -> AppResult<(u64, u64, u64)> +where + S: Stream, +{ + pin_mut!(stream); + if let Some(parent_dir) = std::path::Path::new(&args.output_file).parent() { + tokio::fs::create_dir_all(parent_dir) + .await + .map_err(|e| PipelineError::IoError { source: e })?; + } + if let Some(parent_dir) = std::path::Path::new(&args.excluded_file).parent() { + tokio::fs::create_dir_all(parent_dir) + .await + .map_err(|e| PipelineError::IoError { source: e })?; + } + + let mut parquet_writer_output = ParquetWriter::new(&args.output_file)?; + let mut parquet_writer_excluded = ParquetWriter::new(&args.excluded_file)?; + + let mut results_batch: Vec = Vec::with_capacity(PARQUET_WRITE_BATCH_SIZE); + let mut excluded_batch: Vec = Vec::with_capacity(PARQUET_WRITE_BATCH_SIZE); + let mut outcomes_received_count = 0u64; + let mut success_count = 0u64; + let mut filtered_count = 0u64; + + let aggregation_start_time = Instant::now(); + + while outcomes_received_count < published_count { + aggregation_pb.tick(); + if let Some(outcome) = stream.next().await { + outcomes_received_count += 1; + aggregation_pb.inc(1); + ACTIVE_TASKS_IN_FLIGHT.dec(); + + match outcome { + ProcessingOutcome::Success(doc) => { + success_count += 1; + results_batch.push(doc); + if results_batch.len() >= PARQUET_WRITE_BATCH_SIZE { + parquet_writer_output.write_batch(&results_batch)?; + results_batch.clear(); + } + } + ProcessingOutcome::Filtered { + document, + reason: _, + } => { + filtered_count += 1; + excluded_batch.push(document); + if excluded_batch.len() >= PARQUET_WRITE_BATCH_SIZE { + parquet_writer_excluded.write_batch(&excluded_batch)?; + excluded_batch.clear(); + } + } + ProcessingOutcome::Error { .. } => { + // Count/metrics can go here if needed + } + } + } else { + warn!("Outcome stream closed before all outcomes received."); + break; + } + } + + aggregation_pb.finish_with_message(format!( + "Finished consuming (Received {}/{}) in {}.", + outcomes_received_count, + published_count, + indicatif::HumanDuration(aggregation_start_time.elapsed()) + )); + + if !results_batch.is_empty() { + parquet_writer_output.write_batch(&results_batch)?; + } + if !excluded_batch.is_empty() { + parquet_writer_excluded.write_batch(&excluded_batch)?; + } + + parquet_writer_output.close()?; + parquet_writer_excluded.close()?; + + Ok((outcomes_received_count, success_count, filtered_count)) +} + +pub async fn aggregate_results( + args: &Args, + consume_channel: &CH, + published_count: u64, + aggregation_pb: &ProgressBar, +) -> AppResult<(u64, u64, u64)> { + info!("Starting results aggregation phase..."); + + consume_channel + .queue_declare( + &args.results_queue, + QueueDeclareOptions { + durable: true, + ..Default::default() + }, + FieldTable::default(), + ) + .await?; + + let consumer_tag = format!( + "producer-aggregator-{}-{}", + std::process::id(), + Utc::now().timestamp() + ); + + let consumer = consume_channel + .basic_consume( + &args.results_queue, + &consumer_tag, + BasicConsumeOptions::default(), + FieldTable::default(), + ) + .await?; + + // Map deliveries into ProcessingOutcome values + let mapped_stream = consumer.filter_map(|delivery_result| async { + match delivery_result { + Ok(delivery) => match serde_json::from_slice::(&delivery.data) { + Ok(outcome) => { + let _ = delivery.ack(BasicAckOptions::default()).await; + Some(outcome) + } + Err(err) => { + warn!(error = %err, "Failed to deserialize outcome."); + None + } + }, + Err(err) => { + error!(error = %err, "Failed to receive delivery."); + None + } + } + }); + + aggregate_results_from_stream(args, mapped_stream, published_count, aggregation_pb).await +} diff --git a/src/config.rs b/tests/config_tests.rs similarity index 56% rename from src/config.rs rename to tests/config_tests.rs index 00a07aa..f1c2f6e 100644 --- a/src/config.rs +++ b/tests/config_tests.rs @@ -1,414 +1,10 @@ -// src/config.rs -use crate::error::{PipelineError, Result}; -use serde::Deserialize; -use std::collections::HashSet; -use std::fs; // For reading the file -use std::path::Path; // For path handling // Assuming these are your error types - -// Keep existing config -#[derive(Deserialize, Debug, Clone)] -pub struct ParquetInputConfig { - pub path: String, // Path to the Parquet file or directory - pub text_column: String, // Name of the column containing the main text - pub id_column: Option, // Optional: Name of a column to use as document ID - // Add other column mappings as needed (e.g., for metadata) - pub batch_size: Option, // Optional: Arrow batch size for reading -} - -// --- Pipeline Configuration --- - -/// Represents the overall pipeline configuration read from YAML. -#[derive(Deserialize, Debug, Clone)] -pub struct PipelineConfig { - pub pipeline: Vec, -} - -impl PipelineConfig { - pub fn validate(&self) -> Result<()> { - for step_config in &self.pipeline { - step_config.validate()?; - } - Ok(()) - } -} - -/// Represents a single step in the processing pipeline. -/// Uses Serde's externally tagged enum representation. -#[derive(Deserialize, Debug, Clone)] -#[serde(tag = "type")] // The 'type' field in YAML determines which variant -pub enum StepConfig { - C4QualityFilter(C4QualityParams), - GopherRepetitionFilter(GopherRepetitionParams), - GopherQualityFilter(GopherQualityParams), - C4BadWordsFilter(C4BadWordsParams), // New - LanguageDetectionFilter(LanguageDetectionParams), - FineWebQualityFilter(FineWebQualityFilterParams), // Renamed and new params struct - TokenCounter(TokenCounterParams), - // Add other filter/step types here as needed -} - -impl StepConfig { - /// Returns a string slice representing the name of the step type. - pub fn name(&self) -> &'static str { - match self { - StepConfig::C4QualityFilter(_) => "C4QualityFilter", - StepConfig::GopherRepetitionFilter(_) => "GopherRepetitionFilter", - StepConfig::GopherQualityFilter(_) => "GopherQualityFilter", - StepConfig::C4BadWordsFilter(_) => "C4BadWordsFilter", // New - StepConfig::LanguageDetectionFilter(_) => "LanguageDetectionFilter", - StepConfig::FineWebQualityFilter(_) => "FineWebQualityFilter", // Renamed - StepConfig::TokenCounter(_) => "TokenCounter", // Add cases for other StepConfig variants here - } - } - - pub fn validate(&self) -> Result<()> { - match self { - StepConfig::C4QualityFilter(params) => params.validate(), - StepConfig::GopherRepetitionFilter(params) => params.validate(), - StepConfig::GopherQualityFilter(params) => params.validate(), - StepConfig::C4BadWordsFilter(params) => params.validate(), - StepConfig::LanguageDetectionFilter(params) => params.validate(), - StepConfig::FineWebQualityFilter(params) => params.validate(), - StepConfig::TokenCounter(params) => params.validate(), - } - } -} - -/// Parameters for the C4QualityFilter. -#[derive(Deserialize, Debug, Clone)] -pub struct C4QualityParams { - pub split_paragraph: bool, - pub remove_citations: bool, - pub filter_no_terminal_punct: bool, - pub min_num_sentences: usize, - pub min_words_per_line: usize, - pub max_word_length: usize, - pub filter_lorem_ipsum: bool, - pub filter_javascript: bool, - pub filter_curly_bracket: bool, - pub filter_policy: bool, -} - -impl C4QualityParams { - pub fn validate(&self) -> Result<()> { - if self.min_num_sentences == 0 { - return Err(PipelineError::ConfigValidationError( - "C4QualityParams: min_num_sentences must be greater than 0".to_string(), - )); - } - if self.min_words_per_line == 0 { - return Err(PipelineError::ConfigValidationError( - "C4QualityParams: min_words_per_line must be greater than 0".to_string(), - )); - } - if self.max_word_length == 0 { - return Err(PipelineError::ConfigValidationError( - "C4QualityParams: max_word_length must be greater than 0".to_string(), - )); - } - Ok(()) - } -} - -/// Parameters for the GopherRepetitionFilter. -#[derive(Deserialize, Debug, Clone)] -pub struct GopherRepetitionParams { - // Use Option for optional fields - pub dup_line_frac: Option, - pub dup_para_frac: Option, - pub dup_line_char_frac: Option, - pub dup_para_char_frac: Option, - // Vec<(usize, f64)> represents lists like [[2, 0.2], [3, 0.18]] in YAML - #[serde(default)] // Use default (empty vec) if not specified - pub top_n_grams: Vec<(usize, f64)>, - #[serde(default)] - pub dup_n_grams: Vec<(usize, f64)>, -} - -impl GopherRepetitionParams { - pub fn validate(&self) -> Result<()> { - let fractions = [ - ("dup_line_frac", self.dup_line_frac), - ("dup_para_frac", self.dup_para_frac), - ("dup_line_char_frac", self.dup_line_char_frac), - ("dup_para_char_frac", self.dup_para_char_frac), - ]; - for (name, val) in fractions.iter() { - if let Some(v) = val { - if !(0.0..=1.0).contains(v) { - return Err(PipelineError::ConfigValidationError(format!( - "GopherRepetitionParams: {} must be between 0.0 and 1.0, got {}", - name, v - ))); - } - } - } - - for (name, n_grams) in [ - ("top_n_grams", &self.top_n_grams), - ("dup_n_grams", &self.dup_n_grams), - ] - .iter() - { - for (idx, (size, fraction)) in n_grams.iter().enumerate() { - if *size == 0 { - return Err(PipelineError::ConfigValidationError(format!( - "GopherRepetitionParams: n-gram size in {} at index {} must be greater than 0", - name, idx - ))); - } - if !(0.0..=1.0).contains(fraction) { - return Err(PipelineError::ConfigValidationError(format!( - "GopherRepetitionParams: n-gram fraction in {} at index {} must be between 0.0 and 1.0, got {}", - name, idx, fraction - ))); - } - } - } - Ok(()) - } -} - -/// Parameters for the GopherQualityFilter. -#[derive(Deserialize, Debug, Clone)] -pub struct GopherQualityParams { - // Use Option for optional fields - pub min_doc_words: Option, - pub max_doc_words: Option, - pub min_avg_word_length: Option, - pub max_avg_word_length: Option, - pub max_symbol_word_ratio: Option, - pub max_bullet_lines_ratio: Option, - pub max_ellipsis_lines_ratio: Option, - pub max_non_alpha_words_ratio: Option, - pub min_stop_words: Option, - // Optional list of stop words; if None, the filter's default will be used. - pub stop_words: Option>, -} - -impl GopherQualityParams { - pub fn validate(&self) -> Result<()> { - if let Some(min_doc_words) = self.min_doc_words { - if min_doc_words == 0 { - return Err(PipelineError::ConfigValidationError( - "GopherQualityParams: min_doc_words must be greater than 0".to_string(), - )); - } - } - if let Some(max_doc_words) = self.max_doc_words { - if max_doc_words == 0 { - return Err(PipelineError::ConfigValidationError( - "GopherQualityParams: max_doc_words must be greater than 0".to_string(), - )); - } - } - if let (Some(min_val), Some(max_val)) = (self.min_doc_words, self.max_doc_words) { - if min_val > max_val { - return Err(PipelineError::ConfigValidationError(format!( - "GopherQualityParams: min_doc_words ({}) cannot be greater than max_doc_words ({})", - min_val, max_val - ))); - } - } - - if let Some(min_avg_word_length) = self.min_avg_word_length { - if min_avg_word_length <= 0.0 { - return Err(PipelineError::ConfigValidationError( - "GopherQualityParams: min_avg_word_length must be greater than 0.0".to_string(), - )); - } - } - if let Some(max_avg_word_length) = self.max_avg_word_length { - if max_avg_word_length <= 0.0 { - return Err(PipelineError::ConfigValidationError( - "GopherQualityParams: max_avg_word_length must be greater than 0.0".to_string(), - )); - } - } - if let (Some(min_val), Some(max_val)) = (self.min_avg_word_length, self.max_avg_word_length) - { - if min_val > max_val { - return Err(PipelineError::ConfigValidationError(format!( - "GopherQualityParams: min_avg_word_length ({}) cannot be greater than max_avg_word_length ({})", - min_val, max_val - ))); - } - } - - let ratio_params = [ - ("max_symbol_word_ratio", self.max_symbol_word_ratio), - ("max_bullet_lines_ratio", self.max_bullet_lines_ratio), - ("max_ellipsis_lines_ratio", self.max_ellipsis_lines_ratio), - ("max_non_alpha_words_ratio", self.max_non_alpha_words_ratio), - ]; - for (name, val) in ratio_params.iter() { - if let Some(v) = val { - if *v < 0.0 { - return Err(PipelineError::ConfigValidationError(format!( - "GopherQualityParams: {} must be non-negative, got {}", - name, v - ))); - } - } - } - - // This test is pointless... - // if let Some(min_stop_words) = self.min_stop_words { - // // min_stop_words can be 0, so no check for > 0 needed here. - // // This is valid as per the original description "greater than or equal to 0". - // if min_stop_words < 0 { - // return Err(PipelineError::ConfigValidationError(format!( - // "GopherQualityParams: min_stop_words must be non-negative, got {}", - // min_stop_words - // ))); - // } - // } - - Ok(()) - } -} - -#[derive(Deserialize, Debug, Clone)] -pub struct C4BadWordsParams { - pub keep_fraction: f32, - pub fail_on_missing_language: bool, - pub seed: Option, - pub default_language: String, - #[serde(skip)] // This field will not be deserialized from YAML - pub cache_base_path: Option, -} - -impl C4BadWordsParams { - pub fn validate(&self) -> Result<()> { - if !(0.0..=1.0).contains(&self.keep_fraction) { - return Err(PipelineError::ConfigValidationError(format!( - "C4BadWordsParams: keep_fraction must be between 0.0 and 1.0, got {}", - self.keep_fraction - ))); - } - if self.default_language.is_empty() { - return Err(PipelineError::ConfigValidationError( - "C4BadWordsParams: default_language cannot be empty".to_string(), - )); - } - Ok(()) - } -} - -// Parameters for the LangaugeDetectionFilter -#[derive(Deserialize, Debug, Clone)] -pub struct LanguageDetectionParams { - pub min_confidence: f64, - pub allowed_languages: Vec, -} - -impl LanguageDetectionParams { - pub fn validate(&self) -> Result<()> { - if !(0.0..=1.0).contains(&self.min_confidence) { - return Err(PipelineError::ConfigValidationError(format!( - "LanguageDetectionParams: min_confidence must be between 0.0 and 1.0, got {}", - self.min_confidence - ))); - } - if self.allowed_languages.is_empty() { - return Err(PipelineError::ConfigValidationError( - "LanguageDetectionParams: allowed_languages cannot be empty".to_string(), - )); - } - Ok(()) - } -} - -// Parameters for the FineWebQualityFilter (new filter based on Python logic). -#[derive(Deserialize, Debug, Clone, Default)] // Added Default for easier construction in worker -pub struct FineWebQualityFilterParams { - pub line_punct_thr: f64, - pub line_punct_exclude_zero: bool, - pub stop_chars: Option>, // Will be converted to HashSet in setup - pub short_line_thr: f64, - pub short_line_length: usize, // serde will handle u64 -> usize if value fits - pub char_duplicates_ratio: f64, - pub new_line_ratio: f64, - // pub language: String, -} - -impl FineWebQualityFilterParams { - pub fn validate(&self) -> Result<()> { - let params_to_check = [ - ("line_punct_thr", self.line_punct_thr), - ("short_line_thr", self.short_line_thr), - ("char_duplicates_ratio", self.char_duplicates_ratio), - ("new_line_ratio", self.new_line_ratio), - ]; - - for (name, value) in params_to_check.iter() { - if !(0.0..=1.0).contains(value) { - return Err(PipelineError::ConfigValidationError(format!( - "FineWebQualityFilterParams: {} must be between 0.0 and 1.0, got {}", - name, value - ))); - } - } - - if self.short_line_length == 0 { - return Err(PipelineError::ConfigValidationError( - "FineWebQualityFilterParams: short_line_length must be greater than 0".to_string(), - )); - } - Ok(()) - } -} - -// Parameters for the TokenCounter -#[derive(Deserialize, Debug, Clone)] -pub struct TokenCounterParams { - pub tokenizer_name: String, -} - -impl TokenCounterParams { - pub fn validate(&self) -> Result<()> { - // Add specific validation logic for TokenCounterParams if needed in the future - // For example, check if tokenizer_name is not empty or refers to a known tokenizer - if self.tokenizer_name.is_empty() { - return Err(PipelineError::ConfigValidationError( - "TokenCounterParams: tokenizer_name cannot be empty".to_string(), - )); - } - Ok(()) - } -} - -// {{ Add the new function to load pipeline configuration }} -/// Loads and parses the pipeline configuration YAML file. -pub fn load_pipeline_config>(config_path: P) -> Result { - let path_ref = config_path.as_ref(); - let config_content = fs::read_to_string(path_ref).map_err(|e| { - PipelineError::ConfigError(format!( - "Failed to read pipeline config file '{}': {}", - path_ref.display(), - e - )) - })?; - - let config: PipelineConfig = serde_yaml::from_str(&config_content).map_err(|e| { - PipelineError::ConfigError(format!( - "Failed to parse pipeline config YAML from '{}': {}", - path_ref.display(), - e - )) - })?; - - config.validate()?; // Validate the loaded configuration - - Ok(config) -} - // {{ Add unit tests for load_pipeline_config }} #[cfg(test)] mod tests { - use super::*; use std::io::Write; use tempfile::NamedTempFile; + use TextBlaster::config::pipeline::*; + use TextBlaster::error::PipelineError; // Helper to create a temporary config file with given content fn create_temp_config_file(content: &str) -> NamedTempFile { diff --git a/tests/full_pipeline_test.rs b/tests/full_pipeline_test.rs index 7d2e60d..53931f4 100644 --- a/tests/full_pipeline_test.rs +++ b/tests/full_pipeline_test.rs @@ -12,7 +12,7 @@ use testcontainers::{ ContainerAsync, GenericImage, }; // Added AsyncRunner -use TextBlaster::config::ParquetInputConfig; +use TextBlaster::config::parquet::ParquetInputConfig; use TextBlaster::data_model::TextDocument; use TextBlaster::error::Result; // Assuming this is your crate's Result type use TextBlaster::pipeline::readers::parquet_reader::ParquetReader; diff --git a/tests/parquet_io_test.rs b/tests/parquet_io_test.rs index e66095e..63e7859 100644 --- a/tests/parquet_io_test.rs +++ b/tests/parquet_io_test.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use tempfile::NamedTempFile; // Import necessary items from your crate -use TextBlaster::config::ParquetInputConfig; // Assuming this is public and in src/config.rs +use TextBlaster::config::parquet::ParquetInputConfig; // Assuming this is public and in src/config.rs use TextBlaster::data_model::TextDocument; use TextBlaster::error::Result; use TextBlaster::pipeline::readers::parquet_reader::ParquetReader; diff --git a/tests/producer_tests.rs b/tests/producer_tests.rs new file mode 100644 index 0000000..767a51b --- /dev/null +++ b/tests/producer_tests.rs @@ -0,0 +1,660 @@ +#[cfg(test)] +mod args_tests { + use clap::Parser; + pub use TextBlaster::config::producer::Args; + + #[test] + fn test_parse_all_args() { + let args = Args::parse_from(&[ + "producer", + "-i", + "input.parquet", + "--text-column", + "content", + "--id-column", + "doc_id", + "-a", + "amqp://user:pass@host:port/vhost", + "-q", + "my_tasks", + "-r", + "my_results", + "-o", + "processed.parquet", + "-e", + "errors.parquet", + "--metrics-port", + "9090", + ]); + assert_eq!(args.input_file, "input.parquet"); + assert_eq!(args.text_column, "content"); + assert_eq!(args.id_column, Some("doc_id".to_string())); + assert_eq!(args.amqp_addr, "amqp://user:pass@host:port/vhost"); + assert_eq!(args.task_queue, "my_tasks"); + assert_eq!(args.results_queue, "my_results"); + assert_eq!(args.output_file, "processed.parquet"); + assert_eq!(args.excluded_file, "errors.parquet"); + assert_eq!(args.metrics_port, Some(9090)); + } + #[test] + fn test_parse_required_only() { + let args = Args::parse_from(&["producer", "-i", "input.parquet"]); + assert_eq!(args.input_file, "input.parquet"); + assert_eq!(args.text_column, "text"); + assert_eq!(args.id_column, None); + assert_eq!(args.amqp_addr, "amqp://guest:guest@localhost:5672/%2f"); + assert_eq!(args.task_queue, "task_queue"); + assert_eq!(args.results_queue, "results_queue"); + assert_eq!(args.output_file, "output_processed.parquet"); + assert_eq!(args.excluded_file, "excluded.parquet"); + assert_eq!(args.metrics_port, None); + } + #[test] + fn test_parse_with_optional_id_column() { + let args = Args::parse_from(&[ + "producer", + "-i", + "input.parquet", + "--id-column", + "custom_id", + ]); + assert_eq!(args.id_column, Some("custom_id".to_string())); + } + #[test] + fn test_parse_with_optional_metrics_port() { + let args = Args::parse_from(&["producer", "-i", "input.parquet", "--metrics-port", "8000"]); + assert_eq!(args.metrics_port, Some(8000)); + } + #[test] + fn test_default_values_are_applied() { + let args = Args::parse_from(&["producer", "--input-file", "input.parquet"]); + assert_eq!(args.text_column, "text"); + assert_eq!(args.amqp_addr, "amqp://guest:guest@localhost:5672/%2f"); + } + #[test] + fn test_missing_required_arg_error() { + let result = Args::try_parse_from(&["producer"]); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().kind(), + clap::error::ErrorKind::MissingRequiredArgument + ); + } + #[test] + fn test_invalid_metrics_port_format() { + let result = Args::try_parse_from(&[ + "producer", + "-i", + "input.parquet", + "--metrics-port", + "not_a_port", + ]); + assert!(result.is_err()); + assert_eq!( + result.unwrap_err().kind(), + clap::error::ErrorKind::ValueValidation + ); + } +} + +#[cfg(test)] +mod publish_task_tests { + use arrow::array::StringArray; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use async_trait::async_trait; + use indicatif::ProgressBar; + use lapin::options::{BasicPublishOptions, ConfirmSelectOptions, QueueDeclareOptions}; + use lapin::protocol::basic::AMQPProperties; + use lapin::publisher_confirm::Confirmation; + use lapin::types::FieldTable; + use lapin::Result as LapinResult; + use parquet::arrow::arrow_writer::ArrowWriter; + use parquet::file::properties::WriterProperties; + use std::collections::HashMap; + use std::fs::File; + use std::sync::{Arc, Mutex}; + use tempfile::NamedTempFile; + use TextBlaster::config::producer::Args; + use TextBlaster::data_model::TextDocument; + use TextBlaster::error::PipelineError; + use TextBlaster::producer_logic::*; + use TextBlaster::utils::prometheus_metrics::{ + TASKS_PUBLISHED_TOTAL, TASK_PUBLISH_ERRORS_TOTAL, + }; + + //=============== MOCK SETUP ===============// + + /// Defines the behavior of our mock publisher channel. + #[derive(Clone, Copy)] + enum MockBehavior { + /// Always return a successful ACK. + AlwaysAck, + /// Return a NACK on the first publish attempt. + NackOnFirstPublish, + /// Return a generic LapinError on publish. + FailOnPublish, + } + + /// A mock implementation of the TaskPublisherChannel trait. + /// It allows us to simulate RabbitMQ behavior without a real connection. + struct MockTaskPublisherChannel { + /// Shared state to inspect after the test runs. + state: Arc>, + } + + struct MockState { + /// Stores the payloads that were "published". + published_payloads: Vec>, + /// Controls how the mock responds to publish calls. + behavior: MockBehavior, + } + + impl MockTaskPublisherChannel { + fn new(behavior: MockBehavior) -> Self { + Self { + state: Arc::new(Mutex::new(MockState { + published_payloads: Vec::new(), + behavior, + })), + } + } + } + + #[async_trait] + impl TaskPublisherChannel for MockTaskPublisherChannel { + async fn queue_declare( + &self, + _name: &str, + _options: QueueDeclareOptions, + _arguments: FieldTable, + ) -> LapinResult<()> { + // Return a dummy queue. Its properties don't matter for this test. + Ok(()) + } + + async fn basic_publish( + &self, + _exchange: &str, + _routing_key: &str, + _options: BasicPublishOptions, + payload: &[u8], + _properties: AMQPProperties, + ) -> LapinResult { + let mut state = self.state.lock().unwrap(); + state.published_payloads.push(payload.to_vec()); + + match state.behavior { + MockBehavior::AlwaysAck => Ok(Confirmation::Ack(Default::default())), + MockBehavior::NackOnFirstPublish => Ok(Confirmation::Nack(Default::default())), + MockBehavior::FailOnPublish => { + // CORRECTED LINE: + // We must construct the full AMQPError with a code and text. + let amqp_error = lapin::protocol::AMQPError::new( + lapin::protocol::AMQPErrorKind::Hard( + lapin::protocol::AMQPHardError::INTERNALERROR, + ), + "mock failure".into(), + ); + Err(lapin::Error::ProtocolError(amqp_error)) + } + } + } + + async fn confirm_select(&self, _options: ConfirmSelectOptions) -> LapinResult<()> { + Ok(()) + } + } + + //=============== TEST HELPER FUNCTIONS ===============// + + /// Helper to create a temporary Parquet file with a specified number of documents. + /// Returns the temp file handle (to prevent deletion), the file path, and the original docs. + fn create_test_parquet_file(num_records: usize) -> (NamedTempFile, Vec) { + let temp_file = NamedTempFile::new().unwrap(); + let file_path = temp_file.path().to_str().unwrap().to_string(); + + let ids: Vec = (0..num_records).map(|i| format!("doc_{}", i)).collect(); + let texts: Vec = (0..num_records) + .map(|i| format!("This is text for doc {}.", i)) + .collect(); + + let original_docs: Vec = ids + .iter() + .zip(texts.iter()) + .map(|(id, text)| TextDocument { + id: id.clone(), + source: "test".to_string(), + content: text.clone(), + metadata: HashMap::new(), + }) + .collect(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Utf8, false), + Field::new("text", DataType::Utf8, false), + ])); + + let id_array = StringArray::from_iter_values(ids.iter()); + let text_array = StringArray::from_iter_values(texts.iter()); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(id_array), Arc::new(text_array)], + ) + .unwrap(); + + let file = File::create(&file_path).unwrap(); + let mut writer = + ArrowWriter::try_new(file, schema, Some(WriterProperties::builder().build())).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + (temp_file, original_docs) + } + + fn create_mock_args(input_path: String) -> Args { + Args { + input_file: input_path, + text_column: "text".to_string(), + id_column: Some("id".to_string()), + amqp_addr: "amqp://guest:guest@localhost:5672/%2f".to_string(), + task_queue: "test_task_queue".to_string(), + results_queue: "result_queue".to_string(), + output_file: "output".to_string(), + excluded_file: "excluded".to_string(), + metrics_port: Some(1234), + } + } + + //=============== TEST CASES ===============// + + #[tokio::test] + async fn test_publish_tasks_happy_path() { + // ARRANGE + let num_docs = 5; + let (_temp_file, original_docs) = create_test_parquet_file(num_docs); + let args = create_mock_args(_temp_file.path().to_str().unwrap().to_string()); + let mock_channel = MockTaskPublisherChannel::new(MockBehavior::AlwaysAck); + let pb = ProgressBar::hidden(); + + // Reset metrics for a clean slate + TASKS_PUBLISHED_TOTAL.reset(); + + // ACT + let result = publish_tasks(&args, &mock_channel, &pb).await; + + // ASSERT + assert!(result.is_ok(), "Function should succeed"); + assert_eq!( + result.unwrap(), + num_docs as u64, + "Should report all documents as published" + ); + + // Check metrics + assert_eq!( + TASKS_PUBLISHED_TOTAL.get(), + num_docs as f64, + "Prometheus metric for published tasks should be correct" + ); + + // Check mock state + let state = mock_channel.state.lock().unwrap(); + assert_eq!( + state.published_payloads.len(), + num_docs, + "Exactly 5 messages should have been published" + ); + + // Verify content of a published message + let first_payload = &state.published_payloads[0]; + let deserialized_doc: TextDocument = serde_json::from_slice(first_payload).unwrap(); + assert_eq!( + deserialized_doc.content, original_docs[0].content, + "The content of the published message should match the source document" + ); + } + + #[tokio::test] + async fn test_publish_tasks_stops_on_nack() { + // ARRANGE + let (_temp_file, _) = create_test_parquet_file(5); + let args = create_mock_args(_temp_file.path().to_str().unwrap().to_string()); + let mock_channel = MockTaskPublisherChannel::new(MockBehavior::NackOnFirstPublish); + let pb = ProgressBar::hidden(); + + // Reset metrics + TASK_PUBLISH_ERRORS_TOTAL.reset(); + + // ACT + let result = publish_tasks(&args, &mock_channel, &pb).await; + + // ASSERT + assert!(result.is_err(), "Function should fail on NACK"); + let err = result.unwrap_err(); + assert!( + matches!(err, PipelineError::QueueError(_)), + "Error should be of type QueueError" + ); + assert!( + err.to_string() + .contains("Publish confirmation failed (NACK)"), + "Error message should indicate a NACK" + ); + + // Check metrics + assert_eq!( + TASK_PUBLISH_ERRORS_TOTAL.get(), + 1.0, + "Prometheus metric for publish errors should be incremented" + ); + + // Check mock state: The message was still sent before the NACK was received. + let state = mock_channel.state.lock().unwrap(); + assert_eq!( + state.published_payloads.len(), + 1, + "Only one message should have been attempted before stopping" + ); + } + + #[tokio::test] + async fn test_publish_tasks_propagates_lapin_error() { + // ARRANGE + let (_temp_file, _) = create_test_parquet_file(5); + let args = create_mock_args(_temp_file.path().to_str().unwrap().to_string()); + let mock_channel = MockTaskPublisherChannel::new(MockBehavior::FailOnPublish); + let pb = ProgressBar::hidden(); + + // ACT + let result = publish_tasks(&args, &mock_channel, &pb).await; + + // ASSERT + assert!( + result.is_err(), + "Function should fail if basic_publish returns an error" + ); + let err = result.unwrap_err(); + assert!( + matches!(err, PipelineError::QueueError(_)), + "Error should be a wrapped QueueError" + ); + } + + #[tokio::test] + async fn test_publish_tasks_handles_nonexistent_input_file() { + // ARRANGE + let args = create_mock_args("does_not_exist".to_string()); + // The mock won't even be used, as the failure happens before publishing. + let mock_channel = MockTaskPublisherChannel::new(MockBehavior::AlwaysAck); + let pb = ProgressBar::hidden(); + + // ACT + // Note: The error here is synchronous, as it happens during ParquetReader setup, + // but publish_tasks wraps it in the AppResult. + let result = publish_tasks(&args, &mock_channel, &pb).await; + + // ASSERT + assert!( + result.is_err(), + "Function should fail if input file doesn't exist" + ); + let err = result.unwrap_err(); + assert!( + matches!(err, PipelineError::IoError { source: _ }), + "{}", + format!("Error should be of type ParquetError not {}", err) + ); + } +} + +#[cfg(test)] +mod aggregate_results_tests { + use futures::stream; + use indicatif::ProgressBar; + use std::collections::HashMap; + use tempfile::NamedTempFile; + use TextBlaster::config::producer::Args; + use TextBlaster::data_model::{ProcessingOutcome, TextDocument}; + use TextBlaster::producer_logic::*; + + fn create_mock_args(output_path: String, excluded_path: String) -> Args { + Args { + input_file: "input".to_string(), + text_column: "text".to_string(), + id_column: Some("id".to_string()), + amqp_addr: "amqp://guest:guest@localhost:5672/%2f".to_string(), + task_queue: "test_task_queue".to_string(), + results_queue: "result_queue".to_string(), + output_file: output_path, + excluded_file: excluded_path, + metrics_port: Some(1234), + } + } + + fn setup_args() -> (Args, NamedTempFile, NamedTempFile) { + let output_file = NamedTempFile::new().unwrap(); + let excluded_file = NamedTempFile::new().unwrap(); + + let args = create_mock_args( + output_file.path().to_str().unwrap().to_string(), + excluded_file.path().to_str().unwrap().to_string(), + ); + + (args, output_file, excluded_file) + } + + fn sample_document(id: String) -> TextDocument { + TextDocument { + id, + source: "test".to_string(), + content: "exciting content".to_string(), + metadata: HashMap::new(), + } + } + + #[tokio::test] + async fn test_aggregate_results_success() { + let doc = TextDocument { + id: "test-id".to_string(), + content: "Test document".to_string(), + source: "test-src".to_string(), + metadata: HashMap::new(), + }; + + let outcome = ProcessingOutcome::Success(doc); + let stream = stream::iter(vec![outcome]); // no mocking Delivery or Lapin + + let tmp_output = tempfile::NamedTempFile::new().unwrap(); + let tmp_excluded = tempfile::NamedTempFile::new().unwrap(); + + let args = create_mock_args( + tmp_output.path().to_str().unwrap().to_string(), + tmp_excluded.path().to_str().unwrap().to_string(), + ); + + let pb = ProgressBar::hidden(); + + let (received, success, filtered) = aggregate_results_from_stream(&args, stream, 1, &pb) + .await + .unwrap(); + + assert_eq!(received, 1); + assert_eq!(success, 1); + assert_eq!(filtered, 0); + } + + #[tokio::test] + async fn test_aggregate_single_success() { + let doc = sample_document("success-1".to_string()); + let outcome = ProcessingOutcome::Success(doc.clone()); + let stream = stream::iter(vec![outcome]); + + let (args, _, _) = setup_args(); + let pb = ProgressBar::hidden(); + + let (received, success, filtered) = aggregate_results_from_stream(&args, stream, 1, &pb) + .await + .unwrap(); + + assert_eq!(received, 1); + assert_eq!(success, 1); + assert_eq!(filtered, 0); + + // Optionally verify file content here + let df = tokio::task::spawn_blocking(move || { + polars::prelude::LazyFrame::scan_parquet(args.output_file.clone(), Default::default()) + .unwrap() + .collect() + }) + .await + .unwrap() + .unwrap(); + + assert_eq!(df.shape().0, 1); + assert!(df + .column("id") + .unwrap() + .str() + .unwrap() + .get(0) + .unwrap() + .contains("success-1")); + } + + #[tokio::test] + async fn test_aggregate_single_filtered() { + let doc = sample_document("filtered-1".to_string()); + let outcome = ProcessingOutcome::Filtered { + document: doc.clone(), + reason: "Test filter".into(), + }; + let stream = stream::iter(vec![outcome]); + + let (args, _, _) = setup_args(); + let pb = ProgressBar::hidden(); + + let (received, success, filtered) = aggregate_results_from_stream(&args, stream, 1, &pb) + .await + .unwrap(); + + assert_eq!(received, 1); + assert_eq!(success, 0); + assert_eq!(filtered, 1); + + let df = tokio::task::spawn_blocking(move || { + polars::prelude::LazyFrame::scan_parquet(args.excluded_file.clone(), Default::default()) + .unwrap() + .collect() + }) + .await + .unwrap() + .unwrap(); + + assert_eq!(df.shape().0, 1); + assert!(df + .column("id") + .unwrap() + .str() + .unwrap() + .get(0) + .unwrap() + .contains("filtered-1")); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_aggregate_single_error() { + let doc = sample_document("error-1".to_string()); + let outcome = ProcessingOutcome::Error { + document: doc, + error_message: "Boom".into(), + worker_id: "w123".into(), + }; + let stream = stream::iter(vec![outcome]); + + let (args, _, _) = setup_args(); + let pb = ProgressBar::hidden(); + + let (received, success, filtered) = aggregate_results_from_stream(&args, stream, 1, &pb) + .await + .unwrap(); + + assert_eq!(received, 1); + assert_eq!(success, 0); + assert_eq!(filtered, 0); + + let out_rows = tokio::task::spawn_blocking(move || { + polars::prelude::LazyFrame::scan_parquet(&args.output_file, Default::default()) + .unwrap() + .collect() + }) + .await + .unwrap() + .unwrap() + .height(); + let excl_rows = tokio::task::spawn_blocking(move || { + polars::prelude::LazyFrame::scan_parquet(&args.excluded_file, Default::default()) + .unwrap() + .collect() + }) + .await + .unwrap() + .unwrap() + .height(); + assert_eq!(out_rows, 0); + assert_eq!(excl_rows, 0); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_aggregate_mixed_outcomes() { + let doc1 = sample_document("d1".to_string()); + let doc2 = sample_document("d2".to_string()); + let doc3 = sample_document("d3".to_string()); + + let stream = stream::iter(vec![ + ProcessingOutcome::Success(doc1.clone()), + ProcessingOutcome::Filtered { + document: doc2.clone(), + reason: "Too short".into(), + }, + ProcessingOutcome::Error { + document: doc3, + error_message: "Crash".into(), + worker_id: "w1".into(), + }, + ]); + + let (args, _, _) = setup_args(); + let pb = ProgressBar::hidden(); + + let (received, success, filtered) = aggregate_results_from_stream(&args, stream, 3, &pb) + .await + .unwrap(); + + assert_eq!(received, 3); + assert_eq!(success, 1); + assert_eq!(filtered, 1); + + let out_rows = tokio::task::spawn_blocking(move || { + polars::prelude::LazyFrame::scan_parquet(&args.output_file, Default::default()) + .unwrap() + .collect() + }) + .await + .unwrap() + .unwrap() + .height(); + let excl_rows = tokio::task::spawn_blocking(move || { + polars::prelude::LazyFrame::scan_parquet(&args.excluded_file, Default::default()) + .unwrap() + .collect() + }) + .await + .unwrap() + .unwrap() + .height(); + assert_eq!(out_rows, 1); + assert_eq!(excl_rows, 1); + } +}