From 62125fde974e9ccac98dd608ca0aa370a3741d0a Mon Sep 17 00:00:00 2001 From: LiaCastaneda Date: Mon, 7 Jul 2025 10:44:06 +0200 Subject: [PATCH 1/4] Support for substrait plans --- Cargo.lock | 311 +++++++++++++++---- Cargo.toml | 4 + src/explain.rs | 7 +- src/flight.rs | 20 ++ src/flight_handlers.rs | 21 ++ src/proxy_service.rs | 25 +- src/query_planner.rs | 131 ++++++-- testdata/substrait/select_one.substrait.json | 65 ++++ 8 files changed, 488 insertions(+), 96 deletions(-) create mode 100644 testdata/substrait/select_one.substrait.json diff --git a/Cargo.lock b/Cargo.lock index 7fb94f52..e994af35 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -255,7 +255,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64", + "base64 0.22.1", "chrono", "comfy-table", "half", @@ -308,7 +308,7 @@ dependencies = [ "arrow-schema", "arrow-select", "arrow-string", - "base64", + "base64 0.22.1", "bytes", "futures", "once_cell", @@ -464,6 +464,17 @@ dependencies = [ "zstd-safe", ] +[[package]] +name = "async-recursion" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b43422f69d8ff38f95f1b2bb76517c91589a924d1559a0e935d7c8ce0274c11" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "async-stream" version = "0.3.6" @@ -614,6 +625,12 @@ dependencies = [ "windows-targets 0.52.6", ] +[[package]] +name = "base64" +version = "0.21.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" + [[package]] name = "base64" version = "0.22.1" @@ -773,9 +790,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.27" +version = "1.2.28" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d487aa071b5f64da6f19a3e848e3578944b726ee5a4854b82172f02aa876bfdc" +checksum = "4ad45f4f74e4e20eaa392913b7b33a7091c87e59628f4dd27888205ad888843c" dependencies = [ "jobserver", "libc", @@ -1102,7 +1119,7 @@ dependencies = [ [[package]] name = "datafusion" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "arrow-ipc", @@ -1156,7 +1173,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "async-trait", @@ -1181,7 +1198,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "async-trait", @@ -1203,13 +1220,13 @@ dependencies = [ [[package]] name = "datafusion-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "ahash", "apache-avro", "arrow", "arrow-ipc", - "base64", + "base64 0.22.1", "chrono", "half", "hashbrown 0.14.5", @@ -1230,7 +1247,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "futures", "log", @@ -1240,7 +1257,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "async-compression", @@ -1275,7 +1292,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "apache-avro", "arrow", @@ -1299,7 +1316,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "async-trait", @@ -1323,7 +1340,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "async-trait", @@ -1347,7 +1364,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "async-trait", @@ -1378,12 +1395,12 @@ dependencies = [ [[package]] name = "datafusion-doc" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" [[package]] name = "datafusion-execution" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "dashmap", @@ -1401,7 +1418,7 @@ dependencies = [ [[package]] name = "datafusion-expr" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "async-trait", @@ -1422,7 +1439,7 @@ dependencies = [ [[package]] name = "datafusion-expr-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "datafusion-common", @@ -1434,11 +1451,11 @@ dependencies = [ [[package]] name = "datafusion-functions" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "arrow-buffer", - "base64", + "base64 0.22.1", "blake2", "blake3", "chrono", @@ -1462,7 +1479,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "ahash", "arrow", @@ -1482,7 +1499,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "ahash", "arrow", @@ -1494,7 +1511,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "arrow-ord", @@ -1515,7 +1532,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "async-trait", @@ -1530,7 +1547,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "datafusion-common", @@ -1547,7 +1564,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -1556,7 +1573,7 @@ dependencies = [ [[package]] name = "datafusion-macros" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "datafusion-expr", "quote", @@ -1566,12 +1583,13 @@ dependencies = [ [[package]] name = "datafusion-optimizer" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "chrono", "datafusion-common", "datafusion-expr", + "datafusion-expr-common", "datafusion-physical-expr", "indexmap 2.10.0", "itertools 0.14.0", @@ -1584,7 +1602,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "ahash", "arrow", @@ -1605,7 +1623,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "ahash", "arrow", @@ -1618,7 +1636,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "datafusion-common", @@ -1637,7 +1655,7 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "ahash", "arrow", @@ -1666,7 +1684,7 @@ dependencies = [ [[package]] name = "datafusion-proto" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "chrono", @@ -1681,7 +1699,7 @@ dependencies = [ [[package]] name = "datafusion-proto-common" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "datafusion-common", @@ -1691,7 +1709,7 @@ dependencies = [ [[package]] name = "datafusion-pruning" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "arrow-schema", @@ -1708,7 +1726,7 @@ dependencies = [ [[package]] name = "datafusion-session" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "async-trait", @@ -1731,7 +1749,7 @@ dependencies = [ [[package]] name = "datafusion-sql" version = "48.0.0" -source = "git+https://github.com/apache/datafusion?branch=main#ebf49b4fe585820caaad3b8c04529633190f48be" +source = "git+https://github.com/apache/datafusion?branch=main#a7151730a6c65a62793d293557ca15e975fcb701" dependencies = [ "arrow", "bigdecimal", @@ -1744,6 +1762,24 @@ dependencies = [ "sqlparser", ] +[[package]] +name = "datafusion-substrait" +version = "48.0.0" +source = "git+https://github.com/apache/datafusion?branch=main#85a92514270f8da351496c5d81445ba481cc6f05" +dependencies = [ + "async-recursion", + "async-trait", + "chrono", + "datafusion", + "itertools 0.14.0", + "object_store", + "pbjson-types", + "prost", + "substrait", + "tokio", + "url", +] + [[package]] name = "derive_more" version = "2.0.1" @@ -1798,6 +1834,7 @@ dependencies = [ "clap", "datafusion", "datafusion-proto", + "datafusion-substrait", "env_logger", "futures", "itertools 0.14.0", @@ -1811,6 +1848,8 @@ dependencies = [ "rand 0.8.5", "rustc_version", "rustls", + "serde", + "serde_json", "tempfile", "test-log", "thiserror 1.0.69", @@ -2153,9 +2192,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.10" +version = "0.4.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9421a676d1b147b16b82c9225157dc629087ef8ec4d5e2960f9437a90dac0a5" +checksum = "17da50a276f1e01e0ba6c029e47b7100754904ee8a278f886546e98575380785" dependencies = [ "atomic-waker", "bytes", @@ -2214,7 +2253,7 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "headers-core", "http", @@ -2394,7 +2433,7 @@ version = "0.1.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc2fdfdbff08affe55bb779f33b053aa1fe5dd5b54c257343c17edfa55711bdb" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-core", @@ -2581,6 +2620,17 @@ 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" @@ -2612,6 +2662,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.14.0" @@ -2712,7 +2771,7 @@ version = "0.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "aa60a41b57ae1a0a071af77dbcf89fc9819cfe66edaf2beeb204c34459dcf0b2" dependencies = [ - "base64", + "base64 0.22.1", "chrono", "serde", "serde_json", @@ -2737,7 +2796,7 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cb276b85b6e94ded00ac8ea2c68fcf4697ea0553cb25fddc35d4a0ab718db8d" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "chrono", "either", @@ -3214,7 +3273,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7781f96d79ed0f961a7021424ab01840efbda64ae7a505aaea195efc91eaaec4" dependencies = [ "async-trait", - "base64", + "base64 0.22.1", "bytes", "chrono", "form_urlencoded", @@ -3324,7 +3383,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64", + "base64 0.22.1", "brotli", "bytes", "chrono", @@ -3362,13 +3421,50 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pbjson" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7e6349fa080353f4a597daffd05cb81572a9c031a6d4fff7e504947496fcc68" +dependencies = [ + "base64 0.21.7", + "serde", +] + +[[package]] +name = "pbjson-build" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6eea3058763d6e656105d1403cb04e0a41b7bbac6362d413e7c33be0c32279c9" +dependencies = [ + "heck", + "itertools 0.13.0", + "prost", + "prost-types", +] + +[[package]] +name = "pbjson-types" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e54e5e7bfb1652f95bc361d76f3c780d8e526b134b85417e774166ee941f0887" +dependencies = [ + "bytes", + "chrono", + "pbjson", + "pbjson-build", + "prost", + "prost-build", + "serde", +] + [[package]] name = "pem" version = "3.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "38af38e8470ac9dee3ce1bae1af9c1671fffc44ddfd8bd1d0a3445bf349a8ef3" dependencies = [ - "base64", + "base64 0.22.1", "serde", ] @@ -3930,13 +4026,23 @@ version = "0.8.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2b15c43186be67a4fd63bee50d0303afffcef381492ebe2c5d87f324e1b8815c" +[[package]] +name = "regress" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78ef7fa9ed0256d64a688a3747d0fef7a88851c18a5e1d57f115f38ec2e09366" +dependencies = [ + "hashbrown 0.15.4", + "memchr", +] + [[package]] name = "reqwest" -version = "0.12.20" +version = "0.12.22" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eabf4c97d9130e2bf606614eb937e86edac8292eaa6f422f995d7e8de1eb1813" +checksum = "cbc931937e6ca3a06e3b6c0aa7841849b160a90351d6ab467a8b9b9959767531" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-core", "futures-util", @@ -4227,6 +4333,9 @@ name = "semver" version = "1.0.26" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56e6fa9c48d24d85fb3de5ad847117517440f6beceb7798af16b4a87d616b8d0" +dependencies = [ + "serde", +] [[package]] name = "seq-macro" @@ -4296,6 +4405,18 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_tokenstream" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64060d864397305347a78851c51588fd283767e7e7589829e8121d65512340f1" +dependencies = [ + "proc-macro2", + "quote", + "serde", + "syn", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -4479,6 +4600,31 @@ dependencies = [ "syn", ] +[[package]] +name = "substrait" +version = "0.58.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de6d24c270c6c672a86c183c3a8439ba46c1936f93cf7296aa692de3b0ff0228" +dependencies = [ + "heck", + "pbjson", + "pbjson-build", + "pbjson-types", + "prettyplease", + "prost", + "prost-build", + "prost-types", + "regress", + "schemars", + "semver", + "serde", + "serde_json", + "serde_yaml", + "syn", + "typify", + "walkdir", +] + [[package]] name = "subtle" version = "2.6.1" @@ -4537,9 +4683,9 @@ dependencies = [ [[package]] name = "test-log" -version = "0.2.17" +version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7f46083d221181166e5b6f6b1e5f1d499f3a76888826e6cb1d057554157cd0f" +checksum = "1e33b98a582ea0be1168eba097538ee8dd4bbe0f2b01b22ac92ea30054e5be7b" dependencies = [ "env_logger", "test-log-macros", @@ -4548,9 +4694,9 @@ dependencies = [ [[package]] name = "test-log-macros" -version = "0.2.17" +version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "888d0c3c6db53c0fdab160d2ed5e12ba745383d3e85813f2ea0f2b1475ab553f" +checksum = "451b374529930d7601b1eef8d32bc79ae870b6079b069401709c2a8bf9e75f36" dependencies = [ "proc-macro2", "quote", @@ -4694,17 +4840,19 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.45.1" +version = "1.46.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75ef51a33ef1da925cea3e4eb122833cb377c61439ca401b770f54902b806779" +checksum = "1140bb80481756a8cbe10541f37433b459c5aa1e727b4c020fbfebdc25bf3ec4" dependencies = [ "backtrace", "bytes", + "io-uring", "libc", "mio", "parking_lot", "pin-project-lite", "signal-hook-registry", + "slab", "socket2", "tokio-macros", "windows-sys 0.52.0", @@ -4766,7 +4914,7 @@ dependencies = [ "async-stream", "async-trait", "axum", - "base64", + "base64 0.22.1", "bytes", "h2", "http", @@ -4844,7 +4992,7 @@ version = "0.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "adc82fd73de2a9722ac5da747f12383d2bfdb93591ee6c58486e0097890f05f2" dependencies = [ - "base64", + "base64 0.22.1", "bitflags", "bytes", "futures-util", @@ -4970,6 +5118,53 @@ version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f" +[[package]] +name = "typify" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c6c647a34e851cf0260ccc14687f17cdcb8302ff1a8a687a24b97ca0f82406f" +dependencies = [ + "typify-impl", + "typify-macro", +] + +[[package]] +name = "typify-impl" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "741b7f1e2e1338c0bee5ad5a7d3a9bbd4e24c33765c08b7691810e68d879365d" +dependencies = [ + "heck", + "log", + "proc-macro2", + "quote", + "regress", + "schemars", + "semver", + "serde", + "serde_json", + "syn", + "thiserror 2.0.12", + "unicode-ident", +] + +[[package]] +name = "typify-macro" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7560adf816a1e8dad7c63d8845ef6e31e673e39eab310d225636779230cbedeb" +dependencies = [ + "proc-macro2", + "quote", + "schemars", + "semver", + "serde", + "serde_json", + "serde_tokenstream", + "syn", + "typify-impl", +] + [[package]] name = "ucd-trie" version = "0.1.7" diff --git a/Cargo.toml b/Cargo.toml index 49846aab..b2b00b2e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,6 +44,7 @@ datafusion = { git = "https://github.com/apache/datafusion", branch = "main", fe "avro", ] } datafusion-proto = { git = "https://github.com/apache/datafusion", branch = "main" } +datafusion-substrait = { git = "https://github.com/apache/datafusion", branch = "main" } env_logger = "0.11" futures = "0.3" itertools = "0.14" @@ -53,6 +54,9 @@ 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", diff --git a/src/explain.rs b/src/explain.rs index 4f73d4c1..55073dea 100644 --- a/src/explain.rs +++ b/src/explain.rs @@ -160,9 +160,10 @@ impl ExecutionPlan for DistributedExplainExec { 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))?; + 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)?; diff --git a/src/flight.rs b/src/flight.rs index 97f4b0a3..cd5c290d 100644 --- a/src/flight.rs +++ b/src/flight.rs @@ -134,6 +134,12 @@ pub trait FlightSqlHandler: Send + Sync { request: Request, ) -> Result, Status>; + async fn get_flight_info_substrait( + &self, + query: arrow_flight::sql::CommandStatementSubstraitPlan, + request: Request, + ) -> Result, Status>; + async fn do_get_statement( &self, ticket: arrow_flight::sql::TicketStatementQuery, @@ -163,6 +169,20 @@ impl FlightSqlService for FlightSqlServ { error!("Error in do_flight_info_statement: {:?}", e); }) } + + async fn get_flight_info_substrait_plan( + &self, + query: arrow_flight::sql::CommandStatementSubstraitPlan, + request: Request, + ) -> Result, Status> { + self.handler + .get_flight_info_substrait(query, request) + .await + .inspect_err(|e| { + error!("Error in get_flight_info_substrait: {:?}", e); + }) + } + async fn do_get_statement( &self, ticket: arrow_flight::sql::TicketStatementQuery, diff --git a/src/flight_handlers.rs b/src/flight_handlers.rs index 7454e68f..9b4713d6 100644 --- a/src/flight_handlers.rs +++ b/src/flight_handlers.rs @@ -136,6 +136,27 @@ impl FlightRequestHandler { 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. diff --git a/src/proxy_service.rs b/src/proxy_service.rs index 186c5cfb..6297b89e 100644 --- a/src/proxy_service.rs +++ b/src/proxy_service.rs @@ -19,8 +19,9 @@ use std::sync::Arc; use anyhow::{anyhow, Context}; use arrow_flight::{ - flight_service_server::FlightServiceServer, FlightDescriptor, FlightInfo, Ticket, + flight_service_server::FlightServiceServer, sql::TicketStatementQuery, FlightDescriptor, FlightInfo, Ticket }; +use datafusion_substrait::substrait; use parking_lot::Mutex; use prost::Message; use tokio::{ @@ -74,7 +75,7 @@ impl FlightSqlHandler for DfRayProxyHandler { async fn do_get_statement( &self, - ticket: arrow_flight::sql::TicketStatementQuery, + ticket: TicketStatementQuery, request: Request, ) -> Result, Status> { trace!("do_get_statement"); @@ -98,6 +99,22 @@ impl FlightSqlHandler for DfRayProxyHandler { .await } } + + async fn get_flight_info_substrait( + &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")), + }; + + self.flight_handler.handle_substrait_info_request(plan).await + } } /// DFRayProcessorService is a Arrow Flight service that serves streams of @@ -213,8 +230,7 @@ mod tests { use super::*; // Test-specific imports use arrow_flight::{ - sql::{CommandStatementQuery, TicketStatementQuery}, - FlightDescriptor, Ticket, + sql::{CommandStatementQuery, TicketStatementQuery}, FlightDescriptor, Ticket }; use prost::Message; use tonic::Request; @@ -345,4 +361,5 @@ mod tests { // 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 516c989f..96b842ed 100644 --- a/src/query_planner.rs +++ b/src/query_planner.rs @@ -5,6 +5,7 @@ use arrow::datatypes::SchemaRef; use datafusion::{ logical_expr::LogicalPlan, physical_plan::ExecutionPlan, prelude::SessionContext, }; +use datafusion_substrait::logical_plan::{consumer::from_substrait_plan}; use crate::{ explain::{is_explain_query, DistributedExplainExec}, @@ -65,6 +66,38 @@ impl QueryPlanner { 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 @@ -78,9 +111,8 @@ impl QueryPlanner { 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 - let (distributed_plan, distributed_stages) = - execution_planning(physical_plan.clone(), 8192, Some(2)).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, @@ -92,37 +124,45 @@ impl QueryPlanner { }) } - /// Prepare a distributed query + /// 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 + } - if base_result.distributed_stages.is_empty() { - return Err(anyhow!("No stages generated for query").into()); - } + /// Prepare a distributed query (Substrait entry point) + pub async fn prepare_substrait_query(&self, substrait_plan: datafusion_substrait::substrait::proto::Plan) -> Result { + let base_result = self.prepare_substrait_query_base(substrait_plan, "SUBSTRAIT").await?; + self.dispatch_query_plan(base_result).await + } - let worker_addrs = get_worker_addresses()?; + pub async fn prepare_substrait_query_base( + &self, + substrait_plan: datafusion_substrait::substrait::proto::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}"))?; - // gather some information we need to send back such that - // we can send a ticket 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; + let logical_plan = from_substrait_plan(&ctx.state(), &substrait_plan) + .await + .map_err(|e| anyhow!("Failed to convert DataFusion Logical Plan: {e}"))?; - // distribute the stages to workers, further dividing them up - // into chunks of partitions (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, + let physical_plan = physical_planning(&logical_plan, &ctx).await.map_err(|e| anyhow!("Failed to convert DataFusion Physical Plan: {e}"))?; + + // 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, }) } @@ -206,6 +246,9 @@ impl QueryPlanner { #[cfg(test)] mod tests { use super::*; + use std::{fs::{File}, path::Path}; + use std::io::BufReader; + // ////////////////////////////////////////////////////////////// // Test helper functions @@ -213,10 +256,8 @@ mod tests { /// Set up mock worker environment for testing fn setup_mock_worker_env() { - let mock_addrs = vec![ - ("mock_worker_1".to_string(), "localhost:9001".to_string()), - ("mock_worker_2".to_string(), "localhost:9002".to_string()), - ]; + 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)) @@ -257,6 +298,34 @@ mod tests { } } + #[tokio::test] + async fn test_prepare_substrait_query_base() { + let planner = QueryPlanner::new(); + + // 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"), + )).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")); + } + } + #[tokio::test] async fn test_prepare_explain() { let planner = QueryPlanner::new(); diff --git a/testdata/substrait/select_one.substrait.json b/testdata/substrait/select_one.substrait.json new file mode 100644 index 00000000..d1140d84 --- /dev/null +++ b/testdata/substrait/select_one.substrait.json @@ -0,0 +1,65 @@ +{ + "relations": [ + { + "root": { + "input": { + "project": { + "common": { + "emit": { + "outputMapping": [ + 1 + ] + } + }, + "input": { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "dummy" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_REQUIRED" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "virtualTable": { + "values": [ + { + "fields": [ + { + "i64": "0" + } + ] + } + ] + } + } + }, + "expressions": [ + { + "literal": { + "i64": "1" + } + } + ] + } + }, + "names": [ + "test_col" + ] + } + } + ], + "version": { + "minorNumber": 52 + } +} \ No newline at end of file From 2680f428b745535aaca5fcb9cc6e794e861fd54a Mon Sep 17 00:00:00 2001 From: LiaCastaneda Date: Mon, 7 Jul 2025 15:27:58 +0200 Subject: [PATCH 2/4] Format --- src/explain.rs | 7 +++-- src/flight_handlers.rs | 7 +++-- src/proxy_service.rs | 17 ++++++------ src/query_planner.rs | 60 +++++++++++++++++++++++++++--------------- 4 files changed, 56 insertions(+), 35 deletions(-) diff --git a/src/explain.rs b/src/explain.rs index 55073dea..4f73d4c1 100644 --- a/src/explain.rs +++ b/src/explain.rs @@ -160,10 +160,9 @@ impl ExecutionPlan for DistributedExplainExec { 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))?; + 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)?; diff --git a/src/flight_handlers.rs b/src/flight_handlers.rs index 9b4713d6..c614362f 100644 --- a/src/flight_handlers.rs +++ b/src/flight_handlers.rs @@ -136,7 +136,10 @@ impl FlightRequestHandler { Ok(Response::new(flight_info)) } - pub async fn handle_substrait_info_request(&self, substrait_plan: datafusion_substrait::substrait::proto::Plan) -> Result, Status> { + 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) @@ -150,7 +153,7 @@ impl FlightRequestHandler { query_plan.worker_addresses, query_plan.final_stage_id, query_plan.schema, - None // Regular queries don't have explain data + None, // Regular queries don't have explain data )?; trace!("get_flight_info_statement done"); diff --git a/src/proxy_service.rs b/src/proxy_service.rs index 6297b89e..cadff7e2 100644 --- a/src/proxy_service.rs +++ b/src/proxy_service.rs @@ -19,7 +19,8 @@ use std::sync::Arc; use anyhow::{anyhow, Context}; use arrow_flight::{ - flight_service_server::FlightServiceServer, sql::TicketStatementQuery, FlightDescriptor, FlightInfo, Ticket + flight_service_server::FlightServiceServer, sql::TicketStatementQuery, FlightDescriptor, + FlightInfo, Ticket, }; use datafusion_substrait::substrait; use parking_lot::Mutex; @@ -106,14 +107,14 @@ impl FlightSqlHandler for DfRayProxyHandler { _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}")))? - } + 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")), }; - self.flight_handler.handle_substrait_info_request(plan).await + self.flight_handler + .handle_substrait_info_request(plan) + .await } } @@ -230,7 +231,8 @@ mod tests { use super::*; // Test-specific imports use arrow_flight::{ - sql::{CommandStatementQuery, TicketStatementQuery}, FlightDescriptor, Ticket + sql::{CommandStatementQuery, TicketStatementQuery}, + FlightDescriptor, Ticket, }; use prost::Message; use tonic::Request; @@ -361,5 +363,4 @@ mod tests { // 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 96b842ed..4d1c6057 100644 --- a/src/query_planner.rs +++ b/src/query_planner.rs @@ -5,7 +5,7 @@ use arrow::datatypes::SchemaRef; use datafusion::{ logical_expr::LogicalPlan, physical_plan::ExecutionPlan, prelude::SessionContext, }; -use datafusion_substrait::logical_plan::{consumer::from_substrait_plan}; +use datafusion_substrait::logical_plan::consumer::from_substrait_plan; use crate::{ explain::{is_explain_query, DistributedExplainExec}, @@ -75,8 +75,7 @@ impl QueryPlanner { 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 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; @@ -112,7 +111,8 @@ impl QueryPlanner { 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?; + let (distributed_plan, distributed_stages) = + execution_planning(physical_plan.clone(), 8192, Some(2)).await?; Ok(QueryPlanBase { query_id, @@ -131,8 +131,13 @@ impl QueryPlanner { } /// Prepare a distributed query (Substrait entry point) - pub async fn prepare_substrait_query(&self, substrait_plan: datafusion_substrait::substrait::proto::Plan) -> Result { - let base_result = self.prepare_substrait_query_base(substrait_plan, "SUBSTRAIT").await?; + pub async fn prepare_substrait_query( + &self, + substrait_plan: datafusion_substrait::substrait::proto::Plan, + ) -> Result { + let base_result = self + .prepare_substrait_query_base(substrait_plan, "SUBSTRAIT") + .await?; self.dispatch_query_plan(base_result).await } @@ -141,8 +146,11 @@ impl QueryPlanner { substrait_plan: datafusion_substrait::substrait::proto::Plan, query_type: &str, ) -> Result { - debug!("prepare_substrait_query_base: {} Substrait = {:#?}", query_type, substrait_plan); - + 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}"))?; @@ -150,11 +158,13 @@ impl QueryPlanner { .await .map_err(|e| anyhow!("Failed to convert DataFusion Logical Plan: {e}"))?; - - let physical_plan = physical_planning(&logical_plan, &ctx).await.map_err(|e| anyhow!("Failed to convert DataFusion Physical Plan: {e}"))?; + let physical_plan = physical_planning(&logical_plan, &ctx) + .await + .map_err(|e| anyhow!("Failed to convert DataFusion Physical Plan: {e}"))?; // 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?; + let (distributed_plan, distributed_stages) = + execution_planning(physical_plan.clone(), 8192, Some(2)).await?; Ok(QueryPlanBase { query_id, @@ -246,9 +256,8 @@ impl QueryPlanner { #[cfg(test)] mod tests { use super::*; - use std::{fs::{File}, path::Path}; use std::io::BufReader; - + use std::{fs::File, path::Path}; // ////////////////////////////////////////////////////////////// // Test helper functions @@ -256,8 +265,10 @@ mod tests { /// 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_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)) @@ -303,12 +314,16 @@ mod tests { let planner = QueryPlanner::new(); // 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"), - )).expect("failed to parse json"); - + 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"), + ), + ) + .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 @@ -316,7 +331,10 @@ mod tests { 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))"); + 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 { From 035c45b70137ef33c362ddd68bfbc55597ddd6cc Mon Sep 17 00:00:00 2001 From: LiaCastaneda Date: Tue, 8 Jul 2025 17:07:52 +0200 Subject: [PATCH 3/4] Change method name --- src/flight.rs | 6 +++--- src/proxy_service.rs | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/flight.rs b/src/flight.rs index cd5c290d..26febf2c 100644 --- a/src/flight.rs +++ b/src/flight.rs @@ -134,7 +134,7 @@ pub trait FlightSqlHandler: Send + Sync { request: Request, ) -> Result, Status>; - async fn get_flight_info_substrait( + async fn get_flight_info_substrait_plan( &self, query: arrow_flight::sql::CommandStatementSubstraitPlan, request: Request, @@ -176,10 +176,10 @@ impl FlightSqlService for FlightSqlServ { request: Request, ) -> Result, Status> { self.handler - .get_flight_info_substrait(query, request) + .get_flight_info_substrait_plan(query, request) .await .inspect_err(|e| { - error!("Error in get_flight_info_substrait: {:?}", e); + error!("Error in get_flight_info_substrait_plan: {:?}", e); }) } diff --git a/src/proxy_service.rs b/src/proxy_service.rs index cadff7e2..ad21edb5 100644 --- a/src/proxy_service.rs +++ b/src/proxy_service.rs @@ -101,7 +101,7 @@ impl FlightSqlHandler for DfRayProxyHandler { } } - async fn get_flight_info_substrait( + async fn get_flight_info_substrait_plan( &self, substrait: arrow_flight::sql::CommandStatementSubstraitPlan, _request: Request, From b0dce09869b6031fff41b22fffddc6a44a30f53b Mon Sep 17 00:00:00 2001 From: LiaCastaneda Date: Tue, 8 Jul 2025 18:14:09 +0200 Subject: [PATCH 4/4] minor changes --- src/query_planner.rs | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/query_planner.rs b/src/query_planner.rs index 4d1c6057..035e9e57 100644 --- a/src/query_planner.rs +++ b/src/query_planner.rs @@ -1,11 +1,12 @@ use std::sync::Arc; use anyhow::anyhow; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::{ logical_expr::LogicalPlan, physical_plan::ExecutionPlan, prelude::SessionContext, }; use datafusion_substrait::logical_plan::consumer::from_substrait_plan; +use datafusion_substrait::substrait::proto::Plan; use crate::{ explain::{is_explain_query, DistributedExplainExec}, @@ -131,10 +132,7 @@ impl QueryPlanner { } /// Prepare a distributed query (Substrait entry point) - pub async fn prepare_substrait_query( - &self, - substrait_plan: datafusion_substrait::substrait::proto::Plan, - ) -> Result { + pub async fn prepare_substrait_query(&self, substrait_plan: Plan) -> Result { let base_result = self .prepare_substrait_query_base(substrait_plan, "SUBSTRAIT") .await?; @@ -143,7 +141,7 @@ impl QueryPlanner { pub async fn prepare_substrait_query_base( &self, - substrait_plan: datafusion_substrait::substrait::proto::Plan, + substrait_plan: Plan, query_type: &str, ) -> Result { debug!( @@ -219,7 +217,6 @@ impl QueryPlanner { ); // create the schema for EXPLAIN results - use arrow::datatypes::{DataType, Field, Schema}; let schema = Arc::new(Schema::new(vec![ Field::new("plan_type", DataType::Utf8, false), Field::new("plan", DataType::Utf8, false),