From d82988d698f70df49c815046917d27d413dba50d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 26 Jun 2026 12:23:02 -0600 Subject: [PATCH 1/8] chore(deps): upgrade to DataFusion 54 Bump the workspace to the published datafusion 54 (and arrow/arrow-flight 58.3, object_store 0.13.2, rustyline 18 in ballista-cli), replacing the branch-54 git-rev pin from the earlier draft. API migration to 54: - as_any removed from ExecutionPlan/DataSource/PhysicalExpr: downcast directly on the trait object (and upcast dyn ShuffleWriter to dyn ExecutionPlan); as_any retained where it still exists (Arrow arrays, UserDefinedLogicalNode, ExtensionOptions). - ExecutionPlan::partition_statistics returns Result>. - parse_protobuf_partitioning/_hash_partitioning take a PhysicalPlanDecodeContext; serialize_partitioning takes the codec plus the proto converter. - TaskContext::new and FunctionRegistry gain higher-order-function arguments (wired empty in Ballista). - BatchPartitioner::new_hash_partitioner is fallible. Test updates for 54 behavior changes: - approx_percentile_cont_with_weight result; shuffle-writer tests made robust to hash distribution; execution-graph/dot tests use Ballista session config so the small join stays Partitioned; regenerated AQE plan snapshots (projection folded into HashJoinExec). --- Cargo.lock | 430 ++++++++---------- Cargo.toml | 18 +- ballista-cli/Cargo.toml | 2 +- ballista/client/tests/context_basic.rs | 2 +- ballista/core/src/diagram.rs | 40 +- .../core/src/execution_plans/chaos_exec.rs | 7 +- .../distributed_explain_analyze.rs | 7 - .../src/execution_plans/distributed_query.rs | 6 - .../src/execution_plans/shuffle_reader.rs | 13 +- .../src/execution_plans/shuffle_writer.rs | 47 +- .../execution_plans/sort_shuffle/writer.rs | 10 +- .../src/execution_plans/unresolved_shuffle.rs | 5 - ballista/core/src/planner.rs | 2 - ballista/core/src/registry.rs | 14 + ballista/core/src/serde/mod.rs | 32 +- .../core/src/serde/scheduler/from_proto.rs | 4 + ballista/executor/src/collect.rs | 8 +- ballista/executor/src/execution_engine.rs | 6 +- ballista/executor/src/execution_loop.rs | 2 + ballista/executor/src/executor.rs | 5 - ballista/executor/src/executor_server.rs | 2 + .../src/physical_optimizer/join_selection.rs | 109 +++-- ballista/scheduler/src/planner.rs | 73 ++- ballista/scheduler/src/state/aqe/adapter.rs | 7 +- .../src/state/aqe/execution_plan/adaptive.rs | 5 - .../state/aqe/execution_plan/dynamic_join.rs | 10 +- .../src/state/aqe/execution_plan/exchange.rs | 13 +- .../state/aqe/optimizer_rule/chaos_exec.rs | 3 +- .../aqe/optimizer_rule/coalesce_partitions.rs | 9 +- .../optimizer_rule/distributed_exchange.rs | 127 ++---- .../aqe/optimizer_rule/join_selection.rs | 31 +- .../aqe/optimizer_rule/propagate_empty.rs | 44 +- ballista/scheduler/src/state/aqe/planner.rs | 18 +- .../src/state/aqe/test/alter_stages.rs | 60 ++- .../src/state/aqe/test/plan_to_stages.rs | 5 +- .../src/state/distributed_explain.rs | 2 +- .../scheduler/src/state/execution_graph.rs | 16 +- .../src/state/execution_graph_dot.rs | 77 ++-- .../scheduler/src/state/execution_stage.rs | 4 +- ballista/scheduler/src/test_utils.rs | 7 +- 40 files changed, 525 insertions(+), 757 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index bca6161665..ec247e01fa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -112,7 +112,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -123,7 +123,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -132,35 +132,6 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" -[[package]] -name = "apache-avro" -version = "0.21.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36fa98bc79671c7981272d91a8753a928ff6a1cd8e4f20a44c45bd5d313840bf" -dependencies = [ - "bigdecimal", - "bon", - "bzip2", - "crc32fast", - "digest 0.10.7", - "liblzma", - "log", - "miniz_oxide 0.8.9", - "num-bigint", - "quad-rand", - "rand 0.9.4", - "regex-lite", - "serde", - "serde_bytes", - "serde_json", - "snap", - "strum 0.27.2", - "strum_macros 0.27.2", - "thiserror 2.0.18", - "uuid", - "zstd", -] - [[package]] name = "approx" version = "0.5.1" @@ -260,6 +231,30 @@ dependencies = [ "num-traits", ] +[[package]] +name = "arrow-avro" +version = "58.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "049230728cd6e093088c8d231b4beede184e35cad7777c1505c0d5a8571f4376" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-schema", + "bytes", + "bzip2", + "crc", + "flate2", + "indexmap 2.14.0", + "liblzma", + "rand 0.9.4", + "serde", + "serde_json", + "snap", + "strum_macros", + "uuid", + "zstd", +] + [[package]] name = "arrow-buffer" version = "58.3.0" @@ -583,7 +578,7 @@ dependencies = [ "fastrand", "hex", "http 1.4.2", - "sha1", + "sha1 0.10.6", "time", "tokio", "tracing", @@ -1297,7 +1292,6 @@ dependencies = [ "num-bigint", "num-integer", "num-traits", - "serde", ] [[package]] @@ -1457,31 +1451,6 @@ dependencies = [ "time", ] -[[package]] -name = "bon" -version = "3.9.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2f04f6fef12d70d42a77b1433c9e0f065238479a6cefc4f5bab105e9873a3c3" -dependencies = [ - "bon-macros", - "rustversion", -] - -[[package]] -name = "bon-macros" -version = "3.9.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d0bd4c2f75335ad98052a37efb54f428b492f64340257143b3429c8a508fa7b" -dependencies = [ - "darling 0.23.0", - "ident_case", - "prettyplease", - "proc-macro2", - "quote", - "rustversion", - "syn 2.0.117", -] - [[package]] name = "brotli" version = "8.0.3" @@ -1921,6 +1890,21 @@ dependencies = [ "libc", ] +[[package]] +name = "crc" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-catalog" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" + [[package]] name = "crc32fast" version = "1.5.0" @@ -2211,14 +2195,13 @@ dependencies = [ [[package]] name = "datafusion" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93db0e623840612f7f2cd757f7e8a8922064192363732c88692e0870016e141b" +checksum = "997a31e15872606a49478e670c58302094c97cb96abb0a7d60720f8e92170040" dependencies = [ "arrow", "arrow-schema", "async-trait", - "bytes", "bzip2", "chrono", "datafusion-catalog", @@ -2249,14 +2232,13 @@ dependencies = [ "datafusion-sql", "flate2", "futures", + "indexmap 2.14.0", "itertools 0.14.0", "liblzma", "log", "object_store", "parking_lot", "parquet", - "rand 0.9.4", - "regex", "sqlparser", "tempfile", "tokio", @@ -2267,9 +2249,9 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37cefde60b26a7f4ff61e9d2ff2833322f91df2b568d7238afe67bde5bdffb66" +checksum = "f7dd61161508f8f5fa1107774ea687bd753c22d83a32eebf963549f89de14139" dependencies = [ "arrow", "async-trait", @@ -2292,9 +2274,9 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "17e112307715d6a7a331111a4c2330ff54bc237183511c319e3708a4cff431fb" +checksum = "897c70f871277f9ce99aa38347be0d679bbe3e617156c4d2a8378cec8a2a0891" dependencies = [ "arrow", "async-trait", @@ -2315,9 +2297,9 @@ dependencies = [ [[package]] name = "datafusion-cli" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "84a22c001ad1ac11cda09dab69b151eef5b1a992e23bc524ab0d1e63e5dea327" +checksum = "2a9533845ab888508718c2ad540cff205d1111525a9952def8b0e584b6453784" dependencies = [ "arrow", "async-trait", @@ -2343,17 +2325,17 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d72a11ca44a95e1081870d3abb80c717496e8a7acb467a1d3e932bb636af5cc2" +checksum = "121c9ded5d87d9172319e006f2afdb9928d72dbacd6a90a458d8acb1e3b43a65" dependencies = [ - "ahash", - "apache-avro", "arrow", "arrow-ipc", + "arrow-schema", "chrono", + "foldhash 0.2.0", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "hex", "indexmap 2.14.0", "itertools 0.14.0", @@ -2361,18 +2343,18 @@ dependencies = [ "log", "object_store", "parquet", - "paste", "recursive", "sqlparser", "tokio", + "uuid", "web-time", ] [[package]] name = "datafusion-common-runtime" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89f4afaed29670ec4fd6053643adc749fe3f4bc9d1ce1b8c5679b22c67d12def" +checksum = "981b9dae74f78ee3d9f714fb49b01919eab975461b56149510c3ba9ea11287d1" dependencies = [ "futures", "log", @@ -2381,9 +2363,9 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9fb386e1691355355a96419978a0022b7947b44d4a24a6ea99f00b6b485cbb6" +checksum = "ffd7d295b2ec7c00d8a56562f41ed41062cf0af75549ed891c12a0a09eddfefe" dependencies = [ "arrow", "async-compression", @@ -2407,6 +2389,7 @@ dependencies = [ "liblzma", "log", "object_store", + "parking_lot", "rand 0.9.4", "tokio", "tokio-util", @@ -2416,9 +2399,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ffa6c52cfed0734c5f93754d1c0175f558175248bf686c944fb05c373e5fc096" +checksum = "552b0b3f342f7ec41b3fbd70f6339dc82a30cfd0349e7f280e7852528085349f" dependencies = [ "arrow", "arrow-ipc", @@ -2440,29 +2423,28 @@ dependencies = [ [[package]] name = "datafusion-datasource-avro" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a579c3bd290c66ea4b269493e75e8a3ed42c9c895a651f10210a29538aee50c4" +checksum = "fb517d08967d536284ce70afb5fe8583133779249f2d7b90587d339741a7f195" dependencies = [ - "apache-avro", "arrow", + "arrow-avro", "async-trait", "bytes", "datafusion-common", "datafusion-datasource", - "datafusion-physical-expr-common", + "datafusion-physical-expr-adapter", "datafusion-physical-plan", "datafusion-session", "futures", - "num-traits", "object_store", ] [[package]] name = "datafusion-datasource-csv" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "503f29e0582c1fc189578d665ff57d9300da1f80c282777d7eb67bb79fb8cdca" +checksum = "68850aa426b897e879c8b87e512ea8124f1d0a2869a4e51808ddaaddf1bc0ada" dependencies = [ "arrow", "async-trait", @@ -2483,9 +2465,9 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e33804749abc8d0c8cb7473228483cb8070e524c6f6086ee1b85a64debe2b3d2" +checksum = "402f93242ae08ef99139ee2c528a49d087efe88d5c7b2c3ff5480855a40ce54f" dependencies = [ "arrow", "async-trait", @@ -2500,16 +2482,15 @@ dependencies = [ "datafusion-session", "futures", "object_store", - "serde_json", "tokio", "tokio-stream", ] [[package]] name = "datafusion-datasource-parquet" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32a8e0365e0e08e8ff94d912f0ababcf9065a1a304018ba90b1fc83c855b4997" +checksum = "ffd2499c1bee0eeccf6a57156105700eeeb17bc701899ac719183c4e74231450" dependencies = [ "arrow", "async-trait", @@ -2519,6 +2500,7 @@ dependencies = [ "datafusion-datasource", "datafusion-execution", "datafusion-expr", + "datafusion-functions", "datafusion-functions-aggregate-common", "datafusion-physical-expr", "datafusion-physical-expr-adapter", @@ -2537,20 +2519,19 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8de6ac0df1662b9148ad3c987978b32cbec7c772f199b1d53520c8fa764a87ee" +checksum = "cb9e7e5d11130c48c8bd4e80c79a9772dd28ce6dc330baca9246205d245b9e2e" [[package]] name = "datafusion-execution" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c03c7fbdaefcca4ef6ffe425a5fc2325763bfb426599bb0bf4536466efabe709" +checksum = "37a8643ab852eb68864e1b72ae789e8066282dce48eea6347ffb0aee33d1ccc0" dependencies = [ "arrow", "arrow-buffer", "async-trait", - "chrono", "dashmap", "datafusion-common", "datafusion-expr", @@ -2567,11 +2548,12 @@ dependencies = [ [[package]] name = "datafusion-expr" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "574b9b6977fedbd2a611cbff12e5caf90f31640ad9dc5870f152836d94bad0dd" +checksum = "6932f4d71eed9c8d9341476a2b845aadfabde5495d08dbcd8fc23881f49fa7a0" dependencies = [ "arrow", + "arrow-schema", "async-trait", "chrono", "datafusion-common", @@ -2582,7 +2564,6 @@ dependencies = [ "datafusion-physical-expr-common", "indexmap 2.14.0", "itertools 0.14.0", - "paste", "recursive", "serde_json", "sqlparser", @@ -2590,22 +2571,21 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d7c3adf3db8bf61e92eb90cb659c8e8b734593a8f7c8e12a843c7ddba24b87e" +checksum = "0225491839a31b1f7d2cb8092c2d50792e2fe1c1724e4e6d08e011f5feaf4ed2" dependencies = [ "arrow", "datafusion-common", "indexmap 2.14.0", "itertools 0.14.0", - "paste", ] [[package]] name = "datafusion-functions" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f28aa4e10384e782774b10e72aca4d93ef7b31aa653095d9d4536b0a3dbc51b6" +checksum = "14872c47bfc3d21e53ec82f57074e6987a15941c1e2f43cde4ac6ae2746634e3" dependencies = [ "arrow", "arrow-buffer", @@ -2620,26 +2600,25 @@ dependencies = [ "datafusion-expr", "datafusion-expr-common", "datafusion-macros", + "datafusion-physical-expr-common", "hex", "itertools 0.14.0", "log", - "md-5 0.10.6", + "md-5 0.11.0", "memchr", "num-traits", "rand 0.9.4", "regex", - "sha2 0.10.9", - "unicode-segmentation", + "sha2 0.11.0", "uuid", ] [[package]] name = "datafusion-functions-aggregate" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00aa6217e56098ba84e0a338176fe52f0a84cca398021512c6c8c5eff806d0ad" +checksum = "75a2ca14e1b609be21e657e2d3130b2f446456b08393b377bb721a33952d2e09" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-doc", @@ -2649,19 +2628,18 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr", "datafusion-physical-expr-common", + "foldhash 0.2.0", "half", "log", "num-traits", - "paste", ] [[package]] name = "datafusion-functions-aggregate-common" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b511250349407db7c43832ab2de63f5557b19a20dfd236b39ca2c04468b50d47" +checksum = "1ece74ba09092d2ef9c9b54a38445450aea292a1f8b04faf531936b723a24b3c" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-expr-common", @@ -2670,9 +2648,9 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef13a858e20d50f0a9bb5e96e7ac82b4e7597f247515bccca4fdd2992df0212a" +checksum = "3f3e3f9ee8ca59bf70518802107de6f1b88a9509efdc629fadc5de9d6b2d5ef5" dependencies = [ "arrow", "arrow-ord", @@ -2686,34 +2664,34 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-macros", "datafusion-physical-expr-common", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "itertools 0.14.0", "itoa", "log", - "paste", + "memchr", ] [[package]] name = "datafusion-functions-table" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b40d3f5bbb3905f9ccb1ce9485a9595c77b69758a7c24d3ba79e334ff51e7e" +checksum = "89161dffc22cf2b50f9f4b1bee83b5221d3b4ed7c2e37fd7aa2b22a5297b3a26" dependencies = [ "arrow", "async-trait", "datafusion-catalog", "datafusion-common", "datafusion-expr", + "datafusion-physical-expr", "datafusion-physical-plan", "parking_lot", - "paste", ] [[package]] name = "datafusion-functions-window" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4e88ec9d57c9b685d02f58bfee7be62d72610430ddcedb82a08e5d9925dbfb6" +checksum = "d7339345b226b3874037708bf5023ba1c2de705128f8457a095aae5ae9cb9c78" dependencies = [ "arrow", "datafusion-common", @@ -2724,14 +2702,13 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "log", - "paste", ] [[package]] name = "datafusion-functions-window-common" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8307bb93519b1a91913723a1130cfafeee3f72200d870d88e91a6fc5470ede5c" +checksum = "fa84836dc2392df6f43d6a29d37fb56a8ebdc8b3f4e10ae8dc15861fd20278fb" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -2739,9 +2716,9 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e367e6a71051d0ebdd29b2f85d12059b38b1d1f172c6906e80016da662226bd" +checksum = "587164e03ad68732aa9e7bfe5686e3f25970d4c64fd4bd80790749840892dae5" dependencies = [ "datafusion-doc", "quote", @@ -2750,9 +2727,9 @@ dependencies = [ [[package]] name = "datafusion-optimizer" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e929015451a67f77d9d8b727b2bf3a40c4445fdef6cdc53281d7d97c76888ace" +checksum = "77f20e8cf9e8654d92f4c16b24c487353ee5bf153ffc12d5772cd399ab8cd281" dependencies = [ "arrow", "chrono", @@ -2770,11 +2747,10 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b1e68aba7a4b350401cfdf25a3d6f989ad898a7410164afe9ca52080244cb59" +checksum = "f015a4a82f6f7ff7e1d8d4bf3870a936752fa38b17705dfcc14adef95aa8922c" dependencies = [ - "ahash", "arrow", "datafusion-common", "datafusion-expr", @@ -2782,11 +2758,10 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-physical-expr-common", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "indexmap 2.14.0", "itertools 0.14.0", "parking_lot", - "paste", "petgraph", "recursive", "tokio", @@ -2794,9 +2769,9 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea22315f33cf2e0adc104e8ec42e285f6ed93998d565c65e82fec6a9ee9f9db4" +checksum = "51e6ffff8acdfe54e0ea15ccf38115c4a9184433b0439f42907637928d00a235" dependencies = [ "arrow", "datafusion-common", @@ -2809,26 +2784,26 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-common" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b04b45ea8ad3ac2d78f2ea2a76053e06591c9629c7a603eda16c10649ecf4362" +checksum = "7967a3e171c6a4bf09474b3f7a14f1a3db13ed1714ba12156f33fcce2bba54e8" dependencies = [ - "ahash", "arrow", "chrono", "datafusion-common", "datafusion-expr-common", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "indexmap 2.14.0", "itertools 0.14.0", "parking_lot", + "pin-project", ] [[package]] name = "datafusion-physical-optimizer" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7cb13397809a425918f608dfe8653f332015a3e330004ab191b4404187238b95" +checksum = "59ff803e2a96054cb6d83f35f9e60fd4f42eac515e1932bd1b2dbc91d5fcbf36" dependencies = [ "arrow", "datafusion-common", @@ -2845,12 +2820,13 @@ dependencies = [ [[package]] name = "datafusion-physical-plan" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5edc023675791af9d5fb4cc4c24abf5f7bd3bd4dcf9e5bd90ea1eff6976dcc79" +checksum = "776ee54d47d15bdb126452f9ca17b03761e3b004682914beaedd3f86eb507fbc" dependencies = [ - "ahash", "arrow", + "arrow-data", + "arrow-ipc", "arrow-ord", "arrow-schema", "async-trait", @@ -2865,7 +2841,7 @@ dependencies = [ "datafusion-physical-expr-common", "futures", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "indexmap 2.14.0", "itertools 0.14.0", "log", @@ -2877,9 +2853,9 @@ dependencies = [ [[package]] name = "datafusion-proto" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a387aaef949dc16bb6abc81bd1af850ec7449183aef011214f9724957495738" +checksum = "9dd15a1ba5d3af93808241065c6c44dbca8296a189845e8a587c45c07bf0ffae" dependencies = [ "arrow", "chrono", @@ -2900,14 +2876,13 @@ dependencies = [ "datafusion-proto-common", "object_store", "prost", - "rand 0.9.4", ] [[package]] name = "datafusion-proto-common" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16e614c7c53a9c304c6a850b821010bb492e57300311835f1180613f9d2c63d9" +checksum = "90042982cf9462eb06a0b81f92efa4188dae871e7ea3ab8dc61aa9c9349b2530" dependencies = [ "arrow", "datafusion-common", @@ -2916,9 +2891,9 @@ dependencies = [ [[package]] name = "datafusion-pruning" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac8c76860e355616555081cab5968cec1af7a80701ff374510860bcd567e365a" +checksum = "d5fb9e5774660aa69c3ba93c610f175f75b65cb8c3776edb3626de8f3a4f4ee3" dependencies = [ "arrow", "datafusion-common", @@ -2927,15 +2902,14 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", - "itertools 0.14.0", "log", ] [[package]] name = "datafusion-session" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5412111aa48e2424ba926112e192f7a6b7e4ccb450145d25ce5ede9f19dc491e" +checksum = "15ce715fa2a61f4623cc234bcc14a3ef6a91f189128d5b14b468a6a17cdfc417" dependencies = [ "async-trait", "datafusion-common", @@ -2947,9 +2921,9 @@ dependencies = [ [[package]] name = "datafusion-spark" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e059dcf8544da0d6598d0235be3cc29c209094a5976b2e4822e4a2cf91c2b5c5" +checksum = "390bb0bf37cb2b95ffd65eacd66f60df50793d1f94097799e416f39477a51957" dependencies = [ "arrow", "bigdecimal", @@ -2962,21 +2936,24 @@ dependencies = [ "datafusion-expr", "datafusion-functions", "datafusion-functions-aggregate", + "datafusion-functions-aggregate-common", "datafusion-functions-nested", "log", + "num-traits", "percent-encoding", "rand 0.9.4", "serde_json", - "sha1", - "sha2 0.10.9", + "sha1 0.11.0", + "sha2 0.11.0", + "twox-hash", "url", ] [[package]] name = "datafusion-sql" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa0d133ddf8b9b3b872acac900157f783e7b879fe9a6bccf389abebbfac45ec1" +checksum = "6094ad36a3ed6d7ac87b20b479b2d0b118250f66cf997603829fdc65b44a7099" dependencies = [ "arrow", "bigdecimal", @@ -2993,9 +2970,9 @@ dependencies = [ [[package]] name = "datafusion-substrait" -version = "53.1.0" +version = "54.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98494539a5468979cc42d86c7bc5f0f8cb71ee5c742694c26fc34efdd29dd2e5" +checksum = "3b22c8f8c72d317e54fad6f85c0ef6d1e1da53cc7faadc7eea8daf0f8d86d4f2" dependencies = [ "async-recursion", "async-trait", @@ -3121,7 +3098,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3216,9 +3193,9 @@ dependencies = [ [[package]] name = "endian-type" -version = "0.1.2" +version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d" +checksum = "869b0adbda23651a9c5c0c3d270aac9fcb52e8622a8f2b17e57802d7791962f2" [[package]] name = "env_filter" @@ -3256,7 +3233,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -3306,17 +3283,6 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" -[[package]] -name = "fd-lock" -version = "4.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce92ff622d6dadf7349484f42c93271a0d49b7cc4d466a936405bacbe10aa78" -dependencies = [ - "cfg-if", - "rustix 1.1.4", - "windows-sys 0.59.0", -] - [[package]] name = "ferroid" version = "2.0.0" @@ -4705,9 +4671,9 @@ dependencies = [ [[package]] name = "nix" -version = "0.30.1" +version = "0.31.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" +checksum = "cf20d2fde8ff38632c426f1165ed7436270b44f199fc55284c38276f9db47c3d" dependencies = [ "bitflags 2.13.0", "cfg-if", @@ -4740,7 +4706,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -4765,7 +4731,6 @@ checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9" dependencies = [ "num-integer", "num-traits", - "serde", ] [[package]] @@ -5581,12 +5546,6 @@ dependencies = [ "pulldown-cmark", ] -[[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.39.4" @@ -5682,9 +5641,9 @@ checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" [[package]] name = "radix_trie" -version = "0.2.1" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd" +checksum = "3b4431027dcd37fc2a73ef740b5f233aa805897935b8bce0195e41bbf9a3289a" dependencies = [ "endian-type", "nibble_vec", @@ -5783,7 +5742,7 @@ dependencies = [ "lru 0.18.0", "palette", "serde", - "strum 0.28.0", + "strum", "thiserror 2.0.18", "unicode-segmentation", "unicode-truncate", @@ -5859,7 +5818,7 @@ dependencies = [ "line-clipping", "ratatui-core", "serde", - "strum 0.28.0", + "strum", "time", "unicode-segmentation", "unicode-width 0.2.2", @@ -6181,7 +6140,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -6240,7 +6199,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -6269,24 +6228,23 @@ checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] name = "rustyline" -version = "17.0.2" +version = "18.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e902948a25149d50edc1a8e0141aad50f54e22ba83ff988cf8f7c9ef07f50564" +checksum = "53f6a737db68eb1a8ccff86b584b2fc13eca6a7bb6f78ebc7c529547e3ab9684" dependencies = [ "bitflags 2.13.0", "cfg-if", "clipboard-win", - "fd-lock", "home", "libc", "log", "memchr", - "nix 0.30.1", + "nix 0.31.3", "radix_trie", "unicode-segmentation", "unicode-width 0.2.2", "utf8parse", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -6416,16 +6374,6 @@ dependencies = [ "serde_derive", ] -[[package]] -name = "serde_bytes" -version = "0.11.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a5d440709e79d88e51ac01c4b72fc6cb7314017bb7da9eeff678aa94c10e3ea8" -dependencies = [ - "serde", - "serde_core", -] - [[package]] name = "serde_core" version = "1.0.228" @@ -6463,6 +6411,7 @@ version = "1.0.150" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" dependencies = [ + "indexmap 2.14.0", "itoa", "memchr", "serde", @@ -6572,6 +6521,17 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sha1" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", +] + [[package]] name = "sha2" version = "0.10.9" @@ -6708,14 +6668,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] name = "sqlparser" -version = "0.61.0" +version = "0.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dbf5ea8d4d7c808e1af1cbabebca9a2abe603bcefc22294c5b95018d53200cb7" +checksum = "13c6d1b651dc4edf07eead2a0c6c78016ce971bc2c10da5266861b13f25e7cec" dependencies = [ "log", "recursive", @@ -6749,7 +6709,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -6811,31 +6771,13 @@ dependencies = [ "syn 1.0.109", ] -[[package]] -name = "strum" -version = "0.27.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af23d6f6c1a224baef9d3f61e287d2761385a5b88fdab4eb4c6f11aeb54c4bcf" - [[package]] name = "strum" version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9628de9b8791db39ceda2b119bbe13134770b56c138ec1d3af810d045c04f9bd" dependencies = [ - "strum_macros 0.28.0", -] - -[[package]] -name = "strum_macros" -version = "0.27.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7695ce3845ea4b33927c055a39dc438a45b059f7c1b3d91d38d10355fb8cbca7" -dependencies = [ - "heck 0.5.0", - "proc-macro2", - "quote", - "syn 2.0.117", + "strum_macros", ] [[package]] @@ -6852,11 +6794,12 @@ dependencies = [ [[package]] name = "substrait" -version = "0.62.2" +version = "0.63.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62fc4b483a129b9772ccb9c3f7945a472112fdd9140da87f8a4e7f1d44e045d0" +checksum = "e620ff4d5c02fd6f7752931aa74b16a26af66a63022cc1ad412c77edbe0bab47" dependencies = [ "heck 0.5.0", + "indexmap 2.14.0", "pbjson", "pbjson-build", "pbjson-types", @@ -6977,10 +6920,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -6993,7 +6936,7 @@ dependencies = [ "parking_lot", "rustix 1.1.4", "signal-hook", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -7605,6 +7548,9 @@ name = "twox-hash" version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" +dependencies = [ + "rand 0.9.4", +] [[package]] name = "typenum" @@ -8101,7 +8047,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 57346eaf05..79fc396043 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,20 +26,20 @@ edition = "2024" rust-version = "1.88.0" [workspace.dependencies] -arrow = { version = "58", features = ["ipc_compression"] } -arrow-flight = { version = "58", features = ["flight-sql-experimental"] } +arrow = { version = "58.3", features = ["ipc_compression"] } +arrow-flight = { version = "58.3", features = ["flight-sql-experimental"] } clap = { version = "4.5", features = ["derive", "cargo"] } -datafusion = "53" -datafusion-cli = "53" -datafusion-proto = "53" -datafusion-proto-common = "53" -datafusion-spark = "53" -datafusion-substrait = "53" +datafusion = "54" +datafusion-cli = "54" +datafusion-proto = "54" +datafusion-proto-common = "54" +datafusion-spark = "54" +datafusion-substrait = "54" insta = "1.47" itertools = "0.15" -object_store = "0.13" +object_store = "0.13.2" prost = "0.14" prost-types = "0.14" rstest = { version = "0.26" } diff --git a/ballista-cli/Cargo.toml b/ballista-cli/Cargo.toml index 6796f49eeb..683a3c2453 100644 --- a/ballista-cli/Cargo.toml +++ b/ballista-cli/Cargo.toml @@ -36,7 +36,7 @@ crate-type = ["cdylib", "rlib"] ballista = { path = "../ballista/client", version = "53.0.0", features = ["standalone"], optional = true } datafusion = { workspace = true, optional = true } datafusion-cli = { workspace = true, optional = true } -rustyline = { version = "17.0.1", optional = true } +rustyline = { version = "18.0.0", optional = true } mimalloc = { workspace = true, optional = true } # TUI/web shared deps diff --git a/ballista/client/tests/context_basic.rs b/ballista/client/tests/context_basic.rs index 5e9c29bb38..395f570313 100644 --- a/ballista/client/tests/context_basic.rs +++ b/ballista/client/tests/context_basic.rs @@ -449,7 +449,7 @@ mod basic { "+-------------------------------------------------------------------+", "| approx_percentile_cont_with_weight(test.id,Int64(2),Float64(0.5)) |", "+-------------------------------------------------------------------+", - "| 3 |", + "| 3.5 |", "+-------------------------------------------------------------------+", ]; assert_result_eq(expected, &res); diff --git a/ballista/core/src/diagram.rs b/ballista/core/src/diagram.rs index 5498c1161c..79604e2a13 100644 --- a/ballista/core/src/diagram.rs +++ b/ballista/core/src/diagram.rs @@ -84,43 +84,27 @@ fn build_exec_plan_diagram( id: &mut AtomicUsize, draw_entity: bool, ) -> Result { - let operator_str = if plan.as_any().downcast_ref::().is_some() { + let operator_str = if plan.downcast_ref::().is_some() { "AggregateExec" - } else if plan.as_any().downcast_ref::().is_some() { + } else if plan.downcast_ref::().is_some() { "SortExec" - } else if plan.as_any().downcast_ref::().is_some() { + } else if plan.downcast_ref::().is_some() { "ProjectionExec" - } else if plan.as_any().downcast_ref::().is_some() { + } else if plan.downcast_ref::().is_some() { "HashJoinExec" - } else if plan.as_any().downcast_ref::().is_some() { + } else if plan.downcast_ref::().is_some() { "DataSourceExec" - } else if plan.as_any().downcast_ref::().is_some() { + } else if plan.downcast_ref::().is_some() { "FilterExec" - } else if plan.as_any().downcast_ref::().is_some() { + } else if plan.downcast_ref::().is_some() { "ShuffleWriterExec" - } else if plan - .as_any() - .downcast_ref::() - .is_some() - { + } else if plan.downcast_ref::().is_some() { "SortShuffleWriterExec" - } else if plan - .as_any() - .downcast_ref::() - .is_some() - { + } else if plan.downcast_ref::().is_some() { "UnresolvedShuffleExec" - } else if plan - .as_any() - .downcast_ref::() - .is_some() - { + } else if plan.downcast_ref::().is_some() { "CoalesceBatchesExec" - } else if plan - .as_any() - .downcast_ref::() - .is_some() - { + } else if plan.downcast_ref::().is_some() { "CoalescePartitionsExec" } else { warn!("Unknown: {plan:?}"); @@ -137,7 +121,7 @@ fn build_exec_plan_diagram( )?; } for child in plan.children() { - if let Some(shuffle) = child.as_any().downcast_ref::() { + if let Some(shuffle) = child.downcast_ref::() { if !draw_entity { writeln!( w, diff --git a/ballista/core/src/execution_plans/chaos_exec.rs b/ballista/core/src/execution_plans/chaos_exec.rs index fda436ec58..569f4da724 100644 --- a/ballista/core/src/execution_plans/chaos_exec.rs +++ b/ballista/core/src/execution_plans/chaos_exec.rs @@ -42,7 +42,6 @@ use datafusion::physical_plan::{ use futures::StreamExt; use rand::rngs::StdRng; use rand::{RngExt, SeedableRng}; -use std::any::Any; use std::sync::Arc; /// Physical execution plan node that randomly injects failures for chaos/robustness testing. @@ -124,10 +123,6 @@ impl ExecutionPlan for ChaosExec { "ChaosExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn properties(&self) -> &Arc { self.input.properties() } @@ -218,7 +213,7 @@ impl ExecutionPlan for ChaosExec { self.input.metrics() } - fn partition_statistics(&self, partition: Option) -> Result { + fn partition_statistics(&self, partition: Option) -> Result> { self.input.partition_statistics(partition) } diff --git a/ballista/core/src/execution_plans/distributed_explain_analyze.rs b/ballista/core/src/execution_plans/distributed_explain_analyze.rs index af02b0a22c..3d7ba2c867 100644 --- a/ballista/core/src/execution_plans/distributed_explain_analyze.rs +++ b/ballista/core/src/execution_plans/distributed_explain_analyze.rs @@ -37,7 +37,6 @@ use datafusion::physical_plan::{ }; use datafusion_proto::logical_plan::AsLogicalPlan; use futures::StreamExt; -use std::any::Any; use std::convert::TryInto; use std::marker::PhantomData; use std::sync::Arc; @@ -108,10 +107,6 @@ impl ExecutionPlan for DistributedExplainAnalyzeExec "DistributedExplainAnalyzeExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn properties(&self) -> &Arc { &self.properties } @@ -133,7 +128,6 @@ impl ExecutionPlan for DistributedExplainAnalyzeExec let query_exec = children.pop().unwrap(); if query_exec - .as_any() .downcast_ref::>() .is_some() { @@ -173,7 +167,6 @@ impl ExecutionPlan for DistributedExplainAnalyzeExec } let job_id = query_exec - .as_any() .downcast_ref::>() .ok_or_else(|| { DataFusionError::Internal( diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index 0181678104..a7d79ef824 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -49,7 +49,6 @@ use datafusion_proto::logical_plan::{ use futures::{Stream, StreamExt, TryFutureExt, TryStreamExt}; use log::{debug, error, info}; use parking_lot::Mutex; -use std::any::Any; use std::fmt::Debug; use std::marker::PhantomData; use std::sync::Arc; @@ -177,10 +176,6 @@ impl ExecutionPlan for DistributedQueryExec { "DistributedQueryExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { self.plan.schema().as_arrow().clone().into() } @@ -842,7 +837,6 @@ mod test { let new_exec = exec.clone().with_new_children(vec![]).unwrap(); let new_exec = new_exec - .as_any() .downcast_ref::>() .unwrap(); diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs index 556ccc8090..a9955ac5f5 100644 --- a/ballista/core/src/execution_plans/shuffle_reader.rs +++ b/ballista/core/src/execution_plans/shuffle_reader.rs @@ -47,7 +47,6 @@ use itertools::Itertools; use log::{debug, error, trace}; use rand::prelude::SliceRandom; use rand::rng; -use std::any::Any; use std::collections::HashMap; use std::fmt::Debug; use std::fs::File; @@ -311,10 +310,6 @@ impl ExecutionPlan for ShuffleReaderExec { "ShuffleReaderExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { self.schema.clone() } @@ -426,7 +421,7 @@ impl ExecutionPlan for ShuffleReaderExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result { + fn partition_statistics(&self, partition: Option) -> Result> { if self.broadcast { if let Some(idx) = partition && idx != 0 @@ -445,7 +440,7 @@ impl ExecutionPlan for ShuffleReaderExec { "broadcast shuffle reader at stage {} returned aggregated statistics: {:?}", self.stage_id, stats ); - return Ok(stats); + return Ok(Arc::new(stats)); } if let Some(idx) = partition { let partition_count = self.properties().partitioning.partition_count(); @@ -474,7 +469,7 @@ impl ExecutionPlan for ShuffleReaderExec { "shuffle reader at stage: {} and partition {} returned statistics: {:?}", self.stage_id, idx, stat_for_partition ); - stat_for_partition + stat_for_partition.map(Arc::new) } else { let stats_for_partitions = stats_for_partitions( self.schema.fields().len(), @@ -487,7 +482,7 @@ impl ExecutionPlan for ShuffleReaderExec { "shuffle reader at stage: {} returned statistics for all partitions: {:?}", self.stage_id, stats_for_partitions ); - Ok(stats_for_partitions) + Ok(Arc::new(stats_for_partitions)) } } } diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index a75925d6df..3398be2275 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -24,7 +24,6 @@ use datafusion::arrow::ipc::CompressionType; use datafusion::arrow::ipc::writer::IpcWriteOptions; use datafusion::arrow::ipc::writer::StreamWriter; -use std::any::Any; use std::fmt::Debug; use std::fs; use std::fs::File; @@ -286,7 +285,7 @@ impl ShuffleWriterExec { exprs, num_output_partitions, repart_time, - ); + )?; while let Some(input_batch) = rx.blocking_recv() { partitioner.partition( @@ -438,10 +437,6 @@ impl ExecutionPlan for ShuffleWriterExec { "ShuffleWriterExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { self.plan.schema() } @@ -561,7 +556,7 @@ impl ExecutionPlan for ShuffleWriterExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result { + fn partition_statistics(&self, partition: Option) -> Result> { self.plan.partition_statistics(partition) } } @@ -635,23 +630,22 @@ mod tests { assert_eq!(1, batches.len()); let batch = &batches[0]; assert_eq!(3, batch.num_columns()); - assert_eq!(2, batch.num_rows()); + // One metadata row per non-empty output partition; how many that is + // depends on the hash distribution, so only bound it. + let num_partitions = batch.num_rows(); + assert!((1..=2).contains(&num_partitions)); let path = batch.columns()[1] .as_any() .downcast_ref::() .unwrap(); - - let file0 = path.value(0); - assert!( - file0.ends_with("/jobOne/1/0/data-0.arrow") - || file0.ends_with("\\jobOne\\1\\0\\data-0.arrow") - ); - let file1 = path.value(1); - - assert!( - file1.ends_with("/jobOne/1/1/data-0.arrow") - || file1.ends_with("\\jobOne\\1\\1\\data-0.arrow") - ); + for i in 0..num_partitions { + let f = path.value(i); + assert!( + (0..2).any(|p| f.ends_with(&format!("/jobOne/1/{p}/data-0.arrow")) + || f.ends_with(&format!("\\jobOne\\1\\{p}\\data-0.arrow"))), + "unexpected shuffle file path: {f}" + ); + } let stats = batch.columns()[2] .as_any() @@ -664,8 +658,9 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - assert_eq!(4, num_rows.value(0)); - assert_eq!(4, num_rows.value(1)); + // Total rows are conserved across partitions regardless of the hash split. + let total: u64 = (0..num_partitions).map(|i| num_rows.value(i)).sum(); + assert_eq!(8, total); Ok(()) } @@ -693,7 +688,8 @@ mod tests { assert_eq!(1, batches.len()); let batch = &batches[0]; assert_eq!(3, batch.num_columns()); - assert_eq!(2, batch.num_rows()); + let num_partitions = batch.num_rows(); + assert!((1..=2).contains(&num_partitions)); let stats = batch.columns()[2] .as_any() .downcast_ref::() @@ -704,8 +700,9 @@ mod tests { .as_any() .downcast_ref::() .unwrap(); - assert_eq!(2, num_rows.value(0)); - assert_eq!(2, num_rows.value(1)); + // Total rows are conserved across partitions regardless of the hash split. + let total: u64 = (0..num_partitions).map(|i| num_rows.value(i)).sum(); + assert_eq!(4, total); Ok(()) } diff --git a/ballista/core/src/execution_plans/sort_shuffle/writer.rs b/ballista/core/src/execution_plans/sort_shuffle/writer.rs index 81b201680f..632b5aca07 100644 --- a/ballista/core/src/execution_plans/sort_shuffle/writer.rs +++ b/ballista/core/src/execution_plans/sort_shuffle/writer.rs @@ -21,7 +21,6 @@ //! per input partition, along with an index file mapping partition IDs to //! byte offsets. -use std::any::Any; use std::fs::File; use std::future::Future; use std::io::{BufWriter, Seek, Write}; @@ -539,10 +538,6 @@ impl ExecutionPlan for SortShuffleWriterExec { "SortShuffleWriterExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { self.plan.schema() } @@ -663,7 +658,7 @@ impl ExecutionPlan for SortShuffleWriterExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result { + fn partition_statistics(&self, partition: Option) -> Result> { self.plan.partition_statistics(partition) } } @@ -1164,7 +1159,8 @@ mod tests { // Reference: DataFusion's BatchPartitioner::new_hash_partitioner let mut ref_partitioner = - BatchPartitioner::new_hash_partitioner(exprs.clone(), 4, Time::default()); + BatchPartitioner::new_hash_partitioner(exprs.clone(), 4, Time::default()) + .unwrap(); let mut ref_assignments = [usize::MAX; 10]; ref_partitioner .partition(batch.clone(), |partition, sub_batch| { diff --git a/ballista/core/src/execution_plans/unresolved_shuffle.rs b/ballista/core/src/execution_plans/unresolved_shuffle.rs index 3f3567b6a9..64bfd8fec9 100644 --- a/ballista/core/src/execution_plans/unresolved_shuffle.rs +++ b/ballista/core/src/execution_plans/unresolved_shuffle.rs @@ -15,7 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; use std::sync::Arc; use datafusion::arrow::datatypes::SchemaRef; @@ -192,10 +191,6 @@ impl ExecutionPlan for UnresolvedShuffleExec { "UnresolvedShuffleExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { self.schema.clone() } diff --git a/ballista/core/src/planner.rs b/ballista/core/src/planner.rs index e1ed75998c..ecf41733b8 100644 --- a/ballista/core/src/planner.rs +++ b/ballista/core/src/planner.rs @@ -307,12 +307,10 @@ mod test { assert!(matches!(analyze_df.logical_plan(), LogicalPlan::Analyze(_))); let explain = plan - .as_any() .downcast_ref::>() .unwrap(); assert!( explain.children()[0] - .as_any() .downcast_ref::>() .is_some() ); diff --git a/ballista/core/src/registry.rs b/ballista/core/src/registry.rs index 1d45e7b58a..fe8093cae2 100644 --- a/ballista/core/src/registry.rs +++ b/ballista/core/src/registry.rs @@ -20,6 +20,7 @@ use datafusion::execution::{FunctionRegistry, SessionState}; use datafusion::functions::all_default_functions; use datafusion::functions_aggregate::all_default_aggregate_functions; use datafusion::functions_window::all_default_window_functions; +use datafusion::logical_expr::HigherOrderUDF; use datafusion::logical_expr::planner::ExprPlanner; use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use std::collections::{HashMap, HashSet}; @@ -106,6 +107,19 @@ impl FunctionRegistry for BallistaFunctionRegistry { self.window_functions.keys().cloned().collect() } + fn higher_order_function_names(&self) -> HashSet { + HashSet::new() + } + + fn higher_order_function( + &self, + name: &str, + ) -> datafusion::common::Result> { + Err(DataFusionError::Internal(format!( + "There is no higher-order function named \"{name}\" in the TaskContext" + ))) + } + fn udf(&self, name: &str) -> datafusion::common::Result> { let result = self.scalar_functions.get(name); diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 46e3430c38..0aa1d4feb6 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -36,6 +36,7 @@ use datafusion_proto::physical_plan::from_proto::parse_protobuf_partitioning; use datafusion_proto::physical_plan::to_proto::serialize_partitioning; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, + PhysicalPlanDecodeContext, }; use datafusion_proto::protobuf::proto_error; use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode}; @@ -381,15 +382,15 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { ) })?; let converter = DefaultPhysicalProtoConverter {}; + let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self.default_codec.as_ref()); match ballista_plan { PhysicalPlanType::ShuffleWriter(shuffle_writer) => { let input = inputs[0].clone(); let shuffle_output_partitioning = parse_protobuf_hash_partitioning( shuffle_writer.output_partitioning.as_ref(), - ctx, + &decode_ctx, input.schema().as_ref(), - self.default_codec.as_ref(), &converter, )?; @@ -406,9 +407,8 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { let shuffle_output_partitioning = parse_protobuf_hash_partitioning( sort_shuffle_writer.output_partitioning.as_ref(), - ctx, + &decode_ctx, input.schema().as_ref(), - self.default_codec.as_ref(), &converter, )?; @@ -460,9 +460,8 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { .collect::, DataFusionError>>()?; let partitioning = parse_protobuf_partitioning( shuffle_reader.partitioning.as_ref(), - ctx, + &decode_ctx, schema.as_ref(), - self.default_codec.as_ref(), &converter, )?; let partitioning = partitioning @@ -503,9 +502,8 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { Arc::new(convert_required!(unresolved_shuffle.schema)?); let partitioning = parse_protobuf_partitioning( unresolved_shuffle.partitioning.as_ref(), - ctx, + &decode_ctx, schema.as_ref(), - self.default_codec.as_ref(), &converter, )?; let partitioning = partitioning @@ -557,7 +555,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { node: Arc, buf: &mut Vec, ) -> Result<(), DataFusionError> { - if let Some(exec) = node.as_any().downcast_ref::() { + if let Some(exec) = node.downcast_ref::() { // note that we use shuffle_output_partitioning() rather than output_partitioning() // to get the true output partitioning let output_partitioning = match exec.shuffle_output_partitioning() { @@ -596,7 +594,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { })?; Ok(()) - } else if let Some(exec) = node.as_any().downcast_ref::() { + } else if let Some(exec) = node.downcast_ref::() { let output_partitioning = match exec.shuffle_output_partitioning() { Partitioning::Hash(exprs, partition_count) => { Some(datafusion_proto::protobuf::PhysicalHashRepartition { @@ -639,7 +637,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { })?; Ok(()) - } else if let Some(exec) = node.as_any().downcast_ref::() { + } else if let Some(exec) = node.downcast_ref::() { let stage_id = exec.stage_id as u32; let mut partition = vec![]; for location in &exec.partition { @@ -682,7 +680,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { })?; Ok(()) - } else if let Some(exec) = node.as_any().downcast_ref::() { + } else if let Some(exec) = node.downcast_ref::() { let converter = DefaultPhysicalProtoConverter {}; let partitioning = serialize_partitioning( &exec.properties().partitioning, @@ -708,7 +706,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { })?; Ok(()) - } else if let Some(exec) = node.as_any().downcast_ref::() { + } else if let Some(exec) = node.downcast_ref::() { let proto = protobuf::BallistaPhysicalPlanNode { physical_plan_type: Some(PhysicalPlanType::ChaosExec( protobuf::ChaosExecNode { @@ -834,7 +832,6 @@ mod test { let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap(); let decoded_exec = decoded_plan - .as_any() .downcast_ref::() .expect("Expected UnresolvedShuffleExec"); @@ -871,7 +868,6 @@ mod test { let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap(); let decoded_exec = decoded_plan - .as_any() .downcast_ref::() .expect("Expected ShuffleReaderExec"); @@ -914,7 +910,6 @@ mod test { let ctx = SessionContext::new().task_ctx(); let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap(); let decoded_exec = decoded_plan - .as_any() .downcast_ref::() .expect("Expected ShuffleReaderExec"); @@ -968,7 +963,6 @@ mod test { let ctx = SessionContext::new().task_ctx(); let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap(); let decoded_exec = decoded_plan - .as_any() .downcast_ref::() .expect("Expected ShuffleReaderExec"); @@ -1014,7 +1008,6 @@ mod test { let ctx = SessionContext::new().task_ctx(); let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap(); let decoded_exec = decoded_plan - .as_any() .downcast_ref::() .expect("Expected UnresolvedShuffleExec"); @@ -1065,7 +1058,6 @@ mod test { let ctx = SessionContext::new().task_ctx(); let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap(); let decoded_exec = decoded_plan - .as_any() .downcast_ref::() .expect("Expected ShuffleReaderExec"); @@ -1163,7 +1155,6 @@ mod test { let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap(); let decoded_exec = decoded_plan - .as_any() .downcast_ref::() .expect("Expected UnresolvedShuffleExec"); @@ -1190,7 +1181,6 @@ mod test { let decoded_plan = codec.try_decode(&buf, &[], &ctx).unwrap(); let decoded_exec = decoded_plan - .as_any() .downcast_ref::() .expect("Expected ShuffleReaderExec"); diff --git a/ballista/core/src/serde/scheduler/from_proto.rs b/ballista/core/src/serde/scheduler/from_proto.rs index ba8b9b33a0..34b357ad5e 100644 --- a/ballista/core/src/serde/scheduler/from_proto.rs +++ b/ballista/core/src/serde/scheduler/from_proto.rs @@ -361,6 +361,8 @@ pub fn get_task_definition &dyn Any { - self - } - fn schema(&self) -> SchemaRef { self.plan.schema() } @@ -125,7 +121,7 @@ impl ExecutionPlan for CollectExec { })) } - fn partition_statistics(&self, partition: Option) -> Result { + fn partition_statistics(&self, partition: Option) -> Result> { self.plan.partition_statistics(partition) } } diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index 6ce96fc8be..ea686d43ec 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -113,7 +113,7 @@ impl ExecutionEngine for DefaultExecutionEngine { ) -> Result> { let plan = plan .transform(|p| { - if let Some(reader) = p.as_any().downcast_ref::() { + if let Some(reader) = p.downcast_ref::() { match &self.client_pool { Some(client_pool) => Ok(Transformed::yes(Arc::new( reader @@ -132,7 +132,7 @@ impl ExecutionEngine for DefaultExecutionEngine { // the query plan created by the scheduler always starts with a shuffle writer // (either ShuffleWriterExec or SortShuffleWriterExec) - if let Some(shuffle_writer) = plan.as_any().downcast_ref::() { + if let Some(shuffle_writer) = plan.downcast_ref::() { // recreate the shuffle writer with the correct working directory let exec = ShuffleWriterExec::try_new( job_id, @@ -145,7 +145,7 @@ impl ExecutionEngine for DefaultExecutionEngine { ShuffleWriterVariant::Hash(exec), ))) } else if let Some(sort_shuffle_writer) = - plan.as_any().downcast_ref::() + plan.downcast_ref::() { // recreate the sort shuffle writer with the correct working directory let exec = SortShuffleWriterExec::try_new( diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index f64fae3014..d647ec22a6 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -265,6 +265,8 @@ async fn run_received_task &dyn Any { - self - } - fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } diff --git a/ballista/executor/src/executor_server.rs b/ballista/executor/src/executor_server.rs index 6ef9fa5e4e..1d5d6333b4 100644 --- a/ballista/executor/src/executor_server.rs +++ b/ballista/executor/src/executor_server.rs @@ -398,6 +398,8 @@ impl ExecutorServer Result>> { - let transformed = - if let Some(hash_join) = plan.as_any().downcast_ref::() { - match hash_join.partition_mode() { - PartitionMode::Auto => try_collect_left( - hash_join, - false, - collect_threshold_byte_size, - collect_threshold_num_rows, - )? + let transformed = if let Some(hash_join) = plan.downcast_ref::() { + match hash_join.partition_mode() { + PartitionMode::Auto => try_collect_left( + hash_join, + false, + collect_threshold_byte_size, + collect_threshold_num_rows, + )? + .map_or_else( + || partitioned_hash_join(hash_join).map(Some), + |v| Ok(Some(v)), + )?, + PartitionMode::CollectLeft => try_collect_left(hash_join, true, 0, 0)? .map_or_else( || partitioned_hash_join(hash_join).map(Some), |v| Ok(Some(v)), )?, - PartitionMode::CollectLeft => try_collect_left(hash_join, true, 0, 0)? - .map_or_else( - || partitioned_hash_join(hash_join).map(Some), - |v| Ok(Some(v)), - )?, - PartitionMode::Partitioned => { - let left = hash_join.left(); - let right = hash_join.right(); - if hash_join.join_type().supports_swap() - && should_swap_join_order(&**left, &**right)? - { - hash_join - .swap_inputs(PartitionMode::Partitioned) - .map(Some)? - } else { - None - } + PartitionMode::Partitioned => { + let left = hash_join.left(); + let right = hash_join.right(); + if hash_join.join_type().supports_swap() + && should_swap_join_order(&**left, &**right)? + { + hash_join + .swap_inputs(PartitionMode::Partitioned) + .map(Some)? + } else { + None } } - } else if let Some(cross_join) = plan.as_any().downcast_ref::() { - let left = cross_join.left(); - let right = cross_join.right(); - if right.properties().output_partitioning().partition_count() > 1 { - None - } else if should_swap_join_order(&**left, &**right)? { - cross_join.swap_inputs().map(Some)? - } else { - None - } - } else if let Some(nl_join) = plan.as_any().downcast_ref::() { - let left = nl_join.left(); - let right = nl_join.right(); - // next few lines are different from original datafusion rule - // partition count of right side has to be equal one to be - // able to swap inputs - if right.properties().output_partitioning().partition_count() > 1 { - None - } else if nl_join.join_type().supports_swap() - && should_swap_join_order(&**left, &**right)? - { - nl_join.swap_inputs().map(Some)? - } else { - None - } + } + } else if let Some(cross_join) = plan.downcast_ref::() { + let left = cross_join.left(); + let right = cross_join.right(); + if right.properties().output_partitioning().partition_count() > 1 { + None + } else if should_swap_join_order(&**left, &**right)? { + cross_join.swap_inputs().map(Some)? } else { None - }; + } + } else if let Some(nl_join) = plan.downcast_ref::() { + let left = nl_join.left(); + let right = nl_join.right(); + // next few lines are different from original datafusion rule + // partition count of right side has to be equal one to be + // able to swap inputs + if right.properties().output_partitioning().partition_count() > 1 { + None + } else if nl_join.join_type().supports_swap() + && should_swap_join_order(&**left, &**right)? + { + nl_join.swap_inputs().map(Some)? + } else { + None + } + } else { + None + }; Ok(if let Some(transformed) = transformed { Transformed::yes(transformed) @@ -416,7 +415,7 @@ fn hash_join_convert_symmetric_subrule( config_options: &ConfigOptions, ) -> Result> { // Check if the current plan node is a HashJoinExec. - if let Some(hash_join) = input.as_any().downcast_ref::() { + if let Some(hash_join) = input.downcast_ref::() { let left_unbounded = hash_join.left.boundedness().is_unbounded(); let left_incremental = matches!( hash_join.left.pipeline_behavior(), @@ -556,7 +555,7 @@ pub fn hash_join_swap_subrule( mut input: Arc, _config_options: &ConfigOptions, ) -> Result> { - if let Some(hash_join) = input.as_any().downcast_ref::() + if let Some(hash_join) = input.downcast_ref::() && hash_join.left.boundedness().is_unbounded() && !hash_join.right.boundedness().is_unbounded() && matches!( @@ -830,7 +829,7 @@ mod test { // `swap_inputs` for Inner wraps the join in a ProjectionExec to // restore the output column order. Walk the tree to find the join. fn find_hash_join(plan: &Arc) -> Option<&HashJoinExec> { - if let Some(hj) = plan.as_any().downcast_ref::() { + if let Some(hj) = plan.downcast_ref::() { return Some(hj); } for child in plan.children() { diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index ed396719f4..ce42b26581 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -146,15 +146,13 @@ impl DefaultDistributedPlanner { // Broadcast-join lowering: HashJoinExec(CollectLeft) gets its own // controlled recursion so the build side is written as a broadcast stage. - if let Some(hash_join) = execution_plan.as_any().downcast_ref::() + if let Some(hash_join) = execution_plan.downcast_ref::() && *hash_join.partition_mode() == PartitionMode::CollectLeft { // Build subtree: peel CoalescePartitionsExec if present, then // recurse to lower its internal stages. let mut build = hash_join.left().clone(); - if let Some(coalesce) = - build.as_any().downcast_ref::() - { + if let Some(coalesce) = build.downcast_ref::() { build = coalesce.children()[0].clone(); } let (build, mut stages) = @@ -198,10 +196,7 @@ impl DefaultDistributedPlanner { stages.append(&mut child_stages); } - if let Some(_coalesce) = execution_plan - .as_any() - .downcast_ref::() - { + if let Some(_coalesce) = execution_plan.downcast_ref::() { let input = children[0].clone(); let input = self.optimizer_enforce_sorting.optimize(input, config)?; let shuffle_writer = create_shuffle_writer_with_config( @@ -218,10 +213,9 @@ impl DefaultDistributedPlanner { with_new_children_if_necessary(execution_plan, vec![unresolved_shuffle])?, stages, )) - } else if let Some(_sort_preserving_merge) = execution_plan - .as_any() - .downcast_ref::( - ) { + } else if let Some(_sort_preserving_merge) = + execution_plan.downcast_ref::() + { let shuffle_writer = create_shuffle_writer_with_config( job_id, self.next_stage_id(), @@ -235,9 +229,7 @@ impl DefaultDistributedPlanner { with_new_children_if_necessary(execution_plan, vec![unresolved_shuffle])?, stages, )) - } else if let Some(repart) = - execution_plan.as_any().downcast_ref::() - { + } else if let Some(repart) = execution_plan.downcast_ref::() { match repart.properties().output_partitioning() { Partitioning::Hash(_, _) => { let input = children[0].clone(); @@ -295,7 +287,7 @@ impl DefaultDistributedPlanner { debug!("broadcast check: threshold is 0, broadcast disabled"); return Ok(plan); } - let Some(hash_join) = plan.as_any().downcast_ref::() else { + let Some(hash_join) = plan.downcast_ref::() else { return Ok(plan); }; debug!( @@ -398,7 +390,6 @@ impl DefaultDistributedPlanner { }; let promoted_join = promoted - .as_any() .downcast_ref::() .expect("promoted plan must still be a HashJoinExec"); let new_left: Arc = if promoted_join @@ -429,7 +420,7 @@ fn make_filter_projection_serde_safe( ) -> datafusion::error::Result> { Ok(plan .transform_up(|node| { - if let Some(filter) = node.as_any().downcast_ref::() { + if let Some(filter) = node.downcast_ref::() { let empty_projection = filter .projection() .as_ref() @@ -470,9 +461,7 @@ fn create_unresolved_shuffle( pub fn find_unresolved_shuffles( plan: &Arc, ) -> Result> { - if let Some(unresolved_shuffle) = - plan.as_any().downcast_ref::() - { + if let Some(unresolved_shuffle) = plan.downcast_ref::() { Ok(vec![unresolved_shuffle.clone()]) } else { Ok(plan @@ -495,9 +484,7 @@ pub fn remove_unresolved_shuffles( ) -> Result> { let mut new_children: Vec> = vec![]; for child in stage.children() { - if let Some(unresolved_shuffle) = - child.as_any().downcast_ref::() - { + if let Some(unresolved_shuffle) = child.downcast_ref::() { let p = partition_locations .get(&unresolved_shuffle.stage_id) .ok_or_else(|| { @@ -558,7 +545,7 @@ pub fn rollback_resolved_shuffles( ) -> Result> { let mut new_children: Vec> = vec![]; for child in stage.children() { - if let Some(shuffle_reader) = child.as_any().downcast_ref::() { + if let Some(shuffle_reader) = child.downcast_ref::() { let stage_id = shuffle_reader.stage_id; let unresolved = if shuffle_reader.broadcast { Arc::new(UnresolvedShuffleExec::new_broadcast( @@ -660,11 +647,13 @@ mod test { macro_rules! downcast_exec { ($exec: expr, $ty: ty) => { - $exec.as_any().downcast_ref::<$ty>().expect(&format!( - "Downcast to {} failed. Got {:?}", - stringify!($ty), - $exec - )) + ($exec.as_ref() as &dyn ExecutionPlan) + .downcast_ref::<$ty>() + .expect(&format!( + "Downcast to {} failed. Got {:?}", + stringify!($ty), + $exec + )) }; } @@ -704,7 +693,7 @@ mod test { // FilterExec does not retain an empty projection let mut has_empty_filter_projection = false; serde_safe_plan.apply(|node| { - if let Some(f) = node.as_any().downcast_ref::() + if let Some(f) = node.downcast_ref::() && f.projection() .as_ref() .map(|p| p.is_empty()) @@ -1017,11 +1006,10 @@ order by let mut walker: Vec> = vec![stage.clone() as Arc]; while let Some(node) = walker.pop() { - if let Some(hj) = node.as_any().downcast_ref::() { + if let Some(hj) = node.downcast_ref::() { assert_eq!(*hj.partition_mode(), PartitionMode::CollectLeft); let left = hj.children()[0].clone(); let unresolved = left - .as_any() .downcast_ref::() .expect("left input should be UnresolvedShuffleExec"); assert!(unresolved.broadcast, "left input should be broadcast"); @@ -1062,15 +1050,13 @@ order by let mut walker: Vec> = vec![stage.clone() as Arc]; while let Some(node) = walker.pop() { - if let Some(unresolved) = - node.as_any().downcast_ref::() - { + if let Some(unresolved) = node.downcast_ref::() { assert!( !unresolved.broadcast, "no broadcast reader expected with threshold=0" ); } - if let Some(hj) = node.as_any().downcast_ref::() { + if let Some(hj) = node.downcast_ref::() { assert_ne!( *hj.partition_mode(), PartitionMode::CollectLeft, @@ -1113,8 +1099,7 @@ order by let mut walker: Vec> = vec![stage.clone() as Arc]; while let Some(node) = walker.pop() { - if let Some(unresolved) = - node.as_any().downcast_ref::() + if let Some(unresolved) = node.downcast_ref::() && unresolved.broadcast { max_upstream = max_upstream.max(unresolved.upstream_partition_count); @@ -1184,7 +1169,6 @@ order by let resolved_child = resolved.children()[0].clone(); let reader = resolved_child - .as_any() .downcast_ref::() .expect("expected resolved ShuffleReaderExec"); assert!(reader.broadcast); @@ -1214,7 +1198,6 @@ order by let rolled_back = crate::planner::rollback_resolved_shuffles(parent)?; let child = rolled_back.children()[0].clone(); let unresolved = child - .as_any() .downcast_ref::() .expect("expected rolled-back UnresolvedShuffleExec"); assert!(unresolved.broadcast); @@ -1290,7 +1273,7 @@ order by assert_eq!(2, partitioning.partition_count()); let partition_col = match partitioning { Partitioning::Hash(exprs, 2) => match exprs.as_slice() { - [col] => col.as_any().downcast_ref::(), + [col] => col.downcast_ref::(), _ => None, }, _ => None, @@ -1304,7 +1287,7 @@ order by let window = downcast_exec!(filter.children()[0], BoundedWindowAggExec); let partition_by = window.partition_keys(); let partition_by = match partition_by[..] { - [ref col] => col.as_any().downcast_ref::(), + [ref col] => col.downcast_ref::(), _ => None, }; assert_eq!(Some(&Column::new("l_shipmode", 1)), partition_by); @@ -1321,7 +1304,7 @@ order by ); assert_eq!( Some(&Column::new("l_shipmode", 1)), - expr1.expr.as_any().downcast_ref() + expr1.expr.downcast_ref() ); assert_eq!( SortOptions { @@ -1332,7 +1315,7 @@ order by ); assert_eq!( Some(&Column::new("l_shipdate", 0)), - expr2.expr.as_any().downcast_ref() + expr2.expr.downcast_ref() ); } _ => panic!("invalid sort {sort:?}"), diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index f751d81b97..2040982341 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -45,7 +45,7 @@ impl BallistaAdapter { &mut self, plan: Arc, ) -> datafusion::error::Result>> { - if let Some(exchange) = plan.as_any().downcast_ref::() { + if let Some(exchange) = plan.downcast_ref::() { let schema = exchange.schema().clone(); let partitions = exchange.shuffle_partitions().ok_or_else(|| { DataFusionError::Execution( @@ -119,7 +119,7 @@ impl BallistaAdapter { job_id: &JobId, config: &ConfigOptions, ) -> datafusion::error::Result { - if let Some(root) = plan.as_any().downcast_ref::() { + if let Some(root) = plan.downcast_ref::() { let mut adapter = BallistaAdapter::default(); let plan = root .input() @@ -146,8 +146,7 @@ impl BallistaAdapter { plan: writer, inputs: adapter.inputs, }) - } else if let Some(root) = plan.as_any().downcast_ref::() - { + } else if let Some(root) = plan.downcast_ref::() { let mut adapter = BallistaAdapter::default(); let plan = root .input() diff --git a/ballista/scheduler/src/state/aqe/execution_plan/adaptive.rs b/ballista/scheduler/src/state/aqe/execution_plan/adaptive.rs index 1a1c3962c4..9bc7451ae7 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/adaptive.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/adaptive.rs @@ -22,7 +22,6 @@ use datafusion::{ physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}, }; use parking_lot::Mutex; -use std::any::Any; use std::fmt::Formatter; use std::sync::atomic::AtomicBool; use std::sync::{Arc, atomic::AtomicI64}; @@ -126,10 +125,6 @@ impl ExecutionPlan for AdaptiveDatafusionExec { "AdaptiveDatafusionExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn properties(&self) -> &Arc { self.input.properties() } diff --git a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs index 3c7853c26c..2bac00644b 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -65,10 +65,6 @@ impl ExecutionPlan for DynamicJoinSelectionExec { "DynamicJoinSelectionExec" } - fn as_any(&self) -> &dyn std::any::Any { - self - } - fn properties(&self) -> &Arc { &self.properties } @@ -383,16 +379,14 @@ impl DynamicJoinSelectionExec { /// its children. pub(crate) fn upstream_resolved(&self) -> bool { fn has_join_or_unresolved_exchange(plan: &Arc) -> bool { - if let Some(exchange) = plan.as_any().downcast_ref::() { + if let Some(exchange) = plan.downcast_ref::() { // this should be fine, once we find first resolved // exchange it should not have any unresolved shuffles later // nor dynamic joins return !exchange.shuffle_created(); }; - plan.as_any() - .downcast_ref::() - .is_some() + plan.downcast_ref::().is_some() || plan .children() .iter() diff --git a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs index 28f1011e45..63cc181502 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs @@ -29,7 +29,6 @@ use datafusion::{ }; use log::trace; use parking_lot::Mutex; -use std::any::Any; use std::ops::Deref; use std::sync::{Arc, atomic::AtomicI64}; @@ -326,10 +325,6 @@ impl ExecutionPlan for ExchangeExec { "ExchangeExec" } - fn as_any(&self) -> &dyn Any { - self - } - fn properties(&self) -> &Arc { &self.properties } @@ -381,7 +376,7 @@ impl ExecutionPlan for ExchangeExec { )) } - fn partition_statistics(&self, partition: Option) -> Result { + fn partition_statistics(&self, partition: Option) -> Result> { let schema = self.input.schema(); match self.shuffle_partitions.lock().deref() { // @@ -406,7 +401,7 @@ impl ExecutionPlan for ExchangeExec { "shuffle reader at stage: {:?} and partition {} returned statistics: {:?}", self.stage_id, idx, stat_for_partition ); - stat_for_partition + stat_for_partition.map(Arc::new) } else { let stats_for_partitions = stats_for_partitions( schema.fields().len(), @@ -419,10 +414,10 @@ impl ExecutionPlan for ExchangeExec { "shuffle reader at stage: {:?} returned statistics for all partitions: {:?}", self.stage_id, stats_for_partitions ); - Ok(stats_for_partitions) + Ok(Arc::new(stats_for_partitions)) } } - None => Ok(Statistics::new_unknown(&schema)), + None => Ok(Arc::new(Statistics::new_unknown(&schema))), } } } diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs index a2b8f7b2e1..bef85f3fc4 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs @@ -177,7 +177,7 @@ mod tests { let mut count = 0usize; result .apply(|node| { - if node.as_any().downcast_ref::().is_some() { + if node.downcast_ref::().is_some() { count += 1; } Ok(TreeNodeRecursion::Continue) @@ -193,7 +193,6 @@ mod tests { let plan = leaf_exec(); let result = rule.optimize(plan, &chaos_config(0.42, 42)).unwrap(); let chaos = result - .as_any() .downcast_ref::() .expect("single-node plan always wraps root"); assert!( diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs index 4686fde0fd..5947be4c3d 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -155,7 +155,7 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { ); // Get the subtree below the root. Two root kinds, same outcome. - let input = if let Some(ex) = plan.as_any().downcast_ref::() { + let input = if let Some(ex) = plan.downcast_ref::() { debug!( "[coalesce-rule] root=ExchangeExec plan_id={} stage_id={:?} stage_resolved={}", ex.plan_id, @@ -163,7 +163,7 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { ex.shuffle_partitions().is_some(), ); ex.input().clone() - } else if let Some(adp) = plan.as_any().downcast_ref::() { + } else if let Some(adp) = plan.downcast_ref::() { debug!( "[coalesce-rule] root=AdaptiveDatafusionExec stage_id={:?}", adp.stage_id(), @@ -182,7 +182,7 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { // *this* stage's group, they belong to whatever stage wrote them. let mut leaves: Vec> = Vec::new(); input.apply(|node| { - if node.as_any().is::() { + if node.is::() { leaves.push(node.clone()); Ok(TreeNodeRecursion::Jump) } else { @@ -192,8 +192,7 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { // Helper: downcast each Arc back to &ExchangeExec. fn as_exchange(arc: &Arc) -> &ExchangeExec { - arc.as_any() - .downcast_ref::() + arc.downcast_ref::() .expect("filtered to ExchangeExec above") } diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs index 7b60e1e43d..a07312d14e 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs @@ -45,12 +45,9 @@ impl DistributedExchangeRule { &self, execution_plan: Arc, ) -> datafusion::error::Result>> { - if let Some(coalesce) = execution_plan - .as_any() - .downcast_ref::() - { + if let Some(coalesce) = execution_plan.downcast_ref::() { let input = coalesce.input(); - if input.as_any().downcast_ref::().is_none() + if input.downcast_ref::().is_none() && !matches!(nearest_exchange_status(input), ExchangeStatus::Unresolved) { let exchange_exec = ExchangeExec::new( @@ -63,12 +60,11 @@ impl DistributedExchangeRule { execution_plan.with_new_children(vec![Arc::new(exchange_exec)])?, )); } - } else if let Some(sort_preserving_merge) = execution_plan - .as_any() - .downcast_ref::( - ) { + } else if let Some(sort_preserving_merge) = + execution_plan.downcast_ref::() + { let input = sort_preserving_merge.input(); - if input.as_any().downcast_ref::().is_none() + if input.downcast_ref::().is_none() && !matches!(nearest_exchange_status(input), ExchangeStatus::Unresolved) { let exchange_exec = ExchangeExec::new( @@ -81,8 +77,7 @@ impl DistributedExchangeRule { execution_plan.with_new_children(vec![Arc::new(exchange_exec)])?, )); } - } else if let Some(repartition) = - execution_plan.as_any().downcast_ref::() + } else if let Some(repartition) = execution_plan.downcast_ref::() && let execution_plan::Partitioning::Hash(_, _) = repartition.partitioning() { let input = repartition.input(); @@ -112,7 +107,6 @@ impl PhysicalOptimizerRule for DistributedExchangeRule { if result .data - .as_any() .downcast_ref::() .is_some() { @@ -143,7 +137,7 @@ impl PhysicalOptimizerRule for DistributedExchangeRule { /// (short-circuits), `Resolved` if every branch that has an exchange has a resolved /// one, and `None` if no exchange is found anywhere. fn nearest_exchange_status(plan: &Arc) -> ExchangeStatus { - if let Some(exchange) = plan.as_any().downcast_ref::() { + if let Some(exchange) = plan.downcast_ref::() { if exchange.shuffle_created() && !exchange.inactive_stage { ExchangeStatus::Resolved } else { @@ -241,25 +235,18 @@ mod tests { let result = rule.optimize(coalesce, &config()).unwrap(); - let adaptive = result - .as_any() - .downcast_ref::() - .unwrap(); + let adaptive = result.downcast_ref::().unwrap(); let coalesce_out = adaptive .input() - .as_any() .downcast_ref::() .unwrap(); let child = coalesce_out.children()[0]; assert!( - child.as_any().downcast_ref::().is_some(), + child.downcast_ref::().is_some(), "direct child should remain ExchangeExec" ); assert!( - child.children()[0] - .as_any() - .downcast_ref::() - .is_none(), + child.children()[0].downcast_ref::().is_none(), "ExchangeExec should not wrap another ExchangeExec" ); } @@ -277,18 +264,13 @@ mod tests { let result = rule.optimize(outer, &config()).unwrap(); - let adaptive = result - .as_any() - .downcast_ref::() - .unwrap(); + let adaptive = result.downcast_ref::().unwrap(); let outer_coalesce = adaptive .input() - .as_any() .downcast_ref::() .unwrap(); assert!( outer_coalesce.children()[0] - .as_any() .downcast_ref::() .is_none(), "should not inject ExchangeExec when unresolved exchange is in subtree" @@ -307,18 +289,13 @@ mod tests { let result = rule.optimize(outer, &config()).unwrap(); - let adaptive = result - .as_any() - .downcast_ref::() - .unwrap(); + let adaptive = result.downcast_ref::().unwrap(); let outer_coalesce = adaptive .input() - .as_any() .downcast_ref::() .unwrap(); assert!( outer_coalesce.children()[0] - .as_any() .downcast_ref::() .is_some(), "should inject ExchangeExec when subtree only has resolved exchanges" @@ -334,20 +311,13 @@ mod tests { let result = rule.optimize(input, &config()).unwrap(); - let adaptive = result - .as_any() - .downcast_ref::() - .unwrap(); + let adaptive = result.downcast_ref::().unwrap(); let spm = adaptive .input() - .as_any() .downcast_ref::() .expect("child should be SortPreservingMergeExec"); assert!( - spm.children()[0] - .as_any() - .downcast_ref::() - .is_some(), + spm.children()[0].downcast_ref::().is_some(), "SortPreservingMergeExec should have ExchangeExec injected as its child" ); } @@ -359,23 +329,14 @@ mod tests { let result = rule.optimize(input, &config()).unwrap(); - let adaptive = result - .as_any() - .downcast_ref::() - .unwrap(); + let adaptive = result.downcast_ref::().unwrap(); let spm = adaptive .input() - .as_any() .downcast_ref::() .unwrap(); let child = spm.children()[0]; - assert!(child.as_any().downcast_ref::().is_some()); - assert!( - child.children()[0] - .as_any() - .downcast_ref::() - .is_none() - ); + assert!(child.downcast_ref::().is_some()); + assert!(child.children()[0].downcast_ref::().is_none()); } #[test] @@ -388,20 +349,13 @@ mod tests { let result = rule.optimize(input, &config()).unwrap(); - let adaptive = result - .as_any() - .downcast_ref::() - .unwrap(); + let adaptive = result.downcast_ref::().unwrap(); let spm = adaptive .input() - .as_any() .downcast_ref::() .unwrap(); assert!( - spm.children()[0] - .as_any() - .downcast_ref::() - .is_none(), + spm.children()[0].downcast_ref::().is_none(), "should not inject ExchangeExec when unresolved exchange is in subtree" ); } @@ -419,13 +373,9 @@ mod tests { let result = rule.optimize(repartition, &config()).unwrap(); - let adaptive = result - .as_any() - .downcast_ref::() - .unwrap(); + let adaptive = result.downcast_ref::().unwrap(); let exchange = adaptive .input() - .as_any() .downcast_ref::() .expect("Hash RepartitionExec should be replaced with ExchangeExec"); assert!( @@ -444,16 +394,9 @@ mod tests { let result = rule.optimize(repartition, &config()).unwrap(); - let adaptive = result - .as_any() - .downcast_ref::() - .unwrap(); + let adaptive = result.downcast_ref::().unwrap(); assert!( - adaptive - .input() - .as_any() - .downcast_ref::() - .is_some(), + adaptive.input().downcast_ref::().is_some(), "RoundRobin repartition should be kept as-is (not replaced)" ); } @@ -472,16 +415,9 @@ mod tests { let result = rule.optimize(repartition, &config()).unwrap(); - let adaptive = result - .as_any() - .downcast_ref::() - .unwrap(); + let adaptive = result.downcast_ref::().unwrap(); assert!( - adaptive - .input() - .as_any() - .downcast_ref::() - .is_some(), + adaptive.input().downcast_ref::().is_some(), "Hash repartition should be kept when input has an unresolved exchange" ); } @@ -493,10 +429,7 @@ mod tests { let rule = DistributedExchangeRule::default(); let result = rule.optimize(leaf_exec(), &config()).unwrap(); assert!( - result - .as_any() - .downcast_ref::() - .is_some(), + result.downcast_ref::().is_some(), "optimize should always wrap the result in AdaptiveDatafusionExec" ); } @@ -510,13 +443,11 @@ mod tests { let result = rule.optimize(adaptive, &config()).unwrap(); let outer = result - .as_any() .downcast_ref::() .expect("result should be AdaptiveDatafusionExec"); assert!( outer .input() - .as_any() .downcast_ref::() .is_none(), "existing AdaptiveDatafusionExec should not be wrapped in another one" @@ -609,15 +540,12 @@ mod tests { ) .unwrap(); let exchange1 = result1 - .as_any() .downcast_ref::() .unwrap() .input() - .as_any() .downcast_ref::() .unwrap() .children()[0] - .as_any() .downcast_ref::() .unwrap(); assert_eq!( @@ -632,15 +560,12 @@ mod tests { ) .unwrap(); let exchange2 = result2 - .as_any() .downcast_ref::() .unwrap() .input() - .as_any() .downcast_ref::() .unwrap() .children()[0] - .as_any() .downcast_ref::() .unwrap(); assert_eq!( diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs index 848b0ca0ec..53f44dae61 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs @@ -61,13 +61,11 @@ impl PhysicalOptimizerRule for DelayJoinSelectionRule { .unwrap_or_default(); if bc.adaptive_join_enabled() { let result = plan.transform_up(|node| { - if let Some(hj) = node.as_any().downcast_ref::() { + if let Some(hj) = node.downcast_ref::() { let dynamic = DynamicJoinSelectionExec::from_hash_join(hj, self.plan_id())?; Ok(Transformed::yes(dynamic as Arc)) - } else if let Some(smj) = - node.as_any().downcast_ref::() - { + } else if let Some(smj) = node.downcast_ref::() { let dynamic = DynamicJoinSelectionExec::from_sort_join(smj, self.plan_id())?; Ok(Transformed::yes(dynamic as Arc)) @@ -146,7 +144,7 @@ impl PhysicalOptimizerRule for SelectJoinRule { if bc.adaptive_join_enabled() { let result = plan.transform_up(|node| { if let Some(dynamic_join) = - node.as_any().downcast_ref::() + node.downcast_ref::() { if dynamic_join.upstream_resolved() { match dynamic_join.to_actual_join(config)? { @@ -161,7 +159,6 @@ impl PhysicalOptimizerRule for SelectJoinRule { let right = hash_join_exec.right.clone(); let right = right - .as_any() .downcast_ref::() .ok_or(DataFusionError::Execution( "ExchangeExec is expected at this point (right)" @@ -174,19 +171,17 @@ impl PhysicalOptimizerRule for SelectJoinRule { .with_new_children(vec![left, right])?; let join = join - .as_any() - .downcast_ref::() - .ok_or(DataFusionError::Execution( - "HashJoinExec is expected at this point (right)" - .to_owned(), - ))?; + .downcast_ref::() + .ok_or(DataFusionError::Execution( + "HashJoinExec is expected at this point (right)" + .to_owned(), + ))?; let join = join.swap_inputs(PartitionMode::CollectLeft)?; Ok(Transformed::yes(join)) } else { let left = hash_join_exec .left - .as_any() .downcast_ref::() .ok_or(DataFusionError::Execution( "ExchangeExec is expected at this point (left)" @@ -216,7 +211,7 @@ impl PhysicalOptimizerRule for SelectJoinRule { let exec = plan.transform_up(|p| { if let Some(hash_join) = - p.as_any().downcast_ref::() + p.downcast_ref::() && matches!( hash_join.partition_mode(), PartitionMode::CollectLeft @@ -224,10 +219,9 @@ impl PhysicalOptimizerRule for SelectJoinRule { { if let Some(data_source) = hash_join .left - .as_any() - .downcast_ref::() - && let Ok(Some(left)) = - data_source.repartitioned(1, config) + .downcast_ref::( + ) && let Ok(Some(left)) = + data_source.repartitioned(1, config) { let right = hash_join.right.clone(); let p = @@ -242,7 +236,6 @@ impl PhysicalOptimizerRule for SelectJoinRule { { let left = if let Some(exchange) = hash_join .left - .as_any() .downcast_ref::() { Arc::new( diff --git a/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs b/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs index 86002c316f..87d52a5bd8 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs @@ -41,7 +41,7 @@ use std::sync::Arc; macro_rules! is_empty_exec { ($e:expr) => { - $e.as_any().downcast_ref::().is_some() + $e.downcast_ref::().is_some() }; } @@ -75,45 +75,45 @@ impl PropagateEmptyExecRule { fn transform( plan: Arc, ) -> datafusion::error::Result>> { - if let Some(filter) = plan.as_any().downcast_ref::() + if let Some(filter) = plan.downcast_ref::() && is_empty_exec!(filter.input()) { Ok(Transformed::yes(filter.input().clone())) - } else if let Some(coalesce) = plan.as_any().downcast_ref::() + } else if let Some(coalesce) = plan.downcast_ref::() && is_empty_exec!(coalesce.input()) { Ok(Transformed::yes(coalesce.input().clone())) - } else if let Some(exchange) = plan.as_any().downcast_ref::() + } else if let Some(exchange) = plan.downcast_ref::() && is_empty_exec!(exchange.input()) { Ok(Transformed::yes(exchange.input().clone())) - } else if let Some(projection) = plan.as_any().downcast_ref::() + } else if let Some(projection) = plan.downcast_ref::() && is_empty_exec!(projection.input()) { empty_exec!(projection) - } else if let Some(limit) = plan.as_any().downcast_ref::() + } else if let Some(limit) = plan.downcast_ref::() && is_empty_exec!(limit.input()) { Ok(Transformed::yes(limit.input().clone())) - } else if let Some(limit) = plan.as_any().downcast_ref::() + } else if let Some(limit) = plan.downcast_ref::() && is_empty_exec!(limit.input()) { Ok(Transformed::yes(limit.input().clone())) - } else if let Some(aggregation) = plan.as_any().downcast_ref::() + } else if let Some(aggregation) = plan.downcast_ref::() && is_empty_exec!(aggregation.input()) { empty_exec!(aggregation) - } else if let Some(repartition) = plan.as_any().downcast_ref::() + } else if let Some(repartition) = plan.downcast_ref::() && is_empty_exec!(repartition.input()) { empty_exec!(repartition) } else if let Some(coalesce_partition) = - plan.as_any().downcast_ref::() + plan.downcast_ref::() && is_empty_exec!(coalesce_partition.input()) { empty_exec!(coalesce_partition) } else if let Some(sort_preserving_merge) = - plan.as_any().downcast_ref::() + plan.downcast_ref::() && is_empty_exec!(sort_preserving_merge.input()) { empty_exec!(sort_preserving_merge) @@ -180,7 +180,7 @@ impl PropagateEmptyExecRule { } _ => Ok(Transformed::no(plan)), } - } else if let Some(exchange) = plan.as_any().downcast_ref::() { + } else if let Some(exchange) = plan.downcast_ref::() { let stats = exchange.partition_statistics(None)?; match stats.num_rows { Precision::Exact(0) => empty_exec!(plan), @@ -213,7 +213,7 @@ impl PhysicalOptimizerRule for PropagateEmptyExecRule { /// Extracting the physical join information from the incoming [ExecutionPlan] pub fn as_join(plan: &Arc) -> Option> { - let any = plan.as_any(); + let any = plan.as_ref(); if let Some(join) = any.downcast_ref::() { return Some(JoinInfo { @@ -422,7 +422,7 @@ mod tests { /// Assert the result is an EmptyExec. fn assert_empty(result: &Arc) { assert!( - result.as_any().downcast_ref::().is_some(), + result.downcast_ref::().is_some(), "expected EmptyExec, got {:?}", result.name() ); @@ -431,7 +431,7 @@ mod tests { /// Assert the result is a ProjectionExec (null-padded surviving side). fn assert_projection(result: &Arc) { assert!( - result.as_any().downcast_ref::().is_some(), + result.downcast_ref::().is_some(), "expected ProjectionExec, got {:?}", result.name() ); @@ -441,8 +441,8 @@ mod tests { /// i.e. the join was left untouched. fn assert_untouched(result: &Arc) { assert!( - result.as_any().downcast_ref::().is_none() - && result.as_any().downcast_ref::().is_none(), + result.downcast_ref::().is_none() + && result.downcast_ref::().is_none(), "expected join to be untouched, got {:?}", result.name() ); @@ -732,7 +732,7 @@ mod tests { let plan = hash_join(Arc::clone(&left), empty_stats_exec(), JoinType::LeftAnti); let result = transform(plan); // Should be the left child itself, not wrapped in anything - assert!(result.as_any().downcast_ref::().is_some()); + assert!(result.downcast_ref::().is_some()); } #[test] @@ -761,7 +761,7 @@ mod tests { let right = non_empty_stats_exec(); let plan = hash_join(empty_stats_exec(), Arc::clone(&right), JoinType::RightAnti); let result = transform(plan); - assert!(result.as_any().downcast_ref::().is_some()); + assert!(result.downcast_ref::().is_some()); } #[test] @@ -857,7 +857,7 @@ mod tests { JoinType::LeftAnti, ); let result = transform(plan); - assert!(result.as_any().downcast_ref::().is_some()); + assert!(result.downcast_ref::().is_some()); } #[test] @@ -868,7 +868,7 @@ mod tests { JoinType::RightAnti, ); let result = transform(plan); - assert!(result.as_any().downcast_ref::().is_some()); + assert!(result.downcast_ref::().is_some()); } // ── unknown stats — never optimised ───────────────────────────────────── @@ -945,7 +945,7 @@ mod tests { .optimize(plan.clone(), &ConfigOptions::default()) .unwrap(); - let empty_exec = result.as_any().downcast_ref::(); + let empty_exec = result.downcast_ref::(); assert!(empty_exec.is_some(), "expected EmptyExec"); assert_eq!( empty_exec diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 81afd9d5df..844bcc32b1 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -216,8 +216,8 @@ impl AdaptivePlanner { .as_ref() .map(|stage| { ( - stage.as_any().downcast_ref::(), - stage.as_any().downcast_ref::(), + stage.downcast_ref::(), + stage.downcast_ref::(), ) }) { Some((Some(stage), None)) => { @@ -266,7 +266,6 @@ impl AdaptivePlanner { ), )?; let is_broadcast = stage - .as_any() .downcast_ref::() .map(|e| e.broadcast) .unwrap_or(false); @@ -395,7 +394,7 @@ impl AdaptivePlanner { if !runnable_stages.is_empty() { let mut runnable = Vec::new(); for exec in runnable_stages.into_iter() { - match exec.as_any().downcast_ref::() { + match exec.downcast_ref::() { Some(exchange) if exchange.inactive_stage => continue, Some(exchange) if exchange.stage_id().is_none() => { exchange.set_stage_id(self.stage_id_generator); @@ -418,9 +417,7 @@ impl AdaptivePlanner { } Ok(Some(runnable)) - } else if let Some(root) = - self.plan.as_any().downcast_ref::() - { + } else if let Some(root) = self.plan.downcast_ref::() { // shuffle writer has finished // there is no more runnable stages if root.shuffle_created() { @@ -462,8 +459,7 @@ impl AdaptivePlanner { runnable_stages .into_iter() .map(|exec| { - exec.as_any() - .downcast_ref::() + exec.downcast_ref::() .ok_or_else(|| { datafusion::common::DataFusionError::Plan( "ExchangeExec expected".into(), @@ -569,7 +565,7 @@ impl AdaptivePlanner { node: &Arc, runnable_stages: &mut Vec>, ) -> bool { - if let Some(exchange) = node.as_any().downcast_ref::() + if let Some(exchange) = node.downcast_ref::() && exchange.shuffle_created() { // we found exchange which has partitions resolved or this stage is not @@ -577,7 +573,7 @@ impl AdaptivePlanner { // all runnable children has been run false - } else if let Some(exchange) = node.as_any().downcast_ref::() + } else if let Some(exchange) = node.downcast_ref::() && !exchange.shuffle_created() { // we found exchange which has not been resolved (run) diff --git a/ballista/scheduler/src/state/aqe/test/alter_stages.rs b/ballista/scheduler/src/state/aqe/test/alter_stages.rs index fcf48e9cfa..7757cddd79 100644 --- a/ballista/scheduler/src/state/aqe/test/alter_stages.rs +++ b/ballista/scheduler/src/state/aqe/test/alter_stages.rs @@ -37,7 +37,6 @@ use datafusion::physical_plan::test::exec::StatisticsExec; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, PlanProperties, }; -use std::any::Any; use std::collections::HashSet; use std::fmt::Formatter; use std::sync::Arc; @@ -211,45 +210,42 @@ async fn should_support_join_re_ordering() -> datafusion::error::Result<()> { // join ordering changes as build side is bigger than probe side // after exchange statistic updated. - assert_plan!(planner.current_plan(), @ r" + assert_plan!(planner.current_plan(), @ " AdaptiveDatafusionExec: is_final=false, plan_id=2, stage_id=pending, stage_resolved=false - ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col] - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)] - ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=1, stage_id=1, stage_resolved=true - CooperativeExec - MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] - ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=0, stage_id=0, stage_resolved=true - CooperativeExec - MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)], projection=[big_col@1, big_col@0] + ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=1, stage_id=1, stage_resolved=true + CooperativeExec + MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] + ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=0, stage_id=0, stage_resolved=true + CooperativeExec + MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] "); let stages = planner.runnable_stages()?.unwrap(); assert_eq!(1, stages.len()); - assert_plan!(planner.current_plan(), @ r" + assert_plan!(planner.current_plan(), @ " AdaptiveDatafusionExec: is_final=true, plan_id=2, stage_id=2, stage_resolved=false - ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col] - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)] - ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=1, stage_id=1, stage_resolved=true - CooperativeExec - MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] - ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=0, stage_id=0, stage_resolved=true - CooperativeExec - MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)], projection=[big_col@1, big_col@0] + ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=1, stage_id=1, stage_resolved=true + CooperativeExec + MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] + ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=0, stage_id=0, stage_resolved=true + CooperativeExec + MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] "); planner.finalise_stage_internal(2, small_statistics_exchange())?; - assert_plan!(planner.current_plan(), @ r" + assert_plan!(planner.current_plan(), @ " AdaptiveDatafusionExec: is_final=true, plan_id=2, stage_id=2, stage_resolved=true - ProjectionExec: expr=[big_col@1 as big_col, big_col@0 as big_col] - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)] - ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=1, stage_id=1, stage_resolved=true - CooperativeExec - MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] - ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=0, stage_id=0, stage_resolved=true - CooperativeExec - MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(big_col@0, big_col@0)], projection=[big_col@1, big_col@0] + ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=1, stage_id=1, stage_resolved=true + CooperativeExec + MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] + ExchangeExec: partitioning=Hash([big_col@0], 2), plan_id=0, stage_id=0, stage_resolved=true + CooperativeExec + MockPartitionedScan: num_partitions=2, statistics=[Rows=Exact(262144), Bytes=Exact(2097152), [(Col[0]:)]] "); Ok(()) @@ -649,10 +645,6 @@ impl ExecutionPlan for MockPartitionedScan { "MockPartitionedScan" } - fn as_any(&self) -> &dyn Any { - self - } - fn properties(&self) -> &Arc { &self.plan_properties } @@ -679,7 +671,7 @@ impl ExecutionPlan for MockPartitionedScan { fn partition_statistics( &self, _partition: Option, - ) -> datafusion::common::Result { - Ok(self.statistics.clone()) + ) -> datafusion::common::Result> { + Ok(Arc::new(self.statistics.clone())) } } diff --git a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs index 528a912fb8..b115f273c5 100644 --- a/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs +++ b/ballista/scheduler/src/state/aqe/test/plan_to_stages.rs @@ -25,6 +25,7 @@ use crate::state::aqe::test::{ use ballista_core::execution_plans::SortShuffleWriterExec; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::ColumnStatistics; +use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::Statistics; use datafusion::physical_plan::test::exec::StatisticsExec; use std::collections::HashSet; @@ -475,7 +476,7 @@ async fn should_use_sort_shuffle_when_enabled() -> datafusion::error::Result<()> let plan = stages.first().unwrap().plan.as_ref(); assert!( - plan.as_any() + (plan as &dyn ExecutionPlan) .downcast_ref::() .is_some(), "expected SortShuffleWriterExec when sort shuffle is enabled, got plan: {plan:?}" @@ -505,7 +506,7 @@ async fn should_use_sort_shuffle_by_default() -> datafusion::error::Result<()> { let plan = stages.first().unwrap().plan.as_ref(); assert!( - plan.as_any() + (plan as &dyn ExecutionPlan) .downcast_ref::() .is_some(), "expected SortShuffleWriterExec by default, got plan: {plan:?}" diff --git a/ballista/scheduler/src/state/distributed_explain.rs b/ballista/scheduler/src/state/distributed_explain.rs index 4292950329..c4408be468 100644 --- a/ballista/scheduler/src/state/distributed_explain.rs +++ b/ballista/scheduler/src/state/distributed_explain.rs @@ -195,7 +195,7 @@ pub(crate) async fn handle_explain_plan( plan: Arc, ) -> ballista_core::error::Result> { if let LogicalPlan::Explain(explain_plan) = &logical_plan - && let Some(explain) = plan.as_any().downcast_ref::() + && let Some(explain) = plan.downcast_ref::() { let inner_plan = explain_plan.plan.clone(); let plans = explain.stringified_plans(); diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index 72fd05f8f9..51228a40b6 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -1633,14 +1633,12 @@ impl ExecutionPlanVisitor for ExecutionStageBuilder { plan: &dyn ExecutionPlan, ) -> std::result::Result { // Handle both ShuffleWriterExec and SortShuffleWriterExec - if let Some(shuffle_write) = plan.as_any().downcast_ref::() { + if let Some(shuffle_write) = plan.downcast_ref::() { self.current_stage_id = shuffle_write.stage_id(); - } else if let Some(shuffle_write) = - plan.as_any().downcast_ref::() - { + } else if let Some(shuffle_write) = plan.downcast_ref::() { self.current_stage_id = shuffle_write.stage_id(); } else if let Some(unresolved_shuffle) = - plan.as_any().downcast_ref::() + plan.downcast_ref::() { if let Some(output_links) = self.output_links.get_mut(&unresolved_shuffle.stage_id) @@ -1710,18 +1708,14 @@ impl TaskDescription { /// Returns the number of output partitions this task will produce. pub fn get_output_partition_number(&self) -> usize { // Try ShuffleWriterExec first - if let Some(shuffle_writer) = - self.plan.as_any().downcast_ref::() - { + if let Some(shuffle_writer) = self.plan.downcast_ref::() { return shuffle_writer .shuffle_output_partitioning() .map(|partitioning| partitioning.partition_count()) .unwrap_or(1); } // Try SortShuffleWriterExec - if let Some(shuffle_writer) = - self.plan.as_any().downcast_ref::() - { + if let Some(shuffle_writer) = self.plan.downcast_ref::() { return shuffle_writer .shuffle_output_partitioning() .partition_count(); diff --git a/ballista/scheduler/src/state/execution_graph_dot.rs b/ballista/scheduler/src/state/execution_graph_dot.rs index f43772dd0f..eaef4dc5e9 100644 --- a/ballista/scheduler/src/state/execution_graph_dot.rs +++ b/ballista/scheduler/src/state/execution_graph_dot.rs @@ -148,7 +148,7 @@ fn write_plan_recursive( let node_name = format!("{prefix}_{i}"); let display_name = get_operator_name(plan); - if let Some(reader) = plan.as_any().downcast_ref::() { + if let Some(reader) = plan.downcast_ref::() { for part in &reader.partition { for loc in part { state @@ -156,7 +156,7 @@ fn write_plan_recursive( .insert(node_name.clone(), loc.partition_id.stage_id); } } - } else if let Some(reader) = plan.as_any().downcast_ref::() { + } else if let Some(reader) = plan.downcast_ref::() { state.readers.insert(node_name.clone(), reader.stage_id); } @@ -230,9 +230,9 @@ fn sanitize(str: &str, max_len: Option) -> String { } #[allow(deprecated)] fn get_operator_name(plan: &dyn ExecutionPlan) -> String { - if let Some(exec) = plan.as_any().downcast_ref::() { + if let Some(exec) = plan.downcast_ref::() { format!("Filter: {}", exec.predicate()) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { let expr = exec .expr() .iter() @@ -241,7 +241,7 @@ fn get_operator_name(plan: &dyn ExecutionPlan) -> String { .collect::>() .join(", "); format!("Projection: {}", sanitize_dot_label(&expr)) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { let sort_expr = exec .expr() .iter() @@ -257,7 +257,7 @@ fn get_operator_name(plan: &dyn ExecutionPlan) -> String { .collect::>() .join(", "); format!("Sort: {}", sanitize_dot_label(&sort_expr)) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { let group_exprs_with_alias = exec.group_expr().expr(); let group_expr = group_exprs_with_alias .iter() @@ -277,19 +277,19 @@ aggr=[{}]", sanitize_dot_label(&group_expr), sanitize_dot_label(&aggr_expr) ) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { format!("CoalesceBatches [batchSize={}]", exec.target_batch_size()) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { format!( "CoalescePartitions [{}]", format_partitioning(exec.properties().output_partitioning().clone()) ) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { format!( "RepartitionExec [{}]", format_partitioning(exec.properties().output_partitioning().clone()) ) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { let join_expr = exec .on() .iter() @@ -308,49 +308,46 @@ filter_expr={}", sanitize_dot_label(&join_expr), sanitize_dot_label(&filter_expr) ) - } else if plan.as_any().downcast_ref::().is_some() { + } else if plan.downcast_ref::().is_some() { "CrossJoin".to_string() - } else if plan.as_any().downcast_ref::().is_some() { + } else if plan.downcast_ref::().is_some() { "Union".to_string() - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { format!("UnresolvedShuffleExec [stage_id={}]", exec.stage_id) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { format!("ShuffleReader [{} partitions]", exec.partition.len()) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { format!( "ShuffleWriter [{} partitions]", exec.input_partition_count() ) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { format!( "SortShuffleWriter [{} partitions]", exec.input_partition_count() ) - } else if let Some(exec) = plan.as_any().downcast_ref::() { - let config = if let Some(config) = - exec.data_source().as_any().downcast_ref::() - { - get_file_scan(config) - } else if let Some(_config) = exec - .data_source() - .as_any() - .downcast_ref::() - { - "Memory".to_string() - } else { - "Unknown".to_string() - }; + } else if let Some(exec) = plan.downcast_ref::() { + let config = + if let Some(config) = exec.data_source().downcast_ref::() { + get_file_scan(config) + } else if let Some(_config) = + exec.data_source().downcast_ref::() + { + "Memory".to_string() + } else { + "Unknown".to_string() + }; let parts = exec.properties().output_partitioning().partition_count(); format!("DataSourceExec: ({config}) [{parts} partitions]") - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { format!( "GlobalLimit(skip={}, fetch={:?})", exec.skip(), exec.fetch() ) - } else if let Some(exec) = plan.as_any().downcast_ref::() { + } else if let Some(exec) = plan.downcast_ref::() { format!("LocalLimit({})", exec.fetch()) } else { debug!("Unknown physical operator when producing DOT graph: {plan:?}"); @@ -589,6 +586,14 @@ filter_expr="] .options_mut() .optimizer .enable_round_robin_repartition = false; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold = 0; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold_rows = 0; let ctx = SessionContext::new_with_config(config); let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::UInt32, false), @@ -627,6 +632,14 @@ filter_expr="] .options_mut() .optimizer .enable_round_robin_repartition = false; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold = 0; + config + .options_mut() + .optimizer + .hash_join_single_partition_threshold_rows = 0; let ctx = SessionContext::new_with_config(config); let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, false)])); diff --git a/ballista/scheduler/src/state/execution_stage.rs b/ballista/scheduler/src/state/execution_stage.rs index a9695d3b8b..e9e8f1b07f 100644 --- a/ballista/scheduler/src/state/execution_stage.rs +++ b/ballista/scheduler/src/state/execution_stage.rs @@ -1005,11 +1005,11 @@ impl Debug for FailedStage { /// will be different. Here, we should use the input partition count. fn get_stage_partitions(plan: Arc) -> usize { // Try ShuffleWriterExec first - if let Some(shuffle_writer) = plan.as_any().downcast_ref::() { + if let Some(shuffle_writer) = plan.downcast_ref::() { return shuffle_writer.input_partition_count(); } // Try SortShuffleWriterExec - if let Some(shuffle_writer) = plan.as_any().downcast_ref::() { + if let Some(shuffle_writer) = plan.downcast_ref::() { return shuffle_writer.input_partition_count(); } // Fallback to output partitioning diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index 810983ef2b..4f7b29651c 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -19,7 +19,6 @@ use ballista_core::error::{BallistaError, Result}; use ballista_core::extension::SessionConfigExt; use ballista_core::{JobId, JobStatusSubscriber}; use datafusion::catalog::Session; -use std::any::Any; use std::collections::HashMap; use std::future::Future; use std::sync::Arc; @@ -80,10 +79,6 @@ pub struct ExplodingTableProvider; #[async_trait::async_trait] impl TableProvider for ExplodingTableProvider { - fn as_any(&self) -> &dyn Any { - self - } - fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } @@ -1016,7 +1011,7 @@ pub async fn test_coalesce_plan(partition: usize) -> StaticExecutionGraph { /// Creates a test execution graph with a join operation. pub async fn test_join_plan(partition: usize) -> StaticExecutionGraph { - let mut config = SessionConfig::new().with_target_partitions(partition); + let mut config = SessionConfig::new_with_ballista().with_target_partitions(partition); config .options_mut() .optimizer From 74843142957a1360ce95f17cd049d882f3a14ac8 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 27 Jun 2026 14:32:57 -0600 Subject: [PATCH 2/8] fix: restrict each scan task to its own file group on DataFusion 54 DataFusion 54's DataSourceExec hands file groups to partition streams from a shared work-queue that is only divided across partitions when all partitions of one plan instance are polled concurrently. Ballista executes one partition per task on its own decoded plan instance, so a task that polls a single partition in isolation drains the whole queue and scans the entire table -- producing N-times-inflated results and N-times the work for every task dispatched after the first wave. Restrict each task's DataSourceExec to the file group for its partition_id before execution (other group slots emptied, partition count preserved) so the lone execute(partition_id) reads only its own group. Skip when partition_id is outside the source's file groups (an operator between the scan and the stage output changed the partition count). Closes #1907 --- ballista/executor/src/execution_engine.rs | 53 ++++++++++++++++++++++- 1 file changed, 52 insertions(+), 1 deletion(-) diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index ea686d43ec..fa22edd0b9 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -27,11 +27,16 @@ use ballista_core::execution_plans::{ShuffleReaderExec, ShuffleWriterExec}; use ballista_core::serde::protobuf::ShuffleWritePartition; use ballista_core::{JobId, utils}; use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::datasource::physical_plan::{ + FileGroup, FileScanConfig, FileScanConfigBuilder, +}; +use datafusion::datasource::source::DataSourceExec; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::context::TaskContext; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::MetricsSet; use datafusion::prelude::SessionConfig; +use std::any::Any; use std::fmt::{Debug, Display}; use std::sync::Arc; @@ -101,12 +106,54 @@ impl DefaultExecutionEngine { } } +/// Restrict a `DataSourceExec` to the file group for `partition_id`. +/// +/// DataFusion 54's `DataSourceExec` hands file groups to partition streams from +/// a shared work-queue that is only divided across partitions when all +/// partitions of one plan instance are polled concurrently. Ballista runs one +/// partition per task on its own plan instance, so a task that polls a single +/// partition in isolation would otherwise drain the whole queue and scan the +/// entire table. Keeping only this task's file group (other slots emptied, +/// partition count preserved) makes the lone `execute(partition_id)` read just +/// that group. +/// +/// Returns `None` for any node that is not a file-backed `DataSourceExec`, and +/// for a `partition_id` outside the source's file groups (e.g. when an operator +/// between the scan and the stage output changed the partition count). +fn restrict_scan_to_partition( + plan: &Arc, + partition_id: usize, +) -> Option> { + let exec = plan.downcast_ref::()?; + let source: &dyn Any = exec.data_source().as_ref(); + let config = source.downcast_ref::()?; + if partition_id >= config.file_groups.len() { + return None; + } + let file_groups: Vec = config + .file_groups + .iter() + .enumerate() + .map(|(i, group)| { + if i == partition_id { + group.clone() + } else { + FileGroup::new(vec![]) + } + }) + .collect(); + let config = FileScanConfigBuilder::from(config.clone()) + .with_file_groups(file_groups) + .build(); + Some(DataSourceExec::from_data_source(config)) +} + impl ExecutionEngine for DefaultExecutionEngine { fn create_query_stage_exec( &self, job_id: JobId, stage_id: usize, - _partition_id: usize, + partition_id: usize, plan: Arc, work_dir: &str, _config: &SessionConfig, @@ -124,6 +171,10 @@ impl ExecutionEngine for DefaultExecutionEngine { reader.with_work_dir(work_dir.to_string()), ))), } + } else if let Some(rewritten) = + restrict_scan_to_partition(&p, partition_id) + { + Ok(Transformed::yes(rewritten)) } else { Ok(Transformed::no(p)) } From 7959c70fb4f43187097cef7bc9ca315327c06227 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 27 Jun 2026 14:41:09 -0600 Subject: [PATCH 3/8] fix: decorrelate uncorrelated scalar subqueries for distributed execution DataFusion 54 plans uncorrelated scalar subqueries as a physical ScalarSubqueryExec wrapping a ScalarSubqueryExpr that reads an in-process shared results container. That container cannot cross process or stage boundaries, and datafusion-proto can only deserialize the expr inside its surrounding exec, so when Ballista splits a plan into stages at shuffle boundaries the executor receives a bare ScalarSubqueryExpr and fails to decode it -- hanging queries such as TPC-H q11/q15/q22. Disable datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery in the Ballista restricted configuration so the optimizer rewrites uncorrelated scalar subqueries to joins, which Ballista distributes correctly. Verified via datafusion-cli that the rewrite produces results identical to the physical-execution path. Closes #1909 --- ballista/core/src/extension.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index d53307339c..914de09fe0 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -789,6 +789,22 @@ impl SessionConfigHelperExt for SessionConfig { // // See https://github.com/apache/datafusion-ballista/issues/1648 .set_bool("datafusion.optimizer.prefer_hash_join", false) + // + // DataFusion 54 plans uncorrelated scalar subqueries as a physical + // `ScalarSubqueryExec` wrapping a `ScalarSubqueryExpr` that reads an + // in-process shared results container. That container cannot cross + // process or stage boundaries, and `datafusion-proto` can only + // deserialize the expr inside its surrounding exec, so when Ballista + // splits a plan into stages the expr is serialized without its exec + // and the executor fails to decode it. Disabling this option makes + // the optimizer rewrite uncorrelated scalar subqueries to joins, + // which Ballista distributes correctly. + // + // See https://github.com/apache/datafusion-ballista/issues/1909 + .set_bool( + "datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery", + false, + ) } } From 85805db74ffaa7aca25009099ecd06819163dc22 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sat, 27 Jun 2026 16:22:26 -0600 Subject: [PATCH 4/8] test: cover DataFusion 54 scan restriction and scalar-subquery config Add unit tests for the two DataFusion 54 distributed-execution fixes: - restrict_scan_to_partition keeps only the task's own file group, leaves the partition count intact, returns None for an out-of-range partition, and ignores non-file-backed plans. - the Ballista restricted configuration disables datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery so uncorrelated scalar subqueries are rewritten to joins. --- ballista/core/src/extension.rs | 15 ++++++ ballista/executor/src/execution_engine.rs | 57 +++++++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index 914de09fe0..dd5ec26926 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -1088,6 +1088,21 @@ mod test { ); } + // Uncorrelated scalar subqueries must be rewritten to joins rather than + // planned as a physical `ScalarSubqueryExec`, whose `ScalarSubqueryExpr` + // cannot be deserialized once Ballista splits the plan into stages. See + // #1909. + #[test] + fn should_disable_physical_uncorrelated_scalar_subquery() { + let config = SessionConfig::new().upgrade_for_ballista(); + assert!( + !config + .options() + .optimizer + .enable_physical_uncorrelated_scalar_subquery + ); + } + #[test] fn should_convert_to_key_value_pairs() { // key value pairs should contain datafusion and ballista values diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index fa22edd0b9..f559af2936 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -310,3 +310,60 @@ impl QueryStageExecutor for DefaultQueryStageExec { } } } + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::listing::PartitionedFile; + use datafusion::datasource::physical_plan::ParquetSource; + use datafusion::execution::object_store::ObjectStoreUrl; + use datafusion::physical_plan::empty::EmptyExec; + + /// Build a `DataSourceExec` over `n` file groups, one file each. + fn scan_with_file_groups(n: usize) -> Arc { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let source = Arc::new(ParquetSource::new(schema)); + let mut builder = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source); + for i in 0..n { + builder = + builder.with_file_group(FileGroup::new(vec![PartitionedFile::new( + format!("file{i}.parquet"), + 100, + )])); + } + DataSourceExec::from_data_source(builder.build()) + } + + /// Number of files in each file group of a `DataSourceExec`. + fn group_file_counts(plan: &Arc) -> Vec { + let exec = plan.downcast_ref::().unwrap(); + let source: &dyn Any = exec.data_source().as_ref(); + let config = source.downcast_ref::().unwrap(); + config.file_groups.iter().map(|g| g.len()).collect() + } + + #[test] + fn restrict_scan_keeps_only_its_own_group() { + let plan = scan_with_file_groups(4); + // partition 2 keeps only group 2; the partition count is preserved. + let restricted = restrict_scan_to_partition(&plan, 2).expect("scan rewritten"); + assert_eq!(group_file_counts(&restricted), vec![0, 0, 1, 0]); + } + + #[test] + fn restrict_scan_partition_out_of_range_is_left_untouched() { + let plan = scan_with_file_groups(3); + // an operator between the scan and the stage output may change the + // partition count; in that case we must not rewrite the scan. + assert!(restrict_scan_to_partition(&plan, 3).is_none()); + } + + #[test] + fn restrict_scan_ignores_non_file_scans() { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let plan: Arc = Arc::new(EmptyExec::new(schema)); + assert!(restrict_scan_to_partition(&plan, 0).is_none()); + } +} From e15d7f15b78aeb037947744c1c4a8b4742d0be42 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 29 Jun 2026 09:51:03 -0600 Subject: [PATCH 5/8] fix: register higher-order functions in executor task context DataFusion 54 ships built-in higher-order functions (array_any_match, array_filter, array_transform), but BallistaFunctionRegistry returned an empty set and every TaskContext the executor built passed Default::default() for higher-order functions, so plans using those functions failed to decode on the executor. Carry higher-order functions on BallistaFunctionRegistry (populated from the session state and the defaults) and thread them through get_task_definition, get_task_definition_vec, and all TaskContext construction sites. --- ballista/core/src/registry.rs | 42 +++++++++++++++++-- .../core/src/serde/scheduler/from_proto.rs | 14 ++++--- ballista/executor/src/execution_loop.rs | 5 ++- ballista/executor/src/executor_server.rs | 11 ++++- 4 files changed, 59 insertions(+), 13 deletions(-) diff --git a/ballista/core/src/registry.rs b/ballista/core/src/registry.rs index fe8093cae2..8ec1d97539 100644 --- a/ballista/core/src/registry.rs +++ b/ballista/core/src/registry.rs @@ -19,6 +19,7 @@ use datafusion::common::DataFusionError; use datafusion::execution::{FunctionRegistry, SessionState}; use datafusion::functions::all_default_functions; use datafusion::functions_aggregate::all_default_aggregate_functions; +use datafusion::functions_nested::all_default_higher_order_functions; use datafusion::functions_window::all_default_window_functions; use datafusion::logical_expr::HigherOrderUDF; use datafusion::logical_expr::planner::ExprPlanner; @@ -42,6 +43,8 @@ pub struct BallistaFunctionRegistry { pub aggregate_functions: HashMap>, /// Window user-defined functions. pub window_functions: HashMap>, + /// Higher-order (lambda) user-defined functions. + pub higher_order_functions: HashMap>, } impl Default for BallistaFunctionRegistry { @@ -63,6 +66,12 @@ impl Default for BallistaFunctionRegistry { .map(|f| (f.name().to_string(), f)) .collect(); + let higher_order_functions: HashMap> = + all_default_higher_order_functions() + .into_iter() + .map(|f| (f.name().to_string(), f)) + .collect(); + #[cfg(feature = "spark-compat")] let (scalar_functions, aggregate_functions, window_functions) = { let mut scalar_functions = scalar_functions; @@ -86,6 +95,7 @@ impl Default for BallistaFunctionRegistry { scalar_functions, aggregate_functions, window_functions, + higher_order_functions, } } } @@ -108,16 +118,18 @@ impl FunctionRegistry for BallistaFunctionRegistry { } fn higher_order_function_names(&self) -> HashSet { - HashSet::new() + self.higher_order_functions.keys().cloned().collect() } fn higher_order_function( &self, name: &str, ) -> datafusion::common::Result> { - Err(DataFusionError::Internal(format!( - "There is no higher-order function named \"{name}\" in the TaskContext" - ))) + self.higher_order_functions.get(name).cloned().ok_or_else(|| { + DataFusionError::Internal(format!( + "There is no higher-order function named \"{name}\" in the TaskContext" + )) + }) } fn udf(&self, name: &str) -> datafusion::common::Result> { @@ -156,11 +168,13 @@ impl From<&SessionState> for BallistaFunctionRegistry { let scalar_functions = state.scalar_functions().clone(); let aggregate_functions = state.aggregate_functions().clone(); let window_functions = state.window_functions().clone(); + let higher_order_functions = state.higher_order_functions().clone(); Self { scalar_functions, aggregate_functions, window_functions, + higher_order_functions, } } } @@ -178,6 +192,26 @@ mod tests { assert!(registry.udfs().contains("abs")); } + #[test] + fn test_higher_order_functions_available() { + let registry = BallistaFunctionRegistry::default(); + + // DataFusion ships built-in higher-order functions (e.g. array_filter); + // the registry the executor runs with must carry them or those queries + // fail to plan on the executor. + let names = registry.higher_order_function_names(); + assert!( + !names.is_empty(), + "registry must expose DataFusion's built-in higher-order functions" + ); + for name in &names { + assert!( + registry.higher_order_function(name).is_ok(), + "named higher-order function {name} must be retrievable" + ); + } + } + #[test] #[cfg(not(feature = "spark-compat"))] fn test_spark_functions_unavailable_without_feature() { diff --git a/ballista/core/src/serde/scheduler/from_proto.rs b/ballista/core/src/serde/scheduler/from_proto.rs index 34b357ad5e..e6675194a8 100644 --- a/ballista/core/src/serde/scheduler/from_proto.rs +++ b/ballista/core/src/serde/scheduler/from_proto.rs @@ -19,7 +19,7 @@ use chrono::{TimeZone, Utc}; use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion::execution::TaskContext; -use datafusion::logical_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion::logical_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; use datafusion::physical_plan::metrics::{ Count, Gauge, MetricValue, MetricsSet, PruningMetrics, RatioMetrics, Time, Timestamp, }; @@ -338,6 +338,7 @@ impl Into for protobuf::ExecutorData { /// /// This function deserializes the execution plan from the protobuf representation /// and constructs a complete task definition with the provided runtime configuration. +#[allow(clippy::too_many_arguments)] pub fn get_task_definition( task: protobuf::TaskDefinition, produce_runtime: RuntimeProducer, @@ -345,6 +346,7 @@ pub fn get_task_definition>, aggregate_functions: HashMap>, window_functions: HashMap>, + higher_order_functions: HashMap>, codec: BallistaCodec, ) -> Result { let session_config = session_config.update_from_key_value_pair(&task.props); @@ -354,6 +356,7 @@ pub fn get_task_definition>, aggregate_functions: HashMap>, window_functions: HashMap>, + higher_order_functions: HashMap>, codec: BallistaCodec, ) -> Result, BallistaError> { let session_config = session_config.update_from_key_value_pair(&multi_task.props); @@ -420,6 +424,7 @@ pub fn get_task_definition_vec< scalar_functions: scalar_functions.clone(), aggregate_functions: aggregate_functions.clone(), window_functions: window_functions.clone(), + higher_order_functions: higher_order_functions.clone(), }); let ctx = TaskContext::new( @@ -427,8 +432,7 @@ pub fn get_task_definition_vec< uuid::Uuid::new_v4().to_string(), session_config.clone(), scalar_functions, - // higher-order functions: none registered in Ballista - Default::default(), + higher_order_functions, aggregate_functions, window_functions, runtime.clone(), diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index d647ec22a6..a8b08efd7e 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -256,6 +256,8 @@ async fn run_received_task ExecutorServer ExecutorGrpc self.executor.function_registry.scalar_functions.clone(), self.executor.function_registry.aggregate_functions.clone(), self.executor.function_registry.window_functions.clone(), + self.executor + .function_registry + .higher_order_functions + .clone(), self.codec.clone(), ) .map_err(|e| Status::invalid_argument(format!("{e}")))?, @@ -855,6 +858,10 @@ impl ExecutorGrpc self.executor.function_registry.scalar_functions.clone(), self.executor.function_registry.aggregate_functions.clone(), self.executor.function_registry.window_functions.clone(), + self.executor + .function_registry + .higher_order_functions + .clone(), self.codec.clone(), ) .map_err(|e| Status::invalid_argument(format!("{e}")))?; From 8e2d2733f624906f947d94fa36a3210427b07998 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 29 Jun 2026 09:51:13 -0600 Subject: [PATCH 6/8] fix: warn when an unrecognized scan source is left unrestricted restrict_scan_to_partition only restricts FileScanConfig, the source that hands out file groups from a shared work-queue. Other sources (e.g. MemorySourceConfig) isolate partitions in their own open(partition) and need no restriction, so they are left unchanged silently. Warn for any source type that is neither, so a future source that distributes work across partitions and could over-read in a single-partition task is surfaced rather than failing silently. --- ballista/executor/src/execution_engine.rs | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index f559af2936..e579621c9e 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -27,6 +27,7 @@ use ballista_core::execution_plans::{ShuffleReaderExec, ShuffleWriterExec}; use ballista_core::serde::protobuf::ShuffleWritePartition; use ballista_core::{JobId, utils}; use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::{ FileGroup, FileScanConfig, FileScanConfigBuilder, }; @@ -36,6 +37,7 @@ use datafusion::execution::context::TaskContext; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::MetricsSet; use datafusion::prelude::SessionConfig; +use log::warn; use std::any::Any; use std::fmt::{Debug, Display}; use std::sync::Arc; @@ -120,13 +122,28 @@ impl DefaultExecutionEngine { /// Returns `None` for any node that is not a file-backed `DataSourceExec`, and /// for a `partition_id` outside the source's file groups (e.g. when an operator /// between the scan and the stage output changed the partition count). +/// +/// Only `FileScanConfig` distributes work from the shared queue. Other data +/// sources (e.g. `MemorySourceConfig`) isolate partitions in their own +/// `open(partition)` and need no restriction, so they are left unchanged. An +/// unrecognized source type is warned about, since a future source that +/// distributes work across partitions would over-read here without handling. fn restrict_scan_to_partition( plan: &Arc, partition_id: usize, ) -> Option> { let exec = plan.downcast_ref::()?; let source: &dyn Any = exec.data_source().as_ref(); - let config = source.downcast_ref::()?; + let Some(config) = source.downcast_ref::() else { + if source.downcast_ref::().is_none() { + warn!( + "restrict_scan_to_partition: unrecognized DataSourceExec source type \ + left unrestricted; if it distributes work across partitions from a \ + shared queue, a single-partition task could over-read" + ); + } + return None; + }; if partition_id >= config.file_groups.len() { return None; } From d6c9b2bb48cee5325f324bbd5e787236ab1695f6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 29 Jun 2026 09:52:11 -0600 Subject: [PATCH 7/8] refactor: pass the full physical codec to the decode context The physical plan decode context was given default_codec, which cannot decode Ballista extension nodes nested inside an expression. Pass the Ballista codec itself, matching DataFusion's convention of threading the full extension codec; it delegates to default_codec for non-Ballista nodes, so it is a strict superset. --- ballista/core/src/serde/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 0aa1d4feb6..cc337e7c76 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -382,7 +382,7 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { ) })?; let converter = DefaultPhysicalProtoConverter {}; - let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self.default_codec.as_ref()); + let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self); match ballista_plan { PhysicalPlanType::ShuffleWriter(shuffle_writer) => { let input = inputs[0].clone(); From c7fd20e99c2d524edea1f43c3525d06e2d3e54c0 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 29 Jun 2026 09:52:33 -0600 Subject: [PATCH 8/8] refactor: drop DF53 empty-projection serde workaround datafusion-proto 53.1.0 could not distinguish a FilterExec with an empty projection (Some(vec![])) from a full projection (None), so the planner rewrote empty-projection filters into a ProjectionExec-wrapped equivalent before serialization (#1838). DataFusion 54 encodes None as the full column list and Some(vec![]) as an empty list and reconstructs None only when all columns are present, so the distinction round-trips correctly. Remove make_filter_projection_serde_safe and apply it nowhere; add a physical serde round-trip test asserting a zero-column FilterExec keeps its empty projection. --- ballista/core/src/serde/mod.rs | 49 ++++++++++++++ ballista/scheduler/src/planner.rs | 104 ++---------------------------- 2 files changed, 53 insertions(+), 100 deletions(-) diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index cc337e7c76..0206afe5e8 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -810,6 +810,55 @@ mod test { ])) } + // Regression coverage for #1838 and the removed `make_filter_projection_serde_safe` + // workaround: a `FilterExec` that projects to zero columns must survive physical + // plan serialization with its empty projection intact. datafusion-proto 53.1.0 + // could not distinguish `Some(vec![])` (empty projection) from `None` (full + // projection) and decoded the former back as `None`, shifting column indices. + // DataFusion 54 preserves the distinction, so no Ballista-side rewrite is needed. + #[tokio::test] + async fn filter_exec_empty_projection_survives_physical_serde() { + use datafusion::physical_plan::empty::EmptyExec; + use datafusion::physical_plan::expressions::lit; + use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; + use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, + }; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input: Arc = Arc::new(EmptyExec::new(schema)); + let filter = FilterExecBuilder::new(lit(true), input) + .apply_projection(Some(vec![])) + .unwrap() + .build() + .unwrap(); + let plan: Arc = Arc::new(filter); + assert_eq!( + plan.schema().fields().len(), + 0, + "precondition: filter projects to zero columns" + ); + + let codec = DefaultPhysicalExtensionCodec {}; + let proto = PhysicalPlanNode::try_from_physical_plan(plan, &codec).unwrap(); + let ctx = SessionContext::new().task_ctx(); + let decoded = proto.try_into_physical_plan(&ctx, &codec).unwrap(); + + assert_eq!( + decoded.schema().fields().len(), + 0, + "decoded FilterExec must still project zero columns" + ); + let filter = decoded + .downcast_ref::() + .expect("decoded plan must be a FilterExec"); + assert_eq!( + filter.projection().as_ref().map(|p| p.is_empty()), + Some(true), + "empty projection must round-trip as Some(vec![]), not None" + ); + } + #[tokio::test] async fn test_unresolved_shuffle_exec_roundtrip() { let schema = create_test_schema(); diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index ce42b26581..95fcc6c9b0 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -33,15 +33,11 @@ use ballista_core::{ serde::scheduler::PartitionLocation, }; use datafusion::arrow::datatypes::DataType; -use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::config::ConfigOptions; -use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_optimizer::enforce_sorting::EnforceSorting; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; -use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::{ @@ -111,10 +107,8 @@ impl DistributedPlanner for DefaultDistributedPlanner { config: &ConfigOptions, ) -> Result>> { debug!("Planning query stages for job: [{job_id}]"); - // Workaround until DF 54 migration : - let serde_safe_plan = make_filter_projection_serde_safe(execution_plan)?; let (new_plan, mut stages) = - self.plan_query_stages_internal(job_id, serde_safe_plan, config)?; + self.plan_query_stages_internal(job_id, execution_plan, config)?; stages.push(create_shuffle_writer_with_config( job_id, self.next_stage_id(), @@ -411,40 +405,6 @@ impl DefaultDistributedPlanner { } } -/// Workaround until DF 54 migration: -/// datafusion-proto 53.1.0 encodes both `Some([])` (zero cols) and `None` (all cols) -/// as an empty list, decoding back to `None` and shifting column indices (#1838). -/// Rewrite `FilterExec(Some([]))` → empty `ProjectionExec` wrapped `FilterExec(None)` -fn make_filter_projection_serde_safe( - plan: Arc, -) -> datafusion::error::Result> { - Ok(plan - .transform_up(|node| { - if let Some(filter) = node.downcast_ref::() { - let empty_projection = filter - .projection() - .as_ref() - .map(|p| p.is_empty()) - .unwrap_or(false); - if empty_projection { - let orig_filter_exec = FilterExecBuilder::from(filter) - .apply_projection(None)? - .build()?; - let zero_expr: Vec<(Arc, String)> = vec![]; - let proj = ProjectionExec::try_new( - zero_expr, - Arc::new(orig_filter_exec) as Arc, - )?; - return Ok( - Transformed::yes(Arc::new(proj) as Arc), - ); - } - } - Ok(Transformed::no(node)) - })? - .data) -} - fn create_unresolved_shuffle( shuffle_writer: &dyn ShuffleWriter, ) -> Arc { @@ -623,15 +583,13 @@ mod test { use ballista_core::execution_plans::{SortShuffleWriterExec, UnresolvedShuffleExec}; use ballista_core::serde::BallistaCodec; use datafusion::arrow::compute::SortOptions; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion::execution::TaskContext; use datafusion::physical_expr::expressions::Column; use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; - use datafusion::physical_plan::empty::EmptyExec; - use datafusion::physical_plan::expressions::lit; + use datafusion::physical_plan::filter::FilterExec; - use datafusion::physical_plan::filter::FilterExecBuilder; + use datafusion::physical_plan::joins::HashJoinExec; use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::sorts::sort::SortExec; @@ -657,60 +615,6 @@ mod test { }; } - /// Regression test for issue #1838: a `FilterExec` projecting to zero columns - /// is not serialization-safe in datafusion-proto 53.1.0 (the empty projection - /// decodes back as `None`, shifting downstream column indices). The planner - /// must rewrite it into a serde-safe equivalent that still emits zero columns. - #[test] - fn empty_projection_filter_is_rewritten_serde_safe() -> Result<(), BallistaError> { - // 1-column input feeding a FilterExec that projects to ZERO columns - // (mirrors TPC-DS Q9's `FROM reason WHERE r_reason_sk = 1`). - let schema = Arc::new(Schema::new(vec![Field::new( - "r_reason_sk", - DataType::Int32, - false, - )])); - let input: Arc = Arc::new(EmptyExec::new(schema)); - let filter = FilterExecBuilder::new(lit(true), input) - .apply_projection(Some(vec![]))? - .build()?; - let plan: Arc = Arc::new(filter); - assert_eq!( - plan.schema().fields().len(), - 0, - "precondition: filter projects to zero columns" - ); - - let serde_safe_plan = super::make_filter_projection_serde_safe(plan)?; - - // schema is preserved - assert_eq!( - serde_safe_plan.schema().fields().len(), - 0, - "rewrite must preserve the zero-column output schema" - ); - - // FilterExec does not retain an empty projection - let mut has_empty_filter_projection = false; - serde_safe_plan.apply(|node| { - if let Some(f) = node.downcast_ref::() - && f.projection() - .as_ref() - .map(|p| p.is_empty()) - .unwrap_or(false) - { - has_empty_filter_projection = true; - } - Ok(TreeNodeRecursion::Continue) - })?; - assert!( - !has_empty_filter_projection, - "FilterExec with empty projection must be rewritten to be serde-safe" - ); - - Ok(()) - } - #[tokio::test] async fn distributed_aggregate_plan() -> Result<(), BallistaError> { let ctx = datafusion_test_context("testdata").await?;