diff --git a/Cargo.lock b/Cargo.lock index a5eb81da..a859edd8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -17,12 +17,6 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" -[[package]] -name = "adler32" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aae1277d39aeec15cb388266ecc24b11c80469deae6067e17a1a7aa9e5c1f234" - [[package]] name = "ahash" version = "0.8.12" @@ -138,35 +132,6 @@ version = "1.0.98" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e16d2d3311acee920a9eb8d33b8cbc1787ce4a264e85f964c2404b969bdcd487" -[[package]] -name = "apache-avro" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1aef82843a0ec9f8b19567445ad2421ceeb1d711514384bdd3d49fe37102ee13" -dependencies = [ - "bigdecimal", - "bzip2 0.4.4", - "crc32fast", - "digest", - "libflate", - "log", - "num-bigint", - "quad-rand", - "rand 0.8.5", - "regex-lite", - "serde", - "serde_bytes", - "serde_json", - "snap", - "strum", - "strum_macros", - "thiserror 1.0.69", - "typed-builder", - "uuid", - "xz2", - "zstd", -] - [[package]] name = "arrayref" version = "0.3.9" @@ -194,7 +159,6 @@ dependencies = [ "arrow-ipc", "arrow-json", "arrow-ord", - "arrow-pyarrow", "arrow-row", "arrow-schema", "arrow-select", @@ -330,7 +294,6 @@ dependencies = [ "arrow-schema", "flatbuffers", "lz4_flex", - "zstd", ] [[package]] @@ -368,18 +331,6 @@ dependencies = [ "arrow-select", ] -[[package]] -name = "arrow-pyarrow" -version = "55.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0e55ecf16b9b61d433f6e63c72fc6afcf2597d7db96583de88ebb887d1822268" -dependencies = [ - "arrow-array", - "arrow-data", - "arrow-schema", - "pyo3", -] - [[package]] name = "arrow-row" version = "55.2.0" @@ -399,7 +350,6 @@ version = "55.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "af7686986a3bf2254c9fb130c623cdcb2f8e1f15763e7c71c310f0834da3d292" dependencies = [ - "bitflags", "serde", "serde_json", ] @@ -453,7 +403,7 @@ version = "0.4.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06575e6a9673580f52661c92107baabffbf41e2141373441cbcdc47cb733003c" dependencies = [ - "bzip2 0.5.2", + "bzip2", "flate2", "futures-core", "memchr", @@ -648,7 +598,6 @@ dependencies = [ "num-bigint", "num-integer", "num-traits", - "serde", ] [[package]] @@ -750,16 +699,6 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d71b6127be86fdcfddb610f7182ac57211d4b18a3e9c82eb2d17662f2227ad6a" -[[package]] -name = "bzip2" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bdb116a6ef3f6c3698828873ad02c3014b3c85cadb88496095628e3ef1e347f8" -dependencies = [ - "bzip2-sys", - "libc", -] - [[package]] name = "bzip2" version = "0.5.2" @@ -769,15 +708,6 @@ dependencies = [ "bzip2-sys", ] -[[package]] -name = "bzip2" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bea8dcd42434048e4f7a304411d9273a411f647446c1234a65ce0554923f4cff" -dependencies = [ - "libbz2-rs-sys", -] - [[package]] name = "bzip2-sys" version = "0.1.13+1.0.8" @@ -790,9 +720,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.29" +version = "1.2.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c1599538de2394445747c8cf7935946e3cc27e9625f889d979bfb2aaf569362" +checksum = "d487aa071b5f64da6f19a3e848e3578944b726ee5a4854b82172f02aa876bfdc" dependencies = [ "jobserver", "libc", @@ -991,15 +921,6 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" -[[package]] -name = "core2" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b49ba7ef1ad6107f8824dbe97de947cbaac53c44e7f9756a1fba0d37c1eec505" -dependencies = [ - "memchr", -] - [[package]] name = "cpufeatures" version = "0.2.17" @@ -1096,12 +1017,6 @@ dependencies = [ "syn", ] -[[package]] -name = "dary_heap" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "04d2cd9c18b9f454ed67da600630b021a8a80bf33f8c95896ab33aaf1c26b728" - [[package]] name = "dashmap" version = "6.1.0" @@ -1119,21 +1034,21 @@ dependencies = [ [[package]] name = "datafusion" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc6cb8c2c81eada072059983657d6c9caf3fddefc43b4a65551d243253254a96" dependencies = [ "arrow", "arrow-ipc", "arrow-schema", "async-trait", "bytes", - "bzip2 0.6.0", + "bzip2", "chrono", "datafusion-catalog", "datafusion-catalog-listing", "datafusion-common", "datafusion-common-runtime", "datafusion-datasource", - "datafusion-datasource-avro", "datafusion-datasource-csv", "datafusion-datasource-json", "datafusion-datasource-parquet", @@ -1173,7 +1088,8 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7be8d1b627843af62e447396db08fe1372d882c0eb8d0ea655fd1fbc33120ee" dependencies = [ "arrow", "async-trait", @@ -1198,7 +1114,8 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38ab16c5ae43f65ee525fc493ceffbc41f40dee38b01f643dfcfc12959e92038" dependencies = [ "arrow", "async-trait", @@ -1220,24 +1137,21 @@ dependencies = [ [[package]] name = "datafusion-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3d56b2ac9f476b93ca82e4ef5fb00769c8a3f248d12b4965af7e27635fa7e12" dependencies = [ "ahash", - "apache-avro", "arrow", "arrow-ipc", "base64 0.22.1", - "chrono", "half", "hashbrown 0.14.5", - "hex", "indexmap 2.10.0", "libc", "log", "object_store", "parquet", "paste", - "pyo3", "recursive", "sqlparser", "tokio", @@ -1247,7 +1161,8 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16015071202d6133bc84d72756176467e3e46029f3ce9ad2cb788f9b1ff139b2" dependencies = [ "futures", "log", @@ -1257,13 +1172,14 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b77523c95c89d2a7eb99df14ed31390e04ab29b43ff793e562bdc1716b07e17b" dependencies = [ "arrow", "async-compression", "async-trait", "bytes", - "bzip2 0.6.0", + "bzip2", "chrono", "datafusion-common", "datafusion-common-runtime", @@ -1289,34 +1205,11 @@ dependencies = [ "zstd", ] -[[package]] -name = "datafusion-datasource-avro" -version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" -dependencies = [ - "apache-avro", - "arrow", - "async-trait", - "bytes", - "chrono", - "datafusion-catalog", - "datafusion-common", - "datafusion-datasource", - "datafusion-execution", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "datafusion-physical-plan", - "datafusion-session", - "futures", - "num-traits", - "object_store", - "tokio", -] - [[package]] name = "datafusion-datasource-csv" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40d25c5e2c0ebe8434beeea997b8e88d55b3ccc0d19344293f2373f65bc524fc" dependencies = [ "arrow", "async-trait", @@ -1340,7 +1233,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dc6959e1155741ab35369e1dc7673ba30fc45ed568fad34c01b7cb1daeb4d4c" dependencies = [ "arrow", "async-trait", @@ -1364,7 +1258,8 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7a6afdfe358d70f4237f60eaef26ae5a1ce7cb2c469d02d5fc6c7fd5d84e58b" dependencies = [ "arrow", "async-trait", @@ -1380,7 +1275,6 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-optimizer", "datafusion-physical-plan", - "datafusion-pruning", "datafusion-session", "futures", "itertools 0.14.0", @@ -1395,12 +1289,14 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bcd8a3e3e3d02ea642541be23d44376b5d5c37c2938cce39b3873cdf7186eea" [[package]] name = "datafusion-execution" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670da1d45d045eee4c2319b8c7ea57b26cf48ab77b630aaa50b779e406da476a" dependencies = [ "arrow", "dashmap", @@ -1418,10 +1314,10 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3a577f64bdb7e2cc4043cd97f8901d8c504711fde2dbcb0887645b00d7c660b" dependencies = [ "arrow", - "async-trait", "chrono", "datafusion-common", "datafusion-doc", @@ -1439,7 +1335,8 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51b7916806ace3e9f41884f230f7f38ebf0e955dfbd88266da1826f29a0b9a6a" dependencies = [ "arrow", "datafusion-common", @@ -1451,7 +1348,8 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fb31c9dc73d3e0c365063f91139dc273308f8a8e124adda9898db8085d68357" dependencies = [ "arrow", "arrow-buffer", @@ -1479,7 +1377,8 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebb72c6940697eaaba9bd1f746a697a07819de952b817e3fb841fb75331ad5d4" dependencies = [ "ahash", "arrow", @@ -1499,7 +1398,8 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7fdc54656659e5ecd49bf341061f4156ab230052611f4f3609612a0da259696" dependencies = [ "ahash", "arrow", @@ -1511,7 +1411,8 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fad94598e3374938ca43bca6b675febe557e7a14eb627d617db427d70d65118b" dependencies = [ "arrow", "arrow-ord", @@ -1521,7 +1422,6 @@ dependencies = [ "datafusion-expr", "datafusion-functions", "datafusion-functions-aggregate", - "datafusion-functions-aggregate-common", "datafusion-macros", "datafusion-physical-expr-common", "itertools 0.14.0", @@ -1532,7 +1432,8 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de2fc6c2946da5cab8364fb28b5cac3115f0f3a87960b235ed031c3f7e2e639b" dependencies = [ "arrow", "async-trait", @@ -1547,7 +1448,8 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e5746548a8544870a119f556543adcd88fe0ba6b93723fe78ad0439e0fbb8b4" dependencies = [ "arrow", "datafusion-common", @@ -1564,7 +1466,8 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dcbe9404382cda257c434f22e13577bee7047031dfdb6216dd5e841b9465e6fe" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1573,7 +1476,8 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8dce50e3b637dab0d25d04d2fe79dfdca2b257eabd76790bffd22c7f90d700c8" dependencies = [ "datafusion-expr", "quote", @@ -1583,13 +1487,13 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03cfaacf06445dc3bbc1e901242d2a44f2cae99a744f49f3fefddcee46240058" dependencies = [ "arrow", "chrono", "datafusion-common", "datafusion-expr", - "datafusion-expr-common", "datafusion-physical-expr", "indexmap 2.10.0", "itertools 0.14.0", @@ -1602,7 +1506,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1908034a89d7b2630898e06863583ae4c00a0dd310c1589ca284195ee3f7f8a6" dependencies = [ "ahash", "arrow", @@ -1623,7 +1528,8 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b7a12dd59ea07614b67dbb01d85254fbd93df45bcffa63495e11d3bdf847df" dependencies = [ "ahash", "arrow", @@ -1636,7 +1542,8 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4371cc4ad33978cc2a8be93bd54a232d3f2857b50401a14631c0705f3f910aae" dependencies = [ "arrow", "datafusion-common", @@ -1646,7 +1553,6 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", - "datafusion-pruning", "itertools 0.14.0", "log", "recursive", @@ -1655,7 +1561,8 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc47bc33025757a5c11f2cd094c5b6b5ed87f46fa33c023e6fdfa25fcbfade23" dependencies = [ "ahash", "arrow", @@ -1684,7 +1591,8 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8f5d9acd7d96e3bf2a7bb04818373cab6e51de0356e3694b94905fee7b4e8b6" dependencies = [ "arrow", "chrono", @@ -1699,34 +1607,19 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ecb5ec152c4353b60f7a5635489834391f7a291d2b39a4820cd469e318b78e" dependencies = [ "arrow", "datafusion-common", "prost", ] -[[package]] -name = "datafusion-pruning" -version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" -dependencies = [ - "arrow", - "arrow-schema", - "datafusion-common", - "datafusion-datasource", - "datafusion-expr-common", - "datafusion-physical-expr", - "datafusion-physical-expr-common", - "datafusion-physical-plan", - "itertools 0.14.0", - "log", -] - [[package]] name = "datafusion-session" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7485da32283985d6b45bd7d13a65169dcbe8c869e25d01b2cfbc425254b4b49" dependencies = [ "arrow", "async-trait", @@ -1749,7 +1642,8 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#d359d6496168da59e6ac4bfea30e648674382f87" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a466b15632befddfeac68c125f0260f569ff315c6831538cbb40db754134e0df" dependencies = [ "arrow", "bigdecimal", @@ -1765,7 +1659,8 @@ dependencies = [ [[package]] name = "datafusion-substrait" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#85a92514270f8da351496c5d81445ba481cc6f05" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f3973b1a4f6e9ee7fd99a22d58e1c06e6723a28dc911a60df575974c8339aa" dependencies = [ "async-recursion", "async-trait", @@ -1848,8 +1743,6 @@ dependencies = [ "rand 0.8.5", "rustc_version", "rustls", - "serde", - "serde_json", "tempfile", "test-log", "thiserror 1.0.69", @@ -2192,9 +2085,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.11" +version = "0.4.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17da50a276f1e01e0ba6c029e47b7100754904ee8a278f886546e98575380785" +checksum = "a9421a676d1b147b16b82c9225157dc629087ef8ec4d5e2960f9437a90dac0a5" dependencies = [ "atomic-waker", "bytes", @@ -2608,29 +2501,12 @@ dependencies = [ "hashbrown 0.15.4", ] -[[package]] -name = "indoc" -version = "2.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4c7245a08504955605670dbf141fceab975f15ca21570696aebe9d2e71576bd" - [[package]] name = "integer-encoding" version = "3.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" -[[package]] -name = "io-uring" -version = "0.7.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b86e202f00093dcba4275d4636b93ef9dd75d025ae560d2521b45ea28ab49013" -dependencies = [ - "bitflags", - "cfg-if", - "libc", -] - [[package]] name = "ipnet" version = "2.11.0" @@ -2963,42 +2839,12 @@ dependencies = [ "static_assertions", ] -[[package]] -name = "libbz2-rs-sys" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "775bf80d5878ab7c2b1080b5351a48b2f737d9f6f8b383574eebcc22be0dfccb" - [[package]] name = "libc" version = "0.2.174" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776" -[[package]] -name = "libflate" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45d9dfdc14ea4ef0900c1cddbc8dcd553fbaacd8a4a282cf4018ae9dd04fb21e" -dependencies = [ - "adler32", - "core2", - "crc32fast", - "dary_heap", - "libflate_lz77", -] - -[[package]] -name = "libflate_lz77" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6e0d73b369f386f1c44abd9c570d5318f55ccde816ff4b562fa452e5182863d" -dependencies = [ - "core2", - "hashbrown 0.14.5", - "rle-decode-fast", -] - [[package]] name = "libloading" version = "0.8.8" @@ -3115,15 +2961,6 @@ version = "2.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" -[[package]] -name = "memoffset" -version = "0.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" -dependencies = [ - "autocfg", -] - [[package]] name = "mime" version = "0.3.17" @@ -3204,7 +3041,6 @@ checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" dependencies = [ "num-integer", "num-traits", - "serde", ] [[package]] @@ -3396,7 +3232,6 @@ dependencies = [ "num-bigint", "object_store", "paste", - "ring", "seq-macro", "simdutf8", "snap", @@ -3739,75 +3574,6 @@ dependencies = [ "cc", ] -[[package]] -name = "pyo3" -version = "0.24.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5203598f366b11a02b13aa20cab591229ff0a89fd121a308a5df751d5fc9219" -dependencies = [ - "cfg-if", - "indoc", - "libc", - "memoffset", - "once_cell", - "portable-atomic", - "pyo3-build-config", - "pyo3-ffi", - "pyo3-macros", - "unindent", -] - -[[package]] -name = "pyo3-build-config" -version = "0.24.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "99636d423fa2ca130fa5acde3059308006d46f98caac629418e53f7ebb1e9999" -dependencies = [ - "once_cell", - "target-lexicon", -] - -[[package]] -name = "pyo3-ffi" -version = "0.24.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78f9cf92ba9c409279bc3305b5409d90db2d2c22392d443a87df3a1adad59e33" -dependencies = [ - "libc", - "pyo3-build-config", -] - -[[package]] -name = "pyo3-macros" -version = "0.24.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b999cb1a6ce21f9a6b147dcf1be9ffedf02e0043aec74dc390f3007047cecd9" -dependencies = [ - "proc-macro2", - "pyo3-macros-backend", - "quote", - "syn", -] - -[[package]] -name = "pyo3-macros-backend" -version = "0.24.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "822ece1c7e1012745607d5cf0bcb2874769f0f7cb34c4cde03b9358eb9ef911a" -dependencies = [ - "heck", - "proc-macro2", - "pyo3-build-config", - "quote", - "syn", -] - -[[package]] -name = "quad-rand" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a651516ddc9168ebd67b24afd085a718be02f8858fe406591b013d101ce2f40" - [[package]] name = "quick-xml" version = "0.37.5" @@ -4008,12 +3774,6 @@ dependencies = [ "regex-syntax 0.8.5", ] -[[package]] -name = "regex-lite" -version = "0.1.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "53a49587ad06b26609c52e423de037e7f57f20d53535d66e08c695f347df952a" - [[package]] name = "regex-syntax" version = "0.6.29" @@ -4038,9 +3798,9 @@ dependencies = [ [[package]] name = "reqwest" -version = "0.12.22" +version = "0.12.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbc931937e6ca3a06e3b6c0aa7841849b160a90351d6ab467a8b9b9959767531" +checksum = "eabf4c97d9130e2bf606614eb937e86edac8292eaa6f422f995d7e8de1eb1813" dependencies = [ "base64 0.22.1", "bytes", @@ -4092,12 +3852,6 @@ dependencies = [ "windows-sys 0.52.0", ] -[[package]] -name = "rle-decode-fast" -version = "1.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3582f63211428f83597b51b2ddb88e2a91a9d52d12831f9d08f5e624e8977422" - [[package]] name = "rustc-demangle" version = "0.1.25" @@ -4362,15 +4116,6 @@ dependencies = [ "serde", ] -[[package]] -name = "serde_bytes" -version = "0.11.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8437fd221bde2d4ca316d61b90e337e9e702b3820b87d63caa9ba6c02bd06d96" -dependencies = [ - "serde", -] - [[package]] name = "serde_derive" version = "1.0.219" @@ -4581,30 +4326,11 @@ 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" - -[[package]] -name = "strum_macros" -version = "0.26.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" -dependencies = [ - "heck", - "proc-macro2", - "quote", - "rustversion", - "syn", -] - [[package]] name = "substrait" -version = "0.58.0" +version = "0.56.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "de6d24c270c6c672a86c183c3a8439ba46c1936f93cf7296aa692de3b0ff0228" +checksum = "13de2e20128f2a018dab1cfa30be83ae069219a65968c6f89df66ad124de2397" dependencies = [ "heck", "pbjson", @@ -4662,12 +4388,6 @@ dependencies = [ "syn", ] -[[package]] -name = "target-lexicon" -version = "0.13.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e502f78cdbb8ba4718f566c418c52bc729126ffd16baee5baa718cf25dd5a69a" - [[package]] name = "tempfile" version = "3.20.0" @@ -4683,9 +4403,9 @@ dependencies = [ [[package]] name = "test-log" -version = "0.2.18" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e33b98a582ea0be1168eba097538ee8dd4bbe0f2b01b22ac92ea30054e5be7b" +checksum = "e7f46083d221181166e5b6f6b1e5f1d499f3a76888826e6cb1d057554157cd0f" dependencies = [ "env_logger", "test-log-macros", @@ -4694,9 +4414,9 @@ dependencies = [ [[package]] name = "test-log-macros" -version = "0.2.18" +version = "0.2.17" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "451b374529930d7601b1eef8d32bc79ae870b6079b069401709c2a8bf9e75f36" +checksum = "888d0c3c6db53c0fdab160d2ed5e12ba745383d3e85813f2ea0f2b1475ab553f" dependencies = [ "proc-macro2", "quote", @@ -4840,19 +4560,17 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.46.1" +version = "1.45.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cc3a2344dafbe23a245241fe8b09735b521110d30fcefbbd5feb1797ca35d17" +checksum = "75ef51a33ef1da925cea3e4eb122833cb377c61439ca401b770f54902b806779" dependencies = [ "backtrace", "bytes", - "io-uring", "libc", "mio", "parking_lot", "pin-project-lite", "signal-hook-registry", - "slab", "socket2", "tokio-macros", "windows-sys 0.52.0", @@ -5092,26 +4810,6 @@ version = "2.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b907da542cbced5261bd3256de1b3a1bf340a3d37f93425a07362a1d687de56" -[[package]] -name = "typed-builder" -version = "0.19.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a06fbd5b8de54c5f7c91f6fe4cebb949be2125d7758e630bb58b1d831dbce600" -dependencies = [ - "typed-builder-macro", -] - -[[package]] -name = "typed-builder-macro" -version = "0.19.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9534daa9fd3ed0bd911d462a37f172228077e7abf18c18a5f67199d959205f8" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "typenum" version = "1.18.0" @@ -5189,12 +4887,6 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a1a07cc7db3810833284e8d372ccdc6da29741639ecc70c9ec107df0fa6154c" -[[package]] -name = "unindent" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" - [[package]] name = "unsafe-libyaml" version = "0.2.11" @@ -5238,7 +4930,6 @@ checksum = "3cf4199d1e5d15ddd86a694e4d0dffa9c323ce759fea589f00fef9d81cc1931d" dependencies = [ "getrandom 0.3.3", "js-sys", - "serde", "wasm-bindgen", ] diff --git a/Cargo.toml b/Cargo.toml index b2b00b2e..6f815062 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,12 +39,9 @@ arrow-flight = { version = "55", features = ["flight-sql-experimental"] } async-stream = "0.3" bytes = "1.5" clap = { version = "4.4", features = ["derive"] } -datafusion = { git = "https://github.com/apache/datafusion", branch = "main", features = [ - "pyarrow", - "avro", -] } -datafusion-proto = { git = "https://github.com/apache/datafusion", branch = "main" } -datafusion-substrait = { git = "https://github.com/apache/datafusion", branch = "main" } +datafusion = "48.0.0" +datafusion-proto = "48.0.0" +datafusion-substrait = "48.0.0" env_logger = "0.11" futures = "0.3" itertools = "0.14" @@ -54,9 +51,6 @@ log = "0.4" rand = "0.8" uuid = { version = "1.6", features = ["v4"] } -serde = { version = "1.0", features = ["derive"] } -serde_json = "1.0" - object_store = { version = "0.12.0", features = [ "aws", "gcp", @@ -99,4 +93,4 @@ tonic-build = { version = "0.12", default-features = false, features = [ url = "2" [dev-dependencies] -tempfile = "3.20" \ No newline at end of file +tempfile = "3.20" diff --git a/build.rs b/build.rs index a82e403f..87dc74fd 100644 --- a/build.rs +++ b/build.rs @@ -32,16 +32,16 @@ fn main() -> Result<(), String> { // We don't include the proto files in releases so that downstreams // do not need to have PROTOC included - if Path::new("src/proto/datafusion_ray.proto").exists() { + if Path::new("src/proto/distributed_datafusion.proto").exists() { println!("cargo:rerun-if-changed=src/proto/datafusion_common.proto"); println!("cargo:rerun-if-changed=src/proto/datafusion.proto"); - println!("cargo:rerun-if-changed=src/proto/datafusion_ray.proto"); + println!("cargo:rerun-if-changed=src/proto/distributed_datafusion.proto"); tonic_build::configure() .extern_path(".datafusion", "::datafusion_proto::protobuf") .extern_path(".datafusion_common", "::datafusion_proto::protobuf") - .compile_protos(&["src/proto/datafusion_ray.proto"], &["src/proto"]) + .compile_protos(&["src/proto/distributed_datafusion.proto"], &["src/proto"]) .map_err(|e| format!("protobuf compilation failed: {e}"))?; - let generated_source_path = out.join("datafusion_ray.protobuf.rs"); + let generated_source_path = out.join("distributed_datafusion.protobuf.rs"); let code = std::fs::read_to_string(generated_source_path).unwrap(); let mut file = std::fs::OpenOptions::new() .write(true) diff --git a/scripts/launch_python_arrowflightsql_client.sh b/scripts/launch_python_arrowflightsql_client.sh index c53710cb..44776e48 100755 --- a/scripts/launch_python_arrowflightsql_client.sh +++ b/scripts/launch_python_arrowflightsql_client.sh @@ -63,15 +63,15 @@ DEFAULT_QUERY_PATH="./tpch/queries/" # Parse command line arguments for arg in "$@"; do case $arg in - query_path=*) - TPCH_QUERY_PATH="${arg#*=}" - ;; - *) - echo "Usage: $0 [query_path=/path/to/queries/]" - echo "Example: $0 query_path=./tpch/queries/" - echo "If no argument is provided, default path will be used: $DEFAULT_QUERY_PATH" - exit 1 - ;; + query_path=*) + TPCH_QUERY_PATH="${arg#*=}" + ;; + *) + echo "Usage: $0 [query_path=/path/to/queries/]" + echo "Example: $0 query_path=./tpch/queries/" + echo "If no argument is provided, default path will be used: $DEFAULT_QUERY_PATH" + exit 1 + ;; esac done @@ -103,15 +103,16 @@ fi source .venv/bin/activate # Install required packages if not already installed -if ! pip show adbc_driver_flightsql > /dev/null 2>&1 || ! pip show duckdb > /dev/null 2>&1; then +if ! pip show adbc_driver_flightsql >/dev/null 2>&1 || ! pip show rich >/dev/null 2>&1; then echo "Installing required packages..." - pip install adbc_driver_manager adbc_driver_flightsql duckdb pyarrow + pip install adbc_driver_manager adbc_driver_flightsql rich pyarrow fi # Create a Python startup script -cat > .python_startup.py << 'EOF' +cat >.python_startup.py <<'EOF' import adbc_driver_flightsql.dbapi as dbapi -import duckdb +from rich.console import Console +from rich.table import Table import os import sys @@ -170,53 +171,25 @@ def run_sql(sql_query): """ try: cur.execute(sql_query) - reader = cur.fetch_record_batch() + table = cur.fetch_arrow_table() + + rich_table = Table(show_header=True, header_style="bold magenta") + + # Add columns based on the PyArrow Table schema + for field in table.schema: + rich_table.add_column(field.name) + + # Add rows from the PyArrow Table + for row_index in range(table.num_rows): + row_data = [str(table.column(col_index)[row_index].as_py()) for col_index in range(table.num_columns)] + rich_table.add_row(*row_data) + + console = Console() + console.print(rich_table, markup=False) - # Use basic DuckDB show() - for full output use run_sql_full() or run_sql_raw() - # duckdb.sql("select * from reader").show(max_width=10000) - duckdb.sql("select * from reader").show() except Exception as e: print(f"Error executing SQL query: {str(e)}") -def format_plan(plan_text): - """ - Format the plan text by replacing \n with proper indentation - """ - if not plan_text: - return "" - - # If the plan_text doesn't contain \\n, return as is (it's likely already formatted) - if '\\n' not in plan_text: - return plan_text - - # Split the plan into lines and add proper indentation - lines = plan_text.split('\\n') - formatted_lines = [] - indent_level = 0 - indent_size = 2 - - for line in lines: - # Skip empty lines - if not line.strip(): - continue - - # Count opening and closing parentheses to adjust indentation - open_parens = line.count('(') - close_parens = line.count(')') - - # Adjust indent level based on parentheses - if close_parens > open_parens: - indent_level = max(0, indent_level - (close_parens - open_parens)) - - # Add the line with current indentation - formatted_lines.append(' ' * (indent_level * indent_size) + line.strip()) - - # Update indent level for next line - if open_parens > close_parens: - indent_level += (open_parens - close_parens) - - return '\n'.join(formatted_lines) - def explain_query(query_name): """ Run EXPLAIN for a TPC-H query by name (e.g., 'q5' for EXPLAIN q5.sql) @@ -225,64 +198,12 @@ def explain_query(query_name): query_file = os.path.join(tpch_query_path, f"{query_name}.sql") try: with open(query_file, 'r') as f: - sql = f.read() - print(f"Executing EXPLAIN for query from {query_file}...") - explain_sql(sql) + sql = "explain " + f.read() + run_sql(sql) except FileNotFoundError: print(f"Error: Query file {query_file} not found") except Exception as e: - print(f"Error executing EXPLAIN: {str(e)}") - import traceback - traceback.print_exc() - -def explain_sql(sql_query): - """ - Run EXPLAIN on a given SQL query (passed as a string or variable) and display the formatted output. - """ - try: - print("Executing EXPLAIN...") - cur.execute(f"EXPLAIN {sql_query}") - results = cur.fetchall() - if not results: - print("No explain plan returned") - return - - logical_plan = None - physical_plan = None - distributed_plan = None - distributed_stages = None - for row in results: - if row[0] == 'logical_plan': - logical_plan = row[1] - elif row[0] == 'physical_plan': - physical_plan = row[1] - elif row[0] == 'distributed_plan': - distributed_plan = row[1] - elif row[0] == 'distributed_stages': - distributed_stages = row[1] - formatted_logical = format_plan(logical_plan) if logical_plan else "Logical plan not available" - formatted_physical = format_plan(physical_plan) if physical_plan else "Physical plan not available" - formatted_distributed = format_plan(distributed_plan) if distributed_plan else "Distributed plan not available" - formatted_distributed_stages = format_plan(distributed_stages) if distributed_stages else "Distributed stages not available" - print("\nExecution Plan:") - print("=" * 100) - print("Logical Plan:") - print("-" * 100) - print(formatted_logical) - print("\nPhysical Plan:") - print("-" * 100) - print(formatted_physical) - if distributed_plan: - print("\nDistributed Plan:") - print("-" * 100) - print(formatted_distributed) - if distributed_stages: - print("\nDistributed Stages:") - print("-" * 100) - print(formatted_distributed_stages) - print("=" * 100) - except Exception as e: - print(f"Error executing EXPLAIN: {e}") + print(f"Error executing query: {str(e)}") def explain_analyze_query(query_name): """ @@ -292,9 +213,8 @@ def explain_analyze_query(query_name): query_file = os.path.join(tpch_query_path, f"{query_name}.sql") try: with open(query_file, 'r') as f: - sql = f.read() - print(f"Executing EXPLAIN ANALYZE for query from {query_file}...") - explain_analyze_sql(sql) + sql = "explain analyze " + f.read() + run_sql(sql) except FileNotFoundError: print(f"Error: Query file {query_file} not found") except Exception as e: @@ -302,37 +222,6 @@ def explain_analyze_query(query_name): import traceback traceback.print_exc() -def explain_analyze_sql(sql_query): - """ - Run EXPLAIN ANALYZE on a given SQL query (passed as a string or variable) and display the formatted output with execution statistics. - """ - try: - print("Executing EXPLAIN ANALYZE...") - cur.execute(f"EXPLAIN ANALYZE {sql_query}") - results = cur.fetchall() - if not results: - print("No explain analyze plan returned") - return - - plan_with_metrics = None - - for row in results: - if row[0] == 'Plan with Metrics': - plan_with_metrics = row[1] - - if plan_with_metrics: - # EXPLAIN ANALYZE returns execution statistics - formatted_plan_with_metrics = format_plan(plan_with_metrics) - print("\nExecution Plan with Analysis:") - print("=" * 100) - print("Physical Plan with Execution Statistics:") - print("-" * 100) - print(formatted_plan_with_metrics) - print("=" * 100) - except Exception as e: - print(f"Error executing EXPLAIN ANALYZE: {e}") - -# Add the run_query function to the global namespace sys.ps1 = ">>> " print("\nWelcome to the TPC-H Query Client!") print("Available commands:") @@ -341,11 +230,7 @@ print(" show_query('q5') # Show SQL content of query 5") print(" run_query('q5') # Run TPC-H query 5") print(" explain_query('q5') # Show EXPLAIN plan for TPC-H query 5") print(" explain_analyze_query('q5') # Show EXPLAIN ANALYZE plan for TPC-H query 5") -print(" # Note: EXPLAIN ANALYZE does not work with distributed plans yet") print(" run_sql('select * from nation') # Run SQL query") -print(" explain_sql('select * from nation') # Show EXPLAIN plan for SQL query") -print(" explain_analyze_sql('select * from nation') # Show EXPLAIN ANALYZE plan for SQL query") -print(" # etc...") print("\nConnected to database at grpc://localhost:20200") print("Type 'exit()' to quit\n") EOF @@ -354,4 +239,4 @@ EOF export TPCH_QUERY_PATH=$TPCH_QUERY_PATH # Start Python with the startup script -python3 -i .python_startup.py \ No newline at end of file +python3 -i .python_startup.py diff --git a/scripts/launch_tpch_cluster.sh b/scripts/launch_tpch_cluster.sh index c56e3f73..6cdf0454 100755 --- a/scripts/launch_tpch_cluster.sh +++ b/scripts/launch_tpch_cluster.sh @@ -38,20 +38,20 @@ DEFAULT_LOG_PATH="." # Parse named arguments for arg in "$@"; do case "$arg" in - num_workers=*) - NUM_WORKERS="${arg#*=}" - ;; - tpch_file_path=*) - TPCH_DATA_DIR="${arg#*=}" - ;; - log_file_path=*) - LOG_DIR="${arg#*=}" - ;; - *) - echo "Error: Unknown argument '$arg'" - echo "Usage: $0 [num_workers=N] [tpch_file_path=PATH] [log_file_path=PATH]" - exit 1 - ;; + num_workers=*) + NUM_WORKERS="${arg#*=}" + ;; + tpch_file_path=*) + TPCH_DATA_DIR="${arg#*=}" + ;; + log_file_path=*) + LOG_DIR="${arg#*=}" + ;; + *) + echo "Error: Unknown argument '$arg'" + echo "Usage: $0 [num_workers=N] [tpch_file_path=PATH] [log_file_path=PATH]" + exit 1 + ;; esac done @@ -71,13 +71,14 @@ if [ ! -f "./target/release/distributed-datafusion" ]; then echo "Binary not found, building release version..." echo "This may take a few minutes on first run..." if [ -f "./build.sh" ]; then - ./build.sh --release + ./build.sh #--release else - cargo build --release + cargo build #--release fi - + # Verify the build was successful - if [ ! -f "./target/release/distributed-datafusion" ]; then + #if [ ! -f "./target/release/distributed-datafusion" ]; then + if [ ! -f "./target/debug/distributed-datafusion" ]; then echo "Error: Failed to build distributed-datafusion binary" exit 1 fi @@ -98,8 +99,8 @@ if [ ! -d "$LOG_DIR" ]; then fi # Verify required parquet files exist -required_files=("customer.parquet" "lineitem.parquet" "nation.parquet" "orders.parquet" - "part.parquet" "partsupp.parquet" "region.parquet" "supplier.parquet") +required_files=("customer.parquet" "lineitem.parquet" "nation.parquet" "orders.parquet" + "part.parquet" "partsupp.parquet" "region.parquet" "supplier.parquet") for file in "${required_files[@]}"; do if [ ! -f "${TPCH_DATA_DIR}/${file}" ]; then @@ -125,9 +126,6 @@ partsupp:parquet:${TPCH_DATA_DIR}/partsupp.parquet,\ region:parquet:${TPCH_DATA_DIR}/region.parquet,\ supplier:parquet:${TPCH_DATA_DIR}/supplier.parquet" -# Define views required for TPC-H queries (e.g., q15) -export DFRAY_VIEWS="create view revenue0 (supplier_no, total_revenue) as select l_suppkey, sum(l_extendedprice * (1 - l_discount)) from lineitem where l_shipdate >= date '1996-08-01' and l_shipdate < date '1996-08-01' + interval '3' month group by l_suppkey" - # Array to store worker PIDs and addresses declare -a WORKER_PIDS declare -a WORKER_ADDRESSES @@ -144,28 +142,33 @@ trap cleanup SIGINT SIGTERM EXIT # Start workers echo "Starting workers..." -for ((i=0; i "$LOG_FILE" 2>&1 & + #env DATAFUSION_RAY_LOG_LEVEL="$DATAFUSION_RAY_LOG_LEVEL" DFRAY_TABLES="$DFRAY_TABLES" ./target/release/distributed-datafusion --mode worker --port $PORT >"$LOG_FILE" 2>&1 & + env RUST_BACKTRACE=1 DATAFUSION_RAY_LOG_LEVEL="$DATAFUSION_RAY_LOG_LEVEL" DFRAY_TABLES="$DFRAY_TABLES" ./target/debug/distributed-datafusion --mode worker --port $PORT >"$LOG_FILE" 2>&1 & WORKER_PIDS[$i]=$! - WORKER_ADDRESSES[$i]="${WORKER_NAME}/localhost:${PORT}" + WORKER_ADDRESSES[$i]="localhost:${PORT}" done # Give workers time to start echo "Waiting for workers to initialize..." -sleep 2 +sleep 5 # Construct worker addresses string for proxy -WORKER_ADDRESSES_STR=$(IFS=,; echo "${WORKER_ADDRESSES[*]}") +WORKER_ADDRESSES_STR=$( + IFS=, + echo "${WORKER_ADDRESSES[*]}" +) # Start proxy echo "Starting proxy on port 20200..." echo "Connecting to workers: $WORKER_ADDRESSES_STR" PROXY_LOG="${LOG_DIR}/proxy.log" -env DATAFUSION_RAY_LOG_LEVEL="$DATAFUSION_RAY_LOG_LEVEL" DFRAY_TABLES="$DFRAY_TABLES" DFRAY_VIEWS="$DFRAY_VIEWS" DFRAY_WORKER_ADDRESSES="$WORKER_ADDRESSES_STR" ./target/release/distributed-datafusion --mode proxy --port 20200 > "$PROXY_LOG" 2>&1 & +#env DATAFUSION_RAY_LOG_LEVEL="$DATAFUSION_RAY_LOG_LEVEL" DFRAY_TABLES="$DFRAY_TABLES" DFRAY_WORKER_ADDRESSES="$WORKER_ADDRESSES_STR" ./target/release/distributed-datafusion --mode proxy --port 20200 >"$PROXY_LOG" 2>&1 & +env RUST_BACKTRACE=1 DATAFUSION_RAY_LOG_LEVEL="$DATAFUSION_RAY_LOG_LEVEL" DFRAY_TABLES="$DFRAY_TABLES" DFRAY_WORKER_ADDRESSES="$WORKER_ADDRESSES_STR" ./target/debug/distributed-datafusion --mode proxy --port 20200 >"$PROXY_LOG" 2>&1 & PROXY_PID=$! echo @@ -173,23 +176,23 @@ echo "TPC-H cluster is now running!" echo echo "Cluster Information:" echo " Proxy: localhost:20200" -for ((i=0; i, +} + +impl DistributedAnalyzeExec { + pub fn new(input: Arc, verbose: bool, show_statistics: bool) -> Self { + Self { + input, + verbose, + show_statistics, + } + } + + pub fn annotated_plan(&self) -> String { + DisplayableExecutionPlan::with_metrics(self.input.as_ref()) + .set_show_statistics(self.show_statistics) + .indent(self.verbose) + .to_string() + } +} + +impl DisplayAs for DistributedAnalyzeExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "DistributedAnalyzeExec[verbose = {}, show stats = {}]", + self.verbose, self.show_statistics + ) + } +} + +impl ExecutionPlan for DistributedAnalyzeExec { + fn name(&self) -> &str { + "DistributedAnalyzeExec" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn properties(&self) -> &PlanProperties { + self.input.properties() + } + + fn children(&self) -> Vec<&std::sync::Arc> { + vec![&self.input] + } + + fn with_new_children( + self: std::sync::Arc, + children: Vec>, + ) -> Result> { + // TODO: generalize this + assert_eq!(children.len(), 1); + Ok(Arc::new(Self::new( + children[0].clone(), + self.verbose, + self.show_statistics, + ))) + } + + fn execute( + &self, + partition: usize, + context: std::sync::Arc, + ) -> Result { + self.input.execute(partition, context) + } +} + +#[derive(Debug)] +pub struct DistributedAnalyzeRootExec { + /// Control how much extra to print + pub(crate) verbose: bool, + /// If statistics should be displayed + pub(crate) show_statistics: bool, + /// The input plan (the plan being analyzed) + pub(crate) input: Arc, + /// our plan properties + properties: PlanProperties, +} + +impl DistributedAnalyzeRootExec { + pub fn new(input: Arc, verbose: bool, show_statistics: bool) -> Self { + let field_a = Field::new("Task", DataType::Utf8, false); + let field_b = Field::new("Plan with Metrics", DataType::Utf8, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + let properties = PlanProperties::new( + EquivalenceProperties::new(schema.clone()), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + ); + + Self { + input, + verbose, + show_statistics, + properties, + } + } +} + +impl DisplayAs for DistributedAnalyzeRootExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "DistributedAnalyzeRootExec[verbose = {}, show stats = {}]", + self.verbose, self.show_statistics + ) + } +} + +impl ExecutionPlan for DistributedAnalyzeRootExec { + fn name(&self) -> &str { + "DistributedAnalyzeRootExec" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&std::sync::Arc> { + vec![&self.input] + } + + fn with_new_children( + self: std::sync::Arc, + children: Vec>, + ) -> Result> { + // TODO: generalize this + assert_eq!(children.len(), 1); + Ok(Arc::new(Self::new( + children[0].clone(), + self.verbose, + self.show_statistics, + ))) + } + + fn execute( + &self, + partition: usize, + context: std::sync::Arc, + ) -> Result { + let task_outputs = context + .session_config() + .get_extension::() + .unwrap_or(Arc::new(CtxAnnotatedOutputs::default())) + .0 + .clone(); + + assert!( + partition == 0, + "DistributedAnalyzeRootExec expects only partition 0" + ); + + let host = context + .session_config() + .get_extension::() + .ok_or_else(|| { + DataFusionError::Internal( + "CtxHost not set in session config for DistributedAnalyzeRootExec".to_string(), + ) + })? + .0 + .clone(); + + let stage_id = context + .session_config() + .get_extension::() + .ok_or_else(|| { + DataFusionError::Internal( + "CtxStageId not set in session config for DistributedAnalyzeRootExec" + .to_string(), + ) + })? + .0; + + let partition_group = context + .session_config() + .get_extension::() + .ok_or_else(|| { + DataFusionError::Internal( + "CtxPartitionGroup not set in session config for DistributedAnalyzeRootExec" + .to_string(), + ) + })? + .0 + .clone(); + + // we want to gather all partitions + let coalesce = Arc::new(CoalescePartitionsExec::new(self.input.clone())); + + let mut input_stream = coalesce.execute(partition, context)?; + + let schema_capture = self.schema().clone(); + let input_capture = self.input.clone(); + let show_statistics_capture = self.show_statistics; + let verbose_capture = self.verbose; + + let fmt_plan = move || -> String { + DisplayableExecutionPlan::with_metrics(input_capture.as_ref()) + .set_show_statistics(show_statistics_capture) + .indent(verbose_capture) + .to_string() + }; + + let output = async move { + // consume input, and we do not have to send it downstream as we are the + // root of the distributed analyze so we can discard the results just like + // regular AnalyzeExec + let mut done = false; + while !done { + match input_stream.next().await.transpose() { + Ok(Some(batch)) => { + // we consume the batch, yum. + trace!("consumed {} ", batch.num_rows()); + } + Ok(None) => done = true, + Err(e) => { + return Err(e); + } + } + } + let annotated_plan = fmt_plan(); + let toutput = AnnotatedTaskOutput { + plan: annotated_plan, + host: Some(host), + stage_id, + partition_group, + }; + + let mut tasks = task_outputs.lock(); + tasks.push(toutput); + + tasks.sort_by_key(|t| (t.stage_id, t.partition_group.clone())); + + trace!("sorted tasks: {:?}", tasks); + + let mut task_builder = StringBuilder::with_capacity(1, 1024); + let mut plan_builder = StringBuilder::with_capacity(1, 1024); + + for task_output in tasks.iter() { + task_builder.append_value(format!( + "Task: Stage {}, Partitions {:?}\nHost: {}", + task_output.stage_id, + task_output.partition_group, + task_output + .host + .as_ref() + .map(|h| format!("{} {}", h.name, h.addr)) + .unwrap_or("Unknown".to_string()) + )); + plan_builder.append_value(&task_output.plan); + } + + RecordBatch::try_new( + schema_capture, + vec![ + Arc::new(task_builder.finish()), + Arc::new(plan_builder.finish()), + ], + ) + .map_err(DataFusionError::from) + .inspect(|batch| { + debug!("returning record batch {:?}", batch); + }) + }; + + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + futures::stream::once(output), + ))) + } +} diff --git a/src/codec.rs b/src/codec.rs index 9a7024f3..346b768f 100644 --- a/src/codec.rs +++ b/src/codec.rs @@ -3,7 +3,6 @@ use std::sync::Arc; use arrow::datatypes::Schema; use datafusion::{ common::{internal_datafusion_err, internal_err, Result}, - datasource::source::DataSourceExec, execution::FunctionRegistry, physical_plan::{displayable, ExecutionPlan}, }; @@ -17,14 +16,18 @@ use datafusion_proto::{ use prost::Message; use crate::{ + analyze::{DistributedAnalyzeExec, DistributedAnalyzeRootExec}, isolator::PartitionIsolatorExec, logging::trace, max_rows::MaxRowsExec, protobuf::{ - df_ray_exec_node::Payload, DfRayExecNode, DfRayStageReaderExecNode, MaxRowsExecNode, - PartitionIsolatorExecNode, + df_ray_exec_node::Payload, DfRayExecNode, DfRayStageReaderExecNode, + DistributedAnalyzeExecNode, DistributedAnalyzeRootExecNode, MaxRowsExecNode, + PartitionIsolatorExecNode, RecordBatchExecNode, }, + record_batch_exec::RecordBatchExec, stage_reader::DFRayStageReaderExec, + util::{batch_to_ipc, ipc_to_batch}, }; #[derive(Debug)] @@ -90,12 +93,40 @@ impl PhysicalExtensionCodec for DFRayCodec { ))) } } - Payload::NumpangExec(_) => Err(internal_datafusion_err!( - "NumpangExec not supported in open source version" - )), - Payload::ContextExec(_) => Err(internal_datafusion_err!( - "ContextExec not supported in open source version" - )), + Payload::DistributedAnalyzeExec(distributed_analyze_exec_node) => { + if inputs.len() != 1 { + Err(internal_datafusion_err!( + "DistributedAnalyzeExec requires one input" + )) + } else { + Ok(Arc::new(DistributedAnalyzeExec::new( + inputs[0].clone(), + distributed_analyze_exec_node.verbose, + distributed_analyze_exec_node.show_statistics, + ))) + } + } + Payload::DistributedAnalyzeRootExec(distributed_analyze_root_exec_node) => { + if inputs.len() != 1 { + Err(internal_datafusion_err!( + "DistributedAnalyzeRootExec requires one input" + )) + } else { + Ok(Arc::new(DistributedAnalyzeRootExec::new( + inputs[0].clone(), + distributed_analyze_root_exec_node.verbose, + distributed_analyze_root_exec_node.show_statistics, + ))) + } + } + Payload::RecordBatchExec(rb_exec) => { + // deserialize the record batch stored in the opaque bytes field + let batch = ipc_to_batch(&rb_exec.batch).map_err(|e| { + internal_datafusion_err!("Failed to decode RecordBatch: {:#?}", e) + })?; + + Ok(Arc::new(RecordBatchExec::new(batch))) + } } } else { internal_err!("cannot decode proto extension in dfray codec") @@ -124,7 +155,6 @@ impl PhysicalExtensionCodec for DFRayCodec { Payload::StageReaderExec(pb) } else if let Some(pi) = node.as_any().downcast_ref::() { let pb = PartitionIsolatorExecNode { - dummy: 0.0, partition_count: pi.partition_count as u64, }; @@ -134,8 +164,25 @@ impl PhysicalExtensionCodec for DFRayCodec { max_rows: max.max_rows as u64, }; Payload::MaxRowsExec(pb) - } else if let Some(_exec) = node.as_any().downcast_ref::() { - return internal_err!("DataSourceExec encoding not supported in open source version"); + } else if let Some(exec) = node.as_any().downcast_ref::() { + let pb = DistributedAnalyzeExecNode { + verbose: exec.verbose, + show_statistics: exec.show_statistics, + }; + Payload::DistributedAnalyzeExec(pb) + } else if let Some(exec) = node.as_any().downcast_ref::() { + let pb = DistributedAnalyzeRootExecNode { + verbose: exec.verbose, + show_statistics: exec.show_statistics, + }; + Payload::DistributedAnalyzeRootExec(pb) + } else if let Some(exec) = node.as_any().downcast_ref::() { + let pb = RecordBatchExecNode { + batch: batch_to_ipc(&exec.batch).map_err(|e| { + internal_datafusion_err!("Failed to encode RecordBatch: {:#?}", e) + })?, + }; + Payload::RecordBatchExec(pb) } else { return internal_err!("Not supported node to encode to proto"); }; diff --git a/src/explain.rs b/src/explain.rs index 4f73d4c1..5bf6968c 100644 --- a/src/explain.rs +++ b/src/explain.rs @@ -1,227 +1,73 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you 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::{any::Any, fmt::Formatter, sync::Arc}; - -use arrow::{array::StringArray, datatypes::SchemaRef, record_batch::RecordBatch}; +use anyhow::Context; +use arrow::{ + array::StringArray, + datatypes::{DataType, Field, Schema}, + record_batch::RecordBatch, +}; use datafusion::{ - execution::TaskContext, - physical_expr::EquivalenceProperties, - physical_plan::{ - displayable, - execution_plan::{Boundedness, EmissionType}, - memory::MemoryStream, - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - SendableRecordBatchStream, - }, + logical_expr::LogicalPlan, + physical_plan::{displayable, ExecutionPlan}, + prelude::SessionContext, }; -/// Custom distributed EXPLAIN execution plan that also returns distributed plan and stages -#[derive(Debug)] -pub struct DistributedExplainExec { - schema: SchemaRef, - logical_plan: String, - physical_plan: String, - distributed_plan: String, - distributed_stages: String, - properties: PlanProperties, -} - -impl DistributedExplainExec { - pub fn new( - schema: SchemaRef, - logical_plan: String, - physical_plan: String, - distributed_plan: String, - distributed_stages: String, - ) -> Self { - // properties required by the ExecutionPlan trait - let properties = PlanProperties::new( - EquivalenceProperties::new(schema.clone()), - Partitioning::UnknownPartitioning(1), - EmissionType::Incremental, - Boundedness::Bounded, - ); - - Self { - schema, - logical_plan, - physical_plan, - distributed_plan, - distributed_stages, - properties, - } - } - - pub fn logical_plan(&self) -> &str { - &self.logical_plan - } - - pub fn physical_plan(&self) -> &str { - &self.physical_plan - } - - pub fn distributed_plan(&self) -> &str { - &self.distributed_plan - } - - pub fn distributed_stages(&self) -> &str { - &self.distributed_stages - } - - /// Format distributed stages for display - pub fn format_distributed_stages(stages: &[crate::planning::DFRayStage]) -> String { - let mut result = String::new(); - for (i, stage) in stages.iter().enumerate() { - result.push_str(&format!("Stage {}:\n", stage.stage_id)); - result.push_str(&format!( - " Partition Groups: {:?}\n", - stage.partition_groups - )); - result.push_str(&format!(" Full Partitions: {}\n", stage.full_partitions)); - result.push_str(" Plan:\n"); - let plan_display = format!("{}", displayable(stage.plan.as_ref()).indent(true)); - for line in plan_display.lines() { - result.push_str(&format!(" {}\n", line)); - } - if i < stages.len() - 1 { - result.push('\n'); - } +use crate::{result::Result, util::bytes_to_physical_plan, vocab::DDTask}; + +pub fn format_distributed_tasks(tasks: &[DDTask]) -> Result { + let mut result = String::new(); + for (i, task) in tasks.iter().enumerate() { + let plan = bytes_to_physical_plan(&SessionContext::new(), &task.plan_bytes) + .context("unable to decode task plan for formatted output")?; + + result.push_str(&format!( + "Task: Stage {}, Partitions {:?}\n", + task.stage_id, task.partition_group + )); + result.push_str(&format!(" Full Partitions: {}\n", task.full_partitions)); + result.push_str(" Plan:\n"); + let plan_display = format!("{}", displayable(plan.as_ref()).indent(true)); + for line in plan_display.lines() { + result.push_str(&format!(" {}\n", line)); } - if result.is_empty() { - result.push_str("No distributed stages generated"); + if i < tasks.len() - 1 { + result.push('\n'); } - result - } -} - -impl DisplayAs for DistributedExplainExec { - fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { - write!(f, "DistributedExplainExec") - } -} - -impl ExecutionPlan for DistributedExplainExec { - fn name(&self) -> &str { - "DistributedExplainExec" } - - fn as_any(&self) -> &dyn Any { - self - } - - fn properties(&self) -> &PlanProperties { - &self.properties - } - - fn children(&self) -> Vec<&Arc> { - vec![] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> datafusion::error::Result> { - Ok(self) - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> datafusion::error::Result { - let schema = self.schema.clone(); - - // Create the result data with our 4 plan types - let plan_types = StringArray::from(vec![ - "logical_plan", - "physical_plan", - "distributed_plan", - "distributed_stages", - ]); - let plans = StringArray::from(vec![ - self.logical_plan.as_str(), - self.physical_plan.as_str(), - self.distributed_plan.as_str(), - self.distributed_stages.as_str(), - ]); - - let batch = - RecordBatch::try_new(schema.clone(), vec![Arc::new(plan_types), Arc::new(plans)]) - .map_err(|e| datafusion::error::DataFusionError::ArrowError(Box::new(e), None))?; - - // Use MemoryStream which is designed for DataFusion execution plans - let stream = MemoryStream::try_new(vec![batch], schema, None)?; - - Ok(Box::pin(stream)) - } - - fn schema(&self) -> SchemaRef { - self.schema.clone() + if result.is_empty() { + result.push_str("No distributed tasks generated"); } + Ok(result) } -/// Check if this is an EXPLAIN query (but not EXPLAIN ANALYZE) -/// -/// This function distinguishes between: -/// - EXPLAIN queries (returns true) - show plan information only -/// - EXPLAIN ANALYZE queries (returns false) - execute and show runtime stats -/// - Regular queries (returns false) - normal query execution -pub fn is_explain_query(query: &str) -> bool { - let query_upper = query.trim().to_uppercase(); - // Must start with "EXPLAIN" followed by whitespace or end of string - let is_explain = query_upper.starts_with("EXPLAIN") - && (query_upper.len() == 7 - || query_upper - .chars() - .nth(7) - .is_some_and(|c| c.is_whitespace())); - let is_explain_analyze = query_upper.starts_with("EXPLAIN ANALYZE"); - is_explain && !is_explain_analyze -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_is_explain_query() { - // Test EXPLAIN queries (should return true) - assert!(is_explain_query("EXPLAIN SELECT * FROM table")); - assert!(is_explain_query("explain select * from table")); - assert!(is_explain_query(" EXPLAIN SELECT 1")); - assert!(is_explain_query("EXPLAIN\nSELECT * FROM test")); - - // Test EXPLAIN ANALYZE queries (should return false) - assert!(!is_explain_query("EXPLAIN ANALYZE SELECT * FROM table")); - assert!(!is_explain_query("explain analyze SELECT * FROM table")); - assert!(!is_explain_query(" EXPLAIN ANALYZE SELECT 1")); - - // Test regular queries (should return false) - assert!(!is_explain_query("SELECT * FROM table")); - assert!(!is_explain_query("INSERT INTO table VALUES (1)")); - assert!(!is_explain_query("UPDATE table SET col = 1")); - assert!(!is_explain_query("DELETE FROM table")); - assert!(!is_explain_query("CREATE TABLE test (id INT)")); - - // Test edge cases - assert!(!is_explain_query("")); - assert!(!is_explain_query(" ")); - assert!(!is_explain_query("EXPLAINSELECT")); // No space - assert!(is_explain_query("EXPLAIN")); // Just EXPLAIN - } +pub fn build_explain_batch( + logical_plan: &LogicalPlan, + physical_plan: &Arc, + distributed_plan: &Arc, + distributed_tasks: &[DDTask], +) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("plan_type", DataType::Utf8, false), + Field::new("plan", DataType::Utf8, false), + ])); + + // Create the result data with our 4 plan types + let plan_types = StringArray::from(vec![ + "logical_plan", + "physical_plan", + "distributed_plan", + "distributed_tasks", + ]); + let plans = StringArray::from(vec![ + logical_plan.display_indent().to_string(), + displayable(physical_plan.as_ref()).indent(true).to_string(), + displayable(distributed_plan.as_ref()) + .indent(true) + .to_string(), + format_distributed_tasks(distributed_tasks)?, + ]); + + let batch = RecordBatch::try_new(schema, vec![Arc::new(plan_types), Arc::new(plans)])?; + + Ok(batch) } diff --git a/src/flight.rs b/src/flight.rs index 26febf2c..a46ec241 100644 --- a/src/flight.rs +++ b/src/flight.rs @@ -56,7 +56,8 @@ impl FlightService for FlightServ { request: Request, ) -> Result, Status> { self.handler.do_get(request).await.inspect_err(|e| { - error!("Error in do_get: {:?}", e); + error!("Error in do_get: {:#?}", e); + log::error!("LOG Error in do_get: {:#?}", e); }) } @@ -107,7 +108,7 @@ impl FlightService for FlightServ { request: Request, ) -> Result, Status> { self.handler.do_action(request).await.inspect_err(|e| { - error!("Error in do_action: {:?}", e); + error!("Error in do_action: {:#?}", e); }) } @@ -166,7 +167,7 @@ impl FlightSqlService for FlightSqlServ { .get_flight_info_statement(query, request) .await .inspect_err(|e| { - error!("Error in do_flight_info_statement: {:?}", e); + error!("Error in do_flight_info_statement: {:#?}", e); }) } @@ -192,7 +193,7 @@ impl FlightSqlService for FlightSqlServ { .do_get_statement(ticket, request) .await .inspect_err(|e| { - error!("Error in do_get_statement: {:?}", e); + error!("Error in do_get_statement: {:#?}", e); }) } } diff --git a/src/flight_handlers.rs b/src/flight_handlers.rs deleted file mode 100644 index c614362f..00000000 --- a/src/flight_handlers.rs +++ /dev/null @@ -1,520 +0,0 @@ -use std::sync::Arc; - -use arrow::datatypes::{Schema, SchemaRef}; -use arrow_flight::{ - encode::FlightDataEncoderBuilder, - error::FlightError, - sql::{ProstMessageExt, TicketStatementQuery}, - FlightEndpoint, FlightInfo, Ticket, -}; -use datafusion::physical_plan::{ - coalesce_partitions::CoalescePartitionsExec, ExecutionPlan, Partitioning, -}; -use futures::TryStreamExt; -use prost::Message; -use tonic::{Response, Status}; - -use crate::{ - explain::DistributedExplainExec, - logging::{debug, trace}, - planning::{add_ctx_extentions, get_ctx}, - protobuf::{DistributedExplainExecNode, TicketStatementData}, - query_planner::QueryPlanner, - result::Result, - stage_reader::DFRayStageReaderExec, - util::{display_plan_with_partition_counts, get_addrs}, - vocab::Addrs, -}; - -/// Handler for Arrow Flight SQL requests and responses -pub struct FlightRequestHandler { - pub planner: QueryPlanner, -} - -impl FlightRequestHandler { - pub fn new(planner: QueryPlanner) -> Self { - Self { planner } - } - - /// Create a FlightInfo response with the given ticket data - pub fn create_flight_info_response( - &self, - query_id: String, - final_addrs: Addrs, - final_stage_id: u64, - schema: SchemaRef, - explain_data: Option, - ) -> Result { - let mut flight_info = FlightInfo::new() - .try_with_schema(&schema) - .map_err(|e| Status::internal(format!("Could not create flight info {e:?}")))?; - - let ticket_data = TicketStatementData { - query_id, - stage_id: final_stage_id, - stage_addrs: Some(final_addrs.into()), - schema: Some( - schema - .try_into() - .map_err(|e| Status::internal(format!("Could not convert schema {e:?}")))?, - ), - explain_data, - }; - - let ticket = Ticket::new( - TicketStatementQuery { - statement_handle: ticket_data.encode_to_vec().into(), - } - .as_any() - .encode_to_vec(), - ); - - let endpoint = FlightEndpoint::new().with_ticket(ticket); - flight_info = flight_info.with_endpoint(endpoint); - - Ok(flight_info) - } - - /// Handle EXPLAIN query requests by preparing plans for visualization. - /// - /// EXPLAIN queries return comprehensive plan information including logical, physical, - /// distributed plan, and execution stages for analysis and debugging purposes. - pub async fn handle_explain_request( - &self, - query: &str, - ) -> Result, Status> { - let plans = self - .planner - .prepare_explain(query) - .await - .map_err(|e| Status::internal(format!("Could not prepare EXPLAIN query {e:?}")))?; - - debug!("get flight info: EXPLAIN query id {}", plans.query_id); - - let explain_data = plans.explain_data.map(|data| DistributedExplainExecNode { - schema: data.schema().as_ref().try_into().ok(), - logical_plan: data.logical_plan().to_string(), - physical_plan: data.physical_plan().to_string(), - distributed_plan: data.distributed_plan().to_string(), - distributed_stages: data.distributed_stages().to_string(), - }); - - let flight_info = self.create_flight_info_response( - plans.query_id, - plans.worker_addresses, - plans.final_stage_id, - plans.schema, - explain_data, - )?; - - trace!("get_flight_info_statement done for EXPLAIN"); - Ok(Response::new(flight_info)) - } - - /// Handle query requests by preparing execution plans and stages. - /// - /// Query focus on execution readiness, returning only the essential - /// metadata needed to execute the distributed query plan. - pub async fn handle_query_request(&self, query: &str) -> Result, Status> { - let query_plan = self - .planner - .prepare_query(query) - .await - .map_err(|e| Status::internal(format!("Could not prepare query {e:?}")))?; - - debug!("get flight info: query id {}", query_plan.query_id); - - let flight_info = self.create_flight_info_response( - query_plan.query_id, - query_plan.worker_addresses, - query_plan.final_stage_id, - query_plan.schema, - None, // Regular queries don't have explain data - )?; - - trace!("get_flight_info_statement done"); - Ok(Response::new(flight_info)) - } - - pub async fn handle_substrait_info_request( - &self, - substrait_plan: datafusion_substrait::substrait::proto::Plan, - ) -> Result, Status> { - let query_plan = self - .planner - .prepare_substrait_query(substrait_plan) - .await - .map_err(|e| Status::internal(format!("Could not prepare query {e:?}")))?; - - debug!("get flight info: query id {}", query_plan.query_id); - - let flight_info = self.create_flight_info_response( - query_plan.query_id, - query_plan.worker_addresses, - query_plan.final_stage_id, - query_plan.schema, - None, // Regular queries don't have explain data - )?; - - trace!("get_flight_info_statement done"); - Ok(Response::new(flight_info)) - } - - /// Handle execution of EXPLAIN statement queries. - /// - /// This function does not execute the plan but returns all plans we want to display to the user. - pub async fn handle_explain_statement_execution( - &self, - tsd: TicketStatementData, - remote_addr: &str, - ) -> Result, Status> { - let explain_data = tsd.explain_data.as_ref().ok_or_else(|| { - Status::internal("No explain_data in TicketStatementData for EXPLAIN query") - })?; - - let schema: Schema = explain_data - .schema - .as_ref() - .ok_or_else(|| Status::internal("No schema in ExplainData"))? - .try_into() - .map_err(|e| Status::internal(format!("Cannot convert schema {e}")))?; - - let explain_plan = Arc::new(DistributedExplainExec::new( - Arc::new(schema), - explain_data.logical_plan.clone(), - explain_data.physical_plan.clone(), - explain_data.distributed_plan.clone(), - explain_data.distributed_stages.clone(), - )) as Arc; - - debug!( - "EXPLAIN request for query_id {} from {} explain plan:\n{}", - tsd.query_id, - remote_addr, - display_plan_with_partition_counts(&explain_plan) - ); - - // Create dummy addresses for EXPLAIN execution - let mut dummy_addrs = std::collections::HashMap::new(); - let mut partition_addrs = std::collections::HashMap::new(); - partition_addrs.insert( - 0u64, - vec![("explain_local".to_string(), "local".to_string())], - ); - dummy_addrs.insert(0u64, partition_addrs); - - self.execute_plan_and_build_stream(explain_plan, tsd.query_id, dummy_addrs) - .await - } - - /// Handle execution of regular statement queries - /// - /// This function executes the plan and returns the results to the client. - pub async fn handle_regular_statement_execution( - &self, - tsd: TicketStatementData, - remote_addr: &str, - ) -> Result, Status> { - let schema: Schema = tsd - .schema - .as_ref() - .ok_or_else(|| Status::internal("No schema in TicketStatementData"))? - .try_into() - .map_err(|e| Status::internal(format!("Cannot convert schema {e}")))?; - - // Create an Addrs from the final stage information in the tsd - let stage_addrs = tsd.stage_addrs.ok_or_else(|| { - Status::internal("No stages_addrs in TicketStatementData, cannot proceed") - })?; - - let addrs: Addrs = get_addrs(&stage_addrs).map_err(|e| { - Status::internal(format!("Cannot get addresses from stage_addrs {e:?}")) - })?; - - trace!("calculated addrs: {:?}", addrs); - - // Validate that addrs contains exactly one stage - self.validate_single_stage_addrs(&addrs, tsd.stage_id)?; - - let stage_partition_addrs = addrs.get(&tsd.stage_id).ok_or_else(|| { - Status::internal(format!( - "No partition addresses found for stage_id {}", - tsd.stage_id - )) - })?; - - let plan = Arc::new( - DFRayStageReaderExec::try_new( - Partitioning::UnknownPartitioning(stage_partition_addrs.len()), - Arc::new(schema), - tsd.stage_id, - ) - // TODO: revisit this to allow for consuming a particular partition - .map(|stg| CoalescePartitionsExec::new(Arc::new(stg))) - .map_err(|e| Status::internal(format!("Unexpected error {e}")))?, - ) as Arc; - - debug!( - "request for query_id {} from {} reader plan:\n{}", - tsd.query_id, - remote_addr, - display_plan_with_partition_counts(&plan) - ); - - self.execute_plan_and_build_stream(plan, tsd.query_id, addrs) - .await - } - - /// Validate that addresses contain exactly one stage with the expected stage_id - pub fn validate_single_stage_addrs( - &self, - addrs: &Addrs, - expected_stage_id: u64, - ) -> Result<(), Status> { - if addrs.len() != 1 { - return Err(Status::internal(format!( - "Expected exactly one stage in addrs, got {}", - addrs.len() - ))); - } - if !addrs.contains_key(&expected_stage_id) { - return Err(Status::internal(format!( - "No addresses found for stage_id {} in addrs", - expected_stage_id - ))); - } - Ok(()) - } - - /// Execute a plan and build the response stream. - /// - /// This function handles the common execution logic for both regular queries (which return data) - /// and EXPLAIN queries (which return plan information as text). - pub async fn execute_plan_and_build_stream( - &self, - plan: Arc, - query_id: String, - addrs: Addrs, - ) -> Result, Status> { - let mut ctx = - get_ctx().map_err(|e| Status::internal(format!("Could not create context {e:?}")))?; - - add_ctx_extentions(&mut ctx, "proxy", &query_id, addrs, None) - .map_err(|e| Status::internal(format!("Could not add context extensions {e:?}")))?; - - // TODO: revisit this to allow for consuming a partitular partition - trace!("calling execute plan"); - let partition = 0; - let stream = plan - .execute(partition, ctx.task_ctx()) - .map_err(|e| { - Status::internal(format!( - "Error executing plan for query_id {} partition {}: {e:?}", - query_id, partition - )) - })? - .map_err(|e| FlightError::ExternalError(Box::new(e))); - - let out_stream = FlightDataEncoderBuilder::new() - .build(stream) - .map_err(move |e| Status::internal(format!("Unexpected error building stream {e:?}"))); - - Ok(Response::new(Box::pin(out_stream))) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::test_utils::explain_test_helpers::{ - create_explain_ticket_statement_data, create_test_flight_handler, - verify_explain_stream_results, - }; - use std::collections::HashMap; - - // ////////////////////////////////////////////////////////////// - // Test helper functions - // ////////////////////////////////////////////////////////////// - - /// Create test worker addresses - fn create_test_addrs() -> Addrs { - let mut addrs = HashMap::new(); - let mut stage_addrs = HashMap::new(); - stage_addrs.insert( - 1u64, - vec![ - ("worker1".to_string(), "localhost:8001".to_string()), - ("worker2".to_string(), "localhost:8002".to_string()), - ], - ); - addrs.insert(1u64, stage_addrs); - addrs - } - - // ////////////////////////////////////////////////////////////// - // Unit tests - // ////////////////////////////////////////////////////////////// - - #[test] - fn test_validate_single_stage_addrs_success() { - let handler = create_test_flight_handler(); - let addrs = create_test_addrs(); - - // Should succeed with valid single stage - let result = handler.validate_single_stage_addrs(&addrs, 1); - assert!(result.is_ok()); - } - - #[test] - fn test_validate_single_stage_addrs_multiple_stages() { - let handler = create_test_flight_handler(); - - // Create addresses with multiple stages - let mut addrs = HashMap::new(); - let mut stage1_addrs = HashMap::new(); - stage1_addrs.insert( - 1u64, - vec![("worker1".to_string(), "localhost:8001".to_string())], - ); - let mut stage2_addrs = HashMap::new(); - stage2_addrs.insert( - 2u64, - vec![("worker2".to_string(), "localhost:8002".to_string())], - ); - addrs.insert(1u64, stage1_addrs); - addrs.insert(2u64, stage2_addrs); - - // Should fail with multiple stages - let result = handler.validate_single_stage_addrs(&addrs, 1); - assert!(result.is_err()); - - if let Err(status) = result { - assert!(status.message().contains("Expected exactly one stage")); - } - } - - #[test] - fn test_validate_single_stage_addrs_wrong_stage_id() { - let handler = create_test_flight_handler(); - let addrs = create_test_addrs(); // Contains stage_id 1 - - // Should fail when looking for non-existent stage - let result = handler.validate_single_stage_addrs(&addrs, 999); - assert!(result.is_err()); - - if let Err(status) = result { - assert!(status - .message() - .contains("No addresses found for stage_id 999")); - } - } - - #[test] - fn test_validate_single_stage_addrs_empty() { - let handler = create_test_flight_handler(); - let addrs: Addrs = HashMap::new(); - - // Should fail with empty addresses - let result = handler.validate_single_stage_addrs(&addrs, 1); - assert!(result.is_err()); - - if let Err(status) = result { - assert!(status.message().contains("Expected exactly one stage")); - } - } - - // ////////////////////////////////////////////////////////////// - // Handler core function tests - // ////////////////////////////////////////////////////////////// - - #[tokio::test] - async fn test_handle_explain_request() { - let handler = create_test_flight_handler(); - let query = "EXPLAIN SELECT 1 as test_col, 'hello' as text_col"; - - let result = handler.handle_explain_request(query).await; - assert!(result.is_ok()); - - let response = result.unwrap(); - let flight_info = response.into_inner(); - - // Verify FlightInfo structure - assert!(!flight_info.schema.is_empty()); - assert_eq!(flight_info.endpoint.len(), 1); - assert!(flight_info.endpoint[0].ticket.is_some()); - - // Verify that ticket has content (encoded TicketStatementData) - let ticket = flight_info.endpoint[0].ticket.as_ref().unwrap(); - assert!(!ticket.ticket.is_empty()); - - println!( - "✓ FlightInfo created successfully with {} schema bytes and ticket with {} bytes", - flight_info.schema.len(), - ticket.ticket.len() - ); - } - - #[tokio::test] - async fn test_handle_explain_request_invalid_query() { - let handler = create_test_flight_handler(); - - // Test with EXPLAIN ANALYZE (should fail) - let query = "EXPLAIN ANALYZE SELECT 1"; - let result = handler.handle_explain_request(query).await; - assert!(result.is_err()); - - let error = result.unwrap_err(); - assert_eq!(error.code(), tonic::Code::Internal); - assert!(error.message().contains("Could not prepare EXPLAIN query")); - } - - #[tokio::test] - async fn test_handle_explain_statement_execution() { - let handler = create_test_flight_handler(); - - // First prepare an EXPLAIN query to get the ticket data structure - let query = "EXPLAIN SELECT 1 as test_col"; - let plans = handler.planner.prepare_explain(query).await.unwrap(); - - // Create the TicketStatementData that would be sent to do_get_statement - let tsd = create_explain_ticket_statement_data(plans); - - // Test the execution - let result = handler - .handle_explain_statement_execution(tsd, "test_remote") - .await; - assert!(result.is_ok()); - - let response = result.unwrap(); - let stream = response.into_inner(); - - // Use shared verification function - verify_explain_stream_results(stream).await; - } - - #[tokio::test] - async fn test_handle_explain_statement_execution_missing_explain_data() { - let handler = create_test_flight_handler(); - - // Create TicketStatementData without explain_data (should fail) - let tsd = TicketStatementData { - query_id: "test_query".to_string(), - stage_id: 0, - stage_addrs: None, - schema: None, - explain_data: None, - }; - - let result = handler - .handle_explain_statement_execution(tsd, "test_remote") - .await; - assert!(result.is_err()); - - if let Err(error) = result { - assert_eq!(error.code(), tonic::Code::Internal); - assert!(error - .message() - .contains("No explain_data in TicketStatementData")); - } - } -} diff --git a/src/friendly.rs b/src/friendly.rs index 412ae926..17dda480 100644 --- a/src/friendly.rs +++ b/src/friendly.rs @@ -1619,9 +1619,11 @@ impl Friendly { .lock() .map_err(|e| anyhow!("Failed to lock adjectives: {}", e))?; + let suffix = rng.gen_range(0..1000); + let index = rng.gen_range(0..adjectives.len()); let adj = adjectives.remove(index); - Ok(format!("{}-{}", adj, name)) + Ok(format!("{}-{}-{:0>4}", adj, name, suffix)) } } diff --git a/src/isolator.rs b/src/isolator.rs index a687077c..02fd8754 100644 --- a/src/isolator.rs +++ b/src/isolator.rs @@ -12,7 +12,7 @@ use datafusion::{ use crate::{ logging::{error, trace}, - vocab::{CtxName, CtxPartitionGroup}, + vocab::{CtxHost, CtxPartitionGroup}, }; /// This is a simple execution plan that isolates a partition from the input @@ -108,12 +108,11 @@ impl ExecutionPlan for PartitionIsolatorExec { )); } - let ctx_name = context + let ctx_name = &context .session_config() - .get_extension::() - .ok_or_else(|| internal_datafusion_err!("CtxName not set in session config"))? - .0 - .clone(); + .get_extension::() + .map(|ctx_host| ctx_host.0.to_string()) + .unwrap_or("unknown_context_host!".to_string()); let partitions_in_input = self.input.output_partitioning().partition_count() as u64; diff --git a/src/lib.rs b/src/lib.rs index 29fe2820..e6198dbb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,13 +22,12 @@ use crate::logging::info; mod proto; pub use proto::generated::protobuf; +pub mod analyze; pub mod codec; pub mod explain; pub mod flight; -pub mod flight_handlers; pub mod friendly; pub mod isolator; -pub mod k8s; pub mod logging; pub mod max_rows; pub mod physical; @@ -36,12 +35,13 @@ pub mod planning; pub mod processor_service; pub mod proxy_service; pub mod query_planner; +pub mod record_batch_exec; pub mod result; pub mod stage; pub mod stage_reader; -pub mod test_utils; pub mod util; pub mod vocab; +pub mod worker_discovery; #[cfg(not(target_env = "msvc"))] use tikv_jemallocator::Jemalloc; diff --git a/src/main.rs b/src/main.rs index 0f50127c..21429c94 100644 --- a/src/main.rs +++ b/src/main.rs @@ -33,13 +33,11 @@ async fn main() -> Result<()> { match args.mode.as_str() { "proxy" => { - let mut service = DFRayProxyService::new(args.port); - service.start_up().await?; + let service = DFRayProxyService::new(new_friendly_name()?, args.port).await?; service.serve().await?; } "worker" => { - let mut service = DFRayProcessorService::new(new_friendly_name()?, args.port); - service.start_up().await?; + let service = DFRayProcessorService::new(new_friendly_name()?, args.port).await?; service.serve().await?; } _ => { diff --git a/src/physical.rs b/src/physical.rs index b07ede14..f253a47d 100644 --- a/src/physical.rs +++ b/src/physical.rs @@ -47,7 +47,6 @@ impl Default for DFRayStageOptimizerRule { Self::new() } } - impl DFRayStageOptimizerRule { pub fn new() -> Self { Self {} @@ -72,6 +71,7 @@ impl PhysicalOptimizerRule for DFRayStageOptimizerRule { || plan.as_any().downcast_ref::().is_some() || plan.as_any().downcast_ref::().is_some() { + // insert a stage marker here so we know where to break up the physical plan later let stage = Arc::new(DFRayStageExec::new(plan, stage_counter)); stage_counter += 1; Ok(Transformed::yes(stage as Arc)) @@ -80,7 +80,7 @@ impl PhysicalOptimizerRule for DFRayStageOptimizerRule { } }; - let plan = plan.transform_up(up)?.data; + let plan = plan.clone().transform_up(up)?.data; let final_plan = Arc::new(DFRayStageExec::new(plan, stage_counter)) as Arc; diff --git a/src/planning.rs b/src/planning.rs index 18cc22d5..07923a70 100644 --- a/src/planning.rs +++ b/src/planning.rs @@ -18,9 +18,9 @@ use datafusion::{ logical_expr::LogicalPlan, physical_optimizer::PhysicalOptimizerRule, physical_plan::{ - coalesce_batches::CoalesceBatchesExec, displayable, joins::NestedLoopJoinExec, - repartition::RepartitionExec, sorts::sort::SortExec, ExecutionPlan, - ExecutionPlanProperties, + analyze::AnalyzeExec, coalesce_batches::CoalesceBatchesExec, displayable, + joins::NestedLoopJoinExec, repartition::RepartitionExec, sorts::sort::SortExec, + ExecutionPlan, ExecutionPlanProperties, }, prelude::{SQLOptions, SessionConfig, SessionContext}, }; @@ -29,20 +29,23 @@ use itertools::Itertools; use prost::Message; use crate::{ + analyze::{DistributedAnalyzeExec, DistributedAnalyzeRootExec}, isolator::PartitionIsolatorExec, logging::{debug, error, info, trace}, max_rows::MaxRowsExec, physical::DFRayStageOptimizerRule, - protobuf::{Host, Hosts, PartitionAddrs, StageAddrs, StageData}, result::{DFRayError, Result}, stage::DFRayStageExec, stage_reader::{DFRayStageReaderExec, QueryId}, util::{display_plan_with_partition_counts, get_client, physical_plan_to_bytes, wait_for}, - vocab::{Addrs, CtxName, CtxPartitionGroup, CtxStageAddrs}, + vocab::{ + Addrs, CtxAnnotatedOutputs, CtxHost, CtxPartitionGroup, CtxStageAddrs, CtxStageId, DDTask, + Host, Hosts, PartitionAddrs, StageAddrs, + }, }; #[derive(Debug)] -pub struct DFRayStage { +pub struct DDStage { /// our stage id pub stage_id: u64, /// the physical plan of our stage @@ -56,7 +59,7 @@ pub struct DFRayStage { pub full_partitions: bool, } -impl DFRayStage { +impl DDStage { fn new( stage_id: u64, plan: Arc, @@ -96,7 +99,7 @@ static STATE: LazyLock> = LazyLock::new(|| { pub fn get_ctx() -> Result { match &*STATE { - Ok(ctx) => Ok(SessionContext::new_with_state(ctx.clone())), + Ok(state) => Ok(SessionContext::new_with_state(state.clone())), Err(e) => Err(anyhow!("Context initialization failed: {}", e).into()), } } @@ -129,19 +132,16 @@ async fn make_state() -> Result { .await .context("Failed to add tables from environment")?; - add_views_from_env(&state) - .await - .context("Failed to add views from environment")?; - Ok(state) } pub fn add_ctx_extentions( ctx: &mut SessionContext, - ctx_name: &str, + host: &Host, query_id: &str, + stage_id: u64, stage_addrs: Addrs, - partition_group: Option>, + partition_group: Vec, ) -> Result<()> { let state = ctx.state_ref(); let mut guard = state.write(); @@ -149,15 +149,12 @@ pub fn add_ctx_extentions( config.set_extension(Arc::new(CtxStageAddrs(stage_addrs))); config.set_extension(Arc::new(QueryId(query_id.to_owned()))); - config.set_extension(Arc::new(CtxName(ctx_name.to_owned()))); - - if let Some(pg) = partition_group { - // this only matters if the plan includes an PartitionIsolatorExec, which looks - // for this for this extension and will be ignored otherwise + config.set_extension(Arc::new(CtxHost(host.clone()))); + config.set_extension(Arc::new(CtxStageId(stage_id))); + config.set_extension(Arc::new(CtxAnnotatedOutputs::default())); - trace!("Adding partition group: {:?}", pg); - config.set_extension(Arc::new(CtxPartitionGroup(pg))); - } + trace!("Adding partition group: {:?}", partition_group); + config.set_extension(Arc::new(CtxPartitionGroup(partition_group))); Ok(()) } @@ -210,38 +207,6 @@ pub async fn add_tables_from_env(state: &mut SessionState) -> Result<()> { Ok(()) } -pub async fn add_views_from_env(state: &SessionState) -> Result<()> { - // this string contains CREATE VIEW SQL statements separated by semicolons - let views_str = env::var("DFRAY_VIEWS"); - if views_str.is_err() { - info!("No DFRAY_VIEWS environment variable set, skipping view creation"); - return Ok(()); - } - - let ctx = SessionContext::new_with_state(state.clone()); - - for view_sql in views_str.unwrap().split(';') { - let view_sql = view_sql.trim(); - if view_sql.is_empty() { - continue; - } - - info!("creating view from env: {}", view_sql); - - // Execute the CREATE VIEW statement - match ctx.sql(view_sql).await { - Ok(_) => { - info!("Successfully created view: {}", view_sql); - } - Err(e) => { - return Err(anyhow!("Failed to create view '{}': {}", view_sql, e).into()); - } - } - } - - Ok(()) -} - pub async fn logical_planning(sql: &str, ctx: &SessionContext) -> Result { let options = SQLOptions::new(); let plan = ctx.state().create_logical_plan(sql).await?; @@ -277,7 +242,7 @@ pub async fn execution_planning( physical_plan: Arc, batch_size: usize, partitions_per_worker: Option, -) -> Result<(Arc, Vec)> { +) -> Result<(Arc, Vec)> { let mut stages = vec![]; let mut partition_groups = vec![]; @@ -307,7 +272,7 @@ pub async fn execution_planning( stage_exec.stage_id, )?) as Arc; - let stage = DFRayStage::new( + let stage = DDStage::new( stage_exec.stage_id, input.clone(), partition_groups.clone(), @@ -379,7 +344,7 @@ pub async fn execution_planning( // add coalesce and max rows to last stage let mut last_stage = stages.pop().ok_or(anyhow!("No stages found"))?; - last_stage = DFRayStage::new( + last_stage = DDStage::new( last_stage.stage_id, Arc::new(MaxRowsExec::new( Arc::new(CoalesceBatchesExec::new(last_stage.plan, batch_size)) @@ -393,45 +358,114 @@ pub async fn execution_planning( // done fixing last stage, put it back stages.push(last_stage); + if contains_analyze(stages[stages.len() - 1].plan.as_ref()) { + // if the plan contains an analyze, we need to add the distributed analyze + // stages to the plan + add_distributed_analyze(&mut stages, false, false)?; + } + let txt = stages .iter() .map(|stage| format!("{}", display_plan_with_partition_counts(&stage.plan))) .join(",\n"); - trace!("stages:{}", txt); + trace!("stages:\n{}", txt); Ok((distributed_plan_clone, stages)) } +fn contains_analyze(plan: &dyn ExecutionPlan) -> bool { + trace!( + "checking stage for analyze: {}", + displayable(plan).indent(false) + ); + if plan.as_any().downcast_ref::().is_some() { + true + } else { + for child in plan.children() { + if contains_analyze(child.as_ref()) { + return true; + } + } + false + } +} + +pub fn add_distributed_analyze( + stages: &mut [DDStage], + verbose: bool, + show_statistics: bool, +) -> Result<()> { + trace!("Adding distributed analyze to stages"); + let len = stages.len(); + for (i, stage) in stages.iter_mut().enumerate() { + if i == len - 1 { + let plan_without_analyze = stage + .plan + .clone() + .transform_down(|plan: Arc| { + if let Some(analyze) = plan.as_any().downcast_ref::() { + Ok(Transformed::yes(analyze.input().clone())) + } else { + Ok(Transformed::no(plan)) + } + })? + .data; + + trace!( + "plan without analyze: {}", + displayable(plan_without_analyze.as_ref()).indent(false) + ); + stage.plan = Arc::new(DistributedAnalyzeRootExec::new( + plan_without_analyze, + verbose, + show_statistics, + )) as Arc; + stage.partition_groups = vec![vec![0]]; // accounting for coalesce + } else { + stage.plan = Arc::new(DistributedAnalyzeExec::new( + stage.plan.clone(), + verbose, + show_statistics, + )) as Arc; + } + } + Ok(()) +} + /// Distribute the stages to the workers, assigning each stage to a worker /// Returns an Addrs containing the addresses of the workers that will execute /// final stage only as that's all we care about from the call site pub async fn distribute_stages( query_id: &str, - stages: Vec, - worker_addrs: Vec<(String, String)>, -) -> Result { + stages: Vec, + worker_addrs: Vec, +) -> Result<(Addrs, Vec)> { // map of worker name to address // FIXME: use types over tuples of strings, as we can accidently swap them and // not know - let mut workers: HashMap = worker_addrs.iter().cloned().collect(); + // a map of worker name to host + let mut workers: HashMap = worker_addrs + .iter() + .map(|host| (host.name.clone(), host.clone())) + .collect(); for attempt in 0..3 { // all stages to workers - let (stage_datas, final_addrs) = - assign_to_workers(query_id, &stages, workers.iter().collect())?; + let (task_datas, final_addrs) = + assign_to_workers(query_id, &stages, workers.values().collect())?; // we retry this a few times to ensure that the workers are ready // and can accept the stages - match try_distribute_stages(&stage_datas).await { - Ok(_) => return Ok(final_addrs), + match try_distribute_tasks(&task_datas).await { + Ok(_) => return Ok((final_addrs, task_datas)), Err(DFRayError::WorkerCommunicationError(bad_worker)) => { error!( "distribute stages for query {query_id} attempt {attempt} failed removing \ worker {bad_worker}. Retrying..." ); // if we cannot communicate with a worker, we remove it from the list of workers - workers.remove(&bad_worker); + workers.remove(&bad_worker.name); } Err(e) => return Err(e), } @@ -446,45 +480,34 @@ pub async fn distribute_stages( /// try to distribute the stages to the workers, if we cannot communicate with a /// worker return it as the element in the Err -async fn try_distribute_stages(stage_datas: &[StageData]) -> Result<()> { +async fn try_distribute_tasks(task_datas: &[DDTask]) -> Result<()> { // we can use the stage data to distribute the stages to workers - for stage_data in stage_datas { + for task_data in task_datas { trace!( - "Distributing stage_id {}, pg: {:?} to worker: {:?}", - stage_data.stage_id, - stage_data.partition_group, - stage_data.assigned_addr + "Distributing Task: stage_id {}, pg: {:?} to worker: {:?}", + task_data.stage_id, + task_data.partition_group, + task_data.assigned_host ); // populate its child stages - let mut stage_data = stage_data.clone(); - stage_data.stage_addrs = Some(get_stage_addrs_from_stages( + let mut stage_data = task_data.clone(); + stage_data.stage_addrs = Some(get_stage_addrs_from_tasks( &stage_data.child_stage_ids, - stage_datas, + task_datas, )?); - let mut client = match get_client( - &stage_data - .assigned_addr - .as_ref() - .context("Assigned stage address is missing")? - .name, - &stage_data - .assigned_addr - .as_ref() - .context("Assigned stage address is missing")? - .addr, - ) { + let host = stage_data + .assigned_host + .clone() + .context("Assigned host is missing for task data")?; + + let mut client = match get_client(&host) { Ok(client) => client, Err(e) => { error!("Couldn't not communicate with worker {e:#?}"); return Err(DFRayError::WorkerCommunicationError( - stage_data - .assigned_addr - .as_ref() - .cloned() - .unwrap() // we know we can unwrap it as we checked a few lines up - .name, + host.clone(), // here )); } }; @@ -518,16 +541,24 @@ async fn try_distribute_stages(stage_datas: &[StageData]) -> Result<()> { } // go through our stages, and further divide them into their partition -// groups and produce a StageData for each partition group and assign it +// groups and produce a DDTask for each partition group and assign it // to a worker fn assign_to_workers( query_id: &str, - stages: &[DFRayStage], - worker_addrs: Vec<(&String, &String)>, -) -> Result<(Vec, Addrs)> { - let mut stage_datas = vec![]; + stages: &[DDStage], + worker_addrs: Vec<&Host>, +) -> Result<(Vec, Addrs)> { + let mut task_datas = vec![]; let mut worker_idx = 0; + trace!( + "assigning stages: {:?}", + stages + .iter() + .map(|s| format!("stage_id: {}, pgs:{:?}", s.stage_id, s.partition_groups)) + .join(",\n") + ); + // keep track of which worker has the root of the plan tree (highest stage // number) let mut max_stage_id = -1; @@ -537,10 +568,7 @@ fn assign_to_workers( for partition_group in stage.partition_groups.iter() { let plan_bytes = physical_plan_to_bytes(stage.plan.clone())?; - let addr = Host { - name: worker_addrs[worker_idx].0.to_string(), - addr: worker_addrs[worker_idx].1.to_string(), - }; + let host = worker_addrs[worker_idx].clone(); worker_idx = (worker_idx + 1) % worker_addrs.len(); if stage.stage_id as isize > max_stage_id { @@ -556,11 +584,11 @@ fn assign_to_workers( .or_default() .entry(*part) .or_default() - .push((addr.name.clone(), addr.addr.clone())); + .push(host.clone()); } } - let stage_data = StageData { + let task_data = DDTask { query_id: query_id.to_string(), stage_id: stage.stage_id, plan_bytes, @@ -569,19 +597,16 @@ fn assign_to_workers( stage_addrs: None, // will be calculated and filled in later num_output_partitions: stage.plan.output_partitioning().partition_count() as u64, full_partitions: stage.full_partitions, - assigned_addr: Some(addr), + assigned_host: Some(host), }; - stage_datas.push(stage_data); + task_datas.push(task_data); } } - Ok((stage_datas, final_addrs)) + Ok((task_datas, final_addrs)) } -fn get_stage_addrs_from_stages( - target_stage_ids: &[u64], - stages: &[StageData], -) -> Result { +fn get_stage_addrs_from_tasks(target_stage_ids: &[u64], stages: &[DDTask]) -> Result { let mut stage_addrs = StageAddrs::default(); // this can be more efficient @@ -602,17 +627,12 @@ fn get_stage_addrs_from_stages( .iter() .filter(|s| s.stage_id == stage_id) .map(|s| { - Ok(Host { - name: s - .assigned_addr - .clone() - .context("assigned address missing")? - .name, - addr: s - .assigned_addr - .clone() - .context("assigned address missing")? - .addr, + s.assigned_host.clone().ok_or_else(|| { + anyhow!( + "Assigned address is missing for stage_id: {}", + stage_id + ) + .into() }) }) .collect::>>()?; @@ -622,18 +642,10 @@ fn get_stage_addrs_from_stages( .insert(*part, Hosts { hosts }); } else { // we have the full partition, so this is the only address required - let host = Host { - name: stage - .assigned_addr - .clone() - .context("Assigned address is missing")? - .name, - addr: stage - .assigned_addr - .clone() - .context("Assigned address is missing")? - .addr, - }; + let host = stage + .assigned_host + .clone() + .context("Assigned address is missing")?; partition_addrs .partition_addrs .insert(*part, Hosts { hosts: vec![host] }); diff --git a/src/processor_service.rs b/src/processor_service.rs index 1fb2f1c7..555ff912 100644 --- a/src/processor_service.rs +++ b/src/processor_service.rs @@ -16,7 +16,8 @@ // under the License. use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, + fmt::Debug, sync::Arc, time::{Duration, SystemTime}, }; @@ -29,10 +30,10 @@ use arrow_flight::{ }; use async_stream::stream; use datafusion::{ - physical_plan::{displayable, ExecutionPlan, ExecutionPlanProperties}, + physical_plan::{ExecutionPlan, ExecutionPlanProperties}, prelude::SessionContext, }; -use futures::{Stream, TryStreamExt}; +use futures::{StreamExt, TryStreamExt}; use parking_lot::{Mutex, RwLock}; use prost::Message; use tokio::{ @@ -42,47 +43,68 @@ use tokio::{ use tonic::{async_trait, transport::Server, Request, Response, Status}; use crate::{ + analyze::DistributedAnalyzeExec, flight::{FlightHandler, FlightServ}, logging::{debug, error, info, trace}, planning::{add_ctx_extentions, get_ctx}, - protobuf::{FlightTicketData, StageData}, + protobuf::{ + AnnotatedTaskOutput, AnnotatedTaskOutputs, FlightDataMetadata, FlightTicketData, Host, + }, result::{DFRayError, Result}, util::{ bytes_to_physical_plan, display_plan_with_partition_counts, get_addrs, - register_object_store_for_paths_in_plan, reporting_stream, + register_object_store_for_paths_in_plan, start_up, }, - vocab::{Addrs, CtxName}, + vocab::{Addrs, CtxAnnotatedOutputs, CtxPartitionGroup, DDTask}, }; #[derive(Eq, PartialEq, Hash, Clone, Debug)] -struct PlanKey { +struct StageKey { query_id: String, stage_id: u64, - partition: u64, } -// For each plan key, we may have multiple plans that we might need to hold -// with the same key. -type PlanVec = Vec<(SystemTime, SessionContext, Arc)>; +#[derive(Clone)] +struct Task { + partitions: HashSet, + ctx: SessionContext, + plan: Arc, +} + +impl Debug for Task { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Task") + .field("partitions", &self.partitions) + .finish() + } +} + +#[derive(Debug)] +struct StageTasks { + tasks: Vec, + insert_time: SystemTime, +} /// It only responds to the DoGet Arrow Flight method. struct DfRayProcessorHandler { /// our name, useful for logging name: String, + /// our address string, also useful for logging + addr: String, /// our map of query_id -> (session ctx, execution plan) - #[allow(clippy::type_complexity)] - plans: Arc>>, + stages: Arc>>, done: Arc>, } impl DfRayProcessorHandler { - pub fn new(name: String) -> Self { - let plans: Arc>> = Arc::new(RwLock::new(HashMap::new())); + pub fn new(name: String, addr: String) -> Self { + let stages: Arc>> = + Arc::new(RwLock::new(HashMap::new())); let done = Arc::new(Mutex::new(false)); // start a plan janitor ask to clean up old plans that were not collected for // any reason - let c_plans = plans.clone(); + let c_stages = stages.clone(); let c_done = done.clone(); let c_name = name.clone(); std::thread::spawn(move || { @@ -90,57 +112,62 @@ impl DfRayProcessorHandler { // wait for 10 seconds std::thread::sleep(Duration::from_secs(10)); if *c_done.lock() { - info!("{} plan janitor done", c_name); + info!("{} janitor done", c_name); break; } - trace!("{} plan janitor waking up", c_name); + trace!("{} janitor waking up", c_name); let now = SystemTime::now(); let mut to_remove = vec![]; { - let _guard = c_plans.read(); - for (key, plan_vec) in _guard.iter() { - if plan_vec.is_empty() { - error!("unexpectedly found empty plan vec. removing"); + let _guard = c_stages.read(); + for (key, stage_tasks) in _guard.iter() { + if stage_tasks.tasks.is_empty() { + error!("unexpectedly found empty stage tasks. removing"); to_remove.push(key.clone()); } else { // check if any plan in this vec is older than 1 minute - for (insert_time, _, _) in plan_vec.iter() { - if now - .duration_since(*insert_time) - .map(|d| d.as_secs() > 60) - .inspect_err(|e| { - error!("CANNOT COMPUTE DURATION OR REMOVE PLANS: {e:?}"); - }) - .unwrap_or(false) - { - to_remove.push(key.clone()); - break; - } + if now + .duration_since(stage_tasks.insert_time) + .map(|d| d.as_secs() > 60) + .inspect_err(|e| { + error!("CANNOT COMPUTE DURATION OR REMOVE STAGES: {e:?}"); + // maybe just panic here? + }) + .unwrap_or(false) + { + to_remove.push(key.clone()); + break; } } } } if !to_remove.is_empty() { - let mut _guard = c_plans.write(); + let mut _guard = c_stages.write(); for key in to_remove.iter() { _guard.remove(key); - debug!("{} removed old plan key {:?}", c_name, key); + debug!("{} removed old stage key {:?}", c_name, key); } } } }); - Self { name, plans, done } + Self { + name, + addr, + stages, + done, + } } #[allow(dead_code)] + /// shutdown pub fn all_done(&self) { *self.done.lock() = true; } - pub async fn add_plan( + pub async fn add_task( &self, query_id: String, stage_id: u64, @@ -172,9 +199,10 @@ impl DfRayProcessorHandler { }; trace!( - "{} adding plan for stage {} partitions: {:?} stage_addrs: {:?} plan:\n{}", + "{} adding task for stage {} partition group: {:?} partitions: {:?} stage_addrs: {:?} plan:\n{}", self.name, stage_id, + partition_group, partitions, stage_addrs, display_plan_with_partition_counts(&plan) @@ -184,24 +212,26 @@ impl DfRayProcessorHandler { let now = SystemTime::now(); - for partition in partitions.iter() { - let key = PlanKey { - query_id: query_id.clone(), - stage_id, - partition: *partition, - }; - { - let mut _guard = self.plans.write(); - if let Some(plan_vec) = _guard.get_mut(&key) { - plan_vec.push((now, ctx.clone(), plan.clone())); - } else { - _guard.insert(key.clone(), vec![(now, ctx.clone(), plan.clone())]); - } - trace!("{} added plan for plan key {:?}", self.name, key); - } + let task = Task { + ctx: ctx.clone(), + plan: plan.clone(), + partitions: HashSet::from_iter(partitions.clone()), + }; + + let key = StageKey { + query_id: query_id.clone(), + stage_id, + }; + { + let mut _guard = self.stages.write(); + let stage_tasks = _guard.entry(key.clone()).or_insert_with(|| StageTasks { + tasks: vec![], + insert_time: now, + }); + stage_tasks.tasks.push(task.clone()); + trace!("{} added task for stage key {:?}", self.name, key); } - debug!("{} plans held {:?}", self.name, self.plans.read().len()); Ok(()) } @@ -213,86 +243,297 @@ impl DfRayProcessorHandler { partition_group: Vec, ) -> Result { let mut ctx = get_ctx()?; + let host = Host { + addr: self.addr.clone(), + name: self.name.clone(), + }; add_ctx_extentions( &mut ctx, - &format!("{} stage:{} pg:{:?}", self.name, stage_id, partition_group), + &host, &query_id, + stage_id, stage_addrs.clone(), - Some(partition_group), + partition_group, )?; Ok(ctx) } - fn make_stream( + /// Retrieve the requested ctx and plan to execute. Also return a bool + /// indicating if this is the last partition for this plan + fn get_ctx_and_plan( &self, query_id: &str, stage_id: u64, partition: u64, - ) -> Result> + Send + 'static, Status> { - let key = PlanKey { + ) -> Result<(SessionContext, Arc, bool)> { + let stage_key = StageKey { query_id: query_id.to_string(), stage_id, - partition, }; - let (ctx, plan) = { - let mut _guard = self.plans.write(); - let (plan_key, mut plan_vec) = _guard - .remove_entry(&key) - .ok_or_else(|| { - Status::internal(format!( - "{}, No plan found for plan key {:?}", - self.name, key, - )) - }) - .inspect_err(|e| { - error!( - "{}, No plan found for plan key {:?},{e:?} have keys {:?}", - self.name, - key, - _guard.keys().map(|k| format!("{k:?}")) - ); - })?; + let (ctx, plan, last) = { + let mut _guard = self.stages.write(); + let stage_tasks = _guard.get_mut(&stage_key).context(format!( + "{}, No plan found for stage key{:?}", + self.name, stage_key, + ))?; trace!( - "{} found {} plans for plan key {:?}", + "{} found {} tasks for stage key {:?}", self.name, - plan_vec.len(), - plan_key + stage_tasks.tasks.len(), + stage_key ); - let (_insert_time, ctx, plan) = plan_vec.pop().expect("plan_vec should not be empty"); - if !plan_vec.is_empty() { - _guard.insert(plan_key, plan_vec); + + // of the tasks for this stage, find one that has this partition not yet consumed + let task = stage_tasks + .tasks + .iter_mut() + .find(|t| t.partitions.contains(&partition)) + .context(format!( + "{}, No task found for stage key {:?} partition {}", + self.name, stage_key, partition + ))?; + + // remove this partition from the list of partitions yet to be consumed + if !task.partitions.remove(&partition) { + // it should be in there because we just filtered for it + return Err(anyhow!("UNEXPECTED: partition {partition} not in plan parts").into()); } - (ctx, plan) + // finally, we return the ctx and plan to execute this task, as well as a + // bool indicating if this is the last partition for this task + ( + task.ctx.clone(), + task.plan.clone(), + task.partitions.is_empty(), + ) }; - trace!( - "make_stream for plan {}", - displayable(plan.as_ref()).indent(true) - ); + { + // some house keeping, if there are no more partitions left for this task, + // remove it from the stage tasks. + let mut _guard = self.stages.write(); + + let remove_it = _guard.get_mut(&stage_key).map(|stage_tasks| { + stage_tasks.tasks.retain(|t| !t.partitions.is_empty()); + trace!( + "remaining tasks: for stage {:?}: {:?}", + stage_key, + stage_tasks + ); + + // furthermore, if there are no more tasks left for this stage, + // remove the stage from the map + stage_tasks.tasks.is_empty() + }); - let task_ctx = ctx.task_ctx(); + if remove_it.unwrap_or(false) { + _guard.remove(&stage_key); + debug!("{} removed stage key {:?}", self.name, stage_key); + } + } + + Ok((ctx, plan, last)) + } - let ctx_name = task_ctx - .session_config() - .get_extension::() - .ok_or_else(|| { - Status::internal(format!("{}, CtxName not set in session config", self.name)) - })? - .0 - .clone(); + /// we want to send any additional FlightDataMetadata that we discover while + /// executing the plan further downstream to consumers of our stream. + /// + /// We may discover a FlightDataMetadata as an additional payload incoming + /// on streming requests. The [`RayStageReader`] will look for these and add + /// them to an extension on the context. + /// + /// Our job is to send all record batches from our stream over our flight response + /// but when the stream is exhausted, we'll send an additional message containing the + /// metadata. + /// + /// The reason we have to do it at the end is that the metadata may include an annodated + /// plan with metrics which will only be available after the stream has been + /// fully consumed. + fn make_stream( + &self, + ctx: SessionContext, + plan: Arc, + stage_id: u64, + partition: u64, + last_partition: bool, + ) -> Result { + let task_ctx = ctx.task_ctx(); + // the RecordBatchStream from our plan let stream = plan .execute(partition as usize, task_ctx) - .inspect_err(|e| error!("Could not get partition stream from plan {e:?}")) - .map(|s| reporting_stream(&format!("{ctx_name} s:{stage_id} p:{partition}"), s)) - .map_err(|e| Status::internal(format!("Could not get partition stream from plan {e}")))? + .inspect_err(|e| error!("Could not get partition stream from plan {e:#?}"))? + .inspect(|batch| { + trace!("producing maybe batch {:?}", batch); + }) .map_err(|e| FlightError::from_external_error(Box::new(e))); - info!("{} plans held {}", self.name, self.plans.read().len()); + info!("{} tasks held {}", self.name, self.stages.read().len()); + + fn find_analyze(plan: &dyn ExecutionPlan) -> Option<&DistributedAnalyzeExec> { + if let Some(target) = plan.as_any().downcast_ref::() { + Some(target) + } else { + for child in plan.children() { + if let Some(target) = find_analyze(child.as_ref()) { + return Some(target); + } + } + None + } + } + + let mut flight_data_stream = FlightDataEncoderBuilder::new().build(stream); + let name = self.name.clone(); + let host = Host { + addr: self.addr.clone(), + name: self.name.clone(), + }; + + #[allow(unused_assignments)] // clippy can't understand our assignment to done in the macro + let out_stream = async_stream::stream! { + let mut done = false; + while !done { + + match (done, last_partition, flight_data_stream.next().await) { + (false, false, None) => { + // we finished a partition, but still have more to do, do nothing + done = true; + } + (false, true, None) => { + // no more data in the last partition, so now we yield our additional FlightDataMetadata + debug!("stream exhausted, yielding FlightDataMetadata"); + let task_outputs = ctx.state().config() + .get_extension::() + .unwrap_or(Arc::new(CtxAnnotatedOutputs::default())) + .0 + .clone(); + + let partition_group = ctx.state().config() + .get_extension::() + .expect("CtxPartitionGroup to be set") + .0.clone(); + + + if let Some(analyze) = find_analyze(plan.as_ref()) { + let annotated_plan = analyze.annotated_plan(); + debug!("sending annotated plan: {}", annotated_plan); + + let output = AnnotatedTaskOutput { + plan: annotated_plan, + host: Some(host.clone()), + stage_id, + partition_group, + }; + task_outputs.lock().push(output); + } + + let meta = FlightDataMetadata { + annotated_task_outputs: Some(AnnotatedTaskOutputs { + outputs: task_outputs.lock().clone(), + }), + }; + + + let fake_batch = RecordBatch::new_empty(plan.schema()); + + + let mut fde = FlightDataEncoderBuilder::new() + //.with_schema(plan.schema()) + .with_metadata(meta.encode_to_vec().into()).build(futures::stream::once(async {Ok(fake_batch)})).map_err(|e| { + FlightError::from_external_error(Box::new(e)) + }); + + let flight_data = fde + .next() + .await + .ok_or_else(|| { + Status::internal(format!( + "{}, No FlightDataMetadata from our fde", + name + )) + })??; + + yield Ok(flight_data); + done = true; + + }, + (false, _, Some(Err(e))) => { + yield Err(Status::internal(format!( + "Unexpected error getting flight data stream: {e:?}", + ))); + done = true; + }, + (false, _, Some(Ok(flight_data))) => { + // we have a flight data, so we yield it + trace!("received normal flight data, yielding"); + yield Ok(flight_data); + }, + (true, false, _ ) => { + // we've finished a partition, do nothing + done = true; + }, + (true, true, _ ) => { + yield Err(Status::internal(format!("{name} reached expected unreachable block"))); + }, + } + } + }; + + Ok(Box::pin(out_stream)) + } + fn do_action_get_host(&self) -> Result, Status> { + let addr = self.addr.clone(); + let name = self.name.clone(); + + let out_stream = Box::pin(stream! { + yield Ok::<_, tonic::Status>(arrow_flight::Result { + body: Host{ + addr, + name, + }.encode_to_vec().into() + }); + }) as crate::flight::DoActionStream; + + Ok(Response::new(out_stream)) + } + + async fn do_action_add_plan( + &self, + action: Action, + ) -> Result, Status> { + let task_data = DDTask::decode(action.body.as_ref()).map_err(|e| { + Status::internal(format!( + "{}, Unexpected error decoding StageData: {e:?}", + self.name + )) + })?; + + let addrs = task_data + .stage_addrs + .as_ref() + .context("stage addrs not present") + .map_err(DFRayError::from) + .and_then(get_addrs) + .map_err(|e| Status::internal(format!("{}, {e}", self.name)))?; - Ok(stream) + self.add_task( + task_data.query_id, + task_data.stage_id, + addrs, + task_data.partition_group, + task_data.full_partitions, + &task_data.plan_bytes, + ) + .await + .map_err(|e| Status::internal(format!("{}, Could not add plan: {e:?}", self.name)))?; + + let out_stream = Box::pin(stream! { + yield Ok::<_, tonic::Status>(arrow_flight::Result::default()); + }) as crate::flight::DoActionStream; + + Ok(Response::new(out_stream)) } } @@ -316,31 +557,36 @@ impl FlightHandler for DfRayProcessorHandler { )) })?; - let plan_key = PlanKey { + let task_key = StageKey { query_id: ftd.query_id.clone(), stage_id: ftd.stage_id, - partition: ftd.partition, }; debug!( - "{}, request for plan_key:{:?} from: {},{}", - self.name, plan_key, ftd.requestor_name, remote_addr + "{}, request for task_key:{:?} partition: {} from: {},{}", + self.name, task_key, ftd.partition, ftd.requestor_name, remote_addr ); let name = self.name.clone(); - let stream = self - .make_stream(&ftd.query_id, ftd.stage_id, ftd.partition) + let (ctx, plan, is_last_partition) = self + .get_ctx_and_plan(&ftd.query_id, ftd.stage_id, ftd.partition) .map_err(|e| { - Status::internal(format!("{name} Unexpected error making stream {e:?}")) + Status::internal(format!( + "{name} Could not find task for query_id {} stage {} partition {}: {e:?}", + ftd.query_id, ftd.stage_id, ftd.partition + )) })?; - let out_stream = FlightDataEncoderBuilder::new() - .build(stream) - .map_err(move |e| { - Status::internal(format!("{name} Unexpected error building stream {e:?}")) - }); + let do_get_stream = self + .make_stream(ctx, plan, ftd.stage_id, ftd.partition, is_last_partition) + .map_err(|e| { + Status::internal(format!( + "{name} Could not make stream for query_id {} stage {} partition {}: {e:?}", + ftd.query_id, ftd.stage_id, ftd.partition + )) + })?; - Ok(Response::new(Box::pin(out_stream))) + Ok(Response::new(do_get_stream)) } async fn do_action( @@ -352,44 +598,16 @@ impl FlightHandler for DfRayProcessorHandler { let action_type = action.r#type.as_str(); trace!("{} received action: {}", self.name, action_type); - if action_type != "add_plan" { - return Err(Status::unimplemented(format!( + if action_type == "add_plan" { + self.do_action_add_plan(action).await + } else if action_type == "get_host" { + self.do_action_get_host() + } else { + Err(Status::unimplemented(format!( "{}, Unimplemented action: {}", self.name, action_type - ))); + ))) } - - let stage_data = StageData::decode(action.body.as_ref()).map_err(|e| { - Status::internal(format!( - "{}, Unexpected error decoding StageData: {e:?}", - self.name - )) - })?; - - let addrs = stage_data - .stage_addrs - .as_ref() - .context("stage addrs not present") - .map_err(DFRayError::from) - .and_then(get_addrs) - .map_err(|e| Status::internal(format!("{}, {e}", self.name)))?; - - self.add_plan( - stage_data.query_id, - stage_data.stage_id, - addrs, - stage_data.partition_group, - stage_data.full_partitions, - &stage_data.plan_bytes, - ) - .await - .map_err(|e| Status::internal(format!("{}, Could not add plan: {e:?}", self.name)))?; - - let out_stream = Box::pin(stream! { - yield Ok::<_, tonic::Status>(arrow_flight::Result::default()); - }) as crate::flight::DoActionStream; - - Ok(Response::new(out_stream)) } } @@ -400,62 +618,41 @@ impl FlightHandler for DfRayProcessorHandler { pub struct DFRayProcessorService { #[allow(dead_code)] name: String, - listener: Option, + listener: TcpListener, handler: Arc, - addr: Option, + addr: String, all_done_tx: Arc>>, all_done_rx: Option>, - port: usize, } impl DFRayProcessorService { - pub fn new(name: String, port: usize) -> Self { + pub async fn new(name: String, port: usize) -> Result { let name = format!("[{}]", name); - let listener = None; - let addr = None; let (all_done_tx, all_done_rx) = channel(1); let all_done_tx = Arc::new(Mutex::new(all_done_tx)); - let handler = Arc::new(DfRayProcessorHandler::new(name.clone())); + let listener = start_up(port).await?; - Self { + let addr = format!("{}", listener.local_addr().unwrap()); + + info!("DFRayProcessorService bound to {addr}"); + + let handler = Arc::new(DfRayProcessorHandler::new(name.clone(), addr.clone())); + + Ok(Self { name, listener, handler, addr, all_done_tx, all_done_rx: Some(all_done_rx), - port, - } + }) } - pub async fn start_up(&mut self) -> Result<()> { - let my_host_str = format!("0.0.0.0:{}", self.port); - - self.listener = TcpListener::bind(&my_host_str) - .await - .map(Some) - .context("Could not bind socket to {my_host_str}")?; - - self.addr = Some(format!( - "{}", - self.listener.as_ref().unwrap().local_addr().unwrap() - )); - - info!( - "DFRayProcessorService bound to {}", - self.addr.as_ref().unwrap() - ); - - Ok(()) - } /// get the address of the listing socket for this service - pub fn addr(&self) -> Result { - let addr = self.addr.clone().ok_or(anyhow!( - "DFRayProxyService not started yet, no address available" - ))?; - Ok(addr) + pub fn addr(&self) -> String { + self.addr.clone() } pub async fn all_done(&self) -> Result<()> { @@ -468,8 +665,8 @@ impl DFRayProcessorService { Ok(()) } - /// start the service - pub async fn serve(&mut self) -> Result<()> { + /// start the service, consuming self + pub async fn serve(mut self) -> Result<()> { let mut all_done_rx = self.all_done_rx.take().unwrap(); let signal = async move { @@ -486,12 +683,10 @@ impl DFRayProcessorService { let svc = FlightServiceServer::new(flight_serv); - let listener = self.listener.take().unwrap(); - Server::builder() .add_service(svc) .serve_with_incoming_shutdown( - tokio_stream::wrappers::TcpListenerStream::new(listener), + tokio_stream::wrappers::TcpListenerStream::new(self.listener), signal, ) .await diff --git a/src/proto/datafusion_ray.proto b/src/proto/distributed_datafusion.proto similarity index 59% rename from src/proto/datafusion_ray.proto rename to src/proto/distributed_datafusion.proto index bf4e243b..e6d4a8c3 100644 --- a/src/proto/datafusion_ray.proto +++ b/src/proto/distributed_datafusion.proto @@ -1,13 +1,13 @@ syntax = "proto3"; -package datafusion_ray.protobuf; +package distributed_datafusion.protobuf; import "datafusion.proto"; import "datafusion_common.proto"; option java_multiple_files = true; -option java_outer_classname = "RayDataFusionProto"; -option java_package = "datafusion_ray.protobuf"; +option java_outer_classname = "DistributedDataFusionProto"; +option java_package = "distributed_datafusion.protobuf"; message DFRayStageReaderExecNode { // schema of the stage we will consume @@ -18,37 +18,30 @@ message DFRayStageReaderExecNode { uint64 stage_id = 3; } -// the simplicity of the decoder in src/codec.rs currently requires a different byte -// representation per message. Hence the dummy fields. -// -// I'll come back to this and sort it out. Its not super critical as the plans are -// only exchanged at the start of queries, not during execution. - message MaxRowsExecNode { uint64 max_rows = 1; } message PartitionIsolatorExecNode { - float dummy = 1; uint64 partition_count = 2; } -message NumpangScanExecNode { - datafusion.FileScanExecConf base_conf = 1; +message DistributedAnalyzeExecNode { + // how much data to show + bool verbose = 1; + // if statistics should be displayed + bool show_statistics = 2; } -message ContextScanExecNode { - datafusion.FileScanExecConf base_conf = 1; - repeated datafusion.LogicalExprNode filters = 2; - datafusion_common.Schema schema = 3; +message DistributedAnalyzeRootExecNode { + // how much data to show + bool verbose = 1; + // if statistics should be displayed + bool show_statistics = 2; } -message DistributedExplainExecNode { - datafusion_common.Schema schema = 1; - string logical_plan = 2; - string physical_plan = 3; - string distributed_plan = 4; - string distributed_stages = 5; +message RecordBatchExecNode { + bytes batch = 1; } message DFRayExecNode { @@ -56,11 +49,31 @@ message DFRayExecNode { DFRayStageReaderExecNode stage_reader_exec = 1; MaxRowsExecNode max_rows_exec = 2; PartitionIsolatorExecNode isolator_exec = 3; - NumpangScanExecNode numpang_exec = 4; - ContextScanExecNode context_exec = 5; + DistributedAnalyzeExecNode distributed_analyze_exec = 4; + DistributedAnalyzeRootExecNode distributed_analyze_root_exec = 5; + RecordBatchExecNode record_batch_exec = 6; } } +message AnnotatedTaskOutput { + // the output of the explain analyze + string plan = 1; + // the host who executed this stage + Host host = 2; + // the stage id that was executed + uint64 stage_id = 3; + // the partitions that were executed by this stage + repeated uint64 partition_group = 4; +} + +message AnnotatedTaskOutputs { + repeated AnnotatedTaskOutput outputs = 1; +} + +message FlightDataMetadata { + AnnotatedTaskOutputs annotated_task_outputs = 1; +} + message FlightTicketData { string query_id = 1; uint64 stage_id = 2; @@ -78,8 +91,6 @@ message TicketStatementData { StageAddrs stage_addrs = 3; // the schema of the final stage datafusion_common.Schema schema = 4; - // For EXPLAIN queries, store the explain plan data directly - optional DistributedExplainExecNode explain_data = 5; } message Host { @@ -100,7 +111,7 @@ message StageAddrs { map stage_addrs = 1; } -message StageData { +message DDTask { string query_id = 1; uint64 stage_id = 2; @@ -117,5 +128,5 @@ message StageData { bool full_partitions = 8; - Host assigned_addr = 9; + Host assigned_host = 9; } diff --git a/src/proto/generated/protobuf.rs b/src/proto/generated/protobuf.rs index 270653a9..401aa103 100644 --- a/src/proto/generated/protobuf.rs +++ b/src/proto/generated/protobuf.rs @@ -18,45 +18,35 @@ pub struct MaxRowsExecNode { } #[derive(Clone, Copy, PartialEq, ::prost::Message)] pub struct PartitionIsolatorExecNode { - #[prost(float, tag = "1")] - pub dummy: f32, #[prost(uint64, tag = "2")] pub partition_count: u64, } -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct NumpangScanExecNode { - #[prost(message, optional, tag = "1")] - pub base_conf: ::core::option::Option< - ::datafusion_proto::protobuf::FileScanExecConf, - >, +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct DistributedAnalyzeExecNode { + /// how much data to show + #[prost(bool, tag = "1")] + pub verbose: bool, + /// if statistics should be displayed + #[prost(bool, tag = "2")] + pub show_statistics: bool, } -#[derive(Clone, PartialEq, ::prost::Message)] -pub struct ContextScanExecNode { - #[prost(message, optional, tag = "1")] - pub base_conf: ::core::option::Option< - ::datafusion_proto::protobuf::FileScanExecConf, - >, - #[prost(message, repeated, tag = "2")] - pub filters: ::prost::alloc::vec::Vec<::datafusion_proto::protobuf::LogicalExprNode>, - #[prost(message, optional, tag = "3")] - pub schema: ::core::option::Option<::datafusion_proto::protobuf::Schema>, +#[derive(Clone, Copy, PartialEq, ::prost::Message)] +pub struct DistributedAnalyzeRootExecNode { + /// how much data to show + #[prost(bool, tag = "1")] + pub verbose: bool, + /// if statistics should be displayed + #[prost(bool, tag = "2")] + pub show_statistics: bool, } #[derive(Clone, PartialEq, ::prost::Message)] -pub struct DistributedExplainExecNode { - #[prost(message, optional, tag = "1")] - pub schema: ::core::option::Option<::datafusion_proto::protobuf::Schema>, - #[prost(string, tag = "2")] - pub logical_plan: ::prost::alloc::string::String, - #[prost(string, tag = "3")] - pub physical_plan: ::prost::alloc::string::String, - #[prost(string, tag = "4")] - pub distributed_plan: ::prost::alloc::string::String, - #[prost(string, tag = "5")] - pub distributed_stages: ::prost::alloc::string::String, +pub struct RecordBatchExecNode { + #[prost(bytes = "vec", tag = "1")] + pub batch: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct DfRayExecNode { - #[prost(oneof = "df_ray_exec_node::Payload", tags = "1, 2, 3, 4, 5")] + #[prost(oneof = "df_ray_exec_node::Payload", tags = "1, 2, 3, 4, 5, 6")] pub payload: ::core::option::Option, } /// Nested message and enum types in `DFRayExecNode`. @@ -70,12 +60,39 @@ pub mod df_ray_exec_node { #[prost(message, tag = "3")] IsolatorExec(super::PartitionIsolatorExecNode), #[prost(message, tag = "4")] - NumpangExec(super::NumpangScanExecNode), + DistributedAnalyzeExec(super::DistributedAnalyzeExecNode), #[prost(message, tag = "5")] - ContextExec(super::ContextScanExecNode), + DistributedAnalyzeRootExec(super::DistributedAnalyzeRootExecNode), + #[prost(message, tag = "6")] + RecordBatchExec(super::RecordBatchExecNode), } } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct AnnotatedTaskOutput { + /// the output of the explain analyze + #[prost(string, tag = "1")] + pub plan: ::prost::alloc::string::String, + /// the host who executed this stage + #[prost(message, optional, tag = "2")] + pub host: ::core::option::Option, + /// the stage id that was executed + #[prost(uint64, tag = "3")] + pub stage_id: u64, + /// the partitions that were executed by this stage + #[prost(uint64, repeated, tag = "4")] + pub partition_group: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct AnnotatedTaskOutputs { + #[prost(message, repeated, tag = "1")] + pub outputs: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct FlightDataMetadata { + #[prost(message, optional, tag = "1")] + pub annotated_task_outputs: ::core::option::Option, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct FlightTicketData { #[prost(string, tag = "1")] pub query_id: ::prost::alloc::string::String, @@ -101,9 +118,6 @@ pub struct TicketStatementData { /// the schema of the final stage #[prost(message, optional, tag = "4")] pub schema: ::core::option::Option<::datafusion_proto::protobuf::Schema>, - /// For EXPLAIN queries, store the explain plan data directly - #[prost(message, optional, tag = "5")] - pub explain_data: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct Host { @@ -129,7 +143,7 @@ pub struct StageAddrs { pub stage_addrs: ::std::collections::HashMap, } #[derive(Clone, PartialEq, ::prost::Message)] -pub struct StageData { +pub struct DdTask { #[prost(string, tag = "1")] pub query_id: ::prost::alloc::string::String, #[prost(uint64, tag = "2")] @@ -147,5 +161,5 @@ pub struct StageData { #[prost(bool, tag = "8")] pub full_partitions: bool, #[prost(message, optional, tag = "9")] - pub assigned_addr: ::core::option::Option, + pub assigned_host: ::core::option::Option, } diff --git a/src/proto/mod.rs b/src/proto/mod.rs index 2602738d..a3ddc5fe 100644 --- a/src/proto/mod.rs +++ b/src/proto/mod.rs @@ -16,7 +16,7 @@ // under the License. use crate::{ - protobuf::{Host, Hosts, StageAddrs}, + protobuf::{Hosts, StageAddrs}, vocab::Addrs, }; @@ -29,10 +29,7 @@ impl From for StageAddrs { for (stage_id, partition_addrs) in value { for (part, part_addrs) in partition_addrs { - let host_addrs = part_addrs - .into_iter() - .map(|(name, addr)| Host { name, addr }) - .collect(); + let host_addrs = part_addrs.into_iter().collect(); stage_addrs .stage_addrs .entry(stage_id) diff --git a/src/proxy_service.rs b/src/proxy_service.rs index ad21edb5..e202b13a 100644 --- a/src/proxy_service.rs +++ b/src/proxy_service.rs @@ -17,12 +17,20 @@ use std::sync::Arc; -use anyhow::{anyhow, Context}; +use anyhow::Context; +use arrow::datatypes::Schema; use arrow_flight::{ - flight_service_server::FlightServiceServer, sql::TicketStatementQuery, FlightDescriptor, - FlightInfo, Ticket, + encode::FlightDataEncoderBuilder, + error::FlightError, + flight_service_server::FlightServiceServer, + sql::{ProstMessageExt, TicketStatementQuery}, + FlightDescriptor, FlightEndpoint, FlightInfo, Ticket, +}; +use datafusion::physical_plan::{ + coalesce_partitions::CoalescePartitionsExec, ExecutionPlan, Partitioning, }; use datafusion_substrait::substrait; +use futures::TryStreamExt; use parking_lot::Mutex; use prost::Message; use tokio::{ @@ -32,27 +40,122 @@ use tokio::{ use tonic::{async_trait, transport::Server, Request, Response, Status}; use crate::{ - explain::is_explain_query, flight::{FlightSqlHandler, FlightSqlServ}, - flight_handlers::FlightRequestHandler, - k8s::get_worker_addresses, logging::{debug, info, trace}, + planning::{add_ctx_extentions, get_ctx}, protobuf::TicketStatementData, - query_planner::QueryPlanner, + query_planner::{QueryPlan, QueryPlanner}, result::Result, + stage_reader::DFRayStageReaderExec, + util::{display_plan_with_partition_counts, get_addrs, start_up}, + vocab::{Addrs, Host}, + worker_discovery::get_worker_addresses, }; pub struct DfRayProxyHandler { - pub flight_handler: FlightRequestHandler, + /// our host info, useful for logging + pub host: Host, + + pub planner: QueryPlanner, } impl DfRayProxyHandler { - pub fn new() -> Self { + pub fn new(name: String, addr: String) -> Self { // call this function to bootstrap the worker discovery mechanism get_worker_addresses().expect("Could not get worker addresses upon startup"); + + let host = Host { + name: name.clone(), + addr: addr.clone(), + }; Self { - flight_handler: FlightRequestHandler::new(QueryPlanner::new()), + host: host.clone(), + planner: QueryPlanner::new(), + } + } + + pub fn create_flight_info_response(&self, query_plan: QueryPlan) -> Result { + let mut flight_info = FlightInfo::new() + .try_with_schema(&query_plan.schema) + .map_err(|e| Status::internal(format!("Could not create flight info {e:?}")))?; + + let ticket_data = TicketStatementData { + query_id: query_plan.query_id, + stage_id: query_plan.final_stage_id, + stage_addrs: Some(query_plan.worker_addresses.into()), + schema: Some( + query_plan + .schema + .try_into() + .map_err(|e| Status::internal(format!("Could not convert schema {e:?}")))?, + ), + }; + + let ticket = Ticket::new( + TicketStatementQuery { + statement_handle: ticket_data.encode_to_vec().into(), + } + .as_any() + .encode_to_vec(), + ); + + let endpoint = FlightEndpoint::new().with_ticket(ticket); + flight_info = flight_info.with_endpoint(endpoint); + + Ok(flight_info) + } + + pub async fn execute_plan_and_build_stream( + &self, + plan: Arc, + query_id: String, + stage_id: u64, + addrs: Addrs, + ) -> Result, Status> { + let mut ctx = + get_ctx().map_err(|e| Status::internal(format!("Could not create context {e:?}")))?; + + add_ctx_extentions(&mut ctx, &self.host, &query_id, stage_id, addrs, vec![]) + .map_err(|e| Status::internal(format!("Could not add context extensions {e:?}")))?; + + // TODO: revisit this to allow for consuming a partitular partition + trace!("calling execute plan"); + let partition = 0; + let stream = plan + .execute(partition, ctx.task_ctx()) + .map_err(|e| { + Status::internal(format!( + "Error executing plan for query_id {} partition {}: {e:?}", + query_id, partition + )) + })? + .map_err(|e| FlightError::ExternalError(Box::new(e))); + + let out_stream = FlightDataEncoderBuilder::new() + .build(stream) + .map_err(move |e| Status::internal(format!("Unexpected error building stream {e:?}"))); + + Ok(Response::new(Box::pin(out_stream))) + } + + pub fn validate_single_stage_addrs( + &self, + addrs: &Addrs, + expected_stage_id: u64, + ) -> Result<(), Status> { + if addrs.len() != 1 { + return Err(Status::internal(format!( + "Expected exactly one stage in addrs, got {}", + addrs.len() + ))); + } + if !addrs.contains_key(&expected_stage_id) { + return Err(Status::internal(format!( + "No addresses found for stage_id {} in addrs", + expected_stage_id + ))); } + Ok(()) } } @@ -63,20 +166,44 @@ impl FlightSqlHandler for DfRayProxyHandler { query: arrow_flight::sql::CommandStatementQuery, _request: Request, ) -> Result, Status> { - let is_explain = is_explain_query(&query.query); + let query_plan = self + .planner + .prepare(&query.query) + .await + .map_err(|e| Status::internal(format!("Could not prepare query {e:?}")))?; - if is_explain { - self.flight_handler - .handle_explain_request(&query.query) - .await - } else { - self.flight_handler.handle_query_request(&query.query).await - } + self.create_flight_info_response(query_plan) + .map(Response::new) + .context("Could not create flight info response") + .map_err(|e| Status::internal(format!("Error creating flight info: {e:?}"))) + } + + async fn get_flight_info_substrait_plan( + &self, + substrait: arrow_flight::sql::CommandStatementSubstraitPlan, + _request: Request, + ) -> Result, Status> { + let plan = match &substrait.plan { + Some(substrait_plan) => substrait::proto::Plan::decode(substrait_plan.plan.as_ref()) + .map_err(|e| Status::invalid_argument(format!("Invalid Substrait plan: {e}")))?, + None => return Err(Status::invalid_argument("Missing Substrait plan")), + }; + + let query_plan = self + .planner + .prepare_substrait(plan) + .await + .map_err(|e| Status::internal(format!("Could not prepare query {e:?}")))?; + + self.create_flight_info_response(query_plan) + .map(Response::new) + .context("Could not create flight info response") + .map_err(|e| Status::internal(format!("Error creating flight info: {e:?}"))) } async fn do_get_statement( &self, - ticket: TicketStatementQuery, + ticket: arrow_flight::sql::TicketStatementQuery, request: Request, ) -> Result, Status> { trace!("do_get_statement"); @@ -90,30 +217,53 @@ impl FlightSqlHandler for DfRayProxyHandler { debug!("request for ticket: {:?} from {}", tsd, remote_addr); - if tsd.explain_data.is_some() { - self.flight_handler - .handle_explain_statement_execution(tsd, &remote_addr) - .await - } else { - self.flight_handler - .handle_regular_statement_execution(tsd, &remote_addr) - .await - } - } - - async fn get_flight_info_substrait_plan( - &self, - substrait: arrow_flight::sql::CommandStatementSubstraitPlan, - _request: Request, - ) -> Result, Status> { - let plan = match &substrait.plan { - Some(substrait_plan) => substrait::proto::Plan::decode(substrait_plan.plan.as_ref()) - .map_err(|e| Status::invalid_argument(format!("Invalid Substrait plan: {e}")))?, - None => return Err(Status::invalid_argument("Missing Substrait plan")), - }; + let schema: Schema = tsd + .schema + .as_ref() + .ok_or_else(|| Status::internal("No schema in TicketStatementData"))? + .try_into() + .map_err(|e| Status::internal(format!("Cannot convert schema {e}")))?; + + // Create an Addrs from the final stage information in the tsd + let stage_addrs = tsd.stage_addrs.ok_or_else(|| { + Status::internal("No stages_addrs in TicketStatementData, cannot proceed") + })?; + + let addrs: Addrs = get_addrs(&stage_addrs).map_err(|e| { + Status::internal(format!("Cannot get addresses from stage_addrs {e:?}")) + })?; + + trace!("calculated addrs: {:?}", addrs); + + // Validate that addrs contains exactly one stage + self.validate_single_stage_addrs(&addrs, tsd.stage_id)?; + + let stage_partition_addrs = addrs.get(&tsd.stage_id).ok_or_else(|| { + Status::internal(format!( + "No partition addresses found for stage_id {}", + tsd.stage_id + )) + })?; + + let plan = Arc::new( + DFRayStageReaderExec::try_new( + Partitioning::UnknownPartitioning(stage_partition_addrs.len()), + Arc::new(schema), + tsd.stage_id, + ) + // TODO: revisit this to allow for consuming a particular partition + .map(|stg| CoalescePartitionsExec::new(Arc::new(stg))) + .map_err(|e| Status::internal(format!("Unexpected error {e}")))?, + ) as Arc; + + debug!( + "request for query_id {} from {} reader plan:\n{}", + tsd.query_id, + remote_addr, + display_plan_with_partition_counts(&plan) + ); - self.flight_handler - .handle_substrait_info_request(plan) + self.execute_plan_and_build_stream(plan, tsd.query_id, tsd.stage_id, addrs) .await } } @@ -123,59 +273,40 @@ impl FlightSqlHandler for DfRayProxyHandler { /// /// It only responds to the DoGet Arrow Flight method pub struct DFRayProxyService { - listener: Option, + listener: TcpListener, handler: Arc, - addr: Option, + addr: String, all_done_tx: Arc>>, all_done_rx: Option>, - port: usize, } impl DFRayProxyService { - pub fn new(port: usize) -> Self { + pub async fn new(name: String, port: usize) -> Result { debug!("Creating DFRayProxyService!"); - let listener = None; - let addr = None; let (all_done_tx, all_done_rx) = channel(1); let all_done_tx = Arc::new(Mutex::new(all_done_tx)); - let handler = Arc::new(DfRayProxyHandler::new()); + let listener = start_up(port).await?; - Self { + let addr = format!("{}", listener.local_addr().unwrap()); + + info!("DFRayProcessorService bound to {addr}"); + + let handler = Arc::new(DfRayProxyHandler::new(name, addr.clone())); + + Ok(Self { listener, handler, addr, all_done_tx, all_done_rx: Some(all_done_rx), - port, - } - } - - pub async fn start_up(&mut self) -> Result<()> { - let my_host_str = format!("0.0.0.0:{}", self.port); - - self.listener = TcpListener::bind(&my_host_str) - .await - .map(Some) - .context("Could not bind socket to {my_host_str}")?; - - self.addr = Some(format!( - "{}", - self.listener.as_ref().unwrap().local_addr().unwrap() - )); - - info!("DFRayProxyService bound to {}", self.addr.as_ref().unwrap()); - - Ok(()) + }) } /// get the address of the listing socket for this service pub fn addr(&self) -> Result { - let addr = self.addr.clone().ok_or(anyhow!( - "DFRayProxyService not started yet, no address available" - ))?; - Ok(addr) + Ok(self.addr.clone()) } pub async fn all_done(&self) -> Result<()> { @@ -189,7 +320,7 @@ impl DFRayProxyService { } /// start the service - pub async fn serve(&mut self) -> Result<()> { + pub async fn serve(mut self) -> Result<()> { let mut all_done_rx = self.all_done_rx.take().unwrap(); let signal = async move { @@ -212,12 +343,10 @@ impl DFRayProxyService { //let svc = FlightServiceServer::new(service); let svc = FlightServiceServer::with_interceptor(service, intercept); - let listener = self.listener.take().unwrap(); - Server::builder() .add_service(svc) .serve_with_incoming_shutdown( - tokio_stream::wrappers::TcpListenerStream::new(listener), + tokio_stream::wrappers::TcpListenerStream::new(self.listener), signal, ) .await @@ -225,142 +354,3 @@ impl DFRayProxyService { Ok(()) } } - -#[cfg(test)] -mod tests { - use super::*; - // Test-specific imports - use arrow_flight::{ - sql::{CommandStatementQuery, TicketStatementQuery}, - FlightDescriptor, Ticket, - }; - use prost::Message; - use tonic::Request; - - use crate::test_utils::explain_test_helpers::{ - create_explain_ticket_statement_data, create_test_proxy_handler, - verify_explain_stream_results, - }; - - #[tokio::test] - async fn test_get_flight_info_statement_explain() { - let handler = create_test_proxy_handler(); - - // Test EXPLAIN query - let command = CommandStatementQuery { - query: "EXPLAIN SELECT 1 as test_col".to_string(), - transaction_id: None, - }; - - let request = Request::new(FlightDescriptor::new_cmd(vec![])); - let result = handler.get_flight_info_statement(command, request).await; - - assert!(result.is_ok()); - let response = result.unwrap(); - let flight_info = response.into_inner(); - - // Verify FlightInfo structure - assert!(!flight_info.schema.is_empty()); - assert_eq!(flight_info.endpoint.len(), 1); - assert!(flight_info.endpoint[0].ticket.is_some()); - - // Verify that ticket has content (encoded TicketStatementData) - let ticket = flight_info.endpoint[0].ticket.as_ref().unwrap(); - assert!(!ticket.ticket.is_empty()); - - println!( - "✓ FlightInfo created successfully with {} schema bytes and ticket with {} bytes", - flight_info.schema.len(), - ticket.ticket.len() - ); - } - - #[tokio::test] - async fn test_do_get_statement_explain() { - let handler = create_test_proxy_handler(); - - // First prepare an EXPLAIN query to get proper ticket data - let query = "EXPLAIN SELECT 1 as test_col"; - let plans = handler - .flight_handler - .planner - .prepare_explain(query) - .await - .unwrap(); - - let tsd = create_explain_ticket_statement_data(plans); - - // Create the ticket - let ticket_query = TicketStatementQuery { - statement_handle: tsd.encode_to_vec().into(), - }; - - let request = Request::new(Ticket::new(vec![])); - let result = handler.do_get_statement(ticket_query, request).await; - - assert!(result.is_ok()); - let response = result.unwrap(); - let stream = response.into_inner(); - - // Use shared verification function - verify_explain_stream_results(stream).await; - } - - #[tokio::test] - async fn test_compare_explain_flight_info_responses() { - let handler = create_test_proxy_handler(); - let query = "EXPLAIN SELECT 1 as test_col"; - - // Get FlightInfo from handle_explain_request - let result1 = handler - .flight_handler - .handle_explain_request(query) - .await - .unwrap(); - let flight_info1 = result1.into_inner(); - - // Get FlightInfo from get_flight_info_statement - let command = CommandStatementQuery { - query: query.to_string(), - transaction_id: None, - }; - let request = Request::new(FlightDescriptor::new_cmd(vec![])); - let result2 = handler - .get_flight_info_statement(command, request) - .await - .unwrap(); - let flight_info2 = result2.into_inner(); - - // Compare FlightInfo responses (structure should be identical) - assert_eq!(flight_info1.schema.len(), flight_info2.schema.len()); // Same schema size - assert_eq!(flight_info1.endpoint.len(), flight_info2.endpoint.len()); // Same number of endpoints - assert_eq!(flight_info1.endpoint.len(), 1); // Both should have exactly one endpoint - - // Both should have tickets with content - let ticket1 = flight_info1.endpoint[0].ticket.as_ref().unwrap(); - let ticket2 = flight_info2.endpoint[0].ticket.as_ref().unwrap(); - assert!(!ticket1.ticket.is_empty()); - assert!(!ticket2.ticket.is_empty()); - - println!("✓ Both tests produce FlightInfo with identical structure:"); - println!( - " - Schema bytes: {} vs {}", - flight_info1.schema.len(), - flight_info2.schema.len() - ); - println!( - " - Endpoints: {} vs {}", - flight_info1.endpoint.len(), - flight_info2.endpoint.len() - ); - println!( - " - Ticket bytes: {} vs {}", - ticket1.ticket.len(), - ticket2.ticket.len() - ); - } - - // TODO: Add tests for regular (non-explain) queries - // We might need to create integration or end-to-end test infrastructure for this because - // they need workers -} diff --git a/src/query_planner.rs b/src/query_planner.rs index 035e9e57..14ecc783 100644 --- a/src/query_planner.rs +++ b/src/query_planner.rs @@ -1,38 +1,43 @@ use std::sync::Arc; use anyhow::anyhow; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::{compute::concat_batches, datatypes::SchemaRef}; use datafusion::{ - logical_expr::LogicalPlan, physical_plan::ExecutionPlan, prelude::SessionContext, + logical_expr::LogicalPlan, + physical_plan::{coalesce_partitions::CoalescePartitionsExec, ExecutionPlan}, + prelude::SessionContext, }; -use datafusion_substrait::logical_plan::consumer::from_substrait_plan; -use datafusion_substrait::substrait::proto::Plan; + +use datafusion_substrait::{logical_plan::consumer::from_substrait_plan, substrait::proto::Plan}; +use tokio_stream::StreamExt; use crate::{ - explain::{is_explain_query, DistributedExplainExec}, - k8s::get_worker_addresses, - logging::debug, + explain::build_explain_batch, planning::{ distribute_stages, execution_planning, get_ctx, logical_planning, physical_planning, - DFRayStage, }, + record_batch_exec::RecordBatchExec, result::Result, - vocab::Addrs, + vocab::{Addrs, DDTask}, + worker_discovery::get_worker_addresses, }; -/// Result of base query preparation containing all planning artifacts for both query and its EXPLAIN -pub struct QueryPlanBase { +/// Result of query preparation for execution of both query and its EXPLAIN +pub struct QueryPlan { pub query_id: String, pub session_context: SessionContext, + pub worker_addresses: Addrs, + pub final_stage_id: u64, + pub schema: SchemaRef, pub logical_plan: LogicalPlan, pub physical_plan: Arc, pub distributed_plan: Arc, - pub distributed_stages: Vec, + pub distributed_tasks: Vec, } -impl std::fmt::Debug for QueryPlanBase { +impl std::fmt::Debug for QueryPlan { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("QueryPlanBase") + f.debug_struct("QueryPlan") .field("query_id", &self.query_id) .field("session_context", &"") .field("logical_plan", &self.logical_plan) @@ -43,415 +48,158 @@ impl std::fmt::Debug for QueryPlanBase { ) .field( "stages", - &format!("<{} stages>", self.distributed_stages.len()), + &format!("<{} stages>", self.distributed_tasks.len()), ) .finish() } } -/// Result of query preparation for execution of both query and its EXPLAIN -#[derive(Debug)] -pub struct QueryPlan { - pub query_id: String, - pub worker_addresses: Addrs, - pub final_stage_id: u64, - pub schema: SchemaRef, - pub explain_data: Option, -} - /// Query planner responsible for preparing SQL queries for distributed execution pub struct QueryPlanner; +impl Default for QueryPlanner { + fn default() -> Self { + Self::new() + } +} + impl QueryPlanner { pub fn new() -> Self { Self } - /// Dispatch a distributed query plan to the workers. - async fn dispatch_query_plan(&self, base_result: QueryPlanBase) -> Result { - if base_result.distributed_stages.is_empty() { - return Err(anyhow!("No stages generated for query").into()); - } - - let worker_addrs = get_worker_addresses()?; - - // The last stage produces the data returned to the client. - let final_stage = &base_result.distributed_stages[base_result.distributed_stages.len() - 1]; - let schema = Arc::clone(&final_stage.plan.schema()); - let final_stage_id = final_stage.stage_id; - - // Physically dispatch each stage to the worker pool, further dividing - // them into partition groups. - let final_workers = distribute_stages( - &base_result.query_id, - base_result.distributed_stages, - worker_addrs, - ) - .await?; - - Ok(QueryPlan { - query_id: base_result.query_id, - worker_addresses: final_workers, - final_stage_id, - schema, - explain_data: None, - }) - } - /// Common planning steps shared by both query and its EXPLAIN /// /// Prepare a query by parsing the SQL, planning it, and distributing the /// physical plan into stages that can be executed by workers. - pub async fn prepare_query_base(&self, sql: &str, query_type: &str) -> Result { - debug!("prepare_query_base: {} SQL = {}", query_type, sql); - - let query_id = uuid::Uuid::new_v4().to_string(); + pub async fn prepare(&self, sql: &str) -> Result { let ctx = get_ctx().map_err(|e| anyhow!("Could not create context: {e}"))?; let logical_plan = logical_planning(sql, &ctx).await?; - let physical_plan = physical_planning(&logical_plan, &ctx).await?; - // divide the physical plan into chunks (stages) that we can distribute to workers later in dispatch_query_plan - let (distributed_plan, distributed_stages) = - execution_planning(physical_plan.clone(), 8192, Some(2)).await?; - - Ok(QueryPlanBase { - query_id, - session_context: ctx, - logical_plan, - physical_plan, - distributed_plan, - distributed_stages, - }) + match logical_plan { + p @ LogicalPlan::Explain(_) => self.prepare_explain(p, ctx).await, + // add other logical plans for local execution here following the pattern for explain + p @ LogicalPlan::DescribeTable(_) => self.prepare_local(p, ctx).await, + p => self.prepare_query(p, ctx).await, + } } - /// Prepare a distributed query (SQL entry point) - pub async fn prepare_query(&self, sql: &str) -> Result { - let base_result = self.prepare_query_base(sql, "REGULAR").await?; - self.dispatch_query_plan(base_result).await - } + pub async fn prepare_substrait(&self, substrait_plan: Plan) -> Result { + let ctx = get_ctx().map_err(|e| anyhow!("Could not create context: {e}"))?; - /// Prepare a distributed query (Substrait entry point) - pub async fn prepare_substrait_query(&self, substrait_plan: Plan) -> Result { - let base_result = self - .prepare_substrait_query_base(substrait_plan, "SUBSTRAIT") - .await?; - self.dispatch_query_plan(base_result).await + let logical_plan = from_substrait_plan(&ctx.state(), &substrait_plan).await?; + + match logical_plan { + p @ LogicalPlan::Explain(_) => self.prepare_explain(p, ctx).await, + // add other logical plans for local execution here following the pattern for explain + p @ LogicalPlan::DescribeTable(_) => self.prepare_local(p, ctx).await, + p => self.prepare_query(p, ctx).await, + } } - pub async fn prepare_substrait_query_base( + async fn prepare_query( &self, - substrait_plan: Plan, - query_type: &str, - ) -> Result { - debug!( - "prepare_substrait_query_base: {} Substrait = {:#?}", - query_type, substrait_plan - ); - - let query_id = uuid::Uuid::new_v4().to_string(); - let ctx = get_ctx().map_err(|e| anyhow!("Could not create context: {e}"))?; + logical_plan: LogicalPlan, + ctx: SessionContext, + ) -> Result { + let physical_plan = physical_planning(&logical_plan, &ctx).await?; - let logical_plan = from_substrait_plan(&ctx.state(), &substrait_plan) - .await - .map_err(|e| anyhow!("Failed to convert DataFusion Logical Plan: {e}"))?; + self.send_it(logical_plan, physical_plan, ctx).await + } - let physical_plan = physical_planning(&logical_plan, &ctx) - .await - .map_err(|e| anyhow!("Failed to convert DataFusion Physical Plan: {e}"))?; + async fn prepare_local( + &self, + logical_plan: LogicalPlan, + ctx: SessionContext, + ) -> Result { + let physical_plan = physical_planning(&logical_plan, &ctx).await?; - // divide the physical plan into chunks (stages) that we can distribute to workers later in dispatch_query_plan - let (distributed_plan, distributed_stages) = - execution_planning(physical_plan.clone(), 8192, Some(2)).await?; + // execute it locally + let mut stream = + Arc::new(CoalescePartitionsExec::new(physical_plan)).execute(0, ctx.task_ctx())?; + let mut batches = vec![]; - Ok(QueryPlanBase { - query_id, - session_context: ctx, - logical_plan, - physical_plan, - distributed_plan, - distributed_stages, - }) - } + while let Some(batch) = stream.next().await { + batches.push(batch?); + } - /// Prepare an EXPLAIN query - /// This method only handles EXPLAIN queries (plan only). EXPLAIN ANALYZE queries are handled as regular queries because they need to be executed. - pub async fn prepare_explain(&self, sql: &str) -> Result { - // Validate that this is actually an EXPLAIN query (not EXPLAIN ANALYZE) - if !is_explain_query(sql) { - return Err(anyhow!( - "prepare_explain called with non-EXPLAIN query or EXPLAIN ANALYZE query: {}", - sql - ) - .into()); + if batches.is_empty() { + return Err(anyhow!("No data returned from local execution").into()); } - // Extract the underlying query from the EXPLAIN statement - let underlying_query = sql - .trim() - .strip_prefix("EXPLAIN") - .or_else(|| { - sql.trim() - .to_uppercase() - .strip_prefix("EXPLAIN") - .map(|_| &sql.trim()[7..]) - }) - .unwrap_or(sql) - .trim(); - - let base_result = self.prepare_query_base(underlying_query, "EXPLAIN").await?; - - // generate the plan strings - let logical_plan_string = format!("{}", base_result.logical_plan.display_indent()); - let physical_plan_string = format!( - "{}", - datafusion::physical_plan::displayable(base_result.physical_plan.as_ref()).indent(true) - ); - let distributed_plan_string = format!( - "{}", - datafusion::physical_plan::displayable(base_result.distributed_plan.as_ref()) - .indent(true) - ); - let distributed_stages_string = DistributedExplainExec::format_distributed_stages( - base_result.distributed_stages.as_slice(), - ); - - // create the schema for EXPLAIN results - let schema = Arc::new(Schema::new(vec![ - Field::new("plan_type", DataType::Utf8, false), - Field::new("plan", DataType::Utf8, false), - ])); - - // Create explain data - let explain_data = DistributedExplainExec::new( - Arc::clone(&schema), - logical_plan_string, - physical_plan_string, - distributed_plan_string, - distributed_stages_string, - ); - - // Create dummy addresses for EXPLAIN (no real workers needed) - let mut dummy_addrs = std::collections::HashMap::new(); - let mut partition_addrs = std::collections::HashMap::new(); - partition_addrs.insert( - 0u64, - vec![("explain_local".to_string(), "local".to_string())], - ); - dummy_addrs.insert(0u64, partition_addrs); - - Ok(QueryPlan { - query_id: base_result.query_id, - worker_addresses: dummy_addrs, - final_stage_id: 0, - schema, - explain_data: Some(explain_data), - }) - } -} + let combined_batch = concat_batches(&batches[0].schema(), &batches)?; + let physical_plan = Arc::new(RecordBatchExec::new(combined_batch)); -#[cfg(test)] -mod tests { - use super::*; - use std::io::BufReader; - use std::{fs::File, path::Path}; - - // ////////////////////////////////////////////////////////////// - // Test helper functions - // ////////////////////////////////////////////////////////////// - - /// Set up mock worker environment for testing - fn setup_mock_worker_env() { - let mock_addrs = [ - ("mock_worker_1".to_string(), "localhost:9001".to_string()), - ("mock_worker_2".to_string(), "localhost:9002".to_string()), - ]; - let mock_env_value = mock_addrs - .iter() - .map(|(name, addr)| format!("{}/{}", name, addr)) - .collect::>() - .join(","); - std::env::set_var("DFRAY_WORKER_ADDRESSES", &mock_env_value); + self.send_it(logical_plan, physical_plan, ctx).await } - // ////////////////////////////////////////////////////////////// - // Core function tests - // ////////////////////////////////////////////////////////////// - - #[tokio::test] - async fn test_prepare_query_base() { - let planner = QueryPlanner::new(); - - // Test with a simple SELECT query without the need to read any table - let sql = "SELECT 1 as test_col"; - let result = planner.prepare_query_base(sql, "TEST").await; - - if result.is_ok() { - let query_plan_base = result.unwrap(); - // verify all fields have values - assert!(!query_plan_base.query_id.is_empty()); - assert!(!query_plan_base.distributed_stages.is_empty()); - assert!(!query_plan_base.physical_plan.schema().fields().is_empty()); - // logical plan of select 1 on empty relation - assert_eq!( - query_plan_base.logical_plan.to_string(), - "Projection: Int64(1) AS test_col\n EmptyRelation" - ); - // physical plan of select 1 on empty releation is ProjectionExec - assert_eq!(query_plan_base.physical_plan.name(), "ProjectionExec"); - } else { - // If worker discovery fails, we expect a specific error - let error_msg = format!("{:?}", result.unwrap_err()); - assert!(error_msg.contains("worker") || error_msg.contains("address")); + async fn prepare_explain( + &self, + explain_plan: LogicalPlan, + ctx: SessionContext, + ) -> Result { + let child_plan = explain_plan.inputs(); + if child_plan.len() != 1 { + return Err(anyhow!("EXPLAIN plan must have exactly one child").into()); } - } - #[tokio::test] - async fn test_prepare_substrait_query_base() { - let planner = QueryPlanner::new(); + let logical_plan = child_plan[0]; + + let query_plan = self.prepare_query(logical_plan.clone(), ctx).await?; + + let batch = build_explain_batch( + &query_plan.logical_plan, + &query_plan.physical_plan, + &query_plan.distributed_plan, + &query_plan.distributed_tasks, + )?; + let physical_plan = Arc::new(RecordBatchExec::new(batch)); - // Read the JSON plan and convert to binary Substrait protobuf bytes - let plan = serde_json::from_reader::<_, datafusion_substrait::substrait::proto::Plan>( - BufReader::new( - File::open(Path::new("testdata/substrait/select_one.substrait.json")) - .expect("file not found"), - ), + self.send_it( + query_plan.logical_plan, + physical_plan, + query_plan.session_context, ) - .expect("failed to parse json"); - - let result = planner.prepare_substrait_query_base(plan, "TEST").await; - - if result.is_ok() { - let query_plan_base = result.unwrap(); - // verify all fields have values - assert!(!query_plan_base.query_id.is_empty()); - assert!(!query_plan_base.distributed_stages.is_empty()); - assert!(!query_plan_base.physical_plan.schema().fields().is_empty()); - // logical plan of select 1 on empty relation - assert_eq!( - query_plan_base.logical_plan.to_string(), - "Projection: Int64(1) AS test_col\n Values: (Int64(0))" - ); - // physical plan of select 1 on empty releation is ProjectionExec - assert_eq!(query_plan_base.physical_plan.name(), "ProjectionExec"); - } else { - // If worker discovery fails, we expect a specific error - let error_msg = format!("{:?}", result.unwrap_err()); - assert!(error_msg.contains("worker") || error_msg.contains("address")); - } + .await } + async fn send_it( + &self, + logical_plan: LogicalPlan, + physical_plan: Arc, + ctx: SessionContext, + ) -> Result { + let query_id = uuid::Uuid::new_v4().to_string(); - #[tokio::test] - async fn test_prepare_explain() { - let planner = QueryPlanner::new(); - - // Test with a simple EXPLAIN query - let sql = "EXPLAIN SELECT 1 as test_col"; - let query_plan = planner.prepare_explain(sql).await.unwrap(); - - // EXPLAIN queries should work even without worker discovery since they use dummy addresses - assert!(!query_plan.query_id.is_empty()); - assert_eq!(query_plan.final_stage_id, 0); - assert!(query_plan.explain_data.is_some()); - - // Verify content of the explain data - let explain_data = query_plan.explain_data.unwrap(); - assert_eq!( - explain_data.logical_plan(), - "Projection: Int64(1) AS test_col\n EmptyRelation" - ); - assert_eq!( - explain_data.physical_plan(), - "ProjectionExec: expr=[1 as test_col]\n PlaceholderRowExec\n" - ); - assert_eq!(explain_data.distributed_plan(), - "RayStageExec[0] (output_partitioning=UnknownPartitioning(1))\n ProjectionExec: expr=[1 as test_col]\n PlaceholderRowExec\n"); - assert_eq!(explain_data.distributed_stages(), - "Stage 0:\n Partition Groups: [[0]]\n Full Partitions: false\n Plan:\n MaxRowsExec[max_rows=8192]\n CoalesceBatchesExec: target_batch_size=8192\n ProjectionExec: expr=[1 as test_col]\n PlaceholderRowExec\n"); - - // Should have explain schema (plan_type, plan columns) - assert_eq!(query_plan.schema.fields().len(), 2); - assert_eq!(query_plan.schema.field(0).name(), "plan_type"); - assert_eq!(query_plan.schema.field(1).name(), "plan"); - println!("✓ prepare_explain_query succeeded with proper structure"); - } + // divide the physical plan into chunks (tasks) that we can distribute to workers + let (distributed_plan, distributed_stages) = + execution_planning(physical_plan.clone(), 8192, Some(2)).await?; - #[tokio::test] - async fn test_prepare_explain_invalid_input() { - let planner = QueryPlanner::new(); - - // Test with EXPLAIN ANALYZE (should fail) - let sql = "EXPLAIN ANALYZE SELECT 1"; - let result = planner.prepare_explain(sql).await; - assert!(result.is_err()); - let error_msg = format!("{:?}", result.unwrap_err()); - assert!(error_msg.contains("prepare_explain called with non-EXPLAIN query")); - - // Test with non-EXPLAIN query (should fail) - let sql = "SELECT 1"; - let result = planner.prepare_explain(sql).await; - assert!(result.is_err()); - let error_msg = format!("{:?}", result.unwrap_err()); - assert!(error_msg.contains("prepare_explain called with non-EXPLAIN query")); - } + let worker_addrs = get_worker_addresses()?; - // NOTE: This test is ignored because prepare_query() requires actual worker communication. - // - // 🔍 Root Cause Analysis: - // The issue is NOT with mock worker setup - that works perfectly. The problem is in the - // distribute_stages() retry logic: - // - // 1. ✅ Mock workers are set up correctly: ["mock_worker_1/localhost:9001", "mock_worker_2/localhost:9002"] - // 2. ✅ get_worker_addresses() successfully returns: [("mock_worker_1", "localhost:9001"), ("mock_worker_2", "localhost:9002")] - // 3. ✅ distribute_stages() receives workers and creates HashMap: {"mock_worker_1": "localhost:9001", "mock_worker_2": "localhost:9002"} - // 4. ❌ try_distribute_stages() attempts to create Flight client connections to mock workers (which don't exist) - // 5. ❌ Each connection fails, returning WorkerCommunicationError("mock_worker_X") - // 6. ❌ Retry logic removes "failed" workers: first removes mock_worker_2, then mock_worker_1 - // 7. ❌ After 3 retries, workers HashMap is empty: {} - // 8. ❌ assign_to_workers() panics when trying to access worker_addrs[0] on empty list - // - // 💡 Solutions for proper testing: - // - Mock the Flight client layer (complex, requires significant refactoring) - // - Create a test-only version of distribute_stages() that skips communication - // - Refactor the architecture to use dependency injection for better testability - // - Use integration tests with actual worker processes instead of unit tests - // - // For now, we focus on testing the individual components that don't require worker communication. - #[tokio::test] - #[ignore] - async fn test_prepare_query() { - setup_mock_worker_env(); - let planner = QueryPlanner::new(); - - // Test with a simple SELECT query - let sql = "SELECT 1 as test_col, 'hello' as text_col"; - let result = planner.prepare_query(sql).await; - - match result { - Ok(query_plan) => { - assert!(query_plan.explain_data.is_none()); - assert!(!query_plan.query_id.is_empty()); - assert!(!query_plan.worker_addresses.is_empty()); - assert_eq!(query_plan.schema.fields().len(), 2); - assert_eq!(query_plan.schema.field(0).name(), "test_col"); - assert_eq!(query_plan.schema.field(1).name(), "text_col"); - println!("✓ prepare_query succeeded with proper structure"); - } - Err(e) => { - let error_msg = format!("{:?}", e); - assert!( - error_msg.contains("worker") - || error_msg.contains("address") - || error_msg.contains("DFRAY_WORKER") - || error_msg.contains("index out of bounds"), - "Unexpected error type: {}", - error_msg - ); - println!( - "✓ prepare_query failed with expected worker discovery error: {}", - error_msg - ); - } - } + // gather some information we need to send back such that + // we can send a ticket to the client + let final_stage = &distributed_stages[distributed_stages.len() - 1]; + let schema = Arc::clone(&final_stage.plan.schema()); + let final_stage_id = final_stage.stage_id; + + // distribute the stages to workers, further dividing them up + // into chunks of partitions (partition_groups) + let (final_workers, tasks) = + distribute_stages(&query_id, distributed_stages, worker_addrs).await?; + + let qp = QueryPlan { + query_id, + session_context: ctx, + worker_addresses: final_workers, + final_stage_id, + schema, + logical_plan, + physical_plan, + distributed_plan, + distributed_tasks: tasks, + }; + + Ok(qp) } } diff --git a/src/record_batch_exec.rs b/src/record_batch_exec.rs new file mode 100644 index 00000000..dce27933 --- /dev/null +++ b/src/record_batch_exec.rs @@ -0,0 +1,81 @@ +use std::{fmt::Formatter, sync::Arc}; + +use arrow::array::RecordBatch; +use datafusion::{ + error::Result, + execution::SendableRecordBatchStream, + physical_expr::EquivalenceProperties, + physical_plan::{ + execution_plan::{Boundedness, EmissionType}, + stream::RecordBatchStreamAdapter, + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + }, +}; + +#[derive(Debug)] +pub struct RecordBatchExec { + pub batch: RecordBatch, + pub properties: PlanProperties, +} + +impl RecordBatchExec { + pub fn new(batch: RecordBatch) -> Self { + let properties = PlanProperties::new( + EquivalenceProperties::new(batch.schema()), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + ); + Self { batch, properties } + } +} + +impl DisplayAs for RecordBatchExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "RecordBatchExec") + } +} + +impl ExecutionPlan for RecordBatchExec { + fn name(&self) -> &str { + "RecordBatchExec" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn properties(&self) -> &PlanProperties { + &self.properties + } + + fn children(&self) -> Vec<&std::sync::Arc> { + vec![] + } + + fn with_new_children( + self: std::sync::Arc, + children: Vec>, + ) -> Result> { + // TODO: generalize this + assert_eq!(children.len(), 1); + Ok(Arc::new(Self::new(self.batch.clone()))) + } + + fn execute( + &self, + _partition: usize, + _context: std::sync::Arc, + ) -> Result { + let batch = self.batch.clone(); + + let output = async move { Ok(batch) }; + + let stream = Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + futures::stream::once(output), + )) as SendableRecordBatchStream; + + Ok(stream) + } +} diff --git a/src/result.rs b/src/result.rs index 43536ffb..e549257b 100644 --- a/src/result.rs +++ b/src/result.rs @@ -1,5 +1,8 @@ +use arrow_flight::error::FlightError; use thiserror::Error; +use crate::vocab::Host; + #[derive(Debug, Error)] pub enum DFRayError { #[error("Internal Arrow error: {0}")] @@ -7,10 +10,13 @@ pub enum DFRayError { #[error("Internal DataFusion error: {0}")] DataFusionError(#[from] datafusion::error::DataFusionError), + #[error("Arrow Flight error: {0}")] + FlightError(#[from] FlightError), + #[error("Failed to communicate with worker: {0}")] // the name of the worker // improve this as we swallow the error here - WorkerCommunicationError(String), + WorkerCommunicationError(Host), #[error(transparent)] Other(#[from] anyhow::Error), diff --git a/src/stage_reader.rs b/src/stage_reader.rs index a3c090fe..3f1343a7 100644 --- a/src/stage_reader.rs +++ b/src/stage_reader.rs @@ -1,6 +1,6 @@ use std::{fmt::Formatter, sync::Arc}; -use arrow_flight::Ticket; +use arrow_flight::{decode::FlightRecordBatchStream, error::FlightError, Ticket}; use datafusion::{ arrow::datatypes::SchemaRef, common::{internal_datafusion_err, internal_err}, @@ -12,15 +12,16 @@ use datafusion::{ stream::RecordBatchStreamAdapter, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, }, + prelude::SessionConfig, }; use futures::{stream::TryStreamExt, StreamExt}; use prost::Message; use crate::{ - logging::trace, - protobuf::FlightTicketData, + logging::{error, trace}, + protobuf::{FlightDataMetadata, FlightTicketData}, util::{get_client, CombinedRecordBatchStream}, - vocab::{CtxName, CtxStageAddrs}, + vocab::{CtxAnnotatedOutputs, CtxHost, CtxStageAddrs}, }; pub(crate) struct QueryId(pub String); @@ -124,9 +125,9 @@ impl ExecutionPlan for DFRayStageReaderExec { let ctx_name = &context .session_config() - .get_extension::() - .unwrap_or(Arc::new(CtxName("unknown ctx".to_string()))) - .0; + .get_extension::() + .map(|ctx_host| ctx_host.0.to_string()) + .unwrap_or("unknown_context_host!".to_string()); trace!(" trying to get clients for {:?}", stage_addrs); let clients = stage_addrs @@ -147,7 +148,7 @@ impl ExecutionPlan for DFRayStageReaderExec { )) })?? .iter() - .map(|(name, addr)| get_client(name, addr)) + .map(get_client) .collect::>>()?; trace!("got clients. {name} num clients: {}", clients.len()); @@ -166,6 +167,7 @@ impl ExecutionPlan for DFRayStageReaderExec { let schema = self.schema.clone(); let num_clients = clients.len(); + let ctx_name_capture = ctx_name.clone(); let stream = async_stream::stream! { let mut error = false; @@ -173,19 +175,22 @@ impl ExecutionPlan for DFRayStageReaderExec { let mut streams = vec![]; for (i, mut client) in clients.into_iter().enumerate() { let name = name.clone(); - trace!("{name} Getting flight stream {}/{}",i+1, num_clients); + trace!("{name} - {ctx_name_capture} Getting flight stream {}/{}",i+1, num_clients); match client.do_get(ticket.clone()).await { Ok(flight_stream) => { - trace!("{name} Got flight stream. headers:{:?}", flight_stream.headers()); - let rbr_stream = RecordBatchStreamAdapter::new(schema.clone(), - flight_stream - .map_err(move |e| internal_datafusion_err!("{name} Error consuming flight stream from {}: {e}", client.destination))); + trace!("{name} - {ctx_name_capture} Got flight stream. headers:{:?}", flight_stream.headers()); + let rbr_stream = make_flight_metadata_saver_stream( + ctx_name_capture.clone(), + context.session_config(), + schema.clone(), + flight_stream, + ); - streams.push(Box::pin(rbr_stream) as SendableRecordBatchStream); + streams.push(rbr_stream); }, Err(e) => { error = true; - yield internal_err!("{name} Error getting flight stream from {}: {e}", client.destination); + yield internal_err!("{name} - {ctx_name_capture} Error getting flight stream from {}: {e}", client.host); } } } @@ -205,3 +210,79 @@ impl ExecutionPlan for DFRayStageReaderExec { ))) } } + +fn make_flight_metadata_saver_stream( + name: String, + config: &SessionConfig, + schema: SchemaRef, + stream: FlightRecordBatchStream, +) -> SendableRecordBatchStream { + let mut decoder = stream.into_inner(); + + let task_outputs = config + .get_extension::() + .unwrap_or(Arc::new(CtxAnnotatedOutputs::default())) + .0 + .clone(); + + let name_capture = name.clone(); + + #[allow(unused_assignments)] // clippy can't understand our assignment to done in the macro + let new_stream = async_stream::stream! { + let mut done = false; + while !done { + match (done, decoder.next().await) { + (false, Some(Ok(flight_data))) => { + let app_metadata_bytes = flight_data.app_metadata(); + + if !app_metadata_bytes.is_empty() { + trace!("{name} Received trailing metadata from flight stream"); + // decode the metadata + match FlightDataMetadata::decode(&app_metadata_bytes[..]) + .map(|fdm| fdm.annotated_task_outputs.unwrap_or_default()) { + Ok(outputs) => { + trace!("{name} Decoded flight data metadata annotated outputs: {:?}", outputs); + task_outputs.lock().extend(outputs.outputs); + } + Err(e) => { + yield Err(FlightError::DecodeError(format!( + "{name} Failed to decode flight data metadata: {e:#?}" + ))); + } + } + + // ok we consumed the last flight message and extracted the trailing + // metadata, we don't want to yield this payload as there are no records in + // it + done = true; + } else { + // just normal data, yield to consumer + trace!("received normal data"); + yield Ok(flight_data.inner); + } + }, + (false, Some(Err(e))) => { + // propagate this error + yield Err(e); + } + (false, None) => { + done = true; + } + (true,_) => { + // already done, ignore any further data + error!("{name} Unreachable block. flight stream already done"); + } + } + } + trace!("Done with flight stream {name} - no more data to yield"); + }; + + let trailing_stream = + FlightRecordBatchStream::new_from_flight_data(new_stream).map_err(move |e| { + internal_datafusion_err!( + "{name_capture} Failed to consume from our flight data stream: {e:#?}" + ) + }); + + Box::pin(RecordBatchStreamAdapter::new(schema, trailing_stream)) +} diff --git a/src/test_utils.rs b/src/test_utils.rs deleted file mode 100644 index cd01f1cc..00000000 --- a/src/test_utils.rs +++ /dev/null @@ -1,119 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you 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. - -//! Shared test utilities for DataFusion Distributed tests -//! -//! This module contains common test helper functions that are used across -//! multiple test modules to avoid code duplication. - -#[cfg(test)] -pub mod explain_test_helpers { - use arrow_flight::decode::FlightRecordBatchStream; - use datafusion::physical_plan::ExecutionPlan; - use futures::StreamExt; - - use crate::{ - flight_handlers::FlightRequestHandler, - protobuf::{DistributedExplainExecNode, TicketStatementData}, - proxy_service::DfRayProxyHandler, - query_planner::{QueryPlan, QueryPlanner}, - }; - - /// Create a test FlightRequestHandler for testing - pub fn create_test_flight_handler() -> FlightRequestHandler { - FlightRequestHandler::new(QueryPlanner::new()) - } - - /// Create a test DfRayProxyHandler for testing - bypasses worker discovery initialization - pub fn create_test_proxy_handler() -> DfRayProxyHandler { - // Create the handler directly without calling new() to avoid worker discovery - // during test initialization. - DfRayProxyHandler { - flight_handler: FlightRequestHandler::new(QueryPlanner::new()), - } - } - - /// Create TicketStatementData for EXPLAIN testing from QueryPlan - pub fn create_explain_ticket_statement_data(plans: QueryPlan) -> TicketStatementData { - let explain_data = plans.explain_data.map(|data| DistributedExplainExecNode { - schema: data.schema().as_ref().try_into().ok(), - logical_plan: data.logical_plan().to_string(), - physical_plan: data.physical_plan().to_string(), - distributed_plan: data.distributed_plan().to_string(), - distributed_stages: data.distributed_stages().to_string(), - }); - - TicketStatementData { - query_id: plans.query_id, - stage_id: plans.final_stage_id, - stage_addrs: Some(plans.worker_addresses.into()), - schema: Some(plans.schema.as_ref().try_into().unwrap()), - explain_data, - } - } - - /// Consume a DoGetStream and verify it contains expected EXPLAIN results - pub async fn verify_explain_stream_results(stream: crate::flight::DoGetStream) { - // Convert the stream to a FlightRecordBatchStream and consume it - // Map Status errors to FlightError to match the expected stream type - let mapped_stream = stream - .map(|result| result.map_err(|status| arrow_flight::error::FlightError::from(status))); - let mut flight_stream = FlightRecordBatchStream::new_from_flight_data(mapped_stream); - - let mut batches = Vec::new(); - while let Some(batch_result) = flight_stream.next().await { - let batch = batch_result.expect("Failed to get batch from stream"); - batches.push(batch); - } - - // Verify we got exactly one batch with EXPLAIN results - assert_eq!(batches.len(), 1); - let batch = &batches[0]; - - // Verify schema: should have 2 columns (plan_type, plan) - assert_eq!(batch.num_columns(), 2); - assert_eq!(batch.schema().field(0).name(), "plan_type"); - assert_eq!(batch.schema().field(1).name(), "plan"); - - // Verify we have 4 rows (logical_plan, physical_plan, distributed_plan, distributed_stages) - assert_eq!(batch.num_rows(), 4); - - // Verify the plan_type column contains the expected values - let plan_type_column = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(plan_type_column.value(0), "logical_plan"); - assert_eq!(plan_type_column.value(1), "physical_plan"); - assert_eq!(plan_type_column.value(2), "distributed_plan"); - assert_eq!(plan_type_column.value(3), "distributed_stages"); - - // Verify the plan column contains actual plan content - let plan_column = batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap(); - assert!(plan_column - .value(0) - .contains("Projection: Int64(1) AS test_col")); - assert!(plan_column.value(1).contains("ProjectionExec")); - assert!(plan_column.value(2).contains("RayStageExec")); - assert!(plan_column.value(3).contains("Stage 0:")); - } -} diff --git a/src/util.rs b/src/util.rs index 48c9efd0..f652ac3a 100644 --- a/src/util.rs +++ b/src/util.rs @@ -2,18 +2,25 @@ use std::{ collections::HashMap, fmt::Display, future::Future, + io::Cursor, pin::Pin, sync::{Arc, OnceLock}, task::{Context, Poll}, time::Duration, }; -use anyhow::Context as anyhowctx; +use anyhow::{anyhow, Context as anyhowctx}; use arrow::{ array::RecordBatch, datatypes::SchemaRef, error::ArrowError, - ipc::{convert::fb_to_schema, root_as_message}, + ipc::{ + convert::fb_to_schema, + reader::StreamReader, + root_as_message, + writer::{IpcWriteOptions, StreamWriter}, + MetadataVersion, + }, }; use arrow_flight::{decode::FlightRecordBatchStream, FlightClient, FlightData, Ticket}; use async_stream::stream; @@ -40,6 +47,7 @@ use parking_lot::RwLock; use prost::Message; use tokio::{ macros::support::thread_rng_n, + net::TcpListener, runtime::{Handle, Runtime}, }; use tonic::transport::Channel; @@ -51,7 +59,7 @@ use crate::{ protobuf::StageAddrs, result::Result, stage_reader::DFRayStageReaderExec, - vocab::Addrs, + vocab::{Addrs, Host}, }; struct Spawner { @@ -154,11 +162,7 @@ pub fn get_addrs(stage_addrs: &StageAddrs) -> Result { for (stage_id, partition_addrs) in stage_addrs.stage_addrs.iter() { let mut stage_addrs = HashMap::new(); for (partition, hosts) in partition_addrs.partition_addrs.iter() { - let mut host_addrs = vec![]; - for host in &hosts.hosts { - host_addrs.push((host.name.clone(), host.addr.clone())); - } - stage_addrs.insert(*partition, host_addrs); + stage_addrs.insert(*partition, hosts.hosts.clone()); } addrs.insert(*stage_id, stage_addrs); } @@ -178,6 +182,28 @@ pub fn flight_data_to_schema(flight_data: &FlightData) -> anyhow::Result Result> { + let schema = batch.schema(); + let buffer: Vec = Vec::new(); + let options = IpcWriteOptions::try_new(8, false, MetadataVersion::V5) + .map_err(|e| internal_datafusion_err!("Cannot create ipcwriteoptions {e}"))?; + + let mut stream_writer = StreamWriter::try_new_with_options(buffer, &schema, options)?; + stream_writer.write(batch)?; + let bytes = stream_writer.into_inner()?; + Ok(bytes) +} + +pub fn ipc_to_batch(bytes: &[u8]) -> Result { + let mut stream_reader = StreamReader::try_new_buffered(Cursor::new(bytes), None)?; + + match stream_reader.next() { + Some(Ok(batch_res)) => Ok(batch_res), + Some(Err(e)) => Err(e.into()), + None => Err(anyhow!("Expected a valid batch").into()), + } +} + /// produce a new SendableRecordBatchStream that will respect the rows /// limit in the batches that it produces. /// @@ -300,10 +326,8 @@ pub fn reporting_stream( } pub struct ProcessorClient { - /// The name of the processor we are connecting to - pub destination: String, - /// the address we are connecting to - addr: String, + /// the host we are connecting to + pub(crate) host: Host, /// The flight client to the processor inner: FlightClient, /// the channel cache in the factory @@ -312,14 +336,12 @@ pub struct ProcessorClient { impl ProcessorClient { pub fn new( - name: String, + host: Host, inner: FlightClient, channels: Arc>>, - addr: String, ) -> Self { Self { - destination: name, - addr, + host, inner, channels, } @@ -332,8 +354,8 @@ impl ProcessorClient { let stream = self.inner.do_get(ticket).await .inspect_err(|e| { error!("Error in do_get for processor {}: {e:?}. - Considering this channel poisoned and removing it from ProcessorClientFactory cache", self.destination); - self.channels.write().remove(&self.addr); + Considering this channel poisoned and removing it from ProcessorClientFactory cache", self.host); + self.channels.write().remove(&self.host.addr); })?; Ok(stream) @@ -348,9 +370,9 @@ impl ProcessorClient { "Error in do_action for processor {}: {e:?}. Considering this channel poisoned and removing it from ProcessorClientFactory \ cache", - self.destination + self.host ); - self.channels.write().remove(&self.addr); + self.channels.write().remove(&self.host.addr); })?; Ok(result) @@ -368,7 +390,7 @@ impl ProcessorClientFactory { } } - pub fn get_client(&self, name: &str, addr: &str) -> Result { + pub fn get_client(&self, host: &Host) -> Result { // ideally we want to reuse channels as Tonic encourages cloning them when // you can. This could would allow us to lazily create channels and keep // them around so that on subsequent requests to the same address, we won't @@ -391,18 +413,18 @@ impl ProcessorClientFactory { // higher number of processors, even 2. I'm going to leave in the // functionality for cached channels for now - let url = format!("http://{addr}"); + let url = format!("http://{}", host.addr); - let maybe_chan = self.channels.read().get(addr).cloned(); + let maybe_chan = self.channels.read().get(&host.addr).cloned(); let chan = match maybe_chan { Some(chan) => { - debug!("ProcessorFactory using cached channel for {addr}"); + debug!("ProcessorFactory using cached channel for {host}"); chan } None => { - let addr_c = addr.to_owned(); + let host_str = host.to_string(); let fut = async move { - trace!("ProcessorFactory connecting to {addr_c}"); + trace!("ProcessorFactory connecting to {host_str}"); Channel::from_shared(url.clone()) .map_err(|e| { internal_datafusion_err!("ProcessorFactory invalid url {e:#?}") @@ -421,30 +443,31 @@ impl ProcessorClientFactory { "ProcessorFactory Cannot wait for channel connect future {e:#?}" ) })??; - trace!("ProcessorFactory connected to {addr}"); - self.channels.write().insert(addr.to_string(), chan.clone()); + trace!("ProcessorFactory connected to {host}"); + self.channels + .write() + .insert(host.addr.to_string(), chan.clone()); chan } }; - debug!("ProcessorFactory have channel now for {addr}"); + debug!("ProcessorFactory have channel now for {host}"); let flight_client = FlightClient::new(chan); - debug!("ProcessorFactory made flight client for {addr}"); + debug!("ProcessorFactory made flight client for {host}"); Ok(ProcessorClient::new( - name.to_owned(), + host.clone(), flight_client, self.channels.clone(), - addr.to_owned(), )) } } static FACTORY: OnceLock = OnceLock::new(); -pub fn get_client(name: &str, addr: &str) -> Result { +pub fn get_client(host: &Host) -> Result { let factory = FACTORY.get_or_init(ProcessorClientFactory::new); - factory.get_client(name, addr) + factory.get_client(host) } /// Copied from datafusion_physical_plan::union as its useful and not public @@ -481,7 +504,10 @@ impl Stream for CombinedRecordBatchStream { let stream = self.entries.get_mut(idx).unwrap(); match Pin::new(stream).poll_next(cx) { - Ready(Some(val)) => return Ready(Some(val)), + Ready(Some(val)) => { + trace!("Combined stream got {:?}", val); + return Ready(Some(val)); + } Ready(None) => { // Remove the entry self.entries.swap_remove(idx); @@ -523,7 +549,7 @@ fn print_node(plan: &Arc, indent: usize, output: &mut String) "[ output_partitions: {}]{:>indent$}{}", plan.output_partitioning().partition_count(), "", - displayable(plan.as_ref()).one_line(), + displayable(plan.as_ref()).set_show_schema(true).one_line(), indent = indent )); @@ -617,6 +643,16 @@ pub fn maybe_register_object_store(ctx: &SessionContext, url: &Url) -> Result<() Ok(()) } +pub async fn start_up(port: usize) -> Result { + let my_host_str = format!("0.0.0.0:{}", port); + + let listener = TcpListener::bind(&my_host_str) + .await + .context("Could not bind socket to {my_host_str}")?; + + Ok(listener) +} + #[cfg(test)] mod test { use std::{sync::Arc, vec}; diff --git a/src/vocab.rs b/src/vocab.rs index 1879fdbf..54fd60fb 100644 --- a/src/vocab.rs +++ b/src/vocab.rs @@ -1,4 +1,13 @@ -use std::collections::HashMap; +use std::{collections::HashMap, fmt::Display, sync::Arc}; + +use parking_lot::Mutex; + +pub use crate::protobuf::AnnotatedTaskOutput; +pub use crate::protobuf::DdTask as DDTask; +pub use crate::protobuf::Host; +pub use crate::protobuf::Hosts; +pub use crate::protobuf::PartitionAddrs; +pub use crate::protobuf::StageAddrs; /// a map of stage_id, partition to a list (name,endpoint address) that can /// serve this (stage_id, and partition). It is assumed that to consume a @@ -6,13 +15,25 @@ use std::collections::HashMap; /// merge the results. /// /// This is on a per query basis -pub type Addrs = HashMap>>; +pub type Addrs = HashMap>>; /// used to hold an Addrs as an extenstion for datafusion SessionContext pub(crate) struct CtxStageAddrs(pub Addrs); -/// used to hold a worker name as an extension for datafusion SessionContext -pub(crate) struct CtxName(pub String); +/// used to hold information about how is executing +pub(crate) struct CtxHost(pub Host); /// used to hold a partition group as an extension for datafusion SessionContext -pub struct CtxPartitionGroup(pub Vec); +pub(crate) struct CtxPartitionGroup(pub Vec); + +/// used to hold a stage id as an extension for datafusion SessionContext +pub(crate) struct CtxStageId(pub u64); + +#[derive(Default)] +pub(crate) struct CtxAnnotatedOutputs(pub Arc>>); + +impl Display for Host { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}: {}", self.name, self.addr) + } +} diff --git a/src/k8s.rs b/src/worker_discovery.rs similarity index 72% rename from src/k8s.rs rename to src/worker_discovery.rs index 7303d124..d4a78a93 100644 --- a/src/k8s.rs +++ b/src/worker_discovery.rs @@ -1,9 +1,11 @@ use std::{ collections::HashMap, sync::{Arc, OnceLock}, + time::Duration, }; use anyhow::{anyhow, Context}; +use arrow_flight::{Action, FlightClient}; use futures::{StreamExt, TryStreamExt}; use k8s_openapi::api::{apps::v1::Deployment, core::v1::Pod}; use kube::{ @@ -11,15 +13,18 @@ use kube::{ Client, }; use parking_lot::RwLock; +use prost::Message; +use tonic::transport::Channel; use crate::{ logging::{debug, error, trace}, result::Result, + vocab::Host, }; static WORKER_DISCOVERY: OnceLock> = OnceLock::new(); -pub fn get_worker_addresses() -> Result> { +pub fn get_worker_addresses() -> Result> { match WORKER_DISCOVERY.get_or_init(WorkerDiscovery::new) { Ok(wd) => { let worker_addrs = wd.get_addresses(); @@ -27,7 +32,7 @@ pub fn get_worker_addresses() -> Result> { "Worker addresses found:\n{}", worker_addrs .iter() - .map(|(name, addr)| format!("{}: {}", name, addr)) + .map(|host| format!("{host}")) .collect::>() .join("\n") ); @@ -38,7 +43,7 @@ pub fn get_worker_addresses() -> Result> { } struct WorkerDiscovery { - addresses: Arc>>, + addresses: Arc>>, } impl WorkerDiscovery { @@ -50,12 +55,9 @@ impl WorkerDiscovery { Ok(wd) } - fn get_addresses(&self) -> Vec<(String, String)> { + fn get_addresses(&self) -> Vec { let guard = self.addresses.read(); - guard - .iter() - .map(|(_ip, (name, addr))| (name.clone(), addr.clone())) - .collect() + guard.iter().map(|(_ip, host)| host.clone()).collect() } fn start(&self) -> Result<()> { @@ -64,9 +66,13 @@ impl WorkerDiscovery { let worker_deployment_namespace_env = std::env::var("DFRAY_WORKER_DEPLOYMENT_NAMESPACE"); if worker_addrs_env.is_ok() { - // if the env var is set, use it - self.set_worker_addresses_from_env(worker_addrs_env.unwrap().as_str()) - .context("Failed to set worker addresses from env var")?; + let addresses = self.addresses.clone(); + tokio::spawn(async move { + // if the env var is set, use it + set_worker_addresses_from_env(addresses, worker_addrs_env.unwrap().as_str()) + .await + .expect("Could not set worker addresses from env"); + }); } else if worker_deployment_namespace_env.is_ok() && worker_deployment_env.is_ok() { let addresses = self.addresses.clone(); let deployment = worker_deployment_env.unwrap(); @@ -87,24 +93,24 @@ impl WorkerDiscovery { } Ok(()) } +} - fn set_worker_addresses_from_env(&self, env_str: &str) -> Result<()> { - // get addresss from an env var where addresses are split by comans - // and in the form of name/address,name/address - let mut guard = self.addresses.write(); +async fn set_worker_addresses_from_env( + addresses: Arc>>, + env_str: &str, +) -> Result<()> { + // get addresss from an env var where addresses are split by comans + // and in the form of name/address,name/address - for addr in env_str.split(',') { - let parts: Vec<&str> = addr.split('/').collect(); - if parts.len() != 2 { - return Err(anyhow!("Invalid worker address format: {addr}").into()); - } - let name = parts[0].to_string(); - let address = parts[1].to_string(); - guard.insert(address.clone(), (name, address)); - } - Ok(()) + for addr in env_str.split(',') { + let host = get_worker_host(addr.to_string()) + .await + .context(format!("Failed to get worker host for address: {}", addr))?; + addresses.write().insert(addr.to_owned(), host); } + Ok(()) } + /// Continuously watch for changes to pods in a Kubernetes deployment and call a /// handler function whenever the list of hosts changes. /// @@ -120,7 +126,7 @@ impl WorkerDiscovery { /// Returns an error if there's an issue connecting to the Kubernetes API /// or if the deployment or its pods cannot be found async fn watch_deployment_hosts_continuous( - addresses: Arc>>, + addresses: Arc>>, deployment_name: &str, namespace: &str, ) -> Result<()> { @@ -186,15 +192,14 @@ async fn watch_deployment_hosts_continuous( pod ); if let Some(Some(_ip)) = pod.status.as_ref().map(|s| s.pod_ip.as_ref()) { - let (pod_ip, name_str, host_str) = get_worker_info_from_pod(pod)?; + let (pod_ip, host) = get_worker_info_from_pod(pod).await?; debug!( - "Pod {} has IP address {}, name {}, host {}", + "Pod {} has IP address {}, host {}", pod.name_any(), pod_ip, - name_str, - host_str + host ); - addresses.write().insert(pod_ip, (name_str, host_str)); + addresses.write().insert(pod_ip, host); } else { trace!("Pod {} has no IP address, skipping", pod.name_any()); } @@ -220,7 +225,35 @@ async fn watch_deployment_hosts_continuous( Ok(()) } -fn get_worker_info_from_pod(pod: &Pod) -> Result<(String, String, String)> { +async fn get_worker_host(addr: String) -> Result { + let mut client = Channel::from_shared(format!("http://{addr}")) + .context("Failed to create channel")? + .connect_timeout(Duration::from_secs(2)) + .connect() + .await + .map(FlightClient::new) + .context("Failed to connect to worker")?; + + let action = Action { + r#type: "get_host".to_string(), + body: vec![].into(), + }; + + let mut response = client + .do_action(action) + .await + .context("Failed to send action to worker")?; + + Ok(response + .try_next() + .await + .transpose() + .context("error consuming do_action response")? + .map(Host::decode)? + .context("Failed to decode Host from worker response")?) +} + +async fn get_worker_info_from_pod(pod: &Pod) -> Result<(String, Host)> { let status = pod.status.as_ref().context("Pod has no status")?; let pod_ip = status.pod_ip.as_ref().context("Pod has no IP address")?; @@ -249,7 +282,10 @@ fn get_worker_info_from_pod(pod: &Pod) -> Result<(String, String, String)> { Err(anyhow::anyhow!("Pod {} has no IP address", pod.name_any()).into()) } else { let host_str = format!("{}:{}", pod_ip, port); - let name_str = format!("{}:{}", pod_ip, port); // for now - Ok((pod_ip.to_owned(), name_str, host_str)) + let host = get_worker_host(host_str.clone()).await.context(format!( + "Failed to get worker host for pod {}", + pod.name_any() + ))?; + Ok((pod_ip.to_owned(), host)) } } diff --git a/tests/common/mod.rs b/tests/common/mod.rs index 8ab47cc6..db821a0d 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -220,7 +220,7 @@ except Exception as e: print(f"Error executing distributed query: {{str(e)}}", file=sys.stderr) sys.exit(1) "#, - self.get_proxy_address().split(':').last().unwrap() + self.get_proxy_address().split(':').next_back().unwrap() ); let script_path = Self::write_temp_file( @@ -271,15 +271,13 @@ except Exception as e: for port in &ports { // Find and kill processes using lsof if let Ok(output) = Command::new("lsof") - .args(&["-ti", &format!(":{}", port)]) + .args(["-ti", &format!(":{}", port)]) .output() { let stdout = String::from_utf8_lossy(&output.stdout); for line in stdout.lines() { if let Ok(pid) = line.trim().parse::() { - let _ = Command::new("kill") - .args(&["-9", &pid.to_string()]) - .output(); + let _ = Command::new("kill").args(["-9", &pid.to_string()]).output(); println!(" 🔥 Killed process {} on port {}", pid, port); } } @@ -362,7 +360,7 @@ except Exception as e: // Check if packages are already installed let check_cmd = Command::new("python3") - .args(&["-c", "import adbc_driver_manager; import adbc_driver_flightsql; import duckdb; import pyarrow; print('OK')"]) + .args(["-c", "import adbc_driver_manager; import adbc_driver_flightsql; import duckdb; import pyarrow; print('OK')"]) .output(); if let Ok(output) = check_cmd { @@ -449,7 +447,7 @@ except Exception as e: let worker = Self::spawn_process( binary_path_str, &["--mode", "worker", "--port", &port.to_string()], - &[("DFRAY_TABLES", &tpch_tables), ("DFRAY_VIEWS", &tpch_views)], + &[("DFRAY_TABLES", &tpch_tables), ("DFRAY_VIEWS", tpch_views)], &format!("start worker {}", i + 1), )?; self.worker_processes.push(worker); @@ -473,7 +471,7 @@ except Exception as e: &[ ("DFRAY_WORKER_ADDRESSES", &worker_addresses), ("DFRAY_TABLES", &tpch_tables), - ("DFRAY_VIEWS", &tpch_views), + ("DFRAY_VIEWS", tpch_views), ], "start proxy", )?; @@ -819,7 +817,7 @@ pub fn batches_to_sorted_strings( let mut in_table = false; let mut header_found = false; - for (_line_idx, line) in lines.iter().enumerate() { + for line in lines.iter() { let trimmed = line.trim(); // Detect table boundaries (+---+)