diff --git a/Cargo.lock b/Cargo.lock index 26f07400..cb19233d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11,6 +11,12 @@ dependencies = [ "gimli", ] +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + [[package]] name = "ahash" version = "0.8.12" @@ -34,6 +40,21 @@ dependencies = [ "memchr", ] +[[package]] +name = "alloc-no-stdlib" +version = "2.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc7bb162ec39d46ab1ca8c77bf72e890535becd1751bb45f64c597edb4c8c6b3" + +[[package]] +name = "alloc-stdlib" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94fb8275041c72129eb51b7d0322c29b8387a0386127718b096429201a5d6ece" +dependencies = [ + "alloc-no-stdlib", +] + [[package]] name = "allocator-api2" version = "0.2.21" @@ -111,12 +132,68 @@ version = "1.0.100" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a23eb6b1614318a8071c9b2521f36b424b2c83db5eb3a0fead4a6c0809af6e61" +[[package]] +name = "ar_archive_writer" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7eb93bbb63b9c227414f6eb3a0adfddca591a8ce1e9b60661bb08969b87e340b" +dependencies = [ + "object", +] + [[package]] name = "arbitrary" version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +[[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 = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3f15b4c6b148206ff3a2b35002e08929c2462467b62b9c02036d9c34f9ef994" +dependencies = [ + "arrow-arith", + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-cast 55.2.0", + "arrow-csv", + "arrow-data 55.2.0", + "arrow-ipc 55.2.0", + "arrow-json 55.2.0 (registry+https://github.com/rust-lang/crates.io-index)", + "arrow-ord", + "arrow-row", + "arrow-schema 55.2.0", + "arrow-select 55.2.0", + "arrow-string", +] + +[[package]] +name = "arrow-arith" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30feb679425110209ae35c3fbf82404a39a4c0436bb3ec36164d8bffed2a4ce4" +dependencies = [ + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "chrono", + "num", +] + [[package]] name = "arrow-array" version = "52.2.0" @@ -124,15 +201,32 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "16f4a9468c882dc66862cef4e1fd8423d47e67972377d85d80e022786427768c" dependencies = [ "ahash", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-buffer 52.2.0", + "arrow-data 52.2.0", + "arrow-schema 52.2.0", "chrono", "half", "hashbrown 0.14.5", "num", ] +[[package]] +name = "arrow-array" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70732f04d285d49054a48b72c54f791bb3424abae92d27aafdf776c98af161c8" +dependencies = [ + "ahash", + "arrow-buffer 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "chrono", + "chrono-tz", + "half", + "hashbrown 0.15.5", + "num", +] + [[package]] name = "arrow-buffer" version = "52.2.0" @@ -144,34 +238,93 @@ dependencies = [ "num", ] +[[package]] +name = "arrow-buffer" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "169b1d5d6cb390dd92ce582b06b23815c7953e9dfaaea75556e89d890d19993d" +dependencies = [ + "bytes", + "half", + "num", +] + [[package]] name = "arrow-cast" version = "52.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "da26719e76b81d8bc3faad1d4dbdc1bcc10d14704e63dc17fc9f3e7e1e567c8e" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", - "arrow-select", + "arrow-array 52.2.0", + "arrow-buffer 52.2.0", + "arrow-data 52.2.0", + "arrow-schema 52.2.0", + "arrow-select 52.2.0", + "atoi", + "base64", + "chrono", + "half", + "lexical-core 0.8.5", + "num", + "ryu", +] + +[[package]] +name = "arrow-cast" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4f12eccc3e1c05a766cafb31f6a60a46c2f8efec9b74c6e0648766d30686af8" +dependencies = [ + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "arrow-select 55.2.0", "atoi", "base64", "chrono", + "comfy-table", "half", - "lexical-core", + "lexical-core 1.0.6", "num", "ryu", ] +[[package]] +name = "arrow-csv" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "012c9fef3f4a11573b2c74aec53712ff9fdae4a95f4ce452d1bbf088ee00f06b" +dependencies = [ + "arrow-array 55.2.0", + "arrow-cast 55.2.0", + "arrow-schema 55.2.0", + "chrono", + "csv", + "csv-core", + "regex", +] + [[package]] name = "arrow-data" version = "52.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dd9d6f18c65ef7a2573ab498c374d8ae364b4a4edf67105357491c031f716ca5" dependencies = [ - "arrow-buffer", - "arrow-schema", + "arrow-buffer 52.2.0", + "arrow-schema 52.2.0", + "half", + "num", +] + +[[package]] +name = "arrow-data" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de1ce212d803199684b658fc4ba55fb2d7e87b213de5af415308d2fee3619c2" +dependencies = [ + "arrow-buffer 55.2.0", + "arrow-schema 55.2.0", "half", "num", ] @@ -182,12 +335,96 @@ version = "52.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e786e1cdd952205d9a8afc69397b317cfbb6e0095e445c69cda7e8da5c1eeb0f" dependencies = [ - "arrow-array", - "arrow-buffer", - "arrow-cast", - "arrow-data", - "arrow-schema", - "flatbuffers", + "arrow-array 52.2.0", + "arrow-buffer 52.2.0", + "arrow-cast 52.2.0", + "arrow-data 52.2.0", + "arrow-schema 52.2.0", + "flatbuffers 24.12.23", +] + +[[package]] +name = "arrow-ipc" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9ea5967e8b2af39aff5d9de2197df16e305f47f404781d3230b2dc672da5d92" +dependencies = [ + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "flatbuffers 25.12.19", + "lz4_flex", +] + +[[package]] +name = "arrow-json" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5709d974c4ea5be96d900c01576c7c0b99705f4a3eec343648cb1ca863988a9c" +dependencies = [ + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-cast 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "chrono", + "half", + "indexmap 2.12.1", + "lexical-core 1.0.6", + "memchr", + "num", + "serde", + "serde_json", + "simdutf8", +] + +[[package]] +name = "arrow-json" +version = "55.2.0" +source = "git+https://github.com/ArroyoSystems/arrow-rs?branch=55.2.0%2Fjson#d31f8d8f97c6e1394b52927cd8c23c14fec6ba16" +dependencies = [ + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-cast 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "base64", + "chrono", + "half", + "indexmap 2.12.1", + "lexical-core 1.0.6", + "memchr", + "num", + "serde", + "serde_json", + "simdutf8", +] + +[[package]] +name = "arrow-ord" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6506e3a059e3be23023f587f79c82ef0bcf6d293587e3272d20f2d30b969b5a7" +dependencies = [ + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "arrow-select 55.2.0", +] + +[[package]] +name = "arrow-row" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52bf7393166beaf79b4bed9bfdf19e97472af32ce5b6b48169d321518a08cae2" +dependencies = [ + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "half", ] [[package]] @@ -196,6 +433,16 @@ version = "52.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e972cd1ff4a4ccd22f86d3e53e835c2ed92e0eea6a3e8eadb72b4f1ac802cf8" +[[package]] +name = "arrow-schema" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af7686986a3bf2254c9fb130c623cdcb2f8e1f15763e7c71c310f0834da3d292" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "arrow-select" version = "52.2.0" @@ -203,11 +450,59 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "600bae05d43483d216fb3494f8c32fdbefd8aa4e1de237e790dbb3d9f44690a3" dependencies = [ "ahash", - "arrow-array", - "arrow-buffer", - "arrow-data", - "arrow-schema", + "arrow-array 52.2.0", + "arrow-buffer 52.2.0", + "arrow-data 52.2.0", + "arrow-schema 52.2.0", + "num", +] + +[[package]] +name = "arrow-select" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd2b45757d6a2373faa3352d02ff5b54b098f5e21dccebc45a21806bc34501e5" +dependencies = [ + "ahash", + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "num", +] + +[[package]] +name = "arrow-string" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0377d532850babb4d927a06294314b316e23311503ed580ec6ce6a0158f49d40" +dependencies = [ + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-data 55.2.0", + "arrow-schema 55.2.0", + "arrow-select 55.2.0", + "memchr", "num", + "regex", + "regex-syntax", +] + +[[package]] +name = "async-compression" +version = "0.4.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06575e6a9673580f52661c92107baabffbf41e2141373441cbcdc47cb733003c" +dependencies = [ + "bzip2", + "flate2", + "futures-core", + "memchr", + "pin-project-lite", + "tokio", + "xz2", + "zstd", + "zstd-safe", ] [[package]] @@ -311,19 +606,67 @@ dependencies = [ "tower-service", ] +[[package]] +name = "backtrace" +version = "0.3.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb531853791a215d7c62a30daf0dde835f381ab5de4589cfe7c649d2cbe92bd6" +dependencies = [ + "addr2line", + "cfg-if", + "libc", + "miniz_oxide", + "object", + "rustc-demangle", + "windows-link", +] + +[[package]] +name = "backtrace-ext" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "537beee3be4a18fb023b570f80e3ae28003db9167a751266b259926e25539d50" +dependencies = [ + "backtrace", +] + [[package]] name = "base64" version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "bigdecimal" +version = "0.4.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4d6867f1565b3aad85681f1015055b087fcfd840d6aeee6eee7f2da317603695" +dependencies = [ + "autocfg", + "libm", + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "bincode" -version = "1.3.3" +version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1f45e9417d87227c7a56d22e471c6206462cba514c7590c09aff4cf6d1ddcad" +checksum = "36eaf5d7b090263e8150820482d5d93cd964a81e4019913c972f4edcc6edb740" dependencies = [ + "bincode_derive", "serde", + "unty", +] + +[[package]] +name = "bincode_derive" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf95709a440f45e986983918d0e8a1f30a9b1df04918fc828670606804ac3c09" +dependencies = [ + "virtue", ] [[package]] @@ -386,6 +729,41 @@ dependencies = [ "typenum", ] +[[package]] +name = "bitvec" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc2832c24239b0141d5674bb9174f9d68a8b5b3f2753311927c172ca46f7e9c" +dependencies = [ + "funty", + "radium", + "tap", + "wyz", +] + +[[package]] +name = "blake2" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46502ad458c9a52b69d4d4d32775c788b7a1b85e8bc9d482d92250fc0e3f8efe" +dependencies = [ + "digest", +] + +[[package]] +name = "blake3" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2468ef7d57b3fb7e16b576e8377cdbde2320c60e1491e961d11da40fc4f02a2d" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -395,6 +773,27 @@ dependencies = [ "generic-array", ] +[[package]] +name = "brotli" +version = "8.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bd8b9603c7aa97359dbd97ecf258968c95f3adddd6db2f7e7a5bef101c84560" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor", +] + +[[package]] +name = "brotli-decompressor" +version = "5.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "874bb8112abecc98cbd6d81ea4fa7e94fb9449648c93cc89aa40c81c24d7de03" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "bumpalo" version = "3.19.1" @@ -404,12 +803,27 @@ dependencies = [ "allocator-api2", ] +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + [[package]] name = "bytes" version = "1.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3" +[[package]] +name = "bzip2" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49ecfb22d906f800d4fe833b6282cf4dc1c298f5057ca0b5445e5c209735ca47" +dependencies = [ + "bzip2-sys", +] + [[package]] name = "bzip2-sys" version = "0.1.13+1.0.8" @@ -469,7 +883,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d8144c22e24bbcf26ade86cb6501a0916c46b7e4787abdb0045a467eb1645a1d" dependencies = [ "ambient-authority", - "rand", + "rand 0.8.5", ] [[package]] @@ -538,10 +952,32 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "145052bdd345b87320e369255277e3fb5152762ad123a901ef5c262dd38fe8d2" dependencies = [ "iana-time-zone", + "js-sys", "num-traits", + "wasm-bindgen", "windows-link", ] +[[package]] +name = "chrono-tz" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" +dependencies = [ + "chrono", + "phf 0.12.1", +] + +[[package]] +name = "chumsky" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8eebd66744a15ded14960ab4ccdbfb51ad3b81f51f3f04a80adac98c985396c9" +dependencies = [ + "hashbrown 0.14.5", + "stacker", +] + [[package]] name = "clang-sys" version = "1.8.1" @@ -581,7 +1017,7 @@ version = "4.5.49" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a0b5487afeab2deb2ff4e03a807ad1a03ac532ff5a2cee5d86884440c7f7671" dependencies = [ - "heck", + "heck 0.5.0", "proc-macro2", "quote", "syn", @@ -621,9 +1057,18 @@ dependencies = [ ] [[package]] -name = "colorchoice" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" +name = "codegen_template" +version = "0.1.0" +source = "git+https://github.com/ArroyoSystems/cornucopia?branch=sqlite#6a1a87a8bab82068d4a41525995ed0e715382209" +dependencies = [ + "unicode-xid", + "unscanny", +] + +[[package]] +name = "colorchoice" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" [[package]] @@ -657,12 +1102,59 @@ dependencies = [ "tiny-keccak", ] +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + [[package]] name = "core-foundation-sys" version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cornucopia" +version = "0.9.0" +source = "git+https://github.com/ArroyoSystems/cornucopia?branch=sqlite#6a1a87a8bab82068d4a41525995ed0e715382209" +dependencies = [ + "chumsky", + "clap", + "codegen_template", + "heck 0.4.1", + "indexmap 2.12.1", + "miette", + "postgres", + "postgres-types", + "prettyplease", + "rusqlite", + "syn", + "thiserror 1.0.69", +] + +[[package]] +name = "cornucopia_async" +version = "0.6.0" +source = "git+https://github.com/ArroyoSystems/cornucopia?branch=sqlite#6a1a87a8bab82068d4a41525995ed0e715382209" +dependencies = [ + "async-trait", + "cornucopia_client_core", + "deadpool-postgres", + "rusqlite", + "tokio-postgres", +] + +[[package]] +name = "cornucopia_client_core" +version = "0.4.0" +source = "git+https://github.com/ArroyoSystems/cornucopia?branch=sqlite#6a1a87a8bab82068d4a41525995ed0e715382209" +dependencies = [ + "fallible-iterator 0.2.0", + "postgres-protocol", + "postgres-types", +] + [[package]] name = "cpp_demangle" version = "0.4.5" @@ -757,7 +1249,7 @@ dependencies = [ "cranelift-assembler-x64-meta", "cranelift-codegen-shared", "cranelift-srcgen", - "heck", + "heck 0.5.0", "pulley-interpreter", ] @@ -860,48 +1352,713 @@ dependencies = [ ] [[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.29.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "crossterm" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" +dependencies = [ + "bitflags 2.10.0", + "crossterm_winapi", + "document-features", + "parking_lot", + "rustix 1.1.3", + "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.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "csv" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52cd9d68cf7efc6ddfaaee42e7288d3a99d613d4b50f76ce9827ae0c6e14f938" +dependencies = [ + "csv-core", + "itoa", + "ryu", + "serde_core", +] + +[[package]] +name = "csv-core" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704a3c26996a80471189265814dbc2c257598b96b8a7feae2d31ace646bb9782" +dependencies = [ + "memchr", +] + +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "datafusion" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "arrow-ipc 55.2.0", + "arrow-schema 55.2.0", + "async-trait", + "bytes", + "bzip2", + "chrono", + "datafusion-catalog", + "datafusion-catalog-listing", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-datasource-csv", + "datafusion-datasource-json", + "datafusion-datasource-parquet", + "datafusion-execution", + "datafusion-expr", + "datafusion-expr-common", + "datafusion-functions", + "datafusion-functions-aggregate", + "datafusion-functions-nested", + "datafusion-functions-table", + "datafusion-functions-window", + "datafusion-optimizer", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-optimizer", + "datafusion-physical-plan", + "datafusion-session", + "datafusion-sql", + "flate2", + "futures", + "itertools 0.14.0", + "log", + "object_store", + "parking_lot", + "parquet 55.2.0 (registry+https://github.com/rust-lang/crates.io-index)", + "rand 0.9.2", + "regex", + "sqlparser", + "tempfile", + "tokio", + "url", + "uuid", + "xz2", + "zstd", +] + +[[package]] +name = "datafusion-catalog" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "async-trait", + "dashmap", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-plan", + "datafusion-session", + "datafusion-sql", + "futures", + "itertools 0.14.0", + "log", + "object_store", + "parking_lot", + "tokio", +] + +[[package]] +name = "datafusion-catalog-listing" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "async-trait", + "datafusion-catalog", + "datafusion-common", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-session", + "futures", + "log", + "object_store", + "tokio", +] + +[[package]] +name = "datafusion-common" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "ahash", + "arrow", + "arrow-ipc 55.2.0", + "base64", + "half", + "hashbrown 0.14.5", + "indexmap 2.12.1", + "libc", + "log", + "object_store", + "parquet 55.2.0 (registry+https://github.com/rust-lang/crates.io-index)", + "paste", + "recursive", + "sqlparser", + "tokio", + "web-time", +] + +[[package]] +name = "datafusion-common-runtime" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "futures", + "log", + "tokio", +] + +[[package]] +name = "datafusion-datasource" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "async-compression", + "async-trait", + "bytes", + "bzip2", + "chrono", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-session", + "flate2", + "futures", + "glob", + "itertools 0.14.0", + "log", + "object_store", + "parquet 55.2.0 (registry+https://github.com/rust-lang/crates.io-index)", + "rand 0.9.2", + "tempfile", + "tokio", + "tokio-util", + "url", + "xz2", + "zstd", +] + +[[package]] +name = "datafusion-datasource-csv" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "async-trait", + "bytes", + "datafusion-catalog", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-session", + "futures", + "object_store", + "regex", + "tokio", +] + +[[package]] +name = "datafusion-datasource-json" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "async-trait", + "bytes", + "datafusion-catalog", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "datafusion-session", + "futures", + "object_store", + "serde_json", + "tokio", +] + +[[package]] +name = "datafusion-datasource-parquet" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "async-trait", + "bytes", + "datafusion-catalog", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-datasource", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions-aggregate", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-optimizer", + "datafusion-physical-plan", + "datafusion-session", + "futures", + "itertools 0.14.0", + "log", + "object_store", + "parking_lot", + "parquet 55.2.0 (registry+https://github.com/rust-lang/crates.io-index)", + "rand 0.9.2", + "tokio", +] + +[[package]] +name = "datafusion-doc" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" + +[[package]] +name = "datafusion-execution" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "dashmap", + "datafusion-common", + "datafusion-expr", + "futures", + "log", + "object_store", + "parking_lot", + "rand 0.9.2", + "tempfile", + "url", +] + +[[package]] +name = "datafusion-expr" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "chrono", + "datafusion-common", + "datafusion-doc", + "datafusion-expr-common", + "datafusion-functions-aggregate-common", + "datafusion-functions-window-common", + "datafusion-physical-expr-common", + "indexmap 2.12.1", + "paste", + "recursive", + "serde_json", + "sqlparser", +] + +[[package]] +name = "datafusion-expr-common" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "datafusion-common", + "indexmap 2.12.1", + "itertools 0.14.0", + "paste", +] + +[[package]] +name = "datafusion-functions" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "arrow-buffer 55.2.0", + "base64", + "blake2", + "blake3", + "chrono", + "datafusion-common", + "datafusion-doc", + "datafusion-execution", + "datafusion-expr", + "datafusion-expr-common", + "datafusion-macros", + "hex", + "itertools 0.14.0", + "log", + "md-5", + "rand 0.9.2", + "regex", + "sha2", + "unicode-segmentation", + "uuid", +] + +[[package]] +name = "datafusion-functions-aggregate" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "ahash", + "arrow", + "datafusion-common", + "datafusion-doc", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions-aggregate-common", + "datafusion-macros", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "half", + "log", + "paste", +] + +[[package]] +name = "datafusion-functions-aggregate-common" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "ahash", + "arrow", + "datafusion-common", + "datafusion-expr-common", + "datafusion-physical-expr-common", +] + +[[package]] +name = "datafusion-functions-nested" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "arrow-ord", + "datafusion-common", + "datafusion-doc", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions", + "datafusion-functions-aggregate", + "datafusion-macros", + "datafusion-physical-expr-common", + "itertools 0.14.0", + "log", + "paste", +] + +[[package]] +name = "datafusion-functions-table" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "async-trait", + "datafusion-catalog", + "datafusion-common", + "datafusion-expr", + "datafusion-physical-plan", + "parking_lot", + "paste", +] + +[[package]] +name = "datafusion-functions-window" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "datafusion-common", + "datafusion-doc", + "datafusion-expr", + "datafusion-functions-window-common", + "datafusion-macros", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "log", + "paste", +] + +[[package]] +name = "datafusion-functions-window-common" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "datafusion-common", + "datafusion-physical-expr-common", +] + +[[package]] +name = "datafusion-macros" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "datafusion-expr", + "quote", + "syn", +] + +[[package]] +name = "datafusion-optimizer" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "chrono", + "datafusion-common", + "datafusion-expr", + "datafusion-physical-expr", + "indexmap 2.12.1", + "itertools 0.14.0", + "log", + "recursive", + "regex", + "regex-syntax", +] + +[[package]] +name = "datafusion-physical-expr" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "ahash", + "arrow", + "datafusion-common", + "datafusion-expr", + "datafusion-expr-common", + "datafusion-functions-aggregate-common", + "datafusion-physical-expr-common", + "half", + "hashbrown 0.14.5", + "indexmap 2.12.1", + "itertools 0.14.0", + "log", + "paste", + "petgraph 0.8.3", +] + +[[package]] +name = "datafusion-physical-expr-common" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "ahash", + "arrow", + "datafusion-common", + "datafusion-expr-common", + "hashbrown 0.14.5", + "itertools 0.14.0", +] + +[[package]] +name = "datafusion-physical-optimizer" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "datafusion-common", + "datafusion-execution", + "datafusion-expr", + "datafusion-expr-common", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "datafusion-physical-plan", + "itertools 0.14.0", + "log", + "recursive", +] + +[[package]] +name = "datafusion-physical-plan" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "ahash", + "arrow", + "arrow-ord", + "arrow-schema 55.2.0", + "async-trait", + "chrono", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions-window-common", + "datafusion-physical-expr", + "datafusion-physical-expr-common", + "futures", + "half", + "hashbrown 0.14.5", + "indexmap 2.12.1", + "itertools 0.14.0", + "log", + "parking_lot", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "datafusion-proto" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "chrono", + "datafusion", + "datafusion-common", + "datafusion-expr", + "datafusion-proto-common", + "object_store", + "prost", +] + +[[package]] +name = "datafusion-proto-common" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "datafusion-common", + "prost", +] + +[[package]] +name = "datafusion-session" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" +dependencies = [ + "arrow", + "async-trait", + "dashmap", + "datafusion-common", + "datafusion-common-runtime", + "datafusion-execution", + "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-plan", + "datafusion-sql", + "futures", + "itertools 0.14.0", + "log", + "object_store", + "parking_lot", + "tokio", +] + +[[package]] +name = "datafusion-sql" +version = "48.0.1" +source = "git+https://github.com/ArroyoSystems/arrow-datafusion?branch=48.0.1%2Farroyo#916b45f5c28d94765ae4a6393c5e126b2ea55e1c" dependencies = [ - "bitflags 2.10.0", - "crossterm_winapi", - "document-features", - "parking_lot", - "rustix 1.1.3", - "winapi", + "arrow", + "bigdecimal", + "datafusion-common", + "datafusion-expr", + "indexmap 2.12.1", + "log", + "recursive", + "regex", + "sqlparser", ] [[package]] -name = "crossterm_winapi" -version = "0.9.1" +name = "deadpool" +version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "acdd7c62a3665c7f6830a51635d9ac9b23ed385797f70a83bb8bafe9c572ab2b" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" dependencies = [ - "winapi", + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", ] [[package]] -name = "crunchy" -version = "0.2.4" +name = "deadpool-postgres" +version = "0.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +checksum = "3d697d376cbfa018c23eb4caab1fd1883dd9c906a8c034e8d9a3cb06a7e0bef9" +dependencies = [ + "async-trait", + "deadpool", + "getrandom 0.2.16", + "tokio", + "tokio-postgres", + "tracing", +] [[package]] -name = "crypto-common" -version = "0.1.7" +name = "deadpool-runtime" +version = "0.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" dependencies = [ - "generic-array", - "typenum", + "tokio", ] [[package]] @@ -930,6 +2087,7 @@ checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ "block-buffer", "crypto-common", + "subtle", ] [[package]] @@ -985,6 +2143,12 @@ dependencies = [ "shared_child", ] +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + [[package]] name = "either" version = "1.15.0" @@ -1053,12 +2217,24 @@ version = "3.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dea2df4cf52843e0452895c455a1a2cfbb842a1e7329671acf418fdc53ed4c59" +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + [[package]] name = "fallible-iterator" version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + [[package]] name = "fastrand" version = "2.3.0" @@ -1104,6 +2280,27 @@ dependencies = [ "rustc_version", ] +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35f6839d7b3b98adde531effaf34f0c2badc6f4735d26fe74709d8e513a96ef3" +dependencies = [ + "bitflags 2.10.0", + "rustc_version", +] + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", + "zlib-rs", +] + [[package]] name = "fnv" version = "1.0.7" @@ -1141,27 +2338,51 @@ name = "function-stream" version = "0.6.0" dependencies = [ "anyhow", - "arrow-array", - "arrow-ipc", - "arrow-schema", + "arrow", + "arrow-array 55.2.0", + "arrow-ipc 55.2.0", + "arrow-json 55.2.0 (git+https://github.com/ArroyoSystems/arrow-rs?branch=55.2.0%2Fjson)", + "arrow-schema 55.2.0", "async-trait", "base64", "bincode", + "chrono", "clap", + "cornucopia", + "cornucopia_async", "crossbeam-channel", + "datafusion", + "datafusion-common", + "datafusion-execution", + "datafusion-expr", + "datafusion-functions", + "datafusion-functions-aggregate", + "datafusion-functions-window", + "datafusion-physical-expr", + "datafusion-physical-plan", + "datafusion-proto", + "futures", + "itertools 0.14.0", + "jiter", "log", "lru", "num_cpus", "parking_lot", + "parquet 55.2.0 (git+https://github.com/ArroyoSystems/arrow-rs?branch=55.2.0%2Fparquet)", "pest", "pest_derive", + "petgraph 0.7.1", "proctitle", + "prost", "protocol", "rdkafka", "rocksdb", "serde", "serde_json", + "serde_json_path", "serde_yaml", + "sqlparser", + "strum", "thiserror 2.0.17", "tokio", "tokio-stream", @@ -1169,18 +2390,21 @@ dependencies = [ "tracing", "tracing-appender", "tracing-subscriber", + "typify", + "unicase", "uuid", "wasmtime", "wasmtime-wasi", + "xxhash-rust", ] [[package]] name = "function-stream-cli" version = "0.1.0" dependencies = [ - "arrow-array", - "arrow-ipc", - "arrow-schema", + "arrow-array 52.2.0", + "arrow-ipc 52.2.0", + "arrow-schema 52.2.0", "clap", "comfy-table", "function-stream", @@ -1191,6 +2415,12 @@ dependencies = [ "tonic", ] +[[package]] +name = "funty" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" + [[package]] name = "futures" version = "0.3.31" @@ -1199,6 +2429,7 @@ checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" dependencies = [ "futures-channel", "futures-core", + "futures-executor", "futures-io", "futures-sink", "futures-task", @@ -1221,12 +2452,34 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +[[package]] +name = "futures-executor" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + [[package]] name = "futures-io" version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" +[[package]] +name = "futures-macro" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "futures-sink" version = "0.3.31" @@ -1248,6 +2501,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-io", + "futures-macro", "futures-sink", "futures-task", "memchr", @@ -1287,8 +2541,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" dependencies = [ "cfg-if", + "js-sys", "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", + "wasm-bindgen", ] [[package]] @@ -1309,7 +2565,7 @@ version = "0.32.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" dependencies = [ - "fallible-iterator", + "fallible-iterator 0.3.0", "indexmap 2.12.1", "stable_deref_trait", ] @@ -1357,11 +2613,24 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" +[[package]] +name = "hashbrown" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43a3c133739dddd0d2990f9a4bdf8eb4b21ef50e4851ca85ab661199821d510e" +dependencies = [ + "ahash", +] + [[package]] name = "hashbrown" version = "0.14.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" +dependencies = [ + "ahash", + "allocator-api2", +] [[package]] name = "hashbrown" @@ -1381,6 +2650,21 @@ version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +[[package]] +name = "hashlink" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ba4ff7128dee98c7dc9794b6a411377e1404dba1c97deb8d1a55297bd25d8af" +dependencies = [ + "hashbrown 0.14.5", +] + +[[package]] +name = "heck" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "95505c38b4572b2d910cecb0281560f54b440a19336cbbcb27bf6ce6adc6f5a8" + [[package]] name = "heck" version = "0.5.0" @@ -1393,6 +2677,21 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + [[package]] name = "home" version = "0.5.12" @@ -1649,7 +2948,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "af1955a75fa080c677d3972822ec4bad316169ab1cfc6c257a942c2265dbe5fe" dependencies = [ "bitmaps", - "rand_core", + "rand_core 0.6.4", "rand_xoshiro", "sized-chunks", "typenum", @@ -1678,6 +2977,21 @@ dependencies = [ "serde_core", ] +[[package]] +name = "integer-encoding" +version = "3.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" + +[[package]] +name = "inventory" +version = "0.3.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "009ae045c87e7082cb72dab0ccd01ae075dd00141ddc108f43a0ea150a9e7227" +dependencies = [ + "rustversion", +] + [[package]] name = "io-extras" version = "0.18.4" @@ -1711,6 +3025,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "is_ci" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7655c9839580ee829dfacba1d1278c2b7883e50a277ff7541299489d6bdfdc45" + [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -1761,6 +3081,19 @@ dependencies = [ "cc", ] +[[package]] +name = "jiter" +version = "0.10.0" +source = "git+https://github.com/ArroyoSystems/jiter?branch=disable_python#e5a90990780433a5972031a62eff87555d98884d" +dependencies = [ + "ahash", + "bitvec", + "lexical-parse-float 1.0.6", + "num-bigint", + "num-traits", + "smallvec", +] + [[package]] name = "jobserver" version = "0.1.34" @@ -1811,11 +3144,24 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2cde5de06e8d4c2faabc400238f9ae1c74d5412d03a7bd067645ccbc47070e46" dependencies = [ - "lexical-parse-float", - "lexical-parse-integer", - "lexical-util", - "lexical-write-float", - "lexical-write-integer", + "lexical-parse-float 0.8.5", + "lexical-parse-integer 0.8.6", + "lexical-util 0.8.5", + "lexical-write-float 0.8.5", + "lexical-write-integer 0.8.5", +] + +[[package]] +name = "lexical-core" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d8d125a277f807e55a77304455eb7b1cb52f2b18c143b60e766c120bd64a594" +dependencies = [ + "lexical-parse-float 1.0.6", + "lexical-parse-integer 1.0.6", + "lexical-util 1.0.7", + "lexical-write-float 1.0.6", + "lexical-write-integer 1.0.6", ] [[package]] @@ -1824,21 +3170,40 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683b3a5ebd0130b8fb52ba0bdc718cc56815b6a097e28ae5a6997d0ad17dc05f" dependencies = [ - "lexical-parse-integer", - "lexical-util", + "lexical-parse-integer 0.8.6", + "lexical-util 0.8.5", "static_assertions", ] +[[package]] +name = "lexical-parse-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52a9f232fbd6f550bc0137dcb5f99ab674071ac2d690ac69704593cb4abbea56" +dependencies = [ + "lexical-parse-integer 1.0.6", + "lexical-util 1.0.7", +] + [[package]] name = "lexical-parse-integer" version = "0.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6d0994485ed0c312f6d965766754ea177d07f9c00c9b82a5ee62ed5b47945ee9" dependencies = [ - "lexical-util", + "lexical-util 0.8.5", "static_assertions", ] +[[package]] +name = "lexical-parse-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a7a039f8fb9c19c996cd7b2fcce303c1b2874fe1aca544edc85c4a5f8489b34" +dependencies = [ + "lexical-util 1.0.7", +] + [[package]] name = "lexical-util" version = "0.8.5" @@ -1848,27 +3213,52 @@ dependencies = [ "static_assertions", ] +[[package]] +name = "lexical-util" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2604dd126bb14f13fb5d1bd6a66155079cb9fa655b37f875b3a742c705dbed17" + [[package]] name = "lexical-write-float" version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "accabaa1c4581f05a3923d1b4cfd124c329352288b7b9da09e766b0668116862" dependencies = [ - "lexical-util", - "lexical-write-integer", + "lexical-util 0.8.5", + "lexical-write-integer 0.8.5", "static_assertions", ] +[[package]] +name = "lexical-write-float" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50c438c87c013188d415fbabbb1dceb44249ab81664efbd31b14ae55dabb6361" +dependencies = [ + "lexical-util 1.0.7", + "lexical-write-integer 1.0.6", +] + [[package]] name = "lexical-write-integer" version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e1b6f3d1f4422866b68192d62f77bc5c700bee84f3069f2469d7bc8c77852446" dependencies = [ - "lexical-util", + "lexical-util 0.8.5", "static_assertions", ] +[[package]] +name = "lexical-write-integer" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "409851a618475d2d5796377cad353802345cba92c867d9fbcde9cf4eac4e14df" +dependencies = [ + "lexical-util 1.0.7", +] + [[package]] name = "libc" version = "0.2.179" @@ -1917,6 +3307,17 @@ dependencies = [ "zstd-sys", ] +[[package]] +name = "libsqlite3-sys" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c10584274047cb335c23d3e61bcef8e323adae7c5c8c760540f73610177fc3f" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + [[package]] name = "libz-sys" version = "1.1.23" @@ -1987,6 +3388,26 @@ dependencies = [ "libc", ] +[[package]] +name = "lz4_flex" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "373f5eceeeab7925e0c1098212f2fbc4d416adec9d35051a6ab251e824c1854a" +dependencies = [ + "twox-hash", +] + +[[package]] +name = "lzma-sys" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fda04ab3764e6cde78b9974eec4f779acaba7c4e84b36eca3cf77c581b85d27" +dependencies = [ + "cc", + "libc", + "pkg-config", +] + [[package]] name = "mach2" version = "0.4.3" @@ -2017,6 +3438,16 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4facc753ae494aeb6e3c22f839b158aebd4f9270f55cd3c79906c45476c47ab4" +[[package]] +name = "md-5" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" +dependencies = [ + "cfg-if", + "digest", +] + [[package]] name = "memchr" version = "2.7.6" @@ -2032,6 +3463,38 @@ dependencies = [ "rustix 1.1.3", ] +[[package]] +name = "miette" +version = "5.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59bb584eaeeab6bd0226ccf3509a69d7936d148cf3d036ad350abe35e8c6856e" +dependencies = [ + "backtrace", + "backtrace-ext", + "is-terminal", + "miette-derive", + "once_cell", + "owo-colors", + "supports-color", + "supports-hyperlinks", + "supports-unicode", + "terminal_size", + "textwrap", + "thiserror 1.0.69", + "unicode-width 0.1.14", +] + +[[package]] +name = "miette-derive" +version = "5.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49e7bc1560b95a3c4a25d03de42fe76ca718ab92d1a22a55b9b4cf67b3ae635c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "mime" version = "0.3.17" @@ -2044,6 +3507,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + [[package]] name = "mio" version = "1.1.1" @@ -2051,7 +3524,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" dependencies = [ "libc", - "wasi", + "wasi 0.11.1+wasi-snapshot-preview1", "windows-sys 0.61.2", ] @@ -2225,6 +3698,30 @@ dependencies = [ "memchr", ] +[[package]] +name = "object_store" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbfbfff40aeccab00ec8a910b57ca8ecf4319b335c542f2edcd19dd25a1e2a00" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures", + "http", + "humantime", + "itertools 0.14.0", + "parking_lot", + "percent-encoding", + "thiserror 2.0.17", + "tokio", + "tracing", + "url", + "walkdir", + "wasm-bindgen-futures", + "web-time", +] + [[package]] name = "once_cell" version = "1.21.3" @@ -2249,6 +3746,15 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "ordered-float" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" +dependencies = [ + "num-traits", +] + [[package]] name = "os_pipe" version = "1.2.3" @@ -2259,6 +3765,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "owo-colors" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1b04fb49957986fdce4d6ee7a65027d55d4b6d2265e5848bbb507b58ccfdb6f" + [[package]] name = "parking_lot" version = "0.12.5" @@ -2282,6 +3794,80 @@ dependencies = [ "windows-link", ] +[[package]] +name = "parquet" +version = "55.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b17da4150748086bd43352bc77372efa9b6e3dbd06a04831d2a98c041c225cfa" +dependencies = [ + "ahash", + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-cast 55.2.0", + "arrow-data 55.2.0", + "arrow-ipc 55.2.0", + "arrow-schema 55.2.0", + "arrow-select 55.2.0", + "base64", + "brotli", + "bytes", + "chrono", + "flate2", + "futures", + "half", + "hashbrown 0.15.5", + "lz4_flex", + "num", + "num-bigint", + "object_store", + "paste", + "seq-macro", + "simdutf8", + "snap", + "thrift", + "tokio", + "twox-hash", + "zstd", +] + +[[package]] +name = "parquet" +version = "55.2.0" +source = "git+https://github.com/ArroyoSystems/arrow-rs?branch=55.2.0%2Fparquet#d1d2dd8edf673cddc79ba6403dc6508263a2ddda" +dependencies = [ + "ahash", + "arrow-array 55.2.0", + "arrow-buffer 55.2.0", + "arrow-cast 55.2.0", + "arrow-data 55.2.0", + "arrow-ipc 55.2.0", + "arrow-schema 55.2.0", + "arrow-select 55.2.0", + "base64", + "brotli", + "bytes", + "chrono", + "flate2", + "half", + "hashbrown 0.15.5", + "lz4_flex", + "num", + "num-bigint", + "paste", + "seq-macro", + "simdutf8", + "snap", + "thrift", + "twox-hash", + "zstd", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + [[package]] name = "peeking_take_while" version = "0.1.2" @@ -2357,6 +3943,55 @@ dependencies = [ "indexmap 2.12.1", ] +[[package]] +name = "petgraph" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8701b58ea97060d5e5b155d383a69952a60943f0e6dfe30b04c287beb0b27455" +dependencies = [ + "fixedbitset 0.5.7", + "hashbrown 0.15.5", + "indexmap 2.12.1", + "serde", +] + +[[package]] +name = "phf" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" +dependencies = [ + "phf_shared 0.12.1", +] + +[[package]] +name = "phf" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1562dc717473dbaa4c1f85a36410e03c047b2e7df7f45ee938fbef64ae7fadf" +dependencies = [ + "phf_shared 0.13.1", + "serde", +] + +[[package]] +name = "phf_shared" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981" +dependencies = [ + "siphasher", +] + +[[package]] +name = "phf_shared" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e57fef6bc5981e38c2ce2d63bfa546861309f875b8a75f092d1d54ae2d64f266" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.1.10" @@ -2407,6 +4042,49 @@ dependencies = [ "serde", ] +[[package]] +name = "postgres" +version = "0.19.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c48ece1c6cda0db61b058c1721378da76855140e9214339fa1317decacb176" +dependencies = [ + "bytes", + "fallible-iterator 0.2.0", + "futures-util", + "log", + "tokio", + "tokio-postgres", +] + +[[package]] +name = "postgres-protocol" +version = "0.6.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ee9dd5fe15055d2b6806f4736aa0c9637217074e224bbec46d4041b91bb9491" +dependencies = [ + "base64", + "byteorder", + "bytes", + "fallible-iterator 0.2.0", + "hmac", + "md-5", + "memchr", + "rand 0.9.2", + "sha2", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54b858f82211e84682fecd373f68e1ceae642d8d751a1ebd13f33de6257b3e20" +dependencies = [ + "bytes", + "fallible-iterator 0.2.0", + "postgres-protocol", +] + [[package]] name = "potential_utf" version = "0.1.4" @@ -2486,7 +4164,7 @@ version = "0.13.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ - "heck", + "heck 0.5.0", "itertools 0.14.0", "log", "multimap", @@ -2529,10 +4207,21 @@ dependencies = [ "env_logger", "log", "prost", + "serde", "tonic", "tonic-build", ] +[[package]] +name = "psm" +version = "0.1.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3852766467df634d74f0b2d7819bf8dc483a0eb2e3b0f50f756f9cfe8b0d18d8" +dependencies = [ + "ar_archive_writer", + "cc", +] + [[package]] name = "pulley-interpreter" version = "41.0.3" @@ -2571,6 +4260,12 @@ version = "5.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" +[[package]] +name = "radium" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" + [[package]] name = "radix_trie" version = "0.2.1" @@ -2588,8 +4283,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", ] [[package]] @@ -2599,7 +4304,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", ] [[package]] @@ -2611,13 +4326,22 @@ dependencies = [ "getrandom 0.2.16", ] +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + [[package]] name = "rand_xoshiro" version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa" dependencies = [ - "rand_core", + "rand_core 0.6.4", ] [[package]] @@ -2673,6 +4397,26 @@ dependencies = [ "sasl2-sys", ] +[[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.5.18" @@ -2737,6 +4481,16 @@ version = "0.8.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7a2d987857b319362043e95f5353c0535c1f58eec5336fdfcf626430af7def58" +[[package]] +name = "regress" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82a9ecfa0cb04d0b04dddb99b8ccf4f66bc8dfd23df694b398570bd8ae3a50fb" +dependencies = [ + "hashbrown 0.13.2", + "memchr", +] + [[package]] name = "rocksdb" version = "0.21.0" @@ -2747,6 +4501,21 @@ dependencies = [ "librocksdb-sys", ] +[[package]] +name = "rusqlite" +version = "0.31.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b838eba278d213a8beaf485bd313fd580ca4505a00d5871caeb1457c55322cae" +dependencies = [ + "bitflags 2.10.0", + "fallible-iterator 0.3.0", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "serde_json", + "smallvec", +] + [[package]] name = "rustc-demangle" version = "0.1.26" @@ -2844,6 +4613,15 @@ version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a50f4cf475b65d88e057964e0e9bb1f0aa9bbb2036dc65c64596b42932536984" +[[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 = "sasl2-sys" version = "0.1.22+2.1.28" @@ -2856,6 +4634,30 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "schemars" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fbf2ae1b8bc8e02df939598064d22402220cd5bbcca1c76f7d6a310974d5615" +dependencies = [ + "dyn-clone", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e265784ad618884abaea0600a9adf15393368d840e0222d101a072f3f7534d" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -2872,6 +4674,12 @@ dependencies = [ "serde_core", ] +[[package]] +name = "seq-macro" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" + [[package]] name = "serde" version = "1.0.228" @@ -2902,6 +4710,17 @@ dependencies = [ "syn", ] +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "serde_json" version = "1.0.148" @@ -2915,6 +4734,56 @@ dependencies = [ "zmij", ] +[[package]] +name = "serde_json_path" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b992cea3194eea663ba99a042d61cea4bd1872da37021af56f6a37e0359b9d33" +dependencies = [ + "inventory", + "nom", + "regex", + "serde", + "serde_json", + "serde_json_path_core", + "serde_json_path_macros", + "thiserror 2.0.17", +] + +[[package]] +name = "serde_json_path_core" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dde67d8dfe7d4967b5a95e247d4148368ddd1e753e500adb34b3ffe40c6bc1bc" +dependencies = [ + "inventory", + "serde", + "serde_json", + "thiserror 2.0.17", +] + +[[package]] +name = "serde_json_path_macros" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "517acfa7f77ddaf5c43d5f119c44a683774e130b4247b7d3210f8924506cfac8" +dependencies = [ + "inventory", + "serde_json_path_core", + "serde_json_path_macros_internal", +] + +[[package]] +name = "serde_json_path_macros_internal" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aafbefbe175fa9bf03ca83ef89beecff7d2a95aaacd5732325b90ac8c3bd7b90" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "serde_spanned" version = "1.0.4" @@ -2924,6 +4793,18 @@ dependencies = [ "serde_core", ] +[[package]] +name = "serde_tokenstream" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c49585c52c01f13c5c2ebb333f14f6885d76daa768d8a037d28017ec538c69" +dependencies = [ + "proc-macro2", + "quote", + "serde", + "syn", +] + [[package]] name = "serde_yaml" version = "0.9.34+deprecated" @@ -2996,14 +4877,32 @@ dependencies = [ ] [[package]] -name = "signal-hook-registry" -version = "1.4.8" +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "simd-adler32" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + +[[package]] +name = "siphasher" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" -dependencies = [ - "errno", - "libc", -] +checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" [[package]] name = "sized-chunks" @@ -3030,6 +4929,18 @@ dependencies = [ "serde", ] +[[package]] +name = "smawk" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c388c1b5e93756d0c740965c41e8822f866621d41acbdf6336a6a168f8840c" + +[[package]] +name = "snap" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b6b67fb9a61334225b5b790716f609cd58395f895b3fe8b328786812a40bc3b" + [[package]] name = "socket2" version = "0.5.10" @@ -3050,24 +4961,124 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "sqlparser" +version = "0.55.0" +source = "git+https://github.com/FunctionStream/sqlparser-rs?branch=0.6.0%2Ffunction-sql-parser#7e7cfb6145a426a26a7db12ae5874fed8b9c6b95" +dependencies = [ + "log", + "recursive", + "sqlparser_derive", +] + +[[package]] +name = "sqlparser_derive" +version = "0.3.0" +source = "git+https://github.com/FunctionStream/sqlparser-rs?branch=0.6.0%2Ffunction-sql-parser#7e7cfb6145a426a26a7db12ae5874fed8b9c6b95" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "stacker" +version = "0.1.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d74a23609d509411d10e2176dc2a4346e3b4aea2e7b1869f19fdedbc71c013" +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 = "stringprep" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b4df3d392d81bd458a8a621b8bffbd2302a12ffe288a9d931670948749463b1" +dependencies = [ + "unicode-bidi", + "unicode-normalization", + "unicode-properties", +] + [[package]] name = "strsim" version = "0.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" +[[package]] +name = "strum" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck 0.5.0", + "proc-macro2", + "quote", + "rustversion", + "syn", +] + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "supports-color" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6398cde53adc3c4557306a96ce67b302968513830a77a95b2b17305d9719a89" +dependencies = [ + "is-terminal", + "is_ci", +] + +[[package]] +name = "supports-hyperlinks" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f84231692eb0d4d41e4cdd0cabfdd2e6cd9e255e65f80c9aa7c98dd502b4233d" +dependencies = [ + "is-terminal", +] + +[[package]] +name = "supports-unicode" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f850c19edd184a205e883199a261ed44471c81e39bd95b1357f5febbef00e77a" +dependencies = [ + "is-terminal", +] + [[package]] name = "syn" version = "2.0.113" @@ -3112,6 +5123,12 @@ dependencies = [ "winx", ] +[[package]] +name = "tap" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" + [[package]] name = "target-lexicon" version = "0.13.4" @@ -3140,6 +5157,27 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "terminal_size" +version = "0.1.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "633c1a546cee861a1a6d0dc69ebeca693bf4296661ba7852b9d21d159e0506df" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "textwrap" +version = "0.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7b3e525a49ec206798b40326a44121291b530c963cfb01018f63e135bac543d" +dependencies = [ + "smawk", + "unicode-linebreak", + "unicode-width 0.1.14", +] + [[package]] name = "thiserror" version = "1.0.69" @@ -3189,6 +5227,17 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "thrift" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" +dependencies = [ + "byteorder", + "integer-encoding", + "ordered-float", +] + [[package]] name = "time" version = "0.3.44" @@ -3239,6 +5288,21 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +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 = "tokio" version = "1.49.0" @@ -3267,6 +5331,32 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-postgres" +version = "0.7.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcea47c8f71744367793f16c2db1f11cb859d28f436bdb4ca9193eb1f787ee42" +dependencies = [ + "async-trait", + "byteorder", + "bytes", + "fallible-iterator 0.2.0", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf 0.13.1", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand 0.9.2", + "socket2 0.6.1", + "tokio", + "tokio-util", + "whoami", +] + [[package]] name = "tokio-stream" version = "0.1.18" @@ -3397,7 +5487,7 @@ dependencies = [ "indexmap 1.9.3", "pin-project", "pin-project-lite", - "rand", + "rand 0.8.5", "slab", "tokio", "tokio-util", @@ -3525,24 +5615,104 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "twox-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" + [[package]] name = "typenum" version = "1.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" +[[package]] +name = "typify" +version = "0.0.13" +source = "git+https://github.com/ArroyoSystems/typify.git?branch=arroyo#d14b6fc016bf9d63618d8b43b4d74a648980737b" +dependencies = [ + "typify-impl", + "typify-macro", +] + +[[package]] +name = "typify-impl" +version = "0.0.13" +source = "git+https://github.com/ArroyoSystems/typify.git?branch=arroyo#d14b6fc016bf9d63618d8b43b4d74a648980737b" +dependencies = [ + "heck 0.4.1", + "log", + "proc-macro2", + "quote", + "regress", + "schemars", + "serde_json", + "syn", + "thiserror 1.0.69", + "unicode-ident", +] + +[[package]] +name = "typify-macro" +version = "0.0.13" +source = "git+https://github.com/ArroyoSystems/typify.git?branch=arroyo#d14b6fc016bf9d63618d8b43b4d74a648980737b" +dependencies = [ + "proc-macro2", + "quote", + "schemars", + "serde", + "serde_json", + "serde_tokenstream", + "syn", + "typify-impl", +] + [[package]] name = "ucd-trie" version = "0.1.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" +[[package]] +name = "unicase" +version = "2.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" + +[[package]] +name = "unicode-bidi" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c1cb5db39152898a79168971543b1cb5020dff7fe43c8dc468b0885f5e29df5" + [[package]] name = "unicode-ident" version = "1.0.22" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" +[[package]] +name = "unicode-linebreak" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b09c83c3c29d37506a3e260c08c03743a6bb66a9cd432c6934ab501a190571f" + +[[package]] +name = "unicode-normalization" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-properties" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" + [[package]] name = "unicode-segmentation" version = "1.12.0" @@ -3573,6 +5743,18 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" +[[package]] +name = "unscanny" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9df2af067a7953e9c3831320f35c1cc0600c30d44d9f7a12b01db1cd88d6b47" + +[[package]] +name = "unty" +version = "0.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d49784317cd0d1ee7ec5c716dd598ec5b4483ea832a2dced265471cc0f690ae" + [[package]] name = "url" version = "2.5.7" @@ -3626,6 +5808,22 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "virtue" +version = "0.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "051eb1abcf10076295e815102942cc58f9d5e3b4560e46e53c21e8ff6f3af7b1" + +[[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" @@ -3641,6 +5839,15 @@ version = "0.11.1+wasi-snapshot-preview1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" +[[package]] +name = "wasi" +version = "0.14.7+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "883478de20367e224c0090af9cf5f9fa85bed63a95c1abf3afc5c083ebc06e8c" +dependencies = [ + "wasip2", +] + [[package]] name = "wasip2" version = "1.0.1+wasi-0.2.4" @@ -3650,6 +5857,15 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "wasite" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fe902b4a6b8028a753d5424909b764ccf79b7a209eac9bf97e59cda9f71a42" +dependencies = [ + "wasi 0.14.7+wasi-0.2.4", +] + [[package]] name = "wasm-bindgen" version = "0.2.106" @@ -3663,6 +5879,19 @@ dependencies = [ "wasm-bindgen-shared", ] +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.56" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "836d9622d604feee9e5de25ac10e3ea5f2d65b41eac0d9ce72eb5deae707ce7c" +dependencies = [ + "cfg-if", + "js-sys", + "once_cell", + "wasm-bindgen", + "web-sys", +] + [[package]] name = "wasm-bindgen-macro" version = "0.2.106" @@ -3702,7 +5931,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "af801b6f36459023eaec63fdbaedad2fd5a4ab7dc74ecc110a8b5d375c5775e4" dependencies = [ "anyhow", - "heck", + "heck 0.5.0", "im-rc", "indexmap 2.12.1", "log", @@ -4005,7 +6234,7 @@ checksum = "87acbd416227cdd279565ba49e57cf7f08d112657c3b3f39b70250acdfd094fe" dependencies = [ "anyhow", "bitflags 2.10.0", - "heck", + "heck 0.5.0", "indexmap 2.12.1", "wit-parser", ] @@ -4085,6 +6314,37 @@ dependencies = [ "wast 243.0.0", ] +[[package]] +name = "web-sys" +version = "0.3.83" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b32828d774c412041098d182a8b38b16ea816958e07cf40eec2bc080ae137ac" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "whoami" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fae98cf96deed1b7572272dfc777713c249ae40aa1cf8862e091e8b745f5361" +dependencies = [ + "libredox", + "wasite", + "web-sys", +] + [[package]] name = "wiggle" version = "41.0.3" @@ -4106,7 +6366,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57f773d51c1696bd7d028aa35c884d9fc58f48d79a1176dfbad6c908de314235" dependencies = [ "anyhow", - "heck", + "heck 0.5.0", "proc-macro2", "quote", "syn", @@ -4461,6 +6721,30 @@ version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" +[[package]] +name = "wyz" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05f360fc0b24296329c78fda852a1e9ae82de9cf7b27dae4b7f62f118f77b9ed" +dependencies = [ + "tap", +] + +[[package]] +name = "xxhash-rust" +version = "0.8.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" + +[[package]] +name = "xz2" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "388c44dc09d76f1536602ead6d325eb532f5c122f17782bd57fb47baeeb767e2" +dependencies = [ + "lzma-sys", +] + [[package]] name = "yoke" version = "0.8.1" @@ -4558,6 +6842,12 @@ dependencies = [ "syn", ] +[[package]] +name = "zlib-rs" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513" + [[package]] name = "zmij" version = "1.0.10" diff --git a/Cargo.toml b/Cargo.toml index 4b855aa9..8b38dfe4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ tonic = { version = "0.12", features = ["default"] } async-trait = "0.1" num_cpus = "1.0" protocol = { path = "./protocol" } +prost = "0.13" rdkafka = { version = "0.38", features = ["cmake-build", "ssl", "gssapi"] } crossbeam-channel = "0.5" pest = "2.7" @@ -44,16 +45,49 @@ wasmtime = { version = "41.0.3", features = ["component-model", "async"] } base64 = "0.22" wasmtime-wasi = "41.0.3" rocksdb = { version = "0.21", features = ["multi-threaded-cf", "lz4"] } -bincode = "1.3" +bincode = { version = "2", features = ["serde"] } +chrono = "0.4" tokio-stream = "0.1.18" lru = "0.12" parking_lot = "0.12" -arrow-array = "52" -arrow-ipc = "52" -arrow-schema = "52" +arrow = { version = "55", default-features = false } +arrow-array = "55" +arrow-ipc = "55" +arrow-schema = { version = "55", features = ["serde"] } +futures = "0.3" +serde_json_path = "0.7" +xxhash-rust = { version = "0.8", features = ["xxh3"] } proctitle = "0.1" +unicase = "2.7" +petgraph = "0.7" +itertools = "0.14" +strum = { version = "0.26", features = ["derive"] } +datafusion-functions-aggregate = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} + +typify = { git = 'https://github.com/ArroyoSystems/typify.git', branch = 'arroyo' } +parquet = {git = 'https://github.com/ArroyoSystems/arrow-rs', branch = '55.2.0/parquet'} +arrow-json = {git = 'https://github.com/ArroyoSystems/arrow-rs', branch = '55.2.0/json'} +datafusion = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} +datafusion-common = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} +datafusion-execution = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} +datafusion-expr = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} +datafusion-physical-expr = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} +datafusion-physical-plan = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} +datafusion-proto = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} +datafusion-functions = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} +datafusion-functions-window = {git = 'https://github.com/ArroyoSystems/arrow-datafusion', branch = '48.0.1/arroyo'} + +sqlparser = { git = "https://github.com/FunctionStream/sqlparser-rs", branch = "0.6.0/function-sql-parser" } + +cornucopia_async = { git = "https://github.com/ArroyoSystems/cornucopia", branch = "sqlite" } +cornucopia = { git = "https://github.com/ArroyoSystems/cornucopia", branch = "sqlite" } +jiter = {git = "https://github.com/ArroyoSystems/jiter", branch = "disable_python" } + [features] default = ["incremental-cache", "python"] incremental-cache = ["wasmtime/incremental-cache"] python = [] + +[patch."https://github.com/ArroyoSystems/sqlparser-rs"] +sqlparser = { git = "https://github.com/FunctionStream/sqlparser-rs", branch = "0.6.0/function-sql-parser" } diff --git a/protocol/Cargo.toml b/protocol/Cargo.toml index fde9de52..5fa7d0f0 100644 --- a/protocol/Cargo.toml +++ b/protocol/Cargo.toml @@ -9,6 +9,7 @@ repository = "https://github.com/your-username/rust-function-stream" [dependencies] prost = "0.13" tonic = { version = "0.12", features = ["default"] } +serde = { version = "1.0", features = ["derive"] } log = "0.4" [build-dependencies] diff --git a/protocol/build.rs b/protocol/build.rs index 17e77d30..e258f456 100644 --- a/protocol/build.rs +++ b/protocol/build.rs @@ -10,54 +10,56 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::path::Path; +use std::path::{Path, PathBuf}; fn main() -> Result<(), Box> { - // Initialize logger for build script env_logger::init(); - // Create output directories in the protocol package directory - // Use CARGO_MANIFEST_DIR to get the package root directory let manifest_dir = std::env::var("CARGO_MANIFEST_DIR")?; let out_dir = Path::new(&manifest_dir).join("generated"); - let proto_file = Path::new(&manifest_dir).join("proto/function_stream.proto"); - - // Note: Cargo doesn't directly support cleaning custom directories via cargo clean. - // The generated directory will be automatically regenerated on each build if needed. - // To clean it manually, use: ./clean.sh or make clean or rm -rf protocol/generated log::info!("Generated code will be placed in: {}", out_dir.display()); - log::info!("Proto file: {}", proto_file.display()); - // Create output directories let cli_dir = out_dir.join("cli"); let service_dir = out_dir.join("service"); std::fs::create_dir_all(&cli_dir)?; std::fs::create_dir_all(&service_dir)?; - log::info!( - "Created output directories: {} and {}", - cli_dir.display(), - service_dir.display() - ); - // Generate code for CLI - only client code needed + // 1. function_stream.proto → CLI (client) and Service (server) tonic_build::configure() .out_dir(&cli_dir) - .build_client(true) // Enable client code generation - .build_server(false) // Disable server code generation + .build_client(true) + .build_server(false) .compile_protos(&["proto/function_stream.proto"], &["proto"])?; - // Generate code for Service - only server code needed tonic_build::configure() .out_dir(&service_dir) - .build_client(false) // Disable client code generation - .build_server(true) // Enable server code generation + .build_client(false) + .build_server(true) .compile_protos(&["proto/function_stream.proto"], &["proto"])?; + // 2. fs_api.proto → with file descriptor set + serde for REST/JSON + let api_dir = out_dir.join("api"); + std::fs::create_dir_all(&api_dir)?; + + let descriptor_path = + PathBuf::from(std::env::var("OUT_DIR").unwrap()).join("fs_api_descriptor.bin"); + + tonic_build::configure() + .out_dir(&api_dir) + .protoc_arg("--experimental_allow_proto3_optional") + .file_descriptor_set_path(&descriptor_path) + .type_attribute(".", "#[derive(serde::Serialize, serde::Deserialize)]") + .type_attribute(".", "#[serde(rename_all = \"camelCase\")]") + .build_client(false) + .build_server(false) + .compile_protos(&["proto/fs_api.proto"], &["proto"])?; + log::info!("Protocol Buffers code generated successfully"); println!("cargo:rustc-env=PROTO_GEN_DIR={}", out_dir.display()); - println!("cargo:rerun-if-changed={}", proto_file.display()); + println!("cargo:rerun-if-changed=proto/function_stream.proto"); + println!("cargo:rerun-if-changed=proto/fs_api.proto"); Ok(()) } diff --git a/protocol/proto/fs_api.proto b/protocol/proto/fs_api.proto new file mode 100644 index 00000000..24525583 --- /dev/null +++ b/protocol/proto/fs_api.proto @@ -0,0 +1,289 @@ +// Licensed under the Apache License, Version 2.0 +// Adapted from Arroyo's api.proto for FunctionStream + +syntax = "proto3"; +package fs_api; + +// ─────────────────────── Operators ─────────────────────── + +message ConnectorOp { + string connector = 1; + string config = 2; + string description = 3; +} + +message ProjectionOperator { + string name = 1; + FsSchema input_schema = 2; + FsSchema output_schema = 3; + repeated bytes exprs = 4; +} + +message TumblingWindowAggregateOperator { + string name = 1; + uint64 width_micros = 2; + bytes binning_function = 3; + FsSchema input_schema = 4; + FsSchema partial_schema = 5; + bytes partial_aggregation_plan = 6; + bytes final_aggregation_plan = 7; + optional bytes final_projection = 8; +} + +message SlidingWindowAggregateOperator { + string name = 1; + uint64 width_micros = 2; + uint64 slide_micros = 3; + bytes binning_function = 4; + FsSchema input_schema = 5; + FsSchema partial_schema = 6; + bytes partial_aggregation_plan = 7; + bytes final_aggregation_plan = 8; + bytes final_projection = 9; +} + +message SessionWindowAggregateOperator { + string name = 1; + uint64 gap_micros = 2; + string window_field_name = 3; + uint64 window_index = 4; + FsSchema input_schema = 5; + FsSchema unkeyed_aggregate_schema = 6; + bytes partial_aggregation_plan = 7; + bytes final_aggregation_plan = 8; +} + +message JoinOperator { + string name = 1; + FsSchema left_schema = 2; + FsSchema right_schema = 3; + FsSchema output_schema = 4; + bytes join_plan = 5; + optional uint64 ttl_micros = 6; +} + +message LookupJoinCondition { + bytes left_expr = 1; + string right_key = 2; +} + +message LookupJoinOperator { + FsSchema input_schema = 1; + FsSchema lookup_schema = 2; + ConnectorOp connector = 3; + repeated LookupJoinCondition key_exprs = 4; + JoinType join_type = 5; + optional uint64 ttl_micros = 6; + optional uint64 max_capacity_bytes = 7; +} + +message WindowFunctionOperator { + string name = 1; + FsSchema input_schema = 2; + bytes binning_function = 3; + bytes window_function_plan = 4; +} + +enum AsyncUdfOrdering { + UNORDERED = 0; + ORDERED = 1; +} + +message AsyncUdfOperator { + string name = 1; + DylibUdfConfig udf = 2; + repeated bytes arg_exprs = 3; + repeated bytes final_exprs = 4; + AsyncUdfOrdering ordering = 5; + uint32 max_concurrency = 6; + uint64 timeout_micros = 7; +} + +message UpdatingAggregateOperator { + string name = 1; + FsSchema input_schema = 2; + FsSchema final_schema = 3; + bytes aggregate_exec = 5; + bytes metadata_expr = 6; + uint64 flush_interval_micros = 7; + uint64 ttl_micros = 8; +} + +// ─────────────────────── Watermark ─────────────────────── + +message ExpressionWatermarkConfig { + uint64 period_micros = 1; + optional uint64 idle_time_micros = 2; + FsSchema input_schema = 3; + bytes expression = 4; +} + +// ─────────────────────── Windows ─────────────────────── + +message Window { + oneof window { + SlidingWindow sliding_window = 2; + TumblingWindow tumbling_window = 3; + InstantWindow instant_window = 4; + SessionWindow session_window = 5; + } +} + +message SlidingWindow { + uint64 size_micros = 1; + uint64 slide_micros = 2; +} + +message TumblingWindow { + uint64 size_micros = 1; +} + +message InstantWindow {} + +message SessionWindow { + uint64 gap_micros = 1; +} + +// ─────────────────────── Enums ─────────────────────── + +enum JoinType { + INNER = 0; + LEFT = 1; + RIGHT = 2; + FULL = 3; +} + +enum OffsetMode { + EARLIEST = 0; + LATEST = 1; +} + +enum EdgeType { + UNUSED = 0; + FORWARD = 1; + SHUFFLE = 2; + LEFT_JOIN = 3; + RIGHT_JOIN = 4; +} + +// ─────────────────── Physical Extension Nodes ─────────────────── + +message MemExecNode { + string table_name = 1; + string schema = 2; // json-encoded +} + +message UnnestExecNode { + string schema = 1; // json-encoded +} + +message DebeziumDecodeNode { + string schema = 1; // json-encoded + repeated uint64 primary_keys = 2; +} + +message DebeziumEncodeNode { + string schema = 1; // json-encoded +} + +message FsExecNode { + oneof node { + MemExecNode mem_exec = 1; + UnnestExecNode unnest_exec = 2; + DebeziumDecodeNode debezium_decode = 3; + DebeziumEncodeNode debezium_encode = 4; + } +} + +// ─────────────────── Checkpoints ─────────────────── + +enum TaskCheckpointEventType { + ALIGNMENT_STARTED = 0; + CHECKPOINT_STARTED = 1; + CHECKPOINT_OPERATOR_SETUP_FINISHED = 2; + CHECKPOINT_SYNC_FINISHED = 3; + CHECKPOINT_PRE_COMMIT = 4; +} + +message TaskCheckpointEvent { + uint64 time = 1; + TaskCheckpointEventType event_type = 2; +} + +message TaskCheckpointDetail { + uint32 subtask_index = 1; + uint64 start_time = 2; + optional uint64 finish_time = 3; + optional uint64 bytes = 4; + repeated TaskCheckpointEvent events = 5; +} + +message OperatorCheckpointDetail { + string operator_id = 1; + uint64 start_time = 2; + optional uint64 finish_time = 3; + bool has_state = 4; + optional uint64 started_metadata_write = 6; + map tasks = 5; +} + +// ─────────────────── UDF Config ─────────────────── + +message DylibUdfConfig { + string dylib_path = 1; + repeated bytes arg_types = 2; + bytes return_type = 3; + bool aggregate = 4; + bool is_async = 5; +} + +message PythonUdfConfig { + string name = 1; + repeated bytes arg_types = 2; + bytes return_type = 3; + string definition = 4; +} + +message FsProgramConfig { + map udf_dylibs = 1; + map python_udfs = 2; +} + +// ─────────────────── Arrow Program ─────────────────── + +message FsProgram { + repeated FsNode nodes = 1; + repeated FsEdge edges = 2; + FsProgramConfig program_config = 3; +} + +message FsSchema { + string arrow_schema = 1; // json-encoded Arrow Schema + uint32 timestamp_index = 2; + repeated uint32 key_indices = 3; + bool has_keys = 4; + repeated uint32 routing_key_indices = 5; + bool has_routing_keys = 6; +} + +message ChainedOperator { + string operator_id = 1; + string operator_name = 2; + bytes operator_config = 3; +} + +message FsNode { + int32 node_index = 1; + uint32 node_id = 2; + uint32 parallelism = 3; + string description = 4; + repeated ChainedOperator operators = 5; + repeated FsSchema edges = 6; +} + +message FsEdge { + int32 source = 1; + int32 target = 2; + FsSchema schema = 4; + EdgeType edge_type = 5; +} diff --git a/protocol/src/lib.rs b/protocol/src/lib.rs index b0c6da06..f924a5c6 100644 --- a/protocol/src/lib.rs +++ b/protocol/src/lib.rs @@ -10,25 +10,30 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Protocol Buffers protocol definitions for function stream -// This module exports the generated Protocol Buffers code +// ─────────────── FunctionStream Service (original) ─────────────── -// CLI module - exports client code #[path = "../generated/cli/function_stream.rs"] pub mod cli; -// Service module - exports server code #[path = "../generated/service/function_stream.rs"] pub mod service; -// Re-export commonly used types from both modules -// Data structures are the same in both, so we can re-export from either pub use cli::function_stream_service_client; - -// Re-export client-specific types pub use cli::function_stream_service_client::FunctionStreamServiceClient; - -// Re-export server-specific types pub use service::function_stream_service_server::{ FunctionStreamService, FunctionStreamServiceServer, }; + +// ─────────────── Streaming Pipeline API (fs_api.proto) ─────────────── + +pub mod grpc { + /// Serde-annotated API types for streaming operators, schemas, programs. + #[allow(clippy::all)] + pub mod api { + include!("../generated/api/fs_api.rs"); + } +} + +/// File descriptor set for fs_api.proto (for gRPC reflection / REST gateway). +pub const FS_API_FILE_DESCRIPTOR_SET: &[u8] = + tonic::include_file_descriptor_set!("fs_api_descriptor"); diff --git a/src/api/checkpoints.rs b/src/api/checkpoints.rs new file mode 100644 index 00000000..8462f311 --- /dev/null +++ b/src/api/checkpoints.rs @@ -0,0 +1,96 @@ +use crate::types::to_micros; +use serde::{Deserialize, Serialize}; +use std::time::SystemTime; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Checkpoint { + pub epoch: u32, + pub backend: String, + pub start_time: u64, + pub finish_time: Option, + pub events: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct CheckpointEventSpan { + pub start_time: u64, + pub finish_time: u64, + pub event: String, + pub description: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct SubtaskCheckpointGroup { + pub index: u32, + pub bytes: u64, + pub event_spans: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct OperatorCheckpointGroup { + pub operator_id: String, + pub bytes: u64, + pub started_metadata_write: Option, + pub finish_time: Option, + pub subtasks: Vec, +} + +#[derive(Debug, Copy, Clone, Eq, PartialEq, Serialize, Deserialize)] +pub enum JobCheckpointEventType { + Checkpointing, + CheckpointingOperators, + WritingMetadata, + Compacting, + Committing, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobCheckpointSpan { + pub event: JobCheckpointEventType, + pub start_time: u64, + pub finish_time: Option, +} + +impl JobCheckpointSpan { + pub fn now(event: JobCheckpointEventType) -> Self { + Self { + event, + start_time: to_micros(SystemTime::now()), + finish_time: None, + } + } + + pub fn finish(&mut self) { + if self.finish_time.is_none() { + self.finish_time = Some(to_micros(SystemTime::now())); + } + } +} + +impl From for CheckpointEventSpan { + fn from(value: JobCheckpointSpan) -> Self { + let description = match value.event { + JobCheckpointEventType::Checkpointing => "The entire checkpointing process", + JobCheckpointEventType::CheckpointingOperators => { + "The time spent checkpointing operator states" + } + JobCheckpointEventType::WritingMetadata => "Writing the final checkpoint metadata", + JobCheckpointEventType::Compacting => "Compacting old checkpoints", + JobCheckpointEventType::Committing => { + "Running two-phase commit for transactional connectors" + } + } + .to_string(); + + Self { + start_time: value.start_time, + finish_time: value.finish_time.unwrap_or_default(), + event: format!("{:?}", value.event), + description, + } + } +} diff --git a/src/api/connections.rs b/src/api/connections.rs new file mode 100644 index 00000000..eb69690e --- /dev/null +++ b/src/api/connections.rs @@ -0,0 +1,604 @@ +use crate::types::formats::{BadData, Format, Framing}; +use crate::types::{FsExtensionType, FsSchema}; +use datafusion::arrow::datatypes::{DataType, Field, Fields, TimeUnit}; +use serde::ser::SerializeMap; +use serde::{Deserialize, Serialize, Serializer}; +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::fmt::{Display, Formatter}; +use std::sync::Arc; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Connector { + pub id: String, + pub name: String, + pub icon: String, + pub description: String, + pub table_config: String, + pub enabled: bool, + pub source: bool, + pub sink: bool, + pub custom_schemas: bool, + pub testing: bool, + pub hidden: bool, + pub connection_config: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionProfile { + pub id: String, + pub name: String, + pub connector: String, + pub config: serde_json::Value, + pub description: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionProfilePost { + pub name: String, + pub connector: String, + pub config: serde_json::Value, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[serde(rename_all = "snake_case")] +pub enum ConnectionType { + Source, + Sink, + Lookup, +} + +impl Display for ConnectionType { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + ConnectionType::Source => write!(f, "SOURCE"), + ConnectionType::Sink => write!(f, "SINK"), + ConnectionType::Lookup => write!(f, "LOOKUP"), + } + } +} + +impl TryFrom for ConnectionType { + type Error = String; + + fn try_from(value: String) -> Result { + match value.to_lowercase().as_str() { + "source" => Ok(ConnectionType::Source), + "sink" => Ok(ConnectionType::Sink), + "lookup" => Ok(ConnectionType::Lookup), + _ => Err(format!("Invalid connection type: {value}")), + } + } +} + +// ─────────────────── Field Types ─────────────────── + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum FieldType { + Int32, + Int64, + Uint32, + Uint64, + #[serde(alias = "f32")] + Float32, + #[serde(alias = "f64")] + Float64, + Decimal128(DecimalField), + Bool, + #[serde(alias = "utf8")] + String, + #[serde(alias = "binary")] + Bytes, + Timestamp(TimestampField), + Json, + Struct(StructField), + List(ListField), +} + +impl FieldType { + pub fn sql_type(&self) -> String { + match self { + FieldType::Int32 => "INTEGER".into(), + FieldType::Int64 => "BIGINT".into(), + FieldType::Uint32 => "INTEGER UNSIGNED".into(), + FieldType::Uint64 => "BIGINT UNSIGNED".into(), + FieldType::Float32 => "FLOAT".into(), + FieldType::Float64 => "DOUBLE".into(), + FieldType::Decimal128(f) => format!("DECIMAL({}, {})", f.precision, f.scale), + FieldType::Bool => "BOOLEAN".into(), + FieldType::String => "TEXT".into(), + FieldType::Bytes => "BINARY".into(), + FieldType::Timestamp(t) => format!("TIMESTAMP({})", t.unit.precision()), + FieldType::Json => "JSON".into(), + FieldType::List(item) => format!("{}[]", item.items.field_type.sql_type()), + FieldType::Struct(StructField { fields, .. }) => { + format!( + "STRUCT <{}>", + fields + .iter() + .map(|f| format!("{} {}", f.name, f.field_type.sql_type())) + .collect::>() + .join(", ") + ) + } + } + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Default)] +#[serde(rename_all = "snake_case")] +pub enum TimestampUnit { + #[serde(alias = "s")] + Second, + #[default] + #[serde(alias = "ms")] + Millisecond, + #[serde(alias = "µs", alias = "us")] + Microsecond, + #[serde(alias = "ns")] + Nanosecond, +} + +impl TimestampUnit { + pub fn precision(&self) -> u8 { + match self { + TimestampUnit::Second => 0, + TimestampUnit::Millisecond => 3, + TimestampUnit::Microsecond => 6, + TimestampUnit::Nanosecond => 9, + } + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct TimestampField { + #[serde(default)] + pub unit: TimestampUnit, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct DecimalField { + pub precision: u8, + pub scale: i8, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct StructField { + pub fields: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct ListField { + pub items: Box, +} + +fn default_item_name() -> String { + "item".to_string() +} + +#[derive(Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct ListFieldItem { + #[serde(default = "default_item_name")] + pub name: String, + #[serde(flatten)] + pub field_type: FieldType, + #[serde(default)] + pub required: bool, + #[serde(default)] + pub sql_name: Option, +} + +impl From for Field { + fn from(value: ListFieldItem) -> Self { + SourceField { + name: value.name, + field_type: value.field_type, + required: value.required, + sql_name: None, + metadata_key: None, + } + .into() + } +} + +impl Serialize for ListFieldItem { + fn serialize(&self, s: S) -> Result + where + S: Serializer, + { + let mut f = Serializer::serialize_map(s, None)?; + f.serialize_entry("name", &self.name)?; + serialize_field_type_flat(&self.field_type, &mut f)?; + f.serialize_entry("required", &self.required)?; + f.serialize_entry("sql_name", &self.field_type.sql_type())?; + f.end() + } +} + +impl TryFrom for ListFieldItem { + type Error = String; + + fn try_from(value: Field) -> Result { + let source_field: SourceField = value.try_into()?; + Ok(Self { + name: source_field.name, + field_type: source_field.field_type, + required: source_field.required, + sql_name: None, + }) + } +} + +fn serialize_field_type_flat(ft: &FieldType, map: &mut M) -> Result<(), M::Error> { + let type_tag = match ft { + FieldType::Int32 => "int32", + FieldType::Int64 => "int64", + FieldType::Uint32 => "uint32", + FieldType::Uint64 => "uint64", + FieldType::Float32 => "float32", + FieldType::Float64 => "float64", + FieldType::Decimal128(_) => "decimal128", + FieldType::Bool => "bool", + FieldType::String => "string", + FieldType::Bytes => "bytes", + FieldType::Timestamp(_) => "timestamp", + FieldType::Json => "json", + FieldType::Struct(_) => "struct", + FieldType::List(_) => "list", + }; + map.serialize_entry("type", type_tag)?; + + match ft { + FieldType::Decimal128(d) => { + map.serialize_entry("precision", &d.precision)?; + map.serialize_entry("scale", &d.scale)?; + } + FieldType::Timestamp(t) => { + map.serialize_entry("unit", &t.unit)?; + } + FieldType::Struct(s) => { + map.serialize_entry("fields", &s.fields)?; + } + FieldType::List(l) => { + map.serialize_entry("items", &l.items)?; + } + _ => {} + } + Ok(()) +} + +// ─────────────────── Source Field ─────────────────── + +#[derive(Deserialize, Clone, Debug, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub struct SourceField { + pub name: String, + #[serde(flatten)] + pub field_type: FieldType, + #[serde(default)] + pub required: bool, + #[serde(default)] + pub sql_name: Option, + #[serde(default)] + pub metadata_key: Option, +} + +impl Serialize for SourceField { + fn serialize(&self, s: S) -> Result + where + S: Serializer, + { + let mut f = Serializer::serialize_map(s, None)?; + f.serialize_entry("name", &self.name)?; + serialize_field_type_flat(&self.field_type, &mut f)?; + f.serialize_entry("required", &self.required)?; + if let Some(metadata_key) = &self.metadata_key { + f.serialize_entry("metadata_key", metadata_key)?; + } + f.serialize_entry("sql_name", &self.field_type.sql_type())?; + f.end() + } +} + +impl From for Field { + fn from(f: SourceField) -> Self { + let (t, ext) = match f.field_type { + FieldType::Int32 => (DataType::Int32, None), + FieldType::Int64 => (DataType::Int64, None), + FieldType::Uint32 => (DataType::UInt32, None), + FieldType::Uint64 => (DataType::UInt64, None), + FieldType::Float32 => (DataType::Float32, None), + FieldType::Float64 => (DataType::Float64, None), + FieldType::Bool => (DataType::Boolean, None), + FieldType::String => (DataType::Utf8, None), + FieldType::Bytes => (DataType::Binary, None), + FieldType::Decimal128(d) => (DataType::Decimal128(d.precision, d.scale), None), + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Second, + }) => (DataType::Timestamp(TimeUnit::Second, None), None), + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Millisecond, + }) => (DataType::Timestamp(TimeUnit::Millisecond, None), None), + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Microsecond, + }) => (DataType::Timestamp(TimeUnit::Microsecond, None), None), + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Nanosecond, + }) => (DataType::Timestamp(TimeUnit::Nanosecond, None), None), + FieldType::Json => (DataType::Utf8, Some(FsExtensionType::JSON)), + FieldType::Struct(s) => ( + DataType::Struct(Fields::from( + s.fields + .into_iter() + .map(|t| t.into()) + .collect::>(), + )), + None, + ), + FieldType::List(t) => (DataType::List(Arc::new((*t.items).into())), None), + }; + + FsExtensionType::add_metadata(ext, Field::new(f.name, t, !f.required)) + } +} + +impl TryFrom for SourceField { + type Error = String; + + fn try_from(f: Field) -> Result { + let field_type = match (f.data_type(), FsExtensionType::from_map(f.metadata())) { + (DataType::Boolean, None) => FieldType::Bool, + (DataType::Int32, None) => FieldType::Int32, + (DataType::Int64, None) => FieldType::Int64, + (DataType::UInt32, None) => FieldType::Uint32, + (DataType::UInt64, None) => FieldType::Uint64, + (DataType::Float32, None) => FieldType::Float32, + (DataType::Float64, None) => FieldType::Float64, + (DataType::Decimal128(p, s), None) => FieldType::Decimal128(DecimalField { + precision: *p, + scale: *s, + }), + (DataType::Binary, None) | (DataType::LargeBinary, None) => FieldType::Bytes, + (DataType::Timestamp(TimeUnit::Second, _), None) => { + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Second, + }) + } + (DataType::Timestamp(TimeUnit::Millisecond, _), None) => { + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Millisecond, + }) + } + (DataType::Timestamp(TimeUnit::Microsecond, _), None) => { + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Microsecond, + }) + } + (DataType::Timestamp(TimeUnit::Nanosecond, _), None) => { + FieldType::Timestamp(TimestampField { + unit: TimestampUnit::Nanosecond, + }) + } + (DataType::Utf8, None) => FieldType::String, + (DataType::Utf8, Some(FsExtensionType::JSON)) => FieldType::Json, + (DataType::Struct(fields), None) => { + let fields: Result<_, String> = fields + .into_iter() + .map(|f| (**f).clone().try_into()) + .collect(); + FieldType::Struct(StructField { fields: fields? }) + } + (DataType::List(item), None) => FieldType::List(ListField { + items: Box::new((**item).clone().try_into()?), + }), + dt => return Err(format!("Unsupported data type {dt:?}")), + }; + + Ok(SourceField { + name: f.name().clone(), + field_type, + required: !f.is_nullable(), + sql_name: None, + metadata_key: None, + }) + } +} + +// ─────────────────── Schema Definitions ─────────────────── + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] +#[serde(rename_all = "snake_case", tag = "type")] +pub enum SchemaDefinition { + JsonSchema { + schema: String, + }, + ProtobufSchema { + schema: String, + #[serde(default)] + dependencies: HashMap, + }, + AvroSchema { + schema: String, + }, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionSchema { + pub format: Option, + #[serde(default)] + pub bad_data: Option, + #[serde(default)] + pub framing: Option, + #[serde(default)] + pub fields: Vec, + #[serde(default)] + pub definition: Option, + #[serde(default)] + pub inferred: Option, + #[serde(default)] + pub primary_keys: HashSet, +} + +impl ConnectionSchema { + pub fn try_new( + format: Option, + bad_data: Option, + framing: Option, + fields: Vec, + definition: Option, + inferred: Option, + primary_keys: HashSet, + ) -> anyhow::Result { + let s = ConnectionSchema { + format, + bad_data, + framing, + fields, + definition, + inferred, + primary_keys, + }; + s.validate() + } + + pub fn validate(self) -> anyhow::Result { + let non_metadata_fields: Vec<_> = self + .fields + .iter() + .filter(|f| f.metadata_key.is_none()) + .collect(); + + if let Some(Format::RawString(_)) = &self.format { + if non_metadata_fields.len() != 1 + || non_metadata_fields.first().unwrap().field_type != FieldType::String + || non_metadata_fields.first().unwrap().name != "value" + { + anyhow::bail!( + "raw_string format requires a schema with a single field called `value` of type TEXT" + ); + } + } + + if let Some(Format::Json(json_format)) = &self.format { + if json_format.unstructured + && (non_metadata_fields.len() != 1 + || non_metadata_fields.first().unwrap().field_type != FieldType::Json + || non_metadata_fields.first().unwrap().name != "value") + { + anyhow::bail!( + "json format with unstructured flag enabled requires a schema with a single field called `value` of type JSON" + ); + } + } + + Ok(self) + } + + pub fn fs_schema(&self) -> Arc { + let fields: Vec = self.fields.iter().map(|f| f.clone().into()).collect(); + Arc::new(FsSchema::from_fields(fields)) + } +} + +impl From for FsSchema { + fn from(val: ConnectionSchema) -> Self { + let fields: Vec = val.fields.into_iter().map(|f| f.into()).collect(); + FsSchema::from_fields(fields) + } +} + +// ─────────────────── Connection Table ─────────────────── + +#[derive(Serialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionTable { + #[serde(skip_serializing)] + pub id: i64, + #[serde(rename = "id")] + pub pub_id: String, + pub name: String, + pub created_at: u64, + pub connector: String, + pub connection_profile: Option, + pub table_type: ConnectionType, + pub config: serde_json::Value, + pub schema: ConnectionSchema, + pub consumers: u32, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionTablePost { + pub name: String, + pub connector: String, + pub connection_profile_id: Option, + pub config: serde_json::Value, + pub schema: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConnectionAutocompleteResp { + pub values: BTreeMap>, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct TestSourceMessage { + pub error: bool, + pub done: bool, + pub message: String, +} + +impl TestSourceMessage { + pub fn info(message: impl Into) -> Self { + Self { + error: false, + done: false, + message: message.into(), + } + } + pub fn error(message: impl Into) -> Self { + Self { + error: true, + done: false, + message: message.into(), + } + } + pub fn done(message: impl Into) -> Self { + Self { + error: false, + done: true, + message: message.into(), + } + } + pub fn fail(message: impl Into) -> Self { + Self { + error: true, + done: true, + message: message.into(), + } + } +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConfluentSchema { + pub schema: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ConfluentSchemaQueryParams { + pub endpoint: String, + pub topic: String, +} diff --git a/src/api/metrics.rs b/src/api/metrics.rs new file mode 100644 index 00000000..25d129e5 --- /dev/null +++ b/src/api/metrics.rs @@ -0,0 +1,41 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Copy, Clone, Debug, Hash, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum MetricName { + BytesRecv, + BytesSent, + MessagesRecv, + MessagesSent, + Backpressure, + TxQueueSize, + TxQueueRem, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Metric { + pub time: u64, + pub value: f64, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct SubtaskMetrics { + pub index: u32, + pub metrics: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct MetricGroup { + pub name: MetricName, + pub subtasks: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct OperatorMetricGroup { + pub node_id: u32, + pub metric_groups: Vec, +} diff --git a/src/api/mod.rs b/src/api/mod.rs new file mode 100644 index 00000000..85cbcaaa --- /dev/null +++ b/src/api/mod.rs @@ -0,0 +1,43 @@ +//! REST/RPC API types for the FunctionStream system. +//! +//! Adapted from Arroyo's `arroyo-rpc/src/api_types` and utility modules. + +pub mod checkpoints; +pub mod connections; +pub mod metrics; +pub mod pipelines; +pub mod public_ids; +pub mod schema_resolver; +pub mod udfs; +pub mod var_str; + +use serde::{Deserialize, Serialize}; + +pub use checkpoints::*; +pub use connections::{ + ConnectionProfile, ConnectionSchema, ConnectionType, Connector, FieldType, SchemaDefinition, + SourceField, +}; +pub use metrics::*; +pub use pipelines::*; +pub use udfs::*; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "camelCase")] +pub struct PaginatedCollection { + pub data: Vec, + pub has_more: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "camelCase")] +pub struct NonPaginatedCollection { + pub data: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PaginationQueryParams { + pub starting_after: Option, + pub limit: Option, +} diff --git a/src/api/pipelines.rs b/src/api/pipelines.rs new file mode 100644 index 00000000..3c77ce7a --- /dev/null +++ b/src/api/pipelines.rs @@ -0,0 +1,156 @@ +use super::udfs::Udf; +use crate::types::control::ErrorDomain; +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ValidateQueryPost { + pub query: String, + pub udfs: Option>, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct QueryValidationResult { + pub graph: Option, + pub errors: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelinePost { + pub name: String, + pub query: String, + pub udfs: Option>, + pub parallelism: u64, + pub checkpoint_interval_micros: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PreviewPost { + pub query: String, + pub udfs: Option>, + #[serde(default)] + pub enable_sinks: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelinePatch { + pub parallelism: Option, + pub checkpoint_interval_micros: Option, + pub stop: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelineRestart { + pub force: Option, + pub ignore_state: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Pipeline { + pub id: String, + pub name: String, + pub query: String, + pub udfs: Vec, + pub checkpoint_interval_micros: u64, + pub stop: StopType, + pub created_at: u64, + pub action: Option, + pub action_text: String, + pub action_in_progress: bool, + pub graph: PipelineGraph, + pub preview: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelineGraph { + pub nodes: Vec, + pub edges: Vec, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelineNode { + pub node_id: u32, + pub operator: String, + pub description: String, + pub parallelism: u32, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct PipelineEdge { + pub src_id: u32, + pub dest_id: u32, + pub key_type: String, + pub value_type: String, + pub edge_type: String, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub enum StopType { + None, + Checkpoint, + Graceful, + Immediate, + Force, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct FailureReason { + pub error: String, + pub domain: ErrorDomain, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Job { + pub id: String, + pub running_desired: bool, + pub state: String, + pub run_id: u64, + pub start_time: Option, + pub finish_time: Option, + pub tasks: Option, + pub failure_reason: Option, + pub created_at: u64, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub enum JobLogLevel { + Info, + Warn, + Error, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct JobLogMessage { + pub id: String, + pub created_at: u64, + pub operator_id: Option, + pub task_index: Option, + pub level: JobLogLevel, + pub message: String, + pub details: String, + pub error_domain: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct OutputData { + pub operator_id: String, + pub subtask_idx: u32, + pub timestamps: Vec, + pub start_id: u64, + pub batch: String, +} diff --git a/src/api/public_ids.rs b/src/api/public_ids.rs new file mode 100644 index 00000000..15a9f72e --- /dev/null +++ b/src/api/public_ids.rs @@ -0,0 +1,57 @@ +use std::time::{SystemTime, UNIX_EPOCH}; + +const ID_LENGTH: usize = 10; + +const ALPHABET: &[u8; 62] = b"0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; + +pub enum IdTypes { + ApiKey, + ConnectionProfile, + Schema, + Pipeline, + JobConfig, + Checkpoint, + JobStatus, + ClusterInfo, + JobLogMessage, + ConnectionTable, + ConnectionTablePipeline, + Udf, +} + +/// Generates a unique identifier with a type-specific prefix. +/// +/// Uses a simple time + random approach instead of nanoid to avoid an extra dependency. +pub fn generate_id(id_type: IdTypes) -> String { + let prefix = match id_type { + IdTypes::ApiKey => "ak", + IdTypes::ConnectionProfile => "cp", + IdTypes::Schema => "sch", + IdTypes::Pipeline => "pl", + IdTypes::JobConfig => "job", + IdTypes::Checkpoint => "chk", + IdTypes::JobStatus => "js", + IdTypes::ClusterInfo => "ci", + IdTypes::JobLogMessage => "jlm", + IdTypes::ConnectionTable => "ct", + IdTypes::ConnectionTablePipeline => "ctp", + IdTypes::Udf => "udf", + }; + + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(); + + let mut id = String::with_capacity(ID_LENGTH); + let mut seed = nanos; + for _ in 0..ID_LENGTH { + seed ^= seed + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + let idx = (seed % ALPHABET.len() as u128) as usize; + id.push(ALPHABET[idx] as char); + } + + format!("{prefix}_{id}") +} diff --git a/src/api/schema_resolver.rs b/src/api/schema_resolver.rs new file mode 100644 index 00000000..a9124900 --- /dev/null +++ b/src/api/schema_resolver.rs @@ -0,0 +1,82 @@ +use async_trait::async_trait; + +/// Trait for resolving schemas by ID (e.g., from a schema registry). +#[async_trait] +pub trait SchemaResolver: Send { + async fn resolve_schema(&self, id: u32) -> Result, String>; +} + +/// A resolver that always fails — used when no schema registry is configured. +pub struct FailingSchemaResolver; + +impl Default for FailingSchemaResolver { + fn default() -> Self { + Self + } +} + +#[async_trait] +impl SchemaResolver for FailingSchemaResolver { + async fn resolve_schema(&self, id: u32) -> Result, String> { + Err(format!( + "Schema with id {id} not available, and no schema registry configured" + )) + } +} + +/// A resolver that returns a fixed schema for a known ID. +pub struct FixedSchemaResolver { + id: u32, + schema: String, +} + +impl FixedSchemaResolver { + pub fn new(id: u32, schema: String) -> Self { + FixedSchemaResolver { id, schema } + } +} + +#[async_trait] +impl SchemaResolver for FixedSchemaResolver { + async fn resolve_schema(&self, id: u32) -> Result, String> { + if id == self.id { + Ok(Some(self.schema.clone())) + } else { + Err(format!("Unexpected schema id {}, expected {}", id, self.id)) + } + } +} + +/// A caching wrapper around any `SchemaResolver`. +pub struct CachingSchemaResolver { + inner: R, + cache: tokio::sync::RwLock>, +} + +impl CachingSchemaResolver { + pub fn new(inner: R) -> Self { + Self { + inner, + cache: tokio::sync::RwLock::new(std::collections::HashMap::new()), + } + } +} + +#[async_trait] +impl SchemaResolver for CachingSchemaResolver { + async fn resolve_schema(&self, id: u32) -> Result, String> { + { + let cache = self.cache.read().await; + if let Some(schema) = cache.get(&id) { + return Ok(Some(schema.clone())); + } + } + + let result = self.inner.resolve_schema(id).await?; + if let Some(ref schema) = result { + let mut cache = self.cache.write().await; + cache.insert(id, schema.clone()); + } + Ok(result) + } +} diff --git a/src/api/udfs.rs b/src/api/udfs.rs new file mode 100644 index 00000000..41085168 --- /dev/null +++ b/src/api/udfs.rs @@ -0,0 +1,56 @@ +use serde::{Deserialize, Serialize}; + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct Udf { + pub definition: String, + #[serde(default)] + pub language: UdfLanguage, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct ValidateUdfPost { + pub definition: String, + #[serde(default)] + pub language: UdfLanguage, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct UdfValidationResult { + pub udf_name: Option, + pub errors: Vec, +} + +#[derive(Serialize, Deserialize, Copy, Clone, Debug, Default, Eq, PartialEq)] +#[serde(rename_all = "snake_case")] +pub enum UdfLanguage { + Python, + #[default] + Rust, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct UdfPost { + pub prefix: String, + #[serde(default)] + pub language: UdfLanguage, + pub definition: String, + pub description: Option, +} + +#[derive(Serialize, Deserialize, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub struct GlobalUdf { + pub id: String, + pub prefix: String, + pub name: String, + pub language: UdfLanguage, + pub created_at: u64, + pub updated_at: u64, + pub definition: String, + pub description: Option, + pub dylib_url: Option, +} diff --git a/src/api/var_str.rs b/src/api/var_str.rs new file mode 100644 index 00000000..c4256e38 --- /dev/null +++ b/src/api/var_str.rs @@ -0,0 +1,79 @@ +use serde::{Deserialize, Serialize}; +use std::env; + +/// A string that may contain `{{ VAR }}` placeholders for environment variable substitution. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(transparent)] +pub struct VarStr { + raw_val: String, +} + +impl VarStr { + pub fn new(raw_val: String) -> Self { + VarStr { raw_val } + } + + pub fn raw(&self) -> &str { + &self.raw_val + } + + /// Substitute `{{ VAR_NAME }}` patterns with the corresponding environment variable values. + pub fn sub_env_vars(&self) -> anyhow::Result { + let mut result = self.raw_val.clone(); + let mut start = 0; + + while let Some(open) = result[start..].find("{{") { + let open_abs = start + open; + let Some(close) = result[open_abs..].find("}}") else { + break; + }; + let close_abs = open_abs + close; + + let var_name = result[open_abs + 2..close_abs].trim(); + if var_name.is_empty() { + start = close_abs + 2; + continue; + } + + match env::var(var_name) { + Ok(value) => { + let full_match = &result[open_abs..close_abs + 2]; + let full_match_owned = full_match.to_string(); + result = result.replacen(&full_match_owned, &value, 1); + start = open_abs + value.len(); + } + Err(_) => { + anyhow::bail!("Environment variable {} not found", var_name); + } + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_no_placeholders() { + let input = "This is a test string with no placeholders"; + assert_eq!( + VarStr::new(input.to_string()).sub_env_vars().unwrap(), + input + ); + } + + #[test] + fn test_with_placeholders() { + unsafe { env::set_var("FS_TEST_VAR", "environment variable") }; + let input = "This is a {{ FS_TEST_VAR }}"; + let expected = "This is a environment variable"; + assert_eq!( + VarStr::new(input.to_string()).sub_env_vars().unwrap(), + expected + ); + unsafe { env::remove_var("FS_TEST_VAR") }; + } +} diff --git a/src/coordinator/analyze/analyzer.rs b/src/coordinator/analyze/analyzer.rs index 30552191..297d0de2 100644 --- a/src/coordinator/analyze/analyzer.rs +++ b/src/coordinator/analyze/analyzer.rs @@ -13,8 +13,9 @@ use super::Analysis; use crate::coordinator::execution_context::ExecutionContext; use crate::coordinator::statement::{ - CreateFunction, CreatePythonFunction, DropFunction, ShowFunctions, StartFunction, Statement, - StatementVisitor, StatementVisitorContext, StatementVisitorResult, StopFunction, + CreateFunction, CreatePythonFunction, CreateTable, DropFunction, InsertStatement, + ShowFunctions, StartFunction, Statement, StatementVisitor, StatementVisitorContext, + StatementVisitorResult, StopFunction, StreamingSql, }; use std::fmt; @@ -115,4 +116,29 @@ impl StatementVisitor for Analyzer<'_> { ) -> StatementVisitorResult { StatementVisitorResult::Analyze(Box::new(stmt.clone())) } + + fn visit_create_table( + &self, + stmt: &CreateTable, + _context: &StatementVisitorContext, + ) -> StatementVisitorResult { + StatementVisitorResult::Analyze(Box::new(CreateTable::new(stmt.statement.clone()))) + } + + fn visit_insert_statement( + &self, + stmt: &InsertStatement, + _context: &StatementVisitorContext, + ) -> StatementVisitorResult { + StatementVisitorResult::Analyze(Box::new(InsertStatement::new(stmt.statement.clone()))) + } + + fn visit_streaming_sql( + &self, + stmt: &StreamingSql, + _context: &StatementVisitorContext, + ) -> StatementVisitorResult { + // TODO: add semantic analysis for streaming SQL (schema validation, etc.) + StatementVisitorResult::Analyze(Box::new(StreamingSql::new(stmt.statement.clone()))) + } } diff --git a/src/coordinator/coordinator.rs b/src/coordinator/coordinator.rs index 4ad766d5..378c670b 100644 --- a/src/coordinator/coordinator.rs +++ b/src/coordinator/coordinator.rs @@ -20,6 +20,7 @@ use crate::coordinator::execution::Executor; use crate::coordinator::plan::{LogicalPlanVisitor, LogicalPlanner, PlanNode}; use crate::coordinator::statement::Statement; use crate::runtime::taskexecutor::TaskManager; +use crate::sql::planner::StreamSchemaProvider; use super::execution_context::ExecutionContext; @@ -90,7 +91,8 @@ impl Coordinator { } fn step_build_logical_plan(&self, analysis: &Analysis) -> Result> { - let visitor = LogicalPlanVisitor::new(); + let schema_provider = StreamSchemaProvider::new(); + let visitor = LogicalPlanVisitor::new(schema_provider); let plan = visitor.visit(analysis); Ok(plan) } diff --git a/src/coordinator/execution/executor.rs b/src/coordinator/execution/executor.rs index 7e44217e..dbc76923 100644 --- a/src/coordinator/execution/executor.rs +++ b/src/coordinator/execution/executor.rs @@ -12,8 +12,9 @@ use crate::coordinator::dataset::{ExecuteResult, ShowFunctionsResult, empty_record_batch}; use crate::coordinator::plan::{ - CreateFunctionPlan, CreatePythonFunctionPlan, DropFunctionPlan, PlanNode, PlanVisitor, - PlanVisitorContext, PlanVisitorResult, ShowFunctionsPlan, StartFunctionPlan, StopFunctionPlan, + CreateFunctionPlan, CreatePythonFunctionPlan, CreateTablePlan, DropFunctionPlan, + InsertStatementPlan, PlanNode, PlanVisitor, PlanVisitorContext, PlanVisitorResult, + ShowFunctionsPlan, StartFunctionPlan, StopFunctionPlan, StreamingSqlPlan, }; use crate::coordinator::statement::{ConfigSource, FunctionSource}; use crate::runtime::taskexecutor::TaskManager; @@ -200,4 +201,43 @@ impl PlanVisitor for Executor { PlanVisitorResult::Execute(result) } + + fn visit_create_table_plan( + &self, + plan: &CreateTablePlan, + _context: &PlanVisitorContext, + ) -> PlanVisitorResult { + // TODO: register table in catalog and execute DDL + let result = Err(ExecuteError::Internal(format!( + "CREATE TABLE execution not yet implemented. LogicalPlan:\n{}", + plan.logical_plan.display_indent() + ))); + PlanVisitorResult::Execute(result) + } + + fn visit_insert_statement_plan( + &self, + plan: &InsertStatementPlan, + _context: &PlanVisitorContext, + ) -> PlanVisitorResult { + // TODO: start streaming pipeline for INSERT / anonymous query + let result = Err(ExecuteError::Internal(format!( + "INSERT statement execution not yet implemented. LogicalPlan:\n{}", + plan.logical_plan.display_indent() + ))); + PlanVisitorResult::Execute(result) + } + + fn visit_streaming_sql_plan( + &self, + plan: &StreamingSqlPlan, + _context: &PlanVisitorContext, + ) -> PlanVisitorResult { + // TODO: apply rewrite_plan for streaming transformations, then execute + let result = Err(ExecuteError::Internal(format!( + "Streaming SQL execution not yet implemented. LogicalPlan:\n{}", + plan.logical_plan.display_indent() + ))); + PlanVisitorResult::Execute(result) + } } diff --git a/src/coordinator/mod.rs b/src/coordinator/mod.rs index 0b94d4bf..51b93ca0 100644 --- a/src/coordinator/mod.rs +++ b/src/coordinator/mod.rs @@ -22,6 +22,6 @@ mod statement; pub use coordinator::Coordinator; pub use dataset::{DataSet, ShowFunctionsResult}; pub use statement::{ - CreateFunction, CreatePythonFunction, DropFunction, PythonModule, ShowFunctions, StartFunction, - Statement, StopFunction, + CreateFunction, CreatePythonFunction, CreateTable, DropFunction, InsertStatement, PythonModule, + ShowFunctions, StartFunction, Statement, StopFunction, StreamingSql, }; diff --git a/src/sql/parser/mod.rs b/src/coordinator/plan/create_table_plan.rs similarity index 52% rename from src/sql/parser/mod.rs rename to src/coordinator/plan/create_table_plan.rs index 11f4b18e..450c8813 100644 --- a/src/sql/parser/mod.rs +++ b/src/coordinator/plan/create_table_plan.rs @@ -10,33 +10,23 @@ // See the License for the specific language governing permissions and // limitations under the License. -mod sql_parser; +use datafusion::logical_expr::LogicalPlan; -pub use sql_parser::SqlParser; +use super::{PlanNode, PlanVisitor, PlanVisitorContext, PlanVisitorResult}; #[derive(Debug)] -pub struct ParseError { - pub message: String, +pub struct CreateTablePlan { + pub logical_plan: LogicalPlan, } -impl std::fmt::Display for ParseError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "Parse error: {}", self.message) +impl CreateTablePlan { + pub fn new(logical_plan: LogicalPlan) -> Self { + Self { logical_plan } } } -impl std::error::Error for ParseError {} - -impl From for ParseError { - fn from(message: String) -> Self { - ParseError { message } - } -} - -impl ParseError { - pub fn new(message: impl Into) -> Self { - Self { - message: message.into(), - } +impl PlanNode for CreateTablePlan { + fn accept(&self, visitor: &dyn PlanVisitor, context: &PlanVisitorContext) -> PlanVisitorResult { + visitor.visit_create_table_plan(self, context) } } diff --git a/src/coordinator/plan/insert_statement_plan.rs b/src/coordinator/plan/insert_statement_plan.rs new file mode 100644 index 00000000..e96a2772 --- /dev/null +++ b/src/coordinator/plan/insert_statement_plan.rs @@ -0,0 +1,32 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::logical_expr::LogicalPlan; + +use super::{PlanNode, PlanVisitor, PlanVisitorContext, PlanVisitorResult}; + +#[derive(Debug)] +pub struct InsertStatementPlan { + pub logical_plan: LogicalPlan, +} + +impl InsertStatementPlan { + pub fn new(logical_plan: LogicalPlan) -> Self { + Self { logical_plan } + } +} + +impl PlanNode for InsertStatementPlan { + fn accept(&self, visitor: &dyn PlanVisitor, context: &PlanVisitorContext) -> PlanVisitorResult { + visitor.visit_insert_statement_plan(self, context) + } +} diff --git a/src/coordinator/plan/logical_plan_visitor.rs b/src/coordinator/plan/logical_plan_visitor.rs index 536fec37..fde7f35a 100644 --- a/src/coordinator/plan/logical_plan_visitor.rs +++ b/src/coordinator/plan/logical_plan_visitor.rs @@ -10,22 +10,28 @@ // See the License for the specific language governing permissions and // limitations under the License. +use tracing::debug; + use crate::coordinator::analyze::analysis::Analysis; use crate::coordinator::plan::{ - CreateFunctionPlan, CreatePythonFunctionPlan, DropFunctionPlan, PlanNode, ShowFunctionsPlan, - StartFunctionPlan, StopFunctionPlan, + CreateFunctionPlan, CreatePythonFunctionPlan, CreateTablePlan, DropFunctionPlan, + InsertStatementPlan, PlanNode, ShowFunctionsPlan, StartFunctionPlan, StopFunctionPlan, + StreamingSqlPlan, }; use crate::coordinator::statement::{ - CreateFunction, CreatePythonFunction, DropFunction, ShowFunctions, StartFunction, - StatementVisitor, StatementVisitorContext, StatementVisitorResult, StopFunction, + CreateFunction, CreatePythonFunction, CreateTable, DropFunction, InsertStatement, + ShowFunctions, StartFunction, StatementVisitor, StatementVisitorContext, + StatementVisitorResult, StopFunction, StreamingSql, }; +use crate::sql::planner::StreamSchemaProvider; -#[derive(Debug, Default)] -pub struct LogicalPlanVisitor; +pub struct LogicalPlanVisitor { + schema_provider: StreamSchemaProvider, +} impl LogicalPlanVisitor { - pub fn new() -> Self { - Self + pub fn new(schema_provider: StreamSchemaProvider) -> Self { + Self { schema_provider } } pub fn visit(&self, analysis: &Analysis) -> Box { @@ -51,7 +57,6 @@ impl StatementVisitor for LogicalPlanVisitor { let config_source = stmt.get_config_source().cloned(); let extra_props = stmt.get_extra_properties().clone(); - // Name will be read from config file during execution StatementVisitorResult::Plan(Box::new(CreateFunctionPlan::new( function_source, config_source, @@ -106,4 +111,58 @@ impl StatementVisitor for LogicalPlanVisitor { config_content, ))) } + + fn visit_create_table( + &self, + stmt: &CreateTable, + _context: &StatementVisitorContext, + ) -> StatementVisitorResult { + let sql_to_rel = datafusion::sql::planner::SqlToRel::new(&self.schema_provider); + + match sql_to_rel.sql_statement_to_plan(stmt.statement.clone()) { + Ok(plan) => { + debug!("Create table plan:\n{}", plan.display_graphviz()); + StatementVisitorResult::Plan(Box::new(CreateTablePlan::new(plan))) + } + Err(e) => { + panic!("Failed to convert CREATE TABLE to logical plan: {e}"); + } + } + } + + fn visit_insert_statement( + &self, + stmt: &InsertStatement, + _context: &StatementVisitorContext, + ) -> StatementVisitorResult { + let sql_to_rel = datafusion::sql::planner::SqlToRel::new(&self.schema_provider); + + match sql_to_rel.sql_statement_to_plan(stmt.statement.clone()) { + Ok(plan) => { + debug!("Insert statement plan:\n{}", plan.display_graphviz()); + StatementVisitorResult::Plan(Box::new(InsertStatementPlan::new(plan))) + } + Err(e) => { + panic!("Failed to convert INSERT statement to logical plan: {e}"); + } + } + } + + fn visit_streaming_sql( + &self, + stmt: &StreamingSql, + _context: &StatementVisitorContext, + ) -> StatementVisitorResult { + let sql_to_rel = datafusion::sql::planner::SqlToRel::new(&self.schema_provider); + + match sql_to_rel.sql_statement_to_plan(stmt.statement.clone()) { + Ok(plan) => { + debug!("Logical plan:\n{}", plan.display_graphviz()); + StatementVisitorResult::Plan(Box::new(StreamingSqlPlan::new(plan))) + } + Err(e) => { + panic!("Failed to convert SQL statement to logical plan: {e}"); + } + } + } } diff --git a/src/coordinator/plan/mod.rs b/src/coordinator/plan/mod.rs index 9aa403b5..3d36ec16 100644 --- a/src/coordinator/plan/mod.rs +++ b/src/coordinator/plan/mod.rs @@ -12,22 +12,28 @@ mod create_function_plan; mod create_python_function_plan; +mod create_table_plan; mod drop_function_plan; +mod insert_statement_plan; mod logical_plan_visitor; mod optimizer; mod show_functions_plan; mod start_function_plan; mod stop_function_plan; +mod streaming_sql_plan; mod visitor; pub use create_function_plan::CreateFunctionPlan; pub use create_python_function_plan::CreatePythonFunctionPlan; +pub use create_table_plan::CreateTablePlan; pub use drop_function_plan::DropFunctionPlan; +pub use insert_statement_plan::InsertStatementPlan; pub use logical_plan_visitor::LogicalPlanVisitor; pub use optimizer::LogicalPlanner; pub use show_functions_plan::ShowFunctionsPlan; pub use start_function_plan::StartFunctionPlan; pub use stop_function_plan::StopFunctionPlan; +pub use streaming_sql_plan::StreamingSqlPlan; pub use visitor::{PlanVisitor, PlanVisitorContext, PlanVisitorResult}; use std::fmt; diff --git a/src/coordinator/plan/streaming_sql_plan.rs b/src/coordinator/plan/streaming_sql_plan.rs new file mode 100644 index 00000000..607420a8 --- /dev/null +++ b/src/coordinator/plan/streaming_sql_plan.rs @@ -0,0 +1,32 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::logical_expr::LogicalPlan; + +use super::{PlanNode, PlanVisitor, PlanVisitorContext, PlanVisitorResult}; + +#[derive(Debug)] +pub struct StreamingSqlPlan { + pub logical_plan: LogicalPlan, +} + +impl StreamingSqlPlan { + pub fn new(logical_plan: LogicalPlan) -> Self { + Self { logical_plan } + } +} + +impl PlanNode for StreamingSqlPlan { + fn accept(&self, visitor: &dyn PlanVisitor, context: &PlanVisitorContext) -> PlanVisitorResult { + visitor.visit_streaming_sql_plan(self, context) + } +} diff --git a/src/coordinator/plan/visitor.rs b/src/coordinator/plan/visitor.rs index 44059c67..e3911a8b 100644 --- a/src/coordinator/plan/visitor.rs +++ b/src/coordinator/plan/visitor.rs @@ -11,8 +11,8 @@ // limitations under the License. use super::{ - CreateFunctionPlan, CreatePythonFunctionPlan, DropFunctionPlan, ShowFunctionsPlan, - StartFunctionPlan, StopFunctionPlan, + CreateFunctionPlan, CreatePythonFunctionPlan, CreateTablePlan, DropFunctionPlan, + InsertStatementPlan, ShowFunctionsPlan, StartFunctionPlan, StopFunctionPlan, StreamingSqlPlan, }; /// Context passed to PlanVisitor methods @@ -84,4 +84,22 @@ pub trait PlanVisitor { plan: &CreatePythonFunctionPlan, context: &PlanVisitorContext, ) -> PlanVisitorResult; + + fn visit_create_table_plan( + &self, + plan: &CreateTablePlan, + context: &PlanVisitorContext, + ) -> PlanVisitorResult; + + fn visit_insert_statement_plan( + &self, + plan: &InsertStatementPlan, + context: &PlanVisitorContext, + ) -> PlanVisitorResult; + + fn visit_streaming_sql_plan( + &self, + plan: &StreamingSqlPlan, + context: &PlanVisitorContext, + ) -> PlanVisitorResult; } diff --git a/src/coordinator/statement/create_table.rs b/src/coordinator/statement/create_table.rs new file mode 100644 index 00000000..8aa16bf0 --- /dev/null +++ b/src/coordinator/statement/create_table.rs @@ -0,0 +1,40 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::sql::sqlparser::ast::Statement as DFStatement; + +use super::{Statement, StatementVisitor, StatementVisitorContext, StatementVisitorResult}; + +/// Represents a CREATE TABLE or CREATE VIEW statement. +/// +/// This wraps the raw SQL AST node so the coordinator pipeline can +/// distinguish table/view creation from other streaming SQL operations. +#[derive(Debug)] +pub struct CreateTable { + pub statement: DFStatement, +} + +impl CreateTable { + pub fn new(statement: DFStatement) -> Self { + Self { statement } + } +} + +impl Statement for CreateTable { + fn accept( + &self, + visitor: &dyn StatementVisitor, + context: &StatementVisitorContext, + ) -> StatementVisitorResult { + visitor.visit_create_table(self, context) + } +} diff --git a/src/coordinator/statement/insert_statement.rs b/src/coordinator/statement/insert_statement.rs new file mode 100644 index 00000000..45785251 --- /dev/null +++ b/src/coordinator/statement/insert_statement.rs @@ -0,0 +1,41 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::sql::sqlparser::ast::Statement as DFStatement; + +use super::{Statement, StatementVisitor, StatementVisitorContext, StatementVisitorResult}; + +/// Represents an INSERT INTO or standalone SELECT/query statement. +/// +/// In the streaming SQL context, both INSERT INTO (writing to a sink) +/// and standalone SELECT (anonymous computation) are treated as +/// data-producing operations that feed into the streaming pipeline. +#[derive(Debug)] +pub struct InsertStatement { + pub statement: DFStatement, +} + +impl InsertStatement { + pub fn new(statement: DFStatement) -> Self { + Self { statement } + } +} + +impl Statement for InsertStatement { + fn accept( + &self, + visitor: &dyn StatementVisitor, + context: &StatementVisitorContext, + ) -> StatementVisitorResult { + visitor.visit_insert_statement(self, context) + } +} diff --git a/src/coordinator/statement/mod.rs b/src/coordinator/statement/mod.rs index f887209c..7628b94b 100644 --- a/src/coordinator/statement/mod.rs +++ b/src/coordinator/statement/mod.rs @@ -12,18 +12,24 @@ mod create_function; mod create_python_function; +mod create_table; mod drop_function; +mod insert_statement; mod show_functions; mod start_function; mod stop_function; +mod streaming_sql; mod visitor; pub use create_function::{ConfigSource, CreateFunction, FunctionSource}; pub use create_python_function::{CreatePythonFunction, PythonModule}; +pub use create_table::CreateTable; pub use drop_function::DropFunction; +pub use insert_statement::InsertStatement; pub use show_functions::ShowFunctions; pub use start_function::StartFunction; pub use stop_function::StopFunction; +pub use streaming_sql::StreamingSql; pub use visitor::{StatementVisitor, StatementVisitorContext, StatementVisitorResult}; use std::fmt; diff --git a/src/coordinator/statement/streaming_sql.rs b/src/coordinator/statement/streaming_sql.rs new file mode 100644 index 00000000..1aa49205 --- /dev/null +++ b/src/coordinator/statement/streaming_sql.rs @@ -0,0 +1,39 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::sql::sqlparser::ast::Statement as DFStatement; + +use super::{Statement, StatementVisitor, StatementVisitorContext, StatementVisitorResult}; + +/// Wraps a DataFusion SQL statement (SELECT, INSERT, CREATE TABLE, etc.) +/// so it can flow through the same Statement → StatementVisitor pipeline +/// as FunctionStream DDL commands. +#[derive(Debug)] +pub struct StreamingSql { + pub statement: DFStatement, +} + +impl StreamingSql { + pub fn new(statement: DFStatement) -> Self { + Self { statement } + } +} + +impl Statement for StreamingSql { + fn accept( + &self, + visitor: &dyn StatementVisitor, + context: &StatementVisitorContext, + ) -> StatementVisitorResult { + visitor.visit_streaming_sql(self, context) + } +} diff --git a/src/coordinator/statement/visitor.rs b/src/coordinator/statement/visitor.rs index 13ce2cfc..8de6ffe2 100644 --- a/src/coordinator/statement/visitor.rs +++ b/src/coordinator/statement/visitor.rs @@ -11,7 +11,8 @@ // limitations under the License. use super::{ - CreateFunction, CreatePythonFunction, DropFunction, ShowFunctions, StartFunction, StopFunction, + CreateFunction, CreatePythonFunction, CreateTable, DropFunction, InsertStatement, + ShowFunctions, StartFunction, StopFunction, StreamingSql, }; use crate::coordinator::plan::PlanNode; use crate::coordinator::statement::Statement; @@ -87,4 +88,22 @@ pub trait StatementVisitor { stmt: &CreatePythonFunction, context: &StatementVisitorContext, ) -> StatementVisitorResult; + + fn visit_create_table( + &self, + stmt: &CreateTable, + context: &StatementVisitorContext, + ) -> StatementVisitorResult; + + fn visit_insert_statement( + &self, + stmt: &InsertStatement, + context: &StatementVisitorContext, + ) -> StatementVisitorResult; + + fn visit_streaming_sql( + &self, + stmt: &StreamingSql, + context: &StatementVisitorContext, + ) -> StatementVisitorResult; } diff --git a/src/datastream/logical.rs b/src/datastream/logical.rs new file mode 100644 index 00000000..a6486760 --- /dev/null +++ b/src/datastream/logical.rs @@ -0,0 +1,317 @@ +use itertools::Itertools; + +use crate::datastream::optimizers::Optimizer; +use crate::sql::types::StreamSchema; +use datafusion::arrow::datatypes::DataType; +use petgraph::Direction; +use petgraph::dot::Dot; +use petgraph::graph::DiGraph; +use std::collections::{HashMap, HashSet}; +use std::fmt::{Debug, Display, Formatter}; +use std::sync::Arc; +use strum::{Display, EnumString}; + +#[derive(Clone, Copy, Debug, Eq, PartialEq, EnumString, Display)] +pub enum OperatorName { + ExpressionWatermark, + ArrowValue, + ArrowKey, + Projection, + AsyncUdf, + Join, + InstantJoin, + LookupJoin, + WindowFunction, + TumblingWindowAggregate, + SlidingWindowAggregate, + SessionWindowAggregate, + UpdatingAggregate, + ConnectorSource, + ConnectorSink, +} + +#[derive(Copy, Clone, Debug, Eq, PartialEq, PartialOrd, Ord)] +pub enum LogicalEdgeType { + Forward, + Shuffle, + LeftJoin, + RightJoin, +} + +impl Display for LogicalEdgeType { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self { + LogicalEdgeType::Forward => write!(f, "→"), + LogicalEdgeType::Shuffle => write!(f, "⤨"), + LogicalEdgeType::LeftJoin => write!(f, "-[left]⤨"), + LogicalEdgeType::RightJoin => write!(f, "-[right]⤨"), + } + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct LogicalEdge { + pub edge_type: LogicalEdgeType, + pub schema: Arc, +} + +impl LogicalEdge { + pub fn new(edge_type: LogicalEdgeType, schema: StreamSchema) -> Self { + LogicalEdge { + edge_type, + schema: Arc::new(schema), + } + } + + pub fn project_all(edge_type: LogicalEdgeType, schema: StreamSchema) -> Self { + LogicalEdge { + edge_type, + schema: Arc::new(schema), + } + } +} + +#[derive(Clone, Debug)] +pub struct ChainedLogicalOperator { + pub operator_id: String, + pub operator_name: OperatorName, + pub operator_config: Vec, +} + +#[derive(Clone, Debug)] +pub struct OperatorChain { + pub(crate) operators: Vec, + pub(crate) edges: Vec>, +} + +impl OperatorChain { + pub fn new(operator: ChainedLogicalOperator) -> Self { + Self { + operators: vec![operator], + edges: vec![], + } + } + + pub fn iter( + &self, + ) -> impl Iterator>)> { + self.operators + .iter() + .zip_longest(self.edges.iter()) + .map(|e| e.left_and_right()) + .map(|(l, r)| (l.unwrap(), r)) + } + + pub fn iter_mut( + &mut self, + ) -> impl Iterator>)> { + self.operators + .iter_mut() + .zip_longest(self.edges.iter()) + .map(|e| e.left_and_right()) + .map(|(l, r)| (l.unwrap(), r)) + } + + pub fn first(&self) -> &ChainedLogicalOperator { + &self.operators[0] + } + + pub fn len(&self) -> usize { + self.operators.len() + } + + pub fn is_empty(&self) -> bool { + self.operators.is_empty() + } + + pub fn is_source(&self) -> bool { + self.operators[0].operator_name == OperatorName::ConnectorSource + } + + pub fn is_sink(&self) -> bool { + self.operators[0].operator_name == OperatorName::ConnectorSink + } +} + +#[derive(Clone)] +pub struct LogicalNode { + pub node_id: u32, + pub description: String, + pub operator_chain: OperatorChain, + pub parallelism: usize, +} + +impl LogicalNode { + pub fn single( + id: u32, + operator_id: String, + name: OperatorName, + config: Vec, + description: String, + parallelism: usize, + ) -> Self { + Self { + node_id: id, + description, + operator_chain: OperatorChain { + operators: vec![ChainedLogicalOperator { + operator_id, + operator_name: name, + operator_config: config, + }], + edges: vec![], + }, + parallelism, + } + } +} + +impl Display for LogicalNode { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.description) + } +} + +impl Debug for LogicalNode { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{}[{}]", + self.operator_chain + .operators + .iter() + .map(|op| op.operator_id.clone()) + .collect::>() + .join(" -> "), + self.parallelism + ) + } +} + +pub type LogicalGraph = DiGraph; + +#[derive(Clone, Debug, Eq, PartialEq, Hash, PartialOrd)] +pub struct DylibUdfConfig { + pub dylib_path: String, + pub arg_types: Vec, + pub return_type: DataType, + pub aggregate: bool, + pub is_async: bool, +} + +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +pub struct PythonUdfConfig { + pub arg_types: Vec, + pub return_type: DataType, + pub name: Arc, + pub definition: Arc, +} + +#[derive(Clone, Debug, Default)] +pub struct ProgramConfig { + pub udf_dylibs: HashMap, + pub python_udfs: HashMap, +} + +#[derive(Clone, Debug, Default)] +pub struct LogicalProgram { + pub graph: LogicalGraph, + pub program_config: ProgramConfig, +} + +impl LogicalProgram { + pub fn new(graph: LogicalGraph, program_config: ProgramConfig) -> Self { + Self { + graph, + program_config, + } + } + + pub fn optimize(&mut self, optimizer: &dyn Optimizer) { + optimizer.optimize(&mut self.graph); + } + + pub fn update_parallelism(&mut self, overrides: &HashMap) { + for node in self.graph.node_weights_mut() { + if let Some(p) = overrides.get(&node.node_id) { + node.parallelism = *p; + } + } + } + + pub fn dot(&self) -> String { + format!("{:?}", Dot::with_config(&self.graph, &[])) + } + + pub fn task_count(&self) -> usize { + self.graph.node_weights().map(|nw| nw.parallelism).sum() + } + + pub fn sources(&self) -> HashSet { + self.graph + .externals(Direction::Incoming) + .map(|t| self.graph.node_weight(t).unwrap().node_id) + .collect() + } + + pub fn tasks_per_operator(&self) -> HashMap { + let mut tasks_per_operator = HashMap::new(); + for node in self.graph.node_weights() { + for op in &node.operator_chain.operators { + tasks_per_operator.insert(op.operator_id.clone(), node.parallelism); + } + } + tasks_per_operator + } + + pub fn operator_names_by_id(&self) -> HashMap { + let mut m = HashMap::new(); + for node in self.graph.node_weights() { + for op in &node.operator_chain.operators { + m.insert(op.operator_id.clone(), op.operator_name.to_string()); + } + } + m + } + + pub fn tasks_per_node(&self) -> HashMap { + let mut tasks_per_node = HashMap::new(); + for node in self.graph.node_weights() { + tasks_per_node.insert(node.node_id, node.parallelism); + } + tasks_per_node + } + + pub fn features(&self) -> HashSet { + let mut s = HashSet::new(); + for n in self.graph.node_weights() { + for t in &n.operator_chain.operators { + let feature = match &t.operator_name { + OperatorName::AsyncUdf => "async-udf".to_string(), + OperatorName::ExpressionWatermark + | OperatorName::ArrowValue + | OperatorName::ArrowKey + | OperatorName::Projection => continue, + OperatorName::Join => "join-with-expiration".to_string(), + OperatorName::InstantJoin => "windowed-join".to_string(), + OperatorName::WindowFunction => "sql-window-function".to_string(), + OperatorName::LookupJoin => "lookup-join".to_string(), + OperatorName::TumblingWindowAggregate => { + "sql-tumbling-window-aggregate".to_string() + } + OperatorName::SlidingWindowAggregate => { + "sql-sliding-window-aggregate".to_string() + } + OperatorName::SessionWindowAggregate => { + "sql-session-window-aggregate".to_string() + } + OperatorName::UpdatingAggregate => "sql-updating-aggregate".to_string(), + OperatorName::ConnectorSource => "connector-source".to_string(), + OperatorName::ConnectorSink => "connector-sink".to_string(), + }; + s.insert(feature); + } + } + s + } +} diff --git a/src/datastream/mod.rs b/src/datastream/mod.rs new file mode 100644 index 00000000..994a96b4 --- /dev/null +++ b/src/datastream/mod.rs @@ -0,0 +1,2 @@ +pub mod logical; +pub mod optimizers; diff --git a/src/datastream/optimizers.rs b/src/datastream/optimizers.rs new file mode 100644 index 00000000..2d258aff --- /dev/null +++ b/src/datastream/optimizers.rs @@ -0,0 +1,100 @@ +use crate::datastream::logical::{LogicalEdgeType, LogicalGraph}; +use petgraph::prelude::*; +use petgraph::visit::NodeRef; +use std::mem; + +pub trait Optimizer { + fn optimize_once(&self, plan: &mut LogicalGraph) -> bool; + + fn optimize(&self, plan: &mut LogicalGraph) { + loop { + if !self.optimize_once(plan) { + break; + } + } + } +} + +pub struct ChainingOptimizer {} + +fn remove_in_place(graph: &mut DiGraph, node: NodeIndex) { + let incoming = graph.edges_directed(node, Incoming).next().unwrap(); + + let parent = incoming.source().id(); + let incoming = incoming.id(); + graph.remove_edge(incoming); + + let outgoing: Vec<_> = graph + .edges_directed(node, Outgoing) + .map(|e| (e.id(), e.target().id())) + .collect(); + + for (edge, target) in outgoing { + let weight = graph.remove_edge(edge).unwrap(); + graph.add_edge(parent, target, weight); + } + + graph.remove_node(node); +} + +impl Optimizer for ChainingOptimizer { + fn optimize_once(&self, plan: &mut LogicalGraph) -> bool { + let node_indices: Vec = plan.node_indices().collect(); + + for &node_idx in &node_indices { + let cur = plan.node_weight(node_idx).unwrap(); + + if cur.operator_chain.is_source() { + continue; + } + + let mut successors = plan.edges_directed(node_idx, Outgoing).collect::>(); + + if successors.len() != 1 { + continue; + } + + let edge = successors.remove(0); + let edge_type = edge.weight().edge_type; + + if edge_type != LogicalEdgeType::Forward { + continue; + } + + let successor_idx = edge.target(); + + let successor_node = plan.node_weight(successor_idx).unwrap(); + + if cur.parallelism != successor_node.parallelism + || successor_node.operator_chain.is_sink() + { + continue; + } + + if plan.edges_directed(successor_idx, Incoming).count() > 1 { + continue; + } + + let mut new_cur = cur.clone(); + + new_cur.description = format!("{} -> {}", cur.description, successor_node.description); + + new_cur + .operator_chain + .operators + .extend(successor_node.operator_chain.operators.clone()); + + new_cur + .operator_chain + .edges + .push(edge.weight().schema.clone()); + + mem::swap(&mut new_cur, plan.node_weight_mut(node_idx).unwrap()); + + remove_in_place(plan, successor_idx); + return true; + } + + false + } +} diff --git a/src/lib.rs b/src/lib.rs index a6bb4d28..a41536c5 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,10 +14,13 @@ #![allow(dead_code)] +pub mod api; pub mod config; pub mod coordinator; +pub mod datastream; pub mod logging; pub mod runtime; pub mod server; pub mod sql; pub mod storage; +pub mod types; diff --git a/src/main.rs b/src/main.rs index 562b1526..29935d62 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,13 +12,16 @@ #![allow(dead_code)] +mod api; mod config; mod coordinator; +mod datastream; mod logging; mod runtime; mod server; mod sql; mod storage; +mod types; use anyhow::{Context, Result}; use std::thread; diff --git a/src/server/handler.rs b/src/server/handler.rs index 4721a5a1..45b0cd07 100644 --- a/src/server/handler.rs +++ b/src/server/handler.rs @@ -29,7 +29,7 @@ use crate::coordinator::{ CreateFunction, CreatePythonFunction, DataSet, DropFunction, ShowFunctions, ShowFunctionsResult, StartFunction, Statement, StopFunction, }; -use crate::sql::SqlParser; +use crate::sql::planner::parse::parse_sql; pub struct FunctionStreamServiceImpl { coordinator: Arc, @@ -70,10 +70,10 @@ impl FunctionStreamService for FunctionStreamServiceImpl { let req = request.into_inner(); let parse_start = Instant::now(); - let stmt = match SqlParser::parse(&req.sql) { - Ok(stmt) => { + let parsed = match parse_sql(&req.sql) { + Ok(parsed) => { log::debug!("SQL parsed in {}ms", parse_start.elapsed().as_millis()); - stmt + parsed } Err(e) => { return Ok(TonicResponse::new(Self::build_response( @@ -85,7 +85,7 @@ impl FunctionStreamService for FunctionStreamServiceImpl { }; let exec_start = Instant::now(); - let result = self.coordinator.execute(stmt.as_ref()); + let result = self.coordinator.execute(parsed.as_ref()); log::debug!( "Coordinator execution finished in {}ms", exec_start.elapsed().as_millis() diff --git a/src/sql/catalog/connector.rs b/src/sql/catalog/connector.rs new file mode 100644 index 00000000..01176d47 --- /dev/null +++ b/src/sql/catalog/connector.rs @@ -0,0 +1,59 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::fmt; + +/// Describes the role of a connection in the streaming pipeline. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub enum ConnectionType { + Source, + Sink, + Lookup, +} + +impl fmt::Display for ConnectionType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ConnectionType::Source => write!(f, "source"), + ConnectionType::Sink => write!(f, "sink"), + ConnectionType::Lookup => write!(f, "lookup"), + } + } +} + +/// A connector operation that describes how to interact with an external system. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ConnectorOp { + pub connector: String, + pub config: String, + pub description: String, +} + +impl ConnectorOp { + pub fn new(connector: impl Into, config: impl Into) -> Self { + let connector = connector.into(); + let description = connector.clone(); + Self { + connector, + config: config.into(), + description, + } + } +} + +/// Configuration for a connection profile (e.g., Kafka broker, Pulsar endpoint). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ConnectionProfile { + pub name: String, + pub connector: String, + pub config: std::collections::HashMap, +} diff --git a/src/sql/catalog/connector_table.rs b/src/sql/catalog/connector_table.rs new file mode 100644 index 00000000..8dae1745 --- /dev/null +++ b/src/sql/catalog/connector_table.rs @@ -0,0 +1,199 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; +use std::time::Duration; + +use datafusion::arrow::datatypes::{FieldRef, Schema}; +use datafusion::common::{Result, plan_err}; +use datafusion::logical_expr::Expr; + +use super::connector::{ConnectionType, ConnectorOp}; +use super::field_spec::FieldSpec; +use crate::multifield_partial_ord; +use crate::sql::types::ProcessingMode; + +/// Represents a table backed by an external connector (e.g., Kafka, Pulsar, NATS). +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct ConnectorTable { + pub id: Option, + pub connector: String, + pub name: String, + pub connection_type: ConnectionType, + pub fields: Vec, + pub config: String, + pub description: String, + pub event_time_field: Option, + pub watermark_field: Option, + pub idle_time: Option, + pub primary_keys: Arc>, + pub inferred_fields: Option>, + pub partition_exprs: Arc>>, +} + +multifield_partial_ord!( + ConnectorTable, + id, + connector, + name, + connection_type, + config, + description, + event_time_field, + watermark_field, + idle_time, + primary_keys +); + +impl ConnectorTable { + pub fn new( + name: impl Into, + connector: impl Into, + connection_type: ConnectionType, + ) -> Self { + Self { + id: None, + connector: connector.into(), + name: name.into(), + connection_type, + fields: Vec::new(), + config: String::new(), + description: String::new(), + event_time_field: None, + watermark_field: None, + idle_time: None, + primary_keys: Arc::new(Vec::new()), + inferred_fields: None, + partition_exprs: Arc::new(None), + } + } + + pub fn has_virtual_fields(&self) -> bool { + self.fields.iter().any(|f| f.is_virtual()) + } + + pub fn is_updating(&self) -> bool { + // TODO: check format for debezium/update mode + false + } + + pub fn physical_schema(&self) -> Schema { + Schema::new( + self.fields + .iter() + .filter(|f| !f.is_virtual()) + .map(|f| f.field().clone()) + .collect::>(), + ) + } + + pub fn connector_op(&self) -> ConnectorOp { + ConnectorOp { + connector: self.connector.clone(), + config: self.config.clone(), + description: self.description.clone(), + } + } + + pub fn processing_mode(&self) -> ProcessingMode { + if self.is_updating() { + ProcessingMode::Update + } else { + ProcessingMode::Append + } + } + + pub fn timestamp_override(&self) -> Result> { + if let Some(field_name) = &self.event_time_field { + if self.is_updating() { + return plan_err!("can't use event_time_field with update mode"); + } + let _field = self.get_time_field(field_name)?; + Ok(Some(Expr::Column(datafusion::common::Column::from_name( + field_name, + )))) + } else { + Ok(None) + } + } + + fn get_time_field(&self, field_name: &str) -> Result<&FieldSpec> { + self.fields + .iter() + .find(|f| { + f.field().name() == field_name + && matches!( + f.field().data_type(), + datafusion::arrow::datatypes::DataType::Timestamp(..) + ) + }) + .ok_or_else(|| { + datafusion::error::DataFusionError::Plan(format!( + "field {field_name} not found or not a timestamp" + )) + }) + } + + pub fn watermark_column(&self) -> Result> { + if let Some(field_name) = &self.watermark_field { + let _field = self.get_time_field(field_name)?; + Ok(Some(Expr::Column(datafusion::common::Column::from_name( + field_name, + )))) + } else { + Ok(None) + } + } + + pub fn as_sql_source(&self) -> Result { + match self.connection_type { + ConnectionType::Source => {} + ConnectionType::Sink | ConnectionType::Lookup => { + return plan_err!("cannot read from sink"); + } + } + + if self.is_updating() && self.has_virtual_fields() { + return plan_err!("can't read from a source with virtual fields and update mode"); + } + + let timestamp_override = self.timestamp_override()?; + let watermark_column = self.watermark_column()?; + + Ok(SourceOperator { + name: self.name.clone(), + connector_op: self.connector_op(), + processing_mode: self.processing_mode(), + idle_time: self.idle_time, + struct_fields: self + .fields + .iter() + .filter(|f| !f.is_virtual()) + .map(|f| Arc::new(f.field().clone())) + .collect(), + timestamp_override, + watermark_column, + }) + } +} + +/// A fully resolved source operator ready for execution graph construction. +#[derive(Debug, Clone)] +pub struct SourceOperator { + pub name: String, + pub connector_op: ConnectorOp, + pub processing_mode: ProcessingMode, + pub idle_time: Option, + pub struct_fields: Vec, + pub timestamp_override: Option, + pub watermark_column: Option, +} diff --git a/src/sql/catalog/field_spec.rs b/src/sql/catalog/field_spec.rs new file mode 100644 index 00000000..2fe8a50e --- /dev/null +++ b/src/sql/catalog/field_spec.rs @@ -0,0 +1,52 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::arrow::datatypes::Field; +use datafusion::logical_expr::Expr; + +/// Describes how a field in a connector table should be interpreted. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum FieldSpec { + /// A regular struct field that maps to a column in the data. + Struct(Field), + /// A metadata field extracted from message metadata (e.g., Kafka headers). + Metadata { field: Field, key: String }, + /// A virtual field computed from an expression over other fields. + Virtual { field: Field, expression: Box }, +} + +impl FieldSpec { + pub fn is_virtual(&self) -> bool { + matches!(self, FieldSpec::Virtual { .. }) + } + + pub fn field(&self) -> &Field { + match self { + FieldSpec::Struct(f) => f, + FieldSpec::Metadata { field, .. } => field, + FieldSpec::Virtual { field, .. } => field, + } + } + + pub fn metadata_key(&self) -> Option<&str> { + match self { + FieldSpec::Metadata { key, .. } => Some(key.as_str()), + _ => None, + } + } +} + +impl From for FieldSpec { + fn from(value: Field) -> Self { + FieldSpec::Struct(value) + } +} diff --git a/src/sql/catalog/insert.rs b/src/sql/catalog/insert.rs new file mode 100644 index 00000000..a4a3814a --- /dev/null +++ b/src/sql/catalog/insert.rs @@ -0,0 +1,55 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::common::Result; +use datafusion::logical_expr::{DmlStatement, LogicalPlan, WriteOp}; +use datafusion::sql::sqlparser::ast::Statement; + +use super::optimizer::produce_optimized_plan; +use crate::sql::planner::StreamSchemaProvider; + +/// Represents an INSERT operation in a streaming SQL pipeline. +#[derive(Debug)] +pub enum Insert { + /// Insert into a named sink table. + InsertQuery { + sink_name: String, + logical_plan: LogicalPlan, + }, + /// An anonymous query (no explicit INSERT target). + Anonymous { logical_plan: LogicalPlan }, +} + +impl Insert { + pub fn try_from_statement( + statement: &Statement, + schema_provider: &StreamSchemaProvider, + ) -> Result { + let logical_plan = produce_optimized_plan(statement, schema_provider)?; + + match &logical_plan { + LogicalPlan::Dml(DmlStatement { + table_name, + op: WriteOp::Insert(_), + input, + .. + }) => { + let sink_name = table_name.to_string(); + Ok(Insert::InsertQuery { + sink_name, + logical_plan: (**input).clone(), + }) + } + _ => Ok(Insert::Anonymous { logical_plan }), + } + } +} diff --git a/src/sql/catalog/mod.rs b/src/sql/catalog/mod.rs new file mode 100644 index 00000000..39c7bfcd --- /dev/null +++ b/src/sql/catalog/mod.rs @@ -0,0 +1,25 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +pub mod connector; +pub mod connector_table; +pub mod field_spec; +pub mod insert; +pub mod optimizer; +pub mod table; +pub mod utils; + +pub use connector::{ConnectionType, ConnectorOp}; +pub use connector_table::{ConnectorTable, SourceOperator}; +pub use field_spec::FieldSpec; +pub use insert::Insert; +pub use table::Table; diff --git a/src/sql/catalog/optimizer.rs b/src/sql/catalog/optimizer.rs new file mode 100644 index 00000000..15abe61e --- /dev/null +++ b/src/sql/catalog/optimizer.rs @@ -0,0 +1,95 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use datafusion::common::Result; +use datafusion::common::config::ConfigOptions; +use datafusion::logical_expr::LogicalPlan; +use datafusion::optimizer::OptimizerContext; +use datafusion::optimizer::OptimizerRule; +use datafusion::optimizer::common_subexpr_eliminate::CommonSubexprEliminate; +use datafusion::optimizer::decorrelate_lateral_join::DecorrelateLateralJoin; +use datafusion::optimizer::decorrelate_predicate_subquery::DecorrelatePredicateSubquery; +use datafusion::optimizer::eliminate_cross_join::EliminateCrossJoin; +use datafusion::optimizer::eliminate_duplicated_expr::EliminateDuplicatedExpr; +use datafusion::optimizer::eliminate_filter::EliminateFilter; +use datafusion::optimizer::eliminate_group_by_constant::EliminateGroupByConstant; +use datafusion::optimizer::eliminate_join::EliminateJoin; +use datafusion::optimizer::eliminate_limit::EliminateLimit; +use datafusion::optimizer::eliminate_nested_union::EliminateNestedUnion; +use datafusion::optimizer::eliminate_one_union::EliminateOneUnion; +use datafusion::optimizer::eliminate_outer_join::EliminateOuterJoin; +use datafusion::optimizer::extract_equijoin_predicate::ExtractEquijoinPredicate; +use datafusion::optimizer::filter_null_join_keys::FilterNullJoinKeys; +use datafusion::optimizer::optimizer::Optimizer; +use datafusion::optimizer::propagate_empty_relation::PropagateEmptyRelation; +use datafusion::optimizer::push_down_filter::PushDownFilter; +use datafusion::optimizer::push_down_limit::PushDownLimit; +use datafusion::optimizer::replace_distinct_aggregate::ReplaceDistinctWithAggregate; +use datafusion::optimizer::scalar_subquery_to_join::ScalarSubqueryToJoin; +use datafusion::optimizer::simplify_expressions::SimplifyExpressions; +use datafusion::sql::planner::SqlToRel; +use datafusion::sql::sqlparser::ast::Statement; + +use crate::sql::planner::StreamSchemaProvider; + +/// Converts a SQL statement into an optimized DataFusion logical plan. +/// +/// Applies the DataFusion analyzer followed by a curated set of optimizer rules +/// suitable for streaming SQL (some rules like OptimizeProjections are excluded +/// because they can drop event-time calculation fields). +pub fn produce_optimized_plan( + statement: &Statement, + schema_provider: &StreamSchemaProvider, +) -> Result { + let sql_to_rel = SqlToRel::new(schema_provider); + let plan = sql_to_rel.sql_statement_to_plan(statement.clone())?; + + let analyzed_plan = schema_provider.analyzer.execute_and_check( + plan, + &ConfigOptions::default(), + |_plan, _rule| {}, + )?; + + let rules: Vec> = vec![ + Arc::new(EliminateNestedUnion::new()), + Arc::new(SimplifyExpressions::new()), + Arc::new(ReplaceDistinctWithAggregate::new()), + Arc::new(EliminateJoin::new()), + Arc::new(DecorrelatePredicateSubquery::new()), + Arc::new(ScalarSubqueryToJoin::new()), + Arc::new(DecorrelateLateralJoin::new()), + Arc::new(ExtractEquijoinPredicate::new()), + Arc::new(EliminateDuplicatedExpr::new()), + Arc::new(EliminateFilter::new()), + Arc::new(EliminateCrossJoin::new()), + Arc::new(EliminateLimit::new()), + Arc::new(PropagateEmptyRelation::new()), + Arc::new(EliminateOneUnion::new()), + Arc::new(FilterNullJoinKeys::default()), + Arc::new(EliminateOuterJoin::new()), + Arc::new(PushDownLimit::new()), + Arc::new(PushDownFilter::new()), + Arc::new(EliminateGroupByConstant::new()), + Arc::new(CommonSubexprEliminate::new()), + ]; + + let optimizer = Optimizer::with_rules(rules); + let optimized = optimizer.optimize( + analyzed_plan, + &OptimizerContext::default(), + |_plan, _rule| {}, + )?; + + Ok(optimized) +} diff --git a/src/sql/catalog/table.rs b/src/sql/catalog/table.rs new file mode 100644 index 00000000..b1d60028 --- /dev/null +++ b/src/sql/catalog/table.rs @@ -0,0 +1,202 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use datafusion::arrow::datatypes::FieldRef; +use datafusion::common::{Result, plan_err}; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use datafusion::sql::sqlparser::ast::Statement; + +use super::connector_table::ConnectorTable; +use super::optimizer::produce_optimized_plan; +use crate::sql::planner::StreamSchemaProvider; +use crate::sql::planner::extension::remote_table::RemoteTableExtension; +use crate::sql::planner::plan::rewrite_plan; +use crate::sql::types::DFField; + +/// Represents all table types in the FunctionStream SQL catalog. +#[allow(clippy::enum_variant_names)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum Table { + /// A lookup table backed by an external connector. + LookupTable(ConnectorTable), + /// A source/sink table backed by an external connector. + ConnectorTable(ConnectorTable), + /// An in-memory table with an optional logical plan (for views). + MemoryTable { + name: String, + fields: Vec, + logical_plan: Option, + }, + /// A table defined by a query (CREATE VIEW / CREATE TABLE AS SELECT). + TableFromQuery { + name: String, + logical_plan: LogicalPlan, + }, + /// A preview sink for debugging/inspection. + PreviewSink { logical_plan: LogicalPlan }, +} + +impl Table { + /// Try to construct a Table from a CREATE TABLE or CREATE VIEW statement. + pub fn try_from_statement( + statement: &Statement, + schema_provider: &StreamSchemaProvider, + ) -> Result> { + use datafusion::logical_expr::{CreateMemoryTable, CreateView, DdlStatement}; + use datafusion::sql::sqlparser::ast::CreateTable; + + if let Statement::CreateTable(CreateTable { + name, + columns, + query: None, + .. + }) = statement + { + let name = name.to_string(); + + if columns.is_empty() { + return plan_err!("CREATE TABLE requires at least one column"); + } + + let fields: Vec = columns + .iter() + .map(|col| { + let data_type = crate::sql::types::convert_data_type(&col.data_type) + .map(|(dt, _)| dt) + .unwrap_or(datafusion::arrow::datatypes::DataType::Utf8); + let nullable = !col.options.iter().any(|opt| { + matches!( + opt.option, + datafusion::sql::sqlparser::ast::ColumnOption::NotNull + ) + }); + Arc::new(datafusion::arrow::datatypes::Field::new( + col.name.value.clone(), + data_type, + nullable, + )) + }) + .collect(); + + return Ok(Some(Table::MemoryTable { + name, + fields, + logical_plan: None, + })); + } + + match produce_optimized_plan(statement, schema_provider) { + Ok(LogicalPlan::Ddl(DdlStatement::CreateView(CreateView { name, input, .. }))) + | Ok(LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(CreateMemoryTable { + name, + input, + .. + }))) => { + let rewritten = rewrite_plan(input.as_ref().clone(), schema_provider)?; + let schema = rewritten.schema().clone(); + let remote = RemoteTableExtension { + input: rewritten, + name: name.to_owned(), + schema, + materialize: true, + }; + Ok(Some(Table::TableFromQuery { + name: name.to_string(), + logical_plan: LogicalPlan::Extension(Extension { + node: Arc::new(remote), + }), + })) + } + _ => Ok(None), + } + } + + pub fn name(&self) -> &str { + match self { + Table::MemoryTable { name, .. } | Table::TableFromQuery { name, .. } => name.as_str(), + Table::ConnectorTable(c) | Table::LookupTable(c) => c.name.as_str(), + Table::PreviewSink { .. } => "preview", + } + } + + pub fn get_fields(&self) -> Vec { + match self { + Table::MemoryTable { fields, .. } => fields.clone(), + Table::ConnectorTable(ConnectorTable { + fields, + inferred_fields, + .. + }) + | Table::LookupTable(ConnectorTable { + fields, + inferred_fields, + .. + }) => inferred_fields.clone().unwrap_or_else(|| { + fields + .iter() + .map(|field| field.field().clone().into()) + .collect() + }), + Table::TableFromQuery { logical_plan, .. } => { + logical_plan.schema().fields().iter().cloned().collect() + } + Table::PreviewSink { logical_plan } => { + logical_plan.schema().fields().iter().cloned().collect() + } + } + } + + pub fn set_inferred_fields(&mut self, fields: Vec) -> Result<()> { + let Table::ConnectorTable(t) = self else { + return Ok(()); + }; + + if !t.fields.is_empty() { + return Ok(()); + } + + if let Some(existing) = &t.inferred_fields { + let matches = existing.len() == fields.len() + && existing + .iter() + .zip(&fields) + .all(|(a, b)| a.name() == b.name() && a.data_type() == b.data_type()); + + if !matches { + return plan_err!("all inserts into a table must share the same schema"); + } + } + + let fields: Vec<_> = fields.into_iter().map(|f| f.field().clone()).collect(); + t.inferred_fields.replace(fields); + + Ok(()) + } + + pub fn connector_op(&self) -> Result { + match self { + Table::ConnectorTable(c) | Table::LookupTable(c) => Ok(c.connector_op()), + Table::MemoryTable { .. } => plan_err!("can't write to a memory table"), + Table::TableFromQuery { .. } => plan_err!("can't write to a query-defined table"), + Table::PreviewSink { .. } => Ok(super::connector::ConnectorOp::new("preview", "")), + } + } + + pub fn partition_exprs(&self) -> Option<&Vec> { + match self { + Table::ConnectorTable(c) => (*c.partition_exprs).as_ref(), + _ => None, + } + } +} diff --git a/src/sql/catalog/utils.rs b/src/sql/catalog/utils.rs new file mode 100644 index 00000000..c0b8a7d0 --- /dev/null +++ b/src/sql/catalog/utils.rs @@ -0,0 +1,78 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; +use datafusion::common::{DFSchema, DFSchemaRef, Result as DFResult, TableReference}; + +use crate::sql::types::{DFField, TIMESTAMP_FIELD}; + +/// Returns the Arrow struct type for a window (start, end) pair. +pub fn window_arrow_struct() -> DataType { + DataType::Struct( + vec![ + Arc::new(Field::new( + "start", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )), + Arc::new(Field::new( + "end", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )), + ] + .into(), + ) +} + +/// Adds a `_timestamp` field to a DFSchema if it doesn't already have one. +pub fn add_timestamp_field( + schema: DFSchemaRef, + qualifier: Option, +) -> DFResult { + if has_timestamp_field(&schema) { + return Ok(schema); + } + + let timestamp_field = DFField::new( + qualifier, + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ); + Ok(Arc::new(schema.join(&DFSchema::new_with_metadata( + vec![timestamp_field.into()], + HashMap::new(), + )?)?)) +} + +/// Checks whether a DFSchema contains a `_timestamp` field. +pub fn has_timestamp_field(schema: &DFSchemaRef) -> bool { + schema + .fields() + .iter() + .any(|field| field.name() == TIMESTAMP_FIELD) +} + +/// Adds a `_timestamp` field to an Arrow Schema, returning a new SchemaRef. +pub fn add_timestamp_field_arrow(schema: Schema) -> SchemaRef { + let mut fields = schema.fields().to_vec(); + fields.push(Arc::new(Field::new( + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ))); + Arc::new(Schema::new(fields)) +} diff --git a/src/sql/functions/mod.rs b/src/sql/functions/mod.rs new file mode 100644 index 00000000..84d3c7d4 --- /dev/null +++ b/src/sql/functions/mod.rs @@ -0,0 +1,600 @@ +use crate::sql::planner::StreamSchemaProvider; +use datafusion::arrow::array::{ + Array, ArrayRef, StringArray, UnionArray, + builder::{FixedSizeBinaryBuilder, ListBuilder, StringBuilder}, + cast::{AsArray, as_string_array}, + types::{Float64Type, Int64Type}, +}; +use datafusion::arrow::datatypes::{DataType, Field, UnionFields, UnionMode}; +use datafusion::arrow::row::{RowConverter, SortField}; +use datafusion::common::{DataFusionError, ScalarValue}; +use datafusion::common::{Result, TableReference}; +use datafusion::execution::FunctionRegistry; +use datafusion::logical_expr::expr::{Alias, ScalarFunction}; +use datafusion::logical_expr::{ + ColumnarValue, LogicalPlan, Projection, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignature, Volatility, create_udf, +}; +use datafusion::prelude::{Expr, col}; +use serde_json_path::JsonPath; +use std::any::Any; +use std::collections::HashMap; +use std::fmt::{Debug, Write}; +use std::sync::{Arc, OnceLock}; + +const SERIALIZE_JSON_UNION: &str = "serialize_json_union"; + +/// Borrowed from DataFusion +/// +/// Creates a singleton `ScalarUDF` of the `$UDF` function named `$GNAME` and a +/// function named `$NAME` which returns that function named $NAME. +/// +/// This is used to ensure creating the list of `ScalarUDF` only happens once. +#[macro_export] +macro_rules! make_udf_function { + ($UDF:ty, $GNAME:ident, $NAME:ident) => { + /// Singleton instance of the function + static $GNAME: std::sync::OnceLock> = + std::sync::OnceLock::new(); + + /// Return a [`ScalarUDF`] for [`$UDF`] + /// + /// [`ScalarUDF`]: datafusion_expr::ScalarUDF + pub fn $NAME() -> std::sync::Arc { + $GNAME + .get_or_init(|| { + std::sync::Arc::new(datafusion::logical_expr::ScalarUDF::new_from_impl( + <$UDF>::default(), + )) + }) + .clone() + } + }; +} + +make_udf_function!(MultiHashFunction, MULTI_HASH, multi_hash); + +pub fn register_all(registry: &mut dyn FunctionRegistry) { + registry + .register_udf(Arc::new(create_udf( + "get_first_json_object", + vec![DataType::Utf8, DataType::Utf8], + DataType::Utf8, + Volatility::Immutable, + Arc::new(get_first_json_object), + ))) + .unwrap(); + + registry + .register_udf(Arc::new(create_udf( + "extract_json", + vec![DataType::Utf8, DataType::Utf8], + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))), + Volatility::Immutable, + Arc::new(extract_json), + ))) + .unwrap(); + + registry + .register_udf(Arc::new(create_udf( + "extract_json_string", + vec![DataType::Utf8, DataType::Utf8], + DataType::Utf8, + Volatility::Immutable, + Arc::new(extract_json_string), + ))) + .unwrap(); + + registry + .register_udf(Arc::new(create_udf( + SERIALIZE_JSON_UNION, + vec![DataType::Union(union_fields(), UnionMode::Sparse)], + DataType::Utf8, + Volatility::Immutable, + Arc::new(serialize_json_union), + ))) + .unwrap(); + + registry.register_udf(multi_hash()).unwrap(); +} + +fn parse_path(name: &str, path: &ScalarValue) -> Result> { + let path = match path { + ScalarValue::Utf8(Some(s)) => JsonPath::parse(s) + .map_err(|e| DataFusionError::Execution(format!("Invalid json path '{s}': {e:?}")))?, + ScalarValue::Utf8(None) => { + return Err(DataFusionError::Execution(format!( + "The path argument to {name} cannot be null" + ))); + } + _ => { + return Err(DataFusionError::Execution(format!( + "The path argument to {name} must be of type TEXT" + ))); + } + }; + + Ok(Arc::new(path)) +} + +// Hash function that can take any number of arguments and produces a fast (non-cryptographic) +// 128-bit hash from their string representations +#[derive(Debug)] +pub struct MultiHashFunction { + signature: Signature, +} + +impl MultiHashFunction { + pub fn invoke(&self, args: &[ColumnarValue]) -> Result { + let mut hasher = xxhash_rust::xxh3::Xxh3::new(); + + let all_scalar = args.iter().all(|a| matches!(a, ColumnarValue::Scalar(_))); + + let length = args + .iter() + .map(|t| match t { + ColumnarValue::Scalar(_) => 1, + ColumnarValue::Array(a) => a.len(), + }) + .max() + .ok_or_else(|| { + DataFusionError::Plan("multi_hash must have at least one argument".to_string()) + })?; + + let row_builder = RowConverter::new( + args.iter() + .map(|t| SortField::new(t.data_type().clone())) + .collect(), + )?; + + let arrays = args + .iter() + .map(|c| c.clone().into_array(length)) + .collect::>>()?; + let rows = row_builder.convert_columns(&arrays)?; + + if all_scalar { + hasher.update(rows.row(0).as_ref()); + let result = hasher.digest128().to_be_bytes().to_vec(); + hasher.reset(); + Ok(ColumnarValue::Scalar(ScalarValue::FixedSizeBinary( + size_of::() as i32, + Some(result), + ))) + } else { + let mut builder = + FixedSizeBinaryBuilder::with_capacity(length, size_of::() as i32); + + for row in rows.iter() { + hasher.update(row.as_ref()); + builder.append_value(hasher.digest128().to_be_bytes())?; + hasher.reset(); + } + + Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + } + } +} + +impl Default for MultiHashFunction { + fn default() -> Self { + Self { + signature: Signature::new(TypeSignature::VariadicAny, Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for MultiHashFunction { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "multi_hash" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::FixedSizeBinary(size_of::() as i32)) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + self.invoke(&args.args) + } +} + +fn json_function( + name: &str, + f: F, + to_scalar: ToS, + args: &[ColumnarValue], +) -> Result +where + ArrayT: Array + FromIterator> + 'static, + F: Fn(serde_json::Value, &JsonPath) -> Option, + ToS: Fn(Option) -> ScalarValue, +{ + assert_eq!(args.len(), 2); + Ok(match (&args[0], &args[1]) { + (ColumnarValue::Array(values), ColumnarValue::Scalar(path)) => { + let path = parse_path(name, path)?; + let vs = as_string_array(values); + ColumnarValue::Array(Arc::new( + vs.iter() + .map(|s| s.and_then(|s| f(serde_json::from_str(s).ok()?, &path))) + .collect::(), + ) as ArrayRef) + } + (ColumnarValue::Scalar(value), ColumnarValue::Scalar(path)) => { + let path = parse_path(name, path)?; + let ScalarValue::Utf8(value) = value else { + return Err(DataFusionError::Execution(format!( + "The value argument to {name} must be of type TEXT" + ))); + }; + + let result = value + .as_ref() + .and_then(|v| f(serde_json::from_str(v).ok()?, &path)); + ColumnarValue::Scalar(to_scalar(result)) + } + _ => { + return Err(DataFusionError::Execution( + "The path argument to {name} must be a literal".to_string(), + )); + } + }) +} + +pub fn extract_json(args: &[ColumnarValue]) -> Result { + assert_eq!(args.len(), 2); + + let inner = |s, path: &JsonPath| { + Some( + path.query(&serde_json::from_str(s).ok()?) + .iter() + .map(|v| Some(v.to_string())) + .collect::>>(), + ) + }; + + Ok(match (&args[0], &args[1]) { + (ColumnarValue::Array(values), ColumnarValue::Scalar(path)) => { + let path = parse_path("extract_json", path)?; + let values = as_string_array(values); + + let mut builder = ListBuilder::with_capacity(StringBuilder::new(), values.len()); + + let queried = values.iter().map(|s| s.and_then(|s| inner(s, &path))); + + for v in queried { + builder.append_option(v); + } + + ColumnarValue::Array(Arc::new(builder.finish())) + } + (ColumnarValue::Scalar(value), ColumnarValue::Scalar(path)) => { + let path = parse_path("extract_json", path)?; + let ScalarValue::Utf8(v) = value else { + return Err(DataFusionError::Execution( + "The value argument to extract_json must be of type TEXT".to_string(), + )); + }; + + let mut builder = ListBuilder::with_capacity(StringBuilder::new(), 1); + let result = v.as_ref().and_then(|s| inner(s, &path)); + builder.append_option(result); + + ColumnarValue::Scalar(ScalarValue::List(Arc::new(builder.finish()))) + } + _ => { + return Err(DataFusionError::Execution( + "The path argument to extract_json must be a literal".to_string(), + )); + } + }) +} + +pub fn get_first_json_object(args: &[ColumnarValue]) -> Result { + json_function::( + "get_first_json_object", + |s, path| path.query(&s).first().map(|v| v.to_string()), + |s| s.as_deref().into(), + args, + ) +} + +pub fn extract_json_string(args: &[ColumnarValue]) -> Result { + json_function::( + "extract_json_string", + |s, path| { + path.query(&s) + .first() + .and_then(|v| v.as_str().map(|s| s.to_string())) + }, + |s| s.as_deref().into(), + args, + ) +} + +// This code is vendored from +// https://github.com/datafusion-contrib/datafusion-functions-json/blob/main/src/common_union.rs +// as the `is_json_union` function is not public. It should be kept in sync with that code so +// that we are able to detect JSON unions and rewrite them to serialized JSON for sinks. +pub(crate) fn is_json_union(data_type: &DataType) -> bool { + match data_type { + DataType::Union(fields, UnionMode::Sparse) => fields == &union_fields(), + _ => false, + } +} + +pub(crate) const TYPE_ID_NULL: i8 = 0; +const TYPE_ID_BOOL: i8 = 1; +const TYPE_ID_INT: i8 = 2; +const TYPE_ID_FLOAT: i8 = 3; +const TYPE_ID_STR: i8 = 4; +const TYPE_ID_ARRAY: i8 = 5; +const TYPE_ID_OBJECT: i8 = 6; + +fn union_fields() -> UnionFields { + static FIELDS: OnceLock = OnceLock::new(); + FIELDS + .get_or_init(|| { + let json_metadata: HashMap = + HashMap::from_iter(vec![("is_json".to_string(), "true".to_string())]); + UnionFields::from_iter([ + ( + TYPE_ID_NULL, + Arc::new(Field::new("null", DataType::Null, true)), + ), + ( + TYPE_ID_BOOL, + Arc::new(Field::new("bool", DataType::Boolean, false)), + ), + ( + TYPE_ID_INT, + Arc::new(Field::new("int", DataType::Int64, false)), + ), + ( + TYPE_ID_FLOAT, + Arc::new(Field::new("float", DataType::Float64, false)), + ), + ( + TYPE_ID_STR, + Arc::new(Field::new("str", DataType::Utf8, false)), + ), + ( + TYPE_ID_ARRAY, + Arc::new( + Field::new("array", DataType::Utf8, false) + .with_metadata(json_metadata.clone()), + ), + ), + ( + TYPE_ID_OBJECT, + Arc::new( + Field::new("object", DataType::Utf8, false) + .with_metadata(json_metadata.clone()), + ), + ), + ]) + }) + .clone() +} +// End vendored code + +pub fn serialize_json_union(args: &[ColumnarValue]) -> Result { + assert_eq!(args.len(), 1); + let array = match args.first().unwrap() { + ColumnarValue::Array(a) => a.clone(), + ColumnarValue::Scalar(s) => s.to_array_of_size(1)?, + }; + + let mut b = StringBuilder::with_capacity(array.len(), array.get_array_memory_size()); + + write_union(&mut b, &array)?; + + Ok(ColumnarValue::Array(Arc::new(b.finish()))) +} + +fn write_union(b: &mut StringBuilder, array: &ArrayRef) -> Result<(), std::fmt::Error> { + assert!( + is_json_union(array.data_type()), + "array item is not a valid JSON union" + ); + let json_union = array.as_any().downcast_ref::().unwrap(); + + for i in 0..json_union.len() { + if json_union.is_null(i) { + b.append_null(); + } else { + write_value(b, json_union.type_id(i), &json_union.value(i))?; + b.append_value(""); + } + } + + Ok(()) +} + +fn write_value(b: &mut StringBuilder, id: i8, a: &ArrayRef) -> Result<(), std::fmt::Error> { + match id { + TYPE_ID_NULL => write!(b, "null")?, + TYPE_ID_BOOL => write!(b, "{}", a.as_boolean().value(0))?, + TYPE_ID_INT => write!(b, "{}", a.as_primitive::().value(0))?, + TYPE_ID_FLOAT => write!(b, "{}", a.as_primitive::().value(0))?, + TYPE_ID_STR => { + // assumes that this is already a valid (escaped) json string as the only way to + // construct these values are by parsing (valid) JSON + b.write_char('"')?; + b.write_str(a.as_string::().value(0))?; + b.write_char('"')?; + } + TYPE_ID_ARRAY => { + b.write_str(a.as_string::().value(0))?; + } + TYPE_ID_OBJECT => { + b.write_str(a.as_string::().value(0))?; + } + _ => unreachable!("invalid union type in JSON union: {}", id), + } + + Ok(()) +} + +pub(crate) fn serialize_outgoing_json( + registry: &StreamSchemaProvider, + node: Arc, +) -> LogicalPlan { + let exprs = node + .schema() + .fields() + .iter() + .map(|f| { + if is_json_union(f.data_type()) { + Expr::Alias(Alias::new( + Expr::ScalarFunction(ScalarFunction::new_udf( + registry.udf(SERIALIZE_JSON_UNION).unwrap(), + vec![col(f.name())], + )), + Option::::None, + f.name(), + )) + } else { + col(f.name()) + } + }) + .collect(); + + LogicalPlan::Projection(Projection::try_new(exprs, node).unwrap()) +} + +#[cfg(test)] +mod test { + use datafusion::arrow::array::StringArray; + use datafusion::arrow::array::builder::{ListBuilder, StringBuilder}; + use datafusion::common::ScalarValue; + use std::sync::Arc; + + #[test] + fn test_extract_json() { + let input = Arc::new(StringArray::from(vec![ + r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#, + r#"{"a": 3, "b": 4}"#, + r#"{"a": 5, "b": 6}"#, + ])); + + let path = "$.c.d"; + + let result = super::extract_json(&[ + super::ColumnarValue::Array(input), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let mut expected = ListBuilder::new(StringBuilder::new()); + expected.append_value(vec![Some("\"hello\"".to_string())]); + expected.append_value(Vec::>::new()); + expected.append_value(Vec::>::new()); + if let super::ColumnarValue::Array(result) = result { + assert_eq!(*result, expected.finish()); + } else { + panic!("Expected array, got scalar"); + } + + let result = super::extract_json(&[ + super::ColumnarValue::Scalar(r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#.into()), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let mut expected = ListBuilder::with_capacity(StringBuilder::new(), 1); + expected.append_value(vec![Some("\"hello\"".to_string())]); + + if let super::ColumnarValue::Scalar(ScalarValue::List(result)) = result { + assert_eq!(*result, expected.finish()); + } else { + panic!("Expected scalar list"); + } + } + + #[test] + fn test_get_first_json_object() { + let input = Arc::new(StringArray::from(vec![ + r#"{"a": 1, "b": 2}"#, + r#"{"a": 3}"#, + r#"{"a": 5, "b": 6}"#, + ])); + + let path = "$.b"; + + let result = super::get_first_json_object(&[ + super::ColumnarValue::Array(input), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let expected = StringArray::from(vec![Some("2"), None, Some("6")]); + + if let super::ColumnarValue::Array(result) = result { + assert_eq!(*result, expected); + } else { + panic!("Expected array, got scalar"); + } + + let result = super::get_first_json_object(&[ + super::ColumnarValue::Scalar(r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#.into()), + super::ColumnarValue::Scalar("$.c.d".into()), + ]) + .unwrap(); + + let expected = ScalarValue::Utf8(Some("\"hello\"".to_string())); + + if let super::ColumnarValue::Scalar(result) = result { + assert_eq!(result, expected); + } else { + panic!("Expected scalar"); + } + } + + #[test] + fn test_extract_json_string() { + let input = Arc::new(StringArray::from(vec![ + r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#, + r#"{"a": 3, "b": 4}"#, + r#"{"a": 5, "b": 6}"#, + ])); + + let path = "$.c.d"; + + let result = super::extract_json_string(&[ + super::ColumnarValue::Array(input), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let expected = StringArray::from(vec![Some("hello"), None, None]); + + if let super::ColumnarValue::Array(result) = result { + assert_eq!(*result, expected); + } else { + panic!("Expected array, got scalar"); + } + + let result = super::extract_json_string(&[ + super::ColumnarValue::Scalar(r#"{"a": 1, "b": 2, "c": { "d": "hello" }}"#.into()), + super::ColumnarValue::Scalar(path.into()), + ]) + .unwrap(); + + let expected = ScalarValue::Utf8(Some("hello".to_string())); + + if let super::ColumnarValue::Scalar(result) = result { + assert_eq!(result, expected); + } else { + panic!("Expected scalar"); + } + } +} diff --git a/src/sql/grammar.pest b/src/sql/grammar.pest deleted file mode 100644 index 15f70dd7..00000000 --- a/src/sql/grammar.pest +++ /dev/null @@ -1,134 +0,0 @@ -// ============================================================================= -// FUNCTION SQL Grammar -// -// Using pest PEG syntax, referencing ANTLR style -// ============================================================================= - -// ============================================================================= -// 1. Whitespace (automatically skipped) -// ============================================================================= - -WHITESPACE = _{ " " | "\t" | "\r" | "\n" } - -// ============================================================================= -// 2. Keywords (case-insensitive) -// ============================================================================= - -kw_create = _{ C ~ R ~ E ~ A ~ T ~ E } -kw_drop = _{ D ~ R ~ O ~ P } -kw_start = _{ S ~ T ~ A ~ R ~ T } -kw_stop = _{ S ~ T ~ O ~ P } -kw_show = _{ S ~ H ~ O ~ W } -kw_with = _{ W ~ I ~ T ~ H } -kw_function = _{ F ~ U ~ N ~ C ~ T ~ I ~ O ~ N } -kw_functions = _{ F ~ U ~ N ~ C ~ T ~ I ~ O ~ N ~ S } - -// ============================================================================= -// 3. Operators & Symbols -// ============================================================================= - -LPAREN = _{ "(" } -RPAREN = _{ ")" } -COMMA = _{ "," } -EQ = _{ "=" } -SQUOTE = _{ "'" } -DQUOTE = _{ "\"" } - -// ============================================================================= -// 4. Literals -// ============================================================================= - -// String literal (single or double quotes) -string_literal = @{ - SQUOTE ~ string_inner_single ~ SQUOTE | - DQUOTE ~ string_inner_double ~ DQUOTE -} - -string_inner_single = @{ (!(SQUOTE | "\\") ~ ANY | escape_seq)* } -string_inner_double = @{ (!(DQUOTE | "\\") ~ ANY | escape_seq)* } -escape_seq = @{ "\\" ~ ANY } - -// ============================================================================= -// 5. Identifiers -// ============================================================================= - -// Task name identifier -identifier = @{ (ASCII_ALPHA | "_") ~ (ASCII_ALPHANUMERIC | "_" | "-")* } - -// ============================================================================= -// 6. Statements -// ============================================================================= - -// Entry rule -statement = _{ - SOI ~ ( - create_stmt | - drop_stmt | - start_stmt | - stop_stmt | - show_stmt - ) ~ EOI -} - -// CREATE FUNCTION WITH (...) -// Note: name is read from config file, not from SQL statement -create_stmt = { kw_create ~ kw_function ~ kw_with ~ properties } - -// DROP FUNCTION name -drop_stmt = { kw_drop ~ kw_function ~ identifier } - -// START FUNCTION name -start_stmt = { kw_start ~ kw_function ~ identifier } - -// STOP FUNCTION name -stop_stmt = { kw_stop ~ kw_function ~ identifier } - -// SHOW FUNCTIONS -show_stmt = { kw_show ~ kw_functions } - -// ============================================================================= -// 7. Properties -// ============================================================================= - -// Property list ('key'='value', ...) -properties = { LPAREN ~ property ~ (COMMA ~ property)* ~ RPAREN } - -// Single property 'key'='value' -property = { property_key ~ EQ ~ property_value } - -// Property key (string) -property_key = { string_literal } - -// Property value (string) -property_value = { string_literal } - -// ============================================================================= -// 8. Character Fragments (for case-insensitive matching) -// ============================================================================= - -A = _{ "A" | "a" } -B = _{ "B" | "b" } -C = _{ "C" | "c" } -D = _{ "D" | "d" } -E = _{ "E" | "e" } -F = _{ "F" | "f" } -G = _{ "G" | "g" } -H = _{ "H" | "h" } -I = _{ "I" | "i" } -J = _{ "J" | "j" } -K = _{ "K" | "k" } -L = _{ "L" | "l" } -M = _{ "M" | "m" } -N = _{ "N" | "n" } -O = _{ "O" | "o" } -P = _{ "P" | "p" } -Q = _{ "Q" | "q" } -R = _{ "R" | "r" } -S = _{ "S" | "s" } -T = _{ "T" | "t" } -U = _{ "U" | "u" } -V = _{ "V" | "v" } -W = _{ "W" | "w" } -X = _{ "X" | "x" } -Y = _{ "Y" | "y" } -Z = _{ "Z" | "z" } diff --git a/src/sql/mod.rs b/src/sql/mod.rs index ed3c2e30..e0931530 100644 --- a/src/sql/mod.rs +++ b/src/sql/mod.rs @@ -10,6 +10,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -pub mod parser; +pub mod catalog; +pub mod functions; +pub mod physical; +pub mod planner; +pub mod types; -pub use parser::SqlParser; +pub use planner::StreamSchemaProvider; +pub use planner::parse::parse_sql; +pub use planner::plan::rewrite_plan; +pub use planner::sql_to_plan::statement_to_plan; +pub use planner::{CompiledSql, parse_and_get_arrow_program, parse_sql_statements}; diff --git a/src/sql/parser/sql_parser.rs b/src/sql/parser/sql_parser.rs deleted file mode 100644 index dc110745..00000000 --- a/src/sql/parser/sql_parser.rs +++ /dev/null @@ -1,249 +0,0 @@ -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use pest::Parser; -use pest_derive::Parser; - -use super::ParseError; -use crate::coordinator::{ - CreateFunction, DropFunction, ShowFunctions, StartFunction, Statement, StopFunction, -}; -use std::collections::HashMap; - -#[derive(Parser)] -#[grammar = "src/sql/grammar.pest"] -struct Grammar; - -#[derive(Debug, Default)] -pub struct SqlParser; - -impl SqlParser { - pub fn parse(sql: &str) -> Result, ParseError> { - let pairs = Grammar::parse(Rule::statement, sql) - .map_err(|e| ParseError::new(format!("Parse error: {}", e)))?; - - for pair in pairs { - return match pair.as_rule() { - Rule::create_stmt => { - handle_create_stmt(pair).map(|stmt| stmt as Box) - } - Rule::drop_stmt => handle_drop_stmt(pair).map(|stmt| stmt as Box), - Rule::start_stmt => handle_start_stmt(pair).map(|stmt| stmt as Box), - Rule::stop_stmt => handle_stop_stmt(pair).map(|stmt| stmt as Box), - Rule::show_stmt => handle_show_stmt(pair).map(|stmt| stmt as Box), - _ => continue, - }; - } - - Err(ParseError::new("Unknown statement type")) - } -} - -fn handle_create_stmt( - pair: pest::iterators::Pair, -) -> Result, ParseError> { - let mut inner = pair.into_inner(); - // Note: name is read from config file, not from SQL statement - // Pass empty string here, name will be read from config file later - let properties = inner - .next() - .map(parse_properties) - .ok_or_else(|| ParseError::new("Missing WITH clause"))?; - - Ok(Box::new( - CreateFunction::from_properties(properties).map_err(ParseError::from)?, - )) -} - -fn handle_drop_stmt(pair: pest::iterators::Pair) -> Result, ParseError> { - let mut inner = pair.into_inner(); - let name = inner.next().map(extract_string).unwrap_or_default(); - Ok(Box::new(DropFunction::new(name))) -} - -fn handle_start_stmt(pair: pest::iterators::Pair) -> Result, ParseError> { - let mut inner = pair.into_inner(); - let name = inner.next().map(extract_string).unwrap_or_default(); - Ok(Box::new(StartFunction::new(name))) -} - -fn handle_stop_stmt(pair: pest::iterators::Pair) -> Result, ParseError> { - let mut inner = pair.into_inner(); - let name = inner.next().map(extract_string).unwrap_or_default(); - Ok(Box::new(StopFunction::new(name))) -} - -fn handle_show_stmt(_pair: pest::iterators::Pair) -> Result, ParseError> { - Ok(Box::new(ShowFunctions::new())) -} - -fn extract_string(pair: pest::iterators::Pair) -> String { - match pair.as_rule() { - Rule::string_literal => { - let s = pair.as_str(); - if (s.starts_with('\'') && s.ends_with('\'')) - || (s.starts_with('"') && s.ends_with('"')) - { - unescape_string(&s[1..s.len() - 1]) - } else { - unescape_string(s) - } - } - Rule::identifier => pair.as_str().to_string(), - _ => pair.as_str().to_string(), - } -} - -fn unescape_string(s: &str) -> String { - let mut result = String::with_capacity(s.len()); - let mut chars = s.chars().peekable(); - - while let Some(ch) = chars.next() { - if ch == '\\' { - if let Some(&next) = chars.peek() { - chars.next(); - match next { - 'n' => result.push('\n'), - 't' => result.push('\t'), - 'r' => result.push('\r'), - '\\' => result.push('\\'), - '\'' => result.push('\''), - '"' => result.push('"'), - _ => { - result.push('\\'); - result.push(next); - } - } - } else { - result.push(ch); - } - } else { - result.push(ch); - } - } - - result -} - -fn parse_properties(pair: pest::iterators::Pair) -> HashMap { - let mut properties = HashMap::new(); - - for prop in pair.into_inner() { - if prop.as_rule() == Rule::property { - let mut inner = prop.into_inner(); - if let (Some(key_pair), Some(val_pair)) = (inner.next(), inner.next()) { - let key = key_pair - .into_inner() - .next() - .map(extract_string) - .unwrap_or_default(); - let value = val_pair - .into_inner() - .next() - .map(extract_string) - .unwrap_or_default(); - properties.insert(key, value); - } - } - } - - properties -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_create_function() { - let sql = - "CREATE FUNCTION WITH ('function_path'='./test.wasm', 'config_path'='./config.yml')"; - let _stmt = SqlParser::parse(sql).unwrap(); - } - - #[test] - fn test_create_function_minimal() { - let sql = "CREATE FUNCTION WITH ('function_path'='./processor.wasm')"; - let _stmt = SqlParser::parse(sql).unwrap(); - } - - // Note: SQL only supports Path mode, not Bytes mode - // Bytes mode is only for gRPC requests - - #[test] - fn test_drop_function() { - let sql = "DROP FUNCTION my_task"; - let _stmt = SqlParser::parse(sql).unwrap(); - } - - #[test] - fn test_start_function() { - let sql = "START FUNCTION my_task"; - let _stmt = SqlParser::parse(sql).unwrap(); - } - - #[test] - fn test_stop_function() { - let sql = "STOP FUNCTION my_task"; - let _stmt = SqlParser::parse(sql).unwrap(); - } - - #[test] - fn test_show_functions() { - let sql = "SHOW FUNCTIONS"; - let _stmt = SqlParser::parse(sql).unwrap(); - } - - #[test] - fn test_case_insensitive_keywords() { - let sql1 = "create function with ('function_path'='./test.wasm')"; - let _stmt1 = SqlParser::parse(sql1).unwrap(); - - let sql2 = "Create Function With ('Function_Path'='./test.wasm')"; - let _stmt2 = SqlParser::parse(sql2).unwrap(); - - let sql3 = "show functions"; - let _stmt3 = SqlParser::parse(sql3).unwrap(); - - let sql4 = "start function my_task"; - let _stmt4 = SqlParser::parse(sql4).unwrap(); - } - - #[test] - fn test_case_insensitive_property_keys() { - let sql1 = - "CREATE FUNCTION WITH ('function_path'='./test.wasm', 'config_path'='./config.yml')"; - let _stmt1 = SqlParser::parse(sql1).unwrap(); - - let sql2 = - "CREATE FUNCTION WITH ('Function_Path'='./test.wasm', 'Config_Path'='./config.yml')"; - let _stmt2 = SqlParser::parse(sql2).unwrap(); - - let sql3 = - "CREATE FUNCTION WITH ('FUNCTION_PATH'='./test.wasm', 'CONFIG_PATH'='./config.yml')"; - let _stmt3 = SqlParser::parse(sql3).unwrap(); - - // Note: SQL only supports Path mode (function_path, config_path) - // Bytes mode (function, config) is only for gRPC requests - } - - #[test] - fn test_with_extra_properties() { - let sql = r#"CREATE FUNCTION WITH ( - 'function_path'='./test.wasm', - 'config_path'='./config.yml', - 'parallelism'='4', - 'memory-limit'='256mb' - )"#; - let _stmt = SqlParser::parse(sql).unwrap(); - } -} diff --git a/src/sql/physical/mod.rs b/src/sql/physical/mod.rs new file mode 100644 index 00000000..bfb37f11 --- /dev/null +++ b/src/sql/physical/mod.rs @@ -0,0 +1,1265 @@ +use datafusion::arrow::{ + array::{ + Array, AsArray, BooleanBuilder, PrimitiveArray, RecordBatch, StringArray, StructArray, + TimestampNanosecondArray, TimestampNanosecondBuilder, UInt32Builder, + }, + buffer::NullBuffer, + compute::{concat, take}, + datatypes::{DataType, Field, Fields, Schema, SchemaRef, TimeUnit}, +}; +use datafusion::common::{ + DataFusionError, Result, ScalarValue, Statistics, UnnestOptions, not_impl_err, plan_err, +}; +use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream}; +use datafusion::{ + execution::TaskContext, + physical_plan::{ + DisplayAs, ExecutionPlan, Partitioning, memory::MemoryStream, + stream::RecordBatchStreamAdapter, + }, +}; +use std::collections::HashMap; +use std::{ + any::Any, + mem, + pin::Pin, + sync::{Arc, OnceLock, RwLock}, + task::{Context, Poll}, +}; + +use crate::make_udf_function; +use crate::sql::functions::MultiHashFunction; +use crate::sql::planner::rewrite::UNNESTED_COL; +use crate::sql::planner::schemas::window_arrow_struct; +use crate::types::{TIMESTAMP_FIELD, UPDATING_META_FIELD}; +use datafusion::arrow::datatypes::{TimestampNanosecondType, UInt64Type}; +use datafusion::catalog::memory::MemorySourceConfig; +use datafusion::datasource::memory::DataSourceExec; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, +}; +use datafusion::physical_expr::EquivalenceProperties; +use datafusion::physical_plan::PlanProperties; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec}; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use futures::{ + ready, + stream::{Stream, StreamExt}, +}; +use prost::Message; +use protocol::grpc::api::{ + DebeziumDecodeNode, DebeziumEncodeNode, FsExecNode, MemExecNode, UnnestExecNode, + fs_exec_node::Node, +}; +use std::fmt::Debug; +use tokio::sync::mpsc::UnboundedReceiver; +use tokio_stream::wrappers::UnboundedReceiverStream; + +// ─────────────────── Updating Meta Helpers ─────────────────── + +pub fn updating_meta_fields() -> Fields { + static FIELDS: OnceLock = OnceLock::new(); + FIELDS + .get_or_init(|| { + Fields::from(vec![ + Field::new("is_retract", DataType::Boolean, true), + Field::new("id", DataType::FixedSizeBinary(16), true), + ]) + }) + .clone() +} + +pub fn updating_meta_field() -> Arc { + static FIELD: OnceLock> = OnceLock::new(); + FIELD + .get_or_init(|| { + Arc::new(Field::new( + UPDATING_META_FIELD, + DataType::Struct(updating_meta_fields()), + false, + )) + }) + .clone() +} + +// ─────────────────── WindowFunctionUdf ─────────────────── + +#[derive(Debug)] +pub struct WindowFunctionUdf { + signature: Signature, +} + +impl Default for WindowFunctionUdf { + fn default() -> Self { + Self { + signature: Signature::new( + TypeSignature::Exact(vec![ + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Timestamp(TimeUnit::Nanosecond, None), + ]), + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for WindowFunctionUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "window" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _: &[DataType]) -> Result { + Ok(window_arrow_struct()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let columns = args.args; + if columns.len() != 2 { + return plan_err!( + "window function expected 2 arguments, got {}", + columns.len() + ); + } + if columns[0].data_type() != DataType::Timestamp(TimeUnit::Nanosecond, None) { + return plan_err!( + "window function expected first argument to be a timestamp, got {:?}", + columns[0].data_type() + ); + } + if columns[1].data_type() != DataType::Timestamp(TimeUnit::Nanosecond, None) { + return plan_err!( + "window function expected second argument to be a timestamp, got {:?}", + columns[1].data_type() + ); + } + let fields = vec![ + Arc::new(Field::new( + "start", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )), + Arc::new(Field::new( + "end", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )), + ] + .into(); + + match (&columns[0], &columns[1]) { + (ColumnarValue::Array(start), ColumnarValue::Array(end)) => { + Ok(ColumnarValue::Array(Arc::new(StructArray::new( + fields, + vec![start.clone(), end.clone()], + None, + )))) + } + (ColumnarValue::Array(start), ColumnarValue::Scalar(end)) => { + let end = end.to_array_of_size(start.len())?; + Ok(ColumnarValue::Array(Arc::new(StructArray::new( + fields, + vec![start.clone(), end], + None, + )))) + } + (ColumnarValue::Scalar(start), ColumnarValue::Array(end)) => { + let start = start.to_array_of_size(end.len())?; + Ok(ColumnarValue::Array(Arc::new(StructArray::new( + fields, + vec![start, end.clone()], + None, + )))) + } + (ColumnarValue::Scalar(start), ColumnarValue::Scalar(end)) => { + Ok(ColumnarValue::Scalar(ScalarValue::Struct( + StructArray::new(fields, vec![start.to_array()?, end.to_array()?], None).into(), + ))) + } + } + } +} + +make_udf_function!(WindowFunctionUdf, WINDOW_FUNCTION, window); + +// ─────────────────── Physical Extension Codec ─────────────────── + +#[derive(Debug)] +pub struct FsPhysicalExtensionCodec { + pub context: DecodingContext, +} + +impl Default for FsPhysicalExtensionCodec { + fn default() -> Self { + Self { + context: DecodingContext::None, + } + } +} + +#[derive(Debug)] +pub enum DecodingContext { + None, + Planning, + SingleLockedBatch(Arc>>), + UnboundedBatchStream(Arc>>>), + LockedBatchVec(Arc>>), + LockedJoinPair { + left: Arc>>, + right: Arc>>, + }, + LockedJoinStream { + left: Arc>>>, + right: Arc>>>, + }, +} + +fn make_properties(schema: SchemaRef) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Unbounded { + requires_infinite_memory: false, + }, + ) +} + +impl PhysicalExtensionCodec for FsPhysicalExtensionCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + _registry: &dyn datafusion::execution::FunctionRegistry, + ) -> Result> { + let exec: FsExecNode = Message::decode(buf) + .map_err(|err| DataFusionError::Internal(format!("couldn't deserialize: {err}")))?; + + match exec + .node + .ok_or_else(|| DataFusionError::Internal("exec node is empty".to_string()))? + { + Node::MemExec(mem_exec) => { + let schema: Schema = serde_json::from_str(&mem_exec.schema).map_err(|e| { + DataFusionError::Internal(format!("invalid schema in exec codec: {e:?}")) + })?; + let schema = Arc::new(schema); + match &self.context { + DecodingContext::SingleLockedBatch(single_batch) => Ok(Arc::new( + RwLockRecordBatchReader::new(schema, single_batch.clone()), + )), + DecodingContext::UnboundedBatchStream(unbounded_stream) => Ok(Arc::new( + UnboundedRecordBatchReader::new(schema, unbounded_stream.clone()), + )), + DecodingContext::LockedBatchVec(locked_batches) => Ok(Arc::new( + RecordBatchVecReader::new(schema, locked_batches.clone()), + )), + DecodingContext::Planning => { + Ok(Arc::new(FsMemExec::new(mem_exec.table_name, schema))) + } + DecodingContext::None => Err(DataFusionError::Internal( + "Need an internal context to decode".into(), + )), + DecodingContext::LockedJoinPair { left, right } => { + match mem_exec.table_name.as_str() { + "left" => { + Ok(Arc::new(RwLockRecordBatchReader::new(schema, left.clone()))) + } + "right" => Ok(Arc::new(RwLockRecordBatchReader::new( + schema, + right.clone(), + ))), + _ => Err(DataFusionError::Internal(format!( + "unknown table name {}", + mem_exec.table_name + ))), + } + } + DecodingContext::LockedJoinStream { left, right } => { + match mem_exec.table_name.as_str() { + "left" => Ok(Arc::new(UnboundedRecordBatchReader::new( + schema, + left.clone(), + ))), + "right" => Ok(Arc::new(UnboundedRecordBatchReader::new( + schema, + right.clone(), + ))), + _ => Err(DataFusionError::Internal(format!( + "unknown table name {}", + mem_exec.table_name + ))), + } + } + } + } + Node::UnnestExec(unnest) => { + let schema: Schema = serde_json::from_str(&unnest.schema).map_err(|e| { + DataFusionError::Internal(format!("invalid schema in exec codec: {e:?}")) + })?; + + let column = schema.index_of(UNNESTED_COL).map_err(|_| { + DataFusionError::Internal(format!( + "unnest node schema does not contain {UNNESTED_COL} col" + )) + })?; + + Ok(Arc::new(UnnestExec::new( + inputs + .first() + .ok_or_else(|| { + DataFusionError::Internal("no input for unnest node".to_string()) + })? + .clone(), + vec![ListUnnest { + index_in_input_schema: column, + depth: 1, + }], + vec![], + Arc::new(schema), + UnnestOptions::default(), + ))) + } + Node::DebeziumDecode(debezium) => { + let schema = Arc::new(serde_json::from_str::(&debezium.schema).map_err( + |e| DataFusionError::Internal(format!("invalid schema in exec codec: {e:?}")), + )?); + Ok(Arc::new(DebeziumUnrollingExec { + input: inputs + .first() + .ok_or_else(|| { + DataFusionError::Internal("no input for debezium node".to_string()) + })? + .clone(), + schema: schema.clone(), + properties: make_properties(schema), + primary_keys: debezium + .primary_keys + .into_iter() + .map(|c| c as usize) + .collect(), + })) + } + Node::DebeziumEncode(debezium) => { + let schema = Arc::new(serde_json::from_str::(&debezium.schema).map_err( + |e| DataFusionError::Internal(format!("invalid schema in exec codec: {e:?}")), + )?); + Ok(Arc::new(ToDebeziumExec { + input: inputs + .first() + .ok_or_else(|| { + DataFusionError::Internal("no input for debezium node".to_string()) + })? + .clone(), + schema: schema.clone(), + properties: make_properties(schema), + })) + } + } + } + + fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()> { + let mut proto = None; + + let mem_table: Option<&FsMemExec> = node.as_any().downcast_ref(); + if let Some(table) = mem_table { + proto = Some(FsExecNode { + node: Some(Node::MemExec(MemExecNode { + table_name: table.table_name.clone(), + schema: serde_json::to_string(&table.schema).unwrap(), + })), + }); + } + + let unnest: Option<&UnnestExec> = node.as_any().downcast_ref(); + if let Some(unnest) = unnest { + proto = Some(FsExecNode { + node: Some(Node::UnnestExec(UnnestExecNode { + schema: serde_json::to_string(&unnest.schema()).unwrap(), + })), + }); + } + + let debezium_decode: Option<&DebeziumUnrollingExec> = node.as_any().downcast_ref(); + if let Some(decode) = debezium_decode { + proto = Some(FsExecNode { + node: Some(Node::DebeziumDecode(DebeziumDecodeNode { + schema: serde_json::to_string(&decode.schema).unwrap(), + primary_keys: (*decode.primary_keys).iter().map(|c| *c as u64).collect(), + })), + }); + } + + let debezium_encode: Option<&ToDebeziumExec> = node.as_any().downcast_ref(); + if let Some(encode) = debezium_encode { + proto = Some(FsExecNode { + node: Some(Node::DebeziumEncode(DebeziumEncodeNode { + schema: serde_json::to_string(&encode.schema).unwrap(), + })), + }); + } + + if let Some(node) = proto { + node.encode(buf).map_err(|err| { + DataFusionError::Internal(format!("couldn't serialize exec node {err}")) + })?; + Ok(()) + } else { + Err(DataFusionError::Internal(format!( + "cannot serialize {node:?}" + ))) + } + } +} + +// ─────────────────── RwLockRecordBatchReader ─────────────────── + +#[derive(Debug)] +struct RwLockRecordBatchReader { + schema: SchemaRef, + locked_batch: Arc>>, + properties: PlanProperties, +} + +impl RwLockRecordBatchReader { + fn new(schema: SchemaRef, locked_batch: Arc>>) -> Self { + Self { + schema: schema.clone(), + locked_batch, + properties: make_properties(schema), + } + } +} + +impl DisplayAs for RwLockRecordBatchReader { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "RW Lock RecordBatchReader") + } +} + +impl ExecutionPlan for RwLockRecordBatchReader { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Err(DataFusionError::Internal("not supported".into())) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + let result = self + .locked_batch + .write() + .unwrap() + .take() + .expect("should have set a record batch before calling execute()"); + Ok(Box::pin(MemoryStream::try_new( + vec![result], + self.schema.clone(), + None, + )?)) + } + + fn statistics(&self) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn reset(&self) -> Result<()> { + Ok(()) + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn name(&self) -> &str { + "rw_lock_reader" + } +} + +// ─────────────────── UnboundedRecordBatchReader ─────────────────── + +#[derive(Debug)] +struct UnboundedRecordBatchReader { + schema: SchemaRef, + receiver: Arc>>>, + properties: PlanProperties, +} + +impl UnboundedRecordBatchReader { + fn new( + schema: SchemaRef, + receiver: Arc>>>, + ) -> Self { + Self { + schema: schema.clone(), + receiver, + properties: make_properties(schema), + } + } +} + +impl DisplayAs for UnboundedRecordBatchReader { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "unbounded record batch reader") + } +} + +impl ExecutionPlan for UnboundedRecordBatchReader { + fn name(&self) -> &str { + "unbounded_reader" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Err(DataFusionError::Internal("not supported".into())) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + UnboundedReceiverStream::new( + self.receiver + .write() + .unwrap() + .take() + .expect("unbounded receiver should be present before calling exec"), + ) + .map(Ok), + ))) + } + + fn statistics(&self) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn reset(&self) -> Result<()> { + Ok(()) + } +} + +// ─────────────────── RecordBatchVecReader ─────────────────── + +#[derive(Debug)] +struct RecordBatchVecReader { + schema: SchemaRef, + receiver: Arc>>, + properties: PlanProperties, +} + +impl RecordBatchVecReader { + fn new(schema: SchemaRef, receiver: Arc>>) -> Self { + Self { + schema: schema.clone(), + receiver, + properties: make_properties(schema), + } + } +} + +impl DisplayAs for RecordBatchVecReader { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "record batch vec reader") + } +} + +impl ExecutionPlan for RecordBatchVecReader { + fn name(&self) -> &str { + "vec_reader" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Err(DataFusionError::Internal("not supported".into())) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let memory = MemorySourceConfig::try_new( + &[mem::take(self.receiver.write().unwrap().as_mut())], + self.schema.clone(), + None, + )?; + + DataSourceExec::new(Arc::new(memory)).execute(partition, context) + } + + fn statistics(&self) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn reset(&self) -> Result<()> { + Ok(()) + } +} + +// ─────────────────── FsMemExec ─────────────────── + +#[derive(Debug, Clone)] +pub struct FsMemExec { + pub table_name: String, + pub schema: SchemaRef, + properties: PlanProperties, +} + +impl DisplayAs for FsMemExec { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "EmptyPartitionStream: schema={}", self.schema) + } +} + +impl FsMemExec { + pub fn new(table_name: String, schema: SchemaRef) -> Self { + Self { + schema: schema.clone(), + table_name, + properties: make_properties(schema), + } + } +} + +impl ExecutionPlan for FsMemExec { + fn name(&self) -> &str { + "mem_exec" + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + not_impl_err!("with_new_children is not implemented for mem_exec; should not be called") + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + plan_err!( + "EmptyPartitionStream cannot be executed, this is only used for physical planning before serialization" + ) + } + + fn statistics(&self) -> Result { + Ok(Statistics::new_unknown(&self.schema)) + } + + fn reset(&self) -> Result<()> { + Ok(()) + } +} + +// ─────────────────── DebeziumUnrollingExec ─────────────────── + +#[derive(Debug)] +pub struct DebeziumUnrollingExec { + input: Arc, + schema: SchemaRef, + properties: PlanProperties, + primary_keys: Vec, +} + +impl DebeziumUnrollingExec { + pub fn try_new(input: Arc, primary_keys: Vec) -> Result { + let input_schema = input.schema(); + let before_index = input_schema.index_of("before")?; + let after_index = input_schema.index_of("after")?; + let op_index = input_schema.index_of("op")?; + let _timestamp_index = input_schema.index_of(TIMESTAMP_FIELD)?; + let before_type = input_schema.field(before_index).data_type(); + let after_type = input_schema.field(after_index).data_type(); + if before_type != after_type { + return Err(DataFusionError::Internal( + "before and after columns must have the same type".to_string(), + )); + } + let op_type = input_schema.field(op_index).data_type(); + if *op_type != DataType::Utf8 { + return Err(DataFusionError::Internal( + "op column must be a string".to_string(), + )); + } + let DataType::Struct(fields) = before_type else { + return Err(DataFusionError::Internal( + "before and after columns must be structs".to_string(), + )); + }; + let mut fields = fields.to_vec(); + fields.push(updating_meta_field()); + fields.push(Arc::new(Field::new( + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ))); + + let schema = Arc::new(Schema::new(fields)); + Ok(Self { + input, + schema: schema.clone(), + properties: make_properties(schema), + primary_keys, + }) + } +} + +impl DisplayAs for DebeziumUnrollingExec { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "DebeziumUnrollingExec") + } +} + +impl ExecutionPlan for DebeziumUnrollingExec { + fn name(&self) -> &str { + "debezium_unrolling_exec" + } + + fn as_any(&self) -> &dyn Any { + self as &dyn Any + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(DataFusionError::Internal( + "DebeziumUnrollingExec wrong number of children".to_string(), + )); + } + Ok(Arc::new(DebeziumUnrollingExec { + input: children[0].clone(), + schema: self.schema.clone(), + properties: self.properties.clone(), + primary_keys: self.primary_keys.clone(), + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + Ok(Box::pin(DebeziumUnrollingStream::try_new( + self.input.execute(partition, context)?, + self.schema.clone(), + self.primary_keys.clone(), + )?)) + } + + fn reset(&self) -> Result<()> { + self.input.reset() + } +} + +struct DebeziumUnrollingStream { + input: SendableRecordBatchStream, + schema: SchemaRef, + before_index: usize, + after_index: usize, + op_index: usize, + timestamp_index: usize, + primary_keys: Vec, +} + +impl DebeziumUnrollingStream { + fn try_new( + input: SendableRecordBatchStream, + schema: SchemaRef, + primary_keys: Vec, + ) -> Result { + if primary_keys.is_empty() { + return plan_err!("there must be at least one primary key for a Debezium source"); + } + let input_schema = input.schema(); + let before_index = input_schema.index_of("before")?; + let after_index = input_schema.index_of("after")?; + let op_index = input_schema.index_of("op")?; + let timestamp_index = input_schema.index_of(TIMESTAMP_FIELD)?; + + Ok(Self { + input, + schema, + before_index, + after_index, + op_index, + timestamp_index, + primary_keys, + }) + } + + fn unroll_batch(&self, batch: &RecordBatch) -> Result { + let before = batch.column(self.before_index).as_ref(); + let after = batch.column(self.after_index).as_ref(); + let op = batch + .column(self.op_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| DataFusionError::Internal("op column is not a string".to_string()))?; + + let timestamp = batch + .column(self.timestamp_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal("timestamp column is not a timestamp".to_string()) + })?; + + let num_rows = batch.num_rows(); + let combined_array = concat(&[before, after])?; + let mut take_indices = UInt32Builder::with_capacity(num_rows); + let mut is_retract_builder = BooleanBuilder::with_capacity(num_rows); + + let mut timestamp_builder = TimestampNanosecondBuilder::with_capacity(2 * num_rows); + for i in 0..num_rows { + let op = op.value(i); + match op { + "c" | "r" => { + take_indices.append_value((i + num_rows) as u32); + is_retract_builder.append_value(false); + timestamp_builder.append_value(timestamp.value(i)); + } + "u" => { + take_indices.append_value(i as u32); + is_retract_builder.append_value(true); + timestamp_builder.append_value(timestamp.value(i)); + take_indices.append_value((i + num_rows) as u32); + is_retract_builder.append_value(false); + timestamp_builder.append_value(timestamp.value(i)); + } + "d" => { + take_indices.append_value(i as u32); + is_retract_builder.append_value(true); + timestamp_builder.append_value(timestamp.value(i)); + } + _ => { + return Err(DataFusionError::Internal(format!( + "unexpected op value: {op}" + ))); + } + } + } + let take_indices = take_indices.finish(); + let unrolled_array = take(&combined_array, &take_indices, None)?; + + let mut columns = unrolled_array.as_struct().columns().to_vec(); + + let hash = MultiHashFunction::default().invoke( + &self + .primary_keys + .iter() + .map(|i| ColumnarValue::Array(columns[*i].clone())) + .collect::>(), + )?; + + let ids = hash.into_array(num_rows)?; + + let meta = StructArray::try_new( + updating_meta_fields(), + vec![Arc::new(is_retract_builder.finish()), ids], + None, + )?; + columns.push(Arc::new(meta)); + columns.push(Arc::new(timestamp_builder.finish())); + Ok(RecordBatch::try_new(self.schema.clone(), columns)?) + } +} + +impl Stream for DebeziumUnrollingStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll> { + let result = + ready!(self.input.poll_next_unpin(cx)).map(|result| self.unroll_batch(&result?)); + Poll::Ready(result) + } +} + +impl RecordBatchStream for DebeziumUnrollingStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +// ─────────────────── ToDebeziumExec ─────────────────── + +#[derive(Debug)] +pub struct ToDebeziumExec { + input: Arc, + schema: SchemaRef, + properties: PlanProperties, +} + +impl ToDebeziumExec { + pub fn try_new(input: Arc) -> Result { + let input_schema = input.schema(); + let timestamp_index = input_schema.index_of(TIMESTAMP_FIELD)?; + let struct_fields: Vec<_> = input_schema + .fields() + .into_iter() + .enumerate() + .filter_map(|(index, field)| { + if field.name() == UPDATING_META_FIELD || index == timestamp_index { + None + } else { + Some(field.clone()) + } + }) + .collect(); + let struct_data_type = DataType::Struct(struct_fields.into()); + let before_field = Arc::new(Field::new("before", struct_data_type.clone(), true)); + let after_field = Arc::new(Field::new("after", struct_data_type, true)); + let op_field = Arc::new(Field::new("op", DataType::Utf8, false)); + let timestamp_field = Arc::new(input_schema.field(timestamp_index).clone()); + + let output_schema = Arc::new(Schema::new(vec![ + before_field, + after_field, + op_field, + timestamp_field, + ])); + + Ok(Self { + input, + schema: output_schema.clone(), + properties: make_properties(output_schema), + }) + } +} + +impl DisplayAs for ToDebeziumExec { + fn fmt_as( + &self, + _t: datafusion::physical_plan::DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "ToDebeziumExec") + } +} + +impl ExecutionPlan for ToDebeziumExec { + fn name(&self) -> &str { + "to_debezium_exec" + } + + fn as_any(&self) -> &dyn Any { + self as &dyn Any + } + + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(DataFusionError::Internal( + "ToDebeziumExec wrong number of children".to_string(), + )); + } + Ok(Arc::new(ToDebeziumExec::try_new(children[0].clone())?)) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let updating_meta_index = self.input.schema().index_of(UPDATING_META_FIELD).ok(); + let timestamp_index = self.input.schema().index_of(TIMESTAMP_FIELD)?; + let struct_projection = (0..self.input.schema().fields().len()) + .filter(|index| { + updating_meta_index + .map(|is_retract_index| *index != is_retract_index) + .unwrap_or(true) + && *index != timestamp_index + }) + .collect(); + + Ok(Box::pin(ToDebeziumStream { + input: self.input.execute(partition, context)?, + schema: self.schema.clone(), + updating_meta_index, + timestamp_index, + struct_projection, + })) + } + + fn reset(&self) -> Result<()> { + self.input.reset() + } +} + +struct ToDebeziumStream { + input: SendableRecordBatchStream, + schema: SchemaRef, + updating_meta_index: Option, + timestamp_index: usize, + struct_projection: Vec, +} + +impl ToDebeziumStream { + fn as_debezium_batch(&mut self, batch: &RecordBatch) -> Result { + let value_struct = batch.project(&self.struct_projection)?; + let timestamps = batch + .column(self.timestamp_index) + .as_primitive::(); + + let columns: Vec> = if let Some(metadata_index) = self.updating_meta_index { + let metadata = batch + .column(metadata_index) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Internal("Invalid type for updating_meta column".to_string()) + })?; + + let is_retract = metadata.column(0).as_boolean(); + let id = metadata.column(1).as_fixed_size_binary(); + + let mut id_map: HashMap<&[u8], (usize, usize, bool, bool, i64)> = HashMap::new(); + let mut order = vec![]; + for i in 0..batch.num_rows() { + let row_id = id.value(i); + let is_create = !is_retract.value(i); + let timestamp = timestamps.value(i); + + id_map + .entry(row_id) + .and_modify(|e| { + e.1 = i; + e.3 = is_create; + e.4 = e.4.max(timestamp); + }) + .or_insert_with(|| { + order.push(row_id); + (i, i, is_create, is_create, timestamp) + }); + } + + let mut before = Vec::with_capacity(id_map.len()); + let mut after = Vec::with_capacity(id_map.len()); + let mut op = Vec::with_capacity(id_map.len()); + let mut ts = TimestampNanosecondBuilder::with_capacity(id_map.len()); + + for row_id in order { + let (first_idx, last_idx, first_is_create, last_is_create, timestamp) = + id_map.get(row_id).unwrap(); + + if *first_is_create && *last_is_create { + before.push(None); + after.push(Some(*last_idx)); + op.push("c"); + } else if !(*first_is_create) && !(*last_is_create) { + before.push(Some(*first_idx)); + after.push(None); + op.push("d"); + } else if !(*first_is_create) && *last_is_create { + before.push(Some(*first_idx)); + after.push(Some(*last_idx)); + op.push("u"); + } else { + continue; + } + + ts.append_value(*timestamp); + } + + let before_array = Self::create_output_array(&value_struct, &before)?; + let after_array = Self::create_output_array(&value_struct, &after)?; + let op_array = StringArray::from(op); + + vec![ + Arc::new(before_array), + Arc::new(after_array), + Arc::new(op_array), + Arc::new(ts.finish()), + ] + } else { + let after_array = StructArray::try_new( + value_struct.schema().fields().clone(), + value_struct.columns().to_vec(), + None, + )?; + + let before_array = StructArray::new_null( + value_struct.schema().fields().clone(), + value_struct.num_rows(), + ); + + vec![ + Arc::new(before_array), + Arc::new(after_array), + Arc::new(StringArray::from(vec!["c"; value_struct.num_rows()])), + batch.column(self.timestamp_index).clone(), + ] + }; + + Ok(RecordBatch::try_new(self.schema.clone(), columns)?) + } + + fn create_output_array( + value_struct: &RecordBatch, + indices: &[Option], + ) -> Result { + let mut arrays: Vec> = Vec::with_capacity(value_struct.num_columns()); + for col in value_struct.columns() { + let new_array = take( + col.as_ref(), + &indices + .iter() + .map(|&idx| idx.map(|i| i as u64)) + .collect::>(), + None, + )?; + arrays.push(new_array); + } + + Ok(StructArray::try_new( + value_struct.schema().fields().clone(), + arrays, + Some(NullBuffer::from( + indices.iter().map(|&idx| idx.is_some()).collect::>(), + )), + )?) + } +} + +impl Stream for ToDebeziumStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll> { + let result = + ready!(self.input.poll_next_unpin(cx)).map(|result| self.as_debezium_batch(&result?)); + Poll::Ready(result) + } +} + +impl RecordBatchStream for ToDebeziumStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} diff --git a/src/sql/planner/extension/aggregate.rs b/src/sql/planner/extension/aggregate.rs new file mode 100644 index 00000000..878d3cc5 --- /dev/null +++ b/src/sql/planner/extension/aggregate.rs @@ -0,0 +1,348 @@ +use std::fmt::Formatter; +use std::sync::Arc; +use std::time::Duration; + +use datafusion::arrow::datatypes::DataType; +use datafusion::common::{Column, DFSchemaRef, Result, ScalarValue, internal_err}; +use datafusion::logical_expr; +use datafusion::logical_expr::{ + BinaryExpr, Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore, expr::ScalarFunction, +}; + +use crate::multifield_partial_ord; +use crate::sql::planner::extension::{NamedNode, StreamExtension, TimestampAppendExtension}; +use crate::sql::types::{ + DFField, StreamSchema, TIMESTAMP_FIELD, WindowBehavior, WindowType, fields_with_qualifiers, + schema_from_df_fields, schema_from_df_fields_with_metadata, +}; + +pub(crate) const AGGREGATE_EXTENSION_NAME: &str = "AggregateExtension"; + +/// Extension node for windowed aggregate operations in streaming SQL. +/// Supports tumbling, sliding, session, and instant window aggregations. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct AggregateExtension { + pub(crate) window_behavior: WindowBehavior, + pub(crate) aggregate: LogicalPlan, + pub(crate) schema: DFSchemaRef, + pub(crate) key_fields: Vec, + pub(crate) final_calculation: LogicalPlan, +} + +multifield_partial_ord!(AggregateExtension, aggregate, key_fields, final_calculation); + +impl AggregateExtension { + pub fn new( + window_behavior: WindowBehavior, + aggregate: LogicalPlan, + key_fields: Vec, + ) -> Self { + let final_calculation = + Self::final_projection(&aggregate, window_behavior.clone()).unwrap(); + Self { + window_behavior, + aggregate, + schema: final_calculation.schema().clone(), + key_fields, + final_calculation, + } + } + + /// Build the final projection after aggregation, which adds the window struct + /// and computes the output timestamp based on the window behavior. + pub fn final_projection( + aggregate_plan: &LogicalPlan, + window_behavior: WindowBehavior, + ) -> Result { + let timestamp_field: DFField = aggregate_plan.inputs()[0] + .schema() + .qualified_field_with_unqualified_name(TIMESTAMP_FIELD)? + .into(); + let timestamp_append = LogicalPlan::Extension(Extension { + node: Arc::new(TimestampAppendExtension::new( + aggregate_plan.clone(), + timestamp_field.qualifier().cloned(), + )), + }); + let mut aggregate_fields = fields_with_qualifiers(aggregate_plan.schema()); + let mut aggregate_expressions: Vec<_> = aggregate_fields + .iter() + .map(|field| Expr::Column(field.qualified_column())) + .collect(); + + let (window_field, window_index, width, is_nested) = match window_behavior { + WindowBehavior::InData => return Ok(timestamp_append), + WindowBehavior::FromOperator { + window, + window_field, + window_index, + is_nested, + } => match window { + WindowType::Tumbling { width, .. } | WindowType::Sliding { width, .. } => { + (window_field, window_index, width, is_nested) + } + WindowType::Session { .. } => { + return Ok(LogicalPlan::Extension(Extension { + node: Arc::new(WindowAppendExtension::new( + timestamp_append, + window_field, + window_index, + )), + })); + } + WindowType::Instant => return Ok(timestamp_append), + }, + }; + + if is_nested { + return Self::nested_final_projection( + timestamp_append, + window_field, + window_index, + width, + ); + } + + let timestamp_column = + Column::new(timestamp_field.qualifier().cloned(), timestamp_field.name()); + aggregate_fields.insert(window_index, window_field.clone()); + + let window_expression = Self::build_window_struct_expr(×tamp_column, width); + aggregate_expressions.insert( + window_index, + window_expression + .alias_qualified(window_field.qualifier().cloned(), window_field.name()), + ); + aggregate_fields.push(timestamp_field); + + let bin_end_calculation = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(timestamp_column.clone())), + op: logical_expr::Operator::Plus, + right: Box::new(Expr::Literal( + ScalarValue::IntervalMonthDayNano(Some( + datafusion::arrow::datatypes::IntervalMonthDayNanoType::make_value( + 0, + 0, + (width.as_nanos() - 1) as i64, + ), + )), + None, + )), + }); + aggregate_expressions.push(bin_end_calculation); + + Ok(LogicalPlan::Projection( + logical_expr::Projection::try_new_with_schema( + aggregate_expressions, + Arc::new(timestamp_append), + Arc::new(schema_from_df_fields(&aggregate_fields)?), + )?, + )) + } + + fn build_window_struct_expr(timestamp_column: &Column, width: Duration) -> Expr { + let start_expr = Expr::Column(timestamp_column.clone()); + let end_expr = Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(timestamp_column.clone())), + op: logical_expr::Operator::Plus, + right: Box::new(Expr::Literal( + ScalarValue::IntervalMonthDayNano(Some( + datafusion::arrow::datatypes::IntervalMonthDayNanoType::make_value( + 0, + 0, + width.as_nanos() as i64, + ), + )), + None, + )), + }); + + Expr::ScalarFunction(ScalarFunction { + func: Arc::new(datafusion::logical_expr::ScalarUDF::new_from_impl( + WindowStructUdf {}, + )), + args: vec![start_expr, end_expr], + }) + } + + fn nested_final_projection( + aggregate_plan: LogicalPlan, + window_field: DFField, + window_index: usize, + width: Duration, + ) -> Result { + let timestamp_field: DFField = aggregate_plan + .schema() + .qualified_field_with_unqualified_name(TIMESTAMP_FIELD) + .unwrap() + .into(); + let timestamp_column = + Column::new(timestamp_field.qualifier().cloned(), timestamp_field.name()); + + let mut aggregate_fields = fields_with_qualifiers(aggregate_plan.schema()); + let mut aggregate_expressions: Vec<_> = aggregate_fields + .iter() + .map(|field| Expr::Column(field.qualified_column())) + .collect(); + aggregate_fields.insert(window_index, window_field.clone()); + + let window_expression = Self::build_window_struct_expr(×tamp_column, width); + aggregate_expressions.insert( + window_index, + window_expression + .alias_qualified(window_field.qualifier().cloned(), window_field.name()), + ); + + Ok(LogicalPlan::Projection( + logical_expr::Projection::try_new_with_schema( + aggregate_expressions, + Arc::new(aggregate_plan), + Arc::new(schema_from_df_fields(&aggregate_fields).unwrap()), + ) + .unwrap(), + )) + } +} + +impl UserDefinedLogicalNodeCore for AggregateExtension { + fn name(&self) -> &str { + AGGREGATE_EXTENSION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.aggregate] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "AggregateExtension: {} | window_behavior: {:?}", + self.schema(), + match &self.window_behavior { + WindowBehavior::InData => "InData".to_string(), + WindowBehavior::FromOperator { window, .. } => format!("FromOperator({window:?})"), + } + ) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + if inputs.len() != 1 { + return internal_err!("input size inconsistent"); + } + Ok(Self::new( + self.window_behavior.clone(), + inputs[0].clone(), + self.key_fields.clone(), + )) + } +} + +impl StreamExtension for AggregateExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + let output_schema = (*self.schema).clone().into(); + StreamSchema::from_schema_keys(Arc::new(output_schema), vec![]).unwrap() + } +} + +/// Extension for appending window struct (start, end) to the output +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct WindowAppendExtension { + pub(crate) input: LogicalPlan, + pub(crate) window_field: DFField, + pub(crate) window_index: usize, + pub(crate) schema: DFSchemaRef, +} + +multifield_partial_ord!(WindowAppendExtension, input, window_index); + +impl WindowAppendExtension { + fn new(input: LogicalPlan, window_field: DFField, window_index: usize) -> Self { + let mut fields = fields_with_qualifiers(input.schema()); + fields.insert(window_index, window_field.clone()); + let metadata = input.schema().metadata().clone(); + Self { + input, + window_field, + window_index, + schema: Arc::new(schema_from_df_fields_with_metadata(&fields, metadata).unwrap()), + } + } +} + +impl UserDefinedLogicalNodeCore for WindowAppendExtension { + fn name(&self) -> &str { + "WindowAppendExtension" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "WindowAppendExtension: field {:?} at {}", + self.window_field, self.window_index + ) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Ok(Self::new( + inputs[0].clone(), + self.window_field.clone(), + self.window_index, + )) + } +} + +/// Placeholder UDF to construct the window struct at plan time +#[derive(Debug)] +struct WindowStructUdf; + +impl datafusion::logical_expr::ScalarUDFImpl for WindowStructUdf { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn name(&self) -> &str { + "window" + } + + fn signature(&self) -> &datafusion::logical_expr::Signature { + &datafusion::logical_expr::Signature { + type_signature: datafusion::logical_expr::TypeSignature::Any(2), + volatility: datafusion::logical_expr::Volatility::Immutable, + } + } + + fn return_type(&self, _args: &[DataType]) -> Result { + Ok(crate::sql::planner::schemas::window_arrow_struct()) + } + + fn invoke_with_args( + &self, + _args: datafusion::logical_expr::ScalarFunctionArgs, + ) -> Result { + unimplemented!("WindowStructUdf is a plan-time-only function") + } +} diff --git a/src/sql/planner/extension/debezium.rs b/src/sql/planner/extension/debezium.rs new file mode 100644 index 00000000..1760533c --- /dev/null +++ b/src/sql/planner/extension/debezium.rs @@ -0,0 +1,250 @@ +use std::sync::Arc; + +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::common::{DFSchema, DFSchemaRef, Result, TableReference, plan_err}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; + +use super::{NamedNode, StreamExtension}; +use crate::multifield_partial_ord; +use crate::sql::types::{StreamSchema, TIMESTAMP_FIELD}; + +pub(crate) const DEBEZIUM_UNROLLING_EXTENSION_NAME: &str = "DebeziumUnrollingExtension"; +pub(crate) const TO_DEBEZIUM_EXTENSION_NAME: &str = "ToDebeziumExtension"; + +/// Unrolls a Debezium-formatted (before/after/op) stream into individual rows +/// with an updating metadata column. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct DebeziumUnrollingExtension { + pub(crate) input: LogicalPlan, + pub(crate) schema: DFSchemaRef, + pub primary_keys: Vec, + primary_key_names: Arc>, +} + +multifield_partial_ord!( + DebeziumUnrollingExtension, + input, + primary_keys, + primary_key_names +); + +impl DebeziumUnrollingExtension { + pub(crate) fn as_debezium_schema( + input_schema: &DFSchemaRef, + qualifier: Option, + ) -> Result { + let timestamp_field = if input_schema.has_column_with_unqualified_name(TIMESTAMP_FIELD) { + Some( + input_schema + .field_with_unqualified_name(TIMESTAMP_FIELD)? + .clone(), + ) + } else { + None + }; + let struct_schema: Vec<_> = input_schema + .fields() + .iter() + .filter(|field| field.name() != TIMESTAMP_FIELD) + .cloned() + .collect(); + + let struct_type = DataType::Struct(struct_schema.into()); + + let before = Arc::new(Field::new("before", struct_type.clone(), true)); + let after = Arc::new(Field::new("after", struct_type, true)); + let op = Arc::new(Field::new("op", DataType::Utf8, true)); + let mut fields = vec![before, after, op]; + + if let Some(ts) = timestamp_field { + fields.push(Arc::new(ts)); + } + + let schema = match qualifier { + Some(q) => DFSchema::try_from_qualified_schema(q, &Schema::new(fields))?, + None => DFSchema::try_from(Schema::new(fields))?, + }; + Ok(Arc::new(schema)) + } + + pub fn try_new(input: LogicalPlan, primary_keys: Arc>) -> Result { + let input_schema = input.schema(); + + let Some(before_index) = input_schema.index_of_column_by_name(None, "before") else { + return plan_err!("DebeziumUnrollingExtension requires a before column"); + }; + let Some(after_index) = input_schema.index_of_column_by_name(None, "after") else { + return plan_err!("DebeziumUnrollingExtension requires an after column"); + }; + let Some(op_index) = input_schema.index_of_column_by_name(None, "op") else { + return plan_err!("DebeziumUnrollingExtension requires an op column"); + }; + + let before_type = input_schema.field(before_index).data_type(); + let after_type = input_schema.field(after_index).data_type(); + if before_type != after_type { + return plan_err!( + "before and after columns must have the same type, not {} and {}", + before_type, + after_type + ); + } + + let op_type = input_schema.field(op_index).data_type(); + if *op_type != DataType::Utf8 { + return plan_err!("op column must be a string, not {}", op_type); + } + + let DataType::Struct(fields) = before_type else { + return plan_err!( + "before and after columns must be structs, not {}", + before_type + ); + }; + + let primary_key_idx = primary_keys + .iter() + .map(|pk| fields.find(pk).map(|(i, _)| i)) + .collect::>>() + .ok_or_else(|| { + datafusion::error::DataFusionError::Plan( + "primary key field not found in Debezium schema".to_string(), + ) + })?; + + let qualifier = match ( + input_schema.qualified_field(before_index).0, + input_schema.qualified_field(after_index).0, + ) { + (Some(bq), Some(aq)) => { + if bq != aq { + return plan_err!("before and after columns must have the same alias"); + } + Some(bq.clone()) + } + (None, None) => None, + _ => return plan_err!("before and after columns must both have an alias or neither"), + }; + + let mut out_fields = fields.to_vec(); + + let Some(input_ts_index) = input_schema.index_of_column_by_name(None, TIMESTAMP_FIELD) + else { + return plan_err!("DebeziumUnrollingExtension requires a timestamp field"); + }; + out_fields.push(Arc::new(input_schema.field(input_ts_index).clone())); + + let arrow_schema = Schema::new(out_fields); + let schema = match qualifier { + Some(q) => DFSchema::try_from_qualified_schema(q, &arrow_schema)?, + None => DFSchema::try_from(arrow_schema)?, + }; + + Ok(Self { + input, + schema: Arc::new(schema), + primary_keys: primary_key_idx, + primary_key_names: primary_keys, + }) + } +} + +impl UserDefinedLogicalNodeCore for DebeziumUnrollingExtension { + fn name(&self) -> &str { + DEBEZIUM_UNROLLING_EXTENSION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "DebeziumUnrollingExtension") + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Self::try_new(inputs[0].clone(), self.primary_key_names.clone()) + } +} + +impl StreamExtension for DebeziumUnrollingExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_unkeyed(Arc::new(self.schema.as_ref().into())).unwrap() + } + + fn transparent(&self) -> bool { + true + } +} + +/// Wraps an input stream into Debezium format (before/after/op) for updating sinks. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct ToDebeziumExtension { + pub(crate) input: Arc, + pub(crate) schema: DFSchemaRef, +} + +multifield_partial_ord!(ToDebeziumExtension, input); + +impl ToDebeziumExtension { + pub(crate) fn try_new(input: LogicalPlan) -> Result { + let schema = DebeziumUnrollingExtension::as_debezium_schema(input.schema(), None) + .expect("should be able to create ToDebeziumExtension"); + Ok(Self { + input: Arc::new(input), + schema, + }) + } +} + +impl UserDefinedLogicalNodeCore for ToDebeziumExtension { + fn name(&self) -> &str { + TO_DEBEZIUM_EXTENSION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "ToDebeziumExtension") + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Self::try_new(inputs[0].clone()) + } +} + +impl StreamExtension for ToDebeziumExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_unkeyed(Arc::new(self.schema.as_ref().into())).unwrap() + } + + fn transparent(&self) -> bool { + true + } +} diff --git a/src/sql/planner/extension/join.rs b/src/sql/planner/extension/join.rs new file mode 100644 index 00000000..3857fee7 --- /dev/null +++ b/src/sql/planner/extension/join.rs @@ -0,0 +1,61 @@ +use std::time::Duration; + +use datafusion::common::{DFSchemaRef, Result}; +use datafusion::logical_expr::expr::Expr; +use datafusion::logical_expr::{LogicalPlan, UserDefinedLogicalNodeCore}; + +use crate::sql::planner::extension::{NamedNode, StreamExtension}; +use crate::sql::types::StreamSchema; + +use std::sync::Arc; + +pub(crate) const JOIN_NODE_NAME: &str = "JoinNode"; + +/// Extension node for streaming joins. +/// Supports instant joins (windowed, no state) and updating joins (with TTL-based state). +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] +pub struct JoinExtension { + pub(crate) rewritten_join: LogicalPlan, + pub(crate) is_instant: bool, + pub(crate) ttl: Option, +} + +impl StreamExtension for JoinExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_unkeyed(Arc::new(self.schema().as_ref().into())).unwrap() + } +} + +impl UserDefinedLogicalNodeCore for JoinExtension { + fn name(&self) -> &str { + JOIN_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.rewritten_join] + } + + fn schema(&self) -> &DFSchemaRef { + self.rewritten_join.schema() + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "JoinExtension: {}", self.schema()) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Ok(Self { + rewritten_join: inputs[0].clone(), + is_instant: self.is_instant, + ttl: self.ttl, + }) + } +} diff --git a/src/sql/planner/extension/key_calculation.rs b/src/sql/planner/extension/key_calculation.rs new file mode 100644 index 00000000..c90b6d1d --- /dev/null +++ b/src/sql/planner/extension/key_calculation.rs @@ -0,0 +1,138 @@ +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::arrow::datatypes::{Field, Schema}; +use datafusion::common::{DFSchemaRef, Result, internal_err}; +use datafusion::logical_expr::{Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore}; + +use crate::multifield_partial_ord; +use crate::sql::planner::extension::{NamedNode, StreamExtension}; +use crate::sql::types::{ + StreamSchema, fields_with_qualifiers, schema_from_df_fields_with_metadata, +}; + +pub(crate) const KEY_CALCULATION_NAME: &str = "KeyCalculationExtension"; + +/// Two ways of specifying keys: column indices or expressions to evaluate +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] +pub enum KeysOrExprs { + Keys(Vec), + Exprs(Vec), +} + +/// Calculation for computing keyed data, used for shuffling data to correct nodes +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct KeyCalculationExtension { + pub(crate) name: Option, + pub(crate) input: LogicalPlan, + pub(crate) keys: KeysOrExprs, + pub(crate) schema: DFSchemaRef, +} + +multifield_partial_ord!(KeyCalculationExtension, name, input, keys); + +impl KeyCalculationExtension { + pub fn new_named_and_trimmed(input: LogicalPlan, keys: Vec, name: String) -> Self { + let output_fields: Vec<_> = fields_with_qualifiers(input.schema()) + .into_iter() + .enumerate() + .filter_map(|(index, field)| { + if !keys.contains(&index) { + Some(field.clone()) + } else { + None + } + }) + .collect(); + + let schema = + schema_from_df_fields_with_metadata(&output_fields, input.schema().metadata().clone()) + .unwrap(); + Self { + name: Some(name), + input, + keys: KeysOrExprs::Keys(keys), + schema: Arc::new(schema), + } + } + + pub fn new(input: LogicalPlan, keys: KeysOrExprs) -> Self { + let schema = input.schema().clone(); + Self { + name: None, + input, + keys, + schema, + } + } +} + +impl StreamExtension for KeyCalculationExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + let input_schema = self.input.schema().as_ref(); + match &self.keys { + KeysOrExprs::Keys(keys) => { + StreamSchema::from_schema_keys(Arc::new(input_schema.into()), keys.clone()).unwrap() + } + KeysOrExprs::Exprs(exprs) => { + let mut fields = vec![]; + for (i, e) in exprs.iter().enumerate() { + let (dt, nullable) = e.data_type_and_nullable(input_schema).unwrap(); + fields.push(Field::new(format!("__key_{i}"), dt, nullable).into()); + } + for f in input_schema.fields().iter() { + fields.push(f.clone()); + } + StreamSchema::from_schema_keys( + Arc::new(Schema::new(fields)), + (1..=exprs.len()).collect(), + ) + .unwrap() + } + } + } +} + +impl UserDefinedLogicalNodeCore for KeyCalculationExtension { + fn name(&self) -> &str { + KEY_CALCULATION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "KeyCalculationExtension: {}", self.schema()) + } + + fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> Result { + if inputs.len() != 1 { + return internal_err!("input size inconsistent"); + } + + let keys = match &self.keys { + KeysOrExprs::Keys(k) => KeysOrExprs::Keys(k.clone()), + KeysOrExprs::Exprs(_) => KeysOrExprs::Exprs(exprs), + }; + + Ok(Self { + name: self.name.clone(), + input: inputs[0].clone(), + keys, + schema: self.schema.clone(), + }) + } +} diff --git a/src/sql/planner/extension/lookup.rs b/src/sql/planner/extension/lookup.rs new file mode 100644 index 00000000..daa4b094 --- /dev/null +++ b/src/sql/planner/extension/lookup.rs @@ -0,0 +1,127 @@ +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{Column, DFSchemaRef, JoinType, Result, TableReference, internal_err}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; + +use super::{NamedNode, StreamExtension}; +use crate::multifield_partial_ord; +use crate::sql::catalog::connector_table::ConnectorTable; +use crate::sql::types::StreamSchema; + +pub const SOURCE_EXTENSION_NAME: &str = "LookupSource"; +pub const JOIN_EXTENSION_NAME: &str = "LookupJoin"; + +/// Represents a lookup table source in the streaming plan. +/// Lookup sources provide point-query access to external state. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct LookupSource { + pub(crate) table: ConnectorTable, + pub(crate) schema: DFSchemaRef, +} + +multifield_partial_ord!(LookupSource, table); + +impl UserDefinedLogicalNodeCore for LookupSource { + fn name(&self) -> &str { + SOURCE_EXTENSION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "LookupSource: {}", self.schema) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + if !inputs.is_empty() { + return internal_err!("LookupSource cannot have inputs"); + } + Ok(Self { + table: self.table.clone(), + schema: self.schema.clone(), + }) + } +} + +impl StreamExtension for LookupSource { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_unkeyed(Arc::new(self.schema.as_ref().into())).unwrap() + } +} + +/// Represents a lookup join: a streaming input joined against a lookup table. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct LookupJoin { + pub(crate) input: LogicalPlan, + pub(crate) schema: DFSchemaRef, + pub(crate) connector: ConnectorTable, + pub(crate) on: Vec<(Expr, Column)>, + pub(crate) filter: Option, + pub(crate) alias: Option, + pub(crate) join_type: JoinType, +} + +multifield_partial_ord!(LookupJoin, input, connector, on, filter, alias); + +impl UserDefinedLogicalNodeCore for LookupJoin { + fn name(&self) -> &str { + JOIN_EXTENSION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + let mut e: Vec<_> = self.on.iter().map(|(l, _)| l.clone()).collect(); + if let Some(filter) = &self.filter { + e.push(filter.clone()); + } + e + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "LookupJoinExtension: {}", self.schema) + } + + fn with_exprs_and_inputs(&self, _: Vec, inputs: Vec) -> Result { + Ok(Self { + input: inputs[0].clone(), + schema: self.schema.clone(), + connector: self.connector.clone(), + on: self.on.clone(), + filter: self.filter.clone(), + alias: self.alias.clone(), + join_type: self.join_type, + }) + } +} + +impl StreamExtension for LookupJoin { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_unkeyed(Arc::new(self.schema.as_ref().into())).unwrap() + } +} diff --git a/src/sql/planner/extension/mod.rs b/src/sql/planner/extension/mod.rs new file mode 100644 index 00000000..4de1892e --- /dev/null +++ b/src/sql/planner/extension/mod.rs @@ -0,0 +1,356 @@ +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; +use std::time::Duration; + +use datafusion::arrow::datatypes::{DataType, TimeUnit}; +use datafusion::common::{DFSchemaRef, DataFusionError, Result, TableReference}; +use datafusion::logical_expr::{ + Expr, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, +}; + +use crate::datastream::logical::{LogicalEdge, LogicalNode}; +use crate::sql::planner::schemas::{add_timestamp_field, has_timestamp_field}; +use crate::sql::types::{ + DFField, StreamSchema, TIMESTAMP_FIELD, fields_with_qualifiers, schema_from_df_fields, +}; +use crate::types::FsSchemaRef; + +pub(crate) mod aggregate; +pub(crate) mod debezium; +pub(crate) mod join; +pub(crate) mod key_calculation; +pub(crate) mod lookup; +pub(crate) mod projection; +pub(crate) mod remote_table; +pub(crate) mod sink; +pub(crate) mod table_source; +pub(crate) mod updating_aggregate; +pub(crate) mod watermark_node; +pub(crate) mod window_fn; + +pub(crate) struct NodeWithIncomingEdges { + pub node: LogicalNode, + pub edges: Vec, +} + +pub(crate) trait StreamExtension: Debug { + fn node_name(&self) -> Option; + + fn plan_node( + &self, + _planner: &super::physical_planner::Planner, + _index: usize, + _input_schemas: Vec, + ) -> Result { + Err(DataFusionError::NotImplemented(format!( + "plan_node not yet implemented for {:?}", + self + ))) + } + + fn output_schema(&self) -> StreamSchema; + fn transparent(&self) -> bool { + false + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum NamedNode { + Source(TableReference), + Watermark(TableReference), + RemoteTable(TableReference), + Sink(TableReference), +} + +fn try_from_t( + node: &dyn UserDefinedLogicalNode, +) -> std::result::Result<&dyn StreamExtension, ()> { + node.as_any() + .downcast_ref::() + .map(|t| t as &dyn StreamExtension) + .ok_or(()) +} + +impl<'a> TryFrom<&'a dyn UserDefinedLogicalNode> for &'a dyn StreamExtension { + type Error = DataFusionError; + + fn try_from(node: &'a dyn UserDefinedLogicalNode) -> Result { + use aggregate::AggregateExtension; + use debezium::{DebeziumUnrollingExtension, ToDebeziumExtension}; + use join::JoinExtension; + use key_calculation::KeyCalculationExtension; + use lookup::{LookupJoin, LookupSource}; + use projection::ProjectionExtension; + use remote_table::RemoteTableExtension; + use sink::SinkExtension; + use table_source::TableSourceExtension; + use updating_aggregate::UpdatingAggregateExtension; + use watermark_node::WatermarkNode; + use window_fn::WindowFunctionExtension; + + try_from_t::(node) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .or_else(|_| try_from_t::(node)) + .map_err(|_| DataFusionError::Plan(format!("unexpected node: {}", node.name()))) + } +} + +impl<'a> TryFrom<&'a Arc> for &'a dyn StreamExtension { + type Error = DataFusionError; + + fn try_from(node: &'a Arc) -> Result { + TryFrom::try_from(node.as_ref()) + } +} + +#[macro_export] +macro_rules! multifield_partial_ord { + ($ty:ty, $($field:tt), *) => { + impl PartialOrd for $ty { + fn partial_cmp(&self, other: &Self) -> Option { + $( + let cmp = self.$field.partial_cmp(&other.$field)?; + if cmp != std::cmp::Ordering::Equal { + return Some(cmp); + } + )* + Some(std::cmp::Ordering::Equal) + } + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct TimestampAppendExtension { + pub(crate) input: LogicalPlan, + pub(crate) qualifier: Option, + pub(crate) schema: DFSchemaRef, +} + +impl TimestampAppendExtension { + pub(crate) fn new(input: LogicalPlan, qualifier: Option) -> Self { + if has_timestamp_field(input.schema()) { + unreachable!( + "shouldn't be adding timestamp to a plan that already has it: {:?}", + input.schema() + ); + } + let schema = add_timestamp_field(input.schema().clone(), qualifier.clone()).unwrap(); + Self { + input, + qualifier, + schema, + } + } +} + +multifield_partial_ord!(TimestampAppendExtension, input, qualifier); + +impl UserDefinedLogicalNodeCore for TimestampAppendExtension { + fn name(&self) -> &str { + "TimestampAppendExtension" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!( + f, + "TimestampAppendExtension({:?}): {}", + self.qualifier, + self.schema + .fields() + .iter() + .map(|f| f.name().to_string()) + .collect::>() + .join(", ") + ) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Ok(Self::new(inputs[0].clone(), self.qualifier.clone())) + } +} + +/// Appends an `_updating_meta` and properly qualified `_timestamp` field +/// to the output schema of an updating aggregate. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct IsRetractExtension { + pub(crate) input: LogicalPlan, + pub(crate) schema: DFSchemaRef, + pub(crate) timestamp_qualifier: Option, +} + +multifield_partial_ord!(IsRetractExtension, input, timestamp_qualifier); + +impl IsRetractExtension { + pub(crate) fn new(input: LogicalPlan, timestamp_qualifier: Option) -> Self { + let mut output_fields = fields_with_qualifiers(input.schema()); + + let timestamp_index = output_fields.len() - 1; + output_fields[timestamp_index] = DFField::new( + timestamp_qualifier.clone(), + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ); + let schema = Arc::new(schema_from_df_fields(&output_fields).unwrap()); + Self { + input, + schema, + timestamp_qualifier, + } + } +} + +impl UserDefinedLogicalNodeCore for IsRetractExtension { + fn name(&self) -> &str { + "IsRetractExtension" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "IsRetractExtension") + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Ok(Self::new( + inputs[0].clone(), + self.timestamp_qualifier.clone(), + )) + } +} + +impl StreamExtension for IsRetractExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_unkeyed(Arc::new(self.schema.as_ref().into())).unwrap() + } +} + +pub(crate) const ASYNC_RESULT_FIELD: &str = "__async_result"; + +/// Extension node for async UDF calls in streaming projections. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct AsyncUDFExtension { + pub(crate) input: Arc, + pub(crate) name: String, + pub(crate) arg_exprs: Vec, + pub(crate) final_exprs: Vec, + pub(crate) ordered: bool, + pub(crate) max_concurrency: usize, + pub(crate) timeout: Duration, + pub(crate) final_schema: DFSchemaRef, +} + +multifield_partial_ord!( + AsyncUDFExtension, + input, + name, + arg_exprs, + final_exprs, + ordered, + max_concurrency, + timeout +); + +impl UserDefinedLogicalNodeCore for AsyncUDFExtension { + fn name(&self) -> &str { + "AsyncUDFNode" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.final_schema + } + + fn expressions(&self) -> Vec { + self.arg_exprs + .iter() + .chain(self.final_exprs.iter()) + .cloned() + .collect() + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "AsyncUdfExtension<{}>: {}", self.name, self.final_schema) + } + + fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> Result { + if inputs.len() != 1 { + return Err(DataFusionError::Internal("input size inconsistent".into())); + } + if UserDefinedLogicalNode::expressions(self) != exprs { + return Err(DataFusionError::Internal( + "Tried to recreate async UDF node with different expressions".into(), + )); + } + + Ok(Self { + input: Arc::new(inputs[0].clone()), + name: self.name.clone(), + arg_exprs: self.arg_exprs.clone(), + final_exprs: self.final_exprs.clone(), + ordered: self.ordered, + max_concurrency: self.max_concurrency, + timeout: self.timeout, + final_schema: self.final_schema.clone(), + }) + } +} + +impl StreamExtension for AsyncUDFExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_fields( + self.final_schema + .fields() + .iter() + .map(|f| (**f).clone()) + .collect(), + ) + } +} diff --git a/src/sql/planner/extension/projection.rs b/src/sql/planner/extension/projection.rs new file mode 100644 index 00000000..e6dc8ce7 --- /dev/null +++ b/src/sql/planner/extension/projection.rs @@ -0,0 +1,91 @@ +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result}; +use datafusion::logical_expr::{Expr, ExprSchemable, LogicalPlan, UserDefinedLogicalNodeCore}; + +use crate::multifield_partial_ord; +use crate::sql::planner::extension::{NamedNode, StreamExtension}; +use crate::sql::types::{DFField, StreamSchema, schema_from_df_fields}; + +pub(crate) const PROJECTION_NAME: &str = "ProjectionExtension"; + +/// Projection operations for streaming SQL plans. +/// Handles column projections, shuffles for key-based operations, etc. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct ProjectionExtension { + pub(crate) inputs: Vec, + pub(crate) name: Option, + pub(crate) exprs: Vec, + pub(crate) schema: DFSchemaRef, + pub(crate) shuffle: bool, +} + +multifield_partial_ord!(ProjectionExtension, name, exprs); + +impl ProjectionExtension { + pub(crate) fn new(inputs: Vec, name: Option, exprs: Vec) -> Self { + let input_schema = inputs.first().unwrap().schema(); + let fields: Vec = exprs + .iter() + .map(|e| DFField::from(e.to_field(input_schema).unwrap())) + .collect(); + + let schema = Arc::new(schema_from_df_fields(&fields).unwrap()); + + Self { + inputs, + name, + exprs, + schema, + shuffle: false, + } + } + + pub(crate) fn shuffled(mut self) -> Self { + self.shuffle = true; + self + } +} + +impl StreamExtension for ProjectionExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_unkeyed(Arc::new(self.schema.as_arrow().clone())).unwrap() + } +} + +impl UserDefinedLogicalNodeCore for ProjectionExtension { + fn name(&self) -> &str { + PROJECTION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + self.inputs.iter().collect() + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ProjectionExtension: {}", self.schema()) + } + + fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> Result { + Ok(Self { + name: self.name.clone(), + inputs, + exprs, + schema: self.schema.clone(), + shuffle: self.shuffle, + }) + } +} diff --git a/src/sql/planner/extension/remote_table.rs b/src/sql/planner/extension/remote_table.rs new file mode 100644 index 00000000..2d81cafc --- /dev/null +++ b/src/sql/planner/extension/remote_table.rs @@ -0,0 +1,71 @@ +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result, TableReference, internal_err}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; + +use crate::multifield_partial_ord; +use crate::sql::planner::extension::{NamedNode, StreamExtension}; +use crate::sql::types::StreamSchema; + +pub(crate) const REMOTE_TABLE_NAME: &str = "RemoteTableExtension"; + +/// Lightweight extension that segments the execution graph and enables merging +/// nodes with the same name. Allows materializing intermediate results. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct RemoteTableExtension { + pub(crate) input: LogicalPlan, + pub(crate) name: TableReference, + pub(crate) schema: DFSchemaRef, + pub(crate) materialize: bool, +} + +multifield_partial_ord!(RemoteTableExtension, input, name, materialize); + +impl StreamExtension for RemoteTableExtension { + fn node_name(&self) -> Option { + if self.materialize { + Some(NamedNode::RemoteTable(self.name.to_owned())) + } else { + None + } + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_keys(Arc::new(self.schema.as_ref().into()), vec![]).unwrap() + } +} + +impl UserDefinedLogicalNodeCore for RemoteTableExtension { + fn name(&self) -> &str { + REMOTE_TABLE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "RemoteTableExtension: {}", self.schema) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + if inputs.len() != 1 { + return internal_err!("input size inconsistent"); + } + Ok(Self { + input: inputs[0].clone(), + name: self.name.clone(), + schema: self.schema.clone(), + materialize: self.materialize, + }) + } +} diff --git a/src/sql/planner/extension/sink.rs b/src/sql/planner/extension/sink.rs new file mode 100644 index 00000000..7820925f --- /dev/null +++ b/src/sql/planner/extension/sink.rs @@ -0,0 +1,135 @@ +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result, TableReference, plan_err}; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; + +use super::debezium::ToDebeziumExtension; +use super::remote_table::RemoteTableExtension; +use super::{NamedNode, StreamExtension}; +use crate::multifield_partial_ord; +use crate::sql::catalog::table::Table; +use crate::sql::types::StreamSchema; + +pub(crate) const SINK_NODE_NAME: &str = "SinkExtension"; + +/// Extension node representing a sink (output) in the streaming plan. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct SinkExtension { + pub(crate) name: TableReference, + pub(crate) table: Table, + pub(crate) schema: DFSchemaRef, + pub(crate) inputs: Arc>, +} + +multifield_partial_ord!(SinkExtension, name, inputs); + +impl SinkExtension { + pub fn new( + name: TableReference, + table: Table, + mut schema: DFSchemaRef, + mut input: Arc, + ) -> Result { + match &table { + Table::ConnectorTable(connector_table) => { + if connector_table.is_updating() { + let to_debezium = ToDebeziumExtension::try_new(input.as_ref().clone())?; + input = Arc::new(LogicalPlan::Extension(Extension { + node: Arc::new(to_debezium), + })); + schema = input.schema().clone(); + } + } + Table::LookupTable(..) => return plan_err!("cannot use a lookup table as a sink"), + Table::MemoryTable { .. } => return plan_err!("memory tables not supported as sinks"), + Table::TableFromQuery { .. } => {} + Table::PreviewSink { .. } => { + // preview sinks may also need debezium wrapping for updating inputs + } + } + + Self::add_remote_if_necessary(&schema, &mut input); + + let inputs = Arc::new(vec![(*input).clone()]); + Ok(Self { + name, + table, + schema, + inputs, + }) + } + + pub fn add_remote_if_necessary(schema: &DFSchemaRef, input: &mut Arc) { + if let LogicalPlan::Extension(node) = input.as_ref() { + let Ok(ext): Result<&dyn StreamExtension, _> = (&node.node).try_into() else { + // not a StreamExtension, wrap it + let remote = RemoteTableExtension { + input: input.as_ref().clone(), + name: TableReference::bare("sink projection"), + schema: schema.clone(), + materialize: false, + }; + *input = Arc::new(LogicalPlan::Extension(Extension { + node: Arc::new(remote), + })); + return; + }; + if !ext.transparent() { + return; + } + } + let remote = RemoteTableExtension { + input: input.as_ref().clone(), + name: TableReference::bare("sink projection"), + schema: schema.clone(), + materialize: false, + }; + *input = Arc::new(LogicalPlan::Extension(Extension { + node: Arc::new(remote), + })); + } +} + +impl UserDefinedLogicalNodeCore for SinkExtension { + fn name(&self) -> &str { + SINK_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + self.inputs.iter().collect() + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "SinkExtension({:?}): {}", self.name, self.schema) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Ok(Self { + name: self.name.clone(), + table: self.table.clone(), + schema: self.schema.clone(), + inputs: Arc::new(inputs), + }) + } +} + +impl StreamExtension for SinkExtension { + fn node_name(&self) -> Option { + match &self.table { + Table::PreviewSink { .. } => None, + _ => Some(NamedNode::Sink(self.name.clone())), + } + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_fields(vec![]) + } +} diff --git a/src/sql/planner/extension/table_source.rs b/src/sql/planner/extension/table_source.rs new file mode 100644 index 00000000..cab3ae3d --- /dev/null +++ b/src/sql/planner/extension/table_source.rs @@ -0,0 +1,94 @@ +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result, TableReference}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; + +use super::{NamedNode, StreamExtension}; +use crate::multifield_partial_ord; +use crate::sql::catalog::connector_table::ConnectorTable; +use crate::sql::catalog::field_spec::FieldSpec; +use crate::sql::planner::schemas::add_timestamp_field; +use crate::sql::types::{StreamSchema, schema_from_df_fields}; + +pub(crate) const TABLE_SOURCE_NAME: &str = "TableSourceExtension"; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) struct TableSourceExtension { + pub(crate) name: TableReference, + pub(crate) table: ConnectorTable, + pub(crate) schema: DFSchemaRef, +} + +multifield_partial_ord!(TableSourceExtension, name, table); + +impl TableSourceExtension { + pub fn new(name: TableReference, table: ConnectorTable) -> Self { + let physical_fields = table + .fields + .iter() + .filter_map(|field| match field { + FieldSpec::Struct(field) | FieldSpec::Metadata { field, .. } => { + Some((Some(name.clone()), Arc::new(field.clone())).into()) + } + FieldSpec::Virtual { .. } => None, + }) + .collect::>(); + let base_schema = Arc::new(schema_from_df_fields(&physical_fields).unwrap()); + + let schema = if table.is_updating() { + super::debezium::DebeziumUnrollingExtension::as_debezium_schema( + &base_schema, + Some(name.clone()), + ) + .unwrap() + } else { + base_schema + }; + let schema = add_timestamp_field(schema, Some(name.clone())).unwrap(); + Self { + name, + table, + schema, + } + } +} + +impl UserDefinedLogicalNodeCore for TableSourceExtension { + fn name(&self) -> &str { + TABLE_SOURCE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "TableSourceExtension: {}", self.schema) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, _inputs: Vec) -> Result { + Ok(Self { + name: self.name.clone(), + table: self.table.clone(), + schema: self.schema.clone(), + }) + } +} + +impl StreamExtension for TableSourceExtension { + fn node_name(&self) -> Option { + Some(NamedNode::Source(self.name.clone())) + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_keys(Arc::new(self.schema.as_ref().into()), vec![]).unwrap() + } +} diff --git a/src/sql/planner/extension/updating_aggregate.rs b/src/sql/planner/extension/updating_aggregate.rs new file mode 100644 index 00000000..758edc67 --- /dev/null +++ b/src/sql/planner/extension/updating_aggregate.rs @@ -0,0 +1,89 @@ +use std::sync::Arc; +use std::time::Duration; + +use datafusion::common::{DFSchemaRef, Result, TableReference, plan_err}; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; + +use super::{IsRetractExtension, NamedNode, StreamExtension}; +use crate::sql::types::StreamSchema; + +pub(crate) const UPDATING_AGGREGATE_EXTENSION_NAME: &str = "UpdatingAggregateExtension"; + +/// Extension node for updating (non-windowed) aggregations. +/// Maintains state with TTL and emits retraction/update pairs. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] +pub(crate) struct UpdatingAggregateExtension { + pub(crate) aggregate: LogicalPlan, + pub(crate) key_fields: Vec, + pub(crate) final_calculation: LogicalPlan, + pub(crate) timestamp_qualifier: Option, + pub(crate) ttl: Duration, +} + +impl UpdatingAggregateExtension { + pub fn new( + aggregate: LogicalPlan, + key_fields: Vec, + timestamp_qualifier: Option, + ttl: Duration, + ) -> Result { + let final_calculation = LogicalPlan::Extension(Extension { + node: Arc::new(IsRetractExtension::new( + aggregate.clone(), + timestamp_qualifier.clone(), + )), + }); + + Ok(Self { + aggregate, + key_fields, + final_calculation, + timestamp_qualifier, + ttl, + }) + } +} + +impl UserDefinedLogicalNodeCore for UpdatingAggregateExtension { + fn name(&self) -> &str { + UPDATING_AGGREGATE_EXTENSION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.aggregate] + } + + fn schema(&self) -> &DFSchemaRef { + self.final_calculation.schema() + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "UpdatingAggregateExtension") + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + if inputs.len() != 1 { + return plan_err!("UpdatingAggregateExtension expects exactly one input"); + } + Self::new( + inputs[0].clone(), + self.key_fields.clone(), + self.timestamp_qualifier.clone(), + self.ttl, + ) + } +} + +impl StreamExtension for UpdatingAggregateExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_unkeyed(Arc::new(self.schema().as_ref().into())).unwrap() + } +} diff --git a/src/sql/planner/extension/watermark_node.rs b/src/sql/planner/extension/watermark_node.rs new file mode 100644 index 00000000..a06bdb9a --- /dev/null +++ b/src/sql/planner/extension/watermark_node.rs @@ -0,0 +1,110 @@ +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result, TableReference, internal_err}; +use datafusion::error::DataFusionError; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; + +use crate::multifield_partial_ord; +use crate::sql::planner::extension::{NamedNode, StreamExtension}; +use crate::sql::planner::schemas::add_timestamp_field; +use crate::sql::types::{StreamSchema, TIMESTAMP_FIELD}; + +pub(crate) const WATERMARK_NODE_NAME: &str = "WatermarkNode"; + +/// Represents a watermark node in the streaming query plan. +/// Watermarks track event-time progress and enable time-based operations. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct WatermarkNode { + pub input: LogicalPlan, + pub qualifier: TableReference, + pub watermark_expression: Expr, + pub schema: DFSchemaRef, + timestamp_index: usize, +} + +multifield_partial_ord!( + WatermarkNode, + input, + qualifier, + watermark_expression, + timestamp_index +); + +impl UserDefinedLogicalNodeCore for WatermarkNode { + fn name(&self) -> &str { + WATERMARK_NODE_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.input] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![self.watermark_expression.clone()] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "WatermarkNode({}): {}", self.qualifier, self.schema) + } + + fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> Result { + if inputs.len() != 1 { + return internal_err!("input size inconsistent"); + } + if exprs.len() != 1 { + return internal_err!("expected one expression; found {}", exprs.len()); + } + + let timestamp_index = self + .schema + .index_of_column_by_name(Some(&self.qualifier), TIMESTAMP_FIELD) + .ok_or_else(|| DataFusionError::Plan("missing timestamp column".to_string()))?; + + Ok(Self { + input: inputs[0].clone(), + qualifier: self.qualifier.clone(), + watermark_expression: exprs.into_iter().next().unwrap(), + schema: self.schema.clone(), + timestamp_index, + }) + } +} + +impl StreamExtension for WatermarkNode { + fn node_name(&self) -> Option { + Some(NamedNode::Watermark(self.qualifier.clone())) + } + + fn output_schema(&self) -> StreamSchema { + self.stream_schema() + } +} + +impl WatermarkNode { + pub(crate) fn new( + input: LogicalPlan, + qualifier: TableReference, + watermark_expression: Expr, + ) -> Result { + let schema = add_timestamp_field(input.schema().clone(), Some(qualifier.clone()))?; + let timestamp_index = schema + .index_of_column_by_name(None, TIMESTAMP_FIELD) + .ok_or_else(|| DataFusionError::Plan("missing _timestamp column".to_string()))?; + Ok(Self { + input, + qualifier, + watermark_expression, + schema, + timestamp_index, + }) + } + + pub(crate) fn stream_schema(&self) -> StreamSchema { + StreamSchema::new_unkeyed(Arc::new(self.schema.as_ref().into()), self.timestamp_index) + } +} diff --git a/src/sql/planner/extension/window_fn.rs b/src/sql/planner/extension/window_fn.rs new file mode 100644 index 00000000..95832183 --- /dev/null +++ b/src/sql/planner/extension/window_fn.rs @@ -0,0 +1,62 @@ +use std::sync::Arc; + +use datafusion::common::{DFSchemaRef, Result}; +use datafusion::logical_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; + +use crate::sql::planner::extension::{NamedNode, StreamExtension}; +use crate::sql::types::StreamSchema; + +pub(crate) const WINDOW_FUNCTION_EXTENSION_NAME: &str = "WindowFunctionExtension"; + +/// Extension for window functions (e.g., ROW_NUMBER, RANK) over windowed input. +/// Window functions require already-windowed input and are evaluated per-window. +#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd)] +pub(crate) struct WindowFunctionExtension { + pub(crate) window_plan: LogicalPlan, + pub(crate) key_fields: Vec, +} + +impl WindowFunctionExtension { + pub fn new(window_plan: LogicalPlan, key_fields: Vec) -> Self { + Self { + window_plan, + key_fields, + } + } +} + +impl UserDefinedLogicalNodeCore for WindowFunctionExtension { + fn name(&self) -> &str { + WINDOW_FUNCTION_EXTENSION_NAME + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![&self.window_plan] + } + + fn schema(&self) -> &DFSchemaRef { + self.window_plan.schema() + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "WindowFunction: {}", self.schema()) + } + + fn with_exprs_and_inputs(&self, _exprs: Vec, inputs: Vec) -> Result { + Ok(Self::new(inputs[0].clone(), self.key_fields.clone())) + } +} + +impl StreamExtension for WindowFunctionExtension { + fn node_name(&self) -> Option { + None + } + + fn output_schema(&self) -> StreamSchema { + StreamSchema::from_schema_unkeyed(Arc::new(self.schema().as_ref().clone().into())).unwrap() + } +} diff --git a/src/sql/planner/mod.rs b/src/sql/planner/mod.rs new file mode 100644 index 00000000..d80d3a8d --- /dev/null +++ b/src/sql/planner/mod.rs @@ -0,0 +1,361 @@ +#![allow(clippy::new_without_default)] + +pub(crate) mod extension; +pub mod parse; +pub(crate) mod physical_planner; +pub mod plan; +pub mod rewrite; +pub mod schema_provider; +pub mod schemas; +pub mod sql_to_plan; + +pub(crate) mod mod_prelude { + pub use super::StreamSchemaProvider; +} + +pub use schema_provider::{LogicalBatchInput, StreamSchemaProvider, StreamTable}; + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use datafusion::common::tree_node::TreeNode; +use datafusion::common::{Result, plan_err}; +use datafusion::error::DataFusionError; +use datafusion::execution::SessionStateBuilder; +use datafusion::logical_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion::prelude::SessionConfig; +use datafusion::sql::TableReference; +use datafusion::sql::sqlparser::ast::{OneOrManyWithParens, Statement}; +use datafusion::sql::sqlparser::dialect::FunctionStreamDialect; +use datafusion::sql::sqlparser::parser::Parser; +use tracing::debug; + +use crate::datastream::logical::{LogicalProgram, ProgramConfig}; +use crate::datastream::optimizers::ChainingOptimizer; +use crate::sql::catalog::insert::Insert; +use crate::sql::catalog::table::Table as CatalogTable; +use crate::sql::functions::{is_json_union, serialize_outgoing_json}; +use crate::sql::planner::extension::key_calculation::{KeyCalculationExtension, KeysOrExprs}; +use crate::sql::planner::extension::projection::ProjectionExtension; +use crate::sql::planner::extension::sink::SinkExtension; +use crate::sql::planner::extension::{NamedNode, StreamExtension}; +use crate::sql::planner::plan::rewrite_plan; +use crate::sql::planner::rewrite::{SinkInputRewriter, SourceMetadataVisitor}; +use crate::sql::types::SqlConfig; + +// ── Compilation pipeline ────────────────────────────────────────────── + +#[derive(Clone, Debug)] +pub struct CompiledSql { + pub program: LogicalProgram, + pub connection_ids: Vec, +} + +pub fn parse_sql_statements( + sql: &str, +) -> std::result::Result, datafusion::sql::sqlparser::parser::ParserError> { + Parser::parse_sql(&FunctionStreamDialect {}, sql) +} + +fn try_handle_set_variable( + statement: &Statement, + schema_provider: &mut StreamSchemaProvider, +) -> Result { + if let Statement::SetVariable { + variables, value, .. + } = statement + { + let OneOrManyWithParens::One(opt) = variables else { + return plan_err!("invalid syntax for `SET` call"); + }; + + if opt.to_string() != "updating_ttl" { + return plan_err!( + "invalid option '{}'; supported options are 'updating_ttl'", + opt + ); + } + + if value.len() != 1 { + return plan_err!("invalid `SET updating_ttl` call; expected exactly one expression"); + } + + let duration = duration_from_sql_expr(&value[0])?; + schema_provider.planning_options.ttl = duration; + + return Ok(true); + } + + Ok(false) +} + +fn duration_from_sql_expr( + expr: &datafusion::sql::sqlparser::ast::Expr, +) -> Result { + use datafusion::sql::sqlparser::ast::Expr as SqlExpr; + use datafusion::sql::sqlparser::ast::Value as SqlValue; + use datafusion::sql::sqlparser::ast::ValueWithSpan; + + match expr { + SqlExpr::Interval(interval) => { + let value_str = match interval.value.as_ref() { + SqlExpr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + .. + }) => s.clone(), + other => return plan_err!("expected interval string literal, found {other}"), + }; + + parse_interval_to_duration(&value_str) + } + SqlExpr::Value(ValueWithSpan { + value: SqlValue::SingleQuotedString(s), + .. + }) => parse_interval_to_duration(s), + other => plan_err!("expected an interval expression, found {other}"), + } +} + +fn parse_interval_to_duration(s: &str) -> Result { + let parts: Vec<&str> = s.trim().split_whitespace().collect(); + if parts.len() != 2 { + return plan_err!("invalid interval string '{s}'; expected ' '"); + } + let value: u64 = parts[0] + .parse() + .map_err(|_| DataFusionError::Plan(format!("invalid interval number: {}", parts[0])))?; + match parts[1].to_lowercase().as_str() { + "second" | "seconds" | "s" => Ok(std::time::Duration::from_secs(value)), + "minute" | "minutes" | "min" => Ok(std::time::Duration::from_secs(value * 60)), + "hour" | "hours" | "h" => Ok(std::time::Duration::from_secs(value * 3600)), + "day" | "days" | "d" => Ok(std::time::Duration::from_secs(value * 86400)), + unit => plan_err!("unsupported interval unit '{unit}'"), + } +} + +fn build_sink_inputs(extensions: &[LogicalPlan]) -> HashMap> { + let mut sink_inputs = HashMap::>::new(); + for extension in extensions.iter() { + if let LogicalPlan::Extension(ext) = extension { + if let Some(sink_node) = ext.node.as_any().downcast_ref::() { + if let Some(named_node) = sink_node.node_name() { + let inputs = sink_node + .inputs() + .into_iter() + .cloned() + .collect::>(); + sink_inputs.entry(named_node).or_default().extend(inputs); + } + } + } + } + sink_inputs +} + +fn maybe_add_key_extension_to_sink(plan: LogicalPlan) -> Result { + let LogicalPlan::Extension(ref ext) = plan else { + return Ok(plan); + }; + + let Some(sink) = ext.node.as_any().downcast_ref::() else { + return Ok(plan); + }; + + let Some(partition_exprs) = sink.table.partition_exprs() else { + return Ok(plan); + }; + + if partition_exprs.is_empty() { + return Ok(plan); + } + + let inputs = plan + .inputs() + .into_iter() + .map(|input| { + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(KeyCalculationExtension { + name: Some("key-calc-partition".to_string()), + schema: input.schema().clone(), + input: input.clone(), + keys: KeysOrExprs::Exprs(partition_exprs.clone()), + }), + })) + }) + .collect::>()?; + + use datafusion::prelude::col; + let unkey = LogicalPlan::Extension(Extension { + node: Arc::new( + ProjectionExtension::new( + inputs, + Some("unkey".to_string()), + sink.schema().iter().map(|(_, f)| col(f.name())).collect(), + ) + .shuffled(), + ), + }); + + let node = sink.with_exprs_and_inputs(vec![], vec![unkey])?; + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(node), + })) +} + +fn rewrite_sinks(extensions: Vec) -> Result> { + let mut sink_inputs = build_sink_inputs(&extensions); + let mut new_extensions = vec![]; + for extension in extensions { + let mut rewriter = SinkInputRewriter::new(&mut sink_inputs); + let result = extension.rewrite(&mut rewriter)?; + if !rewriter.was_removed { + new_extensions.push(result.data); + } + } + + new_extensions + .into_iter() + .map(maybe_add_key_extension_to_sink) + .collect() +} + +pub async fn parse_and_get_arrow_program( + query: String, + mut schema_provider: StreamSchemaProvider, + _config: SqlConfig, +) -> Result { + let mut config = SessionConfig::new(); + config + .options_mut() + .optimizer + .enable_round_robin_repartition = false; + config.options_mut().optimizer.repartition_aggregations = false; + config.options_mut().optimizer.repartition_windows = false; + config.options_mut().optimizer.repartition_sorts = false; + config.options_mut().optimizer.repartition_joins = false; + config.options_mut().execution.target_partitions = 1; + + let session_state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .with_physical_optimizer_rules(vec![]) + .build(); + + let mut inserts = vec![]; + for statement in parse_sql_statements(&query)? { + if try_handle_set_variable(&statement, &mut schema_provider)? { + continue; + } + + if let Some(table) = CatalogTable::try_from_statement(&statement, &schema_provider)? { + schema_provider.insert_catalog_table(table); + } else { + inserts.push(Insert::try_from_statement(&statement, &schema_provider)?); + }; + } + + if inserts.is_empty() { + return plan_err!("The provided SQL does not contain a query"); + } + + let mut used_connections = HashSet::new(); + let mut extensions = vec![]; + + for insert in inserts { + let (plan, sink_name) = match insert { + Insert::InsertQuery { + sink_name, + logical_plan, + } => (logical_plan, Some(sink_name)), + Insert::Anonymous { logical_plan } => (logical_plan, None), + }; + + let mut plan_rewrite = rewrite_plan(plan, &schema_provider)?; + + if plan_rewrite + .schema() + .fields() + .iter() + .any(|f| is_json_union(f.data_type())) + { + plan_rewrite = serialize_outgoing_json(&schema_provider, Arc::new(plan_rewrite)); + } + + debug!("Plan = {}", plan_rewrite.display_graphviz()); + + let mut metadata = SourceMetadataVisitor::new(&schema_provider); + plan_rewrite.visit_with_subqueries(&mut metadata)?; + used_connections.extend(metadata.connection_ids.iter()); + + let sink = match sink_name { + Some(sink_name) => { + let table = schema_provider + .get_catalog_table_mut(&sink_name) + .ok_or_else(|| { + DataFusionError::Plan(format!("Connection {sink_name} not found")) + })?; + match table { + CatalogTable::ConnectorTable(c) => { + if let Some(id) = c.id { + used_connections.insert(id); + } + + SinkExtension::new( + TableReference::bare(sink_name), + table.clone(), + plan_rewrite.schema().clone(), + Arc::new(plan_rewrite), + ) + } + CatalogTable::MemoryTable { logical_plan, .. } => { + if logical_plan.is_some() { + return plan_err!("Can only insert into a memory table once"); + } + logical_plan.replace(plan_rewrite); + continue; + } + CatalogTable::LookupTable(_) => { + plan_err!("lookup (temporary) tables cannot be inserted into") + } + CatalogTable::TableFromQuery { .. } => { + plan_err!( + "shouldn't be inserting more data into a table made with CREATE TABLE AS" + ) + } + CatalogTable::PreviewSink { .. } => { + plan_err!("queries shouldn't be able insert into preview sink.") + } + } + } + None => SinkExtension::new( + TableReference::parse_str("preview"), + CatalogTable::PreviewSink { + logical_plan: plan_rewrite.clone(), + }, + plan_rewrite.schema().clone(), + Arc::new(plan_rewrite), + ), + }; + extensions.push(LogicalPlan::Extension(Extension { + node: Arc::new(sink?), + })); + } + + let extensions = rewrite_sinks(extensions)?; + + let mut plan_to_graph_visitor = + physical_planner::PlanToGraphVisitor::new(&schema_provider, &session_state); + for extension in extensions { + plan_to_graph_visitor.add_plan(extension)?; + } + let graph = plan_to_graph_visitor.into_graph(); + + let mut program = LogicalProgram::new(graph, ProgramConfig::default()); + + program.optimize(&ChainingOptimizer {}); + + Ok(CompiledSql { + program, + connection_ids: used_connections.into_iter().collect(), + }) +} diff --git a/src/sql/planner/parse.rs b/src/sql/planner/parse.rs new file mode 100644 index 00000000..4bd8f30e --- /dev/null +++ b/src/sql/planner/parse.rs @@ -0,0 +1,189 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::HashMap; + +use datafusion::common::{Result, plan_err}; +use datafusion::error::DataFusionError; +use datafusion::sql::sqlparser::ast::{SqlOption, Statement as DFStatement}; +use datafusion::sql::sqlparser::dialect::FunctionStreamDialect; +use datafusion::sql::sqlparser::parser::Parser; + +use crate::coordinator::{ + CreateFunction, CreateTable, DropFunction, InsertStatement, ShowFunctions, StartFunction, + Statement as CoordinatorStatement, StopFunction, StreamingSql, +}; + +/// Stage 1: String → Box +/// +/// Parses SQL using FunctionStreamDialect (from sqlparser-rs), then classifies +/// the result into either a FunctionStream DDL statement or a StreamingSql, +/// both unified under the coordinator's Statement trait. +pub fn parse_sql(query: &str) -> Result> { + let trimmed = query.trim(); + if trimmed.is_empty() { + return plan_err!("Query is empty"); + } + + let dialect = FunctionStreamDialect {}; + let mut statements = Parser::parse_sql(&dialect, trimmed) + .map_err(|e| DataFusionError::Plan(format!("SQL parse error: {e}")))?; + + if statements.is_empty() { + return plan_err!("No SQL statements found"); + } + + let stmt = statements.remove(0); + classify_statement(stmt) +} + +/// Classify a parsed DataFusion Statement into the coordinator's Statement type. +/// +/// Statement classification mirrors the analysis flow from `parse_and_get_arrow_program`: +/// - FunctionStream DDL → concrete coordinator types (CreateFunction, DropFunction, etc.) +/// - CREATE TABLE / CREATE VIEW → CreateTable (catalog registration) +/// - INSERT INTO / standalone SELECT → InsertStatement (streaming pipeline) +/// - Everything else → StreamingSql (catch-all) +fn classify_statement(stmt: DFStatement) -> Result> { + match stmt { + DFStatement::CreateFunctionWith { options } => { + let properties = sql_options_to_map(&options); + let create_fn = CreateFunction::from_properties(properties) + .map_err(|e| DataFusionError::Plan(format!("CREATE FUNCTION: {e}")))?; + Ok(Box::new(create_fn)) + } + DFStatement::StartFunction { name } => Ok(Box::new(StartFunction::new(name.to_string()))), + DFStatement::StopFunction { name } => Ok(Box::new(StopFunction::new(name.to_string()))), + DFStatement::DropFunction { func_desc, .. } => { + let name = func_desc + .first() + .map(|d| d.name.to_string()) + .unwrap_or_default(); + Ok(Box::new(DropFunction::new(name))) + } + DFStatement::ShowFunctions { .. } => Ok(Box::new(ShowFunctions::new())), + s @ DFStatement::CreateTable(_) | s @ DFStatement::CreateView { .. } => { + Ok(Box::new(CreateTable::new(s))) + } + s @ DFStatement::Insert(_) => Ok(Box::new(InsertStatement::new(s))), + other => Ok(Box::new(StreamingSql::new(other))), + } +} + +/// Convert Vec (KeyValue pairs) into HashMap. +fn sql_options_to_map(options: &[SqlOption]) -> HashMap { + options + .iter() + .filter_map(|opt| match opt { + SqlOption::KeyValue { key, value } => Some(( + key.value.clone(), + value.to_string().trim_matches('\'').to_string(), + )), + _ => None, + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn is_streaming_sql(stmt: &dyn CoordinatorStatement) -> bool { + let debug = format!("{:?}", stmt); + debug.starts_with("StreamingSql") + } + + fn is_ddl(stmt: &dyn CoordinatorStatement) -> bool { + !is_streaming_sql(stmt) + } + + #[test] + fn test_parse_create_function() { + let sql = + "CREATE FUNCTION WITH ('function_path'='./test.wasm', 'config_path'='./config.yml')"; + let stmt = parse_sql(sql).unwrap(); + assert!(is_ddl(stmt.as_ref())); + } + + #[test] + fn test_parse_create_function_minimal() { + let sql = "CREATE FUNCTION WITH ('function_path'='./processor.wasm')"; + let stmt = parse_sql(sql).unwrap(); + assert!(is_ddl(stmt.as_ref())); + } + + #[test] + fn test_parse_drop_function() { + let sql = "DROP FUNCTION my_task"; + let stmt = parse_sql(sql).unwrap(); + assert!(is_ddl(stmt.as_ref())); + } + + #[test] + fn test_parse_start_function() { + let sql = "START FUNCTION my_task"; + let stmt = parse_sql(sql).unwrap(); + assert!(is_ddl(stmt.as_ref())); + } + + #[test] + fn test_parse_stop_function() { + let sql = "STOP FUNCTION my_task"; + let stmt = parse_sql(sql).unwrap(); + assert!(is_ddl(stmt.as_ref())); + } + + #[test] + fn test_parse_show_functions() { + let sql = "SHOW FUNCTIONS"; + let stmt = parse_sql(sql).unwrap(); + assert!(is_ddl(stmt.as_ref())); + } + + #[test] + fn test_parse_case_insensitive() { + let sql1 = "create function with ('function_path'='./test.wasm')"; + assert!(is_ddl(parse_sql(sql1).unwrap().as_ref())); + + let sql2 = "show functions"; + assert!(is_ddl(parse_sql(sql2).unwrap().as_ref())); + + let sql3 = "start function my_task"; + assert!(is_ddl(parse_sql(sql3).unwrap().as_ref())); + } + + #[test] + fn test_parse_streaming_sql() { + let sql = + "SELECT count(*), tumble(interval '1 minute') as window FROM events GROUP BY window"; + let stmt = parse_sql(sql).unwrap(); + assert!(is_streaming_sql(stmt.as_ref())); + } + + #[test] + fn test_parse_empty() { + assert!(parse_sql("").is_err()); + assert!(parse_sql(" ").is_err()); + } + + #[test] + fn test_parse_with_extra_properties() { + let sql = r#"CREATE FUNCTION WITH ( + 'function_path'='./test.wasm', + 'config_path'='./config.yml', + 'parallelism'='4', + 'memory-limit'='256mb' + )"#; + let stmt = parse_sql(sql).unwrap(); + assert!(is_ddl(stmt.as_ref())); + } +} diff --git a/src/sql/planner/physical_planner.rs b/src/sql/planner/physical_planner.rs new file mode 100644 index 00000000..e7e1cf60 --- /dev/null +++ b/src/sql/planner/physical_planner.rs @@ -0,0 +1,396 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::thread; +use std::time::Duration; + +use datafusion::arrow::datatypes::IntervalMonthDayNanoType; +use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; +use datafusion::common::{ + DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, Spans, plan_err, +}; +use datafusion::execution::context::SessionState; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::functions::datetime::date_bin; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNode}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}; +use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode}; +use datafusion_proto::{ + physical_plan::AsExecutionPlan, + protobuf::{AggregateMode, physical_plan_node::PhysicalPlanType}, +}; +use petgraph::graph::{DiGraph, NodeIndex}; +use prost::Message; +use tokio::runtime::Builder; +use tokio::sync::oneshot; + +use async_trait::async_trait; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr; + +use crate::datastream::logical::{LogicalEdge, LogicalGraph, LogicalNode}; +use crate::sql::physical::{ + DebeziumUnrollingExec, DecodingContext, FsMemExec, FsPhysicalExtensionCodec, ToDebeziumExec, +}; +use crate::sql::planner::StreamSchemaProvider; +use crate::sql::planner::extension::debezium::{ + DEBEZIUM_UNROLLING_EXTENSION_NAME, DebeziumUnrollingExtension, TO_DEBEZIUM_EXTENSION_NAME, +}; +use crate::sql::planner::extension::key_calculation::KeyCalculationExtension; +use crate::sql::planner::extension::{NamedNode, NodeWithIncomingEdges, StreamExtension}; +use crate::sql::planner::schemas::add_timestamp_field_arrow; +use crate::types::{FsSchema, FsSchemaRef}; + +pub(crate) struct PlanToGraphVisitor<'a> { + graph: DiGraph, + output_schemas: HashMap, + named_nodes: HashMap, + traversal: Vec>, + planner: Planner<'a>, +} + +impl<'a> PlanToGraphVisitor<'a> { + pub fn new(schema_provider: &'a StreamSchemaProvider, session_state: &'a SessionState) -> Self { + Self { + graph: Default::default(), + output_schemas: Default::default(), + named_nodes: Default::default(), + traversal: vec![], + planner: Planner::new(schema_provider, session_state), + } + } +} + +pub(crate) struct Planner<'a> { + schema_provider: &'a StreamSchemaProvider, + planner: DefaultPhysicalPlanner, + session_state: &'a SessionState, +} + +impl<'a> Planner<'a> { + pub(crate) fn new( + schema_provider: &'a StreamSchemaProvider, + session_state: &'a SessionState, + ) -> Self { + let planner = + DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new(FsExtensionPlanner {})]); + Self { + schema_provider, + planner, + session_state, + } + } + + pub(crate) fn sync_plan(&self, plan: &LogicalPlan) -> Result> { + let fut = self.planner.create_physical_plan(plan, self.session_state); + let (tx, mut rx) = oneshot::channel(); + thread::scope(|s| { + let _handle = tokio::runtime::Handle::current(); + let builder = thread::Builder::new(); + let builder = if cfg!(debug_assertions) { + builder.stack_size(10_000_000) + } else { + builder + }; + builder + .spawn_scoped(s, move || { + let rt = Builder::new_current_thread().enable_all().build().unwrap(); + rt.block_on(async { + let plan = fut.await; + tx.send(plan).unwrap(); + }); + }) + .unwrap(); + }); + + rx.try_recv().unwrap() + } + + pub(crate) fn create_physical_expr( + &self, + expr: &Expr, + input_dfschema: &DFSchema, + ) -> Result> { + self.planner + .create_physical_expr(expr, input_dfschema, self.session_state) + } + + pub(crate) fn serialize_as_physical_expr( + &self, + expr: &Expr, + schema: &DFSchema, + ) -> Result> { + let physical = self.create_physical_expr(expr, schema)?; + let proto = serialize_physical_expr(&physical, &DefaultPhysicalExtensionCodec {})?; + Ok(proto.encode_to_vec()) + } + + pub(crate) fn split_physical_plan( + &self, + key_indices: Vec, + aggregate: &LogicalPlan, + add_timestamp_field: bool, + ) -> Result { + let physical_plan = self.sync_plan(aggregate)?; + let codec = FsPhysicalExtensionCodec { + context: DecodingContext::Planning, + }; + let mut physical_plan_node = + PhysicalPlanNode::try_from_physical_plan(physical_plan.clone(), &codec)?; + let PhysicalPlanType::Aggregate(mut final_aggregate_proto) = physical_plan_node + .physical_plan_type + .take() + .ok_or_else(|| DataFusionError::Plan("missing physical plan type".to_string()))? + else { + return plan_err!("unexpected physical plan type"); + }; + let AggregateMode::Final = final_aggregate_proto.mode() else { + return plan_err!("unexpected physical plan type"); + }; + + let partial_aggregation_plan = *final_aggregate_proto + .input + .take() + .ok_or_else(|| DataFusionError::Plan("missing input".to_string()))?; + + let partial_aggregation_exec_plan = partial_aggregation_plan.try_into_physical_plan( + self.schema_provider, + &RuntimeEnvBuilder::new().build().unwrap(), + &codec, + )?; + + let partial_schema = partial_aggregation_exec_plan.schema(); + let final_input_table_provider = FsMemExec::new("partial".into(), partial_schema.clone()); + + final_aggregate_proto.input = Some(Box::new(PhysicalPlanNode::try_from_physical_plan( + Arc::new(final_input_table_provider), + &codec, + )?)); + + let finish_plan = PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Aggregate(final_aggregate_proto)), + }; + + let (partial_schema, timestamp_index) = if add_timestamp_field { + ( + add_timestamp_field_arrow((*partial_schema).clone()), + partial_schema.fields().len(), + ) + } else { + (partial_schema.clone(), partial_schema.fields().len() - 1) + }; + + let partial_schema = FsSchema::new_keyed(partial_schema, timestamp_index, key_indices); + + Ok(SplitPlanOutput { + partial_aggregation_plan, + partial_schema, + finish_plan, + }) + } + + pub fn binning_function_proto( + &self, + width: Duration, + input_schema: DFSchemaRef, + ) -> Result { + let date_bin = date_bin().call(vec![ + Expr::Literal( + ScalarValue::IntervalMonthDayNano(Some(IntervalMonthDayNanoType::make_value( + 0, + 0, + width.as_nanos() as i64, + ))), + None, + ), + Expr::Column(datafusion::common::Column { + relation: None, + name: "_timestamp".into(), + spans: Spans::new(), + }), + ]); + + let binning_function = self.create_physical_expr(&date_bin, &input_schema)?; + serialize_physical_expr(&binning_function, &DefaultPhysicalExtensionCodec {}) + } +} + +struct FsExtensionPlanner {} + +#[async_trait] +impl ExtensionPlanner for FsExtensionPlanner { + async fn plan_extension( + &self, + _planner: &dyn PhysicalPlanner, + node: &dyn UserDefinedLogicalNode, + _logical_inputs: &[&LogicalPlan], + physical_inputs: &[Arc], + _session_state: &SessionState, + ) -> Result>> { + let schema = node.schema().as_ref().into(); + if let Ok::<&dyn StreamExtension, _>(stream_extension) = node.try_into() { + if stream_extension.transparent() { + match node.name() { + DEBEZIUM_UNROLLING_EXTENSION_NAME => { + let node = node + .as_any() + .downcast_ref::() + .unwrap(); + let input = physical_inputs[0].clone(); + return Ok(Some(Arc::new(DebeziumUnrollingExec::try_new( + input, + node.primary_keys.clone(), + )?))); + } + TO_DEBEZIUM_EXTENSION_NAME => { + let input = physical_inputs[0].clone(); + return Ok(Some(Arc::new(ToDebeziumExec::try_new(input)?))); + } + _ => return Ok(None), + } + } + }; + let name = + if let Some(key_extension) = node.as_any().downcast_ref::() { + key_extension.name.clone() + } else { + None + }; + Ok(Some(Arc::new(FsMemExec::new( + name.unwrap_or("memory".to_string()), + Arc::new(schema), + )))) + } +} + +impl PlanToGraphVisitor<'_> { + fn add_index_to_traversal(&mut self, index: NodeIndex) { + if let Some(last) = self.traversal.last_mut() { + last.push(index); + } + } + + pub(crate) fn add_plan(&mut self, plan: LogicalPlan) -> Result<()> { + self.traversal.clear(); + plan.visit(self)?; + Ok(()) + } + + pub fn into_graph(self) -> LogicalGraph { + self.graph + } + + pub fn build_extension( + &mut self, + input_nodes: Vec, + extension: &dyn StreamExtension, + ) -> Result<()> { + if let Some(node_name) = extension.node_name() { + if self.named_nodes.contains_key(&node_name) { + return plan_err!( + "extension {:?} has already been planned, shouldn't try again.", + node_name + ); + } + } + + let input_schemas = input_nodes + .iter() + .map(|index| { + Ok(self + .output_schemas + .get(index) + .ok_or_else(|| DataFusionError::Plan("missing input node".to_string()))? + .clone()) + }) + .collect::>>()?; + + let NodeWithIncomingEdges { node, edges } = extension + .plan_node(&self.planner, self.graph.node_count(), input_schemas) + .map_err(|e| e.context(format!("planning operator {extension:?}")))?; + + let node_index = self.graph.add_node(node); + self.add_index_to_traversal(node_index); + + for (source, edge) in input_nodes.into_iter().zip(edges.into_iter()) { + self.graph.add_edge(source, node_index, edge); + } + + self.output_schemas + .insert(node_index, extension.output_schema().into()); + + if let Some(node_name) = extension.node_name() { + self.named_nodes.insert(node_name, node_index); + } + Ok(()) + } +} + +impl TreeNodeVisitor<'_> for PlanToGraphVisitor<'_> { + type Node = LogicalPlan; + + fn f_down(&mut self, node: &Self::Node) -> Result { + let LogicalPlan::Extension(Extension { node }) = node else { + return Ok(TreeNodeRecursion::Continue); + }; + + let stream_extension: &dyn StreamExtension = node + .try_into() + .map_err(|e: DataFusionError| e.context("converting extension"))?; + if stream_extension.transparent() { + return Ok(TreeNodeRecursion::Continue); + } + + if let Some(name) = stream_extension.node_name() { + if let Some(node_index) = self.named_nodes.get(&name) { + self.add_index_to_traversal(*node_index); + return Ok(TreeNodeRecursion::Jump); + } + } + + if !node.inputs().is_empty() { + self.traversal.push(vec![]); + } + + Ok(TreeNodeRecursion::Continue) + } + + fn f_up(&mut self, node: &Self::Node) -> Result { + let LogicalPlan::Extension(Extension { node }) = node else { + return Ok(TreeNodeRecursion::Continue); + }; + + let stream_extension: &dyn StreamExtension = node + .try_into() + .map_err(|e: DataFusionError| e.context("planning extension"))?; + + if stream_extension.transparent() { + return Ok(TreeNodeRecursion::Continue); + } + + if let Some(name) = stream_extension.node_name() { + if self.named_nodes.contains_key(&name) { + return Ok(TreeNodeRecursion::Continue); + } + } + + let input_nodes = if !node.inputs().is_empty() { + self.traversal.pop().unwrap_or_default() + } else { + vec![] + }; + let stream_extension: &dyn StreamExtension = node + .try_into() + .map_err(|e: DataFusionError| e.context("converting extension"))?; + self.build_extension(input_nodes, stream_extension)?; + + Ok(TreeNodeRecursion::Continue) + } +} + +pub(crate) struct SplitPlanOutput { + pub(crate) partial_aggregation_plan: PhysicalPlanNode, + pub(crate) partial_schema: FsSchema, + pub(crate) finish_plan: PhysicalPlanNode, +} diff --git a/src/sql/planner/plan/aggregate.rs b/src/sql/planner/plan/aggregate.rs new file mode 100644 index 00000000..aad17edb --- /dev/null +++ b/src/sql/planner/plan/aggregate.rs @@ -0,0 +1,275 @@ +use std::sync::Arc; + +use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion::common::{DFSchema, Result, not_impl_err, plan_err}; +use datafusion::functions_aggregate::expr_fn::max; +use datafusion::logical_expr; +use datafusion::logical_expr::{Aggregate, Expr, Extension, LogicalPlan}; +use datafusion::prelude::col; +use tracing::debug; + +use crate::sql::planner::StreamSchemaProvider; +use crate::sql::planner::extension::aggregate::AggregateExtension; +use crate::sql::planner::extension::key_calculation::{KeyCalculationExtension, KeysOrExprs}; +use crate::sql::planner::plan::WindowDetectingVisitor; +use crate::sql::types::{ + DFField, TIMESTAMP_FIELD, WindowBehavior, WindowType, fields_with_qualifiers, find_window, + schema_from_df_fields_with_metadata, +}; + +pub(crate) struct AggregateRewriter<'a> { + pub schema_provider: &'a StreamSchemaProvider, +} + +impl AggregateRewriter<'_> { + /// Rewrite a non-windowed aggregate into an updating aggregate with key calculation + pub fn rewrite_non_windowed_aggregate( + input: Arc, + mut key_fields: Vec, + group_expr: Vec, + mut aggr_expr: Vec, + schema: Arc, + _schema_provider: &StreamSchemaProvider, + ) -> Result> { + let key_count = key_fields.len(); + key_fields.extend(fields_with_qualifiers(input.schema())); + + let key_schema = Arc::new(schema_from_df_fields_with_metadata( + &key_fields, + schema.metadata().clone(), + )?); + + let mut key_projection_expressions: Vec<_> = group_expr + .iter() + .zip(key_fields.iter()) + .map(|(expr, f)| expr.clone().alias(f.name().to_string())) + .collect(); + + key_projection_expressions.extend( + fields_with_qualifiers(input.schema()) + .iter() + .map(|field| Expr::Column(field.qualified_column())), + ); + + let key_projection = + LogicalPlan::Projection(logical_expr::Projection::try_new_with_schema( + key_projection_expressions, + input.clone(), + key_schema, + )?); + + let key_plan = LogicalPlan::Extension(Extension { + node: Arc::new(KeyCalculationExtension::new( + key_projection, + KeysOrExprs::Keys((0..key_count).collect()), + )), + }); + + let Ok(timestamp_field) = key_plan + .schema() + .qualified_field_with_unqualified_name(TIMESTAMP_FIELD) + else { + return plan_err!("no timestamp field found in schema"); + }; + + let timestamp_field: DFField = timestamp_field.into(); + let column = timestamp_field.qualified_column(); + aggr_expr.push(max(col(column.clone())).alias(TIMESTAMP_FIELD)); + + let mut output_schema_fields = fields_with_qualifiers(&schema); + output_schema_fields.push(timestamp_field.clone()); + let output_schema = Arc::new(schema_from_df_fields_with_metadata( + &output_schema_fields, + schema.metadata().clone(), + )?); + + let aggregate = Aggregate::try_new_with_schema( + Arc::new(key_plan), + group_expr, + aggr_expr, + output_schema, + )?; + + debug!( + "non-windowed aggregate field names: {:?}", + aggregate + .schema + .fields() + .iter() + .map(|f| f.name()) + .collect::>() + ); + + let final_plan = LogicalPlan::Aggregate(aggregate); + Ok(Transformed::yes(final_plan)) + } +} + +impl TreeNodeRewriter for AggregateRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> Result> { + let LogicalPlan::Aggregate(Aggregate { + input, + mut group_expr, + aggr_expr, + schema, + .. + }) = node + else { + return Ok(Transformed::no(node)); + }; + + let mut window_group_expr: Vec<_> = group_expr + .iter() + .enumerate() + .filter_map(|(i, expr)| { + find_window(expr) + .map(|option| option.map(|inner| (i, inner))) + .transpose() + }) + .collect::>>()?; + + if window_group_expr.len() > 1 { + return not_impl_err!( + "do not support {} window expressions in group by", + window_group_expr.len() + ); + } + + let mut key_fields: Vec = fields_with_qualifiers(&schema) + .iter() + .take(group_expr.len()) + .map(|field| { + DFField::new( + field.qualifier().cloned(), + format!("_key_{}", field.name()), + field.data_type().clone(), + field.is_nullable(), + ) + }) + .collect(); + + let mut window_detecting_visitor = WindowDetectingVisitor::default(); + input.visit_with_subqueries(&mut window_detecting_visitor)?; + + let window = window_detecting_visitor.window; + let window_behavior = match (window.is_some(), !window_group_expr.is_empty()) { + (true, true) => { + let input_window = window.unwrap(); + let (window_index, group_by_window_type) = window_group_expr.pop().unwrap(); + if group_by_window_type != input_window { + return Err(datafusion::error::DataFusionError::NotImplemented( + "window in group by does not match input window".to_string(), + )); + } + let matching_field = window_detecting_visitor.fields.iter().next(); + match matching_field { + Some(field) => { + group_expr[window_index] = Expr::Column(field.qualified_column()); + WindowBehavior::InData + } + None => { + if matches!(input_window, WindowType::Session { .. }) { + return plan_err!("can't reinvoke session window in nested aggregates"); + } + group_expr.remove(window_index); + key_fields.remove(window_index); + let window_field = schema.qualified_field(window_index).into(); + WindowBehavior::FromOperator { + window: input_window, + window_field, + window_index, + is_nested: true, + } + } + } + } + (true, false) => WindowBehavior::InData, + (false, true) => { + let (window_index, window_type) = window_group_expr.pop().unwrap(); + group_expr.remove(window_index); + key_fields.remove(window_index); + let window_field = schema.qualified_field(window_index).into(); + WindowBehavior::FromOperator { + window: window_type, + window_field, + window_index, + is_nested: false, + } + } + (false, false) => { + return Self::rewrite_non_windowed_aggregate( + input, + key_fields, + group_expr, + aggr_expr, + schema, + self.schema_provider, + ); + } + }; + + let key_count = key_fields.len(); + key_fields.extend(fields_with_qualifiers(input.schema())); + + let key_schema = Arc::new(schema_from_df_fields_with_metadata( + &key_fields, + schema.metadata().clone(), + )?); + + let mut key_projection_expressions: Vec<_> = group_expr + .iter() + .zip(key_fields.iter()) + .map(|(expr, f)| expr.clone().alias(f.name().to_string())) + .collect(); + + key_projection_expressions.extend( + fields_with_qualifiers(input.schema()) + .iter() + .map(|field| Expr::Column(field.qualified_column())), + ); + + let key_projection = + LogicalPlan::Projection(logical_expr::Projection::try_new_with_schema( + key_projection_expressions, + input.clone(), + key_schema, + )?); + + let key_plan = LogicalPlan::Extension(Extension { + node: Arc::new(KeyCalculationExtension::new( + key_projection, + KeysOrExprs::Keys((0..key_count).collect()), + )), + }); + + let mut aggregate_schema_fields = fields_with_qualifiers(&schema); + if let WindowBehavior::FromOperator { window_index, .. } = &window_behavior { + aggregate_schema_fields.remove(*window_index); + } + let internal_schema = Arc::new(schema_from_df_fields_with_metadata( + &aggregate_schema_fields, + schema.metadata().clone(), + )?); + + let rewritten_aggregate = Aggregate::try_new_with_schema( + Arc::new(key_plan), + group_expr, + aggr_expr, + internal_schema, + )?; + + let aggregate_extension = AggregateExtension::new( + window_behavior, + LogicalPlan::Aggregate(rewritten_aggregate), + (0..key_count).collect(), + ); + let final_plan = LogicalPlan::Extension(Extension { + node: Arc::new(aggregate_extension), + }); + + WindowDetectingVisitor::get_window(&final_plan)?; + Ok(Transformed::yes(final_plan)) + } +} diff --git a/src/sql/planner/plan/join.rs b/src/sql/planner/plan/join.rs new file mode 100644 index 00000000..04a27e9b --- /dev/null +++ b/src/sql/planner/plan/join.rs @@ -0,0 +1,240 @@ +use std::sync::Arc; + +use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion::common::{ + Column, DataFusionError, JoinConstraint, JoinType, Result, ScalarValue, Spans, TableReference, + not_impl_err, +}; +use datafusion::logical_expr; +use datafusion::logical_expr::expr::Alias; +use datafusion::logical_expr::{ + BinaryExpr, Case, Expr, Extension, Join, LogicalPlan, Projection, build_join_schema, +}; +use datafusion::prelude::coalesce; + +use crate::sql::planner::StreamSchemaProvider; +use crate::sql::planner::extension::join::JoinExtension; +use crate::sql::planner::extension::key_calculation::KeyCalculationExtension; +use crate::sql::planner::plan::WindowDetectingVisitor; +use crate::sql::types::{WindowType, fields_with_qualifiers, schema_from_df_fields_with_metadata}; + +pub(crate) struct JoinRewriter<'a> { + pub schema_provider: &'a StreamSchemaProvider, +} + +impl JoinRewriter<'_> { + fn check_join_windowing(join: &Join) -> Result { + let left_window = WindowDetectingVisitor::get_window(&join.left)?; + let right_window = WindowDetectingVisitor::get_window(&join.right)?; + match (left_window, right_window) { + (None, None) => { + if join.join_type == JoinType::Inner { + Ok(false) + } else { + Err(DataFusionError::NotImplemented( + "can't handle non-inner joins without windows".into(), + )) + } + } + (None, Some(_)) => Err(DataFusionError::NotImplemented( + "can't handle mixed windowing between left (non-windowed) and right (windowed)" + .into(), + )), + (Some(_), None) => Err(DataFusionError::NotImplemented( + "can't handle mixed windowing between left (windowed) and right (non-windowed)" + .into(), + )), + (Some(left_window), Some(right_window)) => { + if left_window != right_window { + return Err(DataFusionError::NotImplemented( + "can't handle mixed windowing between left and right".into(), + )); + } + if let WindowType::Session { .. } = left_window { + return Err(DataFusionError::NotImplemented( + "can't handle session windows in joins".into(), + )); + } + Ok(true) + } + } + } + + fn create_join_key_plan( + input: Arc, + join_expressions: Vec, + name: &'static str, + ) -> Result { + let key_count = join_expressions.len(); + + let join_expressions: Vec<_> = join_expressions + .into_iter() + .enumerate() + .map(|(index, expr)| { + expr.alias_qualified( + Some(TableReference::bare("_stream")), + format!("_key_{index}"), + ) + }) + .chain( + fields_with_qualifiers(input.schema()) + .iter() + .map(|field| Expr::Column(field.qualified_column())), + ) + .collect(); + + let projection = Projection::try_new(join_expressions, input)?; + let key_calculation_extension = KeyCalculationExtension::new_named_and_trimmed( + LogicalPlan::Projection(projection), + (0..key_count).collect(), + name.to_string(), + ); + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(key_calculation_extension), + })) + } + + fn post_join_timestamp_projection(&mut self, input: LogicalPlan) -> Result { + let schema = input.schema().clone(); + let mut schema_with_timestamp = fields_with_qualifiers(&schema); + let timestamp_fields = schema_with_timestamp + .iter() + .filter(|field| field.name() == "_timestamp") + .cloned() + .collect::>(); + + if timestamp_fields.len() != 2 { + return not_impl_err!("join must have two timestamp fields"); + } + + schema_with_timestamp.retain(|field| field.name() != "_timestamp"); + let mut projection_expr = schema_with_timestamp + .iter() + .map(|field| { + Expr::Column(Column { + relation: field.qualifier().cloned(), + name: field.name().to_string(), + spans: Spans::default(), + }) + }) + .collect::>(); + + schema_with_timestamp.push(timestamp_fields[0].clone()); + + let output_schema = Arc::new(schema_from_df_fields_with_metadata( + &schema_with_timestamp, + schema.metadata().clone(), + )?); + + let left_field = ×tamp_fields[0]; + let left_column = Expr::Column(Column { + relation: left_field.qualifier().cloned(), + name: left_field.name().to_string(), + spans: Spans::default(), + }); + let right_field = ×tamp_fields[1]; + let right_column = Expr::Column(Column { + relation: right_field.qualifier().cloned(), + name: right_field.name().to_string(), + spans: Spans::default(), + }); + + let max_timestamp = Expr::Case(Case { + expr: Some(Box::new(Expr::BinaryExpr(BinaryExpr { + left: Box::new(left_column.clone()), + op: logical_expr::Operator::GtEq, + right: Box::new(right_column.clone()), + }))), + when_then_expr: vec![ + ( + Box::new(Expr::Literal(ScalarValue::Boolean(Some(true)), None)), + Box::new(left_column.clone()), + ), + ( + Box::new(Expr::Literal(ScalarValue::Boolean(Some(false)), None)), + Box::new(right_column.clone()), + ), + ], + else_expr: Some(Box::new(coalesce(vec![ + left_column.clone(), + right_column.clone(), + ]))), + }); + + projection_expr.push(Expr::Alias(Alias { + expr: Box::new(max_timestamp), + relation: timestamp_fields[0].qualifier().cloned(), + name: timestamp_fields[0].name().to_string(), + metadata: None, + })); + + Ok(LogicalPlan::Projection(Projection::try_new_with_schema( + projection_expr, + Arc::new(input), + output_schema, + )?)) + } +} + +impl TreeNodeRewriter for JoinRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> Result> { + let LogicalPlan::Join(join) = node else { + return Ok(Transformed::no(node)); + }; + + let is_instant = Self::check_join_windowing(&join)?; + + let Join { + left, + right, + on, + filter, + join_type, + join_constraint: JoinConstraint::On, + schema: _, + null_equals_null: false, + } = join + else { + return not_impl_err!("can't handle join constraint other than ON"); + }; + + if on.is_empty() && !is_instant { + return not_impl_err!("Updating joins must include an equijoin condition"); + } + + let (left_expressions, right_expressions): (Vec<_>, Vec<_>) = + on.clone().into_iter().unzip(); + + let left_input = Self::create_join_key_plan(left, left_expressions, "left")?; + let right_input = Self::create_join_key_plan(right, right_expressions, "right")?; + + let rewritten_join = LogicalPlan::Join(Join { + schema: Arc::new(build_join_schema( + left_input.schema(), + right_input.schema(), + &join_type, + )?), + left: Arc::new(left_input), + right: Arc::new(right_input), + on, + join_type, + join_constraint: JoinConstraint::On, + null_equals_null: false, + filter, + }); + + let final_logical_plan = self.post_join_timestamp_projection(rewritten_join)?; + + let join_extension = JoinExtension { + rewritten_join: final_logical_plan, + is_instant, + ttl: (!is_instant).then_some(self.schema_provider.planning_options.ttl), + }; + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(join_extension), + }))) + } +} diff --git a/src/sql/planner/plan/mod.rs b/src/sql/planner/plan/mod.rs new file mode 100644 index 00000000..d497ca65 --- /dev/null +++ b/src/sql/planner/plan/mod.rs @@ -0,0 +1,449 @@ +use std::collections::HashSet; +use std::sync::Arc; + +use datafusion::common::tree_node::{Transformed, TreeNodeRecursion}; +use datafusion::common::{ + Column, DataFusionError, Result, Spans, TableReference, plan_err, + tree_node::{TreeNode, TreeNodeRewriter, TreeNodeVisitor}, +}; +use datafusion::logical_expr::{ + Aggregate, Expr, Extension, Filter, LogicalPlan, SubqueryAlias, expr::Alias, +}; + +use crate::sql::planner::extension::StreamExtension; +use crate::sql::planner::extension::aggregate::{AGGREGATE_EXTENSION_NAME, AggregateExtension}; +use crate::sql::planner::extension::join::JOIN_NODE_NAME; +use crate::sql::planner::extension::remote_table::RemoteTableExtension; +use crate::sql::planner::schemas::{add_timestamp_field, has_timestamp_field}; +use crate::sql::types::{ + DFField, TIMESTAMP_FIELD, WindowBehavior, WindowType, fields_with_qualifiers, find_window, +}; + +use self::aggregate::AggregateRewriter; +use self::join::JoinRewriter; +use self::window_fn::WindowFunctionRewriter; + +pub(crate) mod aggregate; +pub(crate) mod join; +pub(crate) mod window_fn; + +use super::StreamSchemaProvider; +use tracing::debug; + +/// Stage 3: LogicalPlan → Streaming LogicalPlan +/// +/// Rewrites a standard DataFusion logical plan into one that supports +/// streaming semantics (timestamps, windows, watermarks). +pub fn rewrite_plan( + plan: LogicalPlan, + schema_provider: &StreamSchemaProvider, +) -> Result { + let rewritten_plan = plan.rewrite_with_subqueries(&mut StreamRewriter { schema_provider })?; + + rewritten_plan + .data + .visit_with_subqueries(&mut TimeWindowUdfChecker {})?; + + debug!( + "Streaming logical plan:\n{}", + rewritten_plan.data.display_graphviz() + ); + + Ok(rewritten_plan.data) +} + +/// Visitor that detects window types in a logical plan +#[derive(Debug, Default)] +pub(crate) struct WindowDetectingVisitor { + pub(crate) window: Option, + pub(crate) fields: HashSet, +} + +impl WindowDetectingVisitor { + pub(crate) fn get_window(logical_plan: &LogicalPlan) -> Result> { + let mut visitor = WindowDetectingVisitor { + window: None, + fields: HashSet::new(), + }; + logical_plan.visit_with_subqueries(&mut visitor)?; + Ok(visitor.window.take()) + } +} + +fn extract_column(expr: &Expr) -> Option<&Column> { + match expr { + Expr::Column(column) => Some(column), + Expr::Alias(Alias { expr, .. }) => extract_column(expr), + _ => None, + } +} + +impl TreeNodeVisitor<'_> for WindowDetectingVisitor { + type Node = LogicalPlan; + + fn f_down(&mut self, node: &Self::Node) -> Result { + let LogicalPlan::Extension(Extension { node }) = node else { + return Ok(TreeNodeRecursion::Continue); + }; + + if node.name() == JOIN_NODE_NAME { + let input_windows: HashSet<_> = node + .inputs() + .iter() + .map(|input| Self::get_window(input)) + .collect::>>()?; + if input_windows.len() > 1 { + return Err(DataFusionError::Plan( + "can't handle mixed windowing between left and right".to_string(), + )); + } + self.window = input_windows + .into_iter() + .next() + .expect("join has at least one input"); + return Ok(TreeNodeRecursion::Jump); + } + Ok(TreeNodeRecursion::Continue) + } + + fn f_up(&mut self, node: &Self::Node) -> Result { + match node { + LogicalPlan::Projection(projection) => { + let window_expressions = projection + .expr + .iter() + .enumerate() + .filter_map(|(index, expr)| { + if let Some(column) = extract_column(expr) { + let input_field = projection + .input + .schema() + .field_with_name(column.relation.as_ref(), &column.name); + let input_field = match input_field { + Ok(field) => field, + Err(err) => return Some(Err(err)), + }; + if self.fields.contains( + &(column.relation.clone(), Arc::new(input_field.clone())).into(), + ) { + return self.window.clone().map(|window| Ok((index, window))); + } + } + find_window(expr) + .map(|option| option.map(|inner| (index, inner))) + .transpose() + }) + .collect::>>()?; + self.fields.clear(); + for (index, window) in window_expressions { + if let Some(existing_window) = &self.window { + if *existing_window != window { + return plan_err!( + "can't window by both {:?} and {:?}", + existing_window, + window + ); + } + self.fields + .insert(projection.schema.qualified_field(index).into()); + } else { + return plan_err!( + "can't call a windowing function without grouping by it in an aggregate" + ); + } + } + } + LogicalPlan::SubqueryAlias(subquery_alias) => { + self.fields = self + .fields + .drain() + .map(|field| { + Ok(subquery_alias + .schema + .qualified_field( + subquery_alias + .input + .schema() + .index_of_column(&field.qualified_column())?, + ) + .into()) + }) + .collect::>>()?; + } + LogicalPlan::Aggregate(Aggregate { + input, + group_expr, + aggr_expr: _, + schema, + .. + }) => { + let window_expressions = group_expr + .iter() + .enumerate() + .filter_map(|(index, expr)| { + if let Some(column) = extract_column(expr) { + let input_field = input + .schema() + .field_with_name(column.relation.as_ref(), &column.name); + let input_field = match input_field { + Ok(field) => field, + Err(err) => return Some(Err(err)), + }; + if self + .fields + .contains(&(column.relation.as_ref(), input_field).into()) + { + return self.window.clone().map(|window| Ok((index, window))); + } + } + find_window(expr) + .map(|option| option.map(|inner| (index, inner))) + .transpose() + }) + .collect::>>()?; + self.fields.clear(); + for (index, window) in window_expressions { + if let Some(existing_window) = &self.window { + if *existing_window != window { + return Err(DataFusionError::Plan( + "window expressions do not match".to_string(), + )); + } + } else { + self.window = Some(window); + } + self.fields.insert(schema.qualified_field(index).into()); + } + } + LogicalPlan::Extension(Extension { node }) => { + if node.name() == AGGREGATE_EXTENSION_NAME { + let aggregate_extension = node + .as_any() + .downcast_ref::() + .expect("should be aggregate extension"); + + match &aggregate_extension.window_behavior { + WindowBehavior::FromOperator { + window, + window_field, + window_index: _, + is_nested, + } => { + if self.window.is_some() && !*is_nested { + return Err(DataFusionError::Plan( + "aggregate node should not be recalculating window, as input is windowed.".to_string(), + )); + } + self.window = Some(window.clone()); + self.fields.insert(window_field.clone()); + } + WindowBehavior::InData => { + let input_fields = self.fields.clone(); + self.fields.clear(); + for field in fields_with_qualifiers(node.schema()) { + if input_fields.contains(&field) { + self.fields.insert(field); + } + } + if self.fields.is_empty() { + return Err(DataFusionError::Plan( + "must have window in aggregate. Make sure you are calling one of the windowing functions (hop, tumble, session) or using the window field of the input".to_string(), + )); + } + } + } + } + } + _ => {} + } + Ok(TreeNodeRecursion::Continue) + } +} + +/// Main rewriter for streaming SQL plans. +/// Rewrites standard logical plans into streaming-aware plans with +/// timestamp propagation, window detection, and streaming operator insertion. +pub struct StreamRewriter<'a> { + pub(crate) schema_provider: &'a StreamSchemaProvider, +} + +impl TreeNodeRewriter for StreamRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, mut node: Self::Node) -> Result> { + match node { + LogicalPlan::Projection(ref mut projection) => { + if !has_timestamp_field(&projection.schema) { + let timestamp_field: DFField = projection + .input + .schema() + .qualified_field_with_unqualified_name(TIMESTAMP_FIELD) + .map_err(|_| { + DataFusionError::Plan(format!( + "No timestamp field found in projection input ({})", + projection.input.display() + )) + })? + .into(); + projection.schema = add_timestamp_field( + projection.schema.clone(), + timestamp_field.qualifier().cloned(), + ) + .expect("in projection"); + projection.expr.push(Expr::Column(Column { + relation: timestamp_field.qualifier().cloned(), + name: TIMESTAMP_FIELD.to_string(), + spans: Spans::default(), + })); + } + + // Rewrite row_time() calls to _timestamp column references + let rewritten = projection + .expr + .iter() + .map(|expr| expr.clone().rewrite(&mut RowTimeRewriter {})) + .collect::>>()?; + if rewritten.iter().any(|r| r.transformed) { + projection.expr = rewritten.into_iter().map(|r| r.data).collect(); + } + return Ok(Transformed::yes(node)); + } + LogicalPlan::Aggregate(aggregate) => { + return AggregateRewriter { + schema_provider: self.schema_provider, + } + .f_up(LogicalPlan::Aggregate(aggregate)); + } + LogicalPlan::Join(join) => { + return JoinRewriter { + schema_provider: self.schema_provider, + } + .f_up(LogicalPlan::Join(join)); + } + LogicalPlan::Filter(f) => { + let expr = f + .predicate + .clone() + .rewrite(&mut TimeWindowNullCheckRemover {})?; + return Ok(if expr.transformed { + Transformed::yes(LogicalPlan::Filter(Filter::try_new(expr.data, f.input)?)) + } else { + Transformed::no(LogicalPlan::Filter(f)) + }); + } + LogicalPlan::Window(_) => { + return WindowFunctionRewriter {}.f_up(node); + } + LogicalPlan::Sort(_) => { + return plan_err!( + "ORDER BY is not currently supported in streaming SQL ({})", + node.display() + ); + } + LogicalPlan::Repartition(_) => { + return plan_err!( + "Repartitions are not currently supported ({})", + node.display() + ); + } + LogicalPlan::Union(mut union) => { + union.schema = union.inputs[0].schema().clone(); + for input in union.inputs.iter_mut() { + if let LogicalPlan::Extension(Extension { node }) = input.as_ref() { + let stream_extension: &dyn StreamExtension = node.try_into().unwrap(); + if !stream_extension.transparent() { + continue; + } + } + let remote_table_extension = Arc::new(RemoteTableExtension { + input: input.as_ref().clone(), + name: TableReference::bare("union_input"), + schema: union.schema.clone(), + materialize: false, + }); + *input = Arc::new(LogicalPlan::Extension(Extension { + node: remote_table_extension, + })); + } + return Ok(Transformed::yes(LogicalPlan::Union(union))); + } + LogicalPlan::SubqueryAlias(sa) => { + return Ok(Transformed::yes(LogicalPlan::SubqueryAlias( + SubqueryAlias::try_new(sa.input, sa.alias)?, + ))); + } + LogicalPlan::Limit(_) => { + return plan_err!( + "LIMIT is not currently supported in streaming SQL ({})", + node.display() + ); + } + LogicalPlan::Explain(_) => { + return plan_err!("EXPLAIN is not supported ({})", node.display()); + } + LogicalPlan::Analyze(_) => { + return plan_err!("ANALYZE is not supported ({})", node.display()); + } + _ => {} + } + Ok(Transformed::no(node)) + } +} + +/// Rewrites row_time() function calls to _timestamp column references +struct RowTimeRewriter; + +impl TreeNodeRewriter for RowTimeRewriter { + type Node = Expr; + + fn f_up(&mut self, node: Self::Node) -> Result> { + if let Expr::ScalarFunction(ref func) = node { + if func.func.name() == "row_time" { + return Ok(Transformed::yes(Expr::Column(Column::new_unqualified( + TIMESTAMP_FIELD.to_string(), + )))); + } + } + Ok(Transformed::no(node)) + } +} + +/// Removes IS NOT NULL checks on window expressions that get pushed down incorrectly +pub(crate) struct TimeWindowNullCheckRemover; + +impl TreeNodeRewriter for TimeWindowNullCheckRemover { + type Node = Expr; + + fn f_up(&mut self, node: Self::Node) -> Result> { + if let Expr::IsNotNull(ref inner) = node { + if find_window(inner)?.is_some() { + return Ok(Transformed::yes(Expr::Literal( + datafusion::common::ScalarValue::Boolean(Some(true)), + None, + ))); + } + } + Ok(Transformed::no(node)) + } +} + +/// Checks that window UDFs (tumble/hop/session) are not used outside aggregates +pub(crate) struct TimeWindowUdfChecker; + +impl TreeNodeVisitor<'_> for TimeWindowUdfChecker { + type Node = LogicalPlan; + + fn f_up(&mut self, node: &Self::Node) -> Result { + if let LogicalPlan::Projection(projection) = node { + for expr in &projection.expr { + if let Some(window) = find_window(expr)? { + return plan_err!( + "Window function {:?} can only be used as a GROUP BY expression in an aggregate", + window + ); + } + } + } + Ok(TreeNodeRecursion::Continue) + } +} diff --git a/src/sql/planner/plan/window_fn.rs b/src/sql/planner/plan/window_fn.rs new file mode 100644 index 00000000..66f673d1 --- /dev/null +++ b/src/sql/planner/plan/window_fn.rs @@ -0,0 +1,178 @@ +use std::sync::Arc; + +use datafusion::common::tree_node::Transformed; +use datafusion::common::{Result as DFResult, plan_err, tree_node::TreeNodeRewriter}; +use datafusion::logical_expr; +use datafusion::logical_expr::expr::WindowFunctionParams; +use datafusion::logical_expr::{ + Expr, Extension, LogicalPlan, Projection, Sort, Window, expr::WindowFunction, +}; +use tracing::debug; + +use crate::sql::planner::extension::key_calculation::{KeyCalculationExtension, KeysOrExprs}; +use crate::sql::planner::extension::window_fn::WindowFunctionExtension; +use crate::sql::planner::plan::{WindowDetectingVisitor, extract_column}; +use crate::sql::types::{WindowType, fields_with_qualifiers, schema_from_df_fields}; + +pub(crate) struct WindowFunctionRewriter; + +fn get_window_and_name(expr: &Expr) -> DFResult<(WindowFunction, String)> { + match expr { + Expr::Alias(alias) => { + let (window, _) = get_window_and_name(&alias.expr)?; + Ok((window, alias.name.clone())) + } + Expr::WindowFunction(window_function) => { + Ok((*window_function.clone(), expr.name_for_alias()?)) + } + _ => plan_err!("Expect a column or alias expression, not {:?}", expr), + } +} + +impl TreeNodeRewriter for WindowFunctionRewriter { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DFResult> { + let LogicalPlan::Window(window) = node else { + return Ok(Transformed::no(node)); + }; + + debug!( + "Rewriting window function: {:?}", + LogicalPlan::Window(window.clone()) + ); + + let mut window_detecting_visitor = WindowDetectingVisitor::default(); + window + .input + .visit_with_subqueries(&mut window_detecting_visitor)?; + + let Some(input_window) = window_detecting_visitor.window else { + return plan_err!("Window functions require already windowed input"); + }; + if matches!(input_window, WindowType::Session { .. }) { + return plan_err!("Window functions do not support session windows"); + } + + let input_window_fields = window_detecting_visitor.fields; + + let Window { + input, window_expr, .. + } = window; + + if window_expr.len() != 1 { + return plan_err!("Window functions require exactly one window expression"); + } + + let (WindowFunction { fun, params }, original_name) = get_window_and_name(&window_expr[0])?; + + let mut window_field: Vec<_> = params + .partition_by + .iter() + .enumerate() + .filter_map(|(index, expr)| { + if let Some(column) = extract_column(expr) { + let Ok(input_field) = input + .schema() + .field_with_name(column.relation.as_ref(), &column.name) + else { + return Some(plan_err!( + "Column {} not found in input schema", + column.name + )); + }; + if input_window_fields.contains(&(column.relation.as_ref(), input_field).into()) + { + return Some(Ok((input_field.clone(), index))); + } + } + None + }) + .collect::>()?; + + if window_field.len() != 1 { + return plan_err!( + "Window function requires exactly one window expression in partition_by" + ); + } + + let (_window_field, index) = window_field.pop().unwrap(); + let mut additional_keys = params.partition_by.clone(); + additional_keys.remove(index); + let key_count = additional_keys.len(); + + let params = WindowFunctionParams { + args: params.args, + partition_by: additional_keys.clone(), + order_by: params.order_by, + window_frame: params.window_frame, + null_treatment: params.null_treatment, + }; + + let new_window_func = WindowFunction { fun, params }; + + let mut key_projection_expressions: Vec<_> = additional_keys + .iter() + .enumerate() + .map(|(index, expression)| expression.clone().alias(format!("_key_{index}"))) + .collect(); + + key_projection_expressions.extend( + fields_with_qualifiers(input.schema()) + .iter() + .map(|field| Expr::Column(field.qualified_column())), + ); + + let auto_schema = + Projection::try_new(key_projection_expressions.clone(), input.clone())?.schema; + let mut key_fields = fields_with_qualifiers(&auto_schema) + .iter() + .take(additional_keys.len()) + .cloned() + .collect::>(); + key_fields.extend(fields_with_qualifiers(input.schema())); + let key_schema = Arc::new(schema_from_df_fields(&key_fields)?); + + let key_projection = LogicalPlan::Projection(Projection::try_new_with_schema( + key_projection_expressions, + input.clone(), + key_schema, + )?); + + let key_plan = LogicalPlan::Extension(Extension { + node: Arc::new(KeyCalculationExtension::new( + key_projection, + KeysOrExprs::Keys((0..key_count).collect()), + )), + }); + + let mut sort_expressions: Vec<_> = additional_keys + .iter() + .map(|partition| logical_expr::expr::Sort { + expr: partition.clone(), + asc: true, + nulls_first: false, + }) + .collect(); + sort_expressions.extend(new_window_func.params.order_by.clone()); + + let shuffle = LogicalPlan::Sort(Sort { + expr: sort_expressions, + input: Arc::new(key_plan), + fetch: None, + }); + + let window_expr = + Expr::WindowFunction(Box::new(new_window_func)).alias_if_changed(original_name)?; + + let rewritten_window_plan = + LogicalPlan::Window(Window::try_new(vec![window_expr], Arc::new(shuffle))?); + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(WindowFunctionExtension::new( + rewritten_window_plan, + (0..key_count).collect(), + )), + }))) + } +} diff --git a/src/sql/planner/rewrite/async_udf_rewriter.rs b/src/sql/planner/rewrite/async_udf_rewriter.rs new file mode 100644 index 00000000..def3c4ef --- /dev/null +++ b/src/sql/planner/rewrite/async_udf_rewriter.rs @@ -0,0 +1,118 @@ +use crate::sql::planner::extension::remote_table::RemoteTableExtension; +use crate::sql::planner::extension::{ASYNC_RESULT_FIELD, AsyncUDFExtension}; +use crate::sql::planner::mod_prelude::StreamSchemaProvider; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; +use datafusion::common::{Column, Result as DFResult, TableReference, plan_err}; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; +use std::sync::Arc; +use std::time::Duration; + +type AsyncSplitResult = (String, AsyncOptions, Vec); + +#[derive(Debug, Clone, Copy)] +pub struct AsyncOptions { + pub ordered: bool, + pub max_concurrency: usize, + pub timeout: Duration, +} + +pub struct AsyncUdfRewriter<'a> { + provider: &'a StreamSchemaProvider, +} + +impl<'a> AsyncUdfRewriter<'a> { + pub fn new(provider: &'a StreamSchemaProvider) -> Self { + Self { provider } + } + + fn split_async( + expr: Expr, + provider: &StreamSchemaProvider, + ) -> DFResult<(Expr, Option)> { + let mut found: Option<(String, AsyncOptions, Vec)> = None; + let expr = expr.transform_up(|e| { + if let Expr::ScalarFunction(ScalarFunction { func: udf, args }) = &e { + if let Some(opts) = provider.get_async_udf_options(udf.name()) { + if found + .replace((udf.name().to_string(), opts, args.clone())) + .is_some() + { + return plan_err!( + "multiple async calls in the same expression, which is not allowed" + ); + } + return Ok(Transformed::yes(Expr::Column(Column::new_unqualified( + ASYNC_RESULT_FIELD, + )))); + } + } + Ok(Transformed::no(e)) + })?; + + Ok((expr.data, found)) + } +} + +impl TreeNodeRewriter for AsyncUdfRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DFResult> { + let LogicalPlan::Projection(mut projection) = node else { + for e in node.expressions() { + if let (_, Some((udf, _, _))) = Self::split_async(e.clone(), self.provider)? { + return plan_err!( + "async UDFs are only supported in projections, but {udf} was called in another context" + ); + } + } + return Ok(Transformed::no(node)); + }; + + let mut args = None; + for e in projection.expr.iter_mut() { + let (new_e, Some(udf)) = Self::split_async(e.clone(), self.provider)? else { + continue; + }; + if let Some((prev, _, _)) = args.replace(udf) { + return plan_err!( + "Projection contains multiple async UDFs, which is not supported \ + \n(hint: two async UDF calls, {} and {}, appear in the same SELECT statement)", + prev, + args.unwrap().0 + ); + } + *e = new_e; + } + + let Some((name, opts, arg_exprs)) = args else { + return Ok(Transformed::no(LogicalPlan::Projection(projection))); + }; + + let input = if matches!(*projection.input, LogicalPlan::Projection(..)) { + Arc::new(LogicalPlan::Extension(Extension { + node: Arc::new(RemoteTableExtension { + input: (*projection.input).clone(), + name: TableReference::bare("subquery_projection"), + schema: projection.input.schema().clone(), + materialize: false, + }), + })) + } else { + projection.input + }; + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(AsyncUDFExtension { + input, + name, + arg_exprs, + final_exprs: projection.expr, + ordered: opts.ordered, + max_concurrency: opts.max_concurrency, + timeout: opts.timeout, + final_schema: projection.schema, + }), + }))) + } +} diff --git a/src/sql/planner/rewrite/mod.rs b/src/sql/planner/rewrite/mod.rs new file mode 100644 index 00000000..20b2e9bb --- /dev/null +++ b/src/sql/planner/rewrite/mod.rs @@ -0,0 +1,27 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +pub mod async_udf_rewriter; +pub mod row_time; +pub mod sink_input_rewriter; +pub mod source_metadata_visitor; +pub mod source_rewriter; +pub mod time_window; +pub mod unnest_rewriter; + +pub use async_udf_rewriter::{AsyncOptions, AsyncUdfRewriter}; +pub use row_time::RowTimeRewriter; +pub use sink_input_rewriter::SinkInputRewriter; +pub use source_metadata_visitor::SourceMetadataVisitor; +pub use source_rewriter::SourceRewriter; +pub use time_window::{TimeWindowNullCheckRemover, TimeWindowUdfChecker, is_time_window}; +pub use unnest_rewriter::{UNNESTED_COL, UnnestRewriter}; diff --git a/src/sql/planner/rewrite/row_time.rs b/src/sql/planner/rewrite/row_time.rs new file mode 100644 index 00000000..51309feb --- /dev/null +++ b/src/sql/planner/rewrite/row_time.rs @@ -0,0 +1,39 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion::common::{Column, Result as DFResult}; +use datafusion::logical_expr::Expr; + +use crate::sql::types::TIMESTAMP_FIELD; + +/// Rewrites `row_time()` scalar function calls to a column reference on `_timestamp`. +pub struct RowTimeRewriter {} + +impl TreeNodeRewriter for RowTimeRewriter { + type Node = Expr; + + fn f_down(&mut self, node: Self::Node) -> DFResult> { + if let Expr::ScalarFunction(func) = &node + && func.name() == "row_time" + { + let transformed = Expr::Column(Column { + relation: None, + name: TIMESTAMP_FIELD.to_string(), + spans: Default::default(), + }) + .alias("row_time()"); + return Ok(Transformed::yes(transformed)); + } + Ok(Transformed::no(node)) + } +} diff --git a/src/sql/planner/rewrite/sink_input_rewriter.rs b/src/sql/planner/rewrite/sink_input_rewriter.rs new file mode 100644 index 00000000..e6b6a0bd --- /dev/null +++ b/src/sql/planner/rewrite/sink_input_rewriter.rs @@ -0,0 +1,46 @@ +use crate::sql::planner::extension::sink::SinkExtension; +use crate::sql::planner::extension::{NamedNode, StreamExtension}; +use datafusion::common::Result as DFResult; +use datafusion::common::tree_node::{Transformed, TreeNodeRecursion, TreeNodeRewriter}; +use datafusion::logical_expr::{Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use std::collections::HashMap; +use std::sync::Arc; + +type SinkInputs = HashMap>; + +/// Merges inputs for sinks with the same name to avoid duplicate sinks in the plan. +pub struct SinkInputRewriter<'a> { + sink_inputs: &'a mut SinkInputs, + pub was_removed: bool, +} + +impl<'a> SinkInputRewriter<'a> { + pub fn new(sink_inputs: &'a mut SinkInputs) -> Self { + Self { + sink_inputs, + was_removed: false, + } + } +} + +impl TreeNodeRewriter for SinkInputRewriter<'_> { + type Node = LogicalPlan; + + fn f_down(&mut self, node: Self::Node) -> DFResult> { + if let LogicalPlan::Extension(extension) = &node { + if let Some(sink_node) = extension.node.as_any().downcast_ref::() { + if let Some(named_node) = sink_node.node_name() { + if let Some(inputs) = self.sink_inputs.remove(&named_node) { + let new_node = LogicalPlan::Extension(Extension { + node: Arc::new(sink_node.with_exprs_and_inputs(vec![], inputs)?), + }); + return Ok(Transformed::new(new_node, true, TreeNodeRecursion::Jump)); + } else { + self.was_removed = true; + } + } + } + } + Ok(Transformed::no(node)) + } +} diff --git a/src/sql/planner/rewrite/source_metadata_visitor.rs b/src/sql/planner/rewrite/source_metadata_visitor.rs new file mode 100644 index 00000000..168ff712 --- /dev/null +++ b/src/sql/planner/rewrite/source_metadata_visitor.rs @@ -0,0 +1,57 @@ +use crate::sql::planner::extension::sink::SinkExtension; +use crate::sql::planner::extension::table_source::TableSourceExtension; +use crate::sql::planner::mod_prelude::StreamSchemaProvider; +use datafusion::common::Result as DFResult; +use datafusion::common::tree_node::{TreeNodeRecursion, TreeNodeVisitor}; +use datafusion::logical_expr::{Extension, LogicalPlan}; +use std::collections::HashSet; + +/// Collects connection IDs from source and sink nodes in the logical plan. +pub struct SourceMetadataVisitor<'a> { + schema_provider: &'a StreamSchemaProvider, + pub connection_ids: HashSet, +} + +impl<'a> SourceMetadataVisitor<'a> { + pub fn new(schema_provider: &'a StreamSchemaProvider) -> Self { + Self { + schema_provider, + connection_ids: HashSet::new(), + } + } + + fn get_connection_id(&self, node: &LogicalPlan) -> Option { + let LogicalPlan::Extension(Extension { node }) = node else { + return None; + }; + + let table_name = match node.name() { + "TableSourceExtension" => { + let ext = node.as_any().downcast_ref::()?; + ext.name.to_string() + } + "SinkExtension" => { + let ext = node.as_any().downcast_ref::()?; + ext.name.to_string() + } + _ => return None, + }; + + let table = self.schema_provider.get_catalog_table(&table_name)?; + match table { + crate::sql::catalog::table::Table::ConnectorTable(t) => t.id, + _ => None, + } + } +} + +impl TreeNodeVisitor<'_> for SourceMetadataVisitor<'_> { + type Node = LogicalPlan; + + fn f_down(&mut self, node: &Self::Node) -> DFResult { + if let Some(id) = self.get_connection_id(node) { + self.connection_ids.insert(id); + } + Ok(TreeNodeRecursion::Continue) + } +} diff --git a/src/sql/planner/rewrite/source_rewriter.rs b/src/sql/planner/rewrite/source_rewriter.rs new file mode 100644 index 00000000..209c3288 --- /dev/null +++ b/src/sql/planner/rewrite/source_rewriter.rs @@ -0,0 +1,272 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; +use std::time::Duration; + +use datafusion::common::ScalarValue; +use datafusion::common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion::common::{Column, DataFusionError, Result as DFResult, TableReference, plan_err}; +use datafusion::logical_expr::{ + self, BinaryExpr, Expr, Extension, LogicalPlan, Projection, TableScan, +}; + +use crate::sql::catalog::connector_table::ConnectorTable; +use crate::sql::catalog::field_spec::FieldSpec; +use crate::sql::catalog::table::Table; +use crate::sql::catalog::utils::add_timestamp_field; +use crate::sql::planner::StreamSchemaProvider; +use crate::sql::planner::extension::remote_table::RemoteTableExtension; +use crate::sql::planner::extension::watermark_node::WatermarkNode; +use crate::sql::types::TIMESTAMP_FIELD; + +/// Rewrites table scans into proper source nodes with projections and watermarks. +pub struct SourceRewriter<'a> { + pub(crate) schema_provider: &'a StreamSchemaProvider, +} + +impl SourceRewriter<'_> { + fn watermark_expression(table: &ConnectorTable) -> DFResult { + match table.watermark_field.clone() { + Some(watermark_field) => table + .fields + .iter() + .find_map(|f| { + if f.field().name() == &watermark_field { + return match f { + FieldSpec::Struct(field) | FieldSpec::Metadata { field, .. } => { + Some(Expr::Column(Column { + relation: None, + name: field.name().to_string(), + spans: Default::default(), + })) + } + FieldSpec::Virtual { expression, .. } => Some(*expression.clone()), + }; + } + None + }) + .ok_or_else(|| { + DataFusionError::Plan(format!("Watermark field {watermark_field} not found")) + }), + None => Ok(Expr::BinaryExpr(BinaryExpr { + left: Box::new(Expr::Column(Column { + relation: None, + name: TIMESTAMP_FIELD.to_string(), + spans: Default::default(), + })), + op: logical_expr::Operator::Minus, + right: Box::new(Expr::Literal( + ScalarValue::DurationNanosecond(Some(Duration::from_secs(1).as_nanos() as i64)), + None, + )), + })), + } + } + + fn projection_expressions( + table: &ConnectorTable, + qualifier: &TableReference, + projection: &Option>, + ) -> DFResult> { + let mut expressions: Vec = table + .fields + .iter() + .map(|field| match field { + FieldSpec::Struct(field) | FieldSpec::Metadata { field, .. } => { + Expr::Column(Column { + relation: Some(qualifier.clone()), + name: field.name().to_string(), + spans: Default::default(), + }) + } + FieldSpec::Virtual { field, expression } => expression + .clone() + .alias_qualified(Some(qualifier.clone()), field.name().to_string()), + }) + .collect(); + + if let Some(proj) = projection { + expressions = proj.iter().map(|i| expressions[*i].clone()).collect(); + } + + if let Some(event_time_field) = table.event_time_field.clone() { + let expr = table + .fields + .iter() + .find_map(|f| { + if f.field().name() == &event_time_field { + return match f { + FieldSpec::Struct(field) | FieldSpec::Metadata { field, .. } => { + Some(Expr::Column(Column { + relation: Some(qualifier.clone()), + name: field.name().to_string(), + spans: Default::default(), + })) + } + FieldSpec::Virtual { expression, .. } => Some(*expression.clone()), + }; + } + None + }) + .ok_or_else(|| { + DataFusionError::Plan(format!("Event time field {event_time_field} not found")) + })?; + + expressions + .push(expr.alias_qualified(Some(qualifier.clone()), TIMESTAMP_FIELD.to_string())); + } else { + expressions.push(Expr::Column(Column::new( + Some(qualifier.clone()), + TIMESTAMP_FIELD, + ))); + } + + Ok(expressions) + } + + fn projection(&self, table_scan: &TableScan, table: &ConnectorTable) -> DFResult { + let qualifier = table_scan.table_name.clone(); + + // TODO: replace with TableSourceExtension when available + let source_input = LogicalPlan::TableScan(table_scan.clone()); + + Ok(LogicalPlan::Projection(Projection::try_new( + Self::projection_expressions(table, &qualifier, &table_scan.projection)?, + Arc::new(source_input), + )?)) + } + + fn mutate_connector_table( + &self, + table_scan: &TableScan, + table: &ConnectorTable, + ) -> DFResult> { + let input = self.projection(table_scan, table)?; + + let schema = input.schema().clone(); + let remote = LogicalPlan::Extension(Extension { + node: Arc::new(RemoteTableExtension { + input, + name: table_scan.table_name.to_owned(), + schema, + materialize: true, + }), + }); + + let watermark_node = WatermarkNode::new( + remote, + table_scan.table_name.clone(), + Self::watermark_expression(table)?, + ) + .map_err(|err| { + DataFusionError::Internal(format!("failed to create watermark expression: {err}")) + })?; + + Ok(Transformed::yes(LogicalPlan::Extension(Extension { + node: Arc::new(watermark_node), + }))) + } + + fn mutate_table_from_query( + &self, + table_scan: &TableScan, + logical_plan: &LogicalPlan, + ) -> DFResult> { + let column_expressions: Vec<_> = if let Some(projection) = &table_scan.projection { + logical_plan + .schema() + .columns() + .into_iter() + .enumerate() + .filter_map(|(i, col)| { + if projection.contains(&i) { + Some(Expr::Column(col)) + } else { + None + } + }) + .collect() + } else { + logical_plan + .schema() + .columns() + .into_iter() + .map(Expr::Column) + .collect() + }; + + let target_columns: Vec<_> = table_scan.projected_schema.columns().into_iter().collect(); + + let expressions = column_expressions + .into_iter() + .zip(target_columns) + .map(|(expr, col)| expr.alias_qualified(col.relation, col.name)) + .collect(); + + let projection = LogicalPlan::Projection(Projection::try_new_with_schema( + expressions, + Arc::new(logical_plan.clone()), + table_scan.projected_schema.clone(), + )?); + + Ok(Transformed::yes(projection)) + } +} + +impl TreeNodeRewriter for SourceRewriter<'_> { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DFResult> { + let LogicalPlan::TableScan(mut table_scan) = node else { + return Ok(Transformed::no(node)); + }; + + let table_name = table_scan.table_name.table(); + let table = self + .schema_provider + .get_catalog_table(table_name) + .ok_or_else(|| DataFusionError::Plan(format!("Table {table_name} not found")))?; + + match table { + Table::ConnectorTable(table) => self.mutate_connector_table(&table_scan, table), + Table::LookupTable(_table) => { + // TODO: implement LookupSource extension + plan_err!("Lookup tables are not yet supported") + } + Table::MemoryTable { + name, + fields: _, + logical_plan, + } => { + let Some(logical_plan) = logical_plan else { + return plan_err!( + "Can't query from memory table {} without first inserting into it", + name + ); + }; + table_scan.projected_schema = add_timestamp_field( + table_scan.projected_schema.clone(), + Some(table_scan.table_name.clone()), + )?; + self.mutate_table_from_query(&table_scan, logical_plan) + } + Table::TableFromQuery { + name: _, + logical_plan, + } => self.mutate_table_from_query(&table_scan, logical_plan), + Table::PreviewSink { .. } => Err(DataFusionError::Plan( + "can't select from a preview sink".to_string(), + )), + } + } +} diff --git a/src/sql/planner/rewrite/time_window.rs b/src/sql/planner/rewrite/time_window.rs new file mode 100644 index 00000000..104c0cca --- /dev/null +++ b/src/sql/planner/rewrite/time_window.rs @@ -0,0 +1,83 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use datafusion::common::tree_node::{ + Transformed, TreeNodeRecursion, TreeNodeRewriter, TreeNodeVisitor, +}; +use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue, plan_err}; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{Expr, LogicalPlan}; + +/// Returns the time window function name if the expression is one (tumble/hop/session). +pub fn is_time_window(expr: &Expr) -> Option<&str> { + if let Expr::ScalarFunction(ScalarFunction { func, args: _ }) = expr { + match func.name() { + "tumble" | "hop" | "session" => return Some(func.name()), + _ => {} + } + } + None +} + +struct TimeWindowExprChecker {} + +impl TreeNodeVisitor<'_> for TimeWindowExprChecker { + type Node = Expr; + + fn f_down(&mut self, node: &Self::Node) -> DFResult { + if let Some(w) = is_time_window(node) { + return plan_err!( + "time window function {} is not allowed in this context. \ + Are you missing a GROUP BY clause?", + w + ); + } + Ok(TreeNodeRecursion::Continue) + } +} + +/// Visitor that checks an entire LogicalPlan for misplaced time window UDFs. +pub struct TimeWindowUdfChecker {} + +impl TreeNodeVisitor<'_> for TimeWindowUdfChecker { + type Node = LogicalPlan; + + fn f_down(&mut self, node: &Self::Node) -> DFResult { + use datafusion::common::tree_node::TreeNode; + node.expressions().iter().try_for_each(|expr| { + let mut checker = TimeWindowExprChecker {}; + expr.visit(&mut checker)?; + Ok::<(), DataFusionError>(()) + })?; + Ok(TreeNodeRecursion::Continue) + } +} + +/// Removes `IS NOT NULL` checks wrapping time window functions, +/// replacing them with `true` since time windows are never null. +pub struct TimeWindowNullCheckRemover {} + +impl TreeNodeRewriter for TimeWindowNullCheckRemover { + type Node = Expr; + + fn f_down(&mut self, node: Self::Node) -> DFResult> { + if let Expr::IsNotNull(expr) = &node + && is_time_window(expr).is_some() + { + return Ok(Transformed::yes(Expr::Literal( + ScalarValue::Boolean(Some(true)), + None, + ))); + } + Ok(Transformed::no(node)) + } +} diff --git a/src/sql/planner/rewrite/unnest_rewriter.rs b/src/sql/planner/rewrite/unnest_rewriter.rs new file mode 100644 index 00000000..2a9eabda --- /dev/null +++ b/src/sql/planner/rewrite/unnest_rewriter.rs @@ -0,0 +1,178 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use datafusion::arrow::datatypes::DataType; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRewriter}; +use datafusion::common::{Column, Result as DFResult, plan_err}; +use datafusion::logical_expr::expr::ScalarFunction; +use datafusion::logical_expr::{ColumnUnnestList, Expr, LogicalPlan, Projection, Unnest}; + +use crate::sql::types::{DFField, fields_with_qualifiers, schema_from_df_fields}; + +pub const UNNESTED_COL: &str = "__unnested"; + +/// Rewrites projections containing `unnest()` calls into proper Unnest logical plans. +pub struct UnnestRewriter {} + +impl UnnestRewriter { + fn split_unnest(expr: Expr) -> DFResult<(Expr, Option)> { + let mut captured: Option = None; + + let expr = expr.transform_up(|e| { + if let Expr::ScalarFunction(ScalarFunction { func: udf, args }) = &e + && udf.name() == "unnest" + { + match args.len() { + 1 => { + if captured.replace(args[0].clone()).is_some() { + return plan_err!( + "Multiple unnests in expression, which is not allowed" + ); + } + return Ok(Transformed::yes(Expr::Column(Column::new_unqualified( + UNNESTED_COL, + )))); + } + n => { + panic!("Unnest has wrong number of arguments (expected 1, found {n})"); + } + } + } + Ok(Transformed::no(e)) + })?; + + Ok((expr.data, captured)) + } +} + +impl TreeNodeRewriter for UnnestRewriter { + type Node = LogicalPlan; + + fn f_up(&mut self, node: Self::Node) -> DFResult> { + let LogicalPlan::Projection(projection) = &node else { + if node.expressions().iter().any(|e| { + let e = Self::split_unnest(e.clone()); + e.is_err() || e.unwrap().1.is_some() + }) { + return plan_err!("unnest is only supported in SELECT statements"); + } + return Ok(Transformed::no(node)); + }; + + let mut unnest = None; + let exprs = projection + .expr + .clone() + .into_iter() + .enumerate() + .map(|(i, expr)| { + let (expr, opt) = Self::split_unnest(expr)?; + let is_unnest = if let Some(e) = opt { + if let Some(prev) = unnest.replace((e, i)) + && &prev != unnest.as_ref().unwrap() + { + return plan_err!( + "Projection contains multiple unnests, which is not currently supported" + ); + } + true + } else { + false + }; + + Ok((expr, is_unnest)) + }) + .collect::>>()?; + + if let Some((unnest_inner, unnest_idx)) = unnest { + let produce_list = Arc::new(LogicalPlan::Projection(Projection::try_new( + exprs + .iter() + .cloned() + .map(|(e, is_unnest)| { + if is_unnest { + unnest_inner.clone().alias(UNNESTED_COL) + } else { + e + } + }) + .collect(), + projection.input.clone(), + )?)); + + let unnest_fields = fields_with_qualifiers(produce_list.schema()) + .iter() + .enumerate() + .map(|(i, f)| { + if i == unnest_idx { + let DataType::List(inner) = f.data_type() else { + return plan_err!( + "Argument '{}' to unnest is not a List", + f.qualified_name() + ); + }; + Ok(DFField::new_unqualified( + UNNESTED_COL, + inner.data_type().clone(), + inner.is_nullable(), + )) + } else { + Ok((*f).clone()) + } + }) + .collect::>>()?; + + let unnest_node = LogicalPlan::Unnest(Unnest { + exec_columns: vec![ + DFField::from(produce_list.schema().qualified_field(unnest_idx)) + .qualified_column(), + ], + input: produce_list, + list_type_columns: vec![( + unnest_idx, + ColumnUnnestList { + output_column: Column::new_unqualified(UNNESTED_COL), + depth: 1, + }, + )], + struct_type_columns: vec![], + dependency_indices: vec![], + schema: Arc::new(schema_from_df_fields(&unnest_fields)?), + options: Default::default(), + }); + + let output_node = LogicalPlan::Projection(Projection::try_new( + exprs + .iter() + .enumerate() + .map(|(i, (expr, has_unnest))| { + if *has_unnest { + expr.clone() + } else { + Expr::Column( + DFField::from(unnest_node.schema().qualified_field(i)) + .qualified_column(), + ) + } + }) + .collect(), + Arc::new(unnest_node), + )?); + + Ok(Transformed::yes(output_node)) + } else { + Ok(Transformed::no(LogicalPlan::Projection(projection.clone()))) + } + } +} diff --git a/src/sql/planner/schema_provider.rs b/src/sql/planner/schema_provider.rs new file mode 100644 index 00000000..d860fd6c --- /dev/null +++ b/src/sql/planner/schema_provider.rs @@ -0,0 +1,360 @@ +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use datafusion::arrow::datatypes::{self as datatypes, DataType, Field, Schema}; +use datafusion::common::{Result, plan_err}; +use datafusion::datasource::DefaultTableSource; +use datafusion::error::DataFusionError; +use datafusion::execution::{FunctionRegistry, SessionStateDefaults}; +use datafusion::logical_expr::expr_rewriter::FunctionRewrite; +use datafusion::logical_expr::planner::ExprPlanner; +use datafusion::logical_expr::{ + AggregateUDF, Expr, LogicalPlan, ScalarUDF, TableSource, WindowUDF, +}; +use datafusion::optimizer::Analyzer; +use datafusion::sql::TableReference; +use datafusion::sql::planner::ContextProvider; +use unicase::UniCase; + +use crate::sql::catalog::table::Table as CatalogTable; +use crate::sql::planner::schemas::window_arrow_struct; +use crate::sql::types::{PlaceholderUdf, PlanningOptions}; + +#[derive(Clone, Default)] +pub struct StreamSchemaProvider { + pub source_defs: HashMap, + tables: HashMap, StreamTable>, + catalog_tables: HashMap, CatalogTable>, + pub functions: HashMap>, + pub aggregate_functions: HashMap>, + pub window_functions: HashMap>, + config_options: datafusion::config::ConfigOptions, + pub expr_planners: Vec>, + pub planning_options: PlanningOptions, + pub analyzer: Analyzer, +} + +#[derive(Clone, Debug)] +pub enum StreamTable { + Source { + name: String, + schema: Arc, + event_time_field: Option, + watermark_field: Option, + }, + Sink { + name: String, + schema: Arc, + }, + Memory { + name: String, + logical_plan: Option, + }, +} + +impl StreamTable { + pub fn name(&self) -> &str { + match self { + StreamTable::Source { name, .. } => name, + StreamTable::Sink { name, .. } => name, + StreamTable::Memory { name, .. } => name, + } + } + + pub fn get_fields(&self) -> Vec> { + match self { + StreamTable::Source { schema, .. } => schema.fields().to_vec(), + StreamTable::Sink { schema, .. } => schema.fields().to_vec(), + StreamTable::Memory { .. } => vec![], + } + } +} + +#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] +pub struct LogicalBatchInput { + pub table_name: String, + pub schema: Arc, +} + +#[async_trait::async_trait] +impl datafusion::datasource::TableProvider for LogicalBatchInput { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn schema(&self) -> Arc { + self.schema.clone() + } + + fn table_type(&self) -> datafusion::datasource::TableType { + datafusion::datasource::TableType::Temporary + } + + async fn scan( + &self, + _state: &dyn datafusion::catalog::Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + Ok(Arc::new(crate::sql::physical::FsMemExec::new( + self.table_name.clone(), + self.schema.clone(), + ))) + } +} + +fn create_table(table_name: String, schema: Arc) -> Arc { + let table_provider = LogicalBatchInput { table_name, schema }; + let wrapped = Arc::new(table_provider); + let provider = DefaultTableSource::new(wrapped); + Arc::new(provider) +} + +impl StreamSchemaProvider { + pub fn new() -> Self { + let mut registry = Self { + ..Default::default() + }; + + registry + .register_udf(PlaceholderUdf::with_return( + "hop", + vec![ + DataType::Interval(datatypes::IntervalUnit::MonthDayNano), + DataType::Interval(datatypes::IntervalUnit::MonthDayNano), + ], + window_arrow_struct(), + )) + .unwrap(); + + registry + .register_udf(PlaceholderUdf::with_return( + "tumble", + vec![DataType::Interval(datatypes::IntervalUnit::MonthDayNano)], + window_arrow_struct(), + )) + .unwrap(); + + registry + .register_udf(PlaceholderUdf::with_return( + "session", + vec![DataType::Interval(datatypes::IntervalUnit::MonthDayNano)], + window_arrow_struct(), + )) + .unwrap(); + + registry + .register_udf(PlaceholderUdf::with_return( + "unnest", + vec![DataType::List(Arc::new(Field::new( + "field", + DataType::Utf8, + true, + )))], + DataType::Utf8, + )) + .unwrap(); + + registry + .register_udf(PlaceholderUdf::with_return( + "row_time", + vec![], + DataType::Timestamp(datatypes::TimeUnit::Nanosecond, None), + )) + .unwrap(); + + for p in SessionStateDefaults::default_scalar_functions() { + registry.register_udf(p).unwrap(); + } + for p in SessionStateDefaults::default_aggregate_functions() { + registry.register_udaf(p).unwrap(); + } + for p in SessionStateDefaults::default_window_functions() { + registry.register_udwf(p).unwrap(); + } + for p in SessionStateDefaults::default_expr_planners() { + registry.register_expr_planner(p).unwrap(); + } + + registry + } + + pub fn add_source_table( + &mut self, + name: String, + schema: Arc, + event_time_field: Option, + watermark_field: Option, + ) { + self.tables.insert( + UniCase::new(name.clone()), + StreamTable::Source { + name, + schema, + event_time_field, + watermark_field, + }, + ); + } + + pub fn add_sink_table(&mut self, name: String, schema: Arc) { + self.tables.insert( + UniCase::new(name.clone()), + StreamTable::Sink { name, schema }, + ); + } + + pub fn insert_table(&mut self, table: StreamTable) { + self.tables + .insert(UniCase::new(table.name().to_string()), table); + } + + pub fn get_table(&self, table_name: impl Into) -> Option<&StreamTable> { + self.tables.get(&UniCase::new(table_name.into())) + } + + pub fn get_table_mut(&mut self, table_name: impl Into) -> Option<&mut StreamTable> { + self.tables.get_mut(&UniCase::new(table_name.into())) + } + + pub fn insert_catalog_table(&mut self, table: CatalogTable) { + self.catalog_tables + .insert(UniCase::new(table.name().to_string()), table); + } + + pub fn get_catalog_table(&self, table_name: impl Into) -> Option<&CatalogTable> { + self.catalog_tables.get(&UniCase::new(table_name.into())) + } + + pub fn get_catalog_table_mut( + &mut self, + table_name: impl Into, + ) -> Option<&mut CatalogTable> { + self.catalog_tables + .get_mut(&UniCase::new(table_name.into())) + } + + pub fn get_async_udf_options( + &self, + _name: &str, + ) -> Option { + // TODO: implement async UDF lookup + None + } +} + +impl ContextProvider for StreamSchemaProvider { + fn get_table_source(&self, name: TableReference) -> Result> { + let table = self + .get_table(name.to_string()) + .ok_or_else(|| DataFusionError::Plan(format!("Table {name} not found")))?; + + let fields = table.get_fields(); + let schema = Arc::new(Schema::new_with_metadata( + fields + .iter() + .map(|f| f.as_ref().clone()) + .collect::>(), + HashMap::new(), + )); + Ok(create_table(name.to_string(), schema)) + } + + fn get_function_meta(&self, name: &str) -> Option> { + self.functions.get(name).cloned() + } + + fn get_aggregate_meta(&self, name: &str) -> Option> { + self.aggregate_functions.get(name).cloned() + } + + fn get_variable_type(&self, _variable_names: &[String]) -> Option { + None + } + + fn options(&self) -> &datafusion::config::ConfigOptions { + &self.config_options + } + + fn get_window_meta(&self, name: &str) -> Option> { + self.window_functions.get(name).cloned() + } + + fn udf_names(&self) -> Vec { + self.functions.keys().cloned().collect() + } + + fn udaf_names(&self) -> Vec { + self.aggregate_functions.keys().cloned().collect() + } + + fn udwf_names(&self) -> Vec { + self.window_functions.keys().cloned().collect() + } + + fn get_expr_planners(&self) -> &[Arc] { + &self.expr_planners + } +} + +impl FunctionRegistry for StreamSchemaProvider { + fn udfs(&self) -> HashSet { + self.functions.keys().cloned().collect() + } + + fn udf(&self, name: &str) -> Result> { + if let Some(f) = self.functions.get(name) { + Ok(Arc::clone(f)) + } else { + plan_err!("No UDF with name {name}") + } + } + + fn udaf(&self, name: &str) -> Result> { + if let Some(f) = self.aggregate_functions.get(name) { + Ok(Arc::clone(f)) + } else { + plan_err!("No UDAF with name {name}") + } + } + + fn udwf(&self, name: &str) -> Result> { + if let Some(f) = self.window_functions.get(name) { + Ok(Arc::clone(f)) + } else { + plan_err!("No UDWF with name {name}") + } + } + + fn register_function_rewrite( + &mut self, + rewrite: Arc, + ) -> Result<()> { + self.analyzer.add_function_rewrite(rewrite); + Ok(()) + } + + fn register_udf(&mut self, udf: Arc) -> Result>> { + Ok(self.functions.insert(udf.name().to_string(), udf)) + } + + fn register_udaf(&mut self, udaf: Arc) -> Result>> { + Ok(self + .aggregate_functions + .insert(udaf.name().to_string(), udaf)) + } + + fn register_udwf(&mut self, udwf: Arc) -> Result>> { + Ok(self.window_functions.insert(udwf.name().to_string(), udwf)) + } + + fn register_expr_planner(&mut self, expr_planner: Arc) -> Result<()> { + self.expr_planners.push(expr_planner); + Ok(()) + } + + fn expr_planners(&self) -> Vec> { + self.expr_planners.clone() + } +} diff --git a/src/sql/planner/schemas.rs b/src/sql/planner/schemas.rs new file mode 100644 index 00000000..f903db83 --- /dev/null +++ b/src/sql/planner/schemas.rs @@ -0,0 +1,5 @@ +// Re-export schema utilities from catalog::utils. +// Kept for backward compatibility with existing planner imports. +pub use crate::sql::catalog::utils::{ + add_timestamp_field, add_timestamp_field_arrow, has_timestamp_field, window_arrow_struct, +}; diff --git a/src/sql/planner/sql_to_plan.rs b/src/sql/planner/sql_to_plan.rs new file mode 100644 index 00000000..049cd18e --- /dev/null +++ b/src/sql/planner/sql_to_plan.rs @@ -0,0 +1,22 @@ +use datafusion::common::Result; +use datafusion::logical_expr::LogicalPlan; +use datafusion::sql::sqlparser::ast::Statement; +use tracing::debug; + +use crate::sql::planner::StreamSchemaProvider; + +/// Stage 2: Statement → LogicalPlan +/// +/// Converts a parsed SQL AST statement into a DataFusion logical plan +/// using the StreamSchemaProvider as the catalog context. +pub fn statement_to_plan( + statement: Statement, + schema_provider: &StreamSchemaProvider, +) -> Result { + let sql_to_rel = datafusion::sql::planner::SqlToRel::new(schema_provider); + let plan = sql_to_rel.sql_statement_to_plan(statement)?; + + debug!("Logical plan:\n{}", plan.display_graphviz()); + + Ok(plan) +} diff --git a/src/sql/planner/udafs.rs b/src/sql/planner/udafs.rs new file mode 100644 index 00000000..9685c2d4 --- /dev/null +++ b/src/sql/planner/udafs.rs @@ -0,0 +1,31 @@ +use datafusion::arrow::array::ArrayRef; +use datafusion::error::Result; +use datafusion::physical_plan::Accumulator; +use datafusion::scalar::ScalarValue; +use std::fmt::Debug; + +/// Fake UDAF used just for plan-time placeholder. +#[derive(Debug)] +pub struct EmptyUdaf {} + +impl Accumulator for EmptyUdaf { + fn update_batch(&mut self, _: &[ArrayRef]) -> Result<()> { + unreachable!() + } + + fn evaluate(&self) -> Result { + unreachable!() + } + + fn size(&self) -> usize { + unreachable!() + } + + fn state(&self) -> Result> { + unreachable!() + } + + fn merge_batch(&mut self, _: &[ArrayRef]) -> Result<()> { + unreachable!() + } +} diff --git a/src/sql/types/data_type.rs b/src/sql/types/data_type.rs new file mode 100644 index 00000000..57edc3c9 --- /dev/null +++ b/src/sql/types/data_type.rs @@ -0,0 +1,144 @@ +use std::sync::Arc; + +use datafusion::arrow::datatypes::{ + DECIMAL_DEFAULT_SCALE, DECIMAL128_MAX_PRECISION, DataType, Field, IntervalUnit, TimeUnit, +}; +use datafusion::common::{Result, plan_datafusion_err, plan_err}; + +use crate::types::FsExtensionType; + +pub fn convert_data_type( + sql_type: &datafusion::sql::sqlparser::ast::DataType, +) -> Result<(DataType, Option)> { + use datafusion::sql::sqlparser::ast::ArrayElemTypeDef; + use datafusion::sql::sqlparser::ast::DataType as SQLDataType; + + match sql_type { + SQLDataType::Array(ArrayElemTypeDef::AngleBracket(inner_sql_type)) + | SQLDataType::Array(ArrayElemTypeDef::SquareBracket(inner_sql_type, _)) => { + let (data_type, extension) = convert_simple_data_type(inner_sql_type)?; + + Ok(( + DataType::List(Arc::new(FsExtensionType::add_metadata( + extension, + Field::new("field", data_type, true), + ))), + None, + )) + } + SQLDataType::Array(ArrayElemTypeDef::None) => { + plan_err!("Arrays with unspecified type is not supported") + } + other => convert_simple_data_type(other), + } +} + +fn convert_simple_data_type( + sql_type: &datafusion::sql::sqlparser::ast::DataType, +) -> Result<(DataType, Option)> { + use datafusion::sql::sqlparser::ast::DataType as SQLDataType; + use datafusion::sql::sqlparser::ast::{ExactNumberInfo, TimezoneInfo}; + + if matches!(sql_type, SQLDataType::JSON) { + return Ok((DataType::Utf8, Some(FsExtensionType::JSON))); + } + + let dt = match sql_type { + SQLDataType::Boolean | SQLDataType::Bool => Ok(DataType::Boolean), + SQLDataType::TinyInt(_) => Ok(DataType::Int8), + SQLDataType::SmallInt(_) | SQLDataType::Int2(_) => Ok(DataType::Int16), + SQLDataType::Int(_) | SQLDataType::Integer(_) | SQLDataType::Int4(_) => Ok(DataType::Int32), + SQLDataType::BigInt(_) | SQLDataType::Int8(_) => Ok(DataType::Int64), + SQLDataType::TinyIntUnsigned(_) => Ok(DataType::UInt8), + SQLDataType::SmallIntUnsigned(_) | SQLDataType::Int2Unsigned(_) => Ok(DataType::UInt16), + SQLDataType::IntUnsigned(_) + | SQLDataType::UnsignedInteger + | SQLDataType::Int4Unsigned(_) => Ok(DataType::UInt32), + SQLDataType::BigIntUnsigned(_) | SQLDataType::Int8Unsigned(_) => Ok(DataType::UInt64), + SQLDataType::Float(_) => Ok(DataType::Float32), + SQLDataType::Real | SQLDataType::Float4 => Ok(DataType::Float32), + SQLDataType::Double(_) | SQLDataType::DoublePrecision | SQLDataType::Float8 => { + Ok(DataType::Float64) + } + SQLDataType::Char(_) + | SQLDataType::Varchar(_) + | SQLDataType::Text + | SQLDataType::String(_) => Ok(DataType::Utf8), + SQLDataType::Timestamp(None, TimezoneInfo::None) | SQLDataType::Datetime(_) => { + Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)) + } + SQLDataType::Timestamp(Some(precision), TimezoneInfo::None) => match *precision { + 0 => Ok(DataType::Timestamp(TimeUnit::Second, None)), + 3 => Ok(DataType::Timestamp(TimeUnit::Millisecond, None)), + 6 => Ok(DataType::Timestamp(TimeUnit::Microsecond, None)), + 9 => Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)), + _ => { + return plan_err!( + "unsupported precision {} -- supported precisions are 0 (seconds), \ + 3 (milliseconds), 6 (microseconds), and 9 (nanoseconds)", + precision + ); + } + }, + SQLDataType::Date => Ok(DataType::Date32), + SQLDataType::Time(None, tz_info) => { + if matches!(tz_info, TimezoneInfo::None) + || matches!(tz_info, TimezoneInfo::WithoutTimeZone) + { + Ok(DataType::Time64(TimeUnit::Nanosecond)) + } else { + return plan_err!("Unsupported SQL type {sql_type:?}"); + } + } + SQLDataType::Numeric(exact_number_info) | SQLDataType::Decimal(exact_number_info) => { + let (precision, scale) = match *exact_number_info { + ExactNumberInfo::None => (None, None), + ExactNumberInfo::Precision(precision) => (Some(precision), None), + ExactNumberInfo::PrecisionAndScale(precision, scale) => { + (Some(precision), Some(scale)) + } + }; + make_decimal_type(precision, scale) + } + SQLDataType::Bytea => Ok(DataType::Binary), + SQLDataType::Interval => Ok(DataType::Interval(IntervalUnit::MonthDayNano)), + SQLDataType::Struct(fields, _) => { + let fields: Vec<_> = fields + .iter() + .map(|f| { + Ok::<_, datafusion::error::DataFusionError>(Arc::new(Field::new( + f.field_name + .as_ref() + .ok_or_else(|| { + plan_datafusion_err!("anonymous struct fields are not allowed") + })? + .to_string(), + convert_data_type(&f.field_type)?.0, + true, + ))) + }) + .collect::>()?; + Ok(DataType::Struct(fields.into())) + } + _ => return plan_err!("Unsupported SQL type {sql_type:?}"), + }; + + Ok((dt?, None)) +} + +fn make_decimal_type(precision: Option, scale: Option) -> Result { + let (precision, scale) = match (precision, scale) { + (Some(p), Some(s)) => (p as u8, s as i8), + (Some(p), None) => (p as u8, 0), + (None, Some(_)) => return plan_err!("Cannot specify only scale for decimal data type"), + (None, None) => (DECIMAL128_MAX_PRECISION, DECIMAL_DEFAULT_SCALE), + }; + + if precision == 0 || precision > DECIMAL128_MAX_PRECISION || scale.unsigned_abs() > precision { + plan_err!( + "Decimal(precision = {precision}, scale = {scale}) should satisfy `0 < precision <= 38`, and `scale <= precision`." + ) + } else { + Ok(DataType::Decimal128(precision, scale)) + } +} diff --git a/src/sql/types/df_field.rs b/src/sql/types/df_field.rs new file mode 100644 index 00000000..3797adb2 --- /dev/null +++ b/src/sql/types/df_field.rs @@ -0,0 +1,141 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use datafusion::arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion::common::{Column, DFSchema, Result, TableReference}; + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct DFField { + qualifier: Option, + field: FieldRef, +} + +impl From<(Option, FieldRef)> for DFField { + fn from(value: (Option, FieldRef)) -> Self { + Self { + qualifier: value.0, + field: value.1, + } + } +} + +impl From<(Option<&TableReference>, &Field)> for DFField { + fn from(value: (Option<&TableReference>, &Field)) -> Self { + Self { + qualifier: value.0.cloned(), + field: Arc::new(value.1.clone()), + } + } +} + +impl From for (Option, FieldRef) { + fn from(value: DFField) -> Self { + (value.qualifier, value.field) + } +} + +impl DFField { + pub fn new( + qualifier: Option, + name: impl Into, + data_type: DataType, + nullable: bool, + ) -> Self { + Self { + qualifier, + field: Arc::new(Field::new(name, data_type, nullable)), + } + } + + pub fn new_unqualified(name: &str, data_type: DataType, nullable: bool) -> Self { + DFField { + qualifier: None, + field: Arc::new(Field::new(name, data_type, nullable)), + } + } + + pub fn name(&self) -> &String { + self.field.name() + } + + pub fn data_type(&self) -> &DataType { + self.field.data_type() + } + + pub fn is_nullable(&self) -> bool { + self.field.is_nullable() + } + + pub fn metadata(&self) -> &HashMap { + self.field.metadata() + } + + pub fn qualified_name(&self) -> String { + if let Some(qualifier) = &self.qualifier { + format!("{}.{}", qualifier, self.field.name()) + } else { + self.field.name().to_owned() + } + } + + pub fn qualified_column(&self) -> Column { + Column { + relation: self.qualifier.clone(), + name: self.field.name().to_string(), + spans: Default::default(), + } + } + + pub fn unqualified_column(&self) -> Column { + Column { + relation: None, + name: self.field.name().to_string(), + spans: Default::default(), + } + } + + pub fn qualifier(&self) -> Option<&TableReference> { + self.qualifier.as_ref() + } + + pub fn field(&self) -> &FieldRef { + &self.field + } + + pub fn strip_qualifier(mut self) -> Self { + self.qualifier = None; + self + } + + pub fn with_nullable(mut self, nullable: bool) -> Self { + let f = self.field().as_ref().clone().with_nullable(nullable); + self.field = f.into(); + self + } + + pub fn with_metadata(mut self, metadata: HashMap) -> Self { + let f = self.field().as_ref().clone().with_metadata(metadata); + self.field = f.into(); + self + } +} + +pub fn fields_with_qualifiers(schema: &DFSchema) -> Vec { + schema + .fields() + .iter() + .enumerate() + .map(|(i, f)| (schema.qualified_field(i).0.cloned(), f.clone()).into()) + .collect() +} + +pub fn schema_from_df_fields(fields: &[DFField]) -> Result { + schema_from_df_fields_with_metadata(fields, HashMap::new()) +} + +pub fn schema_from_df_fields_with_metadata( + fields: &[DFField], + metadata: HashMap, +) -> Result { + DFSchema::new_with_metadata(fields.iter().map(|t| t.clone().into()).collect(), metadata) +} diff --git a/src/sql/types/mod.rs b/src/sql/types/mod.rs new file mode 100644 index 00000000..25c67574 --- /dev/null +++ b/src/sql/types/mod.rs @@ -0,0 +1,50 @@ +mod data_type; +mod df_field; +pub(crate) mod placeholder_udf; +mod stream_schema; +mod window; + +use std::time::Duration; + +pub use data_type::convert_data_type; +pub use df_field::{ + DFField, fields_with_qualifiers, schema_from_df_fields, schema_from_df_fields_with_metadata, +}; +pub(crate) use placeholder_udf::PlaceholderUdf; +pub use stream_schema::StreamSchema; +pub(crate) use window::WindowBehavior; +pub use window::{WindowType, find_window, get_duration}; + +pub const TIMESTAMP_FIELD: &str = "_timestamp"; + +#[derive(Clone, Debug, Eq, PartialEq)] +pub enum ProcessingMode { + Append, + Update, +} + +#[derive(Clone, Debug)] +pub struct SqlConfig { + pub default_parallelism: usize, +} + +impl Default for SqlConfig { + fn default() -> Self { + Self { + default_parallelism: 4, + } + } +} + +#[derive(Clone)] +pub struct PlanningOptions { + pub ttl: Duration, +} + +impl Default for PlanningOptions { + fn default() -> Self { + Self { + ttl: Duration::from_secs(24 * 60 * 60), + } + } +} diff --git a/src/sql/types/placeholder_udf.rs b/src/sql/types/placeholder_udf.rs new file mode 100644 index 00000000..5cf96d28 --- /dev/null +++ b/src/sql/types/placeholder_udf.rs @@ -0,0 +1,58 @@ +use std::any::Any; +use std::fmt::{Debug, Formatter}; +use std::sync::Arc; + +use datafusion::arrow::datatypes::DataType; +use datafusion::common::Result; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, +}; + +#[allow(clippy::type_complexity)] +pub(crate) struct PlaceholderUdf { + name: String, + signature: Signature, + return_type: Arc Result + Send + Sync + 'static>, +} + +impl Debug for PlaceholderUdf { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "PlaceholderUDF<{}>", self.name) + } +} + +impl ScalarUDFImpl for PlaceholderUdf { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + &self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, args: &[DataType]) -> Result { + (self.return_type)(args) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + unimplemented!("PlaceholderUdf should never be called at execution time"); + } +} + +impl PlaceholderUdf { + pub fn with_return( + name: impl Into, + args: Vec, + ret: DataType, + ) -> Arc { + Arc::new(ScalarUDF::new_from_impl(PlaceholderUdf { + name: name.into(), + signature: Signature::exact(args, Volatility::Volatile), + return_type: Arc::new(move |_| Ok(ret.clone())), + })) + } +} diff --git a/src/sql/types/stream_schema.rs b/src/sql/types/stream_schema.rs new file mode 100644 index 00000000..e981111b --- /dev/null +++ b/src/sql/types/stream_schema.rs @@ -0,0 +1,76 @@ +use std::sync::Arc; + +use datafusion::arrow::datatypes::{Field, Schema, SchemaRef}; +use datafusion::common::Result; + +use super::TIMESTAMP_FIELD; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct StreamSchema { + pub schema: SchemaRef, + pub timestamp_index: usize, + pub key_indices: Option>, +} + +impl StreamSchema { + pub fn new(schema: SchemaRef, timestamp_index: usize, key_indices: Option>) -> Self { + Self { + schema, + timestamp_index, + key_indices, + } + } + + pub fn new_unkeyed(schema: SchemaRef, timestamp_index: usize) -> Self { + Self { + schema, + timestamp_index, + key_indices: None, + } + } + + pub fn from_fields(fields: Vec) -> Self { + let schema = Arc::new(Schema::new(fields)); + let timestamp_index = schema + .column_with_name(TIMESTAMP_FIELD) + .map(|(i, _)| i) + .unwrap_or(0); + Self { + schema, + timestamp_index, + key_indices: None, + } + } + + pub fn from_schema_keys(schema: SchemaRef, key_indices: Vec) -> Result { + let timestamp_index = schema + .column_with_name(TIMESTAMP_FIELD) + .ok_or_else(|| { + datafusion::error::DataFusionError::Plan(format!( + "no {TIMESTAMP_FIELD} field in schema, schema is {schema:?}" + )) + })? + .0; + Ok(Self { + schema, + timestamp_index, + key_indices: Some(key_indices), + }) + } + + pub fn from_schema_unkeyed(schema: SchemaRef) -> Result { + let timestamp_index = schema + .column_with_name(TIMESTAMP_FIELD) + .ok_or_else(|| { + datafusion::error::DataFusionError::Plan(format!( + "no {TIMESTAMP_FIELD} field in schema" + )) + })? + .0; + Ok(Self { + schema, + timestamp_index, + key_indices: None, + }) + } +} diff --git a/src/sql/types/window.rs b/src/sql/types/window.rs new file mode 100644 index 00000000..9687974a --- /dev/null +++ b/src/sql/types/window.rs @@ -0,0 +1,95 @@ +use std::time::Duration; + +use datafusion::common::{Result, plan_err}; +use datafusion::logical_expr::Expr; + +use super::DFField; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum WindowType { + Tumbling { width: Duration }, + Sliding { width: Duration, slide: Duration }, + Session { gap: Duration }, + Instant, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub(crate) enum WindowBehavior { + FromOperator { + window: WindowType, + window_field: DFField, + window_index: usize, + is_nested: bool, + }, + InData, +} + +pub fn get_duration(expression: &Expr) -> Result { + use datafusion::common::ScalarValue; + + match expression { + Expr::Literal(ScalarValue::IntervalDayTime(Some(val)), _) => { + Ok(Duration::from_secs((val.days as u64) * 24 * 60 * 60) + + Duration::from_millis(val.milliseconds as u64)) + } + Expr::Literal(ScalarValue::IntervalMonthDayNano(Some(val)), _) => { + if val.months != 0 { + return datafusion::common::not_impl_err!( + "Windows do not support durations specified as months" + ); + } + Ok(Duration::from_secs((val.days as u64) * 24 * 60 * 60) + + Duration::from_nanos(val.nanoseconds as u64)) + } + _ => plan_err!( + "unsupported Duration expression, expect duration literal, not {}", + expression + ), + } +} + +pub fn find_window(expression: &Expr) -> Result> { + use datafusion::logical_expr::expr::Alias; + use datafusion::logical_expr::expr::ScalarFunction; + + match expression { + Expr::ScalarFunction(ScalarFunction { func: fun, args }) => match fun.name() { + "hop" => { + if args.len() != 2 { + unreachable!(); + } + let slide = get_duration(&args[0])?; + let width = get_duration(&args[1])?; + if width.as_nanos() % slide.as_nanos() != 0 { + return plan_err!( + "hop() width {:?} must be a multiple of slide {:?}", + width, + slide + ); + } + if slide == width { + Ok(Some(WindowType::Tumbling { width })) + } else { + Ok(Some(WindowType::Sliding { width, slide })) + } + } + "tumble" => { + if args.len() != 1 { + unreachable!("wrong number of arguments for tumble(), expect one"); + } + let width = get_duration(&args[0])?; + Ok(Some(WindowType::Tumbling { width })) + } + "session" => { + if args.len() != 1 { + unreachable!("wrong number of arguments for session(), expected one"); + } + let gap = get_duration(&args[0])?; + Ok(Some(WindowType::Session { gap })) + } + _ => Ok(None), + }, + Expr::Alias(Alias { expr, .. }) => find_window(expr), + _ => Ok(None), + } +} diff --git a/src/storage/task/rocksdb_storage.rs b/src/storage/task/rocksdb_storage.rs index 31709a51..714a9143 100644 --- a/src/storage/task/rocksdb_storage.rs +++ b/src/storage/task/rocksdb_storage.rs @@ -103,11 +103,19 @@ impl TaskStorage for RocksDBTaskStorage { }; let mut batch = WriteBatch::default(); - batch.put_cf(&cf_meta, key, bincode::serialize(&meta)?); + batch.put_cf( + &cf_meta, + key, + bincode::serde::encode_to_vec(&meta, bincode::config::standard())?, + ); batch.put_cf(&cf_conf, key, &task_info.config_bytes); if let Some(ref module) = task_info.module_bytes { - batch.put_cf(&cf_payl, key, bincode::serialize(module)?); + batch.put_cf( + &cf_payl, + key, + bincode::serde::encode_to_vec(module, bincode::config::standard())?, + ); } self.db @@ -124,10 +132,15 @@ impl TaskStorage for RocksDBTaskStorage { .get_cf(&cf, key)? .ok_or_else(|| anyhow!("Task {} not found", task_name))?; - let mut meta: TaskMetadata = bincode::deserialize(&raw)?; + let (mut meta, _): (TaskMetadata, _) = + bincode::serde::decode_from_slice(&raw, bincode::config::standard())?; meta.state = new_state; - self.db.put_cf(&cf, key, bincode::serialize(&meta)?)?; + self.db.put_cf( + &cf, + key, + bincode::serde::encode_to_vec(&meta, bincode::config::standard())?, + )?; Ok(()) } @@ -140,10 +153,15 @@ impl TaskStorage for RocksDBTaskStorage { .get_cf(&cf, key)? .ok_or_else(|| anyhow!("Task {} not found", task_name))?; - let mut meta: TaskMetadata = bincode::deserialize(&raw)?; + let (mut meta, _): (TaskMetadata, _) = + bincode::serde::decode_from_slice(&raw, bincode::config::standard())?; meta.checkpoint_id = checkpoint_id; - self.db.put_cf(&cf, key, bincode::serialize(&meta)?)?; + self.db.put_cf( + &cf, + key, + bincode::serde::encode_to_vec(&meta, bincode::config::standard())?, + )?; Ok(()) } @@ -174,9 +192,17 @@ impl TaskStorage for RocksDBTaskStorage { let module_bytes = self .db .get_cf(&self.get_cf(CF_PAYLOAD)?, key)? - .and_then(|b| bincode::deserialize::(&b).ok()); - - let meta: TaskMetadata = bincode::deserialize(&meta_raw)?; + .and_then(|b| { + bincode::serde::decode_from_slice::( + &b, + bincode::config::standard(), + ) + .ok() + .map(|(v, _)| v) + }); + + let (meta, _): (TaskMetadata, _) = + bincode::serde::decode_from_slice(&meta_raw, bincode::config::standard())?; Ok(StoredTaskInfo { name: task_name.to_string(), diff --git a/src/types/arrow_ext.rs b/src/types/arrow_ext.rs new file mode 100644 index 00000000..701bf8e4 --- /dev/null +++ b/src/types/arrow_ext.rs @@ -0,0 +1,169 @@ +use std::collections::HashMap; +use std::fmt::{Display, Formatter}; +use std::time::SystemTime; + +use datafusion::arrow::datatypes::{DataType, Field, TimeUnit}; + +pub struct DisplayAsSql<'a>(pub &'a DataType); + +impl Display for DisplayAsSql<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + match self.0 { + DataType::Boolean => write!(f, "BOOLEAN"), + DataType::Int8 | DataType::Int16 | DataType::Int32 => write!(f, "INT"), + DataType::Int64 => write!(f, "BIGINT"), + DataType::UInt8 | DataType::UInt16 | DataType::UInt32 => write!(f, "INT UNSIGNED"), + DataType::UInt64 => write!(f, "BIGINT UNSIGNED"), + DataType::Float16 | DataType::Float32 => write!(f, "FLOAT"), + DataType::Float64 => write!(f, "DOUBLE"), + DataType::Timestamp(_, _) => write!(f, "TIMESTAMP"), + DataType::Date32 => write!(f, "DATE"), + DataType::Date64 => write!(f, "DATETIME"), + DataType::Time32(_) => write!(f, "TIME"), + DataType::Time64(_) => write!(f, "TIME"), + DataType::Duration(_) => write!(f, "INTERVAL"), + DataType::Interval(_) => write!(f, "INTERVAL"), + DataType::Binary | DataType::FixedSizeBinary(_) | DataType::LargeBinary => { + write!(f, "BYTEA") + } + DataType::Utf8 | DataType::LargeUtf8 => write!(f, "TEXT"), + DataType::List(inner) => { + write!(f, "{}[]", DisplayAsSql(inner.data_type())) + } + dt => write!(f, "{dt}"), + } + } +} + +/// Arrow extension type markers for FunctionStream-specific semantics. +#[derive(Copy, Clone, Eq, PartialEq, Debug)] +pub enum FsExtensionType { + JSON, +} + +impl FsExtensionType { + pub fn from_map(map: &HashMap) -> Option { + match map.get("ARROW:extension:name")?.as_str() { + "functionstream.json" => Some(Self::JSON), + _ => None, + } + } + + pub fn add_metadata(v: Option, field: Field) -> Field { + if let Some(v) = v { + let mut m = HashMap::new(); + match v { + FsExtensionType::JSON => { + m.insert( + "ARROW:extension:name".to_string(), + "functionstream.json".to_string(), + ); + } + } + field.with_metadata(m) + } else { + field + } + } +} + +pub trait GetArrowType { + fn arrow_type() -> DataType; +} + +pub trait GetArrowSchema { + fn arrow_schema() -> datafusion::arrow::datatypes::Schema; +} + +impl GetArrowType for T +where + T: GetArrowSchema, +{ + fn arrow_type() -> DataType { + DataType::Struct(Self::arrow_schema().fields.clone()) + } +} + +impl GetArrowType for bool { + fn arrow_type() -> DataType { + DataType::Boolean + } +} + +impl GetArrowType for i8 { + fn arrow_type() -> DataType { + DataType::Int8 + } +} + +impl GetArrowType for i16 { + fn arrow_type() -> DataType { + DataType::Int16 + } +} + +impl GetArrowType for i32 { + fn arrow_type() -> DataType { + DataType::Int32 + } +} + +impl GetArrowType for i64 { + fn arrow_type() -> DataType { + DataType::Int64 + } +} + +impl GetArrowType for u8 { + fn arrow_type() -> DataType { + DataType::UInt8 + } +} + +impl GetArrowType for u16 { + fn arrow_type() -> DataType { + DataType::UInt16 + } +} + +impl GetArrowType for u32 { + fn arrow_type() -> DataType { + DataType::UInt32 + } +} + +impl GetArrowType for u64 { + fn arrow_type() -> DataType { + DataType::UInt64 + } +} + +impl GetArrowType for f32 { + fn arrow_type() -> DataType { + DataType::Float32 + } +} + +impl GetArrowType for f64 { + fn arrow_type() -> DataType { + DataType::Float64 + } +} + +impl GetArrowType for String { + fn arrow_type() -> DataType { + DataType::Utf8 + } +} + +impl GetArrowType for Vec { + fn arrow_type() -> DataType { + DataType::Binary + } +} + +impl GetArrowType for SystemTime { + fn arrow_type() -> DataType { + DataType::Timestamp(TimeUnit::Nanosecond, None) + } +} diff --git a/src/types/control.rs b/src/types/control.rs new file mode 100644 index 00000000..efdc754e --- /dev/null +++ b/src/types/control.rs @@ -0,0 +1,152 @@ +use std::collections::HashMap; +use std::time::SystemTime; + +use super::message::CheckpointBarrier; + +/// Control messages sent from the controller to worker tasks. +#[derive(Debug, Clone)] +pub enum ControlMessage { + Checkpoint(CheckpointBarrier), + Stop { + mode: StopMode, + }, + Commit { + epoch: u32, + commit_data: HashMap>>, + }, + LoadCompacted { + compacted: CompactionResult, + }, + NoOp, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub enum StopMode { + Graceful, + Immediate, +} + +#[derive(Debug, Clone)] +pub struct CompactionResult { + pub operator_id: String, + pub compacted_tables: HashMap, +} + +#[derive(Debug, Clone)] +pub struct TableCheckpointMetadata { + pub table_type: TableType, + pub data: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TableType { + GlobalKeyValue, + ExpiringKeyedTimeTable, +} + +/// Responses sent from worker tasks back to the controller. +#[derive(Debug, Clone)] +pub enum ControlResp { + CheckpointEvent(CheckpointEvent), + CheckpointCompleted(CheckpointCompleted), + TaskStarted { + node_id: u32, + task_index: usize, + start_time: SystemTime, + }, + TaskFinished { + node_id: u32, + task_index: usize, + }, + TaskFailed { + node_id: u32, + task_index: usize, + error: TaskError, + }, + Error { + node_id: u32, + operator_id: String, + task_index: usize, + message: String, + details: String, + }, +} + +#[derive(Debug, Clone)] +pub struct CheckpointCompleted { + pub checkpoint_epoch: u32, + pub node_id: u32, + pub operator_id: String, + pub subtask_metadata: SubtaskCheckpointMetadata, +} + +#[derive(Debug, Clone)] +pub struct SubtaskCheckpointMetadata { + pub subtask_index: u32, + pub start_time: u64, + pub finish_time: u64, + pub watermark: Option, + pub bytes: u64, + pub table_metadata: HashMap, + pub table_configs: HashMap, +} + +#[derive(Debug, Clone)] +pub struct TableSubtaskCheckpointMetadata { + pub subtask_index: u32, + pub table_type: TableType, + pub data: Vec, +} + +#[derive(Debug, Clone)] +pub struct TableConfig { + pub table_type: TableType, + pub config: Vec, + pub state_version: u32, +} + +#[derive(Debug, Clone)] +pub struct CheckpointEvent { + pub checkpoint_epoch: u32, + pub node_id: u32, + pub operator_id: String, + pub subtask_index: u32, + pub time: SystemTime, + pub event_type: TaskCheckpointEventType, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TaskCheckpointEventType { + StartedAlignment, + StartedCheckpointing, + FinishedOperatorSetup, + FinishedSync, + FinishedCommit, +} + +#[derive(Debug, Clone)] +pub struct TaskError { + pub job_id: String, + pub node_id: u32, + pub operator_id: String, + pub operator_subtask: u64, + pub error: String, + pub error_domain: ErrorDomain, + pub retry_hint: RetryHint, + pub details: String, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ErrorDomain { + User, + Internal, + External, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RetryHint { + NoRetry, + WithBackoff, +} diff --git a/src/types/date.rs b/src/types/date.rs new file mode 100644 index 00000000..c18e31a7 --- /dev/null +++ b/src/types/date.rs @@ -0,0 +1,70 @@ +use serde::Serialize; +use std::convert::TryFrom; + +#[derive(Debug, Clone, Copy, Eq, PartialEq, PartialOrd, Hash, Serialize)] +pub enum DatePart { + Year, + Month, + Week, + Day, + Hour, + Minute, + Second, + Millisecond, + Microsecond, + Nanosecond, + DayOfWeek, + DayOfYear, +} + +impl TryFrom<&str> for DatePart { + type Error = String; + + fn try_from(value: &str) -> Result { + match value.to_lowercase().as_str() { + "year" => Ok(DatePart::Year), + "month" => Ok(DatePart::Month), + "week" => Ok(DatePart::Week), + "day" => Ok(DatePart::Day), + "hour" => Ok(DatePart::Hour), + "minute" => Ok(DatePart::Minute), + "second" => Ok(DatePart::Second), + "millisecond" => Ok(DatePart::Millisecond), + "microsecond" => Ok(DatePart::Microsecond), + "nanosecond" => Ok(DatePart::Nanosecond), + "dow" => Ok(DatePart::DayOfWeek), + "doy" => Ok(DatePart::DayOfYear), + _ => Err(format!("'{value}' is not a valid DatePart")), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, PartialOrd, Serialize)] +pub enum DateTruncPrecision { + Year, + Quarter, + Month, + Week, + Day, + Hour, + Minute, + Second, +} + +impl TryFrom<&str> for DateTruncPrecision { + type Error = String; + + fn try_from(value: &str) -> Result { + match value.to_lowercase().as_str() { + "year" => Ok(DateTruncPrecision::Year), + "quarter" => Ok(DateTruncPrecision::Quarter), + "month" => Ok(DateTruncPrecision::Month), + "week" => Ok(DateTruncPrecision::Week), + "day" => Ok(DateTruncPrecision::Day), + "hour" => Ok(DateTruncPrecision::Hour), + "minute" => Ok(DateTruncPrecision::Minute), + "second" => Ok(DateTruncPrecision::Second), + _ => Err(format!("'{value}' is not a valid DateTruncPrecision")), + } + } +} diff --git a/src/types/debezium.rs b/src/types/debezium.rs new file mode 100644 index 00000000..3c9f4747 --- /dev/null +++ b/src/types/debezium.rs @@ -0,0 +1,136 @@ +use bincode::{Decode, Encode}; +use serde::{Deserialize, Serialize}; +use std::convert::TryFrom; +use std::fmt::Debug; + +pub trait Key: + Debug + Clone + Encode + Decode<()> + std::hash::Hash + PartialEq + Eq + Send + 'static +{ +} +impl + std::hash::Hash + PartialEq + Eq + Send + 'static> Key + for T +{ +} + +pub trait Data: Debug + Clone + Encode + Decode<()> + Send + 'static {} +impl + Send + 'static> Data for T {} + +#[derive(Debug, Clone, PartialEq, Encode, Decode, Serialize, Deserialize)] +pub enum UpdatingData { + Retract(T), + Update { old: T, new: T }, + Append(T), +} + +impl UpdatingData { + pub fn lower(&self) -> T { + match self { + UpdatingData::Retract(_) => panic!("cannot lower retractions"), + UpdatingData::Update { new, .. } => new.clone(), + UpdatingData::Append(t) => t.clone(), + } + } + + pub fn unwrap_append(&self) -> &T { + match self { + UpdatingData::Append(t) => t, + _ => panic!("UpdatingData is not an append"), + } + } +} + +#[derive(Clone, Encode, Decode, Debug, Serialize, Deserialize, PartialEq)] +#[serde(try_from = "DebeziumShadow")] +pub struct Debezium { + pub before: Option, + pub after: Option, + pub op: DebeziumOp, +} + +#[derive(Clone, Encode, Decode, Debug, Serialize, Deserialize, PartialEq)] +struct DebeziumShadow { + before: Option, + after: Option, + op: DebeziumOp, +} + +impl TryFrom> for Debezium { + type Error = &'static str; + + fn try_from(value: DebeziumShadow) -> Result { + match (value.op, &value.before, &value.after) { + (DebeziumOp::Create, _, None) => { + Err("`after` must be set for Debezium create messages") + } + (DebeziumOp::Update, None, _) => { + Err("`before` must be set for Debezium update messages") + } + (DebeziumOp::Update, _, None) => { + Err("`after` must be set for Debezium update messages") + } + (DebeziumOp::Delete, None, _) => { + Err("`before` must be set for Debezium delete messages") + } + _ => Ok(Debezium { + before: value.before, + after: value.after, + op: value.op, + }), + } + } +} + +#[derive(Copy, Clone, Encode, Decode, Debug, PartialEq)] +pub enum DebeziumOp { + Create, + Update, + Delete, +} + +#[allow(clippy::to_string_trait_impl)] +impl ToString for DebeziumOp { + fn to_string(&self) -> String { + match self { + DebeziumOp::Create => "c", + DebeziumOp::Update => "u", + DebeziumOp::Delete => "d", + } + .to_string() + } +} + +impl<'de> Deserialize<'de> for DebeziumOp { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let s = String::deserialize(deserializer)?; + match s.as_str() { + "c" | "r" => Ok(DebeziumOp::Create), + "u" => Ok(DebeziumOp::Update), + "d" => Ok(DebeziumOp::Delete), + _ => Err(serde::de::Error::custom(format!("Invalid DebeziumOp {s}"))), + } + } +} + +impl Serialize for DebeziumOp { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + match self { + DebeziumOp::Create => serializer.serialize_str("c"), + DebeziumOp::Update => serializer.serialize_str("u"), + DebeziumOp::Delete => serializer.serialize_str("d"), + } + } +} + +#[derive(Copy, Clone, Encode, Decode, Debug, PartialEq, Serialize, Deserialize)] +pub enum JoinType { + Inner, + Left, + Right, + Full, +} diff --git a/src/types/df.rs b/src/types/df.rs new file mode 100644 index 00000000..30b4eb9c --- /dev/null +++ b/src/types/df.rs @@ -0,0 +1,394 @@ +use datafusion::arrow::array::builder::{ArrayBuilder, make_builder}; +use datafusion::arrow::array::{RecordBatch, TimestampNanosecondArray}; +use datafusion::arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaBuilder, TimeUnit}; +use datafusion::arrow::error::ArrowError; +use datafusion::common::{DataFusionError, Result as DFResult}; +use std::sync::Arc; + +use super::TIMESTAMP_FIELD; +use crate::sql::types::StreamSchema; + +pub type FsSchemaRef = Arc; + +/// Core streaming schema with timestamp and key tracking. +/// Analogous to Arroyo's `ArroyoSchema`. +#[derive(Debug, Clone, Eq, PartialEq, Hash)] +pub struct FsSchema { + pub schema: Arc, + pub timestamp_index: usize, + key_indices: Option>, + routing_key_indices: Option>, +} + +impl FsSchema { + pub fn new( + schema: Arc, + timestamp_index: usize, + key_indices: Option>, + routing_key_indices: Option>, + ) -> Self { + Self { + schema, + timestamp_index, + key_indices, + routing_key_indices, + } + } + + pub fn new_unkeyed(schema: Arc, timestamp_index: usize) -> Self { + Self { + schema, + timestamp_index, + key_indices: None, + routing_key_indices: None, + } + } + + pub fn new_keyed(schema: Arc, timestamp_index: usize, key_indices: Vec) -> Self { + Self { + schema, + timestamp_index, + key_indices: Some(key_indices), + routing_key_indices: None, + } + } + + pub fn from_fields(mut fields: Vec) -> Self { + if !fields.iter().any(|f| f.name() == TIMESTAMP_FIELD) { + fields.push(Field::new( + TIMESTAMP_FIELD, + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )); + } + + Self::from_schema_keys(Arc::new(Schema::new(fields)), vec![]).unwrap() + } + + pub fn from_schema_unkeyed(schema: Arc) -> DFResult { + let timestamp_index = schema + .column_with_name(TIMESTAMP_FIELD) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "no {TIMESTAMP_FIELD} field in schema, schema is {schema:?}" + )) + })? + .0; + + Ok(Self { + schema, + timestamp_index, + key_indices: None, + routing_key_indices: None, + }) + } + + pub fn from_schema_keys(schema: Arc, key_indices: Vec) -> DFResult { + let timestamp_index = schema + .column_with_name(TIMESTAMP_FIELD) + .ok_or_else(|| { + DataFusionError::Plan(format!( + "no {TIMESTAMP_FIELD} field in schema, schema is {schema:?}" + )) + })? + .0; + + Ok(Self { + schema, + timestamp_index, + key_indices: Some(key_indices), + routing_key_indices: None, + }) + } + + pub fn schema_without_timestamp(&self) -> Schema { + let mut builder = SchemaBuilder::from(self.schema.fields()); + builder.remove(self.timestamp_index); + builder.finish() + } + + pub fn remove_timestamp_column(&self, batch: &mut RecordBatch) { + batch.remove_column(self.timestamp_index); + } + + pub fn builders(&self) -> Vec> { + self.schema + .fields + .iter() + .map(|f| make_builder(f.data_type(), 8)) + .collect() + } + + pub fn timestamp_column<'a>(&self, batch: &'a RecordBatch) -> &'a TimestampNanosecondArray { + batch + .column(self.timestamp_index) + .as_any() + .downcast_ref::() + .unwrap() + } + + pub fn has_routing_keys(&self) -> bool { + self.routing_keys().map(|k| !k.is_empty()).unwrap_or(false) + } + + pub fn routing_keys(&self) -> Option<&Vec> { + self.routing_key_indices + .as_ref() + .or(self.key_indices.as_ref()) + } + + pub fn storage_keys(&self) -> Option<&Vec> { + self.key_indices.as_ref() + } + + pub fn sort_field_indices(&self, with_timestamp: bool) -> Vec { + let mut indices = vec![]; + if let Some(keys) = &self.key_indices { + indices.extend(keys.iter().copied()); + } + if with_timestamp { + indices.push(self.timestamp_index); + } + indices + } + + pub fn value_indices(&self, with_timestamp: bool) -> Vec { + let field_count = self.schema.fields().len(); + match &self.key_indices { + None => { + let mut indices: Vec = (0..field_count).collect(); + if !with_timestamp { + indices.remove(self.timestamp_index); + } + indices + } + Some(keys) => (0..field_count) + .filter(|index| { + !keys.contains(index) && (with_timestamp || *index != self.timestamp_index) + }) + .collect(), + } + } + + pub fn unkeyed_batch(&self, batch: &RecordBatch) -> Result { + if self.key_indices.is_none() { + return Ok(batch.clone()); + } + let columns: Vec<_> = (0..batch.num_columns()) + .filter(|index| !self.key_indices.as_ref().unwrap().contains(index)) + .collect(); + batch.project(&columns) + } + + pub fn schema_without_keys(&self) -> Result { + if self.key_indices.is_none() { + return Ok(self.clone()); + } + let key_indices = self.key_indices.as_ref().unwrap(); + let unkeyed_schema = Schema::new( + self.schema + .fields() + .iter() + .enumerate() + .filter(|(index, _)| !key_indices.contains(index)) + .map(|(_, field)| field.as_ref().clone()) + .collect::>(), + ); + let timestamp_index = unkeyed_schema.index_of(TIMESTAMP_FIELD)?; + Ok(Self { + schema: Arc::new(unkeyed_schema), + timestamp_index, + key_indices: None, + routing_key_indices: None, + }) + } + + pub fn with_fields(&self, fields: Vec) -> Result { + let schema = Arc::new(Schema::new_with_metadata( + fields, + self.schema.metadata.clone(), + )); + + let timestamp_index = schema.index_of(TIMESTAMP_FIELD)?; + let max_index = *[&self.key_indices, &self.routing_key_indices] + .iter() + .map(|indices| indices.as_ref().and_then(|k| k.iter().max())) + .max() + .flatten() + .unwrap_or(&0); + + if schema.fields.len() - 1 < max_index { + return Err(ArrowError::InvalidArgumentError(format!( + "expected at least {} fields, but were only {}", + max_index + 1, + schema.fields.len() + ))); + } + + Ok(Self { + schema, + timestamp_index, + key_indices: self.key_indices.clone(), + routing_key_indices: self.routing_key_indices.clone(), + }) + } + + pub fn with_additional_fields( + &self, + new_fields: impl Iterator, + ) -> Result { + let mut fields = self.schema.fields.to_vec(); + fields.extend(new_fields.map(Arc::new)); + self.with_fields(fields) + } +} + +/// Proto serialization: convert between FsSchema and the proto `FsSchema` message. +/// +/// Schema is encoded as JSON using Arrow's `SchemaRef` JSON representation. +/// This approach avoids depending on serde for `arrow_schema::Schema` directly. +impl FsSchema { + pub fn to_proto(&self) -> protocol::grpc::api::FsSchema { + let arrow_schema = schema_to_json_string(&self.schema); + let timestamp_index = self.timestamp_index as u32; + + let has_keys = self.key_indices.is_some(); + let key_indices = self + .key_indices + .as_ref() + .map(|ks| ks.iter().map(|i| *i as u32).collect()) + .unwrap_or_default(); + + let has_routing_keys = self.routing_key_indices.is_some(); + let routing_key_indices = self + .routing_key_indices + .as_ref() + .map(|ks| ks.iter().map(|i| *i as u32).collect()) + .unwrap_or_default(); + + protocol::grpc::api::FsSchema { + arrow_schema, + timestamp_index, + key_indices, + has_keys, + routing_key_indices, + has_routing_keys, + } + } + + pub fn from_proto(proto: protocol::grpc::api::FsSchema) -> Result { + let schema = schema_from_json_string(&proto.arrow_schema)?; + let timestamp_index = proto.timestamp_index as usize; + + let key_indices = proto + .has_keys + .then(|| proto.key_indices.into_iter().map(|i| i as usize).collect()); + + let routing_key_indices = proto.has_routing_keys.then(|| { + proto + .routing_key_indices + .into_iter() + .map(|i| i as usize) + .collect() + }); + + Ok(Self { + schema: Arc::new(schema), + timestamp_index, + key_indices, + routing_key_indices, + }) + } +} + +fn schema_to_json_string(schema: &Schema) -> String { + let json_fields: Vec = schema + .fields() + .iter() + .map(|f| { + serde_json::json!({ + "name": f.name(), + "data_type": format!("{:?}", f.data_type()), + "nullable": f.is_nullable(), + }) + }) + .collect(); + serde_json::to_string(&json_fields).unwrap() +} + +fn schema_from_json_string(s: &str) -> Result { + let json_fields: Vec = serde_json::from_str(s) + .map_err(|e| DataFusionError::Plan(format!("Invalid schema JSON: {e}")))?; + + let fields: Vec = json_fields + .into_iter() + .map(|v| { + let name = v["name"] + .as_str() + .ok_or_else(|| DataFusionError::Plan("missing field name".into()))? + .to_string(); + let nullable = v["nullable"].as_bool().unwrap_or(true); + let dt_str = v["data_type"] + .as_str() + .ok_or_else(|| DataFusionError::Plan("missing data_type".into()))?; + let data_type = parse_debug_data_type(dt_str)?; + Ok(Field::new(name, data_type, nullable)) + }) + .collect::>()?; + + Ok(Schema::new(fields)) +} + +fn parse_debug_data_type(s: &str) -> Result { + match s { + "Boolean" => Ok(DataType::Boolean), + "Int8" => Ok(DataType::Int8), + "Int16" => Ok(DataType::Int16), + "Int32" => Ok(DataType::Int32), + "Int64" => Ok(DataType::Int64), + "UInt8" => Ok(DataType::UInt8), + "UInt16" => Ok(DataType::UInt16), + "UInt32" => Ok(DataType::UInt32), + "UInt64" => Ok(DataType::UInt64), + "Float16" => Ok(DataType::Float16), + "Float32" => Ok(DataType::Float32), + "Float64" => Ok(DataType::Float64), + "Utf8" => Ok(DataType::Utf8), + "LargeUtf8" => Ok(DataType::LargeUtf8), + "Binary" => Ok(DataType::Binary), + "LargeBinary" => Ok(DataType::LargeBinary), + "Date32" => Ok(DataType::Date32), + "Date64" => Ok(DataType::Date64), + "Null" => Ok(DataType::Null), + s if s.starts_with("Timestamp(Nanosecond") => { + Ok(DataType::Timestamp(TimeUnit::Nanosecond, None)) + } + s if s.starts_with("Timestamp(Microsecond") => { + Ok(DataType::Timestamp(TimeUnit::Microsecond, None)) + } + s if s.starts_with("Timestamp(Millisecond") => { + Ok(DataType::Timestamp(TimeUnit::Millisecond, None)) + } + s if s.starts_with("Timestamp(Second") => Ok(DataType::Timestamp(TimeUnit::Second, None)), + _ => Err(DataFusionError::Plan(format!( + "Unsupported data type in schema JSON: {s}" + ))), + } +} + +impl From for FsSchema { + fn from(s: StreamSchema) -> Self { + FsSchema { + schema: s.schema, + timestamp_index: s.timestamp_index, + key_indices: s.key_indices, + routing_key_indices: None, + } + } +} + +impl From for Arc { + fn from(s: StreamSchema) -> Self { + Arc::new(FsSchema::from(s)) + } +} diff --git a/src/types/errors.rs b/src/types/errors.rs new file mode 100644 index 00000000..2c425c93 --- /dev/null +++ b/src/types/errors.rs @@ -0,0 +1,67 @@ +use std::fmt; + +/// Unified error type for streaming dataflow operations. +#[derive(Debug)] +pub enum DataflowError { + Arrow(arrow_schema::ArrowError), + DataFusion(datafusion::error::DataFusionError), + Operator(String), + State(String), + Connector(String), + Internal(String), +} + +impl fmt::Display for DataflowError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + DataflowError::Arrow(e) => write!(f, "Arrow error: {e}"), + DataflowError::DataFusion(e) => write!(f, "DataFusion error: {e}"), + DataflowError::Operator(msg) => write!(f, "Operator error: {msg}"), + DataflowError::State(msg) => write!(f, "State error: {msg}"), + DataflowError::Connector(msg) => write!(f, "Connector error: {msg}"), + DataflowError::Internal(msg) => write!(f, "Internal error: {msg}"), + } + } +} + +impl std::error::Error for DataflowError {} + +impl From for DataflowError { + fn from(e: arrow_schema::ArrowError) -> Self { + DataflowError::Arrow(e) + } +} + +impl From for DataflowError { + fn from(e: datafusion::error::DataFusionError) -> Self { + DataflowError::DataFusion(e) + } +} + +/// Macro for creating connector errors. +#[macro_export] +macro_rules! connector_err { + ($($arg:tt)*) => { + $crate::types::errors::DataflowError::Connector(format!($($arg)*)) + }; +} + +/// State-related errors. +#[derive(Debug)] +pub enum StateError { + KeyNotFound(String), + SerializationError(String), + BackendError(String), +} + +impl fmt::Display for StateError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + StateError::KeyNotFound(key) => write!(f, "Key not found: {key}"), + StateError::SerializationError(msg) => write!(f, "Serialization error: {msg}"), + StateError::BackendError(msg) => write!(f, "State backend error: {msg}"), + } + } +} + +impl std::error::Error for StateError {} diff --git a/src/types/formats.rs b/src/types/formats.rs new file mode 100644 index 00000000..25d09a74 --- /dev/null +++ b/src/types/formats.rs @@ -0,0 +1,234 @@ +use serde::{Deserialize, Serialize}; +use std::fmt::{Display, Formatter}; +use std::str::FromStr; + +#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub enum TimestampFormat { + #[default] + #[serde(rename = "rfc3339")] + RFC3339, + UnixMillis, +} + +impl TryFrom<&str> for TimestampFormat { + type Error = (); + + fn try_from(value: &str) -> Result { + match value { + "RFC3339" | "rfc3339" => Ok(TimestampFormat::RFC3339), + "UnixMillis" | "unix_millis" => Ok(TimestampFormat::UnixMillis), + _ => Err(()), + } + } +} + +#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub enum DecimalEncoding { + #[default] + Number, + String, + Bytes, +} + +impl TryFrom<&str> for DecimalEncoding { + type Error = (); + + fn try_from(s: &str) -> Result { + match s { + "number" => Ok(Self::Number), + "string" => Ok(Self::String), + "bytes" => Ok(Self::Bytes), + _ => Err(()), + } + } +} + +#[derive(Serialize, Deserialize, Default, Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub enum JsonCompression { + #[default] + Uncompressed, + Gzip, +} + +impl FromStr for JsonCompression { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + "uncompressed" => Ok(JsonCompression::Uncompressed), + "gzip" => Ok(JsonCompression::Gzip), + _ => Err(format!("invalid json compression '{s}'")), + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Default, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct JsonFormat { + #[serde(default)] + pub confluent_schema_registry: bool, + #[serde(default, alias = "confluent_schema_version")] + pub schema_id: Option, + #[serde(default)] + pub include_schema: bool, + #[serde(default)] + pub debezium: bool, + #[serde(default)] + pub unstructured: bool, + #[serde(default)] + pub timestamp_format: TimestampFormat, + #[serde(default)] + pub decimal_encoding: DecimalEncoding, + #[serde(default)] + pub compression: JsonCompression, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct RawStringFormat {} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct RawBytesFormat {} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct AvroFormat { + #[serde(default)] + pub confluent_schema_registry: bool, + #[serde(default)] + pub raw_datums: bool, + #[serde(default)] + pub into_unstructured_json: bool, + #[serde(default)] + pub schema_id: Option, +} + +impl AvroFormat { + pub fn new( + confluent_schema_registry: bool, + raw_datums: bool, + into_unstructured_json: bool, + ) -> Self { + Self { + confluent_schema_registry, + raw_datums, + into_unstructured_json, + schema_id: None, + } + } +} + +#[derive(Serialize, Deserialize, Copy, Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Default)] +#[serde(rename_all = "snake_case")] +pub enum ParquetCompression { + Uncompressed, + Snappy, + Gzip, + #[default] + Zstd, + Lz4, + Lz4Raw, +} + +impl FromStr for ParquetCompression { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + "uncompressed" => Ok(ParquetCompression::Uncompressed), + "snappy" => Ok(ParquetCompression::Snappy), + "gzip" => Ok(ParquetCompression::Gzip), + "zstd" => Ok(ParquetCompression::Zstd), + "lz4" => Ok(ParquetCompression::Lz4), + "lz4_raw" => Ok(ParquetCompression::Lz4Raw), + _ => Err(format!("invalid parquet compression '{s}'")), + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd, Default)] +#[serde(rename_all = "snake_case")] +pub struct ParquetFormat { + #[serde(default)] + pub compression: ParquetCompression, + #[serde(default)] + pub row_group_bytes: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct ProtobufFormat { + #[serde(default)] + pub into_unstructured_json: bool, + #[serde(default)] + pub message_name: Option, + #[serde(default)] + pub compiled_schema: Option>, + #[serde(default)] + pub confluent_schema_registry: bool, + #[serde(default)] + pub length_delimited: bool, +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case", tag = "type")] +pub enum Format { + Json(JsonFormat), + Avro(AvroFormat), + Protobuf(ProtobufFormat), + Parquet(ParquetFormat), + RawString(RawStringFormat), + RawBytes(RawBytesFormat), +} + +impl Display for Format { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str(self.name()) + } +} + +impl Format { + pub fn name(&self) -> &'static str { + match self { + Format::Json(_) => "json", + Format::Avro(_) => "avro", + Format::Protobuf(_) => "protobuf", + Format::Parquet(_) => "parquet", + Format::RawString(_) => "raw_string", + Format::RawBytes(_) => "raw_bytes", + } + } + + pub fn is_updating(&self) -> bool { + matches!(self, Format::Json(JsonFormat { debezium: true, .. })) + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case", tag = "behavior")] +pub enum BadData { + Fail {}, + Drop {}, +} + +impl Default for BadData { + fn default() -> Self { + BadData::Fail {} + } +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case", tag = "method")] +pub enum Framing { + Newline(NewlineDelimitedFraming), +} + +#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq, Hash, PartialOrd)] +#[serde(rename_all = "snake_case")] +pub struct NewlineDelimitedFraming { + pub max_line_length: Option, +} diff --git a/src/types/hash.rs b/src/types/hash.rs new file mode 100644 index 00000000..8f47a8fa --- /dev/null +++ b/src/types/hash.rs @@ -0,0 +1,88 @@ +use std::ops::RangeInclusive; + +/// Randomly generated seeds for consistent hashing. Changing these breaks existing state. +pub const HASH_SEEDS: [u64; 4] = [ + 5093852630788334730, + 1843948808084437226, + 8049205638242432149, + 17942305062735447798, +]; + +/// Returns the server index (0-based) responsible for the given hash value +/// when distributing across `n` servers. +pub fn server_for_hash(x: u64, n: usize) -> usize { + if n == 1 { + 0 + } else { + let range_size = (u64::MAX / (n as u64)) + 1; + (x / range_size) as usize + } +} + +/// Returns the key range assigned to server `i` out of `n` total servers. +pub fn range_for_server(i: usize, n: usize) -> RangeInclusive { + if n == 1 { + return 0..=u64::MAX; + } + let range_size = (u64::MAX / (n as u64)) + 1; + let start = range_size * (i as u64); + let end = if i + 1 == n { + u64::MAX + } else { + start + range_size - 1 + }; + start..=end +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_range_for_server() { + let n = 6; + + for i in 0..(n - 1) { + let range1 = range_for_server(i, n); + let range2 = range_for_server(i + 1, n); + + assert_eq!(*range1.end() + 1, *range2.start(), "Ranges not adjacent"); + assert_eq!( + i, + server_for_hash(*range1.start(), n), + "start not assigned to range" + ); + assert_eq!( + i, + server_for_hash(*range1.end(), n), + "end not assigned to range" + ); + } + + let last_range = range_for_server(n - 1, n); + assert_eq!( + *last_range.end(), + u64::MAX, + "Last range does not contain u64::MAX" + ); + assert_eq!( + n - 1, + server_for_hash(u64::MAX, n), + "u64::MAX not in last range" + ); + } + + #[test] + fn test_server_for_hash() { + let n = 2; + let x = u64::MAX; + + let server_index = server_for_hash(x, n); + let server_range = range_for_server(server_index, n); + + assert!( + server_range.contains(&x), + "u64::MAX is not in the correct range" + ); + } +} diff --git a/src/types/message.rs b/src/types/message.rs new file mode 100644 index 00000000..29b7f3a5 --- /dev/null +++ b/src/types/message.rs @@ -0,0 +1,42 @@ +use bincode::{Decode, Encode}; +use datafusion::arrow::array::RecordBatch; +use serde::{Deserialize, Serialize}; +use std::time::SystemTime; + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Encode, Decode, Serialize, Deserialize)] +pub enum Watermark { + EventTime(SystemTime), + Idle, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum ArrowMessage { + Data(RecordBatch), + Signal(SignalMessage), +} + +impl ArrowMessage { + pub fn is_end(&self) -> bool { + matches!( + self, + ArrowMessage::Signal(SignalMessage::Stop) + | ArrowMessage::Signal(SignalMessage::EndOfData) + ) + } +} + +#[derive(Debug, Clone, PartialEq, Encode, Decode)] +pub enum SignalMessage { + Barrier(CheckpointBarrier), + Watermark(Watermark), + Stop, + EndOfData, +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq, Encode, Decode, Serialize, Deserialize)] +pub struct CheckpointBarrier { + pub epoch: u32, + pub min_epoch: u32, + pub timestamp: SystemTime, + pub then_stop: bool, +} diff --git a/src/types/mod.rs b/src/types/mod.rs new file mode 100644 index 00000000..ddf7baca --- /dev/null +++ b/src/types/mod.rs @@ -0,0 +1,71 @@ +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Core types shared across the FunctionStream system. +//! +//! This module provides fundamental types used by the runtime, SQL planner, +//! coordinator, and other subsystems — analogous to `arroyo-types` + `arroyo-rpc` in Arroyo. + +pub mod arrow_ext; +pub mod control; +pub mod date; +pub mod debezium; +pub mod df; +pub mod errors; +pub mod formats; +pub mod hash; +pub mod message; +pub mod operator_config; +pub mod task_info; +pub mod time_utils; +pub mod worker; + +// ── Re-exports from existing modules ── +pub use arrow_ext::{DisplayAsSql, FsExtensionType, GetArrowSchema, GetArrowType}; +pub use date::{DatePart, DateTruncPrecision}; +pub use debezium::{Debezium, DebeziumOp, UpdatingData}; +pub use hash::{HASH_SEEDS, range_for_server, server_for_hash}; +pub use message::{ArrowMessage, CheckpointBarrier, SignalMessage, Watermark}; +pub use task_info::{ChainInfo, TaskInfo}; +pub use time_utils::{from_micros, from_millis, from_nanos, to_micros, to_millis, to_nanos}; +pub use worker::{MachineId, WorkerId}; + +// ── Re-exports from new modules ── +pub use control::{ + CheckpointCompleted, CheckpointEvent, CompactionResult, ControlMessage, ControlResp, + ErrorDomain, RetryHint, StopMode, TaskCheckpointEventType, TaskError, +}; +pub use df::{FsSchema, FsSchemaRef}; +pub use errors::DataflowError; +pub use formats::{BadData, Format, Framing, JsonFormat}; +pub use operator_config::{MetadataField, OperatorConfig, RateLimit}; + +// ── Well-known column names ── +pub const TIMESTAMP_FIELD: &str = "_timestamp"; +pub const UPDATING_META_FIELD: &str = "_updating_meta"; + +// ── Environment variables ── +pub const JOB_ID_ENV: &str = "JOB_ID"; +pub const RUN_ID_ENV: &str = "RUN_ID"; + +// ── Metric names ── +pub const MESSAGES_RECV: &str = "fs_worker_messages_recv"; +pub const MESSAGES_SENT: &str = "fs_worker_messages_sent"; +pub const BYTES_RECV: &str = "fs_worker_bytes_recv"; +pub const BYTES_SENT: &str = "fs_worker_bytes_sent"; +pub const BATCHES_RECV: &str = "fs_worker_batches_recv"; +pub const BATCHES_SENT: &str = "fs_worker_batches_sent"; +pub const TX_QUEUE_SIZE: &str = "fs_worker_tx_queue_size"; +pub const TX_QUEUE_REM: &str = "fs_worker_tx_queue_rem"; +pub const DESERIALIZATION_ERRORS: &str = "fs_worker_deserialization_errors"; + +pub const LOOKUP_KEY_INDEX_FIELD: &str = "__lookup_key_index"; diff --git a/src/types/operator_config.rs b/src/types/operator_config.rs new file mode 100644 index 00000000..744dbd85 --- /dev/null +++ b/src/types/operator_config.rs @@ -0,0 +1,30 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use super::formats::{BadData, Format, Framing}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RateLimit { + pub messages_per_second: u32, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MetadataField { + pub field_name: String, + pub key: String, + /// JSON-encoded Arrow DataType string, e.g. `"Utf8"`, `"Int64"`. + #[serde(default)] + pub data_type: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OperatorConfig { + pub connection: Value, + pub table: Value, + pub format: Option, + pub bad_data: Option, + pub framing: Option, + pub rate_limit: Option, + #[serde(default)] + pub metadata_fields: Vec, +} diff --git a/src/types/task_info.rs b/src/types/task_info.rs new file mode 100644 index 00000000..5a31511b --- /dev/null +++ b/src/types/task_info.rs @@ -0,0 +1,80 @@ +use bincode::{Decode, Encode}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::fmt::{Display, Formatter}; +use std::ops::RangeInclusive; + +#[derive(Eq, PartialEq, Hash, Debug, Clone, Encode, Decode, Serialize, Deserialize)] +pub struct TaskInfo { + pub job_id: String, + pub node_id: u32, + pub operator_name: String, + pub operator_id: String, + pub task_index: u32, + pub parallelism: u32, + pub key_range: RangeInclusive, +} + +impl Display for TaskInfo { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!( + f, + "Task_{}-{}/{}", + self.operator_id, self.task_index, self.parallelism + ) + } +} + +impl TaskInfo { + pub fn for_test(job_id: &str, operator_id: &str) -> Self { + Self { + job_id: job_id.to_string(), + node_id: 1, + operator_name: "op".to_string(), + operator_id: operator_id.to_string(), + task_index: 0, + parallelism: 1, + key_range: 0..=u64::MAX, + } + } +} + +pub fn get_test_task_info() -> TaskInfo { + TaskInfo { + job_id: "instance-1".to_string(), + node_id: 1, + operator_name: "test-operator".to_string(), + operator_id: "test-operator-1".to_string(), + task_index: 0, + parallelism: 1, + key_range: 0..=u64::MAX, + } +} + +#[derive(Eq, PartialEq, Hash, Debug, Clone, Encode, Decode, Serialize, Deserialize)] +pub struct ChainInfo { + pub job_id: String, + pub node_id: u32, + pub description: String, + pub task_index: u32, +} + +impl Display for ChainInfo { + fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "TaskChain{}-{} ({})", + self.node_id, self.task_index, self.description + ) + } +} + +impl ChainInfo { + pub fn metric_label_map(&self) -> HashMap { + let mut labels = HashMap::new(); + labels.insert("node_id".to_string(), self.node_id.to_string()); + labels.insert("subtask_idx".to_string(), self.task_index.to_string()); + labels.insert("node_description".to_string(), self.description.to_string()); + labels + } +} diff --git a/src/types/time_utils.rs b/src/types/time_utils.rs new file mode 100644 index 00000000..2ee5a126 --- /dev/null +++ b/src/types/time_utils.rs @@ -0,0 +1,62 @@ +use std::collections::HashMap; +use std::hash::Hash; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +pub fn to_millis(time: SystemTime) -> u64 { + time.duration_since(UNIX_EPOCH).unwrap().as_millis() as u64 +} + +pub fn to_micros(time: SystemTime) -> u64 { + time.duration_since(UNIX_EPOCH).unwrap().as_micros() as u64 +} + +pub fn from_millis(ts: u64) -> SystemTime { + UNIX_EPOCH + Duration::from_millis(ts) +} + +pub fn from_micros(ts: u64) -> SystemTime { + UNIX_EPOCH + Duration::from_micros(ts) +} + +pub fn to_nanos(time: SystemTime) -> u128 { + time.duration_since(UNIX_EPOCH).unwrap().as_nanos() +} + +pub fn from_nanos(ts: u128) -> SystemTime { + UNIX_EPOCH + + Duration::from_secs((ts / 1_000_000_000) as u64) + + Duration::from_nanos((ts % 1_000_000_000) as u64) +} + +pub fn print_time(time: SystemTime) -> String { + chrono::DateTime::::from(time) + .format("%Y-%m-%d %H:%M:%S%.3f") + .to_string() +} + +/// Returns the number of days since the UNIX epoch (for Avro serialization). +pub fn days_since_epoch(time: SystemTime) -> i32 { + time.duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() + .div_euclid(86400) as i32 +} + +pub fn single_item_hash_map, K: Hash + Eq, V>(key: I, value: V) -> HashMap { + let mut map = HashMap::new(); + map.insert(key.into(), value); + map +} + +pub fn string_to_map(s: &str, pair_delimiter: char) -> Option> { + if s.trim().is_empty() { + return Some(HashMap::new()); + } + + s.split(',') + .map(|s| { + let mut kv = s.trim().split(pair_delimiter); + Some((kv.next()?.trim().to_string(), kv.next()?.trim().to_string())) + }) + .collect() +} diff --git a/src/types/worker.rs b/src/types/worker.rs new file mode 100644 index 00000000..c12163ba --- /dev/null +++ b/src/types/worker.rs @@ -0,0 +1,14 @@ +use std::fmt::{Display, Formatter}; +use std::sync::Arc; + +#[derive(Debug, Hash, Eq, PartialEq, Copy, Clone)] +pub struct WorkerId(pub u64); + +#[derive(Debug, Hash, Eq, PartialEq, Clone)] +pub struct MachineId(pub Arc); + +impl Display for MachineId { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +}