diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 86d3a5a203..994cd087da 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -61,7 +61,7 @@ jobs: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} - - uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: enable-cache: true @@ -80,7 +80,7 @@ jobs: with: python-version: "3.12" - - uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: enable-cache: true @@ -103,7 +103,7 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v7.0.1 - - uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: enable-cache: true - name: Generate license file @@ -154,7 +154,7 @@ jobs: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} - - uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: enable-cache: true @@ -203,7 +203,7 @@ jobs: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} - - uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: enable-cache: true @@ -248,7 +248,7 @@ jobs: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} - - uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: enable-cache: true @@ -293,7 +293,7 @@ jobs: version: "27.4" repo-token: ${{ secrets.GITHUB_TOKEN }} - - uses: astral-sh/setup-uv@11f9893b081a58869d3b5fccaea48c9e9e46f990 # v8.3.2 + - uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 with: enable-cache: true diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 6fe3651c99..8e621a1edd 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -48,11 +48,11 @@ jobs: persist-credentials: false - name: Initialize CodeQL - uses: github/codeql-action/init@7188fc363630916deb702c7fdcf4e481b751f97a # v4 + uses: github/codeql-action/init@f205ea1c3313d32999d8d6a48b4f6530d4437b38 # v4 with: languages: actions - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@7188fc363630916deb702c7fdcf4e481b751f97a # v4 + uses: github/codeql-action/analyze@f205ea1c3313d32999d8d6a48b4f6530d4437b38 # v4 with: category: "/language:actions" diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 91459404ca..25a55c0f19 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -238,22 +238,16 @@ jobs: rustup toolchain install stable rustup default stable - uses: swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 #2.9.1 - - name: Install cargo-tomlfmt - uses: taiki-e/install-action@7572810d7dd469b651bb7793945692cf78da5dd7 # v2.85.0 + - name: Install taplo + uses: taiki-e/install-action@67729d5c413db75907f0ad1e39bb04b9c868ff60 # v2.85.7 with: - tool: cargo-tomlfmt@0.2.1 + tool: taplo-cli@0.10.0 - name: Check Cargo.toml formatting run: | - # if you encounter error, try rerun the command below, finally run 'git diff' to - # check which Cargo.toml introduces formatting violation - # - # ignore ./Cargo.toml because putting workspaces in multi-line lists make it easy to read + # if you encounter an error, run 'ci/scripts/rust_toml_fmt.sh --write' + # (or 'taplo format') locally to fix the formatting, then commit the result. ci/scripts/rust_toml_fmt.sh - if test -f "./Cargo.toml.bak"; then - echo "cargo tomlfmt found format violations" - exit 1 - fi datafusion-proto-sync-check: name: Check vendored DataFusion proto is in sync diff --git a/.github/workflows/tpcds.yml b/.github/workflows/tpcds.yml index e153b882e8..fa2e282045 100644 --- a/.github/workflows/tpcds.yml +++ b/.github/workflows/tpcds.yml @@ -52,14 +52,30 @@ on: jobs: tpcds-sf1: - name: TPC-DS SF1 (all queries, static planner) + name: TPC-DS SF1 (static planner, ${{ matrix.label }}) runs-on: ubuntu-latest container: image: amd64/rust - # A passing run takes ~7 minutes. Cap the job well under the 6-hour + # A passing leg takes ~7 minutes. Cap the job well under the 6-hour # default so a hung query fails fast and frees the runner rather than # sitting for hours with no useful signal. timeout-minutes: 45 + strategy: + # One red leg must not cancel the other; the two task-packing modes are + # independent signals. + fail-fast: false + matrix: + include: + - label: "1 partition per task" + task_args: "" + slug: "mpt1" + # 4 rather than a larger cap because the scheduler clamps the slice to + # the executor's free vcores (`budget.vcores.min(cap)`), and the + # executor below runs `--concurrent-tasks 4`. A higher cap would be + # unreachable and the label would overstate what is covered. + - label: "4 partitions per task" + task_args: "-c ballista.scheduler.max_partitions_per_task=4" + slug: "mpt4" steps: - name: Install dependencies run: | @@ -76,6 +92,10 @@ jobs: rust-version: stable - uses: swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 #2.9.1 + with: + # Share the build cache across all matrix legs — the compiled + # binaries are identical; only the per-suite CLI args differ. + shared-key: tpcds-sf1 - name: Build Ballista binaries run: | @@ -157,27 +177,22 @@ jobs: done nc -z 127.0.0.1 50051 || { echo "executor did not start"; exit 1; } - # Run the suite under the default (static) planner. The tpcds binary - # internally loops all non-skipped queries and exits non-zero on any - # failure, so a single invocation covers the whole suite. + # This matrix leg runs the suite under the default (static) planner at + # one task-packing setting. The tpcds binary internally loops all + # non-skipped queries and exits non-zero on any failure, so a single + # invocation covers the whole suite. The other leg runs in parallel. # - # Coverage for the adaptive planner (AQE on) and for single-partition - # execution is being added separately; both currently fail on - # pre-existing bugs. - run_suite() { - local label="$1"; shift - echo "::group::[$label] TPC-DS suite" - ./target/tpch-ci/tpcds \ - --host 127.0.0.1 --port 50050 \ - --path "$DATA_DIR" \ - --partitions 16 \ - --verify \ - "$@" - echo "::endgroup::" - } - - run_suite "static planner" \ - -c datafusion.optimizer.prefer_hash_join=false + # Coverage for the adaptive planner (AQE on) is being added + # separately; it currently fails on pre-existing bugs. + echo "::group::[static planner, ${{ matrix.label }}] TPC-DS suite" + ./target/tpch-ci/tpcds \ + --host 127.0.0.1 --port 50050 \ + --path "$DATA_DIR" \ + --partitions 16 \ + --verify \ + -c datafusion.optimizer.prefer_hash_join=false \ + ${{ matrix.task_args }} + echo "::endgroup::" - name: Upload cluster logs on failure # `!success()` rather than `failure()` so a timed-out or cancelled @@ -186,7 +201,7 @@ jobs: if: ${{ !success() }} uses: actions/upload-artifact@v7 with: - name: tpcds-sf1-cluster-logs + name: tpcds-sf1-cluster-logs-${{ matrix.slug }} retention-days: 14 path: | ${{ runner.temp }}/scheduler.log diff --git a/.github/workflows/tpch.yml b/.github/workflows/tpch.yml index ff9770b2cf..33f4422ec8 100644 --- a/.github/workflows/tpch.yml +++ b/.github/workflows/tpch.yml @@ -98,7 +98,7 @@ jobs: -p ballista-benchmarks - name: Install tpchgen-cli - uses: taiki-e/install-action@7572810d7dd469b651bb7793945692cf78da5dd7 # v2.85.0 + uses: taiki-e/install-action@67729d5c413db75907f0ad1e39bb04b9c868ff60 # v2.85.7 with: tool: tpchgen-cli@2.0.2 diff --git a/.github/workflows/web-tui.yml b/.github/workflows/web-tui.yml index 4e5ad97dcb..d607676c2e 100644 --- a/.github/workflows/web-tui.yml +++ b/.github/workflows/web-tui.yml @@ -68,12 +68,12 @@ jobs: - uses: swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 #2.9.1 - name: Install Trunk - uses: taiki-e/install-action@7572810d7dd469b651bb7793945692cf78da5dd7 # v2.85.0 + uses: taiki-e/install-action@67729d5c413db75907f0ad1e39bb04b9c868ff60 # v2.85.7 with: tool: trunk@0.21.14 - name: Install cargo-get - uses: taiki-e/install-action@7572810d7dd469b651bb7793945692cf78da5dd7 # v2.85.0 + uses: taiki-e/install-action@67729d5c413db75907f0ad1e39bb04b9c868ff60 # v2.85.7 with: tool: cargo-get@1.4.0 diff --git a/Cargo.lock b/Cargo.lock index be8cb8847f..ca8b6ef274 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -557,9 +557,9 @@ checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" [[package]] name = "aws-config" -version = "1.9.0" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "47712fde1909402600ccfbb26e47d482d2e58bb9e9e603d9f17e67cc435a6319" +checksum = "1b180a3c8b55960db3426d8964b8745e652466a1a49fe1a2eda828046d30b5e4" dependencies = [ "aws-credential-types", "aws-runtime", @@ -577,7 +577,7 @@ dependencies = [ "bytes", "fastrand", "hex", - "http 1.4.2", + "http 1.5.0", "sha1 0.10.6", "time", "tokio", @@ -622,9 +622,9 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.8.1" +version = "1.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7816e98ee912159f45d307e5ee6bfea4a335a55aee15f7f3e32f81a6f3000f1d" +checksum = "c9007227e10b5fed2f3e0a2beff489211e2b5604c400b7a9d5d81ca9d64c24bb" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -637,7 +637,7 @@ dependencies = [ "bytes", "bytes-utils", "fastrand", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "percent-encoding", "pin-project-lite", @@ -647,9 +647,9 @@ dependencies = [ [[package]] name = "aws-sdk-sso" -version = "1.103.0" +version = "1.105.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0469f435f645ad2162cfb463b15bde37115966ee3acf2d87fb4871ee309b8401" +checksum = "6ffd0fbe7873cb548a7aa60f9573c268fff94155397fd4f14dc9f1ecaaab8516" dependencies = [ "arc-swap", "aws-credential-types", @@ -666,16 +666,16 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "regex-lite", "tracing", ] [[package]] name = "aws-sdk-ssooidc" -version = "1.105.0" +version = "1.107.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "085faefb253f770655e162b9304321e62a1e71adf7f019ee1f4454228a377b3a" +checksum = "175763eb222a46377df7aa257a3bca980ab3e96703fefc8f4d0b8da6ad2e254c" dependencies = [ "arc-swap", "aws-credential-types", @@ -692,16 +692,16 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "regex-lite", "tracing", ] [[package]] name = "aws-sdk-sts" -version = "1.108.0" +version = "1.110.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c72b08911d8128dd360fe1b22a9fec0fa8b552dde8ec828dcf20ef5ec974e9f" +checksum = "dd8b14781dfbff48984017d57167b6ea0b6471c6920ec52b44a2677c7feb3c13" dependencies = [ "arc-swap", "aws-credential-types", @@ -719,7 +719,7 @@ dependencies = [ "aws-types", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "regex-lite", "tracing", ] @@ -739,7 +739,7 @@ dependencies = [ "hex", "hmac", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "percent-encoding", "sha2 0.11.0", "time", @@ -769,7 +769,7 @@ dependencies = [ "bytes-utils", "futures-core", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "http-body-util", "percent-encoding", @@ -788,7 +788,7 @@ dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", "h2", - "http 1.4.2", + "http 1.5.0", "hyper", "hyper-rustls", "hyper-util", @@ -824,19 +824,22 @@ dependencies = [ [[package]] name = "aws-smithy-query" -version = "0.61.1" +version = "0.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd22a6ba36e3f113cb8d5b3d1fe0ed31c76ee608ef63322d753bb8d2c9479e77" +checksum = "512346c7212ab7436df2d77a16d976a468ae44a418835511d2a69269810aaf62" dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", + "aws-smithy-xml", "urlencoding", ] [[package]] name = "aws-smithy-runtime" -version = "1.12.0" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bea94a9ff8464016338c851e24b472d7131c388c88898a502e781815b2ee6045" +checksum = "07505b34e8f4b3591a4fa69e9792b52289b95488dbbc68c3c0075b7bedb245e1" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -848,7 +851,7 @@ dependencies = [ "bytes", "fastrand", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "http-body 0.4.6", "http-body 1.0.1", "http-body-util", @@ -860,16 +863,16 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.13.0" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ed1ebe6e0a95ea84570225f5a8208dec4b8f77e61a9b0d6f51773fcb4612f0" +checksum = "3b98f2e1fd67ec06618f9c291e5e495a468e60519e44c9c1979cd0521f3affdb" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api-macros", "aws-smithy-types", "bytes", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "pin-project-lite", "tokio", "tracing", @@ -895,7 +898,7 @@ checksum = "7d56e0a4e53127a632224e43633b0fe045fa9e1e3cfc68b9830f1115e103f910" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", - "http 1.4.2", + "http 1.5.0", ] [[package]] @@ -908,7 +911,7 @@ dependencies = [ "bytes", "bytes-utils", "http 0.2.12", - "http 1.4.2", + "http 1.5.0", "http-body 0.4.6", "http-body 1.0.1", "http-body-util", @@ -923,9 +926,9 @@ dependencies = [ [[package]] name = "aws-smithy-xml" -version = "0.61.1" +version = "0.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea3f68eec3607f02acd24067969ce2abc6ba16aa7d5ce59ca450ed2fb5f78957" +checksum = "ce84f71c72fee2cbbadde6e7d082f5fb466e3a84733855295fa7aafd1b31b7d8" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-schema", @@ -935,9 +938,9 @@ dependencies = [ [[package]] name = "aws-types" -version = "1.4.0" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e957a6c6dbce82b7a91f44231c09273159703769f447cbe85e854dfe9cf67f86" +checksum = "eec1cd5469f328c782dc3e33d4153cf118a54e33cbb3356d60d16f89883e1f94" dependencies = [ "aws-credential-types", "aws-smithy-async", @@ -958,7 +961,7 @@ dependencies = [ "bytes", "form_urlencoded", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "http-body-util", "hyper", @@ -989,7 +992,7 @@ checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" dependencies = [ "bytes", "futures-core", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "http-body-util", "mime", @@ -1042,6 +1045,26 @@ dependencies = [ "tokio", ] +[[package]] +name = "ballista-chaos" +version = "0.1.0" +dependencies = [ + "arrow", + "ballista", + "ballista-core", + "ballista-executor", + "ballista-scheduler", + "datafusion", + "env_logger", + "log", + "nix 0.31.3", + "reqwest 0.13.4", + "rstest", + "serde_json", + "tempfile", + "tokio", +] + [[package]] name = "ballista-cli" version = "54.0.0" @@ -1188,11 +1211,12 @@ dependencies = [ "clap 4.6.3", "dashmap", "datafusion", + "datafusion-functions-aggregate-common", "datafusion-proto", "datafusion-substrait", "futures", "graphviz-rust", - "http 1.4.2", + "http 1.5.0", "insta", "itertools 0.15.0", "log", @@ -1403,7 +1427,7 @@ dependencies = [ "futures-util", "hex", "home", - "http 1.4.2", + "http 1.5.0", "http-body-util", "hyper", "hyper-named-pipe", @@ -1533,9 +1557,9 @@ dependencies = [ [[package]] name = "bytesize" -version = "2.4.2" +version = "2.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3d7c8918969267b2932ffd5655509bbbea0833823058c378876953217f5fc50e" +checksum = "7354288c522e7e980fafd2075d63d1285794c3a6a16cdd492f189ea406e5f18b" [[package]] name = "bzip2" @@ -2103,9 +2127,9 @@ dependencies = [ [[package]] name = "ctor" -version = "1.0.10" +version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2e30e509674ef0ec91e21a7735766db37d163d46151b6a361d8b83dd79116bd" +checksum = "2d83cb7e7a873830708d6b02a78cd36a592c6fa14bf267b68725103b85c0d77f" dependencies = [ "link-section", "linktime-proc-macro", @@ -3545,7 +3569,7 @@ dependencies = [ "fnv", "futures-core", "futures-sink", - "http 1.4.2", + "http 1.5.0", "indexmap 2.14.0", "slab", "tokio", @@ -3669,9 +3693,9 @@ dependencies = [ [[package]] name = "http" -version = "1.4.2" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +checksum = "918d3568bebf352712bc2ef3d46a8bcf1a75b373be6539de198e9105cbbf9ce0" dependencies = [ "bytes", "itoa", @@ -3695,7 +3719,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" dependencies = [ "bytes", - "http 1.4.2", + "http 1.5.0", ] [[package]] @@ -3706,7 +3730,7 @@ checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" dependencies = [ "bytes", "futures-core", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "pin-project-lite", ] @@ -3749,7 +3773,7 @@ dependencies = [ "futures-channel", "futures-core", "h2", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "httparse", "httpdate", @@ -3781,7 +3805,7 @@ version = "0.27.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33ca68d021ef39cf6463ab54c1d0f5daf03377b70561305bb89a8f83aab66e0f" dependencies = [ - "http 1.4.2", + "http 1.5.0", "hyper", "hyper-util", "rustls", @@ -3814,7 +3838,7 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "hyper", "ipnet", @@ -4369,15 +4393,15 @@ dependencies = [ [[package]] name = "link-section" -version = "0.19.1" +version = "0.19.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8dc98458dfe90986c5e2f6ddcf68360c7e5c4252600153e06aa4ee8176c0f8d1" +checksum = "5ee1a0d6e252afe82e7bc2db42fba60e02ddf3b1accaf8cb21d96e34ba61f3d4" [[package]] name = "linktime-proc-macro" -version = "0.2.0" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c7b0a3383c2a1002d11349c92c85a666a5fb679e96c79d782cf0dbe557fd6ee" +checksum = "348d0075b1fc163b26d72a7f75fc5141daf2fd1bdf128d873cbaf6785d495bdf" [[package]] name = "linux-raw-sys" @@ -4795,7 +4819,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body-util", "humantime", "hyper", @@ -5548,7 +5572,7 @@ dependencies = [ "once_cell", "socket2", "tracing", - "windows-sys 0.60.2", + "windows-sys 0.59.0", ] [[package]] @@ -5930,7 +5954,7 @@ dependencies = [ "futures-core", "futures-util", "h2", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "http-body-util", "hyper", @@ -5972,7 +5996,7 @@ dependencies = [ "encoding_rs", "futures-core", "h2", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "http-body-util", "hyper", @@ -6093,9 +6117,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.42" +version = "0.23.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" +checksum = "0283386ce02abc0151e1761d08802dfe86c173b0b494af5cbc086574e453da06" dependencies = [ "aws-lc-rs", "log", @@ -6988,7 +7012,7 @@ dependencies = [ "etcetera", "ferroid", "futures", - "http 1.4.2", + "http 1.5.0", "itertools 0.14.0", "log", "memchr", @@ -7264,7 +7288,7 @@ dependencies = [ "base64", "bytes", "h2", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "http-body-util", "hyper", @@ -7350,7 +7374,7 @@ dependencies = [ "bitflags 2.13.0", "bytes", "futures-util", - "http 1.4.2", + "http 1.5.0", "http-body 1.0.1", "pin-project-lite", "tower", @@ -7367,7 +7391,7 @@ checksum = "b11f75e912b0c2be01b63d8cf8057b8c3f97cf34abb3d431a3a4c8675498e233" dependencies = [ "bitflags 2.13.0", "bytes", - "http 1.4.2", + "http 1.5.0", "percent-encoding", "pin-project-lite", "tower-layer", @@ -7642,7 +7666,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e994ba84b0bd1b1b0cf92878b7ef898a5c1760108fe7b6010327e274917a808c" dependencies = [ "base64", - "http 1.4.2", + "http 1.5.0", "httparse", "log", ] @@ -8121,7 +8145,7 @@ version = "0.52.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" dependencies = [ - "windows-targets 0.52.6", + "windows-targets", ] [[package]] @@ -8130,16 +8154,7 @@ version = "0.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" dependencies = [ - "windows-targets 0.52.6", -] - -[[package]] -name = "windows-sys" -version = "0.60.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" -dependencies = [ - "windows-targets 0.53.5", + "windows-targets", ] [[package]] @@ -8157,31 +8172,14 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" dependencies = [ - "windows_aarch64_gnullvm 0.52.6", - "windows_aarch64_msvc 0.52.6", - "windows_i686_gnu 0.52.6", - "windows_i686_gnullvm 0.52.6", - "windows_i686_msvc 0.52.6", - "windows_x86_64_gnu 0.52.6", - "windows_x86_64_gnullvm 0.52.6", - "windows_x86_64_msvc 0.52.6", -] - -[[package]] -name = "windows-targets" -version = "0.53.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" -dependencies = [ - "windows-link", - "windows_aarch64_gnullvm 0.53.1", - "windows_aarch64_msvc 0.53.1", - "windows_i686_gnu 0.53.1", - "windows_i686_gnullvm 0.53.1", - "windows_i686_msvc 0.53.1", - "windows_x86_64_gnu 0.53.1", - "windows_x86_64_gnullvm 0.53.1", - "windows_x86_64_msvc 0.53.1", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", ] [[package]] @@ -8199,96 +8197,48 @@ version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" -[[package]] -name = "windows_aarch64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" - [[package]] name = "windows_aarch64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" -[[package]] -name = "windows_aarch64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" - [[package]] name = "windows_i686_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" -[[package]] -name = "windows_i686_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" - [[package]] name = "windows_i686_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" -[[package]] -name = "windows_i686_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" - [[package]] name = "windows_i686_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" -[[package]] -name = "windows_i686_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" - [[package]] name = "windows_x86_64_gnu" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" -[[package]] -name = "windows_x86_64_gnu" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" - [[package]] name = "windows_x86_64_gnullvm" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" -[[package]] -name = "windows_x86_64_gnullvm" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" - [[package]] name = "windows_x86_64_msvc" version = "0.52.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" -[[package]] -name = "windows_x86_64_msvc" -version = "0.53.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" - [[package]] name = "winnow" version = "1.0.3" diff --git a/Cargo.toml b/Cargo.toml index c42e732a0c..b29fc92dc4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,7 +17,16 @@ [workspace] exclude = ["dev/msrvcheck", "python"] -members = ["ballista-cli", "ballista/client", "ballista/core", "ballista/executor", "ballista/scheduler", "benchmarks", "examples"] +members = [ + "ballista-cli", + "ballista/client", + "ballista/core", + "ballista/executor", + "ballista/scheduler", + "benchmarks", + "chaos-testing", + "examples", +] resolver = "3" [workspace.package] @@ -42,8 +51,10 @@ datafusion-proto-common = "54" datafusion-spark = "54" datafusion-substrait = "54" +ctor = { version = "1.0" } insta = "1.47" itertools = "0.15" +mimalloc = { version = "0.1" } object_store = "0.13.2" prost = "0.14" prost-types = "0.14" @@ -57,22 +68,20 @@ tonic-prost-build = { version = "0.14" } tracing = "0.1" tracing-appender = "0.2.2" tracing-subscriber = { version = "0.3", features = ["env-filter"] } -ctor = { version = "1.0" } -mimalloc = { version = "0.1" } -tokio = { version = "1" } -uuid = { version = "1.23", features = ["v4", "v7"] } -rand = { version = "0.10" } +async-trait = { version = "0.1" } +dashmap = { version = "6.2" } env_logger = { version = "0.11" } futures = { version = "0.3" } log = { version = "0.4" } parking_lot = { version = "0.12" } -tempfile = { version = "3.27" } -dashmap = { version = "6.2" } -async-trait = { version = "0.1" } +rand = { version = "0.10" } serde = { version = "1.0" } +tempfile = { version = "3.27" } +tokio = { version = "1" } tokio-stream = { version = "0.1" } url = { version = "2.5" } +uuid = { version = "1.23", features = ["v4", "v7"] } # cargo build --profile release-lto [profile.release-lto] @@ -120,4 +129,4 @@ codegen-units = 256 lto = false debug = false debug-assertions = false -incremental = false \ No newline at end of file +incremental = false diff --git a/ballista-cli/Cargo.toml b/ballista-cli/Cargo.toml index 0344a3b5c9..573c1de6e8 100644 --- a/ballista-cli/Cargo.toml +++ b/ballista-cli/Cargo.toml @@ -36,8 +36,8 @@ crate-type = ["cdylib", "rlib"] ballista = { path = "../ballista/client", version = "54.0.0", features = ["standalone"], optional = true } datafusion = { workspace = true, optional = true } datafusion-cli = { workspace = true, optional = true } -rustyline = { version = "18.0.0", optional = true } mimalloc = { workspace = true, optional = true } +rustyline = { version = "18.0.0", optional = true } # TUI/web shared deps chrono = { version = "0.4", default-features = false, features = ["std", "clock"], optional = true } @@ -61,8 +61,8 @@ crossterm = { version = "0.29.0", features = ["event-stream"], optional = true } tracing-appender = { version = "0.2", optional = true } # Web-only deps (WASM) -critical-section = { version = "1.2.0", features = ["std"], optional = true } console_error_panic_hook = { version = "0.1", optional = true } +critical-section = { version = "1.2.0", features = ["std"], optional = true } gloo-timers = { version = "0.4", features = ["futures"], optional = true } js-sys = { version = "0.3.95", optional = true } ratzilla = { version = "0.3.1", optional = true } @@ -81,25 +81,57 @@ tokio = { workspace = true, features = ["macros", "sync", "time"] } default = ["cli", "tui"] cli = [ - "dep:ballista", "dep:datafusion", "dep:datafusion-cli", - "dep:rustyline", "dep:mimalloc", + "dep:ballista", + "dep:datafusion", + "dep:datafusion-cli", + "dep:mimalloc", + "dep:rustyline", ] tui = [ - "dep:chrono", "dep:config", "dep:crossterm", "dep:dotparser", "dep:futures", - "dep:percent-encoding", "dep:prometheus-parse", "dep:ratatui", "dep:reqwest", - "dep:serde", "dep:serde_json", "dep:tracing-appender", - "dep:tui-shimmer", "ratatui/crossterm", "ratatui/all-widgets", "ratatui/macros", - "ratatui/layout-cache", "ratatui/underline-color", "ratatui/serde" + "dep:chrono", + "dep:config", + "dep:crossterm", + "dep:dotparser", + "dep:futures", + "dep:percent-encoding", + "dep:prometheus-parse", + "dep:ratatui", + "dep:reqwest", + "dep:serde", + "dep:serde_json", + "dep:tracing-appender", + "dep:tui-shimmer", + "ratatui/crossterm", + "ratatui/all-widgets", + "ratatui/macros", + "ratatui/layout-cache", + "ratatui/underline-color", + "ratatui/serde", ] web = [ - "dep:chrono", "dep:config", "dep:critical-section", "dep:dotparser", "dep:futures", "dep:percent-encoding", - "dep:prometheus-parse", "dep:ratatui", "ratatui/serde", "dep:reqwest", - "dep:serde", "dep:serde_json", "dep:tui-shimmer", - "dep:ratzilla", "dep:wasm-bindgen", "dep:wasm-bindgen-futures", - "dep:gloo-timers", "dep:console_error_panic_hook", - "dep:tracing-web", "dep:web-sys", "dep:js-sys", + "dep:chrono", + "dep:config", + "dep:critical-section", + "dep:dotparser", + "dep:futures", + "dep:percent-encoding", + "dep:prometheus-parse", + "dep:ratatui", + "ratatui/serde", + "dep:reqwest", + "dep:serde", + "dep:serde_json", + "dep:tui-shimmer", + "dep:ratzilla", + "dep:wasm-bindgen", + "dep:wasm-bindgen-futures", + "dep:gloo-timers", + "dep:console_error_panic_hook", + "dep:tracing-web", + "dep:web-sys", + "dep:js-sys", ] [[bin]] diff --git a/ballista-cli/src/tui/domain/jobs.rs b/ballista-cli/src/tui/domain/jobs.rs index 2129efe05b..eb22663709 100644 --- a/ballista-cli/src/tui/domain/jobs.rs +++ b/ballista-cli/src/tui/domain/jobs.rs @@ -117,9 +117,14 @@ impl JobsData { } }), SortColumn::StagesCompleted => jobs.sort_by(|a, b| { - let a_stages = a.completed_stages / a.num_stages; - let b_stages = b.completed_stages / b.num_stages; - let cmp = a_stages.cmp(&b_stages); + let stage_completion = |job: &Job| { + if job.num_stages == 0 { + 0.0 + } else { + job.completed_stages as f64 / job.num_stages as f64 + } + }; + let cmp = stage_completion(a).total_cmp(&stage_completion(b)); if self.sort_order == crate::tui::domain::SortOrder::Descending { cmp.reverse() } else { @@ -588,6 +593,36 @@ mod tests { assert_eq!(refs[2].status, "Completed"); } + #[test] + fn sort_by_stages_completed_uses_completion_ratio() { + let jobs = vec![ + make_job("quarter", "A", "Running", 1, 2, 4, 1, 0), + make_job("three_quarters", "B", "Running", 2, 3, 4, 3, 0), + make_job("half", "C", "Running", 3, 4, 4, 2, 0), + ]; + let data = + make_jobs_data(jobs, SortColumn::StagesCompleted, SortOrder::Ascending); + let mut refs: Vec<&Job> = data.jobs.iter().collect(); + data.sort_jobs(&mut refs); + assert_eq!(refs[0].job_id, "quarter"); + assert_eq!(refs[1].job_id, "half"); + assert_eq!(refs[2].job_id, "three_quarters"); + } + + #[test] + fn sort_by_stages_completed_handles_zero_stages() { + let jobs = vec![ + make_job("zero", "A", "Running", 1, 2, 0, 0, 0), + make_job("half", "B", "Running", 2, 3, 2, 1, 0), + ]; + let data = + make_jobs_data(jobs, SortColumn::StagesCompleted, SortOrder::Ascending); + let mut refs: Vec<&Job> = data.jobs.iter().collect(); + data.sort_jobs(&mut refs); + assert_eq!(refs[0].job_id, "zero"); + assert_eq!(refs[1].job_id, "half"); + } + #[test] fn sort_by_percent_complete_ascending() { let jobs = vec![ diff --git a/ballista/client/Cargo.toml b/ballista/client/Cargo.toml index 062b3435d2..862d4eeb6f 100644 --- a/ballista/client/Cargo.toml +++ b/ballista/client/Cargo.toml @@ -30,7 +30,9 @@ rust-version = { workspace = true } [dependencies] async-trait = { workspace = true } ballista-core = { path = "../core", version = "54.0.0" } -ballista-executor = { path = "../executor", version = "54.0.0", optional = true, default-features = false, features = ["arrow-ipc-optimizations"] } +ballista-executor = { path = "../executor", version = "54.0.0", optional = true, default-features = false, features = [ + "arrow-ipc-optimizations", +] } ballista-scheduler = { path = "../scheduler", version = "54.0.0", optional = true, default-features = false } datafusion = { workspace = true } log = { workspace = true } @@ -46,8 +48,8 @@ datafusion-proto = { workspace = true } env_logger = { workspace = true } rstest = { workspace = true } tempfile = { workspace = true } -uuid = { workspace = true } tonic = { workspace = true } +uuid = { workspace = true } [features] default = ["standalone"] diff --git a/ballista/client/tests/null_aware.rs b/ballista/client/tests/null_aware.rs new file mode 100644 index 0000000000..0815afe5d3 --- /dev/null +++ b/ballista/client/tests/null_aware.rs @@ -0,0 +1,152 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#[cfg(test)] +#[cfg(feature = "standalone")] +mod null_aware { + use std::fs; + use std::path::Path; + + use ballista::prelude::SessionContextExt; + use datafusion::arrow::util::pretty::pretty_format_batches; + use datafusion::prelude::*; + + fn write_tables(dir: &Path, t2_has_null: bool) { + // t1: 4 files -> 4 partitions, values 0..19 + let t1 = dir.join("t1"); + fs::create_dir_all(&t1).unwrap(); + for i in 0..4i32 { + let mut s = String::from("a\n"); + for v in 0..5i32 { + s.push_str(&format!("{}\n", i * 5 + v)); + } + fs::write(t1.join(format!("p{i}.csv")), s).unwrap(); + } + + // t2: 4 files -> 4 partitions. One file optionally holds a NULL key. + // The second column keeps the empty field unambiguous, since a bare + // blank line is skipped by the CSV reader rather than read as NULL. + let t2 = dir.join("t2"); + fs::create_dir_all(&t2).unwrap(); + fs::write(t2.join("p0.csv"), "b,tag\n0,x\n1,x\n").unwrap(); + fs::write( + t2.join("p1.csv"), + if t2_has_null { + "b,tag\n,x\n" + } else { + "b,tag\n2,x\n" + }, + ) + .unwrap(); + fs::write(t2.join("p2.csv"), "b,tag\n3,x\n").unwrap(); + fs::write(t2.join("p3.csv"), "b,tag\n4,x\n").unwrap(); + } + + async fn register(ctx: &SessionContext, dir: &Path) { + for name in ["t1", "t2"] { + ctx.register_csv( + name, + dir.join(name).to_str().unwrap(), + CsvReadOptions::new() + .has_header(true) + .file_extension(".csv"), + ) + .await + .unwrap(); + } + } + + const QUERY: &str = "select a from t1 where a not in (select b from t2) order by a"; + + async fn run_case(case: &str, t2_has_null: bool) { + let dir = std::env::temp_dir().join(format!("ballista_null_aware_{case}")); + let _ = fs::remove_dir_all(&dir); + fs::create_dir_all(&dir).unwrap(); + write_tables(&dir, t2_has_null); + + let df_ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(4), + ); + register(&df_ctx, &dir).await; + let df_batches = df_ctx.sql(QUERY).await.unwrap().collect().await.unwrap(); + let expected = pretty_format_batches(&df_batches).unwrap().to_string(); + + let mut mismatches = vec![]; + // "default" exercises the NOT IN logical rewrite under Ballista's + // stock configuration (sort-merge joins): the rewrite removes the + // subquery before decorrelation, so no null-aware join is planned at + // all. "prefer_hash_join" and "aqe" exercise the rewrite under the + // hash-join planning paths. "rewrite_disabled" turns the rewrite off + // with hash joins preferred, covering the fallback path where the + // null-aware join survives and is lowered to a single task. + for variant in ["default", "prefer_hash_join", "aqe", "rewrite_disabled"] { + let ctx = SessionContext::standalone().await.unwrap(); + if variant == "aqe" { + ctx.sql("SET ballista.planner.adaptive.enabled = true") + .await + .unwrap() + .collect() + .await + .unwrap(); + } + if variant != "default" { + ctx.sql("SET datafusion.optimizer.prefer_hash_join = true") + .await + .unwrap() + .collect() + .await + .unwrap(); + } + if variant == "rewrite_disabled" { + ctx.sql("SET ballista.optimizer.not_in_subquery_rewrite = false") + .await + .unwrap() + .collect() + .await + .unwrap(); + } + register(&ctx, &dir).await; + + match ctx.sql(QUERY).await.unwrap().collect().await { + Ok(batches) => { + let actual = pretty_format_batches(&batches).unwrap().to_string(); + if actual.trim() != expected.trim() { + mismatches.push(format!( + "[{case}/{variant}] expected:\n{expected}\nactual:\n{actual}" + )); + } + } + Err(error) => { + mismatches.push(format!("[{case}/{variant}] failed: {error}")); + } + } + } + + let _ = fs::remove_dir_all(&dir); + assert!(mismatches.is_empty(), "{}", mismatches.join("\n\n")); + } + + #[tokio::test] + async fn not_in_with_null_in_subquery() { + run_case("with_null", true).await; + } + + #[tokio::test] + async fn not_in_without_null_in_subquery() { + run_case("without_null", false).await; + } +} diff --git a/ballista/core/proto/ballista.proto b/ballista/core/proto/ballista.proto index c2471c7640..99fceb91ea 100644 --- a/ballista/core/proto/ballista.proto +++ b/ballista/core/proto/ballista.proto @@ -55,6 +55,7 @@ message BallistaPhysicalPlanNode { BufferExecNode buffer = 7; UnorderedRangeRepartitionExecNode unordered_range_repartition = 8; OrderedRangeRepartitionExecNode ordered_range_repartition = 9; + PerPartitionFilterExecNode per_partition_filter = 10; } } @@ -124,6 +125,14 @@ message UnorderedRangeRepartitionExecNode { uint32 output_partitions = 2; } +// Filter with per-input-partition predicates. `predicates[k]` is the +// boolean expression applied to input partition `k`. Requires +// `predicates.len() == input_partition_count`. The child plan is +// plumbed by the framework as `inputs[0]` during decode. +message PerPartitionFilterExecNode { + repeated datafusion.PhysicalExprNode predicates = 1; +} + message ChaosExecNode { double failure_probability = 1; string fault_type = 2; @@ -592,6 +601,40 @@ message SuccessfulTask { // TODO tasks are currently always shuffle writes but this will not always be the case // so we might want to think about some refactoring of the task definitions repeated ShuffleWritePartition partitions = 2; + // Reports from `RuntimeStatsExec` operators in this task's plan that + // are still valid at the plan's output (walked from the top through + // distribution-preserving nodes only — see + // `range_repartition_common::preserves_distribution`). Empty when the + // plan has no such stats-taps. Currently reports one entry per + // executed `RuntimeStatsExec`; the scheduler groups by `order_by` tag + // to combine reports across tasks/executors. + repeated RuntimeStatsReport runtime_stats = 3; +} + +// One report per `RuntimeStatsExec` in the executed plan. +message RuntimeStatsReport { + // What the operator was sampling. Empty = row-count-only mode. When + // non-empty, the first entry identifies which routing expression this + // sketch describes; the scheduler groups sketches by this tag to + // combine samples across tasks that were sampling the same expression. + repeated datafusion.PhysicalSortExprNode order_by = 1; + // Per-partition observations. Interpretation depends on the operator's + // position in the plan — pre-repartition gets one entry per input + // partition, post-repartition gets one per output sub-partition. The + // scheduler groups by `order_by` tag and aggregates. + repeated RuntimeStatsPartitionEntry partitions = 2; +} + +// One partition's observations from a `RuntimeStatsExec`. +message RuntimeStatsPartitionEntry { + uint32 partition_id = 1; + uint64 row_count = 2; + // Present when the `RuntimeStatsExec` was in sketch mode AND this + // partition observed at least one non-null routing value. + optional QuantileSketchState sketch = 3; + // TODO: `optional MinMaxState min_max` — for a lighter post-repartition + // mode where the bin-packer just needs (min, max, count) per + // sub-partition and a full T-Digest is overkill. } message ExecutionError { diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index fd3060297d..b2d7ed6624 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -97,7 +97,9 @@ pub const BALLISTA_SHUFFLE_SORT_BASED_MEMORY_LIMIT_PER_TASK_BYTES: &str = "ballista.shuffle.sort_based.memory_limit_per_task_bytes"; /// Configuration key for the byte-size threshold below which a hash join's /// smaller side is promoted to `CollectLeft` and lowered via the broadcast -/// pattern in the distributed planner. Set to `0` to disable promotion. +/// pattern in the distributed planner. It also caps null-aware anti joins with +/// a known build size because they require single-task `CollectLeft` execution. +/// Set to `0` to disable promotion and reject null-aware anti joins. pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_BYTES: &str = "ballista.optimizer.broadcast_join_threshold_bytes"; @@ -117,6 +119,15 @@ pub const BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS: &str = pub const BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES: &str = "ballista.optimizer.hash_join_max_build_partition_bytes"; +/// Configuration key controlling the logical rewrite of uncorrelated +/// `NOT IN (subquery)` filter predicates into a plain anti join plus a +/// one-row count aggregate. The rewrite avoids DataFusion's null-aware hash +/// join, which Ballista must otherwise execute in a single task. Enabled by +/// default; set to `false` to keep the null-aware join and its single-task +/// lowering. +pub const BALLISTA_NOT_IN_SUBQUERY_REWRITE: &str = + "ballista.optimizer.not_in_subquery_rewrite"; + /// Configuration key to enable AQE coalesce-shuffle-partitions rule. /// Disabled by default — opt in when the workload benefits from larger /// downstream tasks more than from preserved parallelism. @@ -255,8 +266,10 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| "Byte-size threshold below which a hash join's smaller side is \ promoted to CollectLeft and lowered via the broadcast pattern. \ Governs broadcast selection under both the static distributed \ - planner and adaptive query planning (AQE). Set to 0 to disable \ - promotion.".to_string(), + planner and adaptive query planning (AQE). It also caps \ + null-aware anti joins with a known build size because they require \ + single-task CollectLeft execution. Set to 0 to disable promotion \ + and reject null-aware anti joins.".to_string(), DataType::UInt64, Some((10 * 1024 * 1024).to_string())), ConfigEntry::new(BALLISTA_BROADCAST_JOIN_THRESHOLD_ROWS.to_string(), @@ -274,6 +287,14 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| which makes AQE use a hash join regardless of build size.".to_string(), DataType::UInt64, Some((64 * 1024 * 1024).to_string())), + ConfigEntry::new(BALLISTA_NOT_IN_SUBQUERY_REWRITE.to_string(), + "Rewrites uncorrelated NOT IN (subquery) filter predicates into a \ + plain anti join plus a one-row count aggregate during logical \ + optimization. The rewrite avoids DataFusion's null-aware hash join, \ + which Ballista must otherwise execute in a single task. Set to false \ + to keep the null-aware join and its single-task lowering.".to_string(), + DataType::Boolean, + Some("true".to_string())), ConfigEntry::new(BALLISTA_CLIENT_PULL.to_string(), "Should client employ pull or push job tracking. In pull mode client will make a request to server in the loop, until job finishes. Pull mode is kept for legacy clients.".to_string(), DataType::Boolean, @@ -668,6 +689,13 @@ impl BallistaConfig { self.get_usize_setting(BALLISTA_HASH_JOIN_MAX_BUILD_PARTITION_BYTES) } + /// Whether uncorrelated `NOT IN (subquery)` filter predicates are rewritten + /// into a plain anti join plus a one-row count aggregate during logical + /// optimization, avoiding the single-task null-aware hash join. + pub fn not_in_subquery_rewrite_enabled(&self) -> bool { + self.get_bool_setting(BALLISTA_NOT_IN_SUBQUERY_REWRITE) + } + /// Returns whether the AQE coalesce-shuffle-partitions rule is enabled. pub fn coalesce_enabled(&self) -> bool { self.get_bool_setting(BALLISTA_COALESCE_ENABLED) @@ -938,6 +966,32 @@ mod tests { Ok(()) } + #[test] + fn cluster_config_plumbs_grpc_max_message_size() { + // Sanity check that setting `ballista.client.grpc_max_message_size` + // via `SessionConfig::options_mut().set(...)` — the path Python + // clients' `cluster_config` overrides use — actually reaches the + // reader consulted by `distributed_query`. + use crate::extension::SessionConfigExt; + use datafusion::prelude::SessionConfig; + + let mut config = SessionConfig::new_with_ballista(); + assert_eq!( + config.ballista_config().grpc_client_max_message_size(), + 16 * 1024 * 1024, + ); + + let r = config + .options_mut() + .set("ballista.client.grpc_max_message_size", "67108864"); + assert!(r.is_ok(), "set failed: {r:?}"); + + assert_eq!( + config.ballista_config().grpc_client_max_message_size(), + 64 * 1024 * 1024, + ); + } + // The default must stay non-zero: `0` disables the fit check, and since AQE // no longer consults `prefer_hash_join`, a disabled check would leave the // Partitioned arm unconditionally on hash join and make SortMergeJoin — the diff --git a/ballista/core/src/execution_plans/mod.rs b/ballista/core/src/execution_plans/mod.rs index 7aec3cf47d..5d9240f3a3 100644 --- a/ballista/core/src/execution_plans/mod.rs +++ b/ballista/core/src/execution_plans/mod.rs @@ -23,6 +23,8 @@ mod chaos_exec; mod distributed_explain_analyze; mod distributed_query; mod ordered_range_repartition; +mod per_partition_filter; +pub mod plan_algebra; mod range_repartition_common; mod runtime_stats; mod shuffle_reader; @@ -40,7 +42,14 @@ use datafusion::common::exec_err; pub use distributed_explain_analyze::DistributedExplainAnalyzeExec; pub use distributed_query::{DistributedQueryExec, execute_physical_plan}; pub use ordered_range_repartition::OrderedRangeRepartitionExec; -pub use runtime_stats::{RuntimeStatsExec, sketch_from_proto, sketch_to_proto}; +pub use per_partition_filter::{PerPartitionFilterExec, range_partition_predicates}; +pub use plan_algebra::{preserves_distribution, preserves_partitioning}; +pub use runtime_stats::{ + MergedRuntimeStats, RuntimeStatsExec, TaskRuntimeStats, + collect_reports as collect_runtime_stats_reports, cut_partitions, + log_merged_runtime_stats, merge_reports as merge_runtime_stats_reports, + repartition_routing_expr, sketch_from_proto, sketch_to_proto, +}; pub use shuffle_reader::{CoalescePlan, PartitionGroup, ShuffleReaderExec}; pub use shuffle_reader::{stats_for_partition, stats_for_partitions}; pub use shuffle_writer::DEFAULT_SHUFFLE_CHANNEL_CAPACITY; diff --git a/ballista/core/src/execution_plans/per_partition_filter.rs b/ballista/core/src/execution_plans/per_partition_filter.rs new file mode 100644 index 0000000000..9c93a7a6d1 --- /dev/null +++ b/ballista/core/src/execution_plans/per_partition_filter.rs @@ -0,0 +1,571 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Filter with a distinct predicate per input partition. +//! +//! `FilterExec` in DataFusion carries a single predicate applied to every +//! partition. That's wrong for the range-repartition-consuming reader in +//! the adaptive range shuffle: each downstream partition `k` needs a range +//! predicate `cuts[k-1] <= key < cuts[k]` unique to that partition, so +//! straddling sub-parts from the producer are trimmed to just partition +//! `k`'s slice. +//! +//! One-task-per-downstream-partition + plain `FilterExec` would work but +//! defeats vcore packing (`K` tasks instead of `K / vcores`). This operator +//! keeps packing: `predicates[k]` is applied to `input.execute(k)`, so a +//! single task consuming several partitions still gets each partition's +//! own predicate. +//! +//! Semantics per batch mirror `FilterExec`: evaluate the boolean expr +//! against the batch, then `filter_record_batch`. No projection, no +//! coalescing, no metrics — those can grow later if the wiring warrants. + +use std::fmt::{self, Debug, Formatter}; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::compute::filter_record_batch; +use datafusion::arrow::datatypes::{DataType, SchemaRef}; +use datafusion::common::cast::as_boolean_array; +use datafusion::common::{Result, Statistics, internal_err}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::CardinalityEffect; +use datafusion::physical_plan::stream::{ + EmptyRecordBatchStream, RecordBatchStreamAdapter, +}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, +}; +use futures::{Stream, StreamExt, ready}; + +/// Filter with per-input-partition predicates. +/// +/// `predicates[k]` is applied to `input.execute(k)`. Requires +/// `predicates.len() == input.output_partitioning().partition_count()`. +pub struct PerPartitionFilterExec { + input: Arc, + predicates: Vec>, + properties: Arc, +} + +impl PerPartitionFilterExec { + /// Wrap `input` with a vector of predicates, one per input partition. + /// + /// Fails if the predicate count doesn't match the input partition count + /// or if any predicate does not evaluate to `Boolean` against the input + /// schema. + pub fn try_new( + input: Arc, + predicates: Vec>, + ) -> Result { + let partition_count = input.output_partitioning().partition_count(); + if predicates.len() != partition_count { + return internal_err!( + "PerPartitionFilterExec: predicate count {} does not match input partition count {}", + predicates.len(), + partition_count + ); + } + let schema = input.schema(); + for (k, predicate) in predicates.iter().enumerate() { + let dt = predicate.data_type(&schema)?; + if dt != DataType::Boolean { + return internal_err!( + "PerPartitionFilterExec: predicate[{k}] must evaluate to Boolean, got {dt}" + ); + } + } + let properties = Arc::new(PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + )); + Ok(Self { + input, + predicates, + properties, + }) + } + + /// The per-partition predicates. `predicates()[k]` corresponds to input partition `k`. + pub fn predicates(&self) -> &[Arc] { + &self.predicates + } +} + +impl Debug for PerPartitionFilterExec { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + f.debug_struct("PerPartitionFilterExec") + .field("num_predicates", &self.predicates.len()) + .finish() + } +} + +impl DisplayAs for PerPartitionFilterExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter<'_>) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!( + f, + "PerPartitionFilterExec: predicates=[{}]", + self.predicates + .iter() + .map(|p| p.to_string()) + .collect::>() + .join(", ") + ) + } + DisplayFormatType::TreeRender => { + write!(f, "PerPartitionFilterExec") + } + } + } +} + +impl ExecutionPlan for PerPartitionFilterExec { + fn name(&self) -> &str { + "PerPartitionFilterExec" + } + + fn schema(&self) -> SchemaRef { + self.input.schema() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + let [input] = children.as_slice() else { + return internal_err!( + "PerPartitionFilterExec expects exactly one child, got {}", + children.len() + ); + }; + Ok(Arc::new(PerPartitionFilterExec::try_new( + input.clone(), + self.predicates.clone(), + )?)) + } + + fn required_input_distribution(&self) -> Vec { + vec![Distribution::UnspecifiedDistribution] + } + + fn required_input_ordering(&self) -> Vec> { + vec![None] + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn partition_statistics(&self, _partition: Option) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema()))) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + CardinalityEffect::LowerEqual + } + + fn execute( + &self, + partition: usize, + ctx: Arc, + ) -> Result { + let Some(predicate) = self.predicates.get(partition).cloned() else { + return internal_err!( + "PerPartitionFilterExec: partition {} out of bounds ({} predicates)", + partition, + self.predicates.len() + ); + }; + let schema = self.schema(); + let input = self.input.execute(partition, ctx)?; + let stream = PerPartitionFilterStream { + schema: schema.clone(), + predicate, + input, + }; + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } +} + +struct PerPartitionFilterStream { + schema: SchemaRef, + predicate: Arc, + input: SendableRecordBatchStream, +} + +impl Stream for PerPartitionFilterStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + match ready!(self.input.poll_next_unpin(cx)) { + Some(Ok(batch)) => { + let mask = self + .predicate + .evaluate(&batch) + .and_then(|v| v.into_array(batch.num_rows()))?; + let mask = as_boolean_array(&mask)?; + let filtered = filter_record_batch(&batch, mask)?; + if filtered.num_rows() == 0 { + // Nothing left after filtering; pull the next batch + // rather than emit an empty batch downstream. + continue; + } + return Poll::Ready(Some(Ok(filtered))); + } + Some(Err(e)) => return Poll::Ready(Some(Err(e))), + None => { + // Release the input pipeline's resources on EOS — + // mirrors DataFusion's FilterExec so the input's + // child chain doesn't linger on the heap until the + // outer stream is itself dropped. + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + return Poll::Ready(None); + } + } + } + } +} + +impl RecordBatchStream for PerPartitionFilterStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } +} + +/// Build the `K = cuts.len() + 1` half-open range predicates a +/// `PerPartitionFilterExec` needs to reproduce the range repartition's +/// write-side routing on the read side. +/// +/// Partition `i` receives the predicate +/// +/// ```text +/// i = 0 → routing_expr < cuts[0] +/// 0 < i < K-1 → cuts[i-1] <= routing_expr AND routing_expr < cuts[i] +/// i = K-1 → routing_expr >= cuts[K-2] +/// K = 1 → lit(true) // empty cuts, single-bucket range repartition +/// ``` +/// +/// Consistent with the private `range_repartition_common::split_batch_by_range` +/// helper, which uses the same half-open convention on the write side. +/// Callers pass the range repartition's routing expression verbatim +/// (`CAST(order_by[0] AS Float64)` today). +/// +/// Non-null routing expressions only. Both `UnorderedRangeRepartitionExec` +/// and `OrderedRangeRepartitionExec` refuse nullable routing exprs at +/// `try_new`, so any expression that reaches this helper via +/// `RangeRepartitionRouting` is guaranteed non-null — no `IS NULL` branch +/// needed. +pub fn range_partition_predicates( + routing_expr: Arc, + cuts: &[f64], +) -> Vec> { + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Literal}; + use datafusion::scalar::ScalarValue; + + let partition_count = cuts.len() + 1; + let lit = |v: f64| -> Arc { + Arc::new(Literal::new(ScalarValue::Float64(Some(v)))) + }; + let ge = |lo: f64| -> Arc { + Arc::new(BinaryExpr::new( + routing_expr.clone(), + Operator::GtEq, + lit(lo), + )) + }; + let lt = |hi: f64| -> Arc { + Arc::new(BinaryExpr::new(routing_expr.clone(), Operator::Lt, lit(hi))) + }; + (0..partition_count) + .map(|partition_idx| { + let lo = partition_idx + .checked_sub(1) + .and_then(|cut_idx| cuts.get(cut_idx).copied()); + let hi = cuts.get(partition_idx).copied(); + match (lo, hi) { + (None, None) => { + // K == 1: single bucket covers everything. + Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))) + as Arc + } + (None, Some(hi)) => lt(hi), + (Some(lo), None) => ge(lo), + (Some(lo), Some(hi)) => { + Arc::new(BinaryExpr::new(ge(lo), Operator::And, lt(hi))) + } + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::array::Int64Array; + use datafusion::arrow::datatypes::{Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::physical_plan::repartition::RepartitionExec; + use datafusion::physical_plan::{ExecutionPlan, Partitioning}; + use datafusion::prelude::SessionContext; + use datafusion::scalar::ScalarValue; + use futures::TryStreamExt; + + fn one_col_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])) + } + + /// Memory source with `partitions` partitions, each carrying a single + /// batch of `[start .. start + rows_per)` where `start = k * rows_per`. + fn partitioned_source(partitions: usize, rows_per: usize) -> Arc { + let schema = one_col_schema(); + let mut per_partition: Vec> = Vec::with_capacity(partitions); + for k in 0..partitions { + let start = (k * rows_per) as i64; + let arr = Int64Array::from_iter_values(start..start + rows_per as i64); + let batch = + RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap(); + per_partition.push(vec![batch]); + } + let src = + MemorySourceConfig::try_new(&per_partition, schema, None).expect("mem src"); + Arc::new(DataSourceExec::new(Arc::new(src))) + } + + /// Predicate `v >= lo AND v < hi` against column `v`. + fn range_pred(lo: i64, hi: i64) -> Arc { + let col = Arc::new(Column::new("v", 0)); + let lo_lit = Arc::new(Literal::new(ScalarValue::Int64(Some(lo)))); + let hi_lit = Arc::new(Literal::new(ScalarValue::Int64(Some(hi)))); + let ge: Arc = + Arc::new(BinaryExpr::new(col.clone(), Operator::GtEq, lo_lit)); + let lt: Arc = + Arc::new(BinaryExpr::new(col, Operator::Lt, hi_lit)); + Arc::new(BinaryExpr::new(ge, Operator::And, lt)) + } + + fn ctx() -> Arc { + SessionContext::new().task_ctx() + } + + async fn collect(plan: Arc, partition: usize) -> Result> { + let stream = plan.execute(partition, ctx())?; + let batches: Vec = stream.try_collect().await?; + let mut out = Vec::new(); + for b in batches { + let arr = b + .column(0) + .as_any() + .downcast_ref::() + .expect("Int64Array"); + out.extend(arr.iter().map(|v| v.unwrap())); + } + Ok(out) + } + + /// Each of three partitions carries `[k*100, k*100+100)`. With + /// per-partition predicates that each carve a five-row slice, every + /// partition emits its own five rows and nothing from another + /// partition leaks through. + #[tokio::test] + async fn per_partition_predicate_filters_only_that_partition() -> Result<()> { + let src = partitioned_source(3, 100); + let predicates = vec![ + range_pred(0, 5), // partition 0 → 0..5 + range_pred(105, 110), // partition 1 → 105..110 + range_pred(295, 300), // partition 2 → 295..300 + ]; + let ppf: Arc = + Arc::new(PerPartitionFilterExec::try_new(src, predicates)?); + assert_eq!(collect(ppf.clone(), 0).await?, (0..5).collect::>()); + assert_eq!( + collect(ppf.clone(), 1).await?, + (105..110).collect::>() + ); + assert_eq!(collect(ppf, 2).await?, (295..300).collect::>()); + Ok(()) + } + + /// A predicate that matches nothing yields an empty stream (no zero-row + /// batches surfaced to the caller). Regression pin — an earlier draft + /// forwarded empty batches, which some downstream operators dislike. + #[tokio::test] + async fn empty_predicate_yields_empty_stream() -> Result<()> { + let src = partitioned_source(1, 100); + let predicates = vec![range_pred(1_000_000, 2_000_000)]; + let ppf: Arc = + Arc::new(PerPartitionFilterExec::try_new(src, predicates)?); + assert_eq!(collect(ppf, 0).await?, Vec::::new()); + Ok(()) + } + + /// Predicate-count mismatch is rejected at construction time. + #[test] + fn rejects_predicate_count_mismatch() { + let src = partitioned_source(3, 10); + let err = + PerPartitionFilterExec::try_new(src, vec![range_pred(0, 5)]).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("predicate count 1") && msg.contains("input partition count 3"), + "unexpected error: {msg}" + ); + } + + /// Non-boolean predicate is rejected at construction time (the + /// expression must evaluate to `Boolean` against the input schema). + #[test] + fn rejects_non_boolean_predicate() { + let src = partitioned_source(1, 10); + // Just the column `v` — evaluates to Int64, not Boolean. + let bad: Arc = Arc::new(Column::new("v", 0)); + let err = PerPartitionFilterExec::try_new(src, vec![bad]).unwrap_err(); + assert!( + err.to_string().contains("Boolean"), + "unexpected error: {err}" + ); + } + + /// The K=4 range predicates cover every value under the half-open + /// convention, and each row lands in exactly one predicate. Random + /// probe values are routed through the predicates and expected to + /// match the same partition assignment as the range repartition's + /// write-side `split_batch_by_range` would produce. + #[test] + fn range_partition_predicates_partition_every_value_exactly_once() { + use datafusion::arrow::array::Float64Array; + use datafusion::arrow::datatypes::Field; + use datafusion::physical_expr::expressions::Column; + + let cuts = vec![10.0, 20.0, 30.0]; + let k = cuts.len() + 1; + let routing: Arc = Arc::new(Column::new("v", 0)); + let preds = range_partition_predicates(routing, &cuts); + assert_eq!(preds.len(), k); + + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let values: Vec = vec![ + -5.0, 0.0, 9.999, 10.0, 15.0, 19.999, 20.0, 25.0, 30.0, 100.0, + ]; + let arr = Float64Array::from_iter_values(values.iter().copied()); + let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap(); + + // For each row, find the unique partition whose predicate accepts it. + for (row, &v) in values.iter().enumerate() { + let mut hits = 0; + for pred in &preds { + let mask = pred + .evaluate(&batch) + .and_then(|v| v.into_array(batch.num_rows())) + .unwrap(); + let mask = as_boolean_array(&mask).unwrap(); + if mask.value(row) { + hits += 1; + } + } + assert_eq!( + hits, 1, + "value {v} matched {hits} predicates, expected exactly 1" + ); + } + + // Expected assignment mirrors split_batch_by_range: `partition_point` + // returns the count of cuts `<= key`, which is the partition index + // under the half-open convention. + let expected: Vec = values + .iter() + .map(|v| cuts.partition_point(|&c| c <= *v)) + .collect(); + for (row, want) in expected.iter().enumerate() { + let mask = preds[*want] + .evaluate(&batch) + .and_then(|v| v.into_array(batch.num_rows())) + .unwrap(); + let mask = as_boolean_array(&mask).unwrap(); + assert!( + mask.value(row), + "value {} should have landed in partition {}", + values[row], + want + ); + } + } + + /// Degenerate K=1 (empty cuts) yields a single lit(true) predicate. + #[test] + fn range_partition_predicates_single_bucket_when_cuts_empty() { + use datafusion::physical_expr::expressions::Column; + + let routing: Arc = Arc::new(Column::new("v", 0)); + let preds = range_partition_predicates(routing, &[]); + assert_eq!(preds.len(), 1); + assert_eq!(preds[0].to_string(), "true"); + } + + /// `with_new_children` swaps the input while preserving the predicate + /// vector. Wrapping the original source in a `RepartitionExec` that + /// keeps the partition count (RoundRobin(3)) gives a valid child; the + /// filter still routes partition-`k` rows through `predicates[k]`. + #[tokio::test] + async fn with_new_children_preserves_predicates() -> Result<()> { + let src = partitioned_source(3, 100); + let predicates = + vec![range_pred(0, 3), range_pred(100, 103), range_pred(200, 203)]; + let ppf = Arc::new(PerPartitionFilterExec::try_new( + src.clone(), + predicates.clone(), + )?); + // Wrap the source in RoundRobin(3) — same partition count, different plan. + let repart: Arc = Arc::new(RepartitionExec::try_new( + src, + Partitioning::RoundRobinBatch(3), + )?); + let swapped: Arc = ppf.with_new_children(vec![repart])?; + // Just verify construction succeeded and the operator name survives. + assert_eq!(swapped.name(), "PerPartitionFilterExec"); + Ok(()) + } +} diff --git a/ballista/core/src/execution_plans/plan_algebra.rs b/ballista/core/src/execution_plans/plan_algebra.rs new file mode 100644 index 0000000000..245a8c3771 --- /dev/null +++ b/ballista/core/src/execution_plans/plan_algebra.rs @@ -0,0 +1,70 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Algebraic properties of physical plan nodes — do they preserve +//! partitioning, do they preserve the distribution of row values, etc. +//! +//! DataFusion doesn't expose these as trait methods on `ExecutionPlan` +//! (nice-to-haves like `ExecutionPlan::affects_partitioning()` / +//! `ExecutionPlan::affects_distribution()` that would hopefully land one +//! day), so we downcast against a hand-maintained whitelist. Being +//! conservative is the safety net: unrecognized node → property assumed +//! false → caller falls back to the safer path. + +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; + +use crate::execution_plans::{ + BufferExec, RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec, +}; + +/// Whitelisted ops preserve the routing key's row set, values, and +/// partitioning — an upstream sketch remains valid after the operator. +pub fn preserves_distribution(plan: &dyn ExecutionPlan) -> bool { + // Buffered batches replayed verbatim. + plan.downcast_ref::().is_some() + // Per-partition sort: rows reorder within a partition, row set + // and counts unchanged. `preserve_partitioning=false` collapses + // N→1 (like SortPreservingMergeExec), so gate on the flag. + || plan + .downcast_ref::() + .is_some_and(|sort| sort.preserve_partitioning()) + // Stage-boundary writers: batches to disk unchanged. + || plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + // Pure row-annotation: one input row → one output row with an + // added column (window fn result); values, partitioning, count preserved. + || plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() +} + +/// Looser sibling of [`preserves_distribution`]: partitioning survives, +/// but rows and values within a partition are fair game. +pub fn preserves_partitioning(plan: &dyn ExecutionPlan) -> bool { + // Distribution-preserving is strictly stronger; compose to keep the + // whitelist deduplicated. + preserves_distribution(plan) + // Drops rows, but per-partition — no rows migrate. + || plan.downcast_ref::().is_some() + // Rewrites columns; partition boundaries untouched. + || plan.downcast_ref::().is_some() + // Stats tap; no data mutation. + || plan.downcast_ref::().is_some() +} diff --git a/ballista/core/src/execution_plans/range_repartition_common.rs b/ballista/core/src/execution_plans/range_repartition_common.rs index 30bdcb7236..b8cc948053 100644 --- a/ballista/core/src/execution_plans/range_repartition_common.rs +++ b/ballista/core/src/execution_plans/range_repartition_common.rs @@ -46,15 +46,12 @@ use datafusion::arrow::compute::take_arrays; use datafusion::common::{Result, internal_datafusion_err}; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_plan::sorts::sort::SortExec; -use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use futures::FutureExt; use log::warn; use tokio::sync::mpsc; -use crate::execution_plans::{ - BufferExec, RuntimeStatsExec, ShuffleWriterExec, SortShuffleWriterExec, -}; +use crate::execution_plans::RuntimeStatsExec; +use crate::execution_plans::plan_algebra::preserves_distribution; /// Walk `child`'s subtree for a [`RuntimeStatsExec`] that sketches on our /// routing expression, snapshot its merged T-Digest, and compute `K - 1` @@ -155,54 +152,6 @@ pub(super) fn find_runtime_stats<'a>( find_runtime_stats(only_child, routing_expr) } -/// Whitelist of pass-through operator types the walker will descend through -/// on its way to a matching [`RuntimeStatsExec`]. Unlisted operators might -/// drop rows, duplicate rows, or transform the routing key's value — any of -/// which would make an upstream sketch stale by the time data reaches us. -/// -/// Being conservative is the safety net: unrecognized node → walker gives -/// up → single-bucket fallback. Extending this list requires positive -/// verification that the operator is a distribution-preserving passthrough -/// for the routing key. Absent an upstream `ExecutionPlan::affects_distribution()` -/// method (nice-to-have that hopefully lands one day), we maintain this by hand. -pub(super) fn preserves_distribution(plan: &dyn ExecutionPlan) -> bool { - // Every entry here is a *claim* that the operator (1) doesn't drop - // rows, (2) doesn't duplicate rows, (3) doesn't transform the routing - // key's value, AND (4) doesn't change partitioning (per-partition - // slots downstream still map to the same partitions upstream). Losing - // any of those invalidates the sketch/count on the other side. - // - // Notable *exclusions*: - // • `SortPreservingMergeExec` — collapses N partitions to 1, so - // per-partition slots below it don't align with the single - // partition above. Values are preserved, but the partitioning - // invariant fails. - // • `ProjectionExec` — might compute a new column that shadows or - // replaces the routing key, transforming values invisibly. - // • `FilterExec`, `LimitExec`, joins — drop or duplicate rows. - // • Our own DRRs — repartition by value; that's the whole point. - plan.downcast_ref::().is_some() - // `SortExec` reorders rows within each partition; row set and per- - // partition counts unchanged — but ONLY when `preserve_partitioning` - // is true. The `preserve_partitioning=false` variant collapses N - // partitions to 1 (like `SortPreservingMergeExec`), which would - // invalidate per-partition slot alignment. - || plan - .downcast_ref::() - .is_some_and(|sort| sort.preserve_partitioning()) - // `ShuffleWriterExec` / `SortShuffleWriterExec` sit at the top of - // every stage's plan on the executor side. They write batches to - // disk unchanged — no transformation, no filtering. - || plan.downcast_ref::().is_some() - || plan.downcast_ref::().is_some() - // `BoundedWindowAggExec` / `WindowAggExec` are pure row-annotation: - // each input row emits exactly one output row with a new column - // (the window function's result). The routing key's values, - // partitioning, and row count are all preserved verbatim. - || plan.downcast_ref::().is_some() - || plan.downcast_ref::().is_some() -} - /// Split `batch` into `K = boundaries.len() + 1` sub-batches under the /// half-open convention: partition `p` receives rows where /// `boundaries[p-1] <= key < boundaries[p]` (open at `-∞` on partition 0 diff --git a/ballista/core/src/execution_plans/runtime_stats.rs b/ballista/core/src/execution_plans/runtime_stats.rs index 43760b55e8..949241a163 100644 --- a/ballista/core/src/execution_plans/runtime_stats.rs +++ b/ballista/core/src/execution_plans/runtime_stats.rs @@ -57,7 +57,9 @@ use datafusion::arrow::array::Float64Array; use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::common::{Result, Statistics, internal_datafusion_err, internal_err}; use datafusion::execution::TaskContext; -use datafusion::physical_expr::{Distribution, OrderingRequirements, PhysicalSortExpr}; +use datafusion::physical_expr::{ + Distribution, OrderingRequirements, PhysicalExpr, PhysicalSortExpr, +}; use datafusion::physical_plan::execution_plan::CardinalityEffect; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ @@ -68,6 +70,11 @@ use datafusion_functions_aggregate_common::tdigest::TDigest; use futures::stream::StreamExt; use log::debug; +use crate::serde::protobuf::{ + QuantileSketchState, RuntimeStatsPartitionEntry, RuntimeStatsReport, +}; +use crate::serde::scheduler::PartitionLocation; + /// T-Digest centroid budget. 100 is DataFusion's default and gives ~1% /// quantile error, plenty of margin over the sub-partition counts we /// expect at bin-pack time. @@ -506,15 +513,99 @@ impl Drop for StreamState { } } +/// Walk `plan` and collect one [`RuntimeStatsReport`] +/// per [`RuntimeStatsExec`] that remains valid at the plan's output. +/// "Valid" means reachable through single-child chains of distribution- +/// preserving operators only — see the `preserves_distribution` +/// whitelist in the range-repartition module. A stats-tap sitting +/// *below* an [`super::UnorderedRangeRepartitionExec`] (or any other +/// distribution-changing operator) is excluded automatically because +/// the walker stops at that boundary; its sketch describes data the +/// repartitioner then routed away and is no longer meaningful at the +/// plan's output. +/// +/// Executors call this once per task at completion to package what to +/// return to the scheduler. +pub fn collect_reports(plan: &Arc) -> Result> { + use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, + }; + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DefaultPhysicalProtoConverter {}; + let mut found: Vec<&RuntimeStatsExec> = Vec::new(); + collect_reachable_stats(plan, &mut found); + found + .into_iter() + .map(|stats| stats_to_report(stats, &codec, &converter)) + .collect() +} + +/// DFS `plan` through single-child chains only, descending through +/// distribution-preserving nodes and past any [`RuntimeStatsExec`] found +/// on the way. Stops at any branch, leaf, or non-whitelisted node. +/// Similar in shape to `range_repartition_common::find_runtime_stats` +/// but collects *all* reachable stats rather than returning the first +/// match keyed to a specific routing expression. +fn collect_reachable_stats<'a>( + plan: &'a Arc, + out: &mut Vec<&'a RuntimeStatsExec>, +) { + if let Some(stats) = plan.downcast_ref::() { + out.push(stats); + // Continue descending — a plan could conceivably chain multiple + // stats-taps; `preserves_distribution` still guards the recursion. + } else if !super::plan_algebra::preserves_distribution(plan.as_ref()) { + return; + } + let children = plan.children(); + let [only_child] = children.as_slice() else { + return; + }; + collect_reachable_stats(only_child, out); +} + +fn stats_to_report( + stats: &RuntimeStatsExec, + codec: &dyn datafusion_proto::physical_plan::PhysicalExtensionCodec, + converter: &datafusion_proto::physical_plan::DefaultPhysicalProtoConverter, +) -> Result { + use datafusion_proto::physical_plan::to_proto::serialize_physical_sort_exprs; + let order_by = match stats.order_by() { + Some(order_by) => { + serialize_physical_sort_exprs(order_by.iter().cloned(), codec, converter)? + } + None => Vec::new(), + }; + // Iterate every partition slot the operator holds. Slots the task + // didn't touch have row_count = 0 and an empty sketch; we still emit + // them so the scheduler sees a shape-consistent view. + let partition_count = stats.partition_count(); + let mut partitions = Vec::with_capacity(partition_count); + for partition_id in 0..partition_count { + let row_count = stats.row_count(partition_id)? as u64; + let sketch = match stats.quantile_sketch(partition_id)? { + Some(sk) if sk.count() > 0.0 => Some(sketch_to_proto(&sk)?), + _ => None, + }; + partitions.push(RuntimeStatsPartitionEntry { + partition_id: partition_id as u32, + row_count, + sketch, + }); + } + Ok(RuntimeStatsReport { + order_by, + partitions, + }) +} + /// Serialize a T-Digest to the on-wire -/// [`crate::serde::protobuf::QuantileSketchState`]. +/// [`QuantileSketchState`]. /// /// Wraps `TDigest::to_scalar_state()` — the 6-element canonical form /// `(max_size, sum, count, max, min, centroids_as_list)` — each element /// encoded via `datafusion_proto_common::ScalarValue::try_from`. -pub fn sketch_to_proto( - sketch: &TDigest, -) -> Result { +pub fn sketch_to_proto(sketch: &TDigest) -> Result { let state = sketch.to_scalar_state(); let proto_state = state .iter() @@ -523,18 +614,16 @@ pub fn sketch_to_proto( .map_err(|e| { internal_datafusion_err!("failed to encode TDigest to proto: {e:?}") })?; - Ok(crate::serde::protobuf::QuantileSketchState { state: proto_state }) + Ok(QuantileSketchState { state: proto_state }) } -/// Deserialize a [`crate::serde::protobuf::QuantileSketchState`] into a +/// Deserialize a [`QuantileSketchState`] into a /// T-Digest. /// /// Reverses [`sketch_to_proto`]. Guards against corrupted wire input by /// checking the element count before calling /// `TDigest::from_scalar_state`, which would panic on invalid shape. -pub fn sketch_from_proto( - proto: &crate::serde::protobuf::QuantileSketchState, -) -> Result { +pub fn sketch_from_proto(proto: &QuantileSketchState) -> Result { let scalars = proto .state .iter() @@ -555,6 +644,322 @@ pub fn sketch_from_proto( Ok(TDigest::from_scalar_state(&scalars)) } +/// One group's merged view: sketches from every report sharing the same +/// `order_by` wire tag combined, plus the `partition_count - 1` quantile +/// cuts a globally-informed router would use for `partition_count` +/// output partitions. +#[derive(Debug, Clone)] +pub struct MergedRuntimeStats { + /// How many `PhysicalSortExprNode`s were in the shared `order_by` tag. + pub order_by_len: usize, + /// Number of output partitions the router used (per-report + /// `partitions.len()`, must agree across reports in the group). + pub partition_count: usize, + /// Number of `RuntimeStatsReport`s contributing to this group. + pub task_count: usize, + /// Sum of `row_count` across every partition entry in the group. + pub total_rows: u64, + /// `partition_count - 1` cut points at quantiles `i/partition_count` + /// on the merged T-Digest. Empty when `partition_count < 2` or no + /// non-empty sketches were merged. + pub cuts: Vec, + /// Merged T-Digest's `min()` if at least one non-empty sketch + /// contributed; `None` in row-count-only mode. + pub min: Option, + /// Merged T-Digest's `max()` if at least one non-empty sketch + /// contributed; `None` in row-count-only mode. + pub max: Option, +} + +/// Group `RuntimeStatsReport`s by `order_by` wire tag, merge the T-Digests +/// within each group, and return one [`MergedRuntimeStats`] per group. +pub fn merge_reports(reports: &[RuntimeStatsReport]) -> Result> { + use prost::Message; + use std::collections::HashMap; + + if reports.is_empty() { + return Ok(Vec::new()); + } + + // Group by the bytes of the encoded `order_by`. Prost-encoding each + // `PhysicalSortExprNode` and concatenating gives a stable, cheap + // grouping key without needing `Hash` on the generated proto types. + let mut groups: HashMap, Vec<&RuntimeStatsReport>> = HashMap::new(); + for report in reports { + let mut group_key = Vec::new(); + for expr in &report.order_by { + expr.encode(&mut group_key) + .expect("Vec is an infallible sink for prost::Message::encode"); + } + groups.entry(group_key).or_default().push(report); + } + + let mut merged_groups = Vec::with_capacity(groups.len()); + for group in groups.into_values() { + merged_groups.push(merge_group(&group)?); + } + Ok(merged_groups) +} + +/// Merge one group of reports (all sharing the same `order_by` tag). +/// Kept separate from `merge_reports` so the group iteration reads as a +/// single fallible step per group. +fn merge_group(group: &[&RuntimeStatsReport]) -> Result { + let [first, rest @ ..] = group else { + // `merge_reports` only builds groups from `HashMap::entry().push()`, + // so an empty group is unreachable. Surface as internal error + // rather than panicking. + return internal_err!( + "runtime stats merge: empty group — merge_reports invariant broken" + ); + }; + let partition_count = first.partitions.len(); + let task_count = group.len(); + + // Every task ran the same stage plan, so partition counts must + // agree. Mismatch = internal invariant break. + for report in rest { + if report.partitions.len() != partition_count { + return internal_err!( + "runtime stats merge: order_by_len={} mismatched partition \ + counts across reports ({} vs {})", + first.order_by.len(), + partition_count, + report.partitions.len() + ); + } + } + + let mut total_rows: u64 = 0; + let mut sketches: Vec = Vec::new(); + for report in group { + for entry in &report.partitions { + total_rows = total_rows.saturating_add(entry.row_count); + if let Some(proto_sketch) = entry.sketch.as_ref() { + let sketch = sketch_from_proto(proto_sketch)?; + if sketch.count() > 0.0 { + sketches.push(sketch); + } + } + } + } + + if sketches.is_empty() { + return Ok(MergedRuntimeStats { + order_by_len: first.order_by.len(), + partition_count, + task_count, + total_rows, + cuts: Vec::new(), + min: None, + max: None, + }); + } + + let merged_sketch = TDigest::merge_digests(sketches.iter()); + let cuts: Vec = if partition_count > 1 { + (1..partition_count) + .map(|cut_index| { + merged_sketch.estimate_quantile(cut_index as f64 / partition_count as f64) + }) + .collect() + } else { + Vec::new() + }; + Ok(MergedRuntimeStats { + order_by_len: first.order_by.len(), + partition_count, + task_count, + total_rows, + cuts, + min: Some(merged_sketch.min()), + max: Some(merged_sketch.max()), + }) +} + +/// One producer task's runtime-stats report, kept alongside the +/// `producer_task_id` that emitted it. The scheduler stores these on +/// `RunningStage.runtime_stats_reports` so downstream stages can address +/// individual producer files as `(producer_task_id, partition_id)` pairs — +/// the partition_id inside a report is producer-local (0..K range-repartition +/// sub-parts), so the pair is what uniquely identifies a shuffle file across +/// producers. +#[derive(Debug, Clone)] +pub struct TaskRuntimeStats { + /// Producer task's task_id at the time it emitted the report. Matches + /// the `file_id` stamped on `ShuffleWritePartition` records. + pub producer_task_id: usize, + /// The report itself: per-partition row counts and (in sketch mode) + /// quantile sketches for the routing expression. + pub report: RuntimeStatsReport, +} + +/// Walk `plan` for the first `UnorderedRangeRepartitionExec` or +/// `OrderedRangeRepartitionExec` and return its routing expression +/// (`order_by[0].expr`). `Ok(None)` means no range-repartition operator +/// in the plan; `Err(_)` means one was found but its `order_by` was +/// empty (invariant break — a range repartition without a routing key +/// can't route anything). +pub fn repartition_routing_expr( + plan: &dyn ExecutionPlan, +) -> Result>> { + if let Some(rre) = plan.downcast_ref::() { + return match rre.order_by() { + [first, ..] => Ok(Some(first.expr.clone())), + [] => internal_err!("UnorderedRangeRepartitionExec has empty ORDER BY"), + }; + } + if let Some(rre) = plan.downcast_ref::() { + return match rre.order_by() { + [first, ..] => Ok(Some(first.expr.clone())), + [] => internal_err!("OrderedRangeRepartitionExec has empty ORDER BY"), + }; + } + for child in plan.children() { + if let Some(expr) = repartition_routing_expr(child.as_ref())? { + return Ok(Some(expr)); + } + } + Ok(None) +} + +/// Rebuild a stage's `Vec>` under range-repartition +/// overlap semantics: for each producer file in `original_partitions`, +/// find its sketch (from `reports`), and route the file into every +/// downstream partition whose global cut range overlaps +/// `[sketch.min(), sketch.max()]`. +/// +/// Downstream partition ranges follow the half-open convention: +/// - `k = 0` → `(-∞, cuts[0])` +/// - `0 < k < K - 1` → `[cuts[k-1], cuts[k])` +/// - `k = K - 1` → `[cuts[K-2], +∞)` +/// +/// `[min, max]` overlaps `[lower, upper)` iff `max >= lower AND min < upper`. +/// +/// Files without a corresponding sketch (missing entirely, or present +/// with `count == 0`) are safe to skip only when `partition_stats.num_rows` +/// confirms the file is empty (`Some(0)`). If the file has rows or the +/// row count is unknown (`None`), silently skipping would lose data — +/// error out instead. +pub fn cut_partitions( + original_partitions: Vec>, + reports: &[TaskRuntimeStats], + global_cuts: &[f64], +) -> Result>> { + use std::collections::HashMap; + + // Index sketches by (producer_task_id, sub_part_id). Under + // ShuffleWriter(Passthrough) file_id == task_id, so PartitionLocation's + // (file_id, partition_id.partition_id) is the same pair. + let sketches: HashMap<(usize, u32), &QuantileSketchState> = reports + .iter() + .flat_map(|stats| { + stats.report.partitions.iter().filter_map(move |entry| { + let sketch = entry.sketch.as_ref()?; + Some(((stats.producer_task_id, entry.partition_id), sketch)) + }) + }) + .collect(); + + debug_assert!( + global_cuts.windows(2).all(|w| w[0] <= w[1]), + "global_cuts must be non-decreasing: {global_cuts:?}" + ); + + let partition_count = global_cuts.len() + 1; + let mut remapped: Vec> = vec![Vec::new(); partition_count]; + for partition in original_partitions { + for file in partition { + let Some(task_id) = file.file_id else { + return internal_err!( + "range-repartition remap: missing file_id (partition_id={})", + file.partition_id.partition_id + ); + }; + let sub_part_id = file.partition_id.partition_id as u32; + // Fold "no sketch" and "empty sketch" into one Option — both + // mean "no routing info for this file" + let sketch = sketches + .get(&(task_id as usize, sub_part_id)) + .map(|proto| sketch_from_proto(proto)) + .transpose()? + .filter(|s| s.count() > 0.0); + let Some(sketch) = sketch else { + // No routing info. Safe to skip only if the file has zero rows + if file.partition_stats.num_rows != Some(0) { + return internal_err!( + "range-repartition remap: file has num_rows={:?} but no usable sketch (task_id={task_id}, sub_part_id={sub_part_id})", + file.partition_stats.num_rows + ); + } + continue; + }; + // Bucket i has (lower, upper) = (cuts[i-1], cuts[i]) with ±∞ at the + // ends, and matches iff `sketch_max >= lower && sketch_min < upper`. + // Monotone cuts → the set of matching buckets is a contiguous range + // [b_lo, b_hi], found by two partition_points over `global_cuts`. + let (sketch_min, sketch_max) = (sketch.min(), sketch.max()); + let b_lo = global_cuts.partition_point(|&c| c <= sketch_min); + let b_hi = global_cuts.partition_point(|&c| c <= sketch_max); + for bucket in &mut remapped[b_lo..=b_hi] { + bucket.push(file.clone()); + } + } + } + Ok(remapped) +} + +/// Merge `reports` and log each group's merged view at `debug!` +/// (`RUST_LOG` promotes when needed). Any merge error is logged at +/// `warn!` — the scheduler doesn't want telemetry loss to tank a query +/// whose data was already produced correctly. The scheduler calls this +/// once per stage-attempt at final-success. +pub fn log_merged_runtime_stats( + job_id: &str, + stage_id: usize, + reports: &[TaskRuntimeStats], +) { + let raw: Vec = reports.iter().map(|t| t.report.clone()).collect(); + let merged_groups = match merge_reports(&raw) { + Ok(groups) => groups, + Err(err) => { + log::warn!( + "runtime stats merge failed for job={job_id} stage={stage_id}: {err}" + ); + return; + } + }; + for merged in merged_groups { + match (merged.min, merged.max) { + (Some(min), Some(max)) => log::debug!( + "merged runtime stats: job={} stage={} order_by_len={} \ + partition_count={} task_count={} total_rows={} cuts={:?} \ + min={} max={}", + job_id, + stage_id, + merged.order_by_len, + merged.partition_count, + merged.task_count, + merged.total_rows, + merged.cuts, + min, + max, + ), + _ => log::debug!( + "merged runtime stats: job={} stage={} order_by_len={} \ + partition_count={} task_count={} total_rows={} cuts=[] \ + (no sketches)", + job_id, + stage_id, + merged.order_by_len, + merged.partition_count, + merged.task_count, + merged.total_rows, + ), + } + } +} + #[cfg(test)] mod wire_tests { use super::*; @@ -597,7 +1002,7 @@ mod wire_tests { #[test] fn sketch_from_proto_rejects_wrong_shape() { use datafusion::common::ScalarValue; - let proto = crate::serde::protobuf::QuantileSketchState { + let proto = QuantileSketchState { state: (0..3) .map(|_| { datafusion_proto_common::ScalarValue::try_from(&ScalarValue::Float64( @@ -732,3 +1137,742 @@ mod stream_tests { assert!(stats.merged_quantile_sketch().unwrap().is_none()); } } + +#[cfg(test)] +mod collect_tests { + //! Walker behavior: which `RuntimeStatsExec`s does `collect_reports` + //! see through the whitelist, and what do the emitted reports look + //! like once the plan has been drained? + + use super::*; + use crate::execution_plans::BufferExec; + use crate::execution_plans::buffer::BufferMode; + use datafusion::arrow::array::Float64Array; + use datafusion::arrow::compute::SortOptions; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::common; + use datafusion::physical_plan::expressions::col; + use datafusion::physical_plan::sorts::sort::SortExec; + use datafusion::prelude::SessionContext; + + fn schema_v() -> Arc { + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])) + } + + fn v_batch(schema: &Arc, v: Vec) -> RecordBatch { + RecordBatch::try_new(schema.clone(), vec![Arc::new(Float64Array::from(v))]) + .unwrap() + } + + fn v_input(schema: Arc) -> Arc { + let b1 = v_batch(&schema, vec![1.0, 3.0, 5.0]); + let b2 = v_batch(&schema, vec![2.0, 4.0]); + let memory = + Arc::new(MemorySourceConfig::try_new(&[vec![b1, b2]], schema, None).unwrap()); + Arc::new(DataSourceExec::new(memory)) + } + + fn sort_expr_on_v(schema: &Arc) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: col("v", schema.as_ref()).unwrap(), + options: SortOptions { + descending: false, + nulls_first: false, + }, + } + } + + /// Stats sit at plan root, sketching mode: `collect_reports` returns + /// exactly one report whose partition entry carries the observed + /// row_count and a populated sketch that survives the on-wire round- + /// trip via `sketch_from_proto`. + #[tokio::test] + async fn collect_reports_finds_stats_and_ships_sketch() { + let schema = schema_v(); + let input = v_input(schema.clone()); + let stats = Arc::new( + RuntimeStatsExec::try_new(input, Some(vec![sort_expr_on_v(&schema)])) + .unwrap(), + ); + + // Drive the stream so counters and the sketch actually fill. + let ctx = SessionContext::new().task_ctx(); + let stream = stats.clone().execute(0, ctx).unwrap(); + let _ = common::collect(stream).await.unwrap(); + + let plan: Arc = stats; + let reports = collect_reports(&plan).expect("collect_reports must succeed"); + let [report] = reports.as_slice() else { + panic!( + "expected exactly one report, got {} (order_by tags: {:?})", + reports.len(), + reports.iter().map(|r| r.order_by.len()).collect::>() + ); + }; + assert_eq!(report.order_by.len(), 1, "one sort expr encoded"); + let [entry] = report.partitions.as_slice() else { + panic!( + "expected one partition entry, got {}", + report.partitions.len() + ); + }; + assert_eq!(entry.partition_id, 0); + assert_eq!(entry.row_count, 5); + let proto_sketch = entry.sketch.as_ref().expect("sketch present in wire"); + let round_tripped = sketch_from_proto(proto_sketch).unwrap(); + assert_eq!(round_tripped.count(), 5.0); + assert_eq!(round_tripped.min(), 1.0); + assert_eq!(round_tripped.max(), 5.0); + } + + /// Row-count-only mode: report emitted, but its partition entry + /// carries no sketch. + #[tokio::test] + async fn collect_reports_row_count_only_emits_report_without_sketch() { + let schema = schema_v(); + let input = v_input(schema.clone()); + let stats = Arc::new(RuntimeStatsExec::try_new(input, None).unwrap()); + + let ctx = SessionContext::new().task_ctx(); + let stream = stats.clone().execute(0, ctx).unwrap(); + let _ = common::collect(stream).await.unwrap(); + + let plan: Arc = stats; + let reports = collect_reports(&plan).unwrap(); + let [report] = reports.as_slice() else { + panic!("expected one report, got {}", reports.len()); + }; + assert!(report.order_by.is_empty()); + assert_eq!(report.partitions.len(), 1); + assert!( + report.partitions[0].sketch.is_none(), + "no sketch in row-count-only mode" + ); + assert_eq!(report.partitions[0].row_count, 5); + } + + /// Whitelisted intermediary (`BufferExec` in Dam mode) between plan + /// root and the stats-tap: walker still descends to it. + #[tokio::test] + async fn collect_reports_descends_through_whitelisted_op() { + let schema = schema_v(); + let input = v_input(schema.clone()); + let stats = Arc::new( + RuntimeStatsExec::try_new(input, Some(vec![sort_expr_on_v(&schema)])) + .unwrap(), + ); + let buffer: Arc = + Arc::new(BufferExec::try_new(stats, BufferMode::Dam).unwrap()); + + // Drain via the outer plan so counters fill. + let ctx = SessionContext::new().task_ctx(); + let stream = buffer.clone().execute(0, ctx).unwrap(); + let _ = common::collect(stream).await.unwrap(); + + let reports = collect_reports(&buffer).unwrap(); + assert_eq!(reports.len(), 1, "buffer must not block the walker"); + assert_eq!(reports[0].partitions[0].row_count, 5); + } + + /// A `SortExec` with `preserve_partitioning=false` collapses N→1; + /// the whitelist excludes that variant explicitly. The walker + /// stops at the collapse and doesn't reach the stats below. + #[tokio::test] + async fn collect_reports_stops_at_sort_that_collapses_partitions() { + let schema = schema_v(); + let input = v_input(schema.clone()); + let stats: Arc = Arc::new( + RuntimeStatsExec::try_new(input, Some(vec![sort_expr_on_v(&schema)])) + .unwrap(), + ); + // Default SortExec has preserve_partitioning=false — the + // whitelist path we're testing rejects it. + let sort = SortExec::new( + LexOrdering::new(vec![sort_expr_on_v(&schema)]).unwrap(), + stats, + ); + assert!( + !sort.preserve_partitioning(), + "test fixture assumes N→1 sort" + ); + let plan: Arc = Arc::new(sort); + let reports = collect_reports(&plan).unwrap(); + assert!( + reports.is_empty(), + "N→1 sort must block the walker; got {} reports", + reports.len() + ); + } +} + +#[cfg(test)] +mod merge_tests { + //! Scheduler-side aggregation: given several `RuntimeStatsReport`s + //! sharing an `order_by` tag, verify the merged view (total rows, + //! cuts, min/max) reflects the union of the underlying samples. + + use super::*; + use crate::serde::protobuf::{RuntimeStatsPartitionEntry, RuntimeStatsReport}; + use datafusion_proto::protobuf::PhysicalSortExprNode; + + /// Build a report whose partition slots each carry a sketch made + /// from that slot's `values`. Slot `slot_id` in the resulting + /// report has `row_count = values[slot_id].len()` and a sketch + /// over those values. + fn sketching_report( + order_by: Vec, + values_per_slot: Vec>, + ) -> RuntimeStatsReport { + let partitions = values_per_slot + .into_iter() + .enumerate() + .map(|(slot_id, slot_values)| { + let row_count = slot_values.len() as u64; + let sketch = if slot_values.is_empty() { + None + } else { + let digest = TDigest::new(100).merge_unsorted_f64(slot_values); + Some(sketch_to_proto(&digest).unwrap()) + }; + RuntimeStatsPartitionEntry { + partition_id: slot_id as u32, + row_count, + sketch, + } + }) + .collect(); + RuntimeStatsReport { + order_by, + partitions, + } + } + + fn only_group(reports: &[RuntimeStatsReport]) -> MergedRuntimeStats { + let mut groups = merge_reports(reports).expect("merge should succeed"); + match groups.as_slice() { + [_] => groups.remove(0), + other => panic!("expected exactly one group, got {}", other.len()), + } + } + + /// Two reports over disjoint value ranges — merged sketch spans the + /// union, total_rows sums, and the partition_count=2 midpoint cut + /// falls between the two ranges. + #[test] + fn merge_reports_combines_disjoint_ranges() { + // Both reports share an empty `order_by` — we just need two + // reports that land in the same group. + let low_range = sketching_report(vec![], vec![vec![1.0, 2.0, 3.0], vec![]]); + let high_range = sketching_report(vec![], vec![vec![], vec![10.0, 11.0, 12.0]]); + + let group = only_group(&[low_range, high_range]); + assert_eq!(group.partition_count, 2); + assert_eq!(group.task_count, 2); + assert_eq!(group.total_rows, 6); + let midpoint = match group.cuts.as_slice() { + [midpoint] => *midpoint, + other => panic!("expected exactly one cut, got {other:?}"), + }; + assert!( + (3.0..=10.0).contains(&midpoint), + "midpoint cut should land between ranges (got {midpoint})" + ); + assert_eq!(group.min, Some(1.0)); + assert_eq!(group.max, Some(12.0)); + } + + /// partition_count=4 cuts on a uniform [0, 100) sample land roughly + /// at quartiles — verifies the quantile indices `i / partition_count` + /// for `i in 1..partition_count`. + #[test] + fn merge_reports_partition_count_of_four_produces_three_quartile_cuts() { + let uniform: Vec = (0..100).map(|value| value as f64).collect(); + // Single report, partition_count=4: each slot gets 25 uniform + // samples. + let values_per_slot = vec![ + uniform[0..25].to_vec(), + uniform[25..50].to_vec(), + uniform[50..75].to_vec(), + uniform[75..100].to_vec(), + ]; + let report = sketching_report(vec![], values_per_slot); + + let group = only_group(&[report]); + assert_eq!(group.partition_count, 4); + let (p25, p50, p75) = match group.cuts.as_slice() { + [p25, p50, p75] => (*p25, *p50, *p75), + other => panic!("expected 3 cuts, got {other:?}"), + }; + // Loose bounds — T-Digest quantile estimates aren't exact, but + // must land in the expected quartile bands. + assert!((10.0..40.0).contains(&p25), "p25 near 25, got {p25}"); + assert!((35.0..65.0).contains(&p50), "p50 near 50, got {p50}"); + assert!((60.0..90.0).contains(&p75), "p75 near 75, got {p75}"); + } + + /// Row-count-only reports (no sketches) still produce a group with + /// summed `total_rows` — just empty `cuts` and `None` min/max. + #[test] + fn merge_reports_row_count_only_emits_empty_cuts() { + let make_report = |row_counts: [u64; 2]| RuntimeStatsReport { + order_by: vec![], + partitions: vec![ + RuntimeStatsPartitionEntry { + partition_id: 0, + row_count: row_counts[0], + sketch: None, + }, + RuntimeStatsPartitionEntry { + partition_id: 1, + row_count: row_counts[1], + sketch: None, + }, + ], + }; + let group = only_group(&[make_report([100, 200]), make_report([300, 400])]); + assert_eq!(group.partition_count, 2); + assert_eq!(group.total_rows, 1000); + assert!(group.cuts.is_empty()); + assert!(group.min.is_none()); + assert!(group.max.is_none()); + } + + /// Mismatched partition counts within a group surface as an error — + /// the caller (scheduler / slice-D consumer) sees the invariant + /// break rather than silently getting a partial merge. + #[test] + fn merge_reports_errors_on_mismatched_partition_counts() { + let two_partitions = sketching_report(vec![], vec![vec![1.0], vec![2.0]]); + let one_partition = sketching_report(vec![], vec![vec![3.0]]); + let err = merge_reports(&[two_partitions, one_partition]) + .expect_err("mismatched partition counts must error"); + let message = err.to_string(); + assert!( + message.contains("mismatched partition counts"), + "expected mismatch error, got: {message}" + ); + } + + /// A wire-corrupted sketch — one whose scalar-state length is wrong + /// — surfaces the underlying `sketch_from_proto` error rather than + /// getting silently dropped. + #[test] + fn merge_reports_propagates_sketch_decode_errors() { + use datafusion::common::ScalarValue; + + // Six scalars is the valid shape; three is a corrupted wire. + let corrupt_sketch = QuantileSketchState { + state: (0..3) + .map(|_| { + datafusion_proto_common::ScalarValue::try_from(&ScalarValue::Float64( + Some(0.0), + )) + .unwrap() + }) + .collect(), + }; + let report = RuntimeStatsReport { + order_by: vec![], + partitions: vec![RuntimeStatsPartitionEntry { + partition_id: 0, + row_count: 1, + sketch: Some(corrupt_sketch), + }], + }; + let err = merge_reports(&[report]) + .expect_err("corrupt sketch must surface as an error"); + assert!( + err.to_string().contains("expected 6 elements"), + "expected shape-error propagation, got: {err}" + ); + } + + /// Empty input → empty output; ensures no panics or spurious groups. + #[test] + fn merge_reports_empty_input_is_empty_output() { + assert!(merge_reports(&[]).unwrap().is_empty()); + } +} + +#[cfg(test)] +mod plan_walker_tests { + //! `range_repartition_routing_expr` — recursive walk that recognizes + //! both `UnorderedRangeRepartitionExec` and `OrderedRangeRepartitionExec` + //! at any depth. + + use super::*; + use crate::execution_plans::{ + OrderedRangeRepartitionExec, UnorderedRangeRepartitionExec, + }; + use datafusion::arrow::compute::SortOptions; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_expr::{LexOrdering, PhysicalSortExpr}; + use datafusion::physical_plan::expressions::col; + use datafusion::physical_plan::sorts::sort::SortExec; + + fn v_schema() -> Arc { + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])) + } + + fn v_source() -> Arc { + let schema = v_schema(); + let memory = + Arc::new(MemorySourceConfig::try_new(&[vec![]], schema, None).unwrap()); + Arc::new(DataSourceExec::new(memory)) + } + + fn sort_expr_v() -> PhysicalSortExpr { + let schema = v_schema(); + PhysicalSortExpr { + expr: col("v", schema.as_ref()).unwrap(), + options: SortOptions { + descending: false, + nulls_first: false, + }, + } + } + + fn urre_over_source(k: usize) -> Arc { + Arc::new( + UnorderedRangeRepartitionExec::try_new(v_source(), vec![sort_expr_v()], k) + .unwrap(), + ) + } + + fn orre_over_source(k: usize) -> Arc { + // ORRE demands sorted input. + let sort = Arc::new(SortExec::new( + LexOrdering::new(vec![sort_expr_v()]).unwrap(), + v_source(), + )); + Arc::new( + OrderedRangeRepartitionExec::try_new(sort, vec![sort_expr_v()], k).unwrap(), + ) + } + + #[test] + fn range_repartition_routing_expr_bare_source_returns_none() { + assert!( + repartition_routing_expr(v_source().as_ref()) + .unwrap() + .is_none() + ); + } + + #[test] + fn range_repartition_routing_expr_urre_returns_first_order_by_expr() { + let urre = urre_over_source(4); + let expr = repartition_routing_expr(urre.as_ref()) + .unwrap() + .expect("URRE at root should yield a routing expression"); + // The expression is the `v` column — verify by string equality since + // PhysicalExpr doesn't implement PartialEq. + assert_eq!(format!("{expr}"), format!("{}", sort_expr_v().expr)); + } + + #[test] + fn range_repartition_routing_expr_orre_returns_first_order_by_expr() { + let orre = orre_over_source(4); + let expr = repartition_routing_expr(orre.as_ref()) + .unwrap() + .expect("ORRE at root should yield a routing expression"); + assert_eq!(format!("{expr}"), format!("{}", sort_expr_v().expr)); + } + + #[test] + fn range_repartition_routing_expr_descends_through_stats_wrapper() { + // Canonical shape: RuntimeStatsExec above URRE — the walker must + // recurse through the stats wrapper to find the URRE beneath. + let urre = urre_over_source(4); + let stats: Arc = + Arc::new(RuntimeStatsExec::try_new(urre, Some(vec![sort_expr_v()])).unwrap()); + let expr = repartition_routing_expr(stats.as_ref()) + .unwrap() + .expect("walker must descend past RuntimeStatsExec"); + assert_eq!(format!("{expr}"), format!("{}", sort_expr_v().expr)); + } + + #[test] + fn range_repartition_routing_expr_urre_nested_returns_expr() { + // Two stats-wrapper layers to prove the walk descends more than once. + let urre = urre_over_source(4); + let inner_stats = + Arc::new(RuntimeStatsExec::try_new(urre, Some(vec![sort_expr_v()])).unwrap()); + let outer: Arc = Arc::new( + RuntimeStatsExec::try_new(inner_stats, Some(vec![sort_expr_v()])).unwrap(), + ); + assert!(repartition_routing_expr(outer.as_ref()).unwrap().is_some()); + } + + #[test] + fn range_repartition_routing_expr_orre_nested_returns_expr() { + let orre = orre_over_source(4); + let inner_stats = + Arc::new(RuntimeStatsExec::try_new(orre, Some(vec![sort_expr_v()])).unwrap()); + let outer: Arc = Arc::new( + RuntimeStatsExec::try_new(inner_stats, Some(vec![sort_expr_v()])).unwrap(), + ); + assert!(repartition_routing_expr(outer.as_ref()).unwrap().is_some()); + } +} + +#[cfg(test)] +mod overlap_remap_tests { + //! `overlap_remap_partitions` — takes a passthrough + //! `Vec>` and rewrites it under overlap semantics + //! computed from merged sketches. Straddling sub-parts appear in every + //! downstream partition whose range they touch; bookkeeping mismatches + //! surface as errors rather than silent misroutes. + + use super::*; + use crate::serde::protobuf::{RuntimeStatsPartitionEntry, RuntimeStatsReport}; + use crate::serde::scheduler::{ + ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, + PartitionId, PartitionLocation, PartitionStats, + }; + + /// Build a `PartitionLocation` for a producer file identified by + /// `(producer_task_id, sub_part_id)`. `file_id = producer_task_id` + /// matches the passthrough writer's convention. + fn location(sub_part_id: usize, producer_task_id: usize) -> PartitionLocation { + PartitionLocation { + map_partition_id: 0, + partition_id: PartitionId { + job_id: "test-job".into(), + stage_id: 0, + partition_id: sub_part_id, + }, + executor_meta: ExecutorMetadata { + id: format!("exec-{producer_task_id}"), + host: "".to_string(), + port: 0, + grpc_port: 0, + specification: ExecutorSpecification::default().with_vcores(0), + os_info: ExecutorOperatingSystemSpecification::default(), + }, + // `Some(0)` — helper's default is "empty file". Tests that + // need rows overwrite `.partition_stats` explicitly. + partition_stats: PartitionStats::new(Some(0), None, None), + file_id: Some(producer_task_id as u64), + is_sort_shuffle: false, + } + } + + /// Build a report whose sub-parts each carry a T-Digest sketch over + /// `values`. Slot `sub_part_id` gets a sketch of `values[sub_part_id]`. + fn sketch_report( + producer_task_id: usize, + values_per_sub_part: Vec>, + ) -> TaskRuntimeStats { + let partitions = values_per_sub_part + .into_iter() + .enumerate() + .map(|(sub_part_id, samples)| { + let sketch = if samples.is_empty() { + None + } else { + let digest = TDigest::new(100).merge_unsorted_f64(samples.clone()); + Some(sketch_to_proto(&digest).unwrap()) + }; + RuntimeStatsPartitionEntry { + partition_id: sub_part_id as u32, + row_count: samples.len() as u64, + sketch, + } + }) + .collect(); + TaskRuntimeStats { + producer_task_id, + report: RuntimeStatsReport { + order_by: vec![], + partitions, + }, + } + } + + /// Two producers with disjoint value ranges + one downstream cut → + /// each downstream partition gets exactly one producer's files. + #[test] + fn overlap_remap_disjoint_producers_route_to_single_partition() { + // Producer 100 covers [0, 10); producer 200 covers [20, 30). + let reports = vec![ + sketch_report(100, vec![vec![0.0, 5.0, 9.0]]), + sketch_report(200, vec![vec![20.0, 25.0, 29.0]]), + ]; + // Cut at 15 → partition 0 = (-∞, 15), partition 1 = [15, +∞). + let cuts = vec![15.0]; + // Passthrough map: both producers wrote to sub_part_id=0. + let original_partitions = vec![vec![location(0, 100), location(0, 200)]]; + + let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + assert_eq!(remapped.len(), 2, "K = cuts.len() + 1"); + // Partition 0: only producer 100. + assert_eq!(remapped[0].len(), 1); + assert_eq!(remapped[0][0].file_id, Some(100)); + // Partition 1: only producer 200. + assert_eq!(remapped[1].len(), 1); + assert_eq!(remapped[1][0].file_id, Some(200)); + } + + /// A straddling sub-part — one whose sketched [min, max] spans the cut + /// — appears in BOTH downstream partitions' lists. This is the case + /// PerPartitionFilterExec exists to clean up. + #[test] + fn overlap_remap_straddling_producer_appears_in_both_partitions() { + // Producer 300 covers [5, 25) — straddles the cut at 15. + let reports = vec![sketch_report(300, vec![vec![5.0, 15.0, 25.0]])]; + let cuts = vec![15.0]; + let original_partitions = vec![vec![location(0, 300)]]; + + let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + assert_eq!(remapped.len(), 2); + assert_eq!(remapped[0].len(), 1, "straddler in partition 0"); + assert_eq!(remapped[0][0].file_id, Some(300)); + assert_eq!(remapped[1].len(), 1, "straddler in partition 1"); + assert_eq!(remapped[1][0].file_id, Some(300)); + } + + /// A producer file without `file_id` means the writer that produced it + /// wasn't the passthrough writer — remap can't identify the file, so + /// error rather than silently misroute. + #[test] + fn overlap_remap_missing_file_id_errors() { + let reports = vec![sketch_report(100, vec![vec![1.0, 2.0, 3.0]])]; + let cuts = vec![10.0]; + // Loc has file_id=None — invalid for URRE/ORRE stages. + let mut bad = location(0, 100); + bad.file_id = None; + let original_partitions = vec![vec![bad]]; + + let err = cut_partitions(original_partitions, &reports, &cuts) + .expect_err("missing file_id must surface as an error"); + assert!( + err.to_string().contains("missing file_id"), + "unexpected error: {err}" + ); + } + + /// A report entry with no corresponding PartitionLocation is silently + /// ignored — the walk is driven by `original_partitions`, so orphan + /// reports never get consulted. `SuccessfulTask` bundles the shuffle + /// files and their runtime-stats report together, so this scenario + /// shouldn't happen in practice; the shape just doesn't need to + /// enforce it. + #[test] + fn orphan_report_is_silently_ignored() { + // Report from producer 100, but original_partitions only has + // producer 200 — no file to route. + let reports = vec![sketch_report(100, vec![vec![1.0, 2.0, 3.0]])]; + let cuts = vec![10.0]; + let original_partitions = vec![vec![location(0, 200)]]; + + let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + assert_eq!(remapped.len(), 2); + assert!(remapped[0].is_empty()); + assert!(remapped[1].is_empty()); + } + + /// Empty-sketch entries contribute no locations when the file is + /// itself empty (num_rows == 0). No data lost, no error. + #[test] + fn overlap_remap_empty_sketches_produce_empty_partitions() { + let reports = vec![sketch_report(100, vec![vec![]])]; + let cuts = vec![10.0]; + let original_partitions = vec![vec![location(0, 100)]]; + + let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + assert_eq!(remapped.len(), 2); + assert!(remapped[0].is_empty()); + assert!(remapped[1].is_empty()); + } + + /// A file with rows but no matching sketch cannot be routed by + /// overlap — silently skipping would drop rows. Surface as an error. + #[test] + fn missing_sketch_with_rows_errors() { + let reports = vec![sketch_report(100, vec![vec![1.0, 2.0]])]; + let cuts = vec![10.0]; + // File 200 has 5 rows but no report entry exists for it. + let mut orphan = location(0, 200); + orphan.partition_stats = PartitionStats::new(Some(5), None, None); + let original_partitions = vec![vec![orphan]]; + + let err = cut_partitions(original_partitions, &reports, &cuts) + .expect_err("file with rows but no sketch must error"); + let msg = err.to_string(); + assert!( + msg.contains("num_rows=Some(5)") && msg.contains("no usable sketch"), + "unexpected error: {msg}" + ); + } + + /// Multi-cut layout: with K = 4 buckets and cuts [10, 20, 30], a sketch + /// covering [15, 25] must land in exactly buckets 1 and 2 — verifying the + /// `partition_point` range walk lines up with the original inclusive + /// `lower ≤ sketch_max` / exclusive `sketch_min < upper` semantics. + #[test] + fn overlap_remap_multi_cut_range_matches_bucket_semantics() { + let reports = vec![ + // A: covers only bucket 0 + sketch_report(1, vec![vec![1.0, 5.0, 9.0]]), + // B: straddles cuts[0]=10 → buckets 0 and 1 + sketch_report(2, vec![vec![5.0, 10.0, 15.0]]), + // C: fully inside bucket 1 + sketch_report(3, vec![vec![11.0, 15.0, 19.0]]), + // D: sketch_min == cuts[1] → excluded from bucket 1 (upper is exclusive), + // included in buckets 2 and 3 + sketch_report(4, vec![vec![20.0, 25.0, 35.0]]), + // E: covers only the last bucket + sketch_report(5, vec![vec![31.0, 40.0, 50.0]]), + // F: spans the whole range → every bucket + sketch_report(6, vec![vec![0.0, 20.0, 100.0]]), + ]; + let cuts = vec![10.0, 20.0, 30.0]; + let original_partitions = vec![vec![ + location(0, 1), + location(0, 2), + location(0, 3), + location(0, 4), + location(0, 5), + location(0, 6), + ]]; + + let remapped = cut_partitions(original_partitions, &reports, &cuts).unwrap(); + assert_eq!(remapped.len(), 4); + let ids = |b: &[PartitionLocation]| { + let mut v: Vec = b.iter().map(|l| l.file_id.unwrap()).collect(); + v.sort(); + v + }; + assert_eq!(ids(&remapped[0]), vec![1u64, 2, 6]); + assert_eq!(ids(&remapped[1]), vec![2u64, 3, 6]); + assert_eq!(ids(&remapped[2]), vec![4u64, 6]); + assert_eq!(ids(&remapped[3]), vec![4u64, 5, 6]); + } + + /// A file with an unknown row count (`None`) and no usable sketch + /// also errors — we can't confirm the file is empty, so we can't + /// safely skip it. + #[test] + fn missing_sketch_with_unknown_rows_errors() { + let reports: Vec = vec![]; + let cuts = vec![10.0]; + let mut orphan = location(0, 100); + orphan.partition_stats = PartitionStats::default(); // num_rows = None + let original_partitions = vec![vec![orphan]]; + + let err = cut_partitions(original_partitions, &reports, &cuts) + .expect_err("file with unknown rows but no sketch must error"); + let msg = err.to_string(); + assert!( + msg.contains("num_rows=None") && msg.contains("no usable sketch"), + "unexpected error: {msg}" + ); + } +} diff --git a/ballista/core/src/execution_plans/shuffle_writer.rs b/ballista/core/src/execution_plans/shuffle_writer.rs index 0742dbf213..8854962d1b 100644 --- a/ballista/core/src/execution_plans/shuffle_writer.rs +++ b/ballista/core/src/execution_plans/shuffle_writer.rs @@ -27,8 +27,10 @@ use std::sync::Arc; use std::time::Instant; use crate::JobId; -use crate::execution_plans::SortShuffleWriterExec; -use crate::execution_plans::create_shuffle_path; +use crate::execution_plans::{ + OrderedRangeRepartitionExec, SortShuffleWriterExec, UnorderedRangeRepartitionExec, + create_shuffle_path, +}; use crate::extension::SessionConfigExt; use crate::utils; @@ -73,9 +75,15 @@ enum GlobalPartitionMap { /// The plan collapses to a single output partition (e.g. /// `SortPreservingMergeExec`). Every local index → global partition 0. Collapsed, - /// The plan re-establishes a hash-partition K-space (e.g. - /// `RepartitionExec::Hash(_, K)`). Local index == global (0..K). - HashSpace, + /// The plan re-establishes a fresh K-space of output partitions — every + /// operator that fans rows into K distinct outputs and assigns each output + /// a fresh index. `RepartitionExec::Hash(_, K)` and + /// `RepartitionExec::RoundRobinBatch(K)` qualify; + /// `UnorderedRangeRepartitionExec` (range-routed) qualifies too. Local + /// index == global (0..K) regardless of the routing algorithm — the K-space + /// is per-stage and per-task-slot, `file_id` disambiguates the same + /// `partition_id` across producers. + KSpace, /// Nothing between the writer and the leaves rewrites partitioning — /// local index `i` is `slice[i]` globally. Empty when no slice was /// stamped, in which case local is used as-is (identity). @@ -86,7 +94,7 @@ impl GlobalPartitionMap { fn resolve(&self, local: usize) -> u64 { match self { GlobalPartitionMap::Collapsed => 0, - GlobalPartitionMap::HashSpace => local as u64, + GlobalPartitionMap::KSpace => local as u64, GlobalPartitionMap::PassThrough(slice) => { slice.get(local).copied().unwrap_or(local) as u64 } @@ -98,7 +106,14 @@ impl GlobalPartitionMap { /// determines the output partitioning shape: /// /// - `SortPreservingMergeExec` → `Collapsed` (fan-in to 1). -/// - `RepartitionExec` producing a hash / round-robin K-space → `HashSpace`. +/// - `RepartitionExec(Hash|RoundRobin)` → `KSpace` (fresh K-space by hash / +/// round-robin routing). +/// - `UnorderedRangeRepartitionExec` / `OrderedRangeRepartitionExec` → +/// `KSpace` (fresh K-space by range routing, unordered or order-preserving). +/// Same shape as hash-repartition from the writer's perspective: K distinct +/// outputs, local index == global partition, `file_id` disambiguates across +/// producers. Content-range info (which K-slot holds which value range) +/// travels through the separate `RuntimeStatsExec` sketch-report channel. /// - Otherwise recurse into the sole child (Filter/Sort/Projection/… are /// partitioning-preserving passthroughs). /// - If we hit a leaf or a fan-in without recognising it, treat it as @@ -107,13 +122,13 @@ fn walk_child_partition_mapping( plan: &Arc, global_output_partition_ids: &[usize], ) -> GlobalPartitionMap { - if plan.downcast_ref::().is_some() { + if plan.is::() { return GlobalPartitionMap::Collapsed; } if let Some(repart) = plan.downcast_ref::() { match repart.partitioning() { Partitioning::Hash(_, _) | Partitioning::RoundRobinBatch(_) => { - return GlobalPartitionMap::HashSpace; + return GlobalPartitionMap::KSpace; } Partitioning::UnknownPartitioning(_) => { // RepartitionExec still exchanges rows and freshly numbers @@ -128,6 +143,11 @@ fn walk_child_partition_mapping( } } } + if plan.is::() + || plan.is::() + { + return GlobalPartitionMap::KSpace; + } let children = plan.children(); if children.len() == 1 { return walk_child_partition_mapping(children[0], global_output_partition_ids); @@ -144,7 +164,7 @@ fn walk_child_partition_mapping( /// /// Two cases: /// -/// - `SortShuffleWriter(Hash(K))` — HashSpace by construction; the K-space +/// - `SortShuffleWriter(Hash(K))` — KSpace by construction; the K-space /// `[0..K-1]` is intrinsic. Input ids are irrelevant. /// - `ShuffleWriter` — always passthrough; the child's plan shape decides: /// - `SortPreservingMergeExec` in the child chain → `[0]` (collapse). @@ -164,14 +184,14 @@ pub fn compute_global_output_partition_ids( }; return (0..*k).collect(); } - if stage_plan.downcast_ref::().is_some() { + if stage_plan.is::() { let children = stage_plan.children(); let [child] = children.as_slice() else { unreachable!("ShuffleWriterExec always has exactly one child"); }; return match walk_child_partition_mapping(child, global_input_partition_ids) { GlobalPartitionMap::Collapsed => vec![0], - GlobalPartitionMap::HashSpace => { + GlobalPartitionMap::KSpace => { let k = child.properties().output_partitioning().partition_count(); (0..k).collect() } @@ -494,7 +514,7 @@ impl ShuffleWriterExec { /// operator's output_partitioning). `summary.partition_id` is the /// **global** output partition id downstream will address, computed via /// `walk_child_partition_mapping` over the child plan — either - /// `global_output_partition_ids[local]`, `local` (hash K-space), or `0` + /// `global_output_partition_ids[local]`, `local` (K-space), or `0` /// (collapsed / SPM). pub fn execute_shuffle_write( self, diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index 112ef6db57..2b080c429f 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -319,6 +319,7 @@ impl SessionStateExt for SessionState { .with_cache_factory(Some(Arc::new(BallistaCacheFactory::new()))) .with_runtime_env(Arc::new(runtime_env)) .with_query_planner(Arc::new(planner)) + .with_optimizer_rules(crate::optimizer::ballista_default_optimizer_rules()) .with_scalar_functions(ballista_scalar_functions()) .with_aggregate_functions(ballista_aggregate_functions()) .with_window_functions(ballista_window_functions()) @@ -338,8 +339,12 @@ impl SessionStateExt for SessionState { let ballista_config = session_config.ballista_config(); + let optimizer_rules = + crate::optimizer::with_ballista_optimizer_rules(self.optimizers()); + let builder = SessionStateBuilder::new_from_existing(self) .with_config(session_config) + .with_optimizer_rules(optimizer_rules) .with_cache_factory(Some(Arc::new(BallistaCacheFactory::new()))); let builder = match planner_override { diff --git a/ballista/core/src/lib.rs b/ballista/core/src/lib.rs index 1336c1b825..feb3f05864 100644 --- a/ballista/core/src/lib.rs +++ b/ballista/core/src/lib.rs @@ -69,6 +69,8 @@ pub mod kll; #[cfg(feature = "build-binary")] /// Object store configuration and utilities for distributed file access. pub mod object_store; +/// Ballista-specific logical optimizer rules. +pub mod optimizer; /// Query planning utilities for distributed execution. pub mod planner; /// Runtime registry for codec and function registration. diff --git a/ballista/core/src/optimizer.rs b/ballista/core/src/optimizer.rs new file mode 100644 index 0000000000..d779bdff20 --- /dev/null +++ b/ballista/core/src/optimizer.rs @@ -0,0 +1,475 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Ballista-specific logical optimizer rules. +//! +//! [`NotInSubqueryRewrite`](crate::optimizer::NotInSubqueryRewrite) rewrites +//! uncorrelated `NOT IN (subquery)` filter +//! predicates into a plain anti join combined with a one-row count aggregate, +//! so the query never needs DataFusion's null-aware hash join. That join +//! coordinates probe-side state through in-process atomics, which forces +//! Ballista to execute it in a single task (see the discussion on +//! ). The rewritten +//! plan uses only ordinary joins and aggregates and distributes normally. + +use std::sync::Arc; + +use datafusion::common::tree_node::Transformed; +use datafusion::common::{Column, DFSchemaRef, Result}; +use datafusion::functions_aggregate::expr_fn::count; +use datafusion::logical_expr::expr::InSubquery; +use datafusion::logical_expr::utils::{conjunction, split_conjunction_owned}; +use datafusion::logical_expr::{ + Expr, ExprSchemable, Filter, JoinType, LogicalPlan, LogicalPlanBuilder, +}; +use datafusion::optimizer::optimizer::ApplyOrder; +use datafusion::optimizer::{OptimizerConfig, OptimizerRule}; +use datafusion::prelude::lit; + +use crate::config::BallistaConfig; + +/// Returns `existing` with Ballista's logical optimizer rules prepended. +/// +/// [`NotInSubqueryRewrite`] must run before DataFusion's +/// `DecorrelatePredicateSubquery`, which would otherwise turn eligible `NOT IN` +/// predicates into null-aware anti joins first. A rule with the same name in +/// `existing` is dropped so repeated session upgrades stay idempotent. +pub fn with_ballista_optimizer_rules( + existing: &[Arc], +) -> Vec> { + let ballista_rule = NotInSubqueryRewrite::new(); + let mut rules: Vec> = + Vec::with_capacity(existing.len() + 1); + let name = ballista_rule.name().to_owned(); + rules.push(Arc::new(ballista_rule)); + rules.extend(existing.iter().filter(|rule| rule.name() != name).cloned()); + rules +} + +/// Returns DataFusion's default logical optimizer rules with Ballista's rules +/// prepended. Used by every Ballista session builder. +pub fn ballista_default_optimizer_rules() -> Vec> { + with_ballista_optimizer_rules(&datafusion::optimizer::Optimizer::default().rules) +} + +/// Rewrites `expr NOT IN (subquery)` filter conjuncts into a distributable +/// plan shape before DataFusion's `DecorrelatePredicateSubquery` can turn them +/// into null-aware anti joins. +/// +/// For a filter `... WHERE e NOT IN (SELECT b FROM s)` the rewrite produces: +/// +/// ```text +/// Projection: +/// Filter: __cnt = 0 OR (e IS NOT NULL AND __cnt = __cnt_non_null) +/// CrossJoin +/// Aggregate: count(1) AS __cnt, count(b) AS __cnt_non_null (over s) +/// LeftAnti Join: e = b +/// +/// SubqueryAlias: s +/// ``` +/// +/// which is equivalent under SQL three-valued `NOT IN` semantics evaluated in +/// a `WHERE` context: +/// +/// - `s` empty: every input row passes (`__cnt = 0`). +/// - `e` matches some `b`: the anti join drops the row. +/// - `e` is NULL and `s` is non-empty: the row is dropped. +/// - `s` contains a NULL `b`: every remaining row is dropped +/// (`__cnt <> __cnt_non_null`). +/// +/// The rewrite only fires for predicates where DataFusion would otherwise +/// need null-aware semantics: uncorrelated subqueries with a single output +/// column where either side of the comparison is nullable. Everything else is +/// left for DataFusion's own subquery decorrelation, which already produces +/// distributable plans for those cases. +/// +/// The rule is enabled by default and controlled by the +/// `ballista.optimizer.not_in_subquery_rewrite` configuration key. +#[derive(Debug, Default)] +pub struct NotInSubqueryRewrite {} + +impl NotInSubqueryRewrite { + /// Creates a new instance of the rewrite rule. + pub fn new() -> Self { + Self {} + } +} + +impl OptimizerRule for NotInSubqueryRewrite { + fn name(&self) -> &str { + "not_in_subquery_rewrite" + } + + fn apply_order(&self) -> Option { + Some(ApplyOrder::BottomUp) + } + + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + config: &dyn OptimizerConfig, + ) -> Result> { + let enabled = config + .options() + .extensions + .get::() + .map(|c| c.not_in_subquery_rewrite_enabled()) + .unwrap_or_else(|| { + BallistaConfig::default().not_in_subquery_rewrite_enabled() + }); + if !enabled { + return Ok(Transformed::no(plan)); + } + + let LogicalPlan::Filter(filter) = plan else { + return Ok(Transformed::no(plan)); + }; + rewrite_filter(filter, config) + } +} + +/// Returns the [`InSubquery`] when the expression is a negated `IN` predicate, +/// in either its `expr NOT IN (subquery)` or `NOT (expr IN (subquery))` form. +fn as_negated_in_subquery(expr: &Expr) -> Option<&InSubquery> { + match expr { + Expr::InSubquery(in_subquery) if in_subquery.negated => Some(in_subquery), + Expr::Not(inner) => match inner.as_ref() { + Expr::InSubquery(in_subquery) if !in_subquery.negated => Some(in_subquery), + _ => None, + }, + _ => None, + } +} + +/// Returns the negated `IN` subquery to rewrite, or `None` when the conjunct +/// should be left for DataFusion's own subquery decorrelation: correlated +/// subqueries, multi-column subqueries, volatile probe expressions, and +/// predicates whose keys are provably non-nullable (those never become +/// null-aware joins in the first place). +fn rewritable_in_subquery<'a>( + conjunct: &'a Expr, + input_schema: &DFSchemaRef, +) -> Option<&'a InSubquery> { + let in_subquery = as_negated_in_subquery(conjunct)?; + if !in_subquery.subquery.outer_ref_columns.is_empty() { + return None; + } + let subquery_schema = in_subquery.subquery.subquery.schema(); + if subquery_schema.fields().len() != 1 { + return None; + } + if in_subquery.expr.is_volatile() { + return None; + } + let expr_nullable = in_subquery + .expr + .nullable(input_schema.as_ref()) + .unwrap_or(true); + if !expr_nullable && !subquery_schema.field(0).is_nullable() { + // Both sides are non-nullable: DataFusion plans a plain anti join for + // this case already, so the rewrite would only add an aggregate scan. + return None; + } + Some(in_subquery) +} + +/// Rewrites every eligible `NOT IN` conjunct of `filter`, returning the +/// original filter untouched when there is nothing to do. +fn rewrite_filter( + filter: Filter, + config: &dyn OptimizerConfig, +) -> Result> { + let input_schema = filter.input.schema().clone(); + let has_candidate = split_conjunction_owned(filter.predicate.clone()) + .iter() + .any(|conjunct| rewritable_in_subquery(conjunct, &input_schema).is_some()); + if !has_candidate { + return Ok(Transformed::no(LogicalPlan::Filter(filter))); + } + + let mut current = Arc::unwrap_or_clone(filter.input); + let mut conjuncts = Vec::new(); + for conjunct in split_conjunction_owned(filter.predicate) { + let Some(in_subquery) = rewritable_in_subquery(&conjunct, &input_schema) else { + conjuncts.push(conjunct); + continue; + }; + + let probe_expr = in_subquery.expr.as_ref().clone(); + let subquery_plan = + Arc::unwrap_or_clone(Arc::clone(&in_subquery.subquery.subquery)); + + let alias = config.alias_generator().next("__ballista_not_in"); + let count_all_name = format!("{alias}_cnt"); + let count_non_null_name = format!("{alias}_cnt_non_null"); + + // Anti join drops the rows whose probe expression matches a subquery + // value. The subquery side is aliased so a self-`NOT IN` does not + // produce ambiguous column names in the join predicate. + let aliased_subquery = LogicalPlanBuilder::from(subquery_plan.clone()) + .alias(alias.clone())? + .build()?; + let subquery_column = Expr::Column(Column::new( + Some(alias), + aliased_subquery.schema().field(0).name(), + )); + let anti_join = LogicalPlanBuilder::from(current) + .join_on( + aliased_subquery, + JoinType::LeftAnti, + [probe_expr.clone().eq(subquery_column)], + )? + .build()?; + + // One-row aggregate capturing the two global facts three-valued + // `NOT IN` semantics need: whether the subquery is empty and whether + // it contains a NULL. Placed on the build (left) side of the cross + // join so only a single row is collected. + let count_column = + Expr::Column(Column::from(subquery_plan.schema().qualified_field(0))); + let aggregate = LogicalPlanBuilder::from(subquery_plan) + .aggregate( + Vec::::new(), + vec![ + count(lit(1)).alias(&count_all_name), + count(count_column).alias(&count_non_null_name), + ], + )? + .build()?; + current = LogicalPlanBuilder::from(aggregate) + .cross_join(anti_join)? + .build()?; + + // Rows surviving the anti join pass when the subquery is empty, or + // when the probe value is non-NULL and the subquery contains no NULL. + let count_all = Expr::Column(Column::new_unqualified(&count_all_name)); + let count_non_null = Expr::Column(Column::new_unqualified(&count_non_null_name)); + conjuncts.push( + count_all + .clone() + .eq(lit(0i64)) + .or(probe_expr.is_not_null().and(count_all.eq(count_non_null))), + ); + } + + let predicate = + conjunction(conjuncts).expect("at least one rewritten conjunct must exist"); + let restore_columns: Vec = input_schema + .iter() + .map(|qualified_field| Expr::Column(Column::from(qualified_field))) + .collect(); + let new_plan = LogicalPlanBuilder::from(current) + .filter(predicate)? + .project(restore_columns)? + .build()?; + Ok(Transformed::yes(new_plan)) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::execution::SessionStateBuilder; + use datafusion::logical_expr::builder::LogicalTableSource; + use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; + use datafusion::optimizer::Optimizer; + use datafusion::prelude::{SessionConfig, col, lit, not_in_subquery}; + + use super::NotInSubqueryRewrite; + use crate::config::{BALLISTA_NOT_IN_SUBQUERY_REWRITE, BallistaConfig}; + + fn scan(table: &str, field: &str, nullable: bool) -> LogicalPlan { + let schema = Schema::new(vec![Field::new(field, DataType::Int32, nullable)]); + LogicalPlanBuilder::scan( + table, + Arc::new(LogicalTableSource::new(Arc::new(schema))), + None, + ) + .unwrap() + .build() + .unwrap() + } + + fn optimize_with_config(plan: LogicalPlan, config: BallistaConfig) -> LogicalPlan { + let state = SessionStateBuilder::new() + .with_config(SessionConfig::new().with_option_extension(config)) + .build(); + Optimizer::with_rules(vec![Arc::new(NotInSubqueryRewrite::new())]) + .optimize(plan, &state, |_, _| {}) + .unwrap() + } + + fn optimize(plan: LogicalPlan) -> LogicalPlan { + optimize_with_config(plan, BallistaConfig::default()) + } + + fn not_in_plan(nullable: bool) -> LogicalPlan { + let subquery = scan("t2", "b", nullable); + LogicalPlanBuilder::from(scan("t1", "a", nullable)) + .filter(not_in_subquery(col("a"), Arc::new(subquery))) + .unwrap() + .build() + .unwrap() + } + + #[test] + fn rewrites_nullable_not_in_to_anti_join_with_count_aggregate() { + let plan = not_in_plan(true); + let schema = plan.schema().clone(); + + let optimized = optimize(plan); + let display = format!("{}", optimized.display_indent()); + + assert!( + !display.contains("IN ()"), + "InSubquery must be rewritten away:\n{display}" + ); + assert!( + display.contains("LeftAnti"), + "rewrite must produce an anti join:\n{display}" + ); + assert!( + display.contains("count("), + "rewrite must produce a count aggregate:\n{display}" + ); + assert_eq!( + optimized.schema().as_ref(), + schema.as_ref(), + "rewrite must preserve the plan schema" + ); + } + + #[test] + fn rewrites_negated_in_subquery_wrapped_in_not() { + let subquery = scan("t2", "b", true); + let in_expr = datafusion::prelude::in_subquery(col("a"), Arc::new(subquery)); + let plan = LogicalPlanBuilder::from(scan("t1", "a", true)) + .filter(Expr::Not(Box::new(in_expr))) + .unwrap() + .build() + .unwrap(); + + let display = format!("{}", optimize(plan).display_indent()); + assert!( + !display.contains("IN ()"), + "NOT (IN subquery) must be rewritten away:\n{display}" + ); + assert!(display.contains("LeftAnti"), "{display}"); + } + + #[test] + fn preserves_other_filter_conjuncts() { + let subquery = scan("t2", "b", true); + let plan = LogicalPlanBuilder::from(scan("t1", "a", true)) + .filter( + col("a") + .gt(lit(5)) + .and(not_in_subquery(col("a"), Arc::new(subquery))), + ) + .unwrap() + .build() + .unwrap(); + + let display = format!("{}", optimize(plan).display_indent()); + assert!(!display.contains("IN ()"), "{display}"); + assert!( + display.contains("t1.a > Int32(5)"), + "other conjuncts must survive the rewrite:\n{display}" + ); + } + + #[test] + fn keeps_correlated_not_in_subquery() { + let sub = scan("t2", "b", true); + let mut expr = not_in_subquery(col("a"), Arc::new(sub)); + if let Expr::InSubquery(ref mut in_subquery) = expr { + in_subquery.subquery.outer_ref_columns.push( + datafusion::logical_expr::expr_fn::out_ref_col(DataType::Int32, "t1.a"), + ); + } else { + unreachable!("not_in_subquery must build an InSubquery expression"); + } + let plan = LogicalPlanBuilder::from(scan("t1", "a", true)) + .filter(expr) + .unwrap() + .build() + .unwrap(); + + let display = format!("{}", optimize(plan).display_indent()); + assert!( + display.contains("IN ()"), + "correlated subqueries must be left to DataFusion:\n{display}" + ); + } + + #[test] + fn keeps_not_in_when_keys_are_not_nullable() { + let display = format!("{}", optimize(not_in_plan(false)).display_indent()); + assert!( + display.contains("IN ()"), + "non-nullable keys do not need the rewrite:\n{display}" + ); + } + + #[test] + fn keeps_positive_in_subquery() { + let subquery = scan("t2", "b", true); + let plan = LogicalPlanBuilder::from(scan("t1", "a", true)) + .filter(datafusion::prelude::in_subquery( + col("a"), + Arc::new(subquery), + )) + .unwrap() + .build() + .unwrap(); + + let display = format!("{}", optimize(plan).display_indent()); + assert!( + display.contains("IN ()"), + "positive IN subqueries must not be rewritten:\n{display}" + ); + } + + #[test] + fn disabled_by_configuration() { + use datafusion::config::ExtensionOptions; + + let mut config = BallistaConfig::default(); + config + .set( + BALLISTA_NOT_IN_SUBQUERY_REWRITE + .strip_prefix("ballista.") + .unwrap(), + "false", + ) + .unwrap(); + + let display = format!( + "{}", + optimize_with_config(not_in_plan(true), config).display_indent() + ); + assert!( + display.contains("IN ()"), + "disabling the config key must skip the rewrite:\n{display}" + ); + } +} diff --git a/ballista/core/src/serde/generated/ballista.rs b/ballista/core/src/serde/generated/ballista.rs index ca5b589b7d..5bf6964eb2 100644 --- a/ballista/core/src/serde/generated/ballista.rs +++ b/ballista/core/src/serde/generated/ballista.rs @@ -31,7 +31,7 @@ pub struct LogicalPlanCacheNode { pub struct BallistaPhysicalPlanNode { #[prost( oneof = "ballista_physical_plan_node::PhysicalPlanType", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10" )] pub physical_plan_type: ::core::option::Option< ballista_physical_plan_node::PhysicalPlanType, @@ -59,6 +59,8 @@ pub mod ballista_physical_plan_node { UnorderedRangeRepartition(super::UnorderedRangeRepartitionExecNode), #[prost(message, tag = "9")] OrderedRangeRepartition(super::OrderedRangeRepartitionExecNode), + #[prost(message, tag = "10")] + PerPartitionFilter(super::PerPartitionFilterExecNode), } } /// Value-range router over N locally-sorted overlapping input partitions. @@ -133,6 +135,17 @@ pub struct UnorderedRangeRepartitionExecNode { #[prost(uint32, tag = "2")] pub output_partitions: u32, } +/// Filter with per-input-partition predicates. `predicates\[k\]` is the +/// boolean expression applied to input partition `k`. Requires +/// `predicates.len() == input_partition_count`. The child plan is +/// plumbed by the framework as `inputs\[0\]` during decode. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PerPartitionFilterExecNode { + #[prost(message, repeated, tag = "1")] + pub predicates: ::prost::alloc::vec::Vec< + ::datafusion_proto::protobuf::PhysicalExprNode, + >, +} #[derive(Clone, PartialEq, ::prost::Message)] pub struct ChaosExecNode { #[prost(double, tag = "1")] @@ -871,6 +884,49 @@ pub struct SuccessfulTask { /// so we might want to think about some refactoring of the task definitions #[prost(message, repeated, tag = "2")] pub partitions: ::prost::alloc::vec::Vec, + /// Reports from `RuntimeStatsExec` operators in this task's plan that + /// are still valid at the plan's output (walked from the top through + /// distribution-preserving nodes only — see + /// `range_repartition_common::preserves_distribution`). Empty when the + /// plan has no such stats-taps. Currently reports one entry per + /// executed `RuntimeStatsExec`; the scheduler groups by `order_by` tag + /// to combine reports across tasks/executors. + #[prost(message, repeated, tag = "3")] + pub runtime_stats: ::prost::alloc::vec::Vec, +} +/// One report per `RuntimeStatsExec` in the executed plan. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RuntimeStatsReport { + /// What the operator was sampling. Empty = row-count-only mode. When + /// non-empty, the first entry identifies which routing expression this + /// sketch describes; the scheduler groups sketches by this tag to + /// combine samples across tasks that were sampling the same expression. + #[prost(message, repeated, tag = "1")] + pub order_by: ::prost::alloc::vec::Vec< + ::datafusion_proto::protobuf::PhysicalSortExprNode, + >, + /// Per-partition observations. Interpretation depends on the operator's + /// position in the plan — pre-repartition gets one entry per input + /// partition, post-repartition gets one per output sub-partition. The + /// scheduler groups by `order_by` tag and aggregates. + #[prost(message, repeated, tag = "2")] + pub partitions: ::prost::alloc::vec::Vec, +} +/// One partition's observations from a `RuntimeStatsExec`. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RuntimeStatsPartitionEntry { + #[prost(uint32, tag = "1")] + pub partition_id: u32, + #[prost(uint64, tag = "2")] + pub row_count: u64, + /// Present when the `RuntimeStatsExec` was in sketch mode AND this + /// partition observed at least one non-null routing value. + /// + /// TODO: `optional MinMaxState min_max` — for a lighter post-repartition + /// mode where the bin-packer just needs (min, max, count) per + /// sub-partition and a full T-Digest is overkill. + #[prost(message, optional, tag = "3")] + pub sketch: ::core::option::Option, } #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct ExecutionError {} diff --git a/ballista/core/src/serde/mod.rs b/ballista/core/src/serde/mod.rs index 6be2b29843..3ceaa2dfd6 100644 --- a/ballista/core/src/serde/mod.rs +++ b/ballista/core/src/serde/mod.rs @@ -57,8 +57,9 @@ use std::{convert::TryInto, io::Cursor}; use crate::execution_plans::sort_shuffle::SortShuffleConfig; use crate::execution_plans::{ BufferExec, BufferMode, ChaosExec, CoalescePlan, OrderedRangeRepartitionExec, - PartitionGroup, RuntimeStatsExec, ShuffleReaderExec, ShuffleWriterExec, - SortShuffleWriterExec, UnorderedRangeRepartitionExec, UnresolvedShuffleExec, + PartitionGroup, PerPartitionFilterExec, RuntimeStatsExec, ShuffleReaderExec, + ShuffleWriterExec, SortShuffleWriterExec, UnorderedRangeRepartitionExec, + UnresolvedShuffleExec, }; use crate::serde::protobuf::{ ballista_logical_plan_node::LogicalPlanType, @@ -636,6 +637,31 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { node.output_partitions as usize, )?)) } + PhysicalPlanType::PerPartitionFilter(node) => { + let [input] = inputs else { + return Err(DataFusionError::Internal(format!( + "PerPartitionFilterExec expects exactly 1 input, got {}", + inputs.len() + ))); + }; + let schema = input.schema(); + let predicates = node + .predicates + .iter() + .map(|p| { + datafusion_proto::physical_plan::from_proto::parse_physical_expr( + p, + ctx, + schema.as_ref(), + self, + ) + }) + .collect::, DataFusionError>>()?; + Ok(Arc::new(PerPartitionFilterExec::try_new( + input.clone(), + predicates, + )?)) + } } } @@ -874,6 +900,28 @@ impl PhysicalExtensionCodec for BallistaPhysicalExtensionCodec { )) })?; Ok(()) + } else if let Some(exec) = node.downcast_ref::() { + let predicates = exec + .predicates() + .iter() + .map(|p| { + datafusion_proto::physical_plan::to_proto::serialize_physical_expr( + p, + self.default_codec.as_ref(), + ) + }) + .collect::, DataFusionError>>()?; + let proto = protobuf::BallistaPhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::PerPartitionFilter( + protobuf::PerPartitionFilterExecNode { predicates }, + )), + }; + proto.encode(buf).map_err(|e| { + DataFusionError::Internal(format!( + "failed to encode PerPartitionFilterExec: {e:?}" + )) + })?; + Ok(()) } else { Err(DataFusionError::Internal(format!( "Unsupported plan node, name: [{}] ", @@ -1751,6 +1799,63 @@ mod test { ); } + /// `PerPartitionFilterExec` round-trips through the codec: three input + /// partitions with three distinct range predicates re-materialize + /// on the other side with the same predicate strings in the same + /// order. + #[tokio::test] + async fn test_per_partition_filter_exec_roundtrip() { + use crate::execution_plans::PerPartitionFilterExec; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::physical_plan::repartition::RepartitionExec; + use datafusion::scalar::ScalarValue; + + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let source: Arc = Arc::new( + datafusion::physical_plan::empty::EmptyExec::new(schema.clone()), + ); + let input: Arc = Arc::new( + RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(3)).unwrap(), + ); + let range = |lo: i64, hi: i64| -> Arc { + let c: Arc = Arc::new(Column::new("v", 0)); + let ge: Arc = Arc::new(BinaryExpr::new( + c.clone(), + Operator::GtEq, + Arc::new(Literal::new(ScalarValue::Int64(Some(lo)))), + )); + let lt: Arc = Arc::new(BinaryExpr::new( + c, + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int64(Some(hi)))), + )); + Arc::new(BinaryExpr::new(ge, Operator::And, lt)) + }; + let predicates = vec![range(0, 10), range(10, 20), range(20, 30)]; + let original = + PerPartitionFilterExec::try_new(input.clone(), predicates.clone()).unwrap(); + + let codec = BallistaPhysicalExtensionCodec::default(); + let mut buf: Vec = vec![]; + codec.try_encode(Arc::new(original), &mut buf).unwrap(); + + let ctx = SessionContext::new().task_ctx(); + let decoded_plan = codec.try_decode(&buf, &[input], &ctx).unwrap(); + let decoded = decoded_plan + .downcast_ref::() + .expect("Expected PerPartitionFilterExec"); + assert_eq!(decoded.predicates().len(), 3); + for (k, expected) in predicates.iter().enumerate() { + assert_eq!( + decoded.predicates()[k].to_string(), + expected.to_string(), + "predicate {k} mismatched after roundtrip", + ); + } + } + /// `BufferExec` in `Dam` mode round-trips through the codec. /// The child is decoded from `inputs[0]` by the framework, so the /// wire only carries the mode discriminant — a wrong discriminant diff --git a/ballista/core/src/utils.rs b/ballista/core/src/utils.rs index 34b5b05f10..07f5ca3948 100644 --- a/ballista/core/src/utils.rs +++ b/ballista/core/src/utils.rs @@ -33,11 +33,13 @@ use datafusion::physical_plan::{ExecutionPlan, RecordBatchStream, metrics}; use futures::StreamExt; use log::error; use std::io::BufWriter; +use std::net::SocketAddr; use std::path::Path; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use std::{fs::File, pin::Pin}; use tonic::codegen::StdError; +use tonic::transport::server::TcpIncoming; use tonic::transport::{Channel, Endpoint, Error, Server}; /// Configuration for gRPC client connections. @@ -165,6 +167,7 @@ pub fn default_session_builder( .with_default_features() .with_config(config) .with_runtime_env(Arc::new(RuntimeEnvBuilder::new().build()?)) + .with_optimizer_rules(crate::optimizer::ballista_default_optimizer_rules()) .with_scalar_functions(ballista_scalar_functions()) .with_aggregate_functions(ballista_aggregate_functions()) .with_window_functions(ballista_window_functions()) @@ -355,6 +358,30 @@ pub fn create_grpc_server(config: &GrpcServerConfig) -> Server { ))) } +/// Binds a gRPC server's listening socket, for use with tonic's +/// `serve_with_incoming` / `serve_with_incoming_shutdown`. Unlike tonic's +/// `serve`, which binds lazily inside the future it returns, the socket is +/// listening by the time this returns — so use this whenever a peer may be +/// told to connect as soon as the server is started. +/// +/// tonic ignores the builder's `tcp_nodelay` and `tcp_keepalive` when serving +/// from a pre-bound listener, so this applies the same values that +/// [`create_grpc_server`] sets, for the same reasons. The remaining settings +/// still come from the builder. +/// +/// # Panics +/// +/// The listener is registered with the Tokio reactor, so this must be called +/// from within a Tokio runtime. +pub fn create_grpc_server_incoming( + addr: SocketAddr, + config: &GrpcServerConfig, +) -> Result { + Ok(TcpIncoming::bind(addr)? + .with_nodelay(Some(true)) + .with_keepalive(Some(Duration::from_secs(config.tcp_keepalive_seconds)))) +} + /// Recursively collects metrics from an execution plan and all its children. pub fn collect_plan_metrics(plan: &dyn ExecutionPlan) -> Vec { let mut metrics_array = Vec::::new(); @@ -439,4 +466,37 @@ mod tests { let result = create_grpc_client_endpoint("not a valid url", None); assert!(result.is_err()); } + + /// The point of binding up front is that the port is reachable before + /// anything is served on it, so a peer told to connect back cannot arrive + /// too early. + #[tokio::test] + async fn test_create_grpc_server_incoming_binds_eagerly() { + let incoming = create_grpc_server_incoming( + "127.0.0.1:0".parse().unwrap(), + &GrpcServerConfig::default(), + ) + .expect("bind"); + let addr = incoming.local_addr().expect("local addr"); + + // `incoming` is never handed to a server, and yet: + tokio::net::TcpStream::connect(addr) + .await + .expect("port is already listening"); + } + + /// Binding eagerly means a port conflict surfaces here, as an error, rather + /// than later inside the spawned server task. + #[tokio::test] + async fn test_create_grpc_server_incoming_port_in_use() { + let first = create_grpc_server_incoming( + "127.0.0.1:0".parse().unwrap(), + &GrpcServerConfig::default(), + ) + .expect("bind"); + let addr = first.local_addr().expect("local addr"); + + let result = create_grpc_server_incoming(addr, &GrpcServerConfig::default()); + assert!(result.is_err()); + } } diff --git a/ballista/executor/Cargo.toml b/ballista/executor/Cargo.toml index 44740aaa69..09a708a648 100644 --- a/ballista/executor/Cargo.toml +++ b/ballista/executor/Cargo.toml @@ -34,7 +34,15 @@ required-features = ["build-binary"] [features] arrow-ipc-optimizations = [] -build-binary = ["clap", "tracing-subscriber", "tracing-appender", "tracing", "ballista-core/build-binary", "mimalloc", "dep:axum"] +build-binary = [ + "clap", + "tracing-subscriber", + "tracing-appender", + "tracing", + "ballista-core/build-binary", + "mimalloc", + "dep:axum", +] default = ["arrow-ipc-optimizations", "build-binary"] spark-compat = ["ballista-core/spark-compat"] diff --git a/ballista/executor/src/config.rs b/ballista/executor/src/config.rs index 2b0ffe608d..9dbd4c78cd 100644 --- a/ballista/executor/src/config.rs +++ b/ballista/executor/src/config.rs @@ -169,14 +169,14 @@ pub struct Config { help = "Tracing log rotation policy." )] pub log_rotation_policy: ballista_core::config::LogRotationPolicy, - /// Maximum size of incoming gRPC messages in bytes (default: 16MB). + /// Maximum size of incoming gRPC messages in bytes (default: 16MiB). #[arg( long, default_value_t = 16777216, help = "The maximum size of a decoded message at the grpc server side." )] pub grpc_server_max_decoding_message_size: u32, - /// Maximum size of outgoing gRPC messages in bytes (default: 16MB). + /// Maximum size of outgoing gRPC messages in bytes (default: 16MiB). #[arg( long, default_value_t = 16777216, diff --git a/ballista/executor/src/execution_engine.rs b/ballista/executor/src/execution_engine.rs index 7f779d2c82..4f8e6663d8 100644 --- a/ballista/executor/src/execution_engine.rs +++ b/ballista/executor/src/execution_engine.rs @@ -86,6 +86,18 @@ pub trait QueryStageExecutor: Sync + Send + Debug + Display { /// Collects execution metrics from all operators in the plan. fn collect_plan_metrics(&self) -> Vec; + + /// Collect runtime-stats reports for every `RuntimeStatsExec` still + /// valid at the plan's output (walked through the distribution- + /// preserving whitelist). Called at task completion; whatever comes + /// back rides along in the task's `SuccessfulTask` message to the + /// scheduler. Default returns empty — implementers with real plans + /// override to walk their operator tree. + fn collect_runtime_stats_reports( + &self, + ) -> Vec { + Vec::new() + } } /// Default execution engine using DataFusion's ShuffleWriterExec. @@ -290,6 +302,31 @@ impl QueryStageExecutor for DefaultQueryStageExec { ShuffleWriterVariant::Sort(writer) => utils::collect_plan_metrics(writer), } } + + fn collect_runtime_stats_reports( + &self, + ) -> Vec { + // Walk from the shuffle writer's plan through the whitelist. If + // no `RuntimeStatsExec` sits within reach, we return an empty + // Vec — the majority of plans (anything not on the parallel- + // window path today). Serialization errors are logged and the + // report dropped rather than failing the task; the task's data + // was already produced correctly, telemetry loss shouldn't tank + // the query. + let plan: Arc = match &self.shuffle_writer { + ShuffleWriterVariant::Passthrough(writer) => Arc::new(writer.clone()), + ShuffleWriterVariant::Sort(writer) => Arc::new(writer.clone()), + }; + match ballista_core::execution_plans::collect_runtime_stats_reports(&plan) { + Ok(reports) => reports, + Err(e) => { + log::warn!( + "collect_runtime_stats_reports failed, task will report empty stats: {e}" + ); + Vec::new() + } + } + } } /// Spawn K parallel `plan.execute(N, ctx)` calls against a shuffle writer, diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index e0eea30e41..2390364c1e 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -24,7 +24,7 @@ use crate::cpu_bound_executor::DedicatedExecutor; use crate::executor::Executor; use crate::executor_process::remove_job_data; -use crate::{TaskExecutionTimes, as_task_status}; +use crate::{TaskCompletionExtras, TaskExecutionTimes, as_task_status}; use ballista_core::JobId; use ballista_core::error::BallistaError; use ballista_core::extension::SessionConfigHelperExt; @@ -46,9 +46,16 @@ use std::error::Error; use std::sync::mpsc::{Receiver, Sender, TryRecvError}; use std::time::{Instant, SystemTime, UNIX_EPOCH}; use std::{sync::Arc, time::Duration}; -use tokio::sync::{OwnedSemaphorePermit, Semaphore}; +use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore}; use tonic::codegen::{Body, Bytes, StdError}; +/// Idle sleep between polls when polling is the only way to learn of new work. +const IDLE_POLL_INTERVAL: Duration = Duration::from_millis(50); + +/// Idle sleep when a `poll_now_notify` wake-up is wired and the timer is only +/// a fallback. +const NOTIFIED_IDLE_POLL_INTERVAL: Duration = Duration::from_secs(1); + /// Maximum time the poll loop waits for a free vcore before polling the /// scheduler anyway. `poll_work` doubles as the executor's heartbeat under /// pull-based scheduling, so a fully-busy executor must keep polling (reporting @@ -68,6 +75,12 @@ const HEARTBEAT_POLL_INTERVAL: Duration = Duration::from_secs(5); /// Pass `None` to have the loop create a semaphore sized to the executor's /// configured vcore count. /// +/// `poll_now_notify`, when provided, wakes an idle poll loop immediately +/// (typically wired to the scheduler's `on_work_available` callback) instead +/// of waiting out the idle interval. A notification sent mid-poll is not +/// lost: `Notify` stores the permit and the next `notified().await` returns +/// immediately. +/// /// **Shared semaphores**: when one semaphore is shared across loops that /// connect to different schedulers, each scheduler independently sees the /// current free capacity and may dispatch up to that many tasks. The semaphore @@ -89,6 +102,7 @@ pub async fn poll_loop, executor: Arc, codec: BallistaCodec, + poll_now_notify: Option>, free_vcores: Option>, health: crate::health::ExecutorHealth, ) -> Result<(), BallistaError> @@ -245,8 +259,8 @@ where executor.metadata.id.clone(), task.task_attempt_num as usize, task_key, - None, task_execution_times, + TaskCompletionExtras::default(), )) { warn!("failed to send task status: {error:?}"); }; @@ -263,7 +277,19 @@ where } if !active_job { - tokio::time::sleep(Duration::from_millis(50)).await; + match &poll_now_notify { + Some(notify) => { + tokio::select! { + () = tokio::time::sleep(NOTIFIED_IDLE_POLL_INTERVAL) => {} + () = notify.notified() => { + debug!("Received poll_now notification, polling immediately"); + } + } + } + None => { + tokio::time::sleep(IDLE_POLL_INTERVAL).await; + } + } } } } @@ -392,6 +418,7 @@ async fn run_received_task, BallistaError>>() .ok(); + let runtime_stats = query_stage_exec.collect_runtime_stats_reports(); let end_exec_time = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -409,8 +436,11 @@ async fn run_received_task( ); // 1. Start executor grpc service + // + // The listening socket is bound here rather than inside the spawned task, + // because step 2 registers with the scheduler and the scheduler dials this + // port back to check connectivity. Binding lazily inside the server future + // let that callback lose the race and get ECONNREFUSED, which fails + // registration and takes the executor down with it. let server = { let executor_meta = executor.metadata.clone(); let addr = format!("{}:{}", config.bind_host, executor_meta.grpc_port); let addr = addr.parse().unwrap(); let grpc_server_config = config.grpc_server_config.clone(); + let incoming = create_grpc_server_incoming(addr, &grpc_server_config)?; info!( "Ballista v{BALLISTA_VERSION} Rust Executor Grpc Server listening on {addr:?}" ); @@ -150,7 +159,7 @@ pub async fn startup( let shutdown_signal = grpc_shutdown.recv(); let grpc_server_future = create_grpc_server(&grpc_server_config) .add_service(server) - .serve_with_shutdown(addr, shutdown_signal); + .serve_with_incoming_shutdown(incoming, shutdown_signal); grpc_server_future.await.map_err(|e| { error!("Tonic error, Could not start Executor Grpc Server."); BallistaError::TonicError(e) @@ -159,7 +168,6 @@ pub async fn startup( }; // 2. Do executor registration - // TODO the executor registration should happen only after the executor grpc server started. let executor_server = Arc::new(executor_server); match register_executor(&mut scheduler, executor.clone()).await { Ok(_) => { @@ -451,6 +459,7 @@ impl ExecutorServer, BallistaError>>() .ok(); + let runtime_stats = exec.collect_runtime_stats_reports(); let executor_id = &self.executor.metadata.id; let end_exec_time = SystemTime::now() @@ -468,8 +477,11 @@ impl ExecutorServer ExecutorServer>, + /// Runtime-stats reports harvested from `RuntimeStatsExec` taps in the plan. + pub runtime_stats: Vec, +} + /// Converts a task execution result into a [`TaskStatus`] protobuf message. /// /// This function wraps the outcome of task execution (success or failure) @@ -108,16 +123,22 @@ pub fn as_task_status( executor_id: String, stage_attempt_num: usize, key: TaskKey, - operator_metrics: Option>, execution_times: TaskExecutionTimes, + extras: TaskCompletionExtras, ) -> TaskStatus { + let TaskCompletionExtras { + operator_metrics, + runtime_stats, + } = extras; let metrics = operator_metrics.unwrap_or_default(); let task_id = key.task_id; match execution_result { Ok(partitions) => { debug!( - "Task {task_id} finished with operator_metrics array size {}", - metrics.len() + "Task {task_id} finished with operator_metrics array size {} \ + and {} runtime-stats report(s)", + metrics.len(), + runtime_stats.len(), ); TaskStatus { task_id: task_id as u32, @@ -131,6 +152,7 @@ pub fn as_task_status( status: Some(task_status::Status::Successful(SuccessfulTask { executor_id, partitions, + runtime_stats, })), } } diff --git a/ballista/executor/src/standalone.rs b/ballista/executor/src/standalone.rs index fd9841a0e6..78e7042953 100644 --- a/ballista/executor/src/standalone.rs +++ b/ballista/executor/src/standalone.rs @@ -148,6 +148,7 @@ pub async fn new_standalone_executor_from_builder( executor, codec, None, + None, crate::health::ExecutorHealth::new(), )); Ok(()) diff --git a/ballista/scheduler/Cargo.toml b/ballista/scheduler/Cargo.toml index d2012b3469..6409a87d3d 100644 --- a/ballista/scheduler/Cargo.toml +++ b/ballista/scheduler/Cargo.toml @@ -49,7 +49,6 @@ substrait = ["dep:datafusion-substrait"] arrow-flight = { workspace = true } async-trait = { workspace = true } axum = "0.8.9" -tower-http = { version = "0.7", features = ["cors"] } ballista-core = { path = "../core", version = "54.0.0" } clap = { workspace = true, optional = true } dashmap = { workspace = true } @@ -74,6 +73,7 @@ tokio = { workspace = true, features = ["full"] } tokio-stream = { workspace = true, features = ["net"] } tonic = { workspace = true, features = ["router"] } tonic-prost = { workspace = true, optional = true } +tower-http = { version = "0.7", features = ["cors"] } tracing = { workspace = true, optional = true } tracing-appender = { workspace = true, optional = true } tracing-subscriber = { workspace = true, optional = true } @@ -84,6 +84,7 @@ name = "tpch_plan_stability" path = "tests/tpch_plan_stability/main.rs" [dev-dependencies] +datafusion-functions-aggregate-common = { workspace = true } regex = "1" rstest = { workspace = true } diff --git a/ballista/scheduler/src/cluster/memory.rs b/ballista/scheduler/src/cluster/memory.rs index 51dfa37dce..3baef1dddb 100644 --- a/ballista/scheduler/src/cluster/memory.rs +++ b/ballista/scheduler/src/cluster/memory.rs @@ -493,22 +493,20 @@ impl JobState for InMemoryJobState { async fn fail_unscheduled_job(&self, job_id: &JobId, reason: String) -> Result<()> { if let Some((job_id, (job_name, queued_at))) = self.queued_jobs.remove(job_id) { - self.completed_jobs.insert( - job_id.clone(), - ( - JobStatus { - job_id: job_id.into(), - job_name, - status: Some(Status::Failed(FailedJob { - error: reason, - queued_at, - started_at: 0, - ended_at: timestamp_millis(), - })), - }, - None, - ), - ); + let status = JobStatus { + job_id: job_id.clone().into(), + job_name, + status: Some(Status::Failed(FailedJob { + error: reason, + queued_at, + started_at: 0, + ended_at: timestamp_millis(), + })), + }; + self.completed_jobs + .insert(job_id.clone(), (status.clone(), None)); + self.job_event_sender + .send(&JobStateEvent::JobUpdated { job_id, status }); Ok(()) } else { @@ -533,6 +531,7 @@ mod test { use crate::test_utils::{ test_aggregation_plan, test_join_plan, test_two_aggregations_plan, }; + use ballista_core::JobId; use ballista_core::error::Result; use ballista_core::serde::protobuf::JobStatus; use ballista_core::serde::scheduler::{ @@ -609,6 +608,41 @@ mod test { Ok(()) } + #[tokio::test] + async fn test_in_memory_job_planning_failure_notification() -> Result<()> { + let state = InMemoryJobState::new( + "", + Arc::new(default_session_builder), + Arc::new(default_config_producer), + ); + let mut events = state.job_state_events().await?; + let job_id = JobId::from("job-1"); + + state.accept_job(&job_id, "", 0)?; + state + .fail_unscheduled_job(&job_id, "failed planning".to_owned()) + .await?; + + let event = + tokio::time::timeout(std::time::Duration::from_secs(1), events.next()) + .await + .expect("job state event should arrive") + .expect("job state event stream should remain open"); + + assert!(matches!( + event, + JobStateEvent::JobUpdated { + job_id: event_job_id, + status: JobStatus { + status: Some(ballista_core::serde::protobuf::job_status::Status::Failed(_)), + .. + }, + } if event_job_id == job_id + )); + + Ok(()) + } + #[tokio::test] async fn test_in_memory_job_notification() -> Result<()> { let state = InMemoryJobState::new( diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 940888714b..28320783f4 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -28,12 +28,43 @@ use crate::SessionBuilder; use crate::cluster::DistributionPolicy; use ballista_core::extension::EndpointOverrideFn; -use ballista_core::{ConfigProducer, config::TaskSchedulingPolicy}; +use ballista_core::{ConfigProducer, JobId, config::TaskSchedulingPolicy}; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::physical_plan::PhysicalExtensionCodec; use std::fmt::Display; use std::sync::Arc; +/// Why the scheduler believes new work has become available for executors. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WorkAvailableReason { + /// A job was submitted and its initial tasks are ready to be scheduled. + JobSubmitted { + /// Identifier of the submitted job. + job_id: JobId, + }, + /// Completed tasks resolved downstream stages of a job, and the tasks of + /// those stages are now schedulable. + NewStagesRunnable { + /// Identifier of the job that gained schedulable tasks. + job_id: JobId, + }, +} + +/// Callback invoked when new work becomes available for executors, e.g. to +/// wake idle pull-based executors via the poll loop's `poll_now_notify`. +/// +/// It fires only after the work is visible to a polling executor, so waking +/// one cannot race the scheduler's internal event processing. +/// +/// # Warning +/// +/// The callback runs synchronously inside the scheduler's main event loop. +/// Implementations **must be non-blocking**; offload blocking or long-running +/// work (such as network I/O) to a separate task or thread. +/// +/// `Arc` rather than `Box` so [`SchedulerConfig`] remains [`Clone`]. +pub type OnWorkAvailableFn = Arc; + /// Command-line configuration for the scheduler binary. #[cfg(feature = "build-binary")] #[derive(clap::Parser, Debug)] @@ -159,6 +190,16 @@ pub struct Config { help = "The maximum size of an encoded message at the grpc server side." )] pub grpc_server_max_encoding_message_size: u32, + /// Maximum size of messages sent by the scheduler's outbound gRPC clients + /// (e.g. task assignment to executors). Should be at least as large as the + /// executor's `--grpc-server-max-decoding-message-size` for those RPCs to + /// succeed with big encoded plans. + #[arg( + long, + default_value_t = 16777216, + help = "The maximum size of a message sent by the scheduler's outbound gRPC clients (in bytes)." + )] + pub grpc_client_max_message_size: u32, /// Timeout in seconds before marking an executor as dead. #[arg( long, @@ -173,6 +214,14 @@ pub struct Config { help = "Interval, in seconds, to check expired or dead executors." )] pub expire_dead_executor_interval_seconds: u64, + /// Grace period in seconds to wait for an executor to (re)appear after the + /// cluster has lost its last executor before failing the running jobs. + #[arg( + long, + default_value_t = 30, + help = "Grace period, in seconds, to wait for an executor to (re)register after the last executor is lost before failing running jobs. Prevents jobs from hanging forever when every executor dies, while still tolerating a transient total loss (e.g. a rolling restart). Set to 0 to fail as soon as the loss is observed." + )] + pub no_executors_grace_period_seconds: u64, /// Minimum number of registered executors before /readyz returns 200 #[arg( long, @@ -256,10 +305,18 @@ pub struct SchedulerConfig { pub grpc_server_max_decoding_message_size: u32, /// The maximum size of an encoded message at the grpc server side. pub grpc_server_max_encoding_message_size: u32, + /// The maximum size of a message sent by the scheduler's outbound gRPC clients. + pub grpc_client_max_message_size: u32, /// The executor timeout in seconds. It should be longer than executor's heartbeat intervals. pub executor_timeout_seconds: u64, /// The interval to check expired or dead executors pub expire_dead_executor_interval_seconds: u64, + /// Grace period in seconds to wait for an executor to (re)register after the + /// cluster has lost its last executor before failing the running jobs. This + /// bounds the otherwise-unbounded wait so that a total executor loss fails + /// the affected jobs instead of hanging forever. Set to 0 to fail as soon as + /// the loss is observed. + pub no_executors_grace_period_seconds: u64, /// [ConfigProducer] override option pub override_config_producer: Option, /// [SessionBuilder] override option @@ -287,6 +344,9 @@ pub struct SchedulerConfig { #[cfg(feature = "rest-api")] /// Comma-separated list of allowed methods for CORS pub cors_allowed_methods: String, + /// Callback invoked when new work becomes available for executors. + /// See [`OnWorkAvailableFn`]. + pub on_work_available: Option, } impl Default for SchedulerConfig { @@ -307,8 +367,10 @@ impl Default for SchedulerConfig { scheduler_event_expected_processing_duration: 0, grpc_server_max_decoding_message_size: 16777216, grpc_server_max_encoding_message_size: 16777216, + grpc_client_max_message_size: 16777216, executor_timeout_seconds: 180, expire_dead_executor_interval_seconds: 15, + no_executors_grace_period_seconds: 30, override_config_producer: None, override_session_builder: None, override_logical_codec: None, @@ -324,6 +386,7 @@ impl Default for SchedulerConfig { cors_allowed_origins: String::default(), #[cfg(feature = "rest-api")] cors_allowed_methods: String::default(), + on_work_available: None, } } } @@ -414,6 +477,13 @@ impl SchedulerConfig { self } + /// Sets the grace period, in seconds, to wait for an executor to (re)register + /// after the last executor is lost before failing running jobs. + pub fn with_no_executors_grace_period_seconds(mut self, value: u64) -> Self { + self.no_executors_grace_period_seconds = value; + self + } + /// Sets the maximum gRPC server decoding message size. pub fn with_grpc_server_max_decoding_message_size(mut self, value: u32) -> Self { self.grpc_server_max_decoding_message_size = value; @@ -426,6 +496,12 @@ impl SchedulerConfig { self } + /// Sets the maximum message size for the scheduler's outbound gRPC clients. + pub fn with_grpc_client_max_message_size(mut self, value: u32) -> Self { + self.grpc_client_max_message_size = value; + self + } + /// Sets a custom config producer. pub fn with_override_config_producer( mut self, @@ -459,6 +535,16 @@ impl SchedulerConfig { self.use_tls = use_tls; self } + + /// Sets the callback invoked when new work becomes available for + /// executors. See [`OnWorkAvailableFn`]. + pub fn with_on_work_available( + mut self, + on_work_available: OnWorkAvailableFn, + ) -> Self { + self.on_work_available = Some(on_work_available); + self + } } /// Policy of distributing tasks to available executor slots @@ -539,9 +625,11 @@ impl TryFrom for SchedulerConfig { .grpc_server_max_decoding_message_size, grpc_server_max_encoding_message_size: opt .grpc_server_max_encoding_message_size, + grpc_client_max_message_size: opt.grpc_client_max_message_size, executor_timeout_seconds: opt.executor_timeout_seconds, expire_dead_executor_interval_seconds: opt .expire_dead_executor_interval_seconds, + no_executors_grace_period_seconds: opt.no_executors_grace_period_seconds, override_config_producer: None, override_logical_codec: None, override_physical_codec: None, @@ -557,6 +645,7 @@ impl TryFrom for SchedulerConfig { cors_allowed_origins: opt.cors_allowed_origins, #[cfg(feature = "rest-api")] cors_allowed_methods: opt.cors_allowed_methods, + on_work_available: None, }; Ok(config) diff --git a/ballista/scheduler/src/physical_optimizer/join_selection.rs b/ballista/scheduler/src/physical_optimizer/join_selection.rs index 20d4649283..7e6ffb01c5 100644 --- a/ballista/scheduler/src/physical_optimizer/join_selection.rs +++ b/ballista/scheduler/src/physical_optimizer/join_selection.rs @@ -322,10 +322,10 @@ pub(crate) fn partitioned_hash_join( { hash_join.swap_inputs(PartitionMode::Partitioned) } else { - // Null-aware anti joins must use CollectLeft mode because they track probe-side state - // (probe_side_non_empty, probe_side_has_null) per-partition, but need global knowledge - // for correct null handling. With partitioning, a partition might not see probe rows - // even if the probe side is globally non-empty, leading to incorrect NULL row handling. + // Keep this copied rule aligned with upstream DataFusion. An Auto join + // that cannot be selected by size still requires CollectLeft when it is + // null-aware; Ballista's distributed planner later coalesces the probe + // side so this mode runs in exactly one task. let partition_mode = if hash_join.null_aware { PartitionMode::CollectLeft } else { @@ -368,7 +368,12 @@ fn statistical_join_selection_subrule( PartitionMode::Partitioned => { let left = hash_join.left(); let right = hash_join.right(); + // Keep this branch aligned with upstream DataFusion: a + // null-aware join is not swapped, but its partition mode is not + // changed here. Ballista's distributed planner is responsible + // for lowering it to a single-task CollectLeft join. if hash_join.join_type().supports_swap() + && !hash_join.null_aware && should_swap_join_order(&**left, &**right)? { hash_join @@ -583,6 +588,7 @@ pub fn hash_join_swap_subrule( if let Some(hash_join) = input.downcast_ref::() && hash_join.left.boundedness().is_unbounded() && !hash_join.right.boundedness().is_unbounded() + && !hash_join.null_aware && matches!( *hash_join.join_type(), JoinType::Inner | JoinType::Left | JoinType::LeftSemi | JoinType::LeftAnti @@ -772,6 +778,122 @@ mod test { ); } + #[test] + fn partitioned_null_aware_anti_join_is_not_swapped() { + use datafusion::{ + common::NullEquality, + physical_optimizer::PhysicalOptimizerRule, + physical_plan::joins::{HashJoinExec, PartitionMode}, + }; + + use crate::physical_optimizer::join_selection::JoinSelection; + + // The large left and small right sides would normally be swapped by + // statistical join selection. Keep this rule aligned with upstream + // DataFusion: a null-aware anti join retains both its orientation and + // its existing partition mode. Distributed lowering happens later. + let (big, small) = create_big_and_small(); + let join = Arc::new( + HashJoinExec::try_new( + Arc::clone(&big), + Arc::clone(&small), + vec![( + Arc::new(Column::new("big_col", 0)) as _, + Arc::new(Column::new("small_col", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(), + ) as Arc; + + let optimized = JoinSelection::new() + .optimize(join, &ConfigOptions::new()) + .unwrap(); + let hash_join = optimized + .downcast_ref::() + .expect("null-aware join should remain a HashJoinExec"); + + assert_eq!(*hash_join.join_type(), JoinType::LeftAnti); + assert_eq!(*hash_join.partition_mode(), PartitionMode::Partitioned); + assert!(hash_join.null_aware); + } + + #[test] + fn unbounded_input_rule_does_not_swap_null_aware_anti_join() { + use datafusion::{ + arrow::datatypes::SchemaRef, + common::NullEquality, + execution::{SendableRecordBatchStream, TaskContext}, + physical_plan::{ + EmptyRecordBatchStream, + joins::{HashJoinExec, PartitionMode}, + streaming::{PartitionStream, StreamingTableExec}, + }, + }; + + use crate::physical_optimizer::join_selection::hash_join_swap_subrule; + + #[derive(Debug)] + struct EmptyPartitionStream(SchemaRef); + + impl PartitionStream for EmptyPartitionStream { + fn schema(&self) -> &SchemaRef { + &self.0 + } + + fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { + Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.0))) + } + } + + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, false)])); + let left = Arc::new( + StreamingTableExec::try_new( + Arc::clone(&schema), + vec![Arc::new(EmptyPartitionStream(Arc::clone(&schema)))], + None, + vec![], + true, + None, + ) + .unwrap(), + ) as Arc; + let right = Arc::new(StatisticsExec::new( + Statistics::new_unknown(&schema), + schema.as_ref().clone(), + )) as Arc; + let join = Arc::new( + HashJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + vec![( + Arc::new(Column::new("key", 0)) as _, + Arc::new(Column::new("key", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(), + ) as Arc; + + let optimized = hash_join_swap_subrule(Arc::clone(&join), &ConfigOptions::new()) + .expect( + "the unbounded-input rule must not try to create a null-aware RightAnti", + ); + + assert!(Arc::ptr_eq(&optimized, &join)); + } + fn create_big_and_small() -> (Arc, Arc) { let big = Arc::new(StatisticsExec::new( big_statistics(), diff --git a/ballista/scheduler/src/planner.rs b/ballista/scheduler/src/planner.rs index 7283558991..0b3af8fc00 100644 --- a/ballista/scheduler/src/planner.rs +++ b/ballista/scheduler/src/planner.rs @@ -43,6 +43,7 @@ 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::union::UnionExec; use datafusion::physical_plan::{ ExecutionPlan, Partitioning, with_new_children_if_necessary, }; @@ -51,6 +52,7 @@ use log::debug; use crate::physical_optimizer::join_selection::{ collect_left_broadcast_safe, should_swap_join_order, }; +use crate::state::task_builder::restrict_plan_to_partitions; type PartialQueryStageResult = (Arc, Vec>); @@ -195,6 +197,45 @@ impl DefaultDistributedPlanner { stages.append(&mut child_stages); } + // UnionExec: a branch that cannot be restricted away needs its own + // stage. Per-task restriction gives each branch of a union the + // partitions that task owns and an empty slice to the rest, relying on + // an unowned branch becoming 0-partition so `UnionExec`'s index + // arithmetic routes past it. + // + // That fails when the branch's partition count does not come from + // restrictable leaves — a `CoalescePartitionsExec` or non-preserving + // `SortExec` reports one partition whatever its leaves do, and a + // broadcast `ShuffleReaderExec` is deliberately never pruned. The + // branch keeps reporting a partition and keeps producing all of its + // data, so every task runs every branch: results come back inflated by + // the branch count, or a task executes a branch whose scan was emptied + // (#2184). + // + // Cutting a boundary beneath such a branch turns it into a + // `ShuffleReaderExec`, which restricts to nothing cleanly, and leaves + // the collapse in its own stage where it sees its whole input. + if execution_plan.is::() { + for child in &mut children { + if can_stay_inline(child) { + continue; + } + let writer = create_shuffle_writer_with_config( + job_id, + self.next_stage_id(), + child.clone(), + None, + config, + )?; + *child = create_unresolved_shuffle(writer.as_ref()); + stages.push(writer); + } + return Ok(( + with_new_children_if_necessary(execution_plan, children)?, + stages, + )); + } + if let Some(_coalesce) = execution_plan.downcast_ref::() { let input = children[0].clone(); let input = self.optimizer_enforce_sorting.optimize(input, config)?; @@ -266,6 +307,107 @@ impl DefaultDistributedPlanner { self.next_stage_id } + /// Returns `Some(true/false)` when statistics can determine whether an + /// input fits Ballista's broadcast byte limit, or `None` when its size is + /// unknown. Falls back to a conservative row-width estimate when possible. + fn broadcast_size_under_threshold( + plan: &dyn ExecutionPlan, + threshold: usize, + ) -> Option { + let Ok(stats) = plan.partition_statistics(None) else { + debug!( + "broadcast check: partition_statistics returned error for {}", + plan.name() + ); + return None; + }; + debug!( + "broadcast check: {} total_byte_size={:?} num_rows={:?} threshold={}", + plan.name(), + stats.total_byte_size, + stats.num_rows, + threshold, + ); + if let Some(bytes) = stats.total_byte_size.get_value() + && *bytes != 0 + { + Some(*bytes < threshold) + } else if let Some(rows) = stats.num_rows.get_value() + && *rows != 0 + { + let schema = plan.schema(); + let bytes_per_row: usize = schema + .fields() + .iter() + .map(|f| match f.data_type() { + DataType::Boolean => 1, + DataType::Int8 | DataType::UInt8 => 1, + DataType::Int16 | DataType::UInt16 => 2, + DataType::Int32 | DataType::UInt32 | DataType::Float32 => 4, + DataType::Int64 | DataType::UInt64 | DataType::Float64 => 8, + DataType::Date32 => 4, + DataType::Date64 => 8, + DataType::Decimal128(_, _) => 16, + DataType::Decimal256(_, _) => 32, + _ => 32, // conservative estimate for variable-length types + }) + .sum(); + let estimated_bytes = *rows * bytes_per_row.max(8); + debug!( + "broadcast check: estimated {estimated_bytes} bytes ({rows} rows * {bytes_per_row} bytes/row from {} columns)", + schema.fields().len(), + ); + Some(estimated_bytes < threshold) + } else { + None + } + } + + /// Lowers a null-aware anti join to the only shape supported correctly by + /// DataFusion's in-process hash join: collect the build side and coalesce + /// the probe side so one task owns all shared null/visited state. + fn lower_null_aware_join( + hash_join: &HashJoinExec, + threshold_bytes: usize, + ) -> Result> { + // Unknown size is allowed because file scans commonly lose exact byte + // statistics before this point. Known oversized inputs and an explicit + // threshold of zero fail clearly instead of running an unsafe plan. + if threshold_bytes == 0 + || matches!( + Self::broadcast_size_under_threshold( + &**hash_join.left(), + threshold_bytes + ), + Some(false) + ) + { + return Err(BallistaError::General(format!( + "Null-aware anti join requires single-task execution, but its build side does not fit ballista.optimizer.broadcast_join_threshold_bytes ({threshold_bytes} bytes)" + ))); + } + + // Always keep an explicit coalesce. Some scans report one output + // partition during planning but expand to multiple partitions when the + // distributed stage is built. + let right: Arc = if hash_join + .right() + .downcast_ref::() + .is_some() + { + hash_join.right().clone() + } else { + Arc::new(CoalescePartitionsExec::new(hash_join.right().clone())) + }; + + hash_join + .builder() + .with_partition_mode(PartitionMode::CollectLeft) + .with_new_children(vec![hash_join.left().clone(), right])? + .build_exec() + .map_err(Into::into) + } + /// Reconciles a join's partition mode with the Ballista broadcast /// threshold (`ballista.optimizer.broadcast_join_threshold_bytes`). /// @@ -279,8 +421,9 @@ impl DefaultDistributedPlanner { /// under the Ballista threshold (including a threshold of `0`, which /// disables broadcast joins). This makes the Ballista key authoritative /// even when it is overridden at runtime below the DataFusion session - /// value. Null-aware anti joins are never demoted (they require - /// `CollectLeft`). + /// value. Null-aware anti joins are instead lowered to a single-task + /// `CollectLeft` join unless the build side is known to exceed the + /// Ballista threshold. /// /// Otherwise returns the input unchanged. fn maybe_promote_to_broadcast( @@ -303,6 +446,15 @@ impl DefaultDistributedPlanner { if let Some(hash_join) = plan.downcast_ref::() && *hash_join.partition_mode() == PartitionMode::CollectLeft { + // A null-aware anti join cannot run once per probe partition: + // DataFusion's visited/null state is shared only within one process. + // Collect the build side and coalesce the probe side so the join has + // exactly one task. Reject a build side known to exceed the normal + // broadcast threshold; an unsupported-plan error is safer than a + // wrong answer. + if hash_join.null_aware { + return Self::lower_null_aware_join(hash_join, threshold_bytes); + } // Broadcasting is only correct for probe-driven join types. If the // join type is not broadcast-safe, demote it back to a partitioned // (shuffle) join. Correctness guard, independent of the threshold. @@ -313,19 +465,19 @@ impl DefaultDistributedPlanner { ); return Self::demote_collect_left_to_partitioned(hash_join, config); } - // Null-aware anti joins must stay `CollectLeft`: they track - // probe-side state that a partitioned join cannot reconstruct, so - // never demote them regardless of the threshold. - if hash_join.null_aware { - return Ok(plan); - } // Safe join type: honor the Ballista broadcast threshold. DataFusion // decided `CollectLeft` using its own session threshold, which can // exceed a user's runtime `broadcast_join_threshold_bytes` override. // If broadcasts are disabled (0) or the build (left) side is not // under the Ballista threshold, demote so the Ballista key is // authoritative in the static planner path too. - if threshold_bytes == 0 || !under(&**hash_join.left(), threshold_bytes) { + if threshold_bytes == 0 + || !Self::broadcast_size_under_threshold( + &**hash_join.left(), + threshold_bytes, + ) + .unwrap_or(false) + { debug!( "broadcast check: demoting CollectLeft join to Partitioned; build side not under Ballista threshold={threshold_bytes} (or broadcasts disabled)", ); @@ -334,6 +486,16 @@ impl DefaultDistributedPlanner { return Ok(plan); } + // An already-partitioned null-aware join is not changed by the copied + // DataFusion optimizer rule. Lower it here, where Ballista can also + // enforce single-task distributed execution. + if let Some(hash_join) = plan.downcast_ref::() + && *hash_join.partition_mode() == PartitionMode::Partitioned + && hash_join.null_aware + { + return Self::lower_null_aware_join(hash_join, threshold_bytes); + } + if threshold_bytes == 0 { debug!("broadcast check: threshold is 0, broadcast disabled"); return Ok(plan); @@ -352,63 +514,13 @@ impl DefaultDistributedPlanner { if *hash_join.partition_mode() != PartitionMode::Partitioned { return Ok(plan); } - if hash_join.null_aware { - return Ok(plan); - } - let left = hash_join.left(); let right = hash_join.right(); - fn under(plan: &dyn ExecutionPlan, threshold: usize) -> bool { - let Ok(stats) = plan.partition_statistics(None) else { - debug!( - "broadcast check: partition_statistics returned error for {}", - plan.name() - ); - return false; - }; - debug!( - "broadcast check: {} total_byte_size={:?} num_rows={:?} threshold={}", - plan.name(), - stats.total_byte_size, - stats.num_rows, - threshold, - ); - if let Some(bytes) = stats.total_byte_size.get_value() { - *bytes != 0 && *bytes < threshold - } else if let Some(rows) = stats.num_rows.get_value() { - let schema = plan.schema(); - let bytes_per_row: usize = schema - .fields() - .iter() - .map(|f| { - match f.data_type() { - DataType::Boolean => 1, - DataType::Int8 | DataType::UInt8 => 1, - DataType::Int16 | DataType::UInt16 => 2, - DataType::Int32 | DataType::UInt32 | DataType::Float32 => 4, - DataType::Int64 | DataType::UInt64 | DataType::Float64 => 8, - DataType::Date32 => 4, - DataType::Date64 => 8, - DataType::Decimal128(_, _) => 16, - DataType::Decimal256(_, _) => 32, - _ => 32, // conservative estimate for variable-length types - } - }) - .sum(); - let estimated_bytes = *rows * bytes_per_row.max(8); - debug!( - "broadcast check: estimated {estimated_bytes} bytes ({rows} rows * {bytes_per_row} bytes/row from {} columns)", - schema.fields().len(), - ); - estimated_bytes != 0 && estimated_bytes < threshold - } else { - false - } - } - - let left_under = under(&**left, threshold_bytes); - let right_under = under(&**right, threshold_bytes); + let left_under = Self::broadcast_size_under_threshold(&**left, threshold_bytes) + .unwrap_or(false); + let right_under = Self::broadcast_size_under_threshold(&**right, threshold_bytes) + .unwrap_or(false); if !left_under && !right_under { debug!("broadcast check: neither side under threshold, skipping promotion"); return Ok(plan); @@ -532,6 +644,52 @@ impl DefaultDistributedPlanner { } } +/// Whether a union branch can be left in the union's own stage. +/// +/// A branch already sitting behind a stage boundary needs nothing more, and +/// neither does one that per-task restriction can reduce to zero partitions — +/// that is the property the union's index arithmetic relies on. The second +/// question is put to the real rewriter rather than to a list of operators, so +/// it stays correct as the rewriter's handling of collapses, broadcast readers +/// and leaf types evolves. +fn can_stay_inline(plan: &Arc) -> bool { + if plan.is::() { + return true; + } + !holds_a_broadcast(plan) && restricts_to_nothing(plan) +} + +/// Whether per-task restriction can reduce `plan` to zero partitions. An empty +/// slice means "this task polls nothing here", so a branch that still reports a +/// partition afterwards is one that cannot be restricted away. +fn restricts_to_nothing(plan: &Arc) -> bool { + restrict_plan_to_partitions(plan.clone(), &[]) + .is_ok_and(|p| p.properties().output_partitioning().partition_count() == 0) +} + +/// Whether `plan` contains a broadcast input or a `CollectLeft` join. +/// +/// Asking the rewriter alone is not enough here. A `CollectLeft` join's build +/// and probe sides can still be swapped after stage planning — `JoinSelection` +/// runs on the resolved stage plan (see `ExecutionStage`) — and the two orders +/// restrict differently: the build side is read in full (it must see its whole +/// input), while a broadcast input is never pruned at all. So a branch that +/// looks restrictable now can stop being restrictable by the time the task is +/// built, which is exactly the case that left TPC-DS q5's catalog branch +/// running in every task. +/// +/// Treat any such branch as needing its own stage. The cost is one extra +/// stage; the alternative is a silently inflated result. +fn holds_a_broadcast(plan: &Arc) -> bool { + let here = plan + .downcast_ref::() + .is_some_and(|u| u.broadcast) + || plan + .downcast_ref::() + .is_some_and(|j| *j.partition_mode() == PartitionMode::CollectLeft); + here || plan.children().iter().any(|c| holds_a_broadcast(c)) +} + fn create_unresolved_shuffle( shuffle_writer: &dyn ShuffleWriter, ) -> Arc { @@ -744,9 +902,10 @@ pub(crate) fn create_shuffle_writer_with_config( #[cfg(test)] mod test { + use super::{can_stay_inline, holds_a_broadcast}; use crate::assert_plan; use crate::planner::{DefaultDistributedPlanner, DistributedPlanner}; - use crate::test_utils::datafusion_test_context; + use crate::test_utils::{datafusion_test_context, scan_with_file_groups}; use ballista_core::error::BallistaError; use ballista_core::execution_plans::{SortShuffleWriterExec, UnresolvedShuffleExec}; use ballista_core::serde::BallistaCodec; @@ -755,6 +914,7 @@ mod test { use datafusion::execution::TaskContext; use datafusion::physical_expr::expressions::Column; use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; + use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::filter::FilterExec; @@ -806,6 +966,66 @@ mod test { ); } + /// A plain scan branch restricts away cleanly, so it can stay inline in + /// the union's stage and needs no extra shuffle. + #[test] + fn plain_union_branch_stays_inline() { + assert!(can_stay_inline(&scan_with_file_groups(4))); + } + + /// A branch already sitting behind a stage boundary needs nothing more. + #[test] + fn branch_behind_a_stage_boundary_stays_inline() { + let schema = scan_with_file_groups(1).schema(); + let reader: Arc = Arc::new(UnresolvedShuffleExec::new( + 1, + schema, + Partitioning::UnknownPartitioning(4), + )); + assert!(can_stay_inline(&reader)); + } + + /// A branch that collapses internally reports one partition whatever its + /// leaves do, so restriction cannot empty it and it needs its own stage. + #[test] + fn collapsing_union_branch_needs_its_own_stage() { + let branch: Arc = + Arc::new(CoalescePartitionsExec::new(scan_with_file_groups(4))); + assert_eq!( + branch.properties().output_partitioning().partition_count(), + 1 + ); + assert!( + !can_stay_inline(&branch), + "a collapsing branch cannot be restricted away" + ); + } + + /// The broadcast guard looks through the whole branch, not just its root: + /// a `CollectLeft` join's sides can be swapped after stage planning, and + /// the two orders restrict differently, so any branch holding one is + /// treated as needing its own stage. + #[test] + fn broadcast_is_detected_anywhere_in_a_branch() { + let plain = scan_with_file_groups(4); + assert!( + !holds_a_broadcast(&plain), + "a plain scan holds no broadcast" + ); + + let broadcast: Arc = + Arc::new(UnresolvedShuffleExec::new_broadcast(1, plain.schema(), 4)); + assert!(holds_a_broadcast(&broadcast)); + + // Buried a level down, which is where it actually shows up. + let buried: Arc = + Arc::new(CoalescePartitionsExec::new(broadcast)); + assert!( + holds_a_broadcast(&buried), + "the guard must search the branch, not just its root" + ); + } + macro_rules! downcast_exec { ($exec: expr, $ty: ty) => { ($exec.as_ref() as &dyn ExecutionPlan) @@ -1087,6 +1307,138 @@ order by Ok(()) } + #[test] + fn null_aware_collect_left_join_coalesces_probe_to_one_partition() { + use datafusion::{ + arrow::datatypes::{DataType, Field, Schema}, + common::{ + ColumnStatistics, JoinType, NullEquality, Statistics, stats::Precision, + }, + physical_plan::{ + Partitioning, coalesce_partitions::CoalescePartitionsExec, + joins::PartitionMode, repartition::RepartitionExec, + test::exec::StatisticsExec, + }, + }; + + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); + let stats = Statistics { + num_rows: Precision::Exact(20), + total_byte_size: Precision::Exact(80), + column_statistics: vec![ColumnStatistics::new_unknown()], + }; + let left = Arc::new(StatisticsExec::new(stats.clone(), schema.as_ref().clone())) + as Arc; + let right = Arc::new( + RepartitionExec::try_new( + Arc::new(StatisticsExec::new(stats, schema.as_ref().clone())), + Partitioning::RoundRobinBatch(4), + ) + .unwrap(), + ) as Arc; + let plan = Arc::new( + HashJoinExec::try_new( + left, + right, + vec![( + Arc::new(Column::new("key", 0)) as _, + Arc::new(Column::new("key", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(), + ) as Arc; + + let planned = DefaultDistributedPlanner::maybe_promote_to_broadcast( + plan, + &datafusion::config::ConfigOptions::new(), + ) + .unwrap(); + let hash_join = planned + .downcast_ref::() + .expect("null-aware join should remain a HashJoinExec"); + + assert_eq!(*hash_join.join_type(), JoinType::LeftAnti); + assert_eq!(*hash_join.partition_mode(), PartitionMode::CollectLeft); + assert!(hash_join.null_aware); + assert!( + hash_join + .right() + .downcast_ref::() + .is_some(), + "probe side must be coalesced so the join runs in one task" + ); + assert_eq!( + hash_join + .right() + .properties() + .output_partitioning() + .partition_count(), + 1 + ); + } + + #[test] + fn null_aware_join_rejects_known_oversized_build_side() { + use datafusion::{ + arrow::datatypes::{DataType, Field, Schema}, + common::{ + ColumnStatistics, JoinType, NullEquality, Statistics, stats::Precision, + }, + physical_plan::{joins::PartitionMode, test::exec::StatisticsExec}, + }; + + let schema = + Arc::new(Schema::new(vec![Field::new("key", DataType::Int32, true)])); + let left = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Exact(5_000_000), + total_byte_size: Precision::Exact(20 * 1024 * 1024), + column_statistics: vec![ColumnStatistics::new_unknown()], + }, + schema.as_ref().clone(), + )) as Arc; + let right = Arc::new(StatisticsExec::new( + Statistics::new_unknown(&schema), + schema.as_ref().clone(), + )) as Arc; + let plan = Arc::new( + HashJoinExec::try_new( + left, + right, + vec![( + Arc::new(Column::new("key", 0)) as _, + Arc::new(Column::new("key", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(), + ) as Arc; + + let error = DefaultDistributedPlanner::maybe_promote_to_broadcast( + plan, + &datafusion::config::ConfigOptions::new(), + ) + .unwrap_err(); + assert!( + error.to_string().contains( + "build side does not fit ballista.optimizer.broadcast_join_threshold_bytes" + ), + "{error}" + ); + } + #[tokio::test] async fn distributed_broadcast_join_plan() -> Result<(), BallistaError> { use datafusion::physical_plan::joins::PartitionMode; diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 66ed01b505..ca439dfeb2 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -31,7 +31,7 @@ use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::physical_plan::AsExecutionPlan; -use crate::cluster::BallistaCluster; +use crate::cluster::{BallistaCluster, ClusterStateEventStream, JobStateEventStream}; use crate::config::SchedulerConfig; use crate::metrics::SchedulerMetricsCollector; use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata}; @@ -183,6 +183,16 @@ impl SchedulerServer Result { + self.state.task_manager.job_state_events().await + } + + /// Returns a stream of cluster state events from the configured state backend. + pub async fn cluster_state_events(&self) -> Result { + self.state.executor_manager.cluster_state_events().await + } + /// True when at least `min_ready_executors` executors currently have /// live heartbeats. Embedders can call this from their own health/readiness /// handler and AND it with whatever app-specific state they track. The @@ -451,6 +461,7 @@ pub fn timestamp_millis() -> u64 { #[cfg(test)] mod test { use std::sync::Arc; + use std::time::Duration; use ballista_core::JobId; use ballista_core::extension::SessionConfigExt; @@ -463,7 +474,9 @@ mod test { use datafusion::test_util::scan_empty_with_partitions; use datafusion_proto::protobuf::LogicalPlanNode; use datafusion_proto::protobuf::PhysicalPlanNode; + use futures::StreamExt; + use crate::cluster::ClusterStateEvent; use crate::config::SchedulerConfig; use ballista_core::config::TaskSchedulingPolicy; use ballista_core::error::Result; @@ -487,6 +500,31 @@ mod test { assert_submitted_event, test_cluster_context, }; + #[tokio::test] + async fn test_scheduler_exposes_cluster_state_events() -> Result<()> { + let scheduler = test_scheduler(TaskSchedulingPolicy::PushStaged).await?; + let mut events = scheduler.cluster_state_events().await?; + let (executor_metadata, executor_data) = + test_executors(2).into_iter().next().unwrap(); + let executor_id = executor_metadata.id.clone(); + + scheduler + .state + .executor_manager + .register_executor(executor_metadata, executor_data) + .await?; + + let event = tokio::time::timeout(Duration::from_secs(5), events.next()) + .await + .expect("cluster state event should arrive"); + assert_eq!( + event, + Some(ClusterStateEvent::RegisteredExecutor { executor_id }) + ); + + Ok(()) + } + #[tokio::test] async fn test_pull_scheduling() -> Result<()> { let plan = test_plan(); @@ -569,6 +607,7 @@ mod test { status: Some(task_status::Status::Successful(SuccessfulTask { executor_id: "executor-1".to_owned(), partitions, + runtime_stats: vec![], })), }; diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index ee5174853a..ef3d118e81 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -16,17 +16,18 @@ // under the License. use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; use ballista_core::JobId; -use ballista_core::serde::protobuf::{FailedJob, JobStatus}; +use ballista_core::serde::protobuf::{FailedJob, JobStatus, job_status}; use log::{debug, error, info, trace, warn}; use ballista_core::error::{BallistaError, Result}; use ballista_core::event_loop::{EventAction, EventSender}; use tokio::sync::mpsc::error::TrySendError; -use crate::config::SchedulerConfig; +use crate::config::{SchedulerConfig, WorkAvailableReason}; use crate::metrics::SchedulerMetricsCollector; use crate::scheduler_server::timestamp_millis; use datafusion_proto::logical_plan::AsLogicalPlan; @@ -45,6 +46,11 @@ pub(crate) struct QueryStageScheduler< state: Arc>, metrics_collector: Arc, config: Arc, + /// Guards against arming more than one "all executors lost" grace timer at a + /// time. When a whole cluster dies at once the reaper posts an `ExecutorLost` + /// per executor, and each would otherwise arm its own timer and fail every + /// running job again. See + no_executor_check_pending: Arc, } impl QueryStageScheduler { @@ -57,6 +63,7 @@ impl QueryStageSchedul state, metrics_collector, config, + no_executor_check_pending: Arc::new(AtomicBool::new(false)), } } @@ -197,6 +204,12 @@ impl .post_event(QueryStageSchedulerEvent::ReviveOffers) .await?; } + + // The graph was revived before caching, so the job's tasks + // are already visible to polling executors. + if let Some(callback) = &self.config.on_work_available { + callback(WorkAvailableReason::JobSubmitted { job_id }); + } } QueryStageSchedulerEvent::JobPlanningFailed { job_id, @@ -258,8 +271,19 @@ impl } QueryStageSchedulerEvent::JobUpdated(job_id) => { debug!("Job updated, job_id: [{job_id}]"); - if let Err(e) = self.state.task_manager.update_job(&job_id).await { - error!("Fail to invoke update_job for job {job_id} due to {e:?}"); + match self.state.task_manager.update_job(&job_id).await { + Ok(new_tasks) => { + // update_job revived the graph: the new tasks are + // already visible to polling executors. + if new_tasks > 0 + && let Some(callback) = &self.config.on_work_available + { + callback(WorkAvailableReason::NewStagesRunnable { job_id }); + } + } + Err(e) => { + error!("Fail to invoke update_job for job {job_id} due to {e:?}"); + } } } QueryStageSchedulerEvent::JobCancel(job_id) => { @@ -339,6 +363,88 @@ impl error!("{msg}"); } } + + // If that was the last executor, the running jobs whose tasks were + // just reset can no longer make progress — there is nothing to + // schedule them onto. Rather than hang forever, wait a bounded + // grace period for an executor to (re)register (e.g. a rolling + // restart) and then fail any job still running on an empty cluster. + // Only fires for executors that were actually present, so jobs + // merely queued waiting for their first executor (autoscaling cold + // start) are never affected. + // See https://github.com/apache/datafusion-ballista/issues/2029 + // + // `no_executor_check_pending` collapses the burst of `ExecutorLost` + // events produced when a whole cluster dies at once into a single + // timer, so each running job is failed at most once. + if self.state.executor_manager.get_alive_executors().is_empty() + && !self.no_executor_check_pending.swap(true, Ordering::SeqCst) + { + let state = self.state.clone(); + let sender = event_sender.clone(); + let pending = self.no_executor_check_pending.clone(); + let grace = Duration::from_secs( + state.config.no_executors_grace_period_seconds, + ); + let lost_at = timestamp_millis(); + tokio::spawn(async move { + tokio::time::sleep(grace).await; + + // An executor may have (re)registered during the grace + // window; if so the reset tasks will be scheduled onto it + // and there is nothing to fail. + if state.executor_manager.get_alive_executors().is_empty() { + for job_id in + state.task_manager.get_running_job_cache().keys() + { + // Re-read the live status right before failing: a + // job that finished during the grace window must + // not be failed, and a job planned *after* the + // cluster went empty (started_at > lost_at) has its + // own window and must not inherit this one. + let queued_at = match state + .task_manager + .get_job_status(job_id) + .await + { + Ok(Some(JobStatus { + status: Some(job_status::Status::Running(running)), + .. + })) if running.started_at <= lost_at => { + running.queued_at + } + _ => continue, + }; + + let fail_message = format!( + "all executors were lost and no executor re-registered within {}s; no executors remain to run the tasks for this job", + grace.as_secs() + ); + warn!("Failing job {job_id}: {fail_message}"); + if let Err(e) = sender + .post_event( + QueryStageSchedulerEvent::JobRunningFailed { + job_id: job_id.clone(), + fail_message, + queued_at, + failed_at: timestamp_millis(), + }, + ) + .await + { + error!( + "Fail to post JobRunningFailed for job {job_id}: {e:?}" + ); + } + } + } + + // Cleared last, so the whole burst of `ExecutorLost` events + // that a simultaneous cluster death produces collapses into + // this single check — even when the grace period is 0. + pending.store(false, Ordering::SeqCst); + }); + } } QueryStageSchedulerEvent::CancelTasks(tasks) => { if let Err(e) = self @@ -376,16 +482,36 @@ impl #[cfg(test)] mod tests { - use crate::config::SchedulerConfig; - use crate::test_utils::{SchedulerTest, TestMetricsCollector, await_condition}; + use crate::cluster::JobStateEvent; + use crate::config::{SchedulerConfig, WorkAvailableReason}; + use crate::scheduler_server::SchedulerServer; + use crate::test_utils::{ + SchedulerTest, TestMetricsCollector, await_condition, test_cluster_context, + }; + use ballista_core::BALLISTA_PROTOCOL_VERSION; use ballista_core::config::TaskSchedulingPolicy; use ballista_core::error::Result; + use ballista_core::extension::SessionConfigExt; + use ballista_core::serde::BallistaCodec; + use ballista_core::serde::protobuf::job_status; + use ballista_core::serde::protobuf::scheduler_grpc_server::SchedulerGrpc; + use ballista_core::serde::protobuf::{ + ExecutorRegistration, PollWorkParams, ShuffleWritePartition, SuccessfulTask, + TaskStatus, task_status, + }; + use ballista_core::serde::scheduler::{ + ExecutorOperatingSystemSpecification, ExecutorSpecification, + }; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::functions_aggregate::sum::sum; use datafusion::logical_expr::{LogicalPlan, col}; + use datafusion::prelude::SessionConfig; use datafusion::test_util::scan_empty_with_partitions; - use std::sync::Arc; + use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode}; + use futures::StreamExt; + use std::sync::{Arc, Mutex}; use std::time::Duration; + use tonic::Request; #[tokio::test] async fn test_pending_job_metric() -> Result<()> { @@ -441,6 +567,285 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_scheduler_exposes_job_state_events() -> Result<()> { + let plan = test_plan(1); + let metrics_collector = Arc::new(TestMetricsCollector::default()); + let mut test = SchedulerTest::new( + SchedulerConfig::default() + .with_scheduler_policy(TaskSchedulingPolicy::PushStaged), + metrics_collector, + 1, + 1, + None, + ) + .await?; + + let mut events = test.job_state_events().await?; + let (_, job_id) = test.run("", &plan).await?; + let received_success = tokio::time::timeout(Duration::from_secs(5), async { + while let Some(event) = events.next().await { + if matches!( + event, + JobStateEvent::JobUpdated { + job_id: event_job_id, + status, + } if event_job_id == job_id + && matches!( + status.status, + Some(job_status::Status::Successful(_)) + ) + ) { + return true; + } + } + false + }) + .await + .expect("successful job state event should arrive"); + assert!( + received_success, + "job state event stream closed unexpectedly" + ); + Ok(()) + } + #[tokio::test] + async fn test_on_work_available_callback() -> Result<()> { + let reasons: Arc>> = Arc::default(); + let captured = reasons.clone(); + + let mut scheduler: SchedulerServer = + SchedulerServer::new( + "localhost:50050".to_owned(), + test_cluster_context(), + BallistaCodec::default(), + Arc::new( + SchedulerConfig { + scheduling_policy: TaskSchedulingPolicy::PullStaged, + ..Default::default() + } + .with_on_work_available(Arc::new( + move |reason| { + captured.lock().unwrap().push(reason); + }, + )), + ), + Arc::new(TestMetricsCollector::default()), + ); + scheduler.init().await?; + + let ctx = scheduler + .state + .session_manager + .create_or_update_session( + "session", + &SessionConfig::new_with_ballista().with_target_partitions(2), + ) + .await?; + let job_id = scheduler.submit_job("", ctx, &test_plan(2), None).await?; + + // Job submission runs asynchronously through the event loop. + let submitted = await_condition(Duration::from_millis(10), 100, || { + futures::future::ready(Ok(!reasons.lock().unwrap().is_empty())) + }) + .await?; + assert!(submitted, "JobSubmitted callback never fired"); + assert_eq!( + reasons.lock().unwrap().first(), + Some(&WorkAvailableReason::JobSubmitted { + job_id: job_id.clone() + }) + ); + + // Pull the shuffle stage's tasks; the callback promised they are + // visible by the time it fired. + let exec_meta = ExecutorRegistration { + id: "executor-1".to_owned(), + host: Some("localhost".to_owned()), + port: 50051, + grpc_port: 50052, + specification: Some(ExecutorSpecification::default().with_vcores(2).into()), + os_info: Some(ExecutorOperatingSystemSpecification::default().into()), + ballista_protocol_version: BALLISTA_PROTOCOL_VERSION, + }; + let polled = scheduler + .poll_work(Request::new(PollWorkParams { + metadata: Some(exec_meta.clone()), + num_free_vcores: 2, + task_status: vec![], + })) + .await + .expect("poll_work failed") + .into_inner(); + assert!( + !polled.tasks.is_empty(), + "expected tasks after JobSubmitted" + ); + + // Report the pulled tasks as successful; each writes the plan's two + // shuffle output partitions. + let task_status = polled + .tasks + .iter() + .map(|task| TaskStatus { + task_id: task.task_id, + job_id: task.job_id.clone(), + stage_id: task.stage_id, + stage_attempt_num: task.stage_attempt_num, + launch_time: 0, + start_exec_time: 0, + end_exec_time: 0, + metrics: vec![], + status: Some(task_status::Status::Successful(SuccessfulTask { + executor_id: exec_meta.id.clone(), + partitions: (0..2) + .map(|partition_id| ShuffleWritePartition { + partition_id, + num_batches: 1, + num_rows: 1, + num_bytes: 1, + file_id: None, + is_sort_shuffle: false, + }) + .collect(), + runtime_stats: vec![], + })), + }) + .collect(); + scheduler + .poll_work(Request::new(PollWorkParams { + metadata: Some(exec_meta), + num_free_vcores: 2, + task_status, + })) + .await + .expect("poll_work with task status failed"); + + // Completing the shuffle stage resolves the final stage. + let resolved = await_condition(Duration::from_millis(10), 100, || { + futures::future::ready(Ok(reasons.lock().unwrap().contains( + &WorkAvailableReason::NewStagesRunnable { + job_id: job_id.clone(), + }, + ))) + }) + .await?; + assert!( + resolved, + "expected NewStagesRunnable, got {:?}", + reasons.lock().unwrap() + ); + + Ok(()) + } + + #[tokio::test] + async fn test_running_job_fails_when_all_executors_are_lost() -> Result<()> { + let plan = test_plan(10); + + let metrics_collector = Arc::new(TestMetricsCollector::default()); + + // Grace period of 0 so the job is failed as soon as the loss is observed, + // keeping the test fast. + let mut test = SchedulerTest::new( + SchedulerConfig::default() + .with_scheduler_policy(TaskSchedulingPolicy::PushStaged) + .with_no_executors_grace_period_seconds(0), + metrics_collector.clone(), + 1, + 1, + None, + ) + .await?; + + let job_id = test.submit("", &plan).await?; + + // Wait until the job is actually running with tasks in flight. We + // deliberately never `tick()`, so its tasks never complete. + let job_id_ref = &job_id; + let test_ref = &test; + let running = await_condition(Duration::from_millis(50), 40, || async move { + let status = test_ref.job_status(job_id_ref).await?; + Ok(matches!( + status.and_then(|s| s.status), + Some(job_status::Status::Running(_)) + )) + }) + .await?; + assert!(running, "job should reach the running state"); + + // The only executor is lost. With no executors left, the reset tasks can + // never be scheduled, so the job must fail rather than hang forever + // (#2029). + test.lose_executor("virtual-executor-0").await?; + + let failed = await_condition(Duration::from_millis(100), 50, || async move { + let status = test_ref.job_status(job_id_ref).await?; + Ok(matches!( + status.and_then(|s| s.status), + Some(job_status::Status::Failed(_)) + )) + }) + .await?; + assert!( + failed, + "job should be failed after all executors were lost, but status was {:?}", + test.job_status(&job_id).await? + ); + + Ok(()) + } + + #[tokio::test] + async fn test_running_job_survives_partial_executor_loss() -> Result<()> { + let plan = test_plan(10); + + let metrics_collector = Arc::new(TestMetricsCollector::default()); + + let mut test = SchedulerTest::new( + SchedulerConfig::default() + .with_scheduler_policy(TaskSchedulingPolicy::PushStaged) + .with_no_executors_grace_period_seconds(0), + metrics_collector.clone(), + 2, + 1, + None, + ) + .await?; + + let job_id = test.submit("", &plan).await?; + + let job_id_ref = &job_id; + let test_ref = &test; + let running = await_condition(Duration::from_millis(50), 40, || async move { + let status = test_ref.job_status(job_id_ref).await?; + Ok(matches!( + status.and_then(|s| s.status), + Some(job_status::Status::Running(_)) + )) + }) + .await?; + assert!(running, "job should reach the running state"); + + // Lose only one of two executors. One remains alive, so the job must not + // be failed by the total-loss guard. + test.lose_executor("virtual-executor-0").await?; + + // Give the (grace-0) failure path ample time to fire if it were going to. + tokio::time::sleep(Duration::from_millis(500)).await; + + let status = test.job_status(&job_id).await?; + assert!( + !matches!( + status.as_ref().and_then(|s| s.status.clone()), + Some(job_status::Status::Failed(_)) + ), + "job must not be failed while an executor remains, but status was {status:?}" + ); + + Ok(()) + } + fn test_plan(partitions: usize) -> LogicalPlan { let schema = Schema::new(vec![ Field::new("id", DataType::Utf8, false), diff --git a/ballista/scheduler/src/state/aqe/adapter.rs b/ballista/scheduler/src/state/aqe/adapter.rs index 2040982341..539b45c6b8 100644 --- a/ballista/scheduler/src/state/aqe/adapter.rs +++ b/ballista/scheduler/src/state/aqe/adapter.rs @@ -19,7 +19,9 @@ use crate::planner::create_shuffle_writer_with_config; use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; use crate::state::aqe::planner::AdaptiveStageInfo; use ballista_core::JobId; -use ballista_core::execution_plans::ShuffleReaderExec; +use ballista_core::execution_plans::{ + PerPartitionFilterExec, ShuffleReaderExec, range_partition_predicates, +}; use datafusion::common::exec_err; use datafusion::config::ConfigOptions; use datafusion::error::DataFusionError; @@ -28,6 +30,7 @@ use datafusion::{ common::tree_node::{Transformed, TreeNode}, physical_plan::ExecutionPlan, }; +use log::debug; use std::sync::Arc; #[derive(Debug, Clone, Default)] @@ -105,7 +108,24 @@ impl BallistaAdapter { )?, }; - Ok(Transformed::yes(Arc::new(reader))) + let reader: Arc = Arc::new(reader); + // Without a per-partition filter, straddling sub-parts from a + // range-repartitioned upstream would feed multiple downstream + // partitions and `FinalPartitioned` would split their partial sums. + if let Some(routing) = exchange.range_repartition_routing() { + let predicates = + range_partition_predicates(routing.routing_expr, &routing.cuts); + debug!( + "range-repartition: injecting PerPartitionFilterExec above \ + ShuffleReader for stage {} — {} predicates over {} cuts", + stage_id, + predicates.len(), + routing.cuts.len(), + ); + let filtered = PerPartitionFilterExec::try_new(reader, predicates)?; + return Ok(Transformed::yes(Arc::new(filtered))); + } + Ok(Transformed::yes(reader)) } else { Ok(Transformed::no(plan)) } 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 40ab7e7854..bec08bfc69 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/dynamic_join.rs @@ -19,7 +19,10 @@ use ballista_core::config::BallistaConfig; use datafusion::{ arrow::compute::SortOptions, arrow::datatypes::{DataType, Schema}, - common::{ColumnStatistics, JoinType, NullEquality, Result, exec_err, internal_err}, + common::{ + ColumnStatistics, JoinType, NullEquality, Result, exec_err, internal_err, + plan_err, + }, config::ConfigOptions, execution::{SendableRecordBatchStream, TaskContext}, physical_expr_common::physical_expr::fmt_sql, @@ -248,10 +251,13 @@ impl DynamicJoinSelectionExec { // actually builds from rather than `self.left` unconditionally. // `supports_swap_join_order` is true when the *left* is the larger side, // so a swap moves the build onto `self.right`. - let swap_inputs = SelectJoinRule::supports_swap_join_order( - self.left.as_ref(), - self.right.as_ref(), - )?; + // Null-aware anti joins are only valid as LeftAnti and therefore cannot + // participate in the size-driven input swap. + let swap_inputs = !self.null_aware + && SelectJoinRule::supports_swap_join_order( + self.left.as_ref(), + self.right.as_ref(), + )?; let build_side = if swap_inputs { &self.right } else { &self.left }; let build_max_partition_bytes = max_per_partition_build_bytes(build_side); @@ -284,15 +290,41 @@ impl DynamicJoinSelectionExec { self.join_type }; - let partition_mode = - if under_threshold && collect_left_broadcast_safe(build_side_join_type) { - PartitionMode::CollectLeft - } else { - PartitionMode::Partitioned - }; - let stats_left = self.left.partition_statistics(None)?; let stats_right = self.right.partition_statistics(None)?; + let build_stats = if swap_inputs { + &stats_right + } else { + &stats_left + }; + let build_size_known = match build_stats.total_byte_size.get_value() { + Some(bytes) => *bytes != 0, + None => build_stats + .num_rows + .get_value() + .is_some_and(|rows| *rows != 0), + }; + + // A null-aware anti join must collect its build side and run its probe + // side in one task. Reject a disabled broadcast or a build side known + // to exceed the threshold. Unknown sizes are allowed because file scans + // commonly lose exact statistics before AQE first resolves the join. + if self.null_aware + && (threshold_collect_left_join_bytes == 0 + || (build_size_known && !under_threshold)) + { + return plan_err!( + "Null-aware anti join requires single-task execution, but its build side does not fit ballista.optimizer.broadcast_join_threshold_bytes ({threshold_collect_left_join_bytes} bytes)" + ); + } + + let partition_mode = if self.null_aware + || (under_threshold && collect_left_broadcast_safe(build_side_join_type)) + { + PartitionMode::CollectLeft + } else { + PartitionMode::Partitioned + }; let action = match (&self.selection_state, partition_mode) { (JoinInputState::Unknown, PartitionMode::CollectLeft) => self diff --git a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs index 63cc181502..abb33a8414 100644 --- a/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs +++ b/ballista/scheduler/src/state/aqe/execution_plan/exchange.rs @@ -19,6 +19,7 @@ use ballista_core::execution_plans::{ CoalescePlan, stats_for_partition, stats_for_partitions, }; use ballista_core::serde::scheduler::PartitionLocation; +use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::Statistics; use datafusion::{ error::{DataFusionError, Result}, @@ -32,6 +33,25 @@ use parking_lot::Mutex; use std::ops::Deref; use std::sync::{Arc, atomic::AtomicI64}; +/// Range-partition boundaries recovered from an +/// `UnorderedRangeRepartitionExec` / `OrderedRangeRepartitionExec` upstream +/// of this exchange. Written after the range-repartition-producing stage +/// completes and its runtime-stats sketches are merged; read at +/// task-specialization time to build per-downstream-partition range filters +/// (see `PerPartitionFilterExec`). +/// +/// `cuts` are `K - 1` monotone `f64` boundaries expressed in the value space +/// of `routing_expr`; downstream partition `k` owns `[cuts[k-1], cuts[k])` +/// with virtual `-∞`/`+∞` sentinels on the ends (matching the range +/// repartition's write-side convention). `routing_expr` is the same +/// expression the range repartition routes on so the filter is symmetric +/// with the writer's placement decision. +#[derive(Clone, Debug)] +pub struct RangeRepartitionRouting { + pub cuts: Vec, + pub routing_expr: Arc, +} + /// Execution plan representing an exchange/shuffle boundary used by the /// scheduler during adaptive query execution (AQE). /// @@ -74,6 +94,15 @@ pub struct ExchangeExec { /// transform-rebuilt parent chains. Same pattern as `shuffle_partitions`. coalesce: Arc>>>, + /// Range-partition boundaries recovered at runtime from an upstream + /// range-repartition op (URRE or ORRE). Stored when + /// the range-repartition-producing stage completes and its per-sub-part + /// quantile sketches have been merged. Read at task-specialization time + /// to build `PerPartitionFilterExec` predicates for downstream stage `N+1`. + /// + /// `None` on any exchange that isn't downstream of a range repartition + range_repartition_routing: Arc>>, + /// this disables stage from running even it would be suitable to run. /// /// the main reason for this property this is to allow rules to override @@ -103,6 +132,7 @@ impl ExchangeExec { Arc::new(AtomicI64::new(-1)), Arc::new(Mutex::new(None)), Arc::new(Mutex::new(None)), + Arc::new(Mutex::new(None)), false, false, ) @@ -120,6 +150,7 @@ impl ExchangeExec { Arc::new(AtomicI64::new(-1)), Arc::new(Mutex::new(None)), Arc::new(Mutex::new(None)), + Arc::new(Mutex::new(None)), true, false, ) @@ -133,6 +164,7 @@ impl ExchangeExec { self.stage_id.clone(), self.shuffle_partitions.clone(), self.coalesce.clone(), + self.range_repartition_routing.clone(), true, self.inactive_stage, ) @@ -149,6 +181,7 @@ impl ExchangeExec { stage_id: Arc, stage_partitions: Arc>>>>, coalesce: Arc>>>, + range_repartition_routing: Arc>>, broadcast: bool, inactive_stage: bool, ) -> Self { @@ -173,6 +206,7 @@ impl ExchangeExec { shuffle_partitions: stage_partitions, partitioning, coalesce, + range_repartition_routing, inactive_stage, broadcast, } @@ -257,6 +291,24 @@ impl ExchangeExec { pub fn coalesce(&self) -> Option> { self.coalesce.lock().clone() } + + /// Publishes range-repartition-recovered range boundaries on this + /// exchange. Called from + /// `AdaptiveExecutionGraph::maybe_range_repartition_overlap_remap` once + /// the upstream range-repartition stage completes and its per-sub-part + /// quantile sketches have been merged into `K - 1` monotone cuts. + /// Idempotent overwrite matches the `set_coalesce` pattern. + pub fn resolve_range_repartition_routing(&self, routing: RangeRepartitionRouting) { + self.range_repartition_routing.lock().replace(routing); + } + + /// Returns the range-repartition routing info if + /// `resolve_range_repartition_routing` has fired. Consumers use + /// `Some(_)` as the signal that this exchange is downstream of a range + /// repartition and its tasks need per-partition range filters. + pub fn range_repartition_routing(&self) -> Option { + self.range_repartition_routing.lock().clone() + } } impl DisplayAs for ExchangeExec { @@ -288,6 +340,9 @@ impl DisplayAs for ExchangeExec { cp.upstream_partition_count, )?; } + if let Some(r) = self.range_repartition_routing.lock().as_ref() { + write!(f, ", range_repartition_cuts={}", r.cuts.len())?; + } if self.broadcast { write!(f, ", broadcast=true",)? } @@ -354,6 +409,7 @@ impl ExecutionPlan for ExchangeExec { // doesn't lose the rule's decision. self.shuffle_partitions.clone(), self.coalesce.clone(), + self.range_repartition_routing.clone(), self.broadcast, self.inactive_stage, ); @@ -421,3 +477,94 @@ impl ExecutionPlan for ExchangeExec { } } } + +#[cfg(test)] +mod range_repartition_routing_tests { + //! `RangeRepartitionRouting` parking on `ExchangeExec`. The AQE hook + //! writes here at range-repartition-stage completion; task + //! specialization reads it back at + //! `BallistaAdapter::transform_children` time to wrap the + //! ShuffleReader in a `PerPartitionFilterExec`. Neither side is + //! exercised end-to-end without the URRE-inserting rule (a follow-up + //! PR), so tests here cover the slot itself: roundtrip through + //! `resolve_range_repartition_routing` → `range_repartition_routing()`, + //! and preservation across `with_new_children`. + use super::*; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_plan::expressions::col; + use datafusion::physical_plan::{ExecutionPlan, Partitioning}; + + fn v_source() -> Arc { + let schema = + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])); + let memory = + Arc::new(MemorySourceConfig::try_new(&[vec![]], schema, None).unwrap()); + Arc::new(DataSourceExec::new(memory)) + } + + fn v_routing_expr() -> Arc { + let schema = v_source().schema(); + col("v", schema.as_ref()).unwrap() + } + + fn sample_routing() -> RangeRepartitionRouting { + RangeRepartitionRouting { + cuts: vec![10.0, 20.0, 30.0], + routing_expr: v_routing_expr(), + } + } + + #[test] + fn range_repartition_routing_unresolved_returns_none() { + let exchange = ExchangeExec::new(v_source(), None, 42); + assert!(exchange.range_repartition_routing().is_none()); + } + + #[test] + fn resolve_range_repartition_routing_roundtrips() { + let exchange = ExchangeExec::new(v_source(), None, 42); + exchange.resolve_range_repartition_routing(sample_routing()); + let recovered = exchange + .range_repartition_routing() + .expect("routing must be Some after resolve"); + assert_eq!(recovered.cuts, vec![10.0, 20.0, 30.0]); + } + + #[test] + fn resolve_range_repartition_routing_overwrites_prior_value() { + // Idempotent-overwrite semantics match `set_coalesce`. + let exchange = ExchangeExec::new(v_source(), None, 42); + exchange.resolve_range_repartition_routing(sample_routing()); + exchange.resolve_range_repartition_routing(RangeRepartitionRouting { + cuts: vec![100.0], + routing_expr: v_routing_expr(), + }); + let recovered = exchange.range_repartition_routing().unwrap(); + assert_eq!(recovered.cuts, vec![100.0], "second resolve wins"); + } + + /// `with_new_children` must carry the routing slot through: transform + /// passes that rebuild the parent chain would otherwise silently drop + /// range boundaries the scheduler already parked here. + #[test] + fn with_new_children_preserves_range_repartition_routing() { + let partitioning = Some(Partitioning::UnknownPartitioning(4)); + let exchange = Arc::new(ExchangeExec::new(v_source(), partitioning, 42)); + exchange.resolve_range_repartition_routing(sample_routing()); + + // Rebuild with a fresh (equivalent-schema) child. + let rebuilt = exchange + .clone() + .with_new_children(vec![v_source()]) + .unwrap(); + let rebuilt_exchange = rebuilt + .downcast_ref::() + .expect("with_new_children must return an ExchangeExec"); + let recovered = rebuilt_exchange + .range_repartition_routing() + .expect("routing must survive with_new_children"); + assert_eq!(recovered.cuts, vec![10.0, 20.0, 30.0]); + } +} diff --git a/ballista/scheduler/src/state/aqe/mod.rs b/ballista/scheduler/src/state/aqe/mod.rs index deec100639..0c133db640 100644 --- a/ballista/scheduler/src/state/aqe/mod.rs +++ b/ballista/scheduler/src/state/aqe/mod.rs @@ -18,6 +18,7 @@ use crate::display::print_stage_metrics; use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::scheduler_server::timestamp_millis; +use crate::state::aqe::execution_plan::RangeRepartitionRouting; use crate::state::aqe::planner::AdaptivePlanner; use crate::state::execution_graph::{ ExecutionGraph, ExecutionGraphBox, ExecutionStage, ResolvedStage, RunningTaskInfo, @@ -27,7 +28,9 @@ use crate::state::execution_stage::RunningStage; use crate::state::task_manager::UpdatedStages; use ballista_core::JobId; use ballista_core::error::BallistaError; -use ballista_core::execution_plans::ShuffleWriter; +use ballista_core::execution_plans::{ + ShuffleWriter, cut_partitions, merge_runtime_stats_reports, repartition_routing_expr, +}; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{ @@ -283,6 +286,72 @@ impl AdaptiveExecutionGraph { Ok(events) } + /// Recover the range-repartition routing (cuts + routing expression) + /// from a completed stage's merged runtime-stats sketches. Caller must + /// have already established via `range_repartition_routing_expr` that + /// the stage's plan warrants routing, and passes the recovered expr in. + /// + /// `Ok(None)` means the stage produced no rows (nothing to route + /// through — passthrough is safe). `Err` means the stage's plan says + /// it should route but something went wrong recovering the cuts — + /// an invariant break, not a soft fallback (would misroute real data + /// downstream). + fn repartition_routing( + running_stage: &RunningStage, + routing_expr: Arc, + ) -> ballista_core::error::Result> { + let stage_id = running_stage.stage_id; + if running_stage.runtime_stats_reports.is_empty() { + return Err(BallistaError::General(format!( + "range-repartition stage {stage_id}: no runtime-stats reports" + ))); + } + let reports = running_stage + .runtime_stats_reports + .iter() + .map(|task| task.report.clone()) + .collect::>(); + let merged = merge_runtime_stats_reports(&reports).map_err(|err| { + BallistaError::General(format!( + "range-repartition stage {stage_id}: merge_reports failed: {err}" + )) + })?; + // A range-repartition stage produces exactly one report group (its + // single routing expression); any other count is a shape bug. + let entry = match merged.as_slice() { + [entry] => entry, + groups => { + return Err(BallistaError::General(format!( + "range-repartition stage {stage_id}: merge_reports returned {} groups, expected 1", + groups.len() + ))); + } + }; + if entry.total_rows == 0 { + debug!("range-repartition stage {stage_id}: no rows produced, passthrough"); + return Ok(None); + } + if entry.partition_count < 2 { + // K=1: single output partition — no cuts needed, no routing to + // recover. Everything flows to the one downstream partition. + debug!( + "range-repartition stage {stage_id}: K={}, no routing needed", + entry.partition_count + ); + return Ok(None); + } + if entry.cuts.is_empty() { + return Err(BallistaError::General(format!( + "range-repartition stage {stage_id}: {} rows, K={}, but no cuts (sketch missed)", + entry.total_rows, entry.partition_count + ))); + } + Ok(Some(RangeRepartitionRouting { + cuts: entry.cuts.clone(), + routing_expr, + })) + } + /// Return a Vec of stages to cancel fn update_stage_progress( &mut self, @@ -294,7 +363,29 @@ impl AdaptiveExecutionGraph { .update_exchange_locations(stage_id, locations)?; if is_completed { - let locations = self.planner.finalise_stage(stage_id)?; + let partitions = self.planner.take_stage_output_partitions(stage_id)?; + + // Range-repartition stages need overlap-based remap of their partitions + let maybe_stage = self.stages.get(&stage_id); + let partitions = if let Some(ExecutionStage::Running(stage)) = maybe_stage + && let Some(routing_expr) = repartition_routing_expr(stage.plan.as_ref())? + && let Some(routing) = Self::repartition_routing(stage, routing_expr)? + { + let reports = &stage.runtime_stats_reports; + let remapped = cut_partitions(partitions, reports, &routing.cuts) + .map_err(|err| { + BallistaError::General(format!( + "range-repartition stage {stage_id}: overlap remap failed: {err}" + )) + })?; + // Save boundaries to ExchangeExec so they are there for resolve_stage_partitions + self.planner.set_repartition_routing(stage_id, routing)?; + remapped + } else { + partitions + }; + self.planner + .resolve_stage_partitions(stage_id, partitions.clone())?; let (runnable, stages_to_cancel) = self.planner.actionable_stages()?; @@ -316,7 +407,7 @@ impl AdaptiveExecutionGraph { } else { // There is no more tasks to run // we update output locations - self.output_locations = locations.into_iter().flatten().collect(); + self.output_locations = partitions.into_iter().flatten().collect(); } // marking stages which need cancelling as canceled. // stage ids are returned for task cancellation action @@ -779,13 +870,24 @@ impl ExecutionGraph for AdaptiveExecutionGraph { running_stage .update_task_metrics(task_id, operator_metrics)?; + let ballista_core::serde::protobuf::SuccessfulTask { + partitions, + runtime_stats, + .. + } = successful_task; + debug!( + "append_runtime_stats_reports: job={} stage={} task={} report_count={}", + job_id, + stage_id, + task_id, + runtime_stats.len(), + ); + running_stage + .append_runtime_stats_reports(task_id, runtime_stats); + locations.append( &mut crate::state::execution_graph::partition_to_location( - &job_id, - task_id, - stage_id, - executor, - successful_task.partitions, + &job_id, task_id, stage_id, executor, partitions, ), ); } else { @@ -815,6 +917,11 @@ impl ExecutionGraph for AdaptiveExecutionGraph { stage_metrics, ); } + ballista_core::execution_plans::log_merged_runtime_stats( + job_id.as_str(), + stage_id, + &running_stage.runtime_stats_reports, + ); } let stages_to_cancel = self.update_stage_progress( stage_id, 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 5947be4c3d..b79bdeb36c 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/coalesce_partitions.rs @@ -228,6 +228,16 @@ impl PhysicalOptimizerRule for CoalescePartitionsRule { return Ok(plan); } + // Coalesce and range-repartition are both set, bail so as to not return incorrect results + // TODO: handle this case by coalescing range cuts too (see #2220) + if leaves + .iter() + .any(|arc| as_exchange(arc).range_repartition_routing().is_some()) + { + debug!("[coalesce-rule] range-repartitioned leaf present; bail entire group"); + return Ok(plan); + } + // The alignment-group invariant assumes a shared `M`. In every plan // shape we currently produce, all leaves of one stage subtree are // hash-partitioned by the same target_partitions setting upstream, 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 a07312d14e..80e325566f 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rs @@ -16,9 +16,15 @@ // under the License. use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; +use ballista_core::execution_plans::{ + OrderedRangeRepartitionExec, UnorderedRangeRepartitionExec, preserves_partitioning, +}; +use datafusion::common::plan_err; use datafusion::common::tree_node::{Transformed, TreeNode}; use datafusion::physical_optimizer::PhysicalOptimizerRule; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +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::{ExecutionPlan, execution_plan}; @@ -45,6 +51,40 @@ impl DistributedExchangeRule { &self, execution_plan: Arc, ) -> datafusion::error::Result>> { + // DataFusion's null-aware hash join coordinates visited rows and + // probe-side NULL state with in-process atomics. A CollectLeft join + // running once per probe partition therefore produces duplicate or + // incorrect output. This is the final plan-mutating rule, so enforce an + // explicit probe-side coalesce here after DataFusion's optimizers have + // finished. Add an exchange when the subtree has no existing boundary. + if let Some(hash_join) = execution_plan.downcast_ref::() + && hash_join.null_aware + && *hash_join.partition_mode() == PartitionMode::CollectLeft + && hash_join + .right() + .downcast_ref::() + .is_none() + { + let left = hash_join.left().clone(); + let right = hash_join.right().clone(); + let right = if right.downcast_ref::().is_none() + && !matches!(nearest_exchange_status(&right), ExchangeStatus::Unresolved) + { + Arc::new(ExchangeExec::new( + right, + None, + self.plan_id_generator + .fetch_add(1, std::sync::atomic::Ordering::Relaxed), + )) as Arc + } else { + right + }; + let right = Arc::new(CoalescePartitionsExec::new(right)); + return Ok(Transformed::yes( + execution_plan.with_new_children(vec![left, right])?, + )); + } + if let Some(coalesce) = execution_plan.downcast_ref::() { let input = coalesce.input(); if input.downcast_ref::().is_none() @@ -90,11 +130,83 @@ impl DistributedExchangeRule { ); return Ok(Transformed::yes(Arc::new(exchange_exec))); } + } else if !execution_plan.is::() { + let children = execution_plan.children(); + match children.as_slice() { + [] => {} + [child] => { + if can_be_range_repartitioned(child)? { + let exchange_exec = ExchangeExec::new( + Arc::clone(child), + None, + self.plan_id_generator + .fetch_add(1, std::sync::atomic::Ordering::Relaxed), + ); + return Ok(Transformed::yes( + execution_plan + .with_new_children(vec![Arc::new(exchange_exec)])?, + )); + } + } + many => { + let mut any_range = false; + for c in many { + any_range |= can_be_range_repartitioned(c)?; + } + if any_range { + return plan_err!( + "range-repartitioned child under multi-child parent `{}`: \ + cross-stage cut coordination is not yet implemented", + execution_plan.name() + ); + } + } + } } Ok(Transformed::no(execution_plan)) } } +/// Returns whether a plan should have distributed range-repartitioning added: +/// +/// `Ok(true)` - `plan` has a "range-repartitioned child" that should have an ExchangeExec +/// +/// `Ok(false)` - no range-repartition was present +/// +/// `Err(_)` - there IS a URRE/ORRE below, but an intermediate op +/// disturbs the routing expression enough that we can't safely route +/// through it +fn can_be_range_repartitioned( + plan: &Arc, +) -> datafusion::error::Result { + if !preserves_partitioning(plan.as_ref()) { + // We've hit some other repartitioner, not range-repartitioned + return Ok(false); + } + let children = plan.children(); + let [child] = children.as_slice() else { + // We don't support multi-legged plans for now (SMJ, etc) + return Ok(false); + }; + if !child.is::() + && !child.is::() + { + // Not range-repartitioned + return Ok(false); + } + // We are range repartitioned, but make sure the routing expression survives + if plan.is::() { + // TODO: verify by checking expression itself + return plan_err!( + "range-repartitioned child under `{}`: routing expression \ + cannot be safely remapped to the boundary schema", + plan.name() + ); + } + // We can range-repartition this + Ok(true) +} + impl PhysicalOptimizerRule for DistributedExchangeRule { fn optimize( &self, @@ -110,14 +222,34 @@ impl PhysicalOptimizerRule for DistributedExchangeRule { .downcast_ref::() .is_some() { - Ok(result.data) - } else { - let plan_id = self + return Ok(result.data); + } + + // A range-repartitioned root is never visited as a child by + // `transform_up`, so wrap it here before the outer + // `AdaptiveDatafusionExec` goes on. + // + // TODO: kill this branch — and the range-repart arm in + // `transform()` above — when `ExchangeExec` carries a + // range-cuts partitioning variant the way it carries + // `Partitioning::Hash`. URRE will replace itself with an + // ExchangeExec (like the Hash arm does), `transform_up`'s + // output will already have an ExchangeExec at the range-repart + // position, and the plain `AdaptiveDatafusionExec` wrap below + // handles the root case with no extra ceremony. + let inner = if can_be_range_repartitioned(&result.data)? { + let id = self .plan_id_generator .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Arc::new(ExchangeExec::new(result.data, None, id)) as Arc + } else { + result.data + }; + let plan_id = self + .plan_id_generator + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - Ok(Arc::new(AdaptiveDatafusionExec::new(plan_id, result.data))) - } + Ok(Arc::new(AdaptiveDatafusionExec::new(plan_id, inner))) } fn name(&self) -> &str { @@ -165,6 +297,10 @@ mod tests { use super::*; use crate::assert_plan; use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; + use ballista_core::execution_plans::{ + RuntimeStatsExec, UnorderedRangeRepartitionExec, + }; + use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::{ColumnStatistics, Statistics}; use datafusion::config::ConfigOptions; @@ -208,6 +344,44 @@ mod tests { Arc::new(SortPreservingMergeExec::new(ordering, input)) } + #[test] + fn null_aware_join_coalesces_probe_after_other_optimizers() { + use datafusion::common::{JoinType, NullEquality}; + + let left = leaf_exec(); + let right = leaf_exec(); + let join = Arc::new( + HashJoinExec::try_new( + left, + right, + vec![( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("a", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(), + ) as Arc; + + let result = DistributedExchangeRule::default() + .optimize(join, &config()) + .unwrap(); + + assert_plan!(result.as_ref(), @ r" + AdaptiveDatafusionExec: is_final=false, plan_id=1, stage_id=pending, stage_resolved=false + HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(a@0, a@0)] + StatisticsExec: col_count=1, row_count=Absent + CoalescePartitionsExec + ExchangeExec: partitioning=None, plan_id=0, stage_id=pending, stage_resolved=false + StatisticsExec: col_count=1, row_count=Absent + "); + } + // --- CoalescePartitionsExec --- #[test] @@ -573,4 +747,229 @@ mod tests { "second optimize on same rule: ExchangeExec should get plan_id=2 (counter at 2 after first call used 0 and 1)" ); } + + // --- range-repartition --- + + fn float_leaf_exec() -> Arc { + // URRE routing today requires Float64 non-nullable. + let schema = Schema::new(vec![Field::new("v", DataType::Float64, false)]); + let stats = Statistics { + num_rows: Default::default(), + total_byte_size: Default::default(), + column_statistics: vec![ColumnStatistics::new_unknown()], + }; + Arc::new(StatisticsExec::new(stats, schema)) + } + + fn stats_over_urre_over_leaf() -> Arc { + let sort_expr = PhysicalSortExpr { + expr: Arc::new(Column::new("v", 0)), + options: SortOptions { + descending: false, + nulls_first: false, + }, + }; + let urre: Arc = Arc::new( + UnorderedRangeRepartitionExec::try_new( + float_leaf_exec(), + vec![sort_expr.clone()], + 4, + ) + .unwrap(), + ); + Arc::new(RuntimeStatsExec::try_new(urre, Some(vec![sort_expr])).unwrap()) + } + + fn count_exchanges(plan: &dyn ExecutionPlan) -> usize { + let here = usize::from(plan.is::()); + here + plan + .children() + .iter() + .map(|c| count_exchanges(c.as_ref())) + .sum::() + } + + fn display_plan(plan: &Arc) -> String { + format!( + "{}", + datafusion::physical_plan::displayable(plan.as_ref()).indent(true) + ) + } + + /// A bug in either half of the idempotency guard (the + /// `!.is::()` check on the visited node, or + /// `is_range_repartitioned` accidentally recognising `ExchangeExec`) + /// would double-wrap on every AQE replan. + #[test] + fn range_repartition_optimize_is_idempotent() { + let rule = DistributedExchangeRule::default(); + + let first = rule + .optimize(stats_over_urre_over_leaf(), &config()) + .unwrap(); + let second = rule.optimize(first.clone(), &config()).unwrap(); + + assert_eq!( + count_exchanges(first.as_ref()), + 1, + "first pass must insert exactly one ExchangeExec above the range-repartitioned root" + ); + assert_eq!( + count_exchanges(second.as_ref()), + 1, + "second pass must not insert a second ExchangeExec" + ); + assert_eq!( + display_plan(&first), + display_plan(&second), + "second DER pass over its own output must be a no-op" + ); + } + + /// A range-repartition at the plan root — the canonical shape a + /// range-repartition-inserting rule emits, with nothing above it — + /// must still get an `ExchangeExec` wrapped above it. Without it, + /// `set_repartition_routing` has no parking slot for the recovered + /// cuts and downstream never gets a `PerPartitionFilterExec` to + /// trim straddler duplication. + #[test] + fn range_repartition_at_plan_root_gets_exchange_inserted() { + let rule = DistributedExchangeRule::default(); + let root = stats_over_urre_over_leaf(); + + let result = rule.optimize(root, &config()).unwrap(); + + let adaptive = result + .downcast_ref::() + .expect("DER wraps its output in AdaptiveDatafusionExec"); + let below = adaptive.input(); + assert!( + below.is::(), + "expected an ExchangeExec between the AdaptiveDatafusionExec \ + wrapper and the range-repartitioned subtree; got {}", + below.name() + ); + let rse = below.children()[0]; + assert!(rse.is::()); + assert!(rse.children()[0].is::()); + assert_eq!(count_exchanges(result.as_ref()), 1); + } + + /// A range-repartitioned chain on each branch of a `UnionExec` is + /// rejected at plan time — the machinery needs cross-stage cut + /// coordination (see doc-comment on `is_range_repartitioned`), and + /// silently producing independently-cut shuffles would misroute. + #[test] + fn range_repartition_under_union_is_rejected() { + let rule = DistributedExchangeRule::default(); + let union: Arc = + datafusion::physical_plan::union::UnionExec::try_new(vec![ + stats_over_urre_over_leaf(), + stats_over_urre_over_leaf(), + ]) + .unwrap(); + + let err = rule.optimize(union, &config()).unwrap_err().to_string(); + + assert!( + err.contains("range-repartitioned child under multi-child parent"), + "unexpected error: {err}" + ); + assert!(err.contains("UnionExec"), "unexpected error: {err}"); + } + + /// TODO: SMJ over range-repartitioned inputs is a motivating + /// consumer for cross-stage cut coordination — both sides must + /// share one cut set (derived from both sketches, parked on both + /// boundary exchanges) for equijoin keys to land in the same + /// downstream partition. Until that code exists, DER rejects the + /// shape at plan time. + #[test] + fn range_repartition_under_sort_merge_join_is_rejected() { + use datafusion::common::{JoinType, NullEquality}; + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_plan::joins::SortMergeJoinExec; + + let rule = DistributedExchangeRule::default(); + let on: Vec<(Arc, Arc)> = + vec![(Arc::new(Column::new("v", 0)), Arc::new(Column::new("v", 0)))]; + let smj: Arc = Arc::new( + SortMergeJoinExec::try_new( + stats_over_urre_over_leaf(), + stats_over_urre_over_leaf(), + on, + None, + JoinType::Inner, + vec![SortOptions { + descending: false, + nulls_first: false, + }], + NullEquality::NullEqualsNothing, + ) + .unwrap(), + ); + + let err = rule.optimize(smj, &config()).unwrap_err().to_string(); + + assert!( + err.contains("range-repartitioned child under multi-child parent"), + "unexpected error: {err}" + ); + assert!(err.contains("SortMergeJoinExec"), "unexpected error: {err}"); + } + + /// A `ProjectionExec` between (O/U)RRE and the boundary could + /// reindex, drop, or shadow the routing expression's referenced + /// columns — the read-side `PerPartitionFilterExec` would evaluate + /// against the wrong column and silently misroute. DER rejects the + /// shape at plan time; the fix will be revisited when arbitrary + /// routing expressions replace the current single-column form. + /// + /// See https://github.com/apache/datafusion-ballista/pull/2196#discussion_r3705634907 + #[test] + fn range_repartition_under_projection_is_rejected() { + use datafusion::physical_expr::PhysicalExpr; + use datafusion::physical_plan::projection::ProjectionExec; + + let rule = DistributedExchangeRule::default(); + let schema = Schema::new(vec![ + Field::new("v", DataType::Float64, false), + Field::new("tag", DataType::Float64, false), + ]); + let stats = Statistics { + num_rows: Default::default(), + total_byte_size: Default::default(), + column_statistics: vec![ + ColumnStatistics::new_unknown(), + ColumnStatistics::new_unknown(), + ], + }; + let leaf: Arc = Arc::new(StatisticsExec::new(stats, schema)); + let sort_expr = PhysicalSortExpr { + expr: Arc::new(Column::new("v", 0)), + options: SortOptions { + descending: false, + nulls_first: false, + }, + }; + let urre: Arc = Arc::new( + UnorderedRangeRepartitionExec::try_new(leaf, vec![sort_expr], 4).unwrap(), + ); + // Swap: post-projection output is (tag, v) — routing key `v` + // moves from index 0 to index 1. + let proj_exprs: Vec<(Arc, String)> = vec![ + (Arc::new(Column::new("tag", 1)), "tag".to_string()), + (Arc::new(Column::new("v", 0)), "v".to_string()), + ]; + let proj: Arc = + Arc::new(ProjectionExec::try_new(proj_exprs, urre).unwrap()); + + let err = rule.optimize(proj, &config()).unwrap_err().to_string(); + + assert!( + err.contains("routing expression cannot be safely remapped"), + "unexpected error: {err}" + ); + assert!(err.contains("ProjectionExec"), "unexpected error: {err}"); + } } 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 53f44dae61..e3807ab711 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/join_selection.rs @@ -151,10 +151,12 @@ impl PhysicalOptimizerRule for SelectJoinRule { // at this point we know there are two exchanges // as we added them beforehand JoinSelectionAction::LateCollectLeft(hash_join_exec) => { - if Self::supports_swap_join_order( - hash_join_exec.left.as_ref(), - hash_join_exec.right.as_ref(), - )? { + if !hash_join_exec.null_aware + && Self::supports_swap_join_order( + hash_join_exec.left.as_ref(), + hash_join_exec.right.as_ref(), + )? + { let left = hash_join_exec.left.clone(); let right = hash_join_exec.right.clone(); @@ -199,10 +201,11 @@ impl PhysicalOptimizerRule for SelectJoinRule { } JoinSelectionAction::CollectLeft(hash_join_exec) => { - let plan = if Self::supports_swap_join_order( - hash_join_exec.left.as_ref(), - hash_join_exec.right.as_ref(), - )? { + let plan = if !hash_join_exec.null_aware + && Self::supports_swap_join_order( + hash_join_exec.left.as_ref(), + hash_join_exec.right.as_ref(), + )? { hash_join_exec .swap_inputs(PartitionMode::CollectLeft)? } else { @@ -293,10 +296,11 @@ impl PhysicalOptimizerRule for SelectJoinRule { Ok(Transformed::yes(dynamic_join)) } JoinSelectionAction::Hash(hash_join_exec) => { - let hash_join_exec = if Self::supports_swap_join_order( - hash_join_exec.left.as_ref(), - hash_join_exec.right.as_ref(), - )? { + let hash_join_exec = if !hash_join_exec.null_aware + && Self::supports_swap_join_order( + hash_join_exec.left.as_ref(), + hash_join_exec.right.as_ref(), + )? { hash_join_exec .swap_inputs(*hash_join_exec.partition_mode())? } else { @@ -610,6 +614,130 @@ mod tests { assert_plan!(optimized.as_ref(), @ "DataSourceExec: partitions=1, partition_sizes=[1]"); } + #[test] + fn null_aware_anti_join_is_not_swapped_by_aqe() { + use datafusion::physical_expr::expressions::Column; + use datafusion::{ + arrow::datatypes::{DataType, Field, Schema}, + common::{ + ColumnStatistics, JoinType, NullEquality, Statistics, stats::Precision, + }, + physical_plan::{ + Partitioning, + joins::{HashJoinExec, PartitionMode}, + repartition::RepartitionExec, + test::exec::StatisticsExec, + }, + }; + + fn stats_exec(name: &str, bytes: usize) -> Arc { + Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(bytes / 4), + total_byte_size: Precision::Inexact(bytes), + column_statistics: vec![ColumnStatistics::new_unknown()], + }, + Schema::new(vec![Field::new(name, DataType::Int32, true)]), + )) + } + + // A normal join would swap these inputs to build from the smaller right + // side. Doing that to a null-aware LeftAnti creates an invalid + // null-aware RightAnti join. + let left = stats_exec("big_key", 1024 * 1024); + let right = Arc::new( + RepartitionExec::try_new( + stats_exec("small_key", 1024), + Partitioning::RoundRobinBatch(4), + ) + .unwrap(), + ) as Arc; + let join = HashJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + vec![( + Arc::new(Column::new("big_key", 0)) as _, + Arc::new(Column::new("small_key", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(); + let dynamic = DynamicJoinSelectionExec::from_hash_join(&join, 0).unwrap() + as Arc; + + let resolved = SelectJoinRule::default() + .optimize(dynamic, &ConfigOptions::default()) + .expect("AQE must not swap a null-aware LeftAnti join"); + let hash_join = resolved + .downcast_ref::() + .expect("AQE should resolve to HashJoinExec"); + + assert_eq!(*hash_join.join_type(), JoinType::LeftAnti); + assert_eq!(*hash_join.partition_mode(), PartitionMode::CollectLeft); + assert!(hash_join.null_aware); + assert_eq!(hash_join.left().schema().field(0).name(), "big_key"); + assert_eq!(hash_join.right().schema().field(0).name(), "small_key"); + } + + #[test] + fn null_aware_anti_join_rejects_known_oversized_build_side_in_aqe() { + use datafusion::physical_expr::expressions::Column; + use datafusion::{ + arrow::datatypes::{DataType, Field, Schema}, + common::{ + ColumnStatistics, JoinType, NullEquality, Statistics, stats::Precision, + }, + physical_plan::{ + joins::{HashJoinExec, PartitionMode}, + test::exec::StatisticsExec, + }, + }; + + fn stats_exec(name: &str, bytes: usize) -> Arc { + Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(bytes / 4), + total_byte_size: Precision::Inexact(bytes), + column_statistics: vec![ColumnStatistics::new_unknown()], + }, + Schema::new(vec![Field::new(name, DataType::Int32, true)]), + )) + } + + let join = HashJoinExec::try_new( + stats_exec("big_key", 20 * 1024 * 1024), + stats_exec("small_key", 1024), + vec![( + Arc::new(Column::new("big_key", 0)) as _, + Arc::new(Column::new("small_key", 0)) as _, + )], + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + ) + .unwrap(); + let dynamic = DynamicJoinSelectionExec::from_hash_join(&join, 0).unwrap() + as Arc; + + let error = SelectJoinRule::default() + .optimize(dynamic, &ConfigOptions::default()) + .unwrap_err(); + assert!( + error.to_string().contains( + "build side does not fit ballista.optimizer.broadcast_join_threshold_bytes" + ), + "{error}" + ); + } + /// When `ballista.planner.adaptive_join.enabled = false` the `DelayJoinSelectionRule` /// must be a no-op: a plan containing a `DynamicJoinSelectionExec` node must /// be returned unchanged. 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 87d52a5bd8..87331f4e02 100644 --- a/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs +++ b/ballista/scheduler/src/state/aqe/optimizer_rule/propagate_empty.rs @@ -60,6 +60,9 @@ pub struct JoinInfo<'a> { pub left: &'a Arc, pub right: &'a Arc, pub schema: SchemaRef, + /// True when the join carries an embedded projection, in which case its + /// output schema no longer maps positionally onto the input schemas. + pub has_projection: bool, } /// This is [datafusion::optimizer::propagate_empty_relation::PropagateEmptyRelation] rule with difference @@ -78,7 +81,10 @@ impl PropagateEmptyExecRule { if let Some(filter) = plan.downcast_ref::() && is_empty_exec!(filter.input()) { - Ok(Transformed::yes(filter.input().clone())) + // A FilterExec may carry an embedded projection, so its schema can + // differ from its input's — build a fresh EmptyExec instead of + // reusing the input. + empty_exec!(filter) } else if let Some(coalesce) = plan.downcast_ref::() && is_empty_exec!(coalesce.input()) { @@ -86,7 +92,10 @@ impl PropagateEmptyExecRule { } else if let Some(exchange) = plan.downcast_ref::() && is_empty_exec!(exchange.input()) { - Ok(Transformed::yes(exchange.input().clone())) + // Keep the exchange's output partitioning: parents were planned + // against it, and the stats-based exchange arm below preserves it + // the same way. + empty_exec!(exchange) } else if let Some(projection) = plan.downcast_ref::() && is_empty_exec!(projection.input()) { @@ -101,6 +110,18 @@ impl PropagateEmptyExecRule { Ok(Transformed::yes(limit.input().clone())) } else if let Some(aggregation) = plan.downcast_ref::() && is_empty_exec!(aggregation.input()) + // An aggregate with no GROUP BY emits exactly one row even over zero + // input rows (`sum` -> NULL, `count` -> 0), so it must be preserved. + && !aggregation.group_expr().is_empty() + // Same for the empty grouping set `()` of GROUPING SETS/ROLLUP/CUBE: + // an all-true null mask is the grand-total group, which also emits + // one row over empty input. Mirrors `has_empty_grouping_set` in + // DataFusion's logical PropagateEmptyRelation rule. + && !aggregation + .group_expr() + .groups() + .iter() + .any(|group| group.iter().all(|&null_masked| null_masked)) { empty_exec!(aggregation) } else if let Some(repartition) = plan.downcast_ref::() @@ -123,13 +144,19 @@ impl PropagateEmptyExecRule { let left_field_count = join.left.schema().fields.len(); + // Rewrites that keep one input alive assume the join schema maps + // positionally onto the input schemas; an embedded projection + // breaks that, so those arms must not fire. Arms that produce an + // EmptyExec use the join's own (projected) schema and stay safe. + let no_projection = !join.has_projection; + // Checking whether join would produce an empty result match join.join_type { JoinType::Inner if left_empty || right_empty => empty_exec!(plan), JoinType::Left if left_empty => empty_exec!(plan), // Left Join with empty right: all left rows survive // with NULLs for right columns. - JoinType::Left if right_empty => { + JoinType::Left if right_empty && no_projection => { Ok(Transformed::yes(build_null_padded_projection( Arc::clone(join.left), join.schema, @@ -140,7 +167,7 @@ impl PropagateEmptyExecRule { JoinType::Right if right_empty => empty_exec!(plan), // Right Join with empty left: all right rows survive // with NULLs for left columns. - JoinType::Right if left_empty => { + JoinType::Right if left_empty && no_projection => { Ok(Transformed::yes(build_null_padded_projection( Arc::clone(join.right), join.schema, @@ -151,18 +178,22 @@ impl PropagateEmptyExecRule { JoinType::LeftSemi if left_empty || right_empty => empty_exec!(plan), JoinType::RightSemi if left_empty || right_empty => empty_exec!(plan), JoinType::LeftAnti if left_empty => empty_exec!(plan), - JoinType::LeftAnti if right_empty => { + JoinType::LeftAnti if right_empty && no_projection => { Ok(Transformed::yes((*join.left).clone())) } JoinType::RightAnti if right_empty => empty_exec!(plan), - JoinType::RightAnti if left_empty => { + JoinType::RightAnti if left_empty && no_projection => { Ok(Transformed::yes((*join.right).clone())) } // Return empty if both sides are empty JoinType::Full if left_empty && right_empty => empty_exec!(plan), // For Full Join, if one side is empty, replace with a // Projection that null-pads the empty side's columns. - JoinType::Full if right_empty && is_guaranteed_non_empty(join.left) => { + JoinType::Full + if right_empty + && no_projection + && is_guaranteed_non_empty(join.left) => + { Ok(Transformed::yes(build_null_padded_projection( Arc::clone(join.left), join.schema.clone(), @@ -170,7 +201,11 @@ impl PropagateEmptyExecRule { true, )?)) } - JoinType::Full if left_empty && is_guaranteed_non_empty(join.right) => { + JoinType::Full + if left_empty + && no_projection + && is_guaranteed_non_empty(join.right) => + { Ok(Transformed::yes(build_null_padded_projection( Arc::clone(join.right), join.schema.clone(), @@ -221,6 +256,7 @@ pub fn as_join(plan: &Arc) -> Option> { left: join.left(), right: join.right(), schema: join.schema(), + has_projection: join.contains_projection(), }); } if let Some(join) = any.downcast_ref::() { @@ -229,6 +265,7 @@ pub fn as_join(plan: &Arc) -> Option> { left: join.left(), right: join.right(), schema: join.schema(), + has_projection: false, }); } @@ -327,9 +364,13 @@ mod tests { use datafusion::{ arrow::datatypes::{DataType, Field, Schema}, - common::{ColumnStatistics, JoinType, NullEquality, Statistics}, + common::{ColumnStatistics, JoinType, NullEquality, ScalarValue, Statistics}, + physical_expr::aggregate::AggregateExprBuilder, physical_plan::{ - expressions::Column, joins::PartitionMode, test::exec::StatisticsExec, + aggregates::{AggregateMode, PhysicalGroupBy}, + expressions::{Column, Literal}, + joins::PartitionMode, + test::exec::StatisticsExec, }, }; @@ -871,6 +912,232 @@ mod tests { assert!(result.downcast_ref::().is_some()); } + // ── Aggregate ──────────────────────────────────────────────────────────── + + /// Build an `AggregateExec` over `input` with the given grouping. + fn aggregate( + input: Arc, + group_by: PhysicalGroupBy, + mode: AggregateMode, + ) -> Arc { + let aggr = AggregateExprBuilder::new( + datafusion::functions_aggregate::count::count_udaf(), + vec![Arc::new(Column::new("a", 0))], + ) + .schema(input.schema()) + .alias("count(a)") + .build() + .unwrap(); + + Arc::new( + AggregateExec::try_new( + mode, + group_by, + vec![Arc::new(aggr)], + vec![None], + input, + schema(), + ) + .unwrap(), + ) + } + + #[test] + fn aggregate_with_grouping_over_empty_eliminated() { + // GROUP BY over zero rows produces zero groups — collapsing is correct. + let group_by = PhysicalGroupBy::new_single(vec![( + Arc::new(Column::new("a", 0)), + "a".into(), + )]); + let plan = aggregate( + Arc::new(EmptyExec::new(schema())), + group_by, + AggregateMode::Single, + ); + assert_empty(&transform(plan)); + } + + #[test] + fn aggregate_without_grouping_over_empty_is_preserved() { + // No GROUP BY means exactly one output row even over zero input rows + // (`sum` -> NULL, `count` -> 0). Collapsing to EmptyExec loses that row. + let plan = aggregate( + Arc::new(EmptyExec::new(schema())), + PhysicalGroupBy::default(), + AggregateMode::Single, + ); + assert_untouched(&transform(plan)); + } + + #[test] + fn partial_aggregate_without_grouping_over_empty_is_preserved() { + // Same holds for a partial aggregate: it emits one row of accumulator + // state per partition, which the final aggregate needs. + let plan = aggregate( + Arc::new(EmptyExec::new(schema())), + PhysicalGroupBy::default(), + AggregateMode::Partial, + ); + assert_untouched(&transform(plan)); + } + + #[test] + fn aggregate_grouping_sets_with_empty_subset_over_empty_is_preserved() { + // `ROLLUP(a)` / `GROUPING SETS ((a), ())` lowers to a non-empty `expr` + // list plus a `groups` mask containing the all-null (empty) subset. + // Over zero input rows, that empty subset still emits one all-NULL row + // — matches DataFusion's logical `has_empty_grouping_set` guard. + let group_by = PhysicalGroupBy::new( + vec![(Arc::new(Column::new("a", 0)), "a".into())], + vec![(Arc::new(Literal::new(ScalarValue::Int32(None))), "a".into())], + vec![vec![false], vec![true]], + true, + ); + let plan = aggregate( + Arc::new(EmptyExec::new(schema())), + group_by, + AggregateMode::Single, + ); + assert_untouched(&transform(plan)); + } + + // ── embedded projections — schema no longer matches the inputs ───────── + + fn two_col_schema() -> Schema { + Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ]) + } + + fn non_empty_two_col_stats_exec() -> Arc { + Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Exact(100), + total_byte_size: Precision::Exact(800), + column_statistics: vec![ColumnStatistics::new_unknown(); 2], + }, + two_col_schema(), + )) + } + + fn hash_join_with_projection( + left: Arc, + right: Arc, + join_type: JoinType, + projection: Vec, + ) -> Arc { + Arc::new( + HashJoinExec::try_new( + left, + right, + join_on(), + None, + &join_type, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap() + .with_projection(Some(projection)) + .unwrap(), + ) + } + + #[test] + fn left_join_with_embedded_projection_right_empty_untouched() { + // ProjectionPushdown can embed a projection into HashJoinExec. The join + // schema is then no longer `left fields ++ right fields`, so positional + // null-padding would read the wrong columns. The join must stay as is. + let plan = hash_join_with_projection( + non_empty_two_col_stats_exec(), + empty_stats_exec(), + JoinType::Left, + vec![1], + ); + let join_schema = plan.schema(); + let result = transform(plan); + assert!( + result.downcast_ref::().is_some(), + "expected join to be untouched, got {:?}", + result.name() + ); + assert_eq!(result.schema(), join_schema); + } + + #[test] + fn left_anti_with_embedded_projection_right_empty_untouched() { + // A LeftAnti join with an embedded projection outputs a subset of the + // left columns; replacing it with the raw left child changes the schema. + let plan = hash_join_with_projection( + non_empty_two_col_stats_exec(), + empty_stats_exec(), + JoinType::LeftAnti, + vec![1], + ); + let join_schema = plan.schema(); + let result = transform(plan); + assert_eq!(result.schema(), join_schema); + assert!( + result.downcast_ref::().is_some(), + "expected join to be untouched, got {:?}", + result.name() + ); + } + + #[test] + fn filter_with_embedded_projection_over_empty_preserves_schema() { + // FilterExec can carry an embedded projection, so its schema may differ + // from its input's. The result must use the filter's schema, not the + // input's. + use datafusion::logical_expr::Operator; + use datafusion::physical_plan::expressions::BinaryExpr; + use datafusion::physical_plan::filter::FilterExecBuilder; + + let input = Arc::new(EmptyExec::new(Arc::new(two_col_schema()))); + let predicate: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(0)))), + )); + let plan: Arc = Arc::new( + FilterExecBuilder::new(predicate, input) + .apply_projection(Some(vec![1])) + .unwrap() + .build() + .unwrap(), + ); + let filter_schema = plan.schema(); + let result = transform(plan); + assert_empty(&result); + assert_eq!(result.schema(), filter_schema); + } + + // ── exchange over empty input — partitioning must survive ─────────────── + + #[test] + fn exchange_over_empty_input_preserves_partition_count() { + // Replacing ExchangeExec with its EmptyExec input would drop the + // exchange's output partitioning; parents planned against it would see + // the wrong partition count. Mirrors the stats-based exchange arm, + // which already preserves the partition count. + use datafusion::physical_plan::Partitioning; + + let empty = Arc::new(EmptyExec::new(schema()).with_partitions(1)); + let plan: Arc = Arc::new(ExchangeExec::new( + empty, + Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 16)), + 0, + )); + let result = transform(plan); + assert_empty(&result); + assert_eq!( + result.properties().output_partitioning().partition_count(), + 16 + ); + } + // ── unknown stats — never optimised ───────────────────────────────────── #[test] diff --git a/ballista/scheduler/src/state/aqe/planner.rs b/ballista/scheduler/src/state/aqe/planner.rs index 67d303761a..690a5f2bcd 100644 --- a/ballista/scheduler/src/state/aqe/planner.rs +++ b/ballista/scheduler/src/state/aqe/planner.rs @@ -16,7 +16,9 @@ // under the License. use crate::physical_optimizer::filter_pushdown::FilterPushdown; use crate::state::aqe::adapter::BallistaAdapter; -use crate::state::aqe::execution_plan::{AdaptiveDatafusionExec, ExchangeExec}; +use crate::state::aqe::execution_plan::{ + AdaptiveDatafusionExec, ExchangeExec, RangeRepartitionRouting, +}; use crate::state::aqe::optimizer_rule::chaos_exec::ChaosCreatingRule; use crate::state::aqe::optimizer_rule::{ CoalescePartitionsRule, DelayJoinSelectionRule, DistributedExchangeRule, @@ -199,14 +201,37 @@ impl AdaptivePlanner { let _ = self.runnable_stage_cache.remove(&stage_id); Ok(()) } - /// Resolves a stage by its ID and updates its partitions. - /// - /// # Arguments - /// * `stage_id` - The ID of the stage to resolve. - /// * `partitions` - The resolved partitions for the stage. + + /// Attaches range-repartition-recovered range boundaries to the + /// boundary `ExchangeExec` for `stage_id`. Called right after the + /// completed range-repartition stage's partition mapping is resolved; + /// the routing carries the cuts + routing expression that downstream + /// task specialization needs to build per-partition range filters. /// - /// # Returns - /// A `Result` indicating success or failure. + /// Errors if `stage_id` has no `ExchangeExec` to park on — the caller + /// is responsible for ensuring the stage boundary is an `ExchangeExec`, + /// so downstream read-side filter injection has something to hook into. + pub(super) fn set_repartition_routing( + &mut self, + stage_id: usize, + routing: RangeRepartitionRouting, + ) -> common::Result<()> { + let stage = self.runnable_stage_cache.get(&stage_id).ok_or_else(|| { + DataFusionError::Internal(format!( + "set_repartition_routing: stage {stage_id} not in runnable cache" + )) + })?; + let exchange = stage.downcast_ref::().ok_or_else(|| { + DataFusionError::Internal(format!( + "set_repartition_routing: stage {stage_id} boundary is {}, \ + not ExchangeExec — nowhere to park cuts for downstream filter", + stage.name() + )) + })?; + exchange.resolve_range_repartition_routing(routing); + Ok(()) + } + pub(super) fn finalise_stage_internal( &mut self, stage_id: usize, @@ -261,9 +286,12 @@ impl AdaptivePlanner { } } - /// Once all tasks has been completed marks stage as resolved - /// and returns partition allocations - pub fn finalise_stage( + /// Once all tasks have completed, pop the accumulated stage output as a + /// K-shaped `Vec>` (or the broadcast-shape variant) + /// *without* parking it on the ExchangeExec. Caller can post-process + /// (e.g. range-repartition overlap remap) before calling + /// [`resolve_stage_partitions`](Self::resolve_stage_partitions). + pub fn take_stage_output_partitions( &mut self, stage_id: usize, ) -> common::Result>> { @@ -293,10 +321,18 @@ impl AdaptivePlanner { ))? .partition_locations(output_partition_count) }; - self.finalise_stage_internal(stage_id, stage_output.clone())?; Ok(stage_output) } + /// Save the given partition list on the stage's ExchangeExec and trigger a replan. + pub fn resolve_stage_partitions( + &mut self, + stage_id: usize, + partitions: Vec>, + ) -> common::Result<()> { + self.finalise_stage_internal(stage_id, partitions) + } + /// Replans the stages by applying physical optimizations. /// /// # Returns diff --git a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs index 0ac47ae721..c51cff0a1d 100644 --- a/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs +++ b/ballista/scheduler/src/state/aqe/test/coalesce_rule.rs @@ -203,6 +203,67 @@ async fn should_skip_coalesce_when_rule_disabled() -> datafusion::error::Result< Ok(()) } +/// Range-repartition bail: same happy-path inputs (M=8 @ 50 bytes, target=200), +/// but we park a `RangeRepartitionRouting` slot on the leaf `ExchangeExec` +/// before the downstream optimizer pass fires — mirrors what `SchedulerAqe` +/// does after a range-repartitioned stage completes. Coalesce K and +/// range-repartition K' would independently commit to different partition +/// counts on the same exchange. Contiguous coalescing is compatible with +/// range partitioning in principle (apache/datafusion-ballista#2220), +/// but until we merge cuts alongside groups the rule bails. +#[tokio::test] +async fn should_skip_coalesce_when_leaf_has_range_repartition_routing() +-> datafusion::error::Result<()> { + use crate::state::aqe::execution_plan::RangeRepartitionRouting; + use datafusion::physical_expr::expressions::Column; + + let ctx = coalesce_context(8, true); + ctx.register_batch("t", mock_batch()?)?; + + let plan = ctx + .sql("select min(a) as c0, c as c2 from t group by c") + .await? + .create_physical_plan() + .await?; + let mut planner = + AdaptivePlanner::try_from_plan(ctx.state().config(), plan, "test_job".into())?; + + let _ = planner.runnable_stages()?.unwrap(); + + // Park routing on the upstream exchange while it's still in the + // runnable cache — `finalise_stage_internal` will remove it. Mirrors + // `SchedulerAqe`: `set_repartition_routing` runs before + // `resolve_stage_partitions` at the boundary. 7 cuts → K=8, matching + // the hash-partitioning's M=8 so the adapter's `PerPartitionFilterExec` + // builds cleanly; the point of the test is the rule's bail decision, + // not a realistic end-to-end range-repartition query. + planner.set_repartition_routing( + 0, + RangeRepartitionRouting { + cuts: vec![10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0], + routing_expr: Arc::new(Column::new("c", 0)), + }, + )?; + + // Resolve partitions → triggers `replan_stages`, where + // `CoalescePartitionsRule` fires on the downstream subtree and sees a + // range-repartitioned leaf → bails. + planner.finalise_stage_internal(0, partitions_with_byte_sizes(&[50; 8]))?; + let _ = planner.runnable_stages()?; + + // No `coalesce=` on plan_id=0; `range_repartition_cuts=7` present. + assert_plan!(planner.current_plan(), @ r" + AdaptiveDatafusionExec: is_final=true, plan_id=1, stage_id=1, stage_resolved=false + ProjectionExec: expr=[min(t.a)@1 as c0, c@0 as c2] + AggregateExec: mode=FinalPartitioned, gby=[c@0 as c], aggr=[min(t.a)] + ExchangeExec: partitioning=Hash([c@0], 8), plan_id=0, stage_id=0, stage_resolved=true, range_repartition_cuts=7 + AggregateExec: mode=Partial, gby=[c@1 as c], aggr=[min(t.a)] + DataSourceExec: partitions=1, partition_sizes=[1] + "); + + Ok(()) +} + /// Degenerate K=M path: every partition is already at target. Bin-pack /// flushes after each one (adding the next would exceed target=200), and /// the post-flush merge is rejected (each bucket = 300, neither small nor diff --git a/ballista/scheduler/src/state/aqe/test/mod.rs b/ballista/scheduler/src/state/aqe/test/mod.rs index 52f173d191..0745a90172 100644 --- a/ballista/scheduler/src/state/aqe/test/mod.rs +++ b/ballista/scheduler/src/state/aqe/test/mod.rs @@ -27,6 +27,9 @@ mod job_failure; mod join_selection; /// Tests if plan is going to be split to stages correctly mod plan_to_stages; +/// Regression tests for range-repartition planning end-to-end +/// through `AdaptivePlanner` (DER → routing park → filter injection). +mod range_repartition; /// A table whose statistics are declared rather than measured mod stats_table; diff --git a/ballista/scheduler/src/state/aqe/test/range_repartition.rs b/ballista/scheduler/src/state/aqe/test/range_repartition.rs new file mode 100644 index 0000000000..8d25a8b83e --- /dev/null +++ b/ballista/scheduler/src/state/aqe/test/range_repartition.rs @@ -0,0 +1,254 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Regression tests for range-repartition planning end-to-end +//! through `AdaptivePlanner` — DER inserts the boundary +//! `ExchangeExec`, `set_repartition_routing` parks the recovered +//! cuts on it, and `cut_partitions` duplicates straddlers so +//! downstream can inject a `PerPartitionFilterExec` to trim them. + +use crate::state::aqe::execution_plan::RangeRepartitionRouting; +use crate::state::aqe::planner::AdaptivePlanner; +use ballista_core::execution_plans::{ + RuntimeStatsExec, UnorderedRangeRepartitionExec, cut_partitions, +}; +use ballista_core::extension::SessionConfigExt; +use ballista_core::serde::protobuf::{RuntimeStatsPartitionEntry, RuntimeStatsReport}; +use ballista_core::serde::scheduler::{ + ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification, + PartitionId, PartitionLocation, PartitionStats, +}; +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::physical_expr::PhysicalSortExpr; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::expressions::col; +use datafusion::prelude::SessionConfig; +use std::sync::Arc; + +fn v_schema() -> Arc { + Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)])) +} + +/// The canonical shape a range-repartition rule is told to emit, +/// here at the root of the plan. +fn stats_over_urre_root() -> Arc { + let schema = v_schema(); + let source: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![]], schema.clone(), None).unwrap(), + ))); + let sort_expr = PhysicalSortExpr { + expr: col("v", schema.as_ref()).unwrap(), + options: SortOptions { + descending: false, + nulls_first: false, + }, + }; + let urre: Arc = Arc::new( + UnorderedRangeRepartitionExec::try_new(source, vec![sort_expr.clone()], 2) + .unwrap(), + ); + Arc::new(RuntimeStatsExec::try_new(urre, Some(vec![sort_expr])).unwrap()) +} + +/// Same shape with an ordinary parent above it. A `FilterExec` survives +/// the optimizer pipeline; a `RoundRobinBatch` `RepartitionExec` and +/// `CoalescePartitionsExec` do not (both get stripped), so the parent +/// choice matters for keeping the range-repartition off the plan root. +fn stats_over_urre_with_parent() -> Arc { + let pred: Arc = + Arc::new(datafusion::physical_expr::expressions::BinaryExpr::new( + col("v", v_schema().as_ref()).unwrap(), + datafusion::logical_expr::Operator::Gt, + Arc::new(datafusion::physical_expr::expressions::Literal::new( + datafusion::scalar::ScalarValue::Float64(Some(0.0)), + )), + )); + Arc::new( + datafusion::physical_plan::filter::FilterExec::try_new( + pred, + stats_over_urre_root(), + ) + .unwrap(), + ) +} + +fn location(sub_part_id: usize, producer_task_id: usize, rows: u64) -> PartitionLocation { + PartitionLocation { + map_partition_id: 0, + partition_id: PartitionId { + job_id: "repro-job".into(), + stage_id: 0, + partition_id: sub_part_id, + }, + executor_meta: ExecutorMetadata { + id: format!("exec-{producer_task_id}"), + host: "".to_string(), + port: 0, + grpc_port: 0, + specification: ExecutorSpecification::default().with_vcores(0), + os_info: ExecutorOperatingSystemSpecification::default(), + }, + partition_stats: PartitionStats::new(Some(rows), None, None), + file_id: Some(producer_task_id as u64), + is_sort_shuffle: false, + } +} + +/// End-to-end through `AdaptivePlanner`: with the range-repartition +/// at the plan root, `cut_partitions` duplicates a straddler across +/// both downstream partitions — correct pre-filter shape — and +/// `set_repartition_routing` must park the cuts on a boundary +/// `ExchangeExec` so downstream can inject the read-side filter. +#[tokio::test] +async fn routing_parks_when_range_repartition_is_plan_root() +-> datafusion::error::Result<()> { + let config = SessionConfig::new_with_ballista(); + let mut planner = AdaptivePlanner::try_from_plan( + &config, + stats_over_urre_root(), + "regression-job".into(), + )?; + + let stages = planner.runnable_stages()?; + let stage_id = stages + .as_ref() + .and_then(|s| s.first()) + .map(|e| e.plan.stage_id()) + .expect("a runnable stage must exist"); + + // Producer task 7, sub-part 0, sketched [5, 15, 25], straddles the cut at 15. + let reports = vec![ballista_core::execution_plans::TaskRuntimeStats { + producer_task_id: 7, + report: RuntimeStatsReport { + order_by: vec![], + partitions: vec![RuntimeStatsPartitionEntry { + partition_id: 0, + row_count: 3, + sketch: Some(ballista_core::execution_plans::sketch_to_proto( + &datafusion_functions_aggregate_common::tdigest::TDigest::new(100) + .merge_unsorted_f64(vec![5.0, 15.0, 25.0]), + )?), + }], + }, + }]; + let cuts = vec![15.0]; + let remapped = cut_partitions(vec![vec![location(0, 7, 3)]], &reports, &cuts)?; + + // `cut_partitions` must duplicate the straddler into both partitions — + // the read-side filter is expected to trim on read. + assert_eq!(remapped[0].len(), 1, "straddler routed into partition 0"); + assert_eq!(remapped[1].len(), 1, "straddler routed into partition 1"); + + let routing = RangeRepartitionRouting { + cuts: cuts.clone(), + routing_expr: col("v", v_schema().as_ref()).unwrap(), + }; + planner.set_repartition_routing(stage_id, routing)?; + + let plan_str = format!( + "{}", + datafusion::physical_plan::displayable(planner.current_plan()).indent(true) + ); + assert!( + plan_str.contains("range_repartition_cuts=1"), + "cuts must be parked on the boundary ExchangeExec so downstream \ + gets a PerPartitionFilterExec — actual plan:\n{plan_str}" + ); + + Ok(()) +} + +/// `set_repartition_routing` is only meaningful when the stage's +/// cached boundary is an `ExchangeExec` — otherwise there's nowhere +/// to park cuts and downstream will never inject the read-side +/// filter. Silently returning `Ok(())` in that case was the outlier +/// in a code path that otherwise errors hard on invariant breaks +/// (Andy's review of PR #2196), so make it fail loud. +#[tokio::test] +async fn set_repartition_routing_errs_when_stage_has_no_exchange() +-> datafusion::error::Result<()> { + let config = SessionConfig::new_with_ballista(); + // A bare leaf plan has no `ExchangeExec` anywhere; its only + // runnable stage caches the outer `AdaptiveDatafusionExec` as + // the final-stage wrapper, which is not a parking slot. + let plan: Arc = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![]], v_schema(), None).unwrap(), + ))); + let mut planner = AdaptivePlanner::try_from_plan(&config, plan, "err-path".into())?; + + let stages = planner.runnable_stages()?; + let stage_id = stages + .as_ref() + .and_then(|s| s.first()) + .map(|e| e.plan.stage_id()) + .expect("a runnable stage must exist"); + + let routing = RangeRepartitionRouting { + cuts: vec![0.0], + routing_expr: col("v", v_schema().as_ref()).unwrap(), + }; + let result = planner.set_repartition_routing(stage_id, routing); + assert!( + result.is_err(), + "must fail loud when there's no ExchangeExec to park on; got {result:?}" + ); + + Ok(()) +} + +/// Sibling of `routing_parks_when_range_repartition_is_plan_root`: +/// with an ordinary parent (`FilterExec`) above the range-repartition, +/// the exchange gets inserted and the cuts get parked. Position in +/// the plan should be the only variable between this and the +/// root-level test. +#[tokio::test] +async fn routing_parks_when_range_repartition_has_a_parent() +-> datafusion::error::Result<()> { + let config = SessionConfig::new_with_ballista(); + let mut planner = AdaptivePlanner::try_from_plan( + &config, + stats_over_urre_with_parent(), + "control-job".into(), + )?; + + let stages = planner.runnable_stages()?; + let stage_id = stages + .as_ref() + .and_then(|s| s.first()) + .map(|e| e.plan.stage_id()) + .expect("a runnable stage must exist"); + + let routing = RangeRepartitionRouting { + cuts: vec![15.0], + routing_expr: col("v", v_schema().as_ref()).unwrap(), + }; + planner.set_repartition_routing(stage_id, routing)?; + + let plan_str = format!( + "{}", + datafusion::physical_plan::displayable(planner.current_plan()).indent(true) + ); + assert!( + plan_str.contains("range_repartition_cuts=1"), + "control: the cuts must be parked on the boundary ExchangeExec" + ); + + Ok(()) +} diff --git a/ballista/scheduler/src/state/execution_graph.rs b/ballista/scheduler/src/state/execution_graph.rs index dd5704e769..49f3cab627 100644 --- a/ballista/scheduler/src/state/execution_graph.rs +++ b/ballista/scheduler/src/state/execution_graph.rs @@ -36,7 +36,8 @@ use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::job_status::Status; use ballista_core::serde::protobuf::{FailedJob, ShuffleWritePartition, job_status}; use ballista_core::serde::protobuf::{ - FailedTask, JobStatus, ResultLost, RunningJob, SuccessfulJob, TaskStatus, + FailedTask, JobStatus, ResultLost, RunningJob, SuccessfulJob, SuccessfulTask, + TaskStatus, }; use ballista_core::serde::protobuf::{RunningTask, task_status}; use ballista_core::serde::scheduler::{ @@ -47,6 +48,8 @@ use crate::display::print_stage_metrics; use crate::planner::DistributedPlanner; use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::scheduler_server::timestamp_millis; +use ballista_core::execution_plans::log_merged_runtime_stats; + use crate::state::execution_stage::RunningStage; pub(crate) use crate::state::execution_stage::{ ExecutionStage, ResolvedStage, StageOutput, TaskInfo, UnresolvedStage, @@ -950,12 +953,16 @@ impl ExecutionGraph for StaticExecutionGraph { running_stage .update_task_metrics(task_id, operator_metrics)?; + let SuccessfulTask { + partitions, + runtime_stats, + .. + } = successful_task; + running_stage + .append_runtime_stats_reports(task_id, runtime_stats); + locations.append(&mut partition_to_location( - &job_id, - task_id, - stage_id, - executor, - successful_task.partitions, + &job_id, task_id, stage_id, executor, partitions, )); } else { warn!( @@ -978,6 +985,11 @@ impl ExecutionGraph for StaticExecutionGraph { stage_metrics, ); } + log_merged_runtime_stats( + job_id.as_str(), + stage_id, + &running_stage.runtime_stats_reports, + ); } let output_links = running_stage.output_links.clone(); @@ -1283,13 +1295,10 @@ impl ExecutionGraph for StaticExecutionGraph { /// Convert unresolved stage to be resolved fn resolve_stage(&mut self, stage_id: usize) -> Result { - if let Some(ExecutionStage::UnResolved(stage)) = self.stages.remove(&stage_id) { - self.stages.insert( - stage_id, - ExecutionStage::Resolved( - stage.to_resolved(self.session_config.options())?, - ), - ); + if let Some(ExecutionStage::UnResolved(stage)) = self.stages.get(&stage_id) { + let resolved_stage = stage.to_resolved(self.session_config.options())?; + self.stages + .insert(stage_id, ExecutionStage::Resolved(resolved_stage)); Ok(true) } else { warn!( @@ -1341,7 +1350,7 @@ impl ExecutionGraph for StaticExecutionGraph { stage_id: usize, failure_reasons: HashSet, ) -> Result> { - if let Some(ExecutionStage::Running(stage)) = self.stages.remove(&stage_id) { + if let Some(ExecutionStage::Running(stage)) = self.stages.get(&stage_id) { let running_tasks = stage .running_tasks() .into_iter() @@ -1352,10 +1361,9 @@ impl ExecutionGraph for StaticExecutionGraph { executor_id, }) .collect(); - self.stages.insert( - stage_id, - ExecutionStage::UnResolved(stage.to_unresolved(failure_reasons)?), - ); + let unresolved_stage = stage.to_unresolved(failure_reasons)?; + self.stages + .insert(stage_id, ExecutionStage::UnResolved(unresolved_stage)); Ok(running_tasks) } else { warn!( @@ -1369,9 +1377,10 @@ impl ExecutionGraph for StaticExecutionGraph { /// Convert resolved stage to be unresolved fn rollback_resolved_stage(&mut self, stage_id: usize) -> Result { - if let Some(ExecutionStage::Resolved(stage)) = self.stages.remove(&stage_id) { + if let Some(ExecutionStage::Resolved(stage)) = self.stages.get(&stage_id) { + let unresolved_stage = stage.to_unresolved()?; self.stages - .insert(stage_id, ExecutionStage::UnResolved(stage.to_unresolved()?)); + .insert(stage_id, ExecutionStage::UnResolved(unresolved_stage)); Ok(true) } else { warn!( @@ -1807,6 +1816,7 @@ pub(crate) fn partition_to_location( #[cfg(test)] mod test { use std::collections::HashSet; + use std::sync::Arc; use crate::scheduler_server::event::QueryStageSchedulerEvent; use ballista_core::error::Result; @@ -1814,6 +1824,12 @@ mod test { self, ExecutionError, FailedTask, FetchPartitionError, IoError, JobStatus, TaskKilled, failed_task, job_status, task_status, }; + use datafusion::common::{DataFusionError, Result as DataFusionResult}; + use datafusion::execution::TaskContext; + use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + SendableRecordBatchStream, + }; use crate::state::execution_graph::ExecutionGraph; use crate::state::execution_stage::ExecutionStage; @@ -1825,6 +1841,60 @@ mod test { test_union_all_plan, test_union_plan, }; + #[derive(Debug)] + struct FailingPlanRewriteExec { + input: Arc, + } + + impl DisplayAs for FailingPlanRewriteExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "FailingPlanRewriteExec") + } + } + + impl ExecutionPlan for FailingPlanRewriteExec { + fn name(&self) -> &str { + "FailingPlanRewriteExec" + } + + fn properties(&self) -> &Arc { + self.input.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> DataFusionResult> { + Err(DataFusionError::Internal( + "forced plan rewrite failure".to_owned(), + )) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> DataFusionResult { + self.input.execute(partition, context) + } + } + + fn fail_plan_rewrites(plan: &mut Arc) -> Arc { + let failing_plan: Arc = Arc::new(FailingPlanRewriteExec { + input: Arc::clone(plan), + }); + *plan = Arc::clone(&failing_plan); + failing_plan + } + #[tokio::test] async fn test_intermediate_stage_ids() { // A simple aggregation produces a 2-stage graph: one intermediate @@ -1852,6 +1922,123 @@ mod test { } } + #[tokio::test] + async fn test_resolve_stage_preserves_stage_on_plan_rewrite_error() -> Result<()> { + let mut graph = test_aggregation_plan(4).await; + let stage_id = graph + .stages + .iter() + .find_map(|(stage_id, stage)| { + matches!(stage, ExecutionStage::UnResolved(_)).then_some(*stage_id) + }) + .expect("expected an unresolved stage"); + let stage_count = graph.stage_count(); + let original_plan = match graph.stages.get_mut(&stage_id) { + Some(ExecutionStage::UnResolved(stage)) => { + for input in stage.inputs.values_mut() { + input.complete = true; + } + assert!(stage.resolvable()); + fail_plan_rewrites(&mut stage.plan) + } + _ => unreachable!(), + }; + + assert!(graph.resolve_stage(stage_id).is_err()); + assert_eq!(graph.stage_count(), stage_count); + match graph.stages.get(&stage_id) { + Some(ExecutionStage::UnResolved(stage)) => { + assert!(Arc::ptr_eq(&stage.plan, &original_plan)); + } + _ => panic!("expected the original unresolved stage"), + } + + Ok(()) + } + + #[tokio::test] + async fn test_rollback_resolved_stage_preserves_stage_on_plan_rewrite_error() + -> Result<()> { + let mut graph = test_aggregation_plan(4).await; + revive_graph_and_complete_next_stage(&mut graph)?; + let stage_id = graph + .stages + .iter() + .find_map(|(stage_id, stage)| { + matches!(stage, ExecutionStage::Resolved(stage) if !stage.inputs.is_empty()) + .then_some(*stage_id) + }) + .expect("expected a resolved stage with inputs"); + let stage_count = graph.stage_count(); + let (original_plan, original_attempt) = match graph.stages.get_mut(&stage_id) { + Some(ExecutionStage::Resolved(stage)) => { + (fail_plan_rewrites(&mut stage.plan), stage.stage_attempt_num) + } + _ => unreachable!(), + }; + + assert!(graph.rollback_resolved_stage(stage_id).is_err()); + assert_eq!(graph.stage_count(), stage_count); + match graph.stages.get(&stage_id) { + Some(ExecutionStage::Resolved(stage)) => { + assert_eq!(stage.stage_attempt_num, original_attempt); + assert!(Arc::ptr_eq(&stage.plan, &original_plan)); + } + _ => panic!("expected the original resolved stage"), + } + + Ok(()) + } + + #[tokio::test] + async fn test_rollback_running_stage_preserves_stage_on_plan_rewrite_error() + -> Result<()> { + let mut graph = test_aggregation_plan(4).await; + revive_graph_and_complete_next_stage(&mut graph)?; + let stage_id = graph + .stages + .iter() + .find_map(|(stage_id, stage)| { + matches!(stage, ExecutionStage::Resolved(stage) if !stage.inputs.is_empty()) + .then_some(*stage_id) + }) + .expect("expected a resolved stage with inputs"); + assert!(graph.revive()); + assert!(graph.pop_next_task("executor-id")?.is_some()); + let stage_count = graph.stage_count(); + let (original_plan, original_attempt, original_pending, original_tasks) = + match graph.stages.get_mut(&stage_id) { + Some(ExecutionStage::Running(stage)) => ( + fail_plan_rewrites(&mut stage.plan), + stage.stage_attempt_num, + stage.available_tasks(), + stage.running_tasks(), + ), + _ => unreachable!(), + }; + + assert!( + graph + .rollback_running_stage( + stage_id, + HashSet::from(["executor-id".to_owned()]), + ) + .is_err() + ); + assert_eq!(graph.stage_count(), stage_count); + match graph.stages.get(&stage_id) { + Some(ExecutionStage::Running(stage)) => { + assert_eq!(stage.stage_attempt_num, original_attempt); + assert_eq!(stage.available_tasks(), original_pending); + assert_eq!(stage.running_tasks(), original_tasks); + assert!(Arc::ptr_eq(&stage.plan, &original_plan)); + } + _ => panic!("expected the original running stage"), + } + + Ok(()) + } + #[tokio::test] async fn test_fail_job_sets_end_time_and_failed_metadata() -> Result<()> { let mut graph = test_aggregation_plan(4).await; diff --git a/ballista/scheduler/src/state/execution_stage.rs b/ballista/scheduler/src/state/execution_stage.rs index 2b17a1e16d..4643e3efdd 100644 --- a/ballista/scheduler/src/state/execution_stage.rs +++ b/ballista/scheduler/src/state/execution_stage.rs @@ -32,10 +32,13 @@ use datafusion::prelude::SessionConfig; use log::{debug, warn}; use ballista_core::error::{BallistaError, Result}; -use ballista_core::execution_plans::{ShuffleWriterExec, SortShuffleWriterExec}; +use ballista_core::execution_plans::{ + ShuffleWriterExec, SortShuffleWriterExec, TaskRuntimeStats, +}; use ballista_core::serde::protobuf::failed_task::FailedReason; use ballista_core::serde::protobuf::{ - FailedTask, OperatorMetricsSet, ResultLost, SuccessfulTask, TaskKilled, TaskStatus, + FailedTask, OperatorMetricsSet, ResultLost, RuntimeStatsReport, SuccessfulTask, + TaskKilled, TaskStatus, }; use ballista_core::serde::protobuf::{RunningTask, task_status}; use ballista_core::serde::scheduler::PartitionLocation; @@ -218,6 +221,13 @@ pub struct RunningStage { pub stage_metrics: Option>, /// [SessionConfig] used for this stage pub session_config: Arc, + /// Per-task runtime-stats reports collected from every successful task + /// in this stage attempt. Each entry pairs the producer task's `task_id` + /// (== `file_id` on the emitted shuffle files) with its report — the pair + /// is what uniquely addresses a producer file, since the report's own + /// `partition_id` field is producer-local. Merged and logged once the + /// stage finalizes; dropped when the stage transitions to Successful. + pub runtime_stats_reports: Vec, } /// If a stage finishes successfully, its task statuses and metrics will be finalized @@ -630,6 +640,7 @@ impl RunningStage { task_failure_numbers: vec![0; partitions], stage_metrics: None, session_config, + runtime_stats_reports: Vec::new(), } } @@ -825,6 +836,28 @@ impl RunningStage { true } + /// Accumulate the `RuntimeStatsReport`s a successful task shipped + /// back in its `SuccessfulTask` payload, tagging each with the producer + /// `task_id` so downstream stages can address individual producer files. + /// Consumed at final-success by + /// `ballista_core::execution_plans::log_merged_runtime_stats` and by the + /// step-3 overlap router (`compute_overlapping_locations`); each stage + /// attempt starts with an empty accumulator. + pub fn append_runtime_stats_reports( + &mut self, + producer_task_id: usize, + reports: Vec, + ) { + if reports.is_empty() { + return; + } + self.runtime_stats_reports + .extend(reports.into_iter().map(|report| TaskRuntimeStats { + producer_task_id, + report, + })); + } + /// update and upsert the task metrics to the stage metrics pub fn update_task_metrics( &mut self, @@ -990,6 +1023,11 @@ impl RunningStage { /// Mark the task as lost/killed and push its partitions back to the /// front of `pending` so they are retried on the next bind. Does not /// touch failure counts — those are updated in `update_task_info`. + /// + /// Any runtime-stats reports the task already contributed are also + /// dropped: the retry will produce fresh sketches under a new + /// `task_id`, and leaving the ghost behind would double-count the + /// slice in the stage's merged view. pub fn reset_task_info(&mut self, task_id: usize) { let task = &mut self.task_infos[task_id]; let partitions = task.global_input_partition_ids.clone(); @@ -1000,16 +1038,25 @@ impl RunningStage { failed_reason: Some(FailedReason::TaskKilled(TaskKilled {})), }); self.pending.reschedule(partitions); + self.runtime_stats_reports + .retain(|s| s.producer_task_id != task_id); } /// Reset the running and completed tasks on a given executor by /// marking their `TaskInfo` as `Failed(ResultLost)` and pushing their /// partition slices back to `pending`. Returns the number of tasks /// reset. + /// + /// Runtime-stats reports contributed by any reset producer are + /// dropped — the retries will report fresh sketches under new + /// `task_id`s. Reports are append-only otherwise, so without this + /// purge the stage's merged view would double-count every slice that + /// bounced through a lost executor. pub fn reset_tasks(&mut self, executor: &str) -> usize { let mut reset = 0; let mut to_reschedule: Vec = vec![]; - for task in self.task_infos.iter_mut() { + let mut reset_task_ids: HashSet = HashSet::new(); + for (task_id, task) in self.task_infos.iter_mut().enumerate() { let matches_exec = match &task.task_status { task_status::Status::Running(RunningTask { executor_id }) | task_status::Status::Successful(SuccessfulTask { @@ -1025,10 +1072,13 @@ impl RunningStage { failed_reason: Some(FailedReason::ResultLost(ResultLost {})), }); to_reschedule.extend(task.global_input_partition_ids.iter().copied()); + reset_task_ids.insert(task_id); reset += 1; } } self.pending.reschedule(to_reschedule); + self.runtime_stats_reports + .retain(|s| !reset_task_ids.contains(&s.producer_task_id)); reset } @@ -1132,6 +1182,9 @@ impl SuccessfulStage { .duration_since(UNIX_EPOCH) .unwrap() .as_millis(), + // Fresh attempt: previous attempt's stats are irrelevant. + // Merged-cut logging fires per-attempt on final success. + runtime_stats_reports: Vec::new(), } } @@ -1351,6 +1404,7 @@ mod tests { status: Some(task_status::Status::Successful(SuccessfulTask { executor_id: "executor-1".to_string(), partitions: vec![], + runtime_stats: vec![], })), metrics: vec![], } @@ -1687,4 +1741,83 @@ mod tests { let aggregated = operator_metrics.aggregate_by_name(); assert_eq!(aggregated.output_rows(), Some(2000)); } + + /// Build a bare-bones report tagged with a `partition_id` we can look + /// for in assertions — no sketches, no order-by; the purge cares only + /// about the wrapping `producer_task_id`. + fn make_report(marker_partition_id: u32) -> RuntimeStatsReport { + RuntimeStatsReport { + order_by: vec![], + partitions: vec![ + ballista_core::serde::protobuf::RuntimeStatsPartitionEntry { + partition_id: marker_partition_id, + row_count: 0, + sketch: None, + }, + ], + } + } + + /// When a task is reset for retry, its previously-appended runtime-stats + /// reports must be dropped so the retry's fresh sketches don't merge + /// with ghost data from the original attempt. Reports from *other* tasks + /// must survive. + #[test] + fn test_reset_task_info_purges_runtime_stats_reports() { + let mut stage = make_running_stage(2); + append_running_task(&mut stage, 0, "executor-1", vec![0]); + append_running_task(&mut stage, 1, "executor-1", vec![1]); + + stage.append_runtime_stats_reports(0, vec![make_report(100)]); + stage.append_runtime_stats_reports(1, vec![make_report(200)]); + assert_eq!(stage.runtime_stats_reports.len(), 2); + + stage.reset_task_info(0); + + assert_eq!(stage.runtime_stats_reports.len(), 1); + assert_eq!(stage.runtime_stats_reports[0].producer_task_id, 1); + assert_eq!( + stage.runtime_stats_reports[0].report.partitions[0].partition_id, + 200 + ); + } + + /// Executor loss resets every task the executor was hosting; the + /// runtime-stats reports those (previously-Successful) producers had + /// already contributed must be purged along with the task status. + /// Reports produced on surviving executors must be left alone. + #[test] + fn test_reset_tasks_purges_runtime_stats_reports_for_lost_executor() { + let mut stage = make_running_stage(3); + append_running_task(&mut stage, 0, "executor-1", vec![0]); + append_running_task(&mut stage, 1, "executor-1", vec![1]); + append_running_task(&mut stage, 2, "executor-2", vec![2]); + // Drain the pending queue so reschedules below are visible. + stage.pending.next_slice(3); + + // All three tasks made it to Successful and shipped reports. + for (task_id, executor) in + [(0, "executor-1"), (1, "executor-1"), (2, "executor-2")] + { + stage.task_infos[task_id].task_status = + task_status::Status::Successful(SuccessfulTask { + executor_id: executor.to_string(), + partitions: vec![], + runtime_stats: vec![], + }); + stage.append_runtime_stats_reports( + task_id, + vec![make_report(100 + task_id as u32)], + ); + } + assert_eq!(stage.runtime_stats_reports.len(), 3); + + // Simulate executor-1 heartbeat loss. + let reset_count = stage.reset_tasks("executor-1"); + assert_eq!(reset_count, 2); + + // Only executor-2's producer survives. + assert_eq!(stage.runtime_stats_reports.len(), 1); + assert_eq!(stage.runtime_stats_reports[0].producer_task_id, 2); + } } diff --git a/ballista/scheduler/src/state/executor_manager.rs b/ballista/scheduler/src/state/executor_manager.rs index d8a8331778..1da7565662 100644 --- a/ballista/scheduler/src/state/executor_manager.rs +++ b/ballista/scheduler/src/state/executor_manager.rs @@ -25,7 +25,7 @@ use ballista_core::serde::protobuf::ExecutorMetric; use ballista_core::serde::protobuf::executor_metric::Metric; use log::trace; -use crate::cluster::{BoundTask, ClusterState, ExecutorSlot}; +use crate::cluster::{BoundTask, ClusterState, ClusterStateEventStream, ExecutorSlot}; use crate::config::SchedulerConfig; use crate::state::execution_graph::RunningTaskInfo; @@ -79,13 +79,22 @@ impl ExecutorManager { cluster_state: Arc, config: Arc, ) -> Self { + // Prefer an explicit override_config_producer if the embedder wired one, + // so a full BallistaConfig (with all its grpc-client knobs) still takes + // precedence. Otherwise, use `default()` but override + // `max_message_size` from the scheduler's `grpc_client_max_message_size` + // CLI flag so users can raise the ceiling for outbound task-assignment + // RPCs without having to write a config-producer in Rust. let grpc_client_config = if let Some(config_producer) = &config.override_config_producer { let session_config = config_producer(); let ballista_config = session_config.ballista_config(); GrpcClientConfig::from(&ballista_config) } else { - GrpcClientConfig::default() + GrpcClientConfig { + max_message_size: config.grpc_client_max_message_size as usize, + ..GrpcClientConfig::default() + } }; Self { cluster_state, @@ -103,6 +112,11 @@ impl ExecutorManager { Ok(()) } + /// Returns a stream of cluster state events from the configured state backend. + pub async fn cluster_state_events(&self) -> Result { + self.cluster_state.cluster_state_events().await + } + /// Binds ready-to-run tasks from active jobs to available executor slots. /// /// Returns a list of bound tasks that can be launched on executors. @@ -552,7 +566,13 @@ impl ExecutorManager { } let connection = endpoint.connect().await?; - let client = ExecutorGrpcClient::new(connection); + // Message-size limits are tonic codec settings, not `Endpoint` + // settings, so `create_grpc_client_endpoint` cannot apply them. + // Without this the configured `max_message_size` is silently + // ignored and task assignment falls back to tonic's own defaults. + let client = ExecutorGrpcClient::new(connection) + .max_encoding_message_size(grpc_client_config.max_message_size) + .max_decoding_message_size(grpc_client_config.max_message_size); { self.clients.insert(executor_id.to_owned(), client.clone()); @@ -581,3 +601,49 @@ impl ExecutorManager { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::test_utils::test_cluster_context; + use ballista_core::extension::SessionConfigExt; + use datafusion::prelude::SessionConfig; + + #[test] + fn grpc_client_max_message_size_flag_reaches_client_config() { + let config = Arc::new( + SchedulerConfig::default() + .with_grpc_client_max_message_size(64 * 1024 * 1024), + ); + let manager = + ExecutorManager::new(test_cluster_context().cluster_state(), config); + + assert_eq!( + manager.grpc_client_config.max_message_size, + 64 * 1024 * 1024 + ); + } + + #[test] + fn config_producer_still_wins_over_the_flag() { + let config = Arc::new( + SchedulerConfig::default() + .with_grpc_client_max_message_size(64 * 1024 * 1024) + .with_override_config_producer(Arc::new(|| { + let mut session_config = SessionConfig::new_with_ballista(); + session_config + .options_mut() + .set("ballista.client.grpc_max_message_size", "33554432") + .expect("valid setting"); + session_config + })), + ); + let manager = + ExecutorManager::new(test_cluster_context().cluster_state(), config); + + assert_eq!( + manager.grpc_client_config.max_message_size, + 32 * 1024 * 1024 + ); + } +} diff --git a/ballista/scheduler/src/state/task_builder.rs b/ballista/scheduler/src/state/task_builder.rs index 37da2865c3..68811762e0 100644 --- a/ballista/scheduler/src/state/task_builder.rs +++ b/ballista/scheduler/src/state/task_builder.rs @@ -36,7 +36,8 @@ //! flows from parent to descendants via function arguments, so sibling //! subtrees never share state and there's no traversal-order dependency. -use ballista_core::execution_plans::ShuffleReaderExec; +use ballista_core::execution_plans::{PerPartitionFilterExec, ShuffleReaderExec}; +use datafusion::common::internal_err; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::{ FileGroup, FileScanConfig, FileScanConfigBuilder, @@ -46,7 +47,6 @@ use datafusion::error::Result; use datafusion::physical_expr::Distribution; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::empty::EmptyExec; -use datafusion::physical_plan::joins::{HashJoinExec, NestedLoopJoinExec}; use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::union::UnionExec; @@ -84,6 +84,30 @@ fn restrict( return Ok(rewritten); } + // PerPartitionFilterExec: its `predicates` vec is positionally aligned + // with the child's output partitions (predicates[k] filters + // input.execute(k)). When we restrict the child from K partitions to + // `partitions.len()`, the predicate vec must be sliced by the same + // indices in the same order + if !under_collect && let Some(ppf) = plan.downcast_ref::() { + let children = plan.children(); + let [child] = children.as_slice() else { + return internal_err!( + "PerPartitionFilterExec must have exactly 1 child, got {}", + children.len() + ); + }; + let new_child = restrict((*child).clone(), partitions, false)?; + let new_predicates: Vec<_> = partitions + .iter() + .map(|&part_idx| ppf.predicates()[part_idx].clone()) + .collect(); + return Ok(Arc::new(PerPartitionFilterExec::try_new( + new_child, + new_predicates, + )?)); + } + // UnionExec: parent partition `p` maps to exactly one child's local // partition (`p` minus the sum of preceding children's counts). Split // `partitions` into disjoint per-child sub-slices and recurse; a child @@ -148,11 +172,24 @@ fn union_child_partitions( /// every child (it's sticky — descendants of a collapse all read /// everything). /// - `CoalescePartitionsExec` / `SortPreservingMergeExec`: set for all -/// children. -/// - `HashJoinExec` / `NestedLoopJoinExec`: set per child based on -/// `required_input_distribution()`. Typically only the build side is -/// `SinglePartition`, so the probe side inherits `false` and remains -/// partition-aligned. +/// children. These cannot be derived from the rule below — both declare +/// `UnspecifiedDistribution`, because they consume every input partition +/// without constraining how the input is partitioned. Keep the explicit +/// case. +/// - Anything else: set per child from `required_input_distribution()`. An +/// input declared `SinglePartition` is one the operator consumes whole — +/// a join or cross-join build side, for instance — so restricting it to +/// the task's slice would hand each sibling task a different fraction of +/// it. Typically only the build side is `SinglePartition`, so the probe +/// side inherits `false` and remains partition-aligned. +/// +/// This is asked of every operator rather than a list of join types. +/// `CrossJoinExec` also collects its left input but was not on that list, +/// so each task of a multi-partition stage rebuilt the cross-join's left +/// side from its own slice and the result came back inflated (seen on +/// TPC-DS q77, whose catalog branch is `from cs, cr`). It also brings in +/// a non-preserving `SortExec` and `GlobalLimitExec`, which declare +/// `SinglePartition` for the same reason. /// /// `UnionExec` is handled by the caller before reaching this function — /// its children need per-child *partition* sub-slices, not just per-child @@ -165,14 +202,20 @@ fn child_scopes(plan: &Arc, under_collect: bool) -> Vec if plan.is::() || plan.is::() { return vec![true; children.len()]; } - if plan.is::() || plan.is::() { - return plan - .required_input_distribution() - .into_iter() - .map(|d| matches!(d, Distribution::SinglePartition)) - .collect(); + let required = plan.required_input_distribution(); + if required.len() != children.len() { + // The trait contract is one entry per child, so a mismatch means a + // broken operator. Stay partition-aligned rather than reading whole: + // `true` here is not the safe direction, it makes every task of the + // stage read the entire input, which duplicates the data instead of + // protecting it. `false` is also what every operator outside the + // cases above got before this rule generalized. + return vec![false; children.len()]; } - vec![false; children.len()] + required + .into_iter() + .map(|d| matches!(d, Distribution::SinglePartition)) + .collect() } /// Restrict a plan node to a subset of its output partitions, re-indexed @@ -480,31 +523,9 @@ mod tests { // upstream fix used a per-single-partition helper on the executor side; // this suite ports the same invariants to the scheduler-side rewriter. - use datafusion::arrow::datatypes::SchemaRef; - use datafusion::datasource::listing::PartitionedFile; - use datafusion::datasource::physical_plan::ParquetSource; - use datafusion::execution::object_store::ObjectStoreUrl; + use crate::test_utils::scan_with_file_groups; use datafusion::physical_plan::union::UnionExec; - /// Build a `DataSourceExec` over `n` file groups (one file each), so - /// every group is exactly one partition. The scan's output partition - /// count equals `n`. - fn scan_with_file_groups(n: usize) -> Arc { - let schema: SchemaRef = - 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()) - } - /// File counts per file group of a `DataSourceExec` — the shape a union /// test cares about after restriction. fn group_file_counts(plan: &Arc) -> Vec { @@ -588,4 +609,113 @@ mod tests { let restricted = restrict_plan_to_partitions(plan, &[2]).unwrap(); assert_eq!(group_file_counts(&restricted), vec![1]); } + + /// A `CrossJoinExec` collects its left input, so that side must be read + /// whole. Scoping is taken from `required_input_distribution()` rather + /// than a list of join types, which is what brings cross joins in. + /// + /// Restricting the collected side would give each sibling task a + /// different fraction of it, and the join would emit a different subset + /// of the cartesian product per task — TPC-DS q77's `from cs, cr` branch + /// came back inflated exactly this way. + #[test] + fn cross_join_left_side_is_read_whole() { + use datafusion::physical_plan::joins::CrossJoinExec; + + let join: Arc = Arc::new(CrossJoinExec::new( + scan_with_file_groups(4), + scan_with_file_groups(4), + )); + + // Sanity: the operator really does declare its left input collected. + assert!( + matches!( + join.required_input_distribution().first(), + Some(Distribution::SinglePartition) + ), + "CrossJoinExec must declare SinglePartition on its left input" + ); + + let restricted = restrict_plan_to_partitions(join, &[1]).unwrap(); + let children = restricted.children(); + assert_eq!( + group_file_counts(children[0]), + vec![1, 1, 1, 1], + "collected left side keeps every group" + ); + assert_eq!( + group_file_counts(children[1]), + vec![1], + "streaming right side is pinned to this task's partition" + ); + } + + /// A `PerPartitionFilterExec` restricted to a subset of partitions must + /// slice its `predicates` vector by the same indices, in the same order, + /// as its child. Otherwise the operator's construction invariant + /// (`predicates.len() == child.partition_count()`) breaks and + /// task-local partition `j` would filter through a global predicate + /// that no longer matches. + #[test] + fn per_partition_filter_predicates_are_sliced_with_partitions() { + use ballista_core::execution_plans::PerPartitionFilterExec; + use datafusion::logical_expr::Operator; + use datafusion::physical_expr::expressions::{BinaryExpr, Column, Literal}; + use datafusion::physical_expr::{Partitioning, PhysicalExpr}; + use datafusion::scalar::ScalarValue; + + // 4 upstream partitions, each with its own bespoke predicate so we + // can assert the slice ordering survives. + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let partitions_locs: Vec> = + (0..4).map(|i| vec![create_partition(i)]).collect(); + let reader = ShuffleReaderExec::try_new( + 1, + partitions_locs, + schema.clone(), + Partitioning::UnknownPartitioning(4), + ) + .unwrap(); + let make_pred = |lo: i64| -> Arc { + Arc::new(BinaryExpr::new( + Arc::new(Column::new("v", 0)), + Operator::GtEq, + Arc::new(Literal::new(ScalarValue::Int64(Some(lo)))), + )) + }; + let predicates: Vec> = + (0..4).map(|i| make_pred(i as i64 * 100)).collect(); + let plan: Arc = Arc::new( + PerPartitionFilterExec::try_new( + Arc::new(reader) as Arc, + predicates.clone(), + ) + .unwrap(), + ); + + let restricted = restrict_plan_to_partitions(plan, &[1, 3]).unwrap(); + let ppf = restricted + .downcast_ref::() + .expect("top must remain PerPartitionFilterExec"); + assert_eq!(ppf.predicates().len(), 2); + assert_eq!( + ppf.predicates()[0].to_string(), + predicates[1].to_string(), + "local partition 0 must carry the global-partition-1 predicate" + ); + assert_eq!( + ppf.predicates()[1].to_string(), + predicates[3].to_string(), + "local partition 1 must carry the global-partition-3 predicate" + ); + + // Reader below must have been restricted in the same order. + let child = ppf.children()[0].clone(); + let reader = child + .downcast_ref::() + .expect("child must be a ShuffleReaderExec"); + assert_eq!(reader.partition.len(), 2); + assert_eq!(reader.partition[0][0].partition_id.partition_id, 1); + assert_eq!(reader.partition[1][0].partition_id.partition_id, 3); + } } diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index fe84ef252b..dea5af1004 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::cluster::JobState; +use crate::cluster::{JobState, JobStateEventStream}; use crate::config::SchedulerConfig; use crate::planner::DefaultDistributedPlanner; use crate::scheduler_server::event::{QueryStageSchedulerEvent, SubmitPlan}; @@ -243,6 +243,11 @@ impl TaskManager .count() } + /// Returns a stream of job state events from the configured state backend. + pub async fn job_state_events(&self) -> Result { + self.state.job_state_events().await + } + /// Generate an ExecutionGraph for the job and save it to the persistent state. /// By default, this job will be curated by the scheduler which receives it. /// Then we will also save it to the active execution graph @@ -550,6 +555,7 @@ impl TaskManager let mut job_updates: HashMap> = HashMap::new(); for status in task_status { trace!("Task Update\n{status:?}"); + log_runtime_stats_arrival(executor, &status); let job_id = status.job_id.clone(); let job_task_statuses = job_updates.entry(job_id).or_default(); job_task_statuses.push(status); @@ -999,6 +1005,49 @@ impl From<&ExecutionGraphBox> for JobOverview { } } +/// Log any `RuntimeStatsReport`s that arrived with this task status. For +/// now this is proof-of-life for the wire — the reports flow from +/// executor to scheduler and land here observably. The per-stage +/// accumulator + merged-quantile-cut logging comes in a follow-up commit; +/// this hook is what verifies the plumbing works against a live cluster. +fn log_runtime_stats_arrival( + executor: &ExecutorMetadata, + status: &ballista_core::serde::protobuf::TaskStatus, +) { + use ballista_core::serde::protobuf::task_status::Status; + let Some(Status::Successful(successful)) = status.status.as_ref() else { + return; + }; + if successful.runtime_stats.is_empty() { + return; + } + for (report_idx, report) in successful.runtime_stats.iter().enumerate() { + let non_empty_partitions = + report.partitions.iter().filter(|p| p.row_count > 0).count(); + let total_rows: u64 = report.partitions.iter().map(|p| p.row_count).sum(); + let sketch_count = report + .partitions + .iter() + .filter(|p| p.sketch.is_some()) + .count(); + debug!( + "RuntimeStats arrival: executor={} job={} stage={} task={} \ + report[{}] order_by_len={} partitions={} non_empty={} \ + total_rows={} sketches={}", + executor.id, + status.job_id, + status.stage_id, + status.task_id, + report_idx, + report.order_by.len(), + report.partitions.len(), + non_empty_partitions, + total_rows, + sketch_count, + ); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index be184cd6e6..496aa8d587 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -53,7 +53,7 @@ use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::prelude::{CsvReadOptions, JoinType, col}; use datafusion::test_util::scan_empty_with_partitions; -use crate::cluster::BallistaCluster; +use crate::cluster::{BallistaCluster, JobStateEventStream}; use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::state::execution_graph::{ @@ -308,6 +308,7 @@ pub fn default_task_runner() -> impl TaskRunner { status: Some(task_status::Status::Successful(SuccessfulTask { executor_id: executor_id.clone(), partitions: partitions.clone(), + runtime_stats: vec![], })), }); } @@ -478,6 +479,11 @@ impl SchedulerTest { self.scheduler.running_job_number() } + /// Returns job state events from the underlying scheduler. + pub async fn job_state_events(&self) -> Result { + self.scheduler.job_state_events().await + } + /// Returns the session context for tests. pub async fn ctx(&self) -> Result> { self.scheduler @@ -542,6 +548,32 @@ impl SchedulerTest { .await } + /// Simulates the loss of an executor: deregisters it from the executor + /// manager and posts the `ExecutorLost` event. This mirrors the reaper's + /// `remove_executor` path without waiting out the heartbeat timeout. + pub async fn lose_executor(&self, executor_id: &str) -> Result<()> { + let reason = Some("test: executor lost".to_owned()); + self.scheduler + .state + .executor_manager + .remove_executor(executor_id, reason.clone()) + .await?; + self.post_scheduler_event(QueryStageSchedulerEvent::ExecutorLost( + executor_id.to_owned(), + reason, + )) + .await + } + + /// Returns the current status of a job, if known. + pub async fn job_status(&self, job_id: &JobId) -> Result> { + self.scheduler + .state + .task_manager + .get_job_status(job_id) + .await + } + /// Waits for job completion with a timeout in milliseconds. pub async fn await_completion_timeout( &self, @@ -1209,6 +1241,7 @@ pub fn mock_completed_task(task: TaskDescription, executor_id: &str) -> TaskStat status: Some(task_status::Status::Successful(protobuf::SuccessfulTask { executor_id: executor_id.to_owned(), partitions, + runtime_stats: vec![], })), } } @@ -1243,3 +1276,30 @@ pub fn mock_failed_task(task: TaskDescription, failed_task: FailedTask) -> TaskS status: Some(task_status::Status::Failed(failed_task)), } } + +/// A `DataSourceExec` over `n` single-file groups, so its output partition +/// count is `n` and each group is independently restrictable. +/// +/// Shared by the planner and task-builder tests, both of which need a leaf +/// whose partitions per-task restriction can actually slice. +pub fn scan_with_file_groups(n: usize) -> Arc { + use datafusion::datasource::listing::PartitionedFile; + use datafusion::datasource::physical_plan::{ + FileGroup, FileScanConfigBuilder, ParquetSource, + }; + use datafusion::datasource::source::DataSourceExec; + use datafusion::execution::object_store::ObjectStoreUrl; + + let schema: SchemaRef = + 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()) +} diff --git a/chaos-testing/Cargo.toml b/chaos-testing/Cargo.toml new file mode 100644 index 0000000000..7b03cb0e22 --- /dev/null +++ b/chaos-testing/Cargo.toml @@ -0,0 +1,55 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "ballista-chaos" +description = "Fault-injection harness for testing Ballista high availability" +license = "Apache-2.0" +version = "0.1.0" +edition = { workspace = true } +rust-version = { workspace = true } +publish = false + +[dependencies] +arrow = { workspace = true } +ballista = { path = "../ballista/client" } +ballista-core = { path = "../ballista/core" } +ballista-executor = { path = "../ballista/executor" } +ballista-scheduler = { path = "../ballista/scheduler" } +datafusion = { workspace = true } +env_logger = { workspace = true } +log = { workspace = true } +nix = { version = "0.31", features = ["fs", "signal"] } +reqwest = { version = "0.13", default-features = false, features = ["json"] } +serde_json = "1.0" +tempfile = { workspace = true } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "process", "sync", "time"] } + +[dev-dependencies] +rstest = { workspace = true } + +[lib] +name = "chaos_testing" +path = "src/lib.rs" + +[[bin]] +name = "chaos-executor" +path = "src/bin/chaos-executor.rs" + +[[bin]] +name = "chaos-scheduler" +path = "src/bin/chaos-scheduler.rs" diff --git a/chaos-testing/README.md b/chaos-testing/README.md new file mode 100644 index 0000000000..35cc551d2a --- /dev/null +++ b/chaos-testing/README.md @@ -0,0 +1,351 @@ + + +# ballista-chaos + +A fault-injection harness that runs real, multi-process Ballista clusters and +injects faults into real queries, to exercise Ballista's high-availability +(HA) machinery end to end. + +**This is a bug-hunting harness, not a regression suite in the usual sense.** +Its job is to surface real defects in Ballista's HA behavior. Where it finds +one, the corresponding test reproduces the bug rather than working around it. +Such a test is marked `#[ignore]` with the issue it reproduces, so that it +does not hold CI red on a bug it did not introduce, and is un-ignored — not +rewritten — when that issue is fixed, at which point it becomes the regression +test for the fix. Run them with `cargo test -p ballista-chaos -- --ignored`. +See [Findings](#findings) below for the confirmed bugs this harness has +found so far, each with the test that reproduces it. + +## Why this crate exists + +Ballista's HA state machine — stage/task retry, executor-loss recovery, +map-stage resubmission — lives in +`ballista/scheduler/src/state/execution_graph.rs`. Before this crate, that code +was exercised only by unit tests that hand-construct `TaskStatus` protobufs and +feed them directly into `ExecutionGraph` methods. Those tests are useful for +pinning the state machine's transition logic, but nothing drove it end to end: +no test ran a real query against a real multi-process cluster, killed a real +executor process, and checked that the _result_ was still correct. That gap is +exactly where the bugs in [Findings](#findings) were hiding — they only show up +when a real executor process dies mid-task, a real gRPC connection is refused, +or a real DataFusion error is really propagated through the real serialization +path, none of which a hand-built `TaskStatus` reproduces. + +`ballista-chaos` closes that gap: it spawns a real `ballista-scheduler` and +one or more real `ballista-executor` processes, runs a real multi-stage query +against them through the `ballista` client, and injects faults or kills +processes while the query is in flight. + +## Why fault injection uses UDFs, not `ChaosExec` + +Ballista's AQE planner already has a fault-injection mechanism: +`ChaosCreatingRule` (`ballista/scheduler/src/state/aqe/planner.rs:542`), which +wraps a plan node in `ChaosExec` when `chaos_execution_enabled` is set. It was +deliberately not reused here, because it cannot do what this harness needs: + +- It is wired into the **AQE physical-optimizer pipeline only** + (`plan_preparation_optimizers` in `planner.rs`). It does not run at all when + AQE is off, and every scenario in this crate must run under _both_ AQE + settings — the two planners have materially different join and retry + behavior, and a bug that only reproduces on one side is easy to miss if you + only test the other. +- It picks a **uniformly random plan node** to wrap + (`ballista/scheduler/src/state/aqe/optimizer_rule/chaos_exec.rs`), not a node + the test chooses. A scenario that wants to fault "the scan of `facts`" or + "the shared join build side" specifically has no way to target it. +- It fires **probabilistically** (`chaos_execution_probability`), not + deterministically. A test built on it would need to loop-and-retry until the + fault happened to fire the right number of times, which is exactly the kind + of flakiness this harness is trying to avoid introducing. + +A SQL-level UDF (`chaos_fail`, `chaos_delay`, in `src/udf.rs`) sidesteps all +three problems: it lives in the query text itself, so it plans identically +(modulo AQE's own re-planning) whether AQE is on or off; its `guard` argument +lets a scenario target specific rows (and therefore specific partitions/tasks) +by writing an ordinary predicate; and it fires on every row where the guard is +true, subject only to the fault budget below — no probability, no retries of +the test itself. + +## How determinism works + +Every chaos scenario needs two things to be true: which rows/tasks fault must +be controlled, and how many attempts fault (across the whole cluster, across +retries and executor restarts) must be bounded. Two mechanisms provide these: + +- **The `guard` predicate.** The fixture (`src/fixture.rs`) is a small, fully + deterministic dataset: `facts(key, value)` joined to `dims(key, name)`, with + a known key distribution. A scenario passes a boolean expression over that + data as `chaos_fail`'s/`chaos_delay`'s first argument (e.g. `f.key = 7`); + since the data is fixed, this expression deterministically selects which + partitions the fault can fire in. +- **The filesystem fault budget** (`src/budget.rs`). A budget is a directory + of token files, created with a fixed token count. Consuming a token is + `fs::remove_file`, which is atomic across processes, so a budget of `n` + bounds the fault to firing at most `n` times _cluster-wide_ — across every + executor process, every task attempt, and every retry or restart — not `n` + times per process or per attempt. This is what makes "exactly one retryable + fault, then it must succeed" (Scenario A) and "faults never stop, so retries + must exhaust" (Scenario B) both expressible and deterministic. + +## The `OR TRUE` trap + +`Fixture::chaos_query` splices the injection expression into the query as +`WHERE {injection} IS NOT NULL`, not the more obvious `WHERE {injection} OR +TRUE`. This is deliberate and load-bearing: DataFusion's optimizer +constant-folds `expr OR TRUE` to the literal `TRUE` during logical +optimization, and once the predicate is a literal, the plan no longer +references the UDF call at all — it is dropped, not merely skipped. Every +fault-injection scenario built on `OR TRUE` would silently become a no-op: the +budget would never be consumed, the fault would never fire, and the suite +would report green while testing nothing. + +`chaos_fail`/`chaos_delay` always return `Some(guard)` (never `NULL`), so +`... IS NOT NULL` is always true but is not foldable to a constant without +evaluating the call — the optimizer has no way to know the result is always +non-null without invoking the (volatile) UDF. Two regression tests in +`src/fixture.rs` pin this: + +- `or_true_predicate_is_optimized_away_and_never_fires` proves the bad form + is eliminated from the plan and never consumes a budget token — pinning the + trap so it cannot silently return if someone "simplifies" the predicate back + to `OR TRUE`. +- `chaos_query_predicate_survives_optimization_and_fires` proves the + `IS NOT NULL` form the harness actually uses survives into the physical plan + and does fire. + +## How to run + +```sh +cargo test -p ballista-chaos # everything except the known-bug scenarios +cargo test -p ballista-chaos -- --ignored # the known-bug scenarios; these fail, on purpose +``` + +Every test that spawns a cluster does so through `TestCluster`, which holds +two locks for the cluster's lifetime: a process-wide mutex, and a +machine-wide `flock` (on a file in the system temp dir) for cluster-spawning +processes no in-process lock can see — a second cargo invocation in another +shell, or a runner like nextest that parallelizes test binaries. So scenarios +serialize themselves no matter how the test harness is invoked. This is not +cosmetic: each one starts a whole scheduler-plus-executors cluster, and +concurrent clusters exhaust ports and CPU and fail for reasons unrelated to +the scenario under test. `--test-threads=1` is therefore no longer required +(it will simply make the run marginally less confusing to read). + +CI runners are slow enough that cluster startup alone has blown a 30-second +registration deadline ("timed out waiting for 2 executors to register"); the +deadline is now 120s, and on expiry the error message carries the tail of +every child process log, so a recurrence in CI is diagnosable from the test +output alone. + +The `chaos-scheduler`/`chaos-executor` binaries are spawned as real child +processes rather than run in-process, but `cargo test` builds this crate's bin +targets along with its tests, so no separate build step is needed. The +harness locates them next to the running test executable, which is what makes +it work under any cargo profile (CI uses `--profile ci`, not `dev` or +`release`). + +Unit tests only (fast; the ones that do not spawn a cluster): + +```sh +cargo test -p ballista-chaos --lib +``` + +Each cluster's child-process logs (`scheduler.log`, `executor-0.log`, ...) are +written under that cluster's own temp directory, in a `logs/` subdirectory +(`TestCluster::log_dir()`). When a scenario fails, those logs are the first +place to look for what the scheduler and executors were actually doing. + +## Scenarios + +Every scenario runs under both `ballista.planner.adaptive.enabled=false` (AQE +off, the default, static `DefaultDistributedPlanner`) and `=true` (AQE on, the +experimental dynamic-join-selection planner) — 14 test cases total across the +7 scenarios below, plus a non-lettered `baseline_matches_local_datafusion` +sanity check that every other scenario's assertions depend on. + +| Scenario | Test | What it does | Expected result | +| -------- | ------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------- | +| A | `retryable_fault_is_retried_and_result_is_correct_{aqe_off,aqe_on}` | Injects one retryable IO fault (budget 1); the retry must succeed and match baseline. | **Ignored (both), reproduces the error flattening tracked by [#2027](https://github.com/apache/datafusion-ballista/issues/2027) — Finding 2.** | +| B | `exhausted_retries_fail_the_job_and_leave_the_cluster_healthy` | Injects an inexhaustible IO fault (budget 99 ≫ `task_max_failures`); job must fail, cluster must stay usable after. | Pass (both). | +| C | `panicking_task_fails_the_job_but_the_executor_survives` | Injects a task panic; job must fail non-retryably, both executor processes must survive, cluster must stay usable after. | Pass (both). | +| D | `executor_killed_mid_stage_is_recovered` | SIGKILLs an executor while its tasks are genuinely running (held open by `chaos_delay`); scheduler must reschedule onto the survivor and return the correct result. | **Ignored (both), reproduces [#2027](https://github.com/apache/datafusion-ballista/issues/2027) — Finding 1.** | +| E | `executor_killed_after_shuffle_write_is_recovered` | SIGKILLs the map-side executor _after_ it wrote shuffle output, with a long executor timeout to bias toward the fetch-failure path rather than heartbeat expiry; downstream stage must re-run the map stage. | **Ignored (both), reproduces [#2027](https://github.com/apache/datafusion-ballista/issues/2027) — Finding 1.** | +| F | `restarted_executor_rejoins_and_serves_queries` | Kills an executor, waits for the scheduler to reap it, restarts it, asserts the registered count returns to 2 and the cluster still serves the baseline query. | Pass (both), after the race fix in this crate (see below). | +| G | `killing_every_executor_terminates_the_job` | SIGKILLs every executor mid-query; asserts the job fails with an error naming the executor loss rather than hanging. | Regression test for [#2029](https://github.com/apache/datafusion-ballista/issues/2029) — Finding 3. | + +An ignored scenario above is not a defect in this harness, and its assertions +have not been weakened to make it pass — it reproduces a real Ballista bug and +is ignored only so that CI is not red on a bug this crate did not introduce. +Run the ignored scenarios with `-- --ignored --test-threads=1` to see the +failures; see [Findings](#findings) for what each one proves. + +### A note on Scenario F: the harness race that was fixed here + +`restarted_executor_rejoins_and_serves_queries` used to kill executor 0 and +restart it immediately, then assert `registered_executors() == 2`. That is a +harness bug, not a Ballista bug: SIGKILL does not deregister the executor, so +the scheduler keeps listing it until its heartbeat times out +(`executor_timeout_seconds`, 5s in this harness's defaults). Restarting +immediately races the scheduler's own reap: the assertion could observe three +executors (the dead one, the untouched survivor, and the freshly restarted +one) depending on timing, and failed intermittently with `left: 3, right: 2`. +Ballista was behaving correctly; the test just hadn't waited for the state it +was asserting about. The fix adds `TestCluster::await_executor_count(n)` (an +exact-count analogue of the existing `await_executors(n)`, which only waits for +_at least_ `n` — the right primitive for growing a cluster, but not for +observing a shrink) and has the scenario wait for the count to drop to 1 +before restarting, so the final assertion tests what the scenario name +actually promises. + +## Findings + +Four scenarios above fail because they have found real bugs in Ballista, not +because the harness is broken. Each is `#[ignore]`d against the issue it +reproduces so CI stays green on a tree whose bugs predate this crate, and each +keeps its original assertions: nothing is relaxed to manufacture a pass. When +the issue is fixed, delete the `#[ignore]` — the scenario is then the +regression test for it. One of the findings (#2028) has already gone through +that cycle in reverse: it was fixed on main, but a concurrent refactor +re-broke the same scenario through a different mechanism — see Finding 2. + +### Finding 1 — Shuffle-fetch failures lose their type, so the map-stage resubmit never fires + +Tracked by [#2027](https://github.com/apache/datafusion-ballista/issues/2027). + +**Proven by:** Scenario D (`executor_killed_mid_stage_is_recovered`) and +Scenario E (`executor_killed_after_shuffle_write_is_recovered`), both AQE +settings, `#[ignore]`d against that issue. + +Scenarios D and E are **races**, not deterministic reproducers, and this is +the one place in the crate where that is true. Killing an executor can be +noticed by the scheduler in either of two ways, and they are in a footrace: if +the heartbeat expires first, the `ExecutorLost` path recovers the job +correctly and the scenario passes in a few seconds; if a downstream task tries +to fetch shuffle output from the dead executor first, the bug below bites and +the job fails (or hangs until the scenario's timeout). Scenario D failed +locally on two of three runs; Scenario E passed locally but failed in CI +(`aqe_off`, with the reduce stage reporting a `FetchFailed` flattened inside +`DataFusionError::Shared` — see the log excerpt in #2027). The CI failure also +settles which way `executor_timeout_seconds` biases the race: raising it to +60s (as Scenario E does) _delays_ heartbeat expiry, so the downstream fetch +hits the dead executor first and the scenario exercises the broken +fetch-failure path; it only passes when the kill happens to land after the +reduce tasks already fetched their input. Un-ignoring these scenarios once +#2027 is fixed therefore also means pinning which of the two paths each +exercises — `executor_timeout_seconds` is the knob that decides the race, and +Scenario D currently leaves it at the harness default — otherwise they will be +flaky regression tests. + +The shuffle reader (`ballista/core/src/execution_plans/shuffle_reader.rs`) +correctly produces a typed `BallistaError::FetchFailed(executor_id, +map_stage_id, map_partition_id, desc)` when it cannot reach a dead executor, +and `ballista/core/src/error.rs`'s `impl From for FailedTask` +has a dedicated arm for exactly that variant (around line 205) which produces +`FailedReason::FetchPartitionError` — the signal +`ballista/scheduler/src/state/execution_graph.rs` (around line 826) uses to +resubmit the lost map stage rather than simply failing the job. + +The type does not survive to that point, however. Two real, non-test code +paths erase it before the executor reports its `TaskStatus`: + +- `ballista/executor/src/executor.rs:238-239`, in + `Executor::execute_query_stage`, converts the stage's result with + `result.map_err(|e| BallistaError::DataFusionError(Box::new(e)))` instead of + `BallistaError::from(e)` / `e.into()`. That bypasses the very unwrapping + logic `error.rs`'s `impl From for BallistaError` exists to + provide (`DataFusionError::ArrowError(e, _) => Self::from(*e)`, which would + otherwise recover a `FetchFailed` wrapped inside an `ArrowError::ExternalError`). +- `ballista/core/src/execution_plans/shuffle_writer.rs:245`, in + `ShuffleWriterExec`'s unpartitioned write branch, Debug-formats a + `BallistaError` into an opaque `DataFusionError::Execution(format!("{e:?}"))` + — a conversion that can never be undone by any later `.into()`, because the + original variant no longer exists, only its printed form. + +Either path leaves the executor reporting something like +`BallistaError::DataFusionError(Execution("FetchFailed(\"\", ..., +\"...Connection refused...\")"))` — the `FetchFailed` information is present +only as inert text inside a string. `error.rs`'s `FetchFailed` arm cannot match +a `DataFusionError::Execution`, so the task falls to the catch-all arm +(around line 248) and is marked `retryable: false`, `FailedReason::ExecutionError`. + +**Net effect:** when an executor dies after producing shuffle output that a +downstream stage still needs, Ballista fails the whole query instead of +re-running the map stage that produced it. + +### Finding 2 — Retryable IO errors are misclassified because the shuffle writer flattens them + +Originally tracked by +[#2028](https://github.com/apache/datafusion-ballista/issues/2028) (now +fixed); the surviving flattening mechanism is the one +[#2027](https://github.com/apache/datafusion-ballista/issues/2027) tracks. + +**Proven by:** Scenario A, both cases +(`retryable_fault_is_retried_and_result_is_correct_{aqe_off,aqe_on}`), +`#[ignore]`d against #2027. + +The history matters here because the failure mode moved underneath the +harness. As first found, only the `aqe_on` case failed: an `IoError` raised on +a join's shared broadcast build side arrived wrapped as +`DataFusionError::Shared(Arc)`, and the classifier in +`ballista/core/src/error.rs` matched only a _direct_ +`DataFusionError::IoError`. That was #2028, and it was fixed by classifying on +`find_root()` instead (#2119). + +The sort-shuffle writer refactor (#2038, #2106) then re-broke both cases at an +earlier point in the pipeline: the shuffle write coordinator's error arm +(`ballista/core/src/execution_plans/shuffle_writer.rs`, in the coordinator +fan-out that distributes results to output-partition streams) converts any +task error with `DataFusionError::Execution(format!("{e:?}"))`. The injected +fault now reaches the classifier as +`Execution("IoError(Custom { .. })")` (`aqe_off`) or +`Execution("Shared(IoError(Custom { .. }))")` (`aqe_on`) — the variant exists +only as printed text, so `find_root()` has nothing to unwrap and the task is +marked non-retryable. This is the same type-erasing conversion Finding 1 +describes for `FetchFailed`, which is why both scenarios are ignored against +#2027 rather than the fixed #2028. + +**Net effect:** any genuine transient IO error inside a stage that writes +shuffle output — which is every non-final stage — is classified non-retryable, +turning what should be a retried task into an immediate job failure. + +### Finding 3 — Killing every executor hung the job instead of failing it (fixed) + +Tracked by [#2029](https://github.com/apache/datafusion-ballista/issues/2029), fixed in that issue's PR. + +**Regression test:** Scenario G (`killing_every_executor_terminates_the_job`), +both AQE settings, now enabled (no longer `#[ignore]`d). + +With every executor dead mid-query, there is nothing left to schedule tasks +onto. Previously the job never terminated — the scheduler waited forever rather +than failing the query. The fix makes the scheduler wait a bounded grace period +(`no_executors_grace_period_seconds`) after losing its last executor and then +fail the job with a clear error. The scenario turns that grace down via the +cluster builder and asserts the query fails with an error naming the executor +loss. + +### For comparison: the heartbeat-expiry path does recover + +Ballista's HA recovery is not uniformly broken. Whenever the kill in Scenario +D or E happens to be noticed by heartbeat expiry (`ExecutorLost`) before any +downstream task fetches from the dead executor, the job recovers and matches +the baseline — that is why both scenarios pass on some runs. What Finding 1 +breaks is specifically the fetch-failure path, and both scenarios are ignored +against #2027 because whether a given run exercises that path is a timing +race, not something the harness currently controls. diff --git a/chaos-testing/src/bin/chaos-executor.rs b/chaos-testing/src/bin/chaos-executor.rs new file mode 100644 index 0000000000..a2efe82a72 --- /dev/null +++ b/chaos-testing/src/bin/chaos-executor.rs @@ -0,0 +1,62 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A Ballista executor whose function registry includes the chaos UDFs. +//! +//! Configured entirely from the environment, because `TestCluster` spawns it as +//! a child process. Mirrors `examples/examples/custom-executor.rs`. + +use ballista_executor::executor_process::{ + ExecutorProcessConfig, start_executor_process, +}; +use chaos_testing::registry::chaos_function_registry; +use std::sync::Arc; + +fn env_u16(key: &str) -> u16 { + std::env::var(key) + .unwrap_or_else(|_| panic!("{key} must be set")) + .parse() + .unwrap_or_else(|e| panic!("{key} must be a u16: {e}")) +} + +#[tokio::main] +async fn main() -> ballista_core::error::Result<()> { + env_logger::init(); + + let config = ExecutorProcessConfig { + bind_host: "127.0.0.1".to_string(), + port: env_u16("CHAOS_EXECUTOR_PORT"), + grpc_port: env_u16("CHAOS_EXECUTOR_GRPC_PORT"), + scheduler_host: "127.0.0.1".to_string(), + scheduler_port: env_u16("CHAOS_SCHEDULER_PORT"), + vcores: std::env::var("CHAOS_CONCURRENT_TASKS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(4), + work_dir: std::env::var("CHAOS_WORK_DIR").ok(), + // The default is 60s. Executor-loss scenarios need the scheduler to see a + // missing heartbeat within seconds, not minutes. + executor_heartbeat_interval_seconds: std::env::var("CHAOS_HEARTBEAT_SECONDS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(1), + override_function_registry: Some(chaos_function_registry()), + ..Default::default() + }; + + start_executor_process(Arc::new(config)).await +} diff --git a/chaos-testing/src/bin/chaos-scheduler.rs b/chaos-testing/src/bin/chaos-scheduler.rs new file mode 100644 index 0000000000..460582060a --- /dev/null +++ b/chaos-testing/src/bin/chaos-scheduler.rs @@ -0,0 +1,74 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A Ballista scheduler whose session builder includes the chaos UDFs, with +//! executor-loss detection tuned down from minutes to seconds. +//! +//! Mirrors `examples/examples/custom-scheduler.rs`. + +use ballista_core::error::BallistaError; +use ballista_scheduler::cluster::BallistaCluster; +use ballista_scheduler::config::SchedulerConfig; +use ballista_scheduler::scheduler_process::start_server; +use chaos_testing::registry::chaos_session_state; +use std::net::AddrParseError; +use std::sync::Arc; + +fn env_parsed(key: &str, default: T) -> T { + std::env::var(key) + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(default) +} + +#[tokio::main] +async fn main() -> ballista_core::error::Result<()> { + env_logger::init(); + + let config = SchedulerConfig { + bind_host: "127.0.0.1".to_string(), + bind_port: std::env::var("CHAOS_SCHEDULER_PORT") + .expect("CHAOS_SCHEDULER_PORT must be set") + .parse() + .expect("CHAOS_SCHEDULER_PORT must be a u16"), + // Defaults are 180s / 15s, which would make every executor-kill scenario + // take three minutes. Tests override these per scenario. + executor_timeout_seconds: env_parsed("CHAOS_EXECUTOR_TIMEOUT_SECONDS", 5), + expire_dead_executor_interval_seconds: env_parsed( + "CHAOS_EXPIRE_INTERVAL_SECONDS", + 1, + ), + task_max_failures: env_parsed("CHAOS_TASK_MAX_FAILURES", 4), + stage_max_failures: env_parsed("CHAOS_STAGE_MAX_FAILURES", 4), + // The default is 30s. Turn it down so the total-executor-loss scenario + // fails the job a second or so after the reap instead of waiting it out. + no_executors_grace_period_seconds: env_parsed( + "CHAOS_NO_EXECUTORS_GRACE_SECONDS", + 1, + ), + override_session_builder: Some(Arc::new(chaos_session_state)), + ..Default::default() + }; + + let addr = format!("{}:{}", config.bind_host, config.bind_port); + let addr = addr + .parse() + .map_err(|e: AddrParseError| BallistaError::Configuration(e.to_string()))?; + + let cluster = BallistaCluster::new_from_config(&config).await?; + start_server(cluster, addr, Arc::new(config)).await +} diff --git a/chaos-testing/src/budget.rs b/chaos-testing/src/budget.rs new file mode 100644 index 0000000000..4b5e3ace21 --- /dev/null +++ b/chaos-testing/src/budget.rs @@ -0,0 +1,126 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::path::{Path, PathBuf}; + +/// A filesystem-backed budget of injectable faults, shared across processes. +/// +/// A budget of `n` is a directory containing `n` token files. Consuming a token +/// is `fs::remove_file`, which is atomic across processes: exactly one caller +/// can succeed for a given token, whichever executor it runs in. The directory +/// outlives task retries and executor restarts, so the budget bounds the total +/// number of injected faults for the whole run rather than per attempt. +#[derive(Debug, Clone)] +pub struct FaultBudget { + dir: PathBuf, +} + +impl FaultBudget { + /// Create the budget directory with `tokens` tokens, replacing any existing one. + pub fn create(dir: &Path, tokens: usize) -> std::io::Result { + let _ = std::fs::remove_dir_all(dir); + std::fs::create_dir_all(dir)?; + for i in 0..tokens { + std::fs::write(dir.join(format!("token-{i}")), b"")?; + } + Ok(Self { + dir: dir.to_path_buf(), + }) + } + + /// Open an existing budget directory. Used by executor processes, which only + /// ever consume; a missing directory simply means no faults are available. + pub fn open(dir: &Path) -> Self { + Self { + dir: dir.to_path_buf(), + } + } + + pub fn dir(&self) -> &Path { + &self.dir + } + + pub fn remaining(&self) -> usize { + std::fs::read_dir(&self.dir) + .map(|entries| entries.flatten().count()) + .unwrap_or(0) + } + + /// Attempt to consume one token. Returns true iff this caller won the token. + /// + /// `remove_file` is the atomicity primitive: if two executors race for the + /// last token, exactly one `remove_file` returns Ok and the other errors. + pub fn try_consume(&self) -> bool { + let Ok(entries) = std::fs::read_dir(&self.dir) else { + return false; + }; + for entry in entries.flatten() { + if std::fs::remove_file(entry.path()).is_ok() { + return true; + } + } + false + } +} + +#[cfg(test)] +mod tests { + use crate::budget::FaultBudget; + use std::path::PathBuf; + + fn temp_dir(name: &str) -> PathBuf { + let dir = std::env::temp_dir().join(format!("ballista-chaos-{name}")); + let _ = std::fs::remove_dir_all(&dir); + dir + } + + #[test] + fn consumes_exactly_the_budget() { + let dir = temp_dir("budget-exact"); + let budget = FaultBudget::create(&dir, 2).unwrap(); + + assert_eq!(budget.remaining(), 2); + assert!(budget.try_consume()); + assert!(budget.try_consume()); + assert!( + !budget.try_consume(), + "third consume must fail: budget was 2" + ); + assert_eq!(budget.remaining(), 0); + } + + #[test] + fn zero_budget_never_consumes() { + let dir = temp_dir("budget-zero"); + let budget = FaultBudget::create(&dir, 0).unwrap(); + assert!(!budget.try_consume()); + } + + #[test] + fn open_sees_tokens_created_by_another_handle() { + // This models a separate executor process reading the same budget dir. + let dir = temp_dir("budget-shared"); + let creator = FaultBudget::create(&dir, 1).unwrap(); + let other = FaultBudget::open(&dir); + + assert!(other.try_consume(), "second handle must see the token"); + assert!( + !creator.try_consume(), + "token already consumed by the other handle" + ); + } +} diff --git a/chaos-testing/src/cluster.rs b/chaos-testing/src/cluster.rs new file mode 100644 index 0000000000..a4c1143823 --- /dev/null +++ b/chaos-testing/src/cluster.rs @@ -0,0 +1,726 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use nix::fcntl::{Flock, FlockArg}; +use std::fs::{File, OpenOptions}; +use std::net::TcpListener; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command, Stdio}; +use std::sync::{Arc, OnceLock}; +use std::time::{Duration, Instant}; +use tokio::sync::{Mutex, OwnedMutexGuard}; + +/// Reserve a free TCP port by binding to :0 and immediately releasing it. +/// +/// Inherently racy, but adequate here: the child binds within milliseconds and +/// the tests are the only thing running. +fn free_port() -> u16 { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind ephemeral port"); + listener.local_addr().expect("local addr").port() +} + +/// Open a child process log file for append. +/// +/// Appending (rather than truncating) matters for Task 6's kill/restart +/// scenarios: a restarted executor reuses the same log path, and the prior +/// process's output is the evidence of why it died. It must not be wiped out +/// by the replacement process starting up. +fn open_log(path: &Path) -> std::io::Result { + OpenOptions::new().create(true).append(true).open(path) +} + +/// Locate a binary built by this crate. +/// +/// The profile directory is taken from the running test executable +/// (`//deps/`) rather than inferred from the profile's +/// settings. Inferring it from `cfg!(debug_assertions)` only works for the +/// stock `dev` and `release` profiles: CI builds under `--profile ci`, which +/// inherits `dev` but turns debug assertions off, so the binaries land in +/// `target/ci/` while the inference points at `target/release/`. Deriving the +/// directory from `current_exe` is correct for any profile and honours +/// `CARGO_TARGET_DIR` for free. +fn binary(name: &str) -> PathBuf { + let mut path = std::env::current_exe().expect("locate the test executable"); + path.pop(); // deps/ + path.pop(); // // + path.push(name); + assert!( + path.exists(), + "{} not found at {}. Run `cargo build -p ballista-chaos --bins` first.", + name, + path.display() + ); + path +} + +/// One supervised executor process. +/// +/// `child` is used by this task's `kill_executor`/`executor_is_alive`. `port`, +/// `grpc_port`, and `work_dir` are still not read anywhere yet — no code in +/// this task needed to target an executor by its network address or inspect +/// its working directory — so they keep a narrower `#[allow(dead_code)]` than +/// the struct-wide one Task 5 left; a later scenario that needs to address a +/// specific executor's port or inspect its shuffle files can drop it then. +pub(crate) struct ExecutorHandle { + pub(crate) child: Child, + #[allow(dead_code)] + pub(crate) port: u16, + #[allow(dead_code)] + pub(crate) grpc_port: u16, + #[allow(dead_code)] + pub(crate) work_dir: PathBuf, +} + +pub struct TestClusterBuilder { + executors: usize, + executor_timeout_seconds: u64, + expire_interval_seconds: u64, + task_max_failures: usize, + stage_max_failures: usize, + concurrent_tasks: usize, + no_executors_grace_seconds: u64, +} + +impl Default for TestClusterBuilder { + fn default() -> Self { + Self { + executors: 2, + // Ballista's defaults are 180s/15s, which would make an executor-kill + // scenario take three minutes to even notice the death. + executor_timeout_seconds: 5, + expire_interval_seconds: 1, + task_max_failures: 4, + stage_max_failures: 4, + concurrent_tasks: 4, + // Ballista's default is 30s. A short grace makes the total-loss + // scenario fail the job a second or so after the reap. + no_executors_grace_seconds: 1, + } + } +} + +impl TestClusterBuilder { + pub fn executors(mut self, n: usize) -> Self { + self.executors = n; + self + } + + /// How long the scheduler waits on a missing heartbeat before declaring the + /// executor lost. Scenario E raises this deliberately to isolate the + /// FetchPartitionError path from the ExecutorLost path. + pub fn executor_timeout_seconds(mut self, seconds: u64) -> Self { + self.executor_timeout_seconds = seconds; + self + } + + /// How long the scheduler waits, after losing its last executor, for one to + /// (re)register before failing the running jobs (see #2029). + pub fn no_executors_grace_seconds(mut self, seconds: u64) -> Self { + self.no_executors_grace_seconds = seconds; + self + } + + pub fn task_max_failures(mut self, n: usize) -> Self { + self.task_max_failures = n; + self + } + + pub async fn start(self) -> Result { + // Held for the cluster's whole lifetime, so only one cluster exists in + // this process at a time. `--test-threads=1` gives the same guarantee, + // but nothing forces a caller (CI runs a plain `cargo test` over the + // whole workspace) to pass it, and a dozen concurrent clusters exhaust + // ports and CPU and fail for reasons that have nothing to do with the + // scenario under test. Making the harness enforce its own requirement + // is more robust than documenting a flag. + let cluster_lock = cluster_lock().lock_owned().await; + + // The mutex above only serializes within one process. `cargo test` + // runs test binaries sequentially, but nothing else does: a second + // cargo invocation in another shell, an IDE's background test run, or + // a runner like nextest (which parallelizes binaries) would put two + // multi-process clusters on the machine at once, and on a two-core CI + // runner that is enough for executors to miss the registration + // deadline. flock is advisory and machine-wide, so the second process + // blocks here until the first one's cluster is gone. + let machine_lock = tokio::task::spawn_blocking(machine_lock) + .await + .map_err(|e| format!("acquire machine-wide cluster lock: {e}"))??; + + let temp = tempfile::tempdir().map_err(|e| e.to_string())?; + let log_dir = temp.path().join("logs"); + std::fs::create_dir_all(&log_dir).map_err(|e| e.to_string())?; + let scheduler_port = free_port(); + + let scheduler_log = log_dir.join("scheduler.log"); + let scheduler_stdout = open_log(&scheduler_log).map_err(|e| { + format!("open scheduler log {}: {e}", scheduler_log.display()) + })?; + let scheduler_stderr = open_log(&scheduler_log).map_err(|e| { + format!("open scheduler log {}: {e}", scheduler_log.display()) + })?; + + let mut scheduler = Command::new(binary("chaos-scheduler")); + scheduler + .env("CHAOS_SCHEDULER_PORT", scheduler_port.to_string()) + .env( + "CHAOS_EXECUTOR_TIMEOUT_SECONDS", + self.executor_timeout_seconds.to_string(), + ) + .env( + "CHAOS_EXPIRE_INTERVAL_SECONDS", + self.expire_interval_seconds.to_string(), + ) + .env( + "CHAOS_TASK_MAX_FAILURES", + self.task_max_failures.to_string(), + ) + .env( + "CHAOS_STAGE_MAX_FAILURES", + self.stage_max_failures.to_string(), + ) + .env( + "CHAOS_NO_EXECUTORS_GRACE_SECONDS", + self.no_executors_grace_seconds.to_string(), + ) + .env( + "RUST_LOG", + std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into()), + ) + .stdout(Stdio::from(scheduler_stdout)) + .stderr(Stdio::from(scheduler_stderr)); + let scheduler = scheduler + .spawn() + .map_err(|e| format!("spawn scheduler: {e}"))?; + + let mut cluster = TestCluster { + scheduler, + scheduler_port, + executors: Vec::new(), + temp, + log_dir, + builder: self, + _cluster_lock: cluster_lock, + _machine_lock: machine_lock, + }; + + for i in 0..cluster.builder.executors { + cluster.spawn_executor(i)?; + } + + let n = cluster.builder.executors; + cluster.await_executors(n).await?; + Ok(cluster) + } +} + +/// A multi-process Ballista cluster under test. Every child is killed on drop. +pub struct TestCluster { + scheduler: Child, + scheduler_port: u16, + pub(crate) executors: Vec, + temp: tempfile::TempDir, + log_dir: PathBuf, + builder: TestClusterBuilder, + /// Serializes clusters across the whole test process; see + /// [`TestClusterBuilder::start`]. `Drop for TestCluster` reaps every child + /// before this field is dropped, so the next cluster never starts until the + /// previous one's processes are gone. + _cluster_lock: OwnedMutexGuard<()>, + /// Serializes clusters across test *processes* on the same machine; see + /// [`TestClusterBuilder::start`]. + _machine_lock: Flock, +} + +/// The process-wide lock guaranteeing one cluster at a time. +fn cluster_lock() -> Arc> { + static LOCK: OnceLock>> = OnceLock::new(); + LOCK.get_or_init(|| Arc::new(Mutex::new(()))).clone() +} + +/// The machine-wide lock guaranteeing one cluster at a time across processes. +/// +/// Blocks until acquired, so call it from a blocking-friendly context. The +/// kernel releases a flock when its file handle closes, so a SIGKILLed test +/// process cannot leave the lock stuck. +fn machine_lock() -> Result, String> { + let path = std::env::temp_dir().join("ballista-chaos-cluster.lock"); + let file = OpenOptions::new() + .create(true) + .write(true) + .truncate(false) + .open(&path) + .map_err(|e| format!("open lock file {}: {e}", path.display()))?; + Flock::lock(file, FlockArg::LockExclusive) + .map_err(|(_, e)| format!("flock {}: {e}", path.display())) +} + +impl TestCluster { + pub fn builder() -> TestClusterBuilder { + TestClusterBuilder::default() + } + + /// The Ballista client URL. + pub fn scheduler_url(&self) -> String { + format!("df://127.0.0.1:{}", self.scheduler_port) + } + + /// The scheduler REST base URL. gRPC and REST share one port. + pub fn rest_url(&self) -> String { + format!("http://127.0.0.1:{}", self.scheduler_port) + } + + /// The shared directory for fixtures and fault budgets. Every executor can + /// read it, which is what makes the fault budget cluster-wide. + pub fn shared_dir(&self) -> &std::path::Path { + self.temp.path() + } + + /// Directory containing each child process's stdout/stderr log + /// (`scheduler.log`, `executor-{index}.log`). When a scenario fails, these + /// logs are the evidence of what the scheduler and executors were doing. + pub fn log_dir(&self) -> &Path { + &self.log_dir + } + + pub(crate) fn spawn_executor(&mut self, index: usize) -> Result<(), String> { + let port = free_port(); + let grpc_port = free_port(); + let work_dir = self.temp.path().join(format!("executor-{index}")); + std::fs::create_dir_all(&work_dir).map_err(|e| e.to_string())?; + + // Appends rather than truncates: a respawn at this same index (Task 6's + // kill/restart scenarios) must not erase the log of the process that + // just died. + let executor_log = self.log_dir.join(format!("executor-{index}.log")); + let executor_stdout = open_log(&executor_log).map_err(|e| { + format!("open executor {index} log {}: {e}", executor_log.display()) + })?; + let executor_stderr = open_log(&executor_log).map_err(|e| { + format!("open executor {index} log {}: {e}", executor_log.display()) + })?; + + let child = Command::new(binary("chaos-executor")) + .env("CHAOS_EXECUTOR_PORT", port.to_string()) + .env("CHAOS_EXECUTOR_GRPC_PORT", grpc_port.to_string()) + .env("CHAOS_SCHEDULER_PORT", self.scheduler_port.to_string()) + .env( + "CHAOS_CONCURRENT_TASKS", + self.builder.concurrent_tasks.to_string(), + ) + .env("CHAOS_WORK_DIR", work_dir.display().to_string()) + .env("CHAOS_HEARTBEAT_SECONDS", "1") + .env( + "RUST_LOG", + std::env::var("RUST_LOG").unwrap_or_else(|_| "info".into()), + ) + .stdout(Stdio::from(executor_stdout)) + .stderr(Stdio::from(executor_stderr)) + .spawn() + .map_err(|e| format!("spawn executor {index}: {e}"))?; + + if self.executors.len() > index { + self.executors[index] = ExecutorHandle { + child, + port, + grpc_port, + work_dir, + }; + } else { + self.executors.push(ExecutorHandle { + child, + port, + grpc_port, + work_dir, + }); + } + Ok(()) + } + + /// Block until `n` executors have registered with the scheduler. + /// + /// The deadline is generous because a loaded CI runner starts these child + /// processes slowly; a healthy cluster returns in a couple of seconds + /// regardless. On timeout the error carries every child's log tail — the + /// only evidence of a startup failure CI would otherwise throw away. + pub async fn await_executors(&self, n: usize) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(120); + loop { + if let Ok(count) = self.registered_executors().await + && count >= n + { + return Ok(()); + } + if Instant::now() > deadline { + return Err(format!( + "timed out waiting for {n} executors to register\n{}", + self.log_tails() + )); + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + } + + /// The last lines of every child process log, for timeout diagnostics. + fn log_tails(&self) -> String { + let mut out = String::new(); + let mut paths: Vec = std::fs::read_dir(&self.log_dir) + .map(|d| d.filter_map(|e| e.ok().map(|e| e.path())).collect()) + .unwrap_or_default(); + paths.sort(); + for path in paths { + let content = std::fs::read_to_string(&path).unwrap_or_default(); + let tail: Vec<&str> = content.lines().rev().take(20).collect(); + out.push_str(&format!( + "--- {} (last {} lines) ---\n", + path.display(), + tail.len() + )); + for line in tail.into_iter().rev() { + out.push_str(line); + out.push('\n'); + } + } + out + } + + /// Block until the scheduler considers exactly `n` executors registered. + /// + /// Unlike `await_executors` (which waits for *at least* `n`, the right + /// condition when growing a cluster), a SIGKILLed executor is not dropped + /// from `/api/executors` the instant it dies — the scheduler keeps + /// listing it until its heartbeat times out (`executor_timeout_seconds`). + /// A scenario that kills an executor and wants to observe the scheduler + /// actually reaping it (rather than just transiently over-counting) needs + /// to wait for the count to come down to `n` exactly, not merely reach it. + pub async fn await_executor_count(&self, n: usize) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(120); + loop { + if let Ok(count) = self.registered_executors().await + && count == n + { + return Ok(()); + } + if Instant::now() > deadline { + return Err(format!( + "timed out waiting for exactly {n} registered executors" + )); + } + tokio::time::sleep(Duration::from_millis(200)).await; + } + } + + /// How many executors the scheduler currently considers registered. + pub async fn registered_executors(&self) -> Result { + let body: serde_json::Value = + reqwest::get(format!("{}/api/executors", self.rest_url())) + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string())?; + Ok(body.as_array().map(|a| a.len()).unwrap_or(0)) + } + + /// The id of the single job the scheduler currently knows about. + /// + /// The harness runs one query at a time, so "the running job" is unambiguous. + pub async fn running_job_id(&self) -> Result { + let deadline = Instant::now() + Duration::from_secs(30); + loop { + let body: serde_json::Value = + reqwest::get(format!("{}/api/jobs", self.rest_url())) + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string())?; + + if let Some(job) = body.as_array().and_then(|jobs| jobs.first()) + && let Some(id) = job.get("job_id").and_then(|v| v.as_str()) + { + return Ok(id.to_string()); + } + if Instant::now() > deadline { + return Err("timed out waiting for a job to appear".to_string()); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + + async fn stages(&self, job_id: &str) -> Result { + reqwest::get(format!("{}/api/job/{job_id}/stages", self.rest_url())) + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string()) + } + + /// Block until `stage_id` has at least one task in state Running. + /// + /// This is what lets a kill land *while the stage is genuinely executing*, + /// rather than after an arbitrary sleep that may fire too early or too late. + pub async fn await_stage_running( + &self, + job_id: &str, + stage_id: usize, + ) -> Result<(), String> { + self.await_stage_task_state(job_id, stage_id, "Running") + .await + } + + /// Block until any task in any stage is Running. + /// + /// Planner-agnostic sync point: the static and adaptive (AQE) planners + /// number and materialize stages differently, so rather than target a + /// specific stage id we wait until the job is genuinely executing a task + /// somewhere. Used where the scenario only needs a kill to land mid-flight. + pub async fn await_any_stage_running(&self, job_id: &str) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(60); + loop { + let stages = self.stages(job_id).await?; + let running = + stages + .get("stages") + .and_then(|s| s.as_array()) + .is_some_and(|stages| { + stages.iter().any(|stage| { + stage.get("tasks").and_then(|t| t.as_array()).is_some_and( + |tasks| { + tasks.iter().any(|t| { + t.get("status").and_then(|s| s.as_str()) + == Some("Running") + }) + }, + ) + }) + }); + if running { + return Ok(()); + } + if Instant::now() > deadline { + return Err( + "timed out waiting for any stage to start running".to_string() + ); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + + /// Block until every task in `stage_id` is Successful. + pub async fn await_stage_successful( + &self, + job_id: &str, + stage_id: usize, + ) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(60); + loop { + let stages = self.stages(job_id).await?; + if let Some(stage) = find_stage(&stages, stage_id) + && let Some(tasks) = stage.get("tasks").and_then(|t| t.as_array()) + { + let all_ok = !tasks.is_empty() + && tasks.iter().all(|t| { + t.get("status").and_then(|s| s.as_str()) == Some("Successful") + }); + if all_ok { + return Ok(()); + } + } + if Instant::now() > deadline { + return Err(format!("timed out waiting for stage {stage_id} to succeed")); + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + } + + async fn await_stage_task_state( + &self, + job_id: &str, + stage_id: usize, + state: &str, + ) -> Result<(), String> { + let deadline = Instant::now() + Duration::from_secs(60); + loop { + let stages = self.stages(job_id).await?; + if let Some(stage) = find_stage(&stages, stage_id) + && let Some(tasks) = stage.get("tasks").and_then(|t| t.as_array()) + { + let hit = tasks + .iter() + .any(|t| t.get("status").and_then(|s| s.as_str()) == Some(state)); + if hit { + return Ok(()); + } + } + if Instant::now() > deadline { + return Err(format!( + "timed out waiting for a {state} task in stage {stage_id}" + )); + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + } + + /// The scheduler's short categorical job status: "Queued", "Running", + /// "Completed", "Failed", or "Invalid". + /// + /// The brief's original version read the JSON field literally named + /// `job_status`, but that field (`JobResponse::job_status` / + /// `handlers::format_job_status`'s second return value) is actually a long + /// human-readable sentence, e.g. "Completed. Produced 1 partition + /// containing 50 rows. Elapsed time: 49 ms." — never `"Successful"` as the + /// original doc comment claimed (a completed job reports `"Completed"`, + /// not `"Successful"`), and not stable/matchable for assertions. The short + /// categorical value the doc comment actually promises lives in the + /// sibling `status` field, so this reads that one instead. Verified against + /// a live cluster: `{"status": "Completed", "job_status": "Completed. + /// Produced 1 partition containing 50 rows. Elapsed time: 49 ms.", ...}`. + pub async fn job_status(&self, job_id: &str) -> Result { + let body: serde_json::Value = + reqwest::get(format!("{}/api/job/{job_id}", self.rest_url())) + .await + .map_err(|e| e.to_string())? + .json() + .await + .map_err(|e| e.to_string())?; + Ok(body + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or("unknown") + .to_string()) + } + + /// SIGKILL an executor. Not SIGTERM: a graceful shutdown would let the + /// executor deregister, which is a different (and much easier) code path + /// than the crash we are trying to test. + pub fn kill_executor(&mut self, index: usize) -> Result<(), String> { + use nix::sys::signal::{Signal, kill}; + use nix::unistd::Pid; + + let pid = self.executors[index].child.id(); + kill(Pid::from_raw(pid as i32), Signal::SIGKILL).map_err(|e| e.to_string())?; + let _ = self.executors[index].child.wait(); + Ok(()) + } + + /// Start a fresh executor process in the given slot and wait for it to register. + pub async fn restart_executor(&mut self, index: usize) -> Result<(), String> { + let expected = self.executors.len(); + self.spawn_executor(index)?; + self.await_executors(expected).await + } + + /// Whether an executor process is still alive. + pub fn executor_is_alive(&mut self, index: usize) -> bool { + matches!(self.executors[index].child.try_wait(), Ok(None)) + } +} + +/// Stage ids come back from the REST API as strings. +fn find_stage(stages: &serde_json::Value, stage_id: usize) -> Option<&serde_json::Value> { + stages.get("stages")?.as_array()?.iter().find(|s| { + s.get("stage_id").and_then(|v| v.as_str()) == Some(stage_id.to_string().as_str()) + }) +} + +/// If `child` already exited with a non-zero status, log the path of its +/// output so a human investigating a failed scenario knows where to look. +/// Then make sure it is actually gone. +fn reap(child: &mut Child, log_path: &Path) { + if let Ok(Some(status)) = child.try_wait() + && !status.success() + { + log::warn!( + "process exited with {status}; see log at {}", + log_path.display() + ); + } + let _ = child.kill(); + let _ = child.wait(); +} + +impl Drop for TestCluster { + fn drop(&mut self) { + for (index, executor) in self.executors.iter_mut().enumerate() { + let log_path = self.log_dir.join(format!("executor-{index}.log")); + reap(&mut executor.child, &log_path); + } + let scheduler_log = self.log_dir.join("scheduler.log"); + reap(&mut self.scheduler, &scheduler_log); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn cluster_starts_with_the_requested_executors_registered() { + let cluster = TestCluster::builder() + .executors(2) + .start() + .await + .expect("cluster must start"); + + // The scheduler's own view is the source of truth: if the executors did + // not register, every later scenario would silently run single-executor. + let executors: serde_json::Value = + reqwest::get(format!("{}/api/executors", cluster.rest_url())) + .await + .unwrap() + .json() + .await + .unwrap(); + + assert_eq!( + executors.as_array().map(|a| a.len()), + Some(2), + "expected 2 registered executors, got {executors:?}" + ); + } + + #[tokio::test] + async fn killed_executor_is_reaped_and_can_be_restarted() { + let mut cluster = TestCluster::builder() + .executors(2) + .executor_timeout_seconds(5) + .start() + .await + .unwrap(); + + assert_eq!(cluster.registered_executors().await.unwrap(), 2); + + cluster.kill_executor(0).unwrap(); + + // The scheduler must notice the missing heartbeat and drop the executor. + // With the defaults (180s timeout, 60s heartbeat) this would never happen + // inside a test; it works only because the harness turns both down. + cluster + .await_executor_count(1) + .await + .expect("scheduler never reaped the killed executor"); + + cluster.restart_executor(0).await.unwrap(); + assert_eq!( + cluster.registered_executors().await.unwrap(), + 2, + "restarted executor must re-register" + ); + } +} diff --git a/chaos-testing/src/fixture.rs b/chaos-testing/src/fixture.rs new file mode 100644 index 0000000000..3871fce4ef --- /dev/null +++ b/chaos-testing/src/fixture.rs @@ -0,0 +1,373 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::arrow::array::{Int64Array, StringArray}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::error::Result; +use datafusion::prelude::SessionContext; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +/// Rows per `facts` partition file. +const FACT_ROWS: i64 = 2_000; +/// Number of `facts` partition files. Multiple files give the scan enough +/// parallelism to spread across executors. +const FACT_FILES: i64 = 8; +/// Number of distinct keys in `dims`. The join on `key` forces a shuffle, which +/// is what creates the multi-stage plan the HA paths need. +const DIM_KEYS: i64 = 50; + +/// A small deterministic Parquet dataset: `facts(key, value)` joined to +/// `dims(key, name)`. +pub struct Fixture { + facts_dir: PathBuf, + dims_dir: PathBuf, +} + +impl Fixture { + pub async fn write(dir: &Path) -> Result { + let ctx = SessionContext::new(); + let facts_dir = dir.join("facts"); + let dims_dir = dir.join("dims"); + std::fs::create_dir_all(&facts_dir)?; + std::fs::create_dir_all(&dims_dir)?; + + let fact_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("value", DataType::Int64, false), + ])); + + for file in 0..FACT_FILES { + let keys: Vec = (0..FACT_ROWS) + .map(|r| (r + file * FACT_ROWS) % DIM_KEYS) + .collect(); + let values: Vec = (0..FACT_ROWS).map(|r| r + file * FACT_ROWS).collect(); + let batch = RecordBatch::try_new( + fact_schema.clone(), + vec![ + Arc::new(Int64Array::from(keys)), + Arc::new(Int64Array::from(values)), + ], + )?; + let df = ctx.read_batch(batch)?; + df.write_parquet( + facts_dir + .join(format!("part-{file}.parquet")) + .to_str() + .unwrap(), + datafusion::dataframe::DataFrameWriteOptions::new(), + None, + ) + .await?; + } + + let dim_schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("name", DataType::Utf8, false), + ])); + let batch = RecordBatch::try_new( + dim_schema, + vec![ + Arc::new(Int64Array::from((0..DIM_KEYS).collect::>())), + Arc::new(StringArray::from( + (0..DIM_KEYS) + .map(|k| format!("dim-{k}")) + .collect::>(), + )), + ], + )?; + ctx.read_batch(batch)? + .write_parquet( + dims_dir.join("part-0.parquet").to_str().unwrap(), + datafusion::dataframe::DataFrameWriteOptions::new(), + None, + ) + .await?; + + Ok(Self { + facts_dir, + dims_dir, + }) + } + + /// `CREATE EXTERNAL TABLE` statements, to run against any SessionContext + /// (local for the baseline, Ballista for the cluster run). + pub fn register_sql(&self) -> Vec { + vec![ + format!( + "CREATE EXTERNAL TABLE facts STORED AS PARQUET LOCATION '{}'", + self.facts_dir.display() + ), + format!( + "CREATE EXTERNAL TABLE dims STORED AS PARQUET LOCATION '{}'", + self.dims_dir.display() + ), + ] + } + + /// The chaos-free query. A join plus a grouped aggregate: at least two + /// stages, with a shuffle between them. + pub fn baseline_query() -> &'static str { + "SELECT d.name, COUNT(*) AS n, SUM(f.value) AS total \ + FROM facts f JOIN dims d ON f.key = d.key \ + GROUP BY d.name ORDER BY d.name" + } + + /// The same query with a chaos UDF spliced into the WHERE clause. + /// + /// `injection` is a complete `chaos_*(...)` call returning BOOLEAN. Its own + /// guard argument selects which rows fault (and hence which partitions and + /// tasks). The predicate is written so every row passes through regardless: + /// the guard decides only *where the fault fires*, never which rows survive, + /// so a chaos run must return exactly the baseline result. + /// + /// The predicate uses `IS NOT NULL` rather than `OR TRUE`: DataFusion's + /// optimizer constant-folds `expr OR TRUE` to the literal `TRUE` and drops + /// the volatile UDF call entirely (verified empirically in the + /// `or_true_predicate_is_optimized_away_and_never_fires` and + /// `chaos_query_predicate_survives_optimization_and_fires` tests below), + /// which would silently disarm every fault injection. `chaos_fail`/ + /// `chaos_delay` always return `Some(guard)` (see udf.rs), so + /// `... IS NOT NULL` is always true without being foldable to a constant, + /// and the optimizer must still evaluate the call to determine nullness. + /// + /// Example: `chaos_query("chaos_fail(f.key = 7, 'io', '/tmp/b')")` + pub fn chaos_query(injection: &str) -> String { + format!( + "SELECT d.name, COUNT(*) AS n, SUM(f.value) AS total \ + FROM facts f JOIN dims d ON f.key = d.key \ + WHERE {injection} IS NOT NULL \ + GROUP BY d.name ORDER BY d.name" + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::SessionContext; + + #[tokio::test] + async fn baseline_query_is_deterministic_and_non_empty() { + let dir = tempfile::tempdir().unwrap(); + let fixture = Fixture::write(dir.path()).await.unwrap(); + + let ctx = SessionContext::new(); + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + let rows = ctx + .sql(Fixture::baseline_query()) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let total: usize = rows.iter().map(|b| b.num_rows()).sum(); + assert!(total > 0, "baseline query must return rows"); + + // Determinism: the same query over the same data must give the same answer. + let rows2 = ctx + .sql(Fixture::baseline_query()) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!( + datafusion::arrow::util::pretty::pretty_format_batches(&rows) + .unwrap() + .to_string(), + datafusion::arrow::util::pretty::pretty_format_batches(&rows2) + .unwrap() + .to_string(), + ); + } + + /// Empirically confirms the risk called out on `chaos_query`: DataFusion's + /// constant-folding simplifies `expr OR TRUE` to the literal `TRUE` during + /// logical optimization, and once the predicate is a literal the join/filter + /// no longer references the UDF call at all, so it is never invoked. If this + /// test ever starts failing (budget still consumed), the optimizer's folding + /// behavior changed and `chaos_query`'s `IS NOT NULL` predicate should be + /// re-verified instead of reverting to `OR TRUE`. + #[tokio::test] + async fn or_true_predicate_is_optimized_away_and_never_fires() { + let dir = tempfile::tempdir().unwrap(); + let fixture = Fixture::write(dir.path()).await.unwrap(); + let budget_dir = dir.path().join("budget"); + let budget = crate::budget::FaultBudget::create(&budget_dir, 1).unwrap(); + + let ctx = SessionContext::new(); + ctx.register_udf(crate::udf::chaos_fail_udf().as_ref().clone()); + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + // key = 7 genuinely exists in the generated data (keys cycle 0..DIM_KEYS), + // so an empty-partition short-circuit is not what's suppressing the call. + let sql = format!( + "SELECT d.name, COUNT(*) AS n, SUM(f.value) AS total \ + FROM facts f JOIN dims d ON f.key = d.key \ + WHERE chaos_fail(f.key = 7, 'io', '{}') OR TRUE \ + GROUP BY d.name ORDER BY d.name", + budget_dir.display() + ); + + let explain = ctx + .sql(&format!("EXPLAIN {sql}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + let explain_str = + datafusion::arrow::util::pretty::pretty_format_batches(&explain) + .unwrap() + .to_string(); + println!("EXPLAIN (OR TRUE form):\n{explain_str}"); + assert!( + !explain_str.contains("chaos_fail"), + "expected the optimizer to eliminate the chaos_fail call from the plan, but it is still present:\n{explain_str}" + ); + + let result = ctx.sql(&sql).await.unwrap().collect().await; + assert!( + result.is_ok(), + "expected the query to succeed because the fault never fires, got {result:?}" + ); + assert_eq!( + budget.remaining(), + 1, + "the token must be untouched: `OR TRUE` is constant-folded away, so chaos_fail is never invoked" + ); + } + + /// The predicate form `Fixture::chaos_query` actually uses. Unlike `OR TRUE`, + /// `chaos_fail(...) IS NOT NULL` is not foldable to a constant (the UDF's + /// return type/value is not known to the optimizer without invoking it), so + /// the call must survive into the physical plan and the fault must fire. + #[tokio::test] + async fn chaos_query_predicate_survives_optimization_and_fires() { + let dir = tempfile::tempdir().unwrap(); + let fixture = Fixture::write(dir.path()).await.unwrap(); + let budget_dir = dir.path().join("budget"); + let budget = crate::budget::FaultBudget::create(&budget_dir, 1).unwrap(); + + let ctx = SessionContext::new(); + ctx.register_udf(crate::udf::chaos_fail_udf().as_ref().clone()); + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + let injection = + format!("chaos_fail(f.key = 7, 'io', '{}')", budget_dir.display()); + let sql = Fixture::chaos_query(&injection); + + let explain = ctx + .sql(&format!("EXPLAIN {sql}")) + .await + .unwrap() + .collect() + .await + .unwrap(); + let explain_str = + datafusion::arrow::util::pretty::pretty_format_batches(&explain) + .unwrap() + .to_string(); + println!("EXPLAIN (IS NOT NULL form):\n{explain_str}"); + assert!( + explain_str.contains("chaos_fail"), + "expected the chaos_fail call to survive into the plan, but it is missing:\n{explain_str}" + ); + + let err = ctx.sql(&sql).await.unwrap().collect().await.unwrap_err(); + // The join fans the same filtered stream out to multiple consumers + // (CollectLeft build side plus the probe side), so DataFusion wraps the + // propagated error in `DataFusionError::Shared` rather than surfacing the + // bare `IoError` directly; `find_root` unwraps that layer. + assert!( + matches!( + err.find_root(), + datafusion::error::DataFusionError::IoError(_) + ), + "expected an IoError (possibly Shared-wrapped) from the injected fault, got {err:?}" + ); + assert_eq!( + budget.remaining(), + 0, + "the token must be consumed: the predicate must force chaos_fail to be evaluated" + ); + } + + /// CRITICAL INVARIANT: When a chaos fault cannot fire (budget exhausted), + /// `chaos_query()` must return EXACTLY the same result as `baseline_query()`. + /// + /// This invariant underpins every HA scenario: we detect when a re-run stage + /// duplicates or drops partitions by asserting that the result after a fault + /// is identical to the baseline. The chaos_query predicate `WHERE ... IS NOT + /// NULL` preserves the row set only because chaos_fail/chaos_delay always + /// return `Some(guard)` (never NULL). If a future change ever returned NULL + /// from either UDF, `IS NOT NULL` would silently drop rows and no test would + /// catch it — until a real distributed run failed mysteriously. This test + /// pins that invariant before any such refactor happens. + #[tokio::test] + async fn chaos_query_without_a_firing_fault_equals_baseline() { + let dir = tempfile::tempdir().unwrap(); + let fixture = Fixture::write(dir.path()).await.unwrap(); + let budget_dir = dir.path().join("budget"); + // 0 tokens: the fault cannot fire, no matter how many times it is called. + let _budget = crate::budget::FaultBudget::create(&budget_dir, 0).unwrap(); + + let ctx = SessionContext::new(); + ctx.register_udf(crate::udf::chaos_fail_udf().as_ref().clone()); + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + let injection = + format!("chaos_fail(f.key = 7, 'io', '{}')", budget_dir.display()); + let chaos_sql = Fixture::chaos_query(&injection); + let baseline_sql = Fixture::baseline_query(); + + let chaos_rows = ctx.sql(&chaos_sql).await.unwrap().collect().await.unwrap(); + let baseline_rows = ctx + .sql(baseline_sql) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let chaos_str = + datafusion::arrow::util::pretty::pretty_format_batches(&chaos_rows) + .unwrap() + .to_string(); + let baseline_str = + datafusion::arrow::util::pretty::pretty_format_batches(&baseline_rows) + .unwrap() + .to_string(); + + assert_eq!( + chaos_str, baseline_str, + "chaos_query with a non-firing fault must return the same rows as baseline_query" + ); + } +} diff --git a/chaos-testing/src/lib.rs b/chaos-testing/src/lib.rs new file mode 100644 index 0000000000..5c1c527643 --- /dev/null +++ b/chaos-testing/src/lib.rs @@ -0,0 +1,27 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Fault-injection harness for testing Ballista high availability. +//! +//! See `specs/2026-07-13-ha-chaos-integration-test.md` in the ballista-ai-dev +//! repository for the design. + +pub mod budget; +pub mod cluster; +pub mod fixture; +pub mod registry; +pub mod udf; diff --git a/chaos-testing/src/registry.rs b/chaos-testing/src/registry.rs new file mode 100644 index 0000000000..8f4df2946f --- /dev/null +++ b/chaos-testing/src/registry.rs @@ -0,0 +1,103 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::udf::{chaos_delay_udf, chaos_fail_udf}; +use ballista_core::registry::BallistaFunctionRegistry; +use datafusion::execution::FunctionRegistry; +use datafusion::execution::session_state::{SessionState, SessionStateBuilder}; +use datafusion::prelude::SessionConfig; +use std::sync::Arc; + +/// The executor's function registry, extended with the chaos UDFs. +/// +/// Starts from `BallistaFunctionRegistry::default()` (the DataFusion built-ins) +/// so overriding the registry does not silently drop the standard functions. +pub fn chaos_function_registry() -> Arc { + let mut registry = BallistaFunctionRegistry::default(); + for udf in [chaos_fail_udf(), chaos_delay_udf()] { + registry + .scalar_functions + .insert(udf.name().to_string(), udf); + } + Arc::new(registry) +} + +/// The scheduler's session state, extended with the chaos UDFs. +/// +/// The scheduler needs them to plan a logical plan that references them; the +/// executor needs them to decode the physical plan. Both must agree. +/// +/// Note: `SessionStateBuilder::with_scalar_functions` *replaces* the builder's +/// scalar function list rather than appending to it, so calling it after +/// `with_default_features()` would silently drop every built-in (including +/// `abs`). To avoid that footgun, the chaos UDFs are registered on the already +/// built `SessionState` via `FunctionRegistry::register_udf`, which inserts +/// into the existing map instead of replacing it. +pub fn chaos_session_state( + config: SessionConfig, +) -> datafusion::error::Result { + let mut state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .build(); + + for udf in [chaos_fail_udf(), chaos_delay_udf()] { + state.register_udf(udf)?; + } + + Ok(state) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn registry_resolves_the_chaos_udfs_by_name() { + // The executor resolves UDFs by name at plan-decode time. If these names + // are absent, every chaos query fails to deserialize on the executor. + let registry = chaos_function_registry(); + assert!(registry.udf("chaos_fail").is_ok()); + assert!(registry.udf("chaos_delay").is_ok()); + } + + #[test] + fn registry_retains_the_standard_functions() { + // Regression guard: overriding the registry must not drop the built-ins, + // or ordinary query operators stop working on the executor. + let registry = chaos_function_registry(); + assert!(registry.udf("abs").is_ok()); + } + + #[test] + fn session_state_resolves_the_chaos_udfs_by_name() { + // Mirrors registry_resolves_the_chaos_udfs_by_name, but for the + // scheduler's session state rather than the executor's registry. + let state = chaos_session_state(SessionConfig::new()).unwrap(); + assert!(state.udf("chaos_fail").is_ok()); + assert!(state.udf("chaos_delay").is_ok()); + } + + #[test] + fn session_state_retains_the_standard_functions() { + // Regression guard for the with_scalar_functions replace-vs-append + // footgun described on chaos_session_state: the scheduler must still be + // able to plan ordinary queries that use built-in functions. + let state = chaos_session_state(SessionConfig::new()).unwrap(); + assert!(state.udf("abs").is_ok()); + } +} diff --git a/chaos-testing/src/udf.rs b/chaos-testing/src/udf.rs new file mode 100644 index 0000000000..afd7987997 --- /dev/null +++ b/chaos-testing/src/udf.rs @@ -0,0 +1,328 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::budget::FaultBudget; +use arrow::array::{Array, BooleanArray}; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::cast::as_boolean_array; +use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, + TypeSignature, Volatility, +}; +use std::path::Path; +use std::sync::Arc; + +/// `chaos_fail(guard BOOLEAN, mode UTF8, budget_dir UTF8) -> BOOLEAN` +/// +/// Pass-through: returns `guard` unchanged. If any row in the batch has +/// `guard = true` and a fault token is available in `budget_dir`, consumes one +/// token and injects a fault: +/// +/// - `io` -> `DataFusionError::IoError` (Ballista: retryable, counts to failures) +/// - `exec` -> `DataFusionError::Execution` (Ballista: non-retryable) +/// - `panic` -> `panic!` (Ballista: caught, becomes non-retryable Internal) +/// +/// Volatile so DataFusion neither constant-folds nor CSEs it away. +#[derive(Debug, PartialEq, Eq, Hash)] +struct ChaosFail { + signature: Signature, +} + +impl Default for ChaosFail { + fn default() -> Self { + Self { + signature: Signature::new( + TypeSignature::Exact(vec![ + DataType::Boolean, + DataType::Utf8, + DataType::Utf8, + ]), + Volatility::Volatile, + ), + } + } +} + +impl ScalarUDFImpl for ChaosFail { + fn name(&self) -> &str { + "chaos_fail" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let guard = args.args[0].clone().into_array(args.number_rows)?; + let mode = scalar_utf8(&args.args[1], "mode")?; + let budget_dir = scalar_utf8(&args.args[2], "budget_dir")?; + + let guard = as_boolean_array(&guard)?; + let selected = (0..guard.len()).any(|i| !guard.is_null(i) && guard.value(i)); + + if selected && FaultBudget::open(Path::new(&budget_dir)).try_consume() { + match mode.as_str() { + "io" => { + let msg = "chaos_fail: injected retryable IO fault"; + log::error!("{msg}"); + return Err(DataFusionError::IoError(std::io::Error::other(msg))); + } + "exec" => { + let msg = + "chaos_fail: injected non-retryable execution fault".to_string(); + log::error!("{msg}"); + return Err(DataFusionError::Execution(msg)); + } + "panic" => { + log::error!("chaos_fail: injected panic"); + panic!("chaos_fail: injected panic"); + } + other => { + return Err(DataFusionError::Configuration(format!( + "chaos_fail: unknown mode {other:?} (expected io, exec, or panic)" + ))); + } + } + } + + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from( + (0..guard.len()) + .map(|i| (!guard.is_null(i)).then(|| guard.value(i))) + .collect::>>(), + )))) + } +} + +/// `chaos_delay(guard BOOLEAN, ms INT64) -> BOOLEAN` +/// +/// Pass-through: returns `guard` unchanged, sleeping `ms` milliseconds per batch +/// in which any row has `guard = true`. Used to hold a stage open long enough for +/// the harness to kill an executor while the stage is genuinely running. +#[derive(Debug, PartialEq, Eq, Hash)] +struct ChaosDelay { + signature: Signature, +} + +impl Default for ChaosDelay { + fn default() -> Self { + Self { + signature: Signature::new( + TypeSignature::Exact(vec![DataType::Boolean, DataType::Int64]), + Volatility::Volatile, + ), + } + } +} + +impl ScalarUDFImpl for ChaosDelay { + fn name(&self) -> &str { + "chaos_delay" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Boolean) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let guard = args.args[0].clone().into_array(args.number_rows)?; + let ms = match &args.args[1] { + ColumnarValue::Scalar(s) => match s { + datafusion::scalar::ScalarValue::Int64(Some(v)) => *v as u64, + other => { + return Err(DataFusionError::Configuration(format!( + "chaos_delay: ms must be a non-null INT64 literal, got {other:?}" + ))); + } + }, + ColumnarValue::Array(_) => { + return Err(DataFusionError::Configuration( + "chaos_delay: ms must be a literal, not a column".to_string(), + )); + } + }; + + let guard = as_boolean_array(&guard)?; + let selected = (0..guard.len()).any(|i| !guard.is_null(i) && guard.value(i)); + if selected { + std::thread::sleep(std::time::Duration::from_millis(ms)); + } + + Ok(ColumnarValue::Array(Arc::new(BooleanArray::from( + (0..guard.len()) + .map(|i| (!guard.is_null(i)).then(|| guard.value(i))) + .collect::>>(), + )))) + } +} + +fn scalar_utf8(value: &ColumnarValue, arg: &str) -> Result { + match value { + ColumnarValue::Scalar(datafusion::scalar::ScalarValue::Utf8(Some(s))) => { + Ok(s.clone()) + } + other => Err(DataFusionError::Configuration(format!( + "chaos_fail: {arg} must be a non-null UTF8 literal, got {other:?}" + ))), + } +} + +/// The `chaos_fail` UDF. +pub fn chaos_fail_udf() -> Arc { + Arc::new(ScalarUDF::from(ChaosFail::default())) +} + +/// The `chaos_delay` UDF. +pub fn chaos_delay_udf() -> Arc { + Arc::new(ScalarUDF::from(ChaosDelay::default())) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::budget::FaultBudget; + use arrow::array::BooleanArray; + use datafusion::arrow::array::RecordBatch; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::prelude::SessionContext; + use std::sync::Arc; + + fn temp_dir(name: &str) -> std::path::PathBuf { + let dir = std::env::temp_dir().join(format!("ballista-chaos-udf-{name}")); + let _ = std::fs::remove_dir_all(&dir); + dir + } + + /// Build a one-column table `t(guard BOOLEAN)` with the given values. + async fn ctx_with_guards(guards: Vec) -> SessionContext { + let ctx = SessionContext::new(); + ctx.register_udf(chaos_fail_udf().as_ref().clone()); + ctx.register_udf(chaos_delay_udf().as_ref().clone()); + + let schema = Arc::new(Schema::new(vec![Field::new( + "guard", + DataType::Boolean, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(BooleanArray::from(guards))], + ) + .unwrap(); + ctx.register_batch("t", batch).unwrap(); + ctx + } + + #[tokio::test] + async fn io_mode_errors_when_a_token_is_available() { + let dir = temp_dir("io-fires"); + FaultBudget::create(&dir, 1).unwrap(); + let ctx = ctx_with_guards(vec![true]).await; + + let sql = format!("SELECT chaos_fail(guard, 'io', '{}') FROM t", dir.display()); + let err = ctx.sql(&sql).await.unwrap().collect().await.unwrap_err(); + + // Must be an IoError: that is the only variant Ballista treats as retryable. + assert!( + matches!(err, datafusion::error::DataFusionError::IoError(_)), + "expected IoError, got {err:?}" + ); + } + + #[tokio::test] + async fn passes_through_when_budget_is_exhausted() { + let dir = temp_dir("io-exhausted"); + FaultBudget::create(&dir, 0).unwrap(); + let ctx = ctx_with_guards(vec![true, false]).await; + + let sql = format!( + "SELECT chaos_fail(guard, 'io', '{}') AS g FROM t", + dir.display() + ); + let batches = ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + + // Pass-through: output must equal the input guard column. + let col = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(col.value(0)); + assert!(!col.value(1)); + } + + #[tokio::test] + async fn does_not_fire_when_no_guard_row_is_true() { + let dir = temp_dir("io-no-guard"); + FaultBudget::create(&dir, 1).unwrap(); + let ctx = ctx_with_guards(vec![false, false]).await; + + let sql = format!("SELECT chaos_fail(guard, 'io', '{}') FROM t", dir.display()); + ctx.sql(&sql).await.unwrap().collect().await.unwrap(); + + // The token must be untouched: the guard never selected this data. + assert_eq!(FaultBudget::open(&dir).remaining(), 1); + } + + /// Tests that panic mode injects a panic when a guard row is true. + /// The panic crosses the async stream boundary. The authoritative end-to-end + /// panic test is Scenario C in Task 8, where the executor's `catch_unwind` + /// converts the panic into a `FailedTask`. + #[should_panic(expected = "chaos_fail: injected panic")] + #[tokio::test] + async fn panic_mode_panics() { + let dir = temp_dir("panic-fires"); + FaultBudget::create(&dir, 1).unwrap(); + let ctx = ctx_with_guards(vec![true]).await; + + let sql = format!( + "SELECT chaos_fail(guard, 'panic', '{}') FROM t", + dir.display() + ); + let _ = ctx.sql(&sql).await.unwrap().collect().await; + } + + #[tokio::test] + async fn delay_sleeps_and_passes_through() { + let ctx = ctx_with_guards(vec![true]).await; + let start = std::time::Instant::now(); + + let batches = ctx + .sql("SELECT chaos_delay(guard, 150) AS g FROM t") + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert!(start.elapsed() >= std::time::Duration::from_millis(150)); + let col = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(col.value(0)); + } +} diff --git a/chaos-testing/tests/common/mod.rs b/chaos-testing/tests/common/mod.rs new file mode 100644 index 0000000000..1a2755d3d2 --- /dev/null +++ b/chaos-testing/tests/common/mod.rs @@ -0,0 +1,131 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use ballista::prelude::{SessionConfigExt, SessionContextExt}; +use ballista_core::config::BALLISTA_ADAPTIVE_PLANNER_ENABLED; +use chaos_testing::budget::FaultBudget; +use chaos_testing::cluster::TestCluster; +use chaos_testing::fixture::Fixture; +use datafusion::arrow::util::pretty::pretty_format_batches; +use datafusion::execution::SessionStateBuilder; +use datafusion::prelude::{SessionConfig, SessionContext}; + +/// One cluster plus its fixture, wired for a single scenario. +pub struct ChaosRun { + pub cluster: TestCluster, + pub fixture: Fixture, + ctx: SessionContext, +} + +impl ChaosRun { + pub async fn start(aqe: bool, executors: usize) -> Self { + Self::start_with(aqe, executors, 5).await + } + + /// `executor_timeout_seconds` selects which HA mechanism a kill surfaces + /// through: short biases toward ExecutorLost (heartbeat expiry), long biases + /// toward FetchPartitionError (a downstream fetch from a dead executor). + pub async fn start_with( + aqe: bool, + executors: usize, + executor_timeout_seconds: u64, + ) -> Self { + let _ = env_logger::builder().is_test(true).try_init(); + + let cluster = TestCluster::builder() + .executors(executors) + .executor_timeout_seconds(executor_timeout_seconds) + .start() + .await + .expect("cluster must start"); + + let fixture = Fixture::write(cluster.shared_dir()) + .await + .expect("fixture must be written"); + + let config = SessionConfig::new_with_ballista() + .set_bool(BALLISTA_ADAPTIVE_PLANNER_ENABLED, aqe); + let state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .build(); + + let ctx = SessionContext::remote_with_state(&cluster.scheduler_url(), state) + .await + .expect("client must connect to the scheduler"); + + // Registered *after* `remote_with_state`, not baked into the initial + // SessionState: `SessionStateExt::upgrade_for_ballista` (ballista/core) + // rebuilds the state via `with_scalar_functions(ballista_scalar_functions())`, + // which replaces rather than merges the scalar-function map, silently + // dropping any custom UDF registered beforehand. `register_udf` mutates + // the already-upgraded state's registry directly, so it survives. + ctx.register_udf(chaos_testing::udf::chaos_fail_udf().as_ref().clone()); + ctx.register_udf(chaos_testing::udf::chaos_delay_udf().as_ref().clone()); + + for stmt in fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + + Self { + cluster, + fixture, + ctx, + } + } + + /// A fault budget in the cluster's shared directory, visible to every executor. + pub fn budget(&self, name: &str, tokens: usize) -> FaultBudget { + let dir = self.cluster.shared_dir().join(format!("budget-{name}")); + FaultBudget::create(&dir, tokens).expect("budget must be created") + } + + /// Run a query on the cluster, returning the formatted result. + pub async fn sql(&self, sql: &str) -> Result { + let df = self.ctx.sql(sql).await.map_err(|e| e.to_string())?; + let batches = df.collect().await.map_err(|e| e.to_string())?; + Ok(pretty_format_batches(&batches) + .map_err(|e| e.to_string())? + .to_string()) + } + + /// The expected answer, computed by plain local DataFusion. + pub async fn local_baseline(&self) -> String { + let ctx = SessionContext::new(); + for stmt in self.fixture.register_sql() { + ctx.sql(&stmt).await.unwrap().collect().await.unwrap(); + } + let batches = ctx + .sql(Fixture::baseline_query()) + .await + .unwrap() + .collect() + .await + .unwrap(); + pretty_format_batches(&batches).unwrap().to_string() + } + + /// A clone of the session context, for running a query concurrently with a fault. + /// + /// Not used until Task 9's executor-loss scenarios, which submit a query in + /// a spawned task while the main task polls the scheduler and kills an + /// executor. + #[allow(dead_code)] + pub fn clone_ctx(&self) -> SessionContext { + self.ctx.clone() + } +} diff --git a/chaos-testing/tests/ha.rs b/chaos-testing/tests/ha.rs new file mode 100644 index 0000000000..fe9f93d11f --- /dev/null +++ b/chaos-testing/tests/ha.rs @@ -0,0 +1,418 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end high-availability scenarios against a real multi-process cluster. +//! +//! Every scenario runs under both AQE settings. The AQE-on axis is where bugs are +//! expected: a resubmitted stage under AQE is re-planned against runtime +//! statistics, so a re-run map stage can come back with a different plan than the +//! one whose output was lost. +//! +//! Every test in this file spawns a whole multi-process cluster, so this file +//! must always be run with `--test-threads=1` or ports and CPU will be +//! exhausted by concurrent clusters. + +mod common; + +use common::ChaosRun; +use rstest::rstest; + +/// The cluster must agree with local DataFusion before any fault is injected. +/// Every recovery scenario asserts against this baseline, so if it is wrong, +/// every other assertion is meaningless. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn baseline_matches_local_datafusion(#[case] aqe: bool) { + let run = ChaosRun::start(aqe, 2).await; + + let expected = run.local_baseline().await; + let actual = run + .sql(ballista_chaos_query_baseline()) + .await + .expect("baseline query must succeed on the cluster"); + + assert_eq!( + actual, expected, + "cluster result must match local DataFusion" + ); +} + +fn ballista_chaos_query_baseline() -> &'static str { + chaos_testing::fixture::Fixture::baseline_query() +} + +use chaos_testing::fixture::Fixture; + +/// Scenario A: one retryable (IO) fault, budget 1. +/// +/// A single task attempt anywhere in the cluster faults; the budget is then +/// exhausted, so the retry must succeed. The load-bearing assertion is that the +/// result still equals the baseline: a retried stage is exactly where duplicated +/// or dropped partitions would show up. +/// +/// This scenario is split into one test per AQE setting rather than being an +/// `rstest` over both, because historically only the AQE-on case failed and +/// `rstest` cannot ignore an individual case. +/// +/// Both cases originally passed or reproduced #2028 (the AQE-on case: +/// `Shared(IoError)` misses the classifier's shallow match). #2028 was fixed +/// by classifying on `find_root()` (#2119), but the sort-shuffle writer +/// refactor (#2038/#2106) then made both cases fail the same way: the shuffle +/// write coordinator flattens any task error into +/// `DataFusionError::Execution(format!("{e:?}"))` +/// (`ballista/core/src/execution_plans/shuffle_writer.rs`, error arm of the +/// coordinator fan-out), so the injected `IoError` reaches the classifier as +/// inert text inside an `Execution` string, `find_root()` has nothing to see +/// through, and the task is marked non-retryable. That is the same +/// type-erasing mechanism tracked for fetch failures by #2027. +/// +/// Ignored, not deleted or weakened. Un-ignore both as the regression tests +/// when #2027's error-flattening is fixed; see chaos-testing/README.md, +/// Finding 2. +#[tokio::test] +#[ignore = "reproduces #2027's error flattening: the shuffle writer Debug-formats the IoError into an opaque Execution string, so it is misclassified as non-retryable"] +async fn retryable_fault_is_retried_and_result_is_correct_aqe_off() { + retryable_fault_is_retried_and_result_is_correct(false).await; +} + +/// See `retryable_fault_is_retried_and_result_is_correct_aqe_off` above; the +/// AQE-on case fails identically (the error arrives as +/// `Execution("Shared(IoError(..))")` — stringified before the classifier, +/// which is what distinguishes this from the fixed #2028). +#[tokio::test] +#[ignore = "reproduces #2027's error flattening: the shuffle writer Debug-formats the IoError into an opaque Execution string, so it is misclassified as non-retryable"] +async fn retryable_fault_is_retried_and_result_is_correct_aqe_on() { + retryable_fault_is_retried_and_result_is_correct(true).await; +} + +async fn retryable_fault_is_retried_and_result_is_correct(aqe: bool) { + let run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + let budget = run.budget("scenario-a", 1); + let sql = Fixture::chaos_query(&format!( + "chaos_fail(f.key = 7, 'io', '{}')", + budget.dir().display() + )); + + let actual = run + .sql(&sql) + .await + .expect("query must recover from one IO fault"); + + assert_eq!( + actual, expected, + "result after retry must match the baseline" + ); + assert_eq!(budget.remaining(), 0, "the fault must actually have fired"); +} + +/// Scenario B: an inexhaustible retryable fault. +/// +/// The budget far exceeds task_max_failures (4), so every attempt faults. The job +/// must fail rather than retry forever, and the cluster must remain usable +/// afterwards — a scheduler that wedges after a failed job is an HA bug. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn exhausted_retries_fail_the_job_and_leave_the_cluster_healthy(#[case] aqe: bool) { + let run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + let budget = run.budget("scenario-b", 99); + let sql = Fixture::chaos_query(&format!( + "chaos_fail(f.key = 7, 'io', '{}')", + budget.dir().display() + )); + + let err = run + .sql(&sql) + .await + .expect_err("the job must fail once retries are exhausted"); + assert!(!err.is_empty(), "the failure must carry an error message"); + + // The cluster must still serve queries. A chaos-free query proves the + // scheduler and both executors survived the failed job. + let after = run + .sql(Fixture::baseline_query()) + .await + .expect("cluster must still be healthy after a failed job"); + assert_eq!(after, expected); +} + +/// Scenario C: a panicking task. +/// +/// The executor catches the panic (executor.rs:237) and turns it into a +/// non-retryable Internal error, so the job fails immediately with no retry. This +/// test encodes *today's* behaviour, not necessarily the desired behaviour: if we +/// later decide panics should be retryable, this is the test that changes. +/// +/// The second assertion is the important one: the executor process must survive. +/// A panic in one task must not take down the whole executor and every other task +/// running on it. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn panicking_task_fails_the_job_but_the_executor_survives(#[case] aqe: bool) { + let mut run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + let budget = run.budget("scenario-c", 1); + let sql = Fixture::chaos_query(&format!( + "chaos_fail(f.key = 7, 'panic', '{}')", + budget.dir().display() + )); + + let err = run + .sql(&sql) + .await + .expect_err("a panicking task must fail the job"); + assert!(!err.is_empty()); + assert_eq!(budget.remaining(), 0, "the panic must actually have fired"); + + // Both executor processes must still be alive. + assert!( + run.cluster.executor_is_alive(0), + "executor 0 died on a task panic" + ); + assert!( + run.cluster.executor_is_alive(1), + "executor 1 died on a task panic" + ); + + // And the cluster must still serve queries. + let after = run + .sql(Fixture::baseline_query()) + .await + .expect("cluster must still be healthy after a panicking task"); + assert_eq!(after, expected); +} + +use std::time::Duration; + +/// Scenario D: SIGKILL an executor while it is running tasks. +/// +/// `chaos_delay` holds stage 1 open so the kill lands while tasks are genuinely +/// in flight. The scheduler must detect the loss, reschedule the dead executor's +/// tasks onto the survivor, and still return the correct result. +/// +/// Both cases reproduce #2027: the shuffle reader's typed +/// `BallistaError::FetchFailed` is flattened into an opaque +/// `DataFusionError::Execution` before the executor reports its `TaskStatus`, +/// so the scheduler never sees the `FetchPartitionError` that would make it +/// resubmit the lost map stage, and fails the query instead of recovering it. +/// +/// Ignored, not deleted or weakened. Un-ignore it as the regression test when +/// #2027 is fixed; see chaos-testing/README.md, Finding 1. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +#[ignore = "reproduces #2027: FetchFailed loses its type, so the map stage is never resubmitted"] +async fn executor_killed_mid_stage_is_recovered(#[case] aqe: bool) { + let mut run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + // Delay every scan task by 300ms per batch so the stage stays running long + // enough to kill an executor inside it. + let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 300)"); + + // Submit the query concurrently, then kill executor 0 once stage 1 is running. + let query = tokio::spawn({ + let ctx = run.clone_ctx(); + let sql = sql.clone(); + async move { ctx.sql(&sql).await?.collect().await } + }); + + let job_id = run.cluster.running_job_id().await.expect("job must appear"); + run.cluster + .await_stage_running(&job_id, 1) + .await + .expect("stage 1 must start running"); + run.cluster.kill_executor(0).expect("kill executor 0"); + + let batches = tokio::time::timeout(Duration::from_secs(120), query) + .await + .expect("query must not hang after an executor is killed") + .expect("query task must not panic") + .expect("query must recover from the lost executor"); + + let actual = datafusion::arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + assert_eq!( + actual, expected, + "result after executor loss must match the baseline" + ); +} + +/// Scenario E: SIGKILL a map-side executor after it wrote shuffle output. +/// +/// The downstream stage must fetch shuffle partitions from an executor that no +/// longer exists. Recovery requires re-running the map stage. The executor +/// timeout is raised to 60s to bias the failure toward the FetchPartitionError +/// path rather than the heartbeat-expiry ExecutorLost path; both are valid +/// recoveries, so the assertion is on correctness, and the path that actually +/// fired is only recorded. +/// +/// Reproduces #2027: whenever the reduce stage genuinely has to fetch from the +/// dead executor, the typed `FetchFailed` arrives flattened inside +/// `DataFusionError::Shared`, the scheduler never resubmits the map stage, and +/// the job fails. The scenario only passes when the kill happens to land after +/// the reduce tasks have already fetched their partitions, which makes it +/// timing-dependent (it failed in CI, aqe_off case). +/// +/// Ignored, not deleted or weakened. Un-ignore it as the regression test when +/// #2027 is fixed; see chaos-testing/README.md, Finding 1. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +#[ignore = "reproduces #2027: FetchFailed loses its type, so the map stage is never resubmitted"] +async fn executor_killed_after_shuffle_write_is_recovered(#[case] aqe: bool) { + let mut run = ChaosRun::start_with(aqe, 2, 60).await; + let expected = run.local_baseline().await; + + // Delay the *aggregate* side so the reduce stage is slow, giving us a window + // between "stage 1 succeeded" and "stage 2 has finished fetching". + let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 50)"); + + let query = tokio::spawn({ + let ctx = run.clone_ctx(); + let sql = sql.clone(); + async move { ctx.sql(&sql).await?.collect().await } + }); + + let job_id = run.cluster.running_job_id().await.expect("job must appear"); + run.cluster + .await_stage_successful(&job_id, 1) + .await + .expect("stage 1 must complete before we kill its executor"); + run.cluster.kill_executor(0).expect("kill executor 0"); + + let batches = tokio::time::timeout(Duration::from_secs(180), query) + .await + .expect("query must not hang after shuffle output is lost") + .expect("query task must not panic") + .expect("query must recover by re-running the map stage"); + + let actual = datafusion::arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string(); + assert_eq!( + actual, expected, + "result after shuffle-output loss must match the baseline" + ); +} + +/// Scenario F: an executor is killed and restarted; the cluster must reabsorb it. +/// +/// The kill and the restart are separated by a wait for the scheduler to +/// actually reap the dead executor (`registered_executors` dropping to 1), +/// rather than restarting immediately. SIGKILL does not deregister the +/// executor: the scheduler keeps listing it until its heartbeat expires +/// (`executor_timeout_seconds`), so a restart fired immediately after the +/// kill lands while the scheduler still counts *three* executors (the dead +/// one, the survivor, and the freshly restarted one) — a harness race, not a +/// Ballista bug, that used to make this scenario fail with `left: 3, right: +/// 2`. Waiting for the reap first means the final assertion is actually +/// testing what the scenario name promises: that a restarted executor +/// rejoins a cluster that has already noticed it was gone. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn restarted_executor_rejoins_and_serves_queries(#[case] aqe: bool) { + let mut run = ChaosRun::start(aqe, 2).await; + let expected = run.local_baseline().await; + + run.cluster.kill_executor(0).expect("kill executor 0"); + run.cluster + .await_executor_count(1) + .await + .expect("scheduler must reap the killed executor before we restart it"); + + run.cluster + .restart_executor(0) + .await + .expect("restarted executor must re-register"); + + assert_eq!( + run.cluster.registered_executors().await.unwrap(), + 2, + "both executors must be registered after the restart" + ); + + let actual = run + .sql(Fixture::baseline_query()) + .await + .expect("cluster must serve queries after an executor restart"); + assert_eq!(actual, expected); +} + +/// Scenario G: every executor is killed mid-query. +/// +/// There is no executor left to recover onto, so the job cannot succeed. Once +/// the last executor is reaped, the scheduler waits a bounded grace period for an +/// executor to (re)register and then fails the job, rather than waiting forever +/// for tasks that can never be scheduled. Both cases assert the query terminates +/// with an error that names the executor loss, instead of hanging. +/// +/// Regression test for #2029. The grace period is turned down via the cluster +/// builder so the job fails a second or so after the reap instead of waiting out +/// the 30s default. +#[rstest] +#[case::aqe_off(false)] +#[case::aqe_on(true)] +#[tokio::test] +async fn killing_every_executor_terminates_the_job(#[case] aqe: bool) { + let mut run = ChaosRun::start(aqe, 2).await; + + let sql = Fixture::chaos_query("chaos_delay(f.key >= 0, 300)"); + + let query = tokio::spawn({ + let ctx = run.clone_ctx(); + let sql = sql.clone(); + async move { ctx.sql(&sql).await?.collect().await } + }); + + let job_id = run.cluster.running_job_id().await.expect("job must appear"); + run.cluster + .await_any_stage_running(&job_id) + .await + .expect("the job must start running a task before we kill its executors"); + + run.cluster.kill_executor(0).expect("kill executor 0"); + run.cluster.kill_executor(1).expect("kill executor 1"); + + // Once the executors are reaped and the grace period elapses, the scheduler + // must fail the job rather than hang forever (#2029). + let result = tokio::time::timeout(Duration::from_secs(120), query) + .await + .expect("job must terminate, not hang, after every executor is killed") + .expect("query task should not panic"); + let err = result.expect_err("query must fail once every executor is lost"); + let msg = err.to_string().to_lowercase(); + assert!( + msg.contains("executor"), + "failure should name the executor loss, got: {err}" + ); +} diff --git a/ci/scripts/rust_toml_fmt.sh b/ci/scripts/rust_toml_fmt.sh index 3ce50ace0a..ad79505373 100755 --- a/ci/scripts/rust_toml_fmt.sh +++ b/ci/scripts/rust_toml_fmt.sh @@ -17,5 +17,24 @@ # specific language governing permissions and limitations # under the License. -set -ex -find . -mindepth 2 -name 'Cargo.toml' -exec cargo tomlfmt -k -p {} \; +# Checks Cargo.toml formatting using taplo. The formatting rules live in +# taplo.toml at the repository root. +# +# Usage: +# ci/scripts/rust_toml_fmt.sh # check formatting (default, used in CI) +# ci/scripts/rust_toml_fmt.sh --write # reformat files in place + +set -e + +if ! command -v taplo &> /dev/null; then + echo "Installing taplo using cargo" + cargo install taplo-cli --version 0.10.0 --locked +fi + +if [ "${1:-}" = "--write" ]; then + taplo format +else + # `taplo format --check` exits non-zero if any file is not correctly + # formatted. Run `ci/scripts/rust_toml_fmt.sh --write` to fix violations. + taplo format --check +fi diff --git a/dev/rust_lint.sh b/dev/rust_lint.sh index f634266010..6149237ec9 100755 --- a/dev/rust_lint.sh +++ b/dev/rust_lint.sh @@ -17,11 +17,7 @@ # specific language governing permissions and limitations # under the License. set -e -if ! command -v cargo-tomlfmt &> /dev/null; then - echo "Installing cargo-tomlfmt using cargo" - cargo install cargo-tomlfmt -fi ci/scripts/rust_fmt.sh ci/scripts/rust_clippy.sh -ci/scripts/rust_toml_fmt.sh \ No newline at end of file +ci/scripts/rust_toml_fmt.sh diff --git a/docs/source/upgrading/55.0.0.md b/docs/source/upgrading/55.0.0.md index ccbabfe89e..8c2432bd77 100644 --- a/docs/source/upgrading/55.0.0.md +++ b/docs/source/upgrading/55.0.0.md @@ -71,3 +71,47 @@ behavior. If you relied on it being `true` and want small joins to broadcast, se `HashJoinExec`, which remains eligible for broadcast promotion. Note this opts you into the non-spilling hash join for _all_ joins in the session, not only the small ones. + +### API changes + +#### `ballista_executor::as_task_status` signature reshaped + +`as_task_status` now takes a `TaskCompletionExtras` struct in place of the +`operator_metrics: Option>` positional parameter. The +struct also carries a new `runtime_stats: Vec` field used to +transport `RuntimeStatsExec` reports back to the scheduler. It is marked +`#[non_exhaustive]` with a `Default` impl, so future additions to the struct +are non-breaking for callers that construct via `..Default::default()`. + +The parameter order also changed: `execution_times` now precedes the extras +struct. + +**Action required:** external callers building a `TaskStatus` via +`as_task_status` should migrate from: + +```rust +as_task_status( + execution_result, + executor_id, + stage_attempt_num, + key, + operator_metrics, + execution_times, +) +``` + +to: + +```rust +as_task_status( + execution_result, + executor_id, + stage_attempt_num, + key, + execution_times, + TaskCompletionExtras { + operator_metrics, + ..Default::default() + }, +) +``` diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index c40f1a582c..c0c721cf09 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -119,9 +119,10 @@ standard DataFusion settings. | ballista.client.pull | Boolean | false | Should client employ pull or push job tracking. In pull mode client will make a request to server in the loop, until job finishes. Pull mode is kept for legacy clients. | | ballista.client.use_tls | Boolean | false | Should connection between client, scheduler, and executors use TLS. | | ballista.job.name | Utf8 | (none) | Sets the job name that will appear in the web user interface for any submitted jobs | -| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 | Byte-size threshold below which a hash join's smaller side is promoted to CollectLeft and lowered via the broadcast pattern. Governs broadcast selection under both the static distributed planner and adaptive query planning (AQE). Set to 0 to disable promotion. | +| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 | Byte-size threshold below which a hash join's smaller side is promoted to CollectLeft and lowered via the broadcast pattern. Governs broadcast selection under both the static distributed planner and adaptive query planning (AQE). It also caps null-aware anti joins with a known build size because they require single-task CollectLeft execution. Set to 0 to disable promotion and reject null-aware anti joins. | | ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000 | Row-count threshold below which a hash join's smaller side is promoted to CollectLeft and lowered via the broadcast pattern, used as a fallback when byte-size statistics are unavailable. Applies to adaptive query planning (AQE). Set to 0 to disable promotion via the row-count path. | | ballista.optimizer.hash_join_max_build_partition_bytes | UInt64 | 67108864 | Maximum per-partition hash-join build-side bytes for a Partitioned hash join under AQE. A build partition larger than this falls back to SortMergeJoin (spillable). Defaults to 64 MiB; 0 disables the check, which makes AQE use a hash join regardless of build size. | +| ballista.optimizer.not_in_subquery_rewrite | Boolean | true | Rewrites uncorrelated NOT IN (subquery) filter predicates into a plain anti join plus a one-row count aggregate during logical optimization. The rewrite avoids DataFusion's null-aware hash join, which Ballista must otherwise execute in a single task. Set to false to keep the null-aware join and its single-task lowering. | | ballista.planner.adaptive.enabled | Boolean | false | Enables Adaptive Query Planning (EXPERIMENTAL) | | ballista.planner.adaptive_join.enabled | Boolean | true | Enables the AQE dynamic join-selection rule (SelectJoinRule). When true (default), DynamicJoinSelectionExec nodes are resolved to concrete HashJoin or CollectLeft join implementations at runtime. Disable only for debugging. | | ballista.planner.coalesce.enabled | Boolean | false | Enables the AQE coalesce-shuffle-partitions rule. Disabled by default — opt in when fewer/larger downstream tasks matter more than parallelism. | diff --git a/docs/source/user-guide/tuning-guide.md b/docs/source/user-guide/tuning-guide.md index 6e0ba0c511..0ce634d548 100644 --- a/docs/source/user-guide/tuning-guide.md +++ b/docs/source/user-guide/tuning-guide.md @@ -253,11 +253,11 @@ Adaptive Query Planning is EXPERIMENTAL, should be used for testing purposes onl ### Configuration -| key | type | default | description | -| ------------------------------------------------- | ------- | -------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| ballista.planner.adaptive.enabled | Boolean | false | Enables the adaptive planner. Experimental. | -| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 | Byte-size threshold below which a hash join's smaller side is broadcast (`CollectLeft`). Governs both the static planner and AQE. Set to 0 to disable broadcast joins. | -| ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000 | Row-count fallback threshold used when byte-size statistics are unavailable. Applies to AQE. Set to 0 to disable promotion via the row-count path. | +| key | type | default | description | +| ------------------------------------------------- | ------- | -------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| ballista.planner.adaptive.enabled | Boolean | false | Enables the adaptive planner. Experimental. | +| ballista.optimizer.broadcast_join_threshold_bytes | UInt64 | 10485760 | Byte-size threshold below which a hash join's smaller side is broadcast (`CollectLeft`). Also caps null-aware anti joins with known build sizes because they run in one task. Set to 0 to disable broadcasts and reject null-aware anti joins. | +| ballista.optimizer.broadcast_join_threshold_rows | UInt64 | 1000000 | Row-count fallback threshold used when byte-size statistics are unavailable. Applies to AQE. Set to 0 to disable promotion via the row-count path. | ### What AQE does today @@ -271,6 +271,13 @@ implemented: - **Broadcast join selection.** When a join input's runtime size falls under `ballista.optimizer.broadcast_join_threshold_bytes` (or the row-count fallback), the smaller side is broadcast (`CollectLeft`) instead of shuffled. + Null-aware anti joins use `CollectLeft` with a single probe task because their + state cannot be coordinated across executors. A known oversized build side is + rejected instead of producing an incorrect distributed result. In practice + most `NOT IN (subquery)` predicates never plan a null-aware join at all: + `ballista.optimizer.not_in_subquery_rewrite` (enabled by default) rewrites + them into a fully distributable anti join plus a one-row count aggregate + during logical optimization. - **Empty stage elimination.** When a completed stage produces zero rows, its downstream exchange is replaced with an empty execution node, and emptiness is propagated up the plan so downstream stages are skipped entirely. diff --git a/examples/Cargo.toml b/examples/Cargo.toml index 531ce36565..a10926fb36 100644 --- a/examples/Cargo.toml +++ b/examples/Cargo.toml @@ -76,7 +76,7 @@ tokio = { workspace = true, features = [ "rt", "rt-multi-thread", "sync", - "parking_lot" + "parking_lot", ] } tonic = { workspace = true } url = { workspace = true } diff --git a/examples/examples/mtls-cluster.rs b/examples/examples/mtls-cluster.rs index 609935d6ba..08c4690b04 100644 --- a/examples/examples/mtls-cluster.rs +++ b/examples/examples/mtls-cluster.rs @@ -433,7 +433,7 @@ async fn run_executor() -> Result<(), Box> { info!("Starting execution poll loop..."); let health = ballista_executor::health::ExecutorHealth::new(); let poll_handle = tokio::spawn(async move { - execution_loop::poll_loop(scheduler, executor, codec, None, health).await + execution_loop::poll_loop(scheduler, executor, codec, None, None, health).await }); tokio::select! { diff --git a/python/uv.lock b/python/uv.lock index 54493c14b0..81ee8f45bb 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -270,62 +270,59 @@ wheels = [ [[package]] name = "cryptography" -version = "48.0.1" +version = "50.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/12/45/870e7f4bef50e5f53b9f51d4428aee5290eedf58ba443f16b1ebb7ab8e66/cryptography-48.0.1.tar.gz", hash = "sha256:266f4ee051abb2f725b74ef8072b521ce1feacf685a3364fa6a6b45548db791a", size = 832989, upload-time = "2026-06-09T22:32:31.8Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1b/bc/ee4137cbbe105652c0ee4252792b78fc8e7afa4b8e61d9d5dc05a7f45731/cryptography-48.0.1-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:3e4a1a3232eef2e6c732827d5722db29a0cc8b27af2a4d865b094cf954be9ca1", size = 8008324, upload-time = "2026-06-09T22:31:00.702Z" }, - { url = "https://files.pythonhosted.org/packages/d5/85/6379d42181bfc713094f081360fc5784d6c816b599d45e7f082502d173ce/cryptography-48.0.1-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:32143b24adb918f078134e1e230f1eb8cc04886b92c28b5f0041aaf3e5699225", size = 4696243, upload-time = "2026-06-09T22:32:33.446Z" }, - { url = "https://files.pythonhosted.org/packages/9c/87/c85d147b53323c7eb4d850920c8901377323c2a0ff8d79c262d4fee89aa2/cryptography-48.0.1-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f0d27a5696721ef7a672b8c810f6aded391058e0b9486e63e6d93baf765da691", size = 4713235, upload-time = "2026-06-09T22:31:40.141Z" }, - { url = "https://files.pythonhosted.org/packages/79/58/67cbf8cf1ee7c54b439ca07bbecf8362c07afc11a3724fea70f745784add/cryptography-48.0.1-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:eb86ce1af36fe65041b6db9a8bb064ee621a7e5fded0f80d475ec243477cd242", size = 4702323, upload-time = "2026-06-09T22:31:42.191Z" }, - { url = "https://files.pythonhosted.org/packages/89/c6/24266ac10c47f6cd2a865f4446062b466da1d1f10b27189eac00e61bf0c9/cryptography-48.0.1-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:b024e784ad6c077ee0147b35ea9cbfc1e34e1fd4c1dcca214c2794d73a12df08", size = 5300085, upload-time = "2026-06-09T22:31:58.703Z" }, - { url = "https://files.pythonhosted.org/packages/d2/bb/cc4b78784f97efc8c5874c2a9743708d172be6663024b34a0467885ae0c8/cryptography-48.0.1-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:3752f2dbc8f07a30aad2932c986cea495b03bb554887828225da104f732852b6", size = 4746137, upload-time = "2026-06-09T22:31:31.01Z" }, - { url = "https://files.pythonhosted.org/packages/1f/52/0c44de3f5267f8fbe8e835138017522a333436166e406f0db9b9e6e3033f/cryptography-48.0.1-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:bd81490cd5801d755cf97bb68ac191f14b708470b1c7cf4580f669b9c9264cd8", size = 4333867, upload-time = "2026-06-09T22:32:28.096Z" }, - { url = "https://files.pythonhosted.org/packages/9a/2e/772d7adbfa931537bc401640b7cac9976bff689bda187833e5d63b428e49/cryptography-48.0.1-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:66fd0771e7b9c6dcd44cf1120690d2338d16d72795cf40cae2786a39eba65429", size = 4701805, upload-time = "2026-06-09T22:31:38.284Z" }, - { url = "https://files.pythonhosted.org/packages/f8/a3/b06844f303873493c963caf581c04df31c7035e0c1b0f02c4814d319ec80/cryptography-48.0.1-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:3fd2ca57062b241c856670b073487d2e86c4637937ca5601e48f97bf8e11fc8f", size = 5258461, upload-time = "2026-06-09T22:31:04.187Z" }, - { url = "https://files.pythonhosted.org/packages/9f/13/8b765e2e12b07c74941caadb9d1c8fdc006c4dfbf2b8f2d610519758954d/cryptography-48.0.1-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:0ee6ea481db1ab889cba043ec1eda17bb9c1ea79db6722f779c3667f9f70322f", size = 4745488, upload-time = "2026-06-09T22:32:30.07Z" }, - { url = "https://files.pythonhosted.org/packages/2e/aa/48972bce55049b32a94f4907eda4d75fa385aad8a39506cc2fc72196ecf0/cryptography-48.0.1-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:f2ceef93cb096aa3c4cc4b5c94ca6131f9196d28c64d6111533402a9b2054d41", size = 4830256, upload-time = "2026-06-09T22:31:43.868Z" }, - { url = "https://files.pythonhosted.org/packages/47/a2/e5079a032fb85cf6005046ca92bbd78b0c82dad2b5751ab8c311659da06f/cryptography-48.0.1-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:9bd3f92d76217892b15df84ca256c2c113d386fdda7a7d8691aeeced976507c6", size = 4979117, upload-time = "2026-06-09T22:31:05.845Z" }, - { url = "https://files.pythonhosted.org/packages/b7/a0/8f50cae9c74e718ed769d63ed5c74bd0ea830c9550a74629cebd1b9c7bc7/cryptography-48.0.1-cp311-abi3-win32.whl", hash = "sha256:b9a32b876490d66c8bcc9963ef220199569748434ab01a9d6aaeabf88e7f5158", size = 3304154, upload-time = "2026-06-09T22:32:16.845Z" }, - { url = "https://files.pythonhosted.org/packages/c5/69/0572c77dbace6fef72f33755bd52ea399c71367250d366237f8691826b9e/cryptography-48.0.1-cp311-abi3-win_amd64.whl", hash = "sha256:39489bfca54c7a1f6b297efcd8bc608ab92d16c4ca631b0cad4da46724588b24", size = 3817138, upload-time = "2026-06-09T22:32:00.388Z" }, - { url = "https://files.pythonhosted.org/packages/42/06/3e768b4c3bc78201583fa35a0e18f640dd782ff41afba88f8545481a8874/cryptography-48.0.1-cp314-cp314t-macosx_10_9_universal2.whl", hash = "sha256:f817adc181390bd54f2f700107a7419040fb7c1bdf2fc26f36551a06a68c3345", size = 7989830, upload-time = "2026-06-09T22:31:07.8Z" }, - { url = "https://files.pythonhosted.org/packages/8a/13/6476736484b94041110c8340a3eb63962fea4975baea8cb4a512adb44d4d/cryptography-48.0.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d5d30989c6917b478b5817902e85fddaea2261efa8648383d965381ccb9e1ac4", size = 4689201, upload-time = "2026-06-09T22:31:09.745Z" }, - { url = "https://files.pythonhosted.org/packages/79/62/65a87f34d2a431546e2509b85d55e8c90df86d668f6731da64d538512ac2/cryptography-48.0.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:df637c05205ea7c1d7fbcbe54bbfea648a52951155f997af13d895d0ecc96991", size = 4702822, upload-time = "2026-06-09T22:32:24.409Z" }, - { url = "https://files.pythonhosted.org/packages/7f/59/810b5204b0a9b10f4b6bc06bd551a8b609803cd931806bc3b71884b225e5/cryptography-48.0.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:869c3b8a53bfe27147832df48b32adadf558249d50e76cb3769d40e986b13265", size = 4694875, upload-time = "2026-06-09T22:32:08.737Z" }, - { url = "https://files.pythonhosted.org/packages/24/dc/d8ca05ffea724eec6d232ea6f18e74c269eb6bdfdcc9bfba689790d1325f/cryptography-48.0.1-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:e361afba8918070d376df76f408a4f67fec0ee9cff81a99e48fe9a233ef59e17", size = 5290385, upload-time = "2026-06-09T22:31:15.212Z" }, - { url = "https://files.pythonhosted.org/packages/03/8c/3be6cb4da181f5bb6c19cf560c2359d60644a6b5fc5b57854e528f47b296/cryptography-48.0.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:d069066deead00ac7f090be101be875a06855908f7ec004c27b8fefb4acfb411", size = 4737082, upload-time = "2026-06-09T22:32:22.66Z" }, - { url = "https://files.pythonhosted.org/packages/aa/f6/d5f60a5a1434dbfd949e227fd0065d194c7e6b6ac526b17f5c06152b8231/cryptography-48.0.1-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:09f73a725d582cef64b91281a322cd798d14a33b2b6f2b7ad9531dc336d84c02", size = 4325328, upload-time = "2026-06-09T22:32:10.777Z" }, - { url = "https://files.pythonhosted.org/packages/17/b7/ba75dd947a14b6ad907b01ae8f6b5b348cdd1b48142f0063dee9e20c1d9d/cryptography-48.0.1-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:15254441469dd6bf027039453288e2072124f8b6603563f5d759e1c9b69273fa", size = 4694530, upload-time = "2026-06-09T22:31:53.105Z" }, - { url = "https://files.pythonhosted.org/packages/62/29/50d6b9e8aff12d8b67afaeb3569335e32dc83a5723e3bbded24fdac9f809/cryptography-48.0.1-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:8ace4507d1e6533c125f4fac754f8bb8b6a74c08e92179dabd7e16571a3efbf3", size = 5245046, upload-time = "2026-06-09T22:31:25.774Z" }, - { url = "https://files.pythonhosted.org/packages/9f/04/618f4115cfc0add0838c82507aa18a346089428da8653ad38b3ff36f5cb3/cryptography-48.0.1-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:b4e391975f038e66432328639620a4aff2d307513b004f1ca06d6225bced815c", size = 4736660, upload-time = "2026-06-09T22:32:12.676Z" }, - { url = "https://files.pythonhosted.org/packages/24/9c/06e062462a0de28a3b3911322eded4c16deb9f441b1b7575d3dc59488ab5/cryptography-48.0.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:42fcd8e26fe555d9b3577a135f5091fefa0aa4e99129c23fb56787a1bd4ada72", size = 4822229, upload-time = "2026-06-09T22:31:17.062Z" }, - { url = "https://files.pythonhosted.org/packages/f4/be/0561971eaaee4b8a0e7d5113c536921063ab91aaf23278ac374eaf881e11/cryptography-48.0.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:c1400da5e32a43253392277eac7490a60e497d810a63dd5608d71bbd7af507c9", size = 4966364, upload-time = "2026-06-09T22:31:32.842Z" }, - { url = "https://files.pythonhosted.org/packages/a4/27/728c77876f12b000820b69ae490f3c4083775e79e07827e9e60be07ad209/cryptography-48.0.1-cp314-cp314t-win32.whl", hash = "sha256:0df56b056bc17c1b7d6821dfa65216e62bd232d8ab05eb3db44e71d235651471", size = 3278498, upload-time = "2026-06-09T22:31:29.154Z" }, - { url = "https://files.pythonhosted.org/packages/06/e3/79a612c6d7b1e6ee0edd43633d53035bec2cfb78c82b76f7864f39e36f34/cryptography-48.0.1-cp314-cp314t-win_amd64.whl", hash = "sha256:9de21387aa95e2a895823d0745b430bed4f33503ba9ab5e0b5311f33e37d66d2", size = 3798790, upload-time = "2026-06-09T22:31:56.697Z" }, - { url = "https://files.pythonhosted.org/packages/ca/6c/00fa2a95997164c8b2072ce327c23d4ab20809ccc323ea5fab91e53a4bba/cryptography-48.0.1-cp39-abi3-macosx_10_9_universal2.whl", hash = "sha256:4fdc69f8e4316bcf0c8c8ec1f26f285d12e8142d88d96c876a59a03be3f6ae67", size = 7987408, upload-time = "2026-06-09T22:32:20.777Z" }, - { url = "https://files.pythonhosted.org/packages/b0/d9/45f309a7e4e5f3f8f121d6d3be9e94024a7726ec598d6e08ae04edb2f04d/cryptography-48.0.1-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:48fe40804d4caa2288f24e70ca8c64c42dd826da0ad7e4f1b41b2128d679e6c8", size = 4690196, upload-time = "2026-06-09T22:31:54.74Z" }, - { url = "https://files.pythonhosted.org/packages/5f/9f/a1bc8bcc798811b8527eb374bbccf30a3f3e806829d967118222bf1125eb/cryptography-48.0.1-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:86be3b1b0b6bf09482fb50a979c508d2950ed95f5621ec77f4e385962006b83a", size = 4696782, upload-time = "2026-06-09T22:31:45.615Z" }, - { url = "https://files.pythonhosted.org/packages/66/c2/81a4fb4e4373c500bb526bc337ac5719dd31dd15b970b84a238168c6aa08/cryptography-48.0.1-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:4ab0a343c807bbcd90c971cd1ecf072937cd01847a9e002bef88fb47ac6be577", size = 4696618, upload-time = "2026-06-09T22:31:11.564Z" }, - { url = "https://files.pythonhosted.org/packages/e5/0b/aa68b221dde92d09cb29a024ede17550ee21e77a404e59fc093c82bb51e1/cryptography-48.0.1-cp39-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:9621de99d2da096006b629979efd8ae7eb2d8b822488d0c89ee4000c306c59b1", size = 5289970, upload-time = "2026-06-09T22:31:20.368Z" }, - { url = "https://files.pythonhosted.org/packages/78/13/fba657f958d2af66ea959a4ba01212632089249d34af1ae48054136344d7/cryptography-48.0.1-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:88c852a0ae366e262e5a1744b685e6a433dc8788dd2a277e418bf4904203609d", size = 4731873, upload-time = "2026-06-09T22:31:22.253Z" }, - { url = "https://files.pythonhosted.org/packages/4c/4c/9a964756d24a26b3e34dfcb16f961b89838786e6700b635b0d1e3adff4b6/cryptography-48.0.1-cp39-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:43c5835e2cb98c8733d86f57d6fc879b613f5c3478607281c3e36daffc6dd8a6", size = 4330804, upload-time = "2026-06-09T22:31:36.56Z" }, - { url = "https://files.pythonhosted.org/packages/4b/0f/a10f3a6eb12950a10e3a874070283aa2dd5875b2bfd15fad8a3e17b3f13e/cryptography-48.0.1-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:fe0180af5bf9236518a087e35bf2d9a347d5f5f51e63c579d683ddff424e3d46", size = 4696217, upload-time = "2026-06-09T22:31:13.351Z" }, - { url = "https://files.pythonhosted.org/packages/f3/6f/5cd12f951165ea73ef85266775d97e4c763b2474ccfd816dd69d3a18d6f8/cryptography-48.0.1-cp39-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:b7a2d1a937a738a881737cec135a38bb61470589b17515b9f73f571d0ae10401", size = 5245252, upload-time = "2026-06-09T22:32:02.193Z" }, - { url = "https://files.pythonhosted.org/packages/68/ab/8aaa12e4516ec4464033ab79b6f3b592bd5a92102467c4ace8a0d970203f/cryptography-48.0.1-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:b74ca3b8e5ecdd833bf6a002ca41b4793bb27fb8f1c06ffaf2643c9e9140e31b", size = 4731388, upload-time = "2026-06-09T22:32:04.019Z" }, - { url = "https://files.pythonhosted.org/packages/1b/24/50027ea4dca85ec1f40688f3c24fb32ccacd520583c9592c3cc95628e6fb/cryptography-48.0.1-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:2c37f2461406063b417837f5f3daab668652acd82423efcd7f0a9f04be972de1", size = 4824186, upload-time = "2026-06-09T22:32:18.707Z" }, - { url = "https://files.pythonhosted.org/packages/52/41/04cb5eb17085ade6f50cc611fb657df6a0f5885350de8764ece89c050197/cryptography-48.0.1-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:86fe77abb1bd87afb251d4d02ada7ecf53a32cee9b67d976abb2e45a13297475", size = 4964539, upload-time = "2026-06-09T22:31:18.793Z" }, - { url = "https://files.pythonhosted.org/packages/36/bf/ed70785c496e89d7e73b7cda2d21f2447fd6d4e821714b8d04ff217fed92/cryptography-48.0.1-cp39-abi3-win32.whl", hash = "sha256:6b2c0c3e6ccf3ade7750f836ef3ee36eea250cc467d45c256895573ac08cc6f1", size = 3282307, upload-time = "2026-06-09T22:30:53.162Z" }, - { url = "https://files.pythonhosted.org/packages/b3/ff/371ea7d252656ee1eb6d83eeeef3d1d0c6baf1d6497687d081ea03814670/cryptography-48.0.1-cp39-abi3-win_amd64.whl", hash = "sha256:9a49ca6c81417f6a5edb50375a60cccdd70fa0a91a5211829dbea74eba94d2ac", size = 3793408, upload-time = "2026-06-09T22:32:15.191Z" }, - { url = "https://files.pythonhosted.org/packages/a9/d3/eb4e394e587341fdad09a09101fa76478ead3a78b0ad63e55c22f0d75c02/cryptography-48.0.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:08a597acce1ff37f347400087776599e2348a3a8bc53b44120e463cd274efe4a", size = 3951747, upload-time = "2026-06-09T22:31:23.871Z" }, - { url = "https://files.pythonhosted.org/packages/e0/4a/3f43451b4f858bfceaaaffc649e6e787e8d4fb332a1d443af39ab02cc8f1/cryptography-48.0.1-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:735824ec41b7f74a7c45fb1591349333e4c696cb6c044e5f46356e560143e4cd", size = 4641226, upload-time = "2026-06-09T22:31:02.532Z" }, - { url = "https://files.pythonhosted.org/packages/73/4e/855584c2c23b09e4ce2d3b9c30e983e679cd60b068c513c6bbdb91e11782/cryptography-48.0.1-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:92a46e1d638daa264ba2971c0b0489c9409787943efae4d60ffda3d091ef832c", size = 4668958, upload-time = "2026-06-09T22:32:06.213Z" }, - { url = "https://files.pythonhosted.org/packages/42/3b/d35750e41d803d1e516fd6d6011f065424924da7af1748cef4cc9cb3ede1/cryptography-48.0.1-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:7e234ac052af99f2700826a5c29ea99d9c1b1f80341cde62d11c8154dc8e0bd9", size = 4640793, upload-time = "2026-06-09T22:32:26.331Z" }, - { url = "https://files.pythonhosted.org/packages/ca/aa/cdb7181fe865285e87e96825aaab239400f1de0c3bfba9bd9769b79f1a92/cryptography-48.0.1-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:33842cf0888951cef5bc7ac724ab844a42044c1727b967b7f8997289a0464f92", size = 4668505, upload-time = "2026-06-09T22:31:27.534Z" }, - { url = "https://files.pythonhosted.org/packages/5d/8c/ce3823c06c2804f194f9e64f0d67fa3f4094a39f2bb1a990cd03603af8fc/cryptography-48.0.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:6184ca7b174f28d7c703f1290d4b297217c45355f77a98f67e9b7f14549ac54a", size = 3742204, upload-time = "2026-06-09T22:31:34.773Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/de/41/6cbdcf9142d00fe82836fbb51e503e58088575cf7a0fe1dbff6695bf0840/cryptography-50.0.0.tar.gz", hash = "sha256:eeac2acb5a20ed25e0ad6d1df9891a520b78b404266b6d11778f25d5d691a6c9", size = 880201, upload-time = "2026-07-31T14:25:10.11Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c5/5c/59086b4aac5e879d38ddbcf74e4be7ade89cebc3eb199a55da998c3bb46a/cryptography-50.0.0-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:031e2d5dd4bb9caa3ca9c82e5a197fd8ae680232cee62603d1a813f3f07e3d03", size = 4001252, upload-time = "2026-07-31T14:23:33.331Z" }, + { url = "https://files.pythonhosted.org/packages/57/ef/8f2df13c7216bcad3e1c74e07f6e193d93e998e114f524a53877c9af27ad/cryptography-50.0.0-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fd9192b7b70c573d7f214eb1ae35e00d359f6f5e4b27c7e21e30de1fc6204645", size = 4719554, upload-time = "2026-07-31T14:23:35.611Z" }, + { url = "https://files.pythonhosted.org/packages/d9/41/029086c34d91052fc3b88bcc8056f709a7c915c7a23b235a54eb800b1c97/cryptography-50.0.0-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:06a32a980526a6ab9a4b9bf8f7385800791e2bb960903cb6b530e4817509a3b7", size = 4702130, upload-time = "2026-07-31T14:23:37.635Z" }, + { url = "https://files.pythonhosted.org/packages/7d/ff/b6ce0954962e7f7b969f850a883744197bb3910bdfd7b6da162eab7d9f68/cryptography-50.0.0-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:a1b30560f2acc95aa8b2e06e716a13dbfc97314747b80d9707e307f77b40d6b3", size = 4725244, upload-time = "2026-07-31T14:23:39.471Z" }, + { url = "https://files.pythonhosted.org/packages/06/1e/63a1027cb7fec360a182208e1b7767d5aa1fe57be3d6aa856e69a321edc0/cryptography-50.0.0-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:8d89f3976b10b4ce31118de72329025f70d2c6ead14a8217c5514dd2c6d5a78f", size = 5342265, upload-time = "2026-07-31T14:23:41.286Z" }, + { url = "https://files.pythonhosted.org/packages/6b/72/a1116d683a6d7ece94590013882515de087edf9ef0e6292aae615a44df73/cryptography-50.0.0-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b42a28c1844fd9de8f3f7d540e36b66f3a9c83fceac7170ebc7a6a19edd9dcae", size = 4734609, upload-time = "2026-07-31T14:23:43.139Z" }, + { url = "https://files.pythonhosted.org/packages/15/37/36a9c479bbe49acea2636c7fd3360d20f7b7e079c300352011c44850b181/cryptography-50.0.0-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:900131fafd8aead39ac7dd3a7e833be754c17a95cfd91221636949fe4eb0aa8a", size = 4356517, upload-time = "2026-07-31T14:23:44.939Z" }, + { url = "https://files.pythonhosted.org/packages/32/98/8a151d64367204cbc63ec65d37502f1d9c53cf4bfc6ec3c532614dbec60d/cryptography-50.0.0-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:07949c449a1abcf60d1ee6e88956d89404c7df3c8258f46589e912988e551987", size = 4724529, upload-time = "2026-07-31T14:23:46.93Z" }, + { url = "https://files.pythonhosted.org/packages/22/f6/ec13b470172126464a86bf54d2294a46d29837fc51ba3e45d4047946fb5e/cryptography-50.0.0-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:f89831ef99dd7dd169ab06d63a831adb9e20a87aac6d380266bbda5823349169", size = 5299852, upload-time = "2026-07-31T14:23:48.851Z" }, + { url = "https://files.pythonhosted.org/packages/da/3a/f05e32c99d440c9bb891ea0e36c9091891e36be5a9a87ab2ee6ea20729f6/cryptography-50.0.0-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:82148ec5bddac30b51a5b3c1945075f896fa022cb93f8e4a01e9f6ee95292c5f", size = 4734462, upload-time = "2026-07-31T14:23:50.861Z" }, + { url = "https://files.pythonhosted.org/packages/ca/dc/bd72b26be8953f80625f63151efd38eee71c76ca6cf591c08ff34615a79e/cryptography-50.0.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1489e263a8048bb8b6a8bac662eb2d402ea5d2b7b4699b72f385f1e2772db105", size = 4852708, upload-time = "2026-07-31T14:23:52.715Z" }, + { url = "https://files.pythonhosted.org/packages/27/20/c930314a2ab476d15dec966ec87e2e9637bb02b06106b12c0396c57bb603/cryptography-50.0.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7cec5b856506da6defb290f30c9ee687d5f5e8cb0bd3f6459dde43b0b4fa40ef", size = 5004179, upload-time = "2026-07-31T14:23:54.887Z" }, + { url = "https://files.pythonhosted.org/packages/32/2e/c9db68a0c4bfa28e310707527c0ee3a2bd254104d2e02e68f368e197aa4c/cryptography-50.0.0-cp311-abi3-win_amd64.whl", hash = "sha256:bd1c592e4d5974f0d08d4888e432157adba757c66da0246918e43677fafa2d30", size = 3840395, upload-time = "2026-07-31T14:23:56.677Z" }, + { url = "https://files.pythonhosted.org/packages/c3/fb/951032a3bf22a5697c83183fb6294a4843772947a70e616c57b3ff5f522e/cryptography-50.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:49e7d93abdbd2990caced757e5fade25302f719c3c8fb6e6fff2dde98999fc41", size = 3989258, upload-time = "2026-07-31T14:23:58.881Z" }, + { url = "https://files.pythonhosted.org/packages/d4/67/91eb047e69c5e845f2f14b8a2e4a1aab0f283cb885531e9e22c8adb176bc/cryptography-50.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:19736989797678c6af1e55cd49055cdbcb55d8f6b5583ac5335f933aba9101dc", size = 4700648, upload-time = "2026-07-31T14:24:00.702Z" }, + { url = "https://files.pythonhosted.org/packages/30/82/85f0f7425c856b9f96459411eb12e74ef72df9caf6f8f15bf23a33ff131f/cryptography-50.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:80b63928fa35083b33966ce1efb70e5b9607181e49dcd1c22c8c005e319f667f", size = 4682442, upload-time = "2026-07-31T14:24:02.538Z" }, + { url = "https://files.pythonhosted.org/packages/1a/28/b555a365adff1cca2fbe7b9e487d68a40de6bc67ff2cb587473eb43de0e7/cryptography-50.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:d58c3db7cd6eed54e6c06744db55456b65ebd7492ddeae9c1e93cfca7aa857d3", size = 4707596, upload-time = "2026-07-31T14:24:04.394Z" }, + { url = "https://files.pythonhosted.org/packages/72/d8/f52538140cc719df62a01cf87d1c7142318d235817109d6f4054d7c352d6/cryptography-50.0.0-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:df2a58a472f332225671c35b0a830208b86d004f82baa8530fa3782c85646533", size = 5314552, upload-time = "2026-07-31T14:24:06.31Z" }, + { url = "https://files.pythonhosted.org/packages/38/14/6120e5bd7c5aa022ad15424ba4d5c5269d0d9448ed4d55e492ea91e3c1c4/cryptography-50.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:11b74db56cdbe3cdee6e3f6982ecb70334fa10dce99ed58bf7894aaaa3b2a037", size = 4717113, upload-time = "2026-07-31T14:24:08.349Z" }, + { url = "https://files.pythonhosted.org/packages/fa/71/190bf38c3ee2e0f8efc9860ae100c9df4169742eef274b91e7aa1cb133b9/cryptography-50.0.0-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:f59e38625469987d7ef6d495323c55e7db6c212eaf6112267e0d3b565a2e9c9f", size = 4338580, upload-time = "2026-07-31T14:24:10.227Z" }, + { url = "https://files.pythonhosted.org/packages/3a/63/504ccfbbe61fd8aa983f7f146399cdf034c72c2fc55f5b2dfdcdcdb20c99/cryptography-50.0.0-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:ecfed7367f965a0328cfbdd70da860f15441f002f613185668c6e6ebf5a0ac11", size = 4707038, upload-time = "2026-07-31T14:24:12.169Z" }, + { url = "https://files.pythonhosted.org/packages/01/77/2cf79bbfc4d12ca106437a6e170d6aaa01a373e93093118aaaef0e801bd4/cryptography-50.0.0-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:9aa87839c383bdbab6ef865787a1fb877af8dd03464c4400322726feaaadfc6d", size = 5273110, upload-time = "2026-07-31T14:24:14.38Z" }, + { url = "https://files.pythonhosted.org/packages/e5/45/8aae2972c520145377ea3559a605a899bebe227bf070b33cdb445929a9b9/cryptography-50.0.0-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:6ba6a53445bd3cfa809ef3ef5f1589aa6ba08784a1d962bf47d0940e871dab1c", size = 4716439, upload-time = "2026-07-31T14:24:16.415Z" }, + { url = "https://files.pythonhosted.org/packages/7b/20/4fe50b619a48c2525cc46e2dbc1ac490708d704be5d467bdaac6dc955682/cryptography-50.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:3f5735ffe4996d28b809371756219f5354864902a3b9e7c0b9ee87041209fc9c", size = 4837383, upload-time = "2026-07-31T14:24:18.553Z" }, + { url = "https://files.pythonhosted.org/packages/92/91/3a31366e183343d3703f8995c095f5734676bd6938118047e50fcf279eb4/cryptography-50.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:1b4a266766514614f8aa60416e71f2fc6e575d36e7bdc90f644fadb2f4b75b95", size = 4985772, upload-time = "2026-07-31T14:24:20.385Z" }, + { url = "https://files.pythonhosted.org/packages/74/9a/02ffe35b2853d121689871eb5dce862092562b3a1ed5cc98f1aaed441506/cryptography-50.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:12b9c6996425c76ea6c457ace4f3073e715b8c545add07cd1a8f3a4f90691269", size = 3816291, upload-time = "2026-07-31T14:24:22.125Z" }, + { url = "https://files.pythonhosted.org/packages/03/37/73d005be173aff344af30e9fd2a576575cb2391a7101d9cd3842e1fa8cce/cryptography-50.0.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:ccdc4a71a4dabae05de219404f9f4abc38e3b58422177ff93d0da05967dafa07", size = 4036009, upload-time = "2026-07-31T14:24:24.122Z" }, + { url = "https://files.pythonhosted.org/packages/ff/c6/7a6202a534e32103a285b7834a120869557fe198d51d7cfe59754c8bda9c/cryptography-50.0.0-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:910e1d2668e7de9648f2bcee30e180db2a6b15c30f887d7c4c93ddf96e3992e3", size = 4745252, upload-time = "2026-07-31T14:24:26.118Z" }, + { url = "https://files.pythonhosted.org/packages/85/4f/0fa8c2f4428198f15d9ff8d63400e27afbf94ce833f6108da1eb3753f945/cryptography-50.0.0-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a91296cb61e8df6f86d0c19cc4068228da256bf59bf86049fbd821084565327f", size = 4728939, upload-time = "2026-07-31T14:24:27.994Z" }, + { url = "https://files.pythonhosted.org/packages/d1/63/54dd723490ba2dc09b299682c10b38db38f159728bcaae8c591b8af2f22d/cryptography-50.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:e722f16708d854fe924790e051061f6704a472c3bac347b6fd88033ea8dd0dc5", size = 4748483, upload-time = "2026-07-31T14:24:30.254Z" }, + { url = "https://files.pythonhosted.org/packages/1d/dd/7c77d26285cc7f6991efce64a0f5b4f9383bfa5dd8c5033003eaf7db4cdb/cryptography-50.0.0-cp39-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:d764dcf130c428ef66786f866dd750f53182bc608813489915e9fc106bb0c82f", size = 5367599, upload-time = "2026-07-31T14:24:32.457Z" }, + { url = "https://files.pythonhosted.org/packages/46/c9/f60aed34c013f317f92817b6c171c2d22a78270fa41109bd4b08af26b194/cryptography-50.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:105110f43a471dbd0060b9c9516cb8a6a79233631a04cc2ba16f28323ac6e025", size = 4762647, upload-time = "2026-07-31T14:24:34.599Z" }, + { url = "https://files.pythonhosted.org/packages/be/f3/f9a0173b139372c3a48ed98154b45cc6b9de17c789d5ab552e621c293609/cryptography-50.0.0-cp39-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:828743d939e9629bc267b8e2d08d8bb67cd4319c771a33d4b18b22dd8fb7440a", size = 4385197, upload-time = "2026-07-31T14:24:36.647Z" }, + { url = "https://files.pythonhosted.org/packages/d8/36/83bb81f6e569bc38e1e4a7bc80f29b46bb9601920bc455fc8e888f5d5742/cryptography-50.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:2a8183b489dc1f7f80f135780fadc1108f14b31b8a40411c7a5b17425f65f28b", size = 4748095, upload-time = "2026-07-31T14:24:39.493Z" }, + { url = "https://files.pythonhosted.org/packages/6b/16/d3008eff98c764979865834c3d386d4fd041b5f52e7f34fc29ac1a5eb515/cryptography-50.0.0-cp39-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:6e7d61120573a7f2cd94cc095f9e81f6967c61ccdf194285aa143ecec8e0b708", size = 5325948, upload-time = "2026-07-31T14:24:41.556Z" }, + { url = "https://files.pythonhosted.org/packages/9c/f8/d97f9603efda3888187bfdb893f26c41be4735c10631d05d284ee6b047c4/cryptography-50.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:37fdb0d0111f1e2ff07139dfb79f1b49531f8e213c46f1163dd7642979b58c47", size = 4762400, upload-time = "2026-07-31T14:24:43.636Z" }, + { url = "https://files.pythonhosted.org/packages/64/a2/4615c8f7d81a00b1d6e6afe19f694e1543582349fb5f4076f6cb5dc36485/cryptography-50.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:c87f62a3d3b9888ed0fdde100ec06aa61ca9cd44bad9057d1dff9a516b5f5bb9", size = 4878208, upload-time = "2026-07-31T14:24:45.522Z" }, + { url = "https://files.pythonhosted.org/packages/d2/1a/efcfb02f91407149a0dacffffab791f7e19bf6385f63b3666dc8b5e5c9c8/cryptography-50.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:65c2c3add92b45fd0709db8594536aea39c2a67af0e27ffcf049c498501140b7", size = 5037050, upload-time = "2026-07-31T14:24:47.697Z" }, + { url = "https://files.pythonhosted.org/packages/57/30/4a22984d4f1bdfb8c054f07a92bc176b97a3134cc1d6c4b3bffb1f3688b4/cryptography-50.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:d24fead1d4d076e1bfb006dcec392074a3cd8d7b4fc8a595aa64073b2b7a96ba", size = 3874135, upload-time = "2026-07-31T14:24:50.085Z" }, + { url = "https://files.pythonhosted.org/packages/9d/3e/e54cde8c01631a5a8226ccd617eab9e57fd5cfdad90f1a9e6bb570794631/cryptography-50.0.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:5e34edd123674534acd70147f0ca331eaa2c74e6325fb2028c886aa26ba0b68c", size = 3963170, upload-time = "2026-07-31T14:24:51.968Z" }, + { url = "https://files.pythonhosted.org/packages/01/b6/0b9e125e90f3d2dcf599a218a899cda7326a3158cfa258723f0b398b08f6/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:8eb5e1172eb569ea8a872796576e6a67c276351728b6455d5beb01242b027c6a", size = 4692441, upload-time = "2026-07-31T14:24:53.743Z" }, + { url = "https://files.pythonhosted.org/packages/53/c9/a5151588710785a96d7bc4de27d4cd62f263bbbcb203cfe29df537eb6505/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:910d11e1a385c654bf738bf3e6b8e6ed5de0f5610fcae2be9e5b398d8081d20e", size = 4699810, upload-time = "2026-07-31T14:24:55.746Z" }, + { url = "https://files.pythonhosted.org/packages/c7/1a/15b92b25eb6ce3089cd49377ae990a0f3ad485a510f968aed1f19dbdcdf2/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:62598a8a57f815db4c6259a4e97d857dab56697e7de8e8ab02352ab74da1995d", size = 4691924, upload-time = "2026-07-31T14:24:58.082Z" }, + { url = "https://files.pythonhosted.org/packages/62/15/219075012ab13e8905f3cd572204f4acb4b111df787104346b9bc0cea789/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:07479a1cb08219ab719147e742e76090c9c773321959bb94946fffdd397a6437", size = 4699593, upload-time = "2026-07-31T14:24:59.951Z" }, + { url = "https://files.pythonhosted.org/packages/8e/b5/c2c5fce26f0ee40d21bafe7f191d29a34b35a65ac4fe8a1191d1983612e9/cryptography-50.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c99c003e088647b8a5b7c145d6f78c335f6348332b62e142d411c4b63d1460b9", size = 3813796, upload-time = "2026-07-31T14:25:02.298Z" }, ] [[package]] diff --git a/taplo.toml b/taplo.toml new file mode 100644 index 0000000000..47b33161c3 --- /dev/null +++ b/taplo.toml @@ -0,0 +1,63 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +## https://taplo.tamasfe.dev/configuration/file.html + +include = ["**/Cargo.toml"] +exclude = ["target/*"] + +[formatting] +# Align consecutive entries vertically. +align_entries = false +# Append trailing commas for multi-line arrays. +array_trailing_comma = true +# Expand arrays to multiple lines that exceed the maximum column width. +array_auto_expand = true +# Collapse arrays that don't exceed the maximum column width and don't contain comments. +array_auto_collapse = false +# Omit white space padding from single-line arrays +compact_arrays = true +# Omit white space padding from the start and end of inline tables. +compact_inline_tables = false +# Maximum column width in characters, affects array expansion and collapse, this doesn't take whitespace into account. +# Note that this is not set in stone, and works on a best-effort basis. +column_width = 120 +# Indent based on tables and arrays of tables and their subtables, subtables out of order are not indented. +indent_tables = false +# The substring that is used for indentation, should be tabs or spaces (but technically can be anything). +indent_string = ' ' +# Add trailing newline at the end of the file if not present. +trailing_newline = true +# Alphabetically reorder keys that are not separated by empty lines. +reorder_keys = false +# Maximum amount of allowed consecutive blank lines. This does not affect the whitespace at the end of the document, as it is always stripped. +allowed_blank_lines = 1 +# Use CRLF for line endings. +crlf = false + +[[rule]] +keys = [ + "build-dependencies", + "dependencies", + "dev-dependencies", + "workspace.dependencies", +] +formatting = { reorder_keys = true } + +[[rule]] +keys = ["package"] +formatting = { reorder_keys = false }