diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index e71d245c7530..2c740dedeeb8 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -25,3 +25,23 @@ If there are user-facing changes then we may require documentation to be updated + +# Does this PR break compatibility with Ballista? + + \ No newline at end of file diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index f0f17c0a94ab..060224968c51 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -66,9 +66,8 @@ jobs: # # ignore subproject CHANGELOG.md because they are machine generated npx prettier@2.3.2 --write \ - '{ballista,datafusion,datafusion-examples,docs,python}/**/*.md' \ - '!{ballista,datafusion,python}/CHANGELOG.md' \ + '{datafusion,datafusion-examples,docs,python}/**/*.md' \ + '!{datafusion,python}/CHANGELOG.md' \ README.md \ - CONTRIBUTING.md \ - 'ballista/**/*.{ts,tsx}' + CONTRIBUTING.md git diff --exit-code diff --git a/.github/workflows/dev_pr/labeler.yml b/.github/workflows/dev_pr/labeler.yml index 5b956a18a43e..9fdafc640da1 100644 --- a/.github/workflows/dev_pr/labeler.yml +++ b/.github/workflows/dev_pr/labeler.yml @@ -20,10 +20,6 @@ datafusion: - datafusion-cli/**/* - datafusion-examples/**/* -ballista: - - ballista/**/* - - ballista-examples/**/* - python: - python/**/* @@ -41,5 +37,4 @@ documentation: - README.md - ./**/README.md - DEVELOPERS.md - - ballista/docs/**.* - datafusion/docs/**.* diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index ce1bc32f8feb..d2142bd7b7c8 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -121,8 +121,6 @@ jobs: cargo run --example csv_sql cargo run --example parquet_sql cargo run --example avro_sql --features=datafusion/avro - cd ../ballista-examples - cargo run --example test_sql --features=ballista/standalone env: CARGO_HOME: "/github/home/.cargo" CARGO_TARGET_DIR: "/github/home/target" @@ -146,6 +144,9 @@ jobs: - uses: actions/checkout@v2 with: submodules: true + - uses: actions/setup-python@v3 + with: + python-version: '3.x' - name: Cache Cargo uses: actions/cache@v2 with: @@ -162,16 +163,9 @@ jobs: uses: ./.github/actions/setup-builder with: rust-version: ${{ matrix.rust }} - # Ballista is currently not part of the main workspace so requires a separate test step - - name: Run Ballista tests + - name: Run tests run: | - export ARROW_TEST_DATA=$(pwd)/testing/data - export PARQUET_TEST_DATA=$(pwd)/parquet-testing/data - cd ballista/rust - # snmalloc requires cmake so build without default features - cargo test --no-default-features --features sled - # Ensure also compiles in standalone mode - cargo test --no-default-features --features standalone + ./dev/build-arrow-ballista.sh env: CARGO_HOME: "/github/home/.cargo" CARGO_TARGET_DIR: "/github/home/target" @@ -237,8 +231,6 @@ jobs: POSTGRES_PASSWORD: postgres - name: Build datafusion-cli run: (cd datafusion-cli && cargo build) - - name: Build ballista-cli - run: (cd ballista-cli && cargo build) - name: Test Psql Parity run: python -m pytest -v integration-tests/test_psql_parity.py env: diff --git a/.gitignore b/.gitignore index 1cfeb1f0b623..529c590f73d1 100644 --- a/.gitignore +++ b/.gitignore @@ -96,3 +96,6 @@ venv/* # apache release artifacts dev/dist + +# CI +arrow-ballista diff --git a/Cargo.toml b/Cargo.toml index ef7dd30c2704..813d9e94f5b3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,13 +27,8 @@ members = [ "datafusion/row", "datafusion-examples", "benchmarks", - "ballista/rust/client", - "ballista/rust/core", - "ballista/rust/executor", - "ballista/rust/scheduler", - "ballista-examples", ] -exclude = ["ballista-cli", "datafusion-cli"] +exclude = ["datafusion-cli"] [profile.release] codegen-units = 1 diff --git a/ballista-cli/Cargo.lock b/ballista-cli/Cargo.lock deleted file mode 100644 index 528148c9b43f..000000000000 --- a/ballista-cli/Cargo.lock +++ /dev/null @@ -1,2557 +0,0 @@ -# This file is automatically @generated by Cargo. -# It is not intended for manual editing. -version = 3 - -[[package]] -name = "adler" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" - -[[package]] -name = "ahash" -version = "0.7.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fcb51a0695d8f838b1ee009b3fbf66bda078cd64590202a864a8f3e8c4315c47" -dependencies = [ - "getrandom", - "once_cell", - "version_check", -] - -[[package]] -name = "aho-corasick" -version = "0.7.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e37cfd5e7657ada45f742d6e99ca5788580b5c529dc78faf11ece6dc702656f" -dependencies = [ - "memchr", -] - -[[package]] -name = "alloc-no-stdlib" -version = "2.0.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "35ef4730490ad1c4eae5c4325b2a95f521d023e5c885853ff7aca0a6a1631db3" - -[[package]] -name = "alloc-stdlib" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "697ed7edc0f1711de49ce108c541623a0af97c6c60b2f6e2b65229847ac843c2" -dependencies = [ - "alloc-no-stdlib", -] - -[[package]] -name = "anyhow" -version = "1.0.57" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08f9b8508dccb7687a1d6c4ce66b2b0ecef467c94667de27d8d7fe1f8d2a9cdc" - -[[package]] -name = "arrayref" -version = "0.3.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4c527152e37cf757a3f78aae5a06fbeefdb07ccc535c980a3208ee3060dd544" - -[[package]] -name = "arrayvec" -version = "0.7.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8da52d66c7071e2e3fa2a1e5c6d088fec47b593032b254f5e980de8ea54454d6" - -[[package]] -name = "arrow" -version = "13.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c6bee230122beb516ead31935a61f683715f987c6f003eff44ad6986624105a" -dependencies = [ - "bitflags", - "chrono", - "comfy-table", - "csv", - "flatbuffers", - "half", - "hex", - "indexmap", - "lazy_static", - "lexical-core", - "multiversion", - "num", - "rand", - "regex", - "serde", - "serde_derive", - "serde_json", -] - -[[package]] -name = "arrow-flight" -version = "13.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0a3666d2dbc637fa979d1f0bf3031d39a80e709f3b9ec88e3d573c1d666bf553" -dependencies = [ - "arrow", - "base64", - "bytes", - "futures", - "proc-macro2", - "prost", - "prost-derive", - "tokio", - "tonic", - "tonic-build", -] - -[[package]] -name = "async-stream" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dad5c83079eae9969be7fadefe640a1c566901f05ff91ab221de4b6f68d9507e" -dependencies = [ - "async-stream-impl", - "futures-core", -] - -[[package]] -name = "async-stream-impl" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10f203db73a71dfa2fb6dd22763990fa26f3d2625a6da2da900d23b87d26be27" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "async-trait" -version = "0.1.53" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed6aa3524a2dfcf9fe180c51eae2b58738348d819517ceadf95789c51fff7600" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "atty" -version = "0.2.14" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" -dependencies = [ - "hermit-abi", - "libc", - "winapi", -] - -[[package]] -name = "autocfg" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" - -[[package]] -name = "axum" -version = "0.5.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4af7447fc1214c1f3a1ace861d0216a6c8bb13965b64bbad9650f375b67689a" -dependencies = [ - "async-trait", - "axum-core", - "bitflags", - "bytes", - "futures-util", - "http", - "http-body", - "hyper", - "itoa 1.0.1", - "matchit", - "memchr", - "mime", - "percent-encoding", - "pin-project-lite", - "serde", - "sync_wrapper", - "tokio", - "tower", - "tower-http", - "tower-layer", - "tower-service", -] - -[[package]] -name = "axum-core" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da31c0ed7b4690e2c78fe4b880d21cd7db04a346ebc658b4270251b695437f17" -dependencies = [ - "async-trait", - "bytes", - "futures-util", - "http", - "http-body", - "mime", -] - -[[package]] -name = "ballista" -version = "0.6.0" -dependencies = [ - "ballista-core", - "datafusion", - "futures", - "log", - "parking_lot", - "sqlparser", - "tempfile", - "tokio", -] - -[[package]] -name = "ballista-cli" -version = "0.6.0" -dependencies = [ - "arrow", - "ballista", - "clap", - "datafusion", - "datafusion-cli", - "dirs", - "env_logger", - "mimalloc", - "rustyline", - "tokio", -] - -[[package]] -name = "ballista-core" -version = "0.6.0" -dependencies = [ - "ahash", - "arrow-flight", - "async-trait", - "chrono", - "clap", - "datafusion", - "datafusion-proto", - "futures", - "hashbrown 0.12.1", - "libloading", - "log", - "once_cell", - "parking_lot", - "parse_arg", - "prost", - "prost-types", - "rustc_version", - "serde", - "sqlparser", - "tokio", - "tonic", - "tonic-build", - "uuid", - "walkdir", -] - -[[package]] -name = "base64" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "904dfeac50f3cdaba28fc6f57fdcddb75f49ed61346676a78c4ffe55877802fd" - -[[package]] -name = "bitflags" -version = "1.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" - -[[package]] -name = "blake2" -version = "0.10.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9cf849ee05b2ee5fba5e36f97ff8ec2533916700fc0758d40d92136a42f3388" -dependencies = [ - "digest", -] - -[[package]] -name = "blake3" -version = "1.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a08e53fc5a564bb15bfe6fae56bd71522205f1f91893f9c0116edad6496c183f" -dependencies = [ - "arrayref", - "arrayvec", - "cc", - "cfg-if", - "constant_time_eq", - "digest", -] - -[[package]] -name = "block-buffer" -version = "0.10.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bf7fe51849ea569fd452f37822f606a5cabb684dc918707a0193fd4664ff324" -dependencies = [ - "generic-array", -] - -[[package]] -name = "brotli" -version = "3.3.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1a0b1dbcc8ae29329621f8d4f0d835787c1c38bb1401979b49d13b0b305ff68" -dependencies = [ - "alloc-no-stdlib", - "alloc-stdlib", - "brotli-decompressor", -] - -[[package]] -name = "brotli-decompressor" -version = "2.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59ad2d4653bf5ca36ae797b1f4bb4dbddb60ce49ca4aed8a2ce4829f60425b80" -dependencies = [ - "alloc-no-stdlib", - "alloc-stdlib", -] - -[[package]] -name = "bstr" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba3569f383e8f1598449f1a423e72e99569137b47740b1da11ef19af3d5c3223" -dependencies = [ - "lazy_static", - "memchr", - "regex-automata", - "serde", -] - -[[package]] -name = "byteorder" -version = "1.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" - -[[package]] -name = "bytes" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4872d67bab6358e59559027aa3b9157c53d9358c51423c17554809a8858e0f8" - -[[package]] -name = "cc" -version = "1.0.73" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2fff2a6927b3bb87f9595d67196a70493f627687a71d87a0d692242c33f58c11" -dependencies = [ - "jobserver", -] - -[[package]] -name = "cfg-if" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" - -[[package]] -name = "chrono" -version = "0.4.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "670ad68c9088c2a963aaa298cb369688cf3f9465ce5e2d4ca10e6e0098a1ce73" -dependencies = [ - "libc", - "num-integer", - "num-traits", - "winapi", -] - -[[package]] -name = "clap" -version = "3.1.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85a35a599b11c089a7f49105658d089b8f2cf0882993c17daf6de15285c2c35d" -dependencies = [ - "atty", - "bitflags", - "clap_derive", - "clap_lex", - "indexmap", - "lazy_static", - "strsim", - "termcolor", - "textwrap", -] - -[[package]] -name = "clap_derive" -version = "3.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3aab4734e083b809aaf5794e14e756d1c798d2c69c7f7de7a09a2f5214993c1" -dependencies = [ - "heck 0.4.0", - "proc-macro-error", - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "clap_lex" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a37c35f1112dad5e6e0b1adaff798507497a18fceeb30cceb3bae7d1427b9213" -dependencies = [ - "os_str_bytes", -] - -[[package]] -name = "clipboard-win" -version = "4.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f3e1238132dc01f081e1cbb9dace14e5ef4c3a51ee244bd982275fb514605db" -dependencies = [ - "error-code", - "str-buf", - "winapi", -] - -[[package]] -name = "cmake" -version = "0.1.48" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e8ad8cef104ac57b68b89df3208164d228503abbdce70f6880ffa3d970e7443a" -dependencies = [ - "cc", -] - -[[package]] -name = "comfy-table" -version = "5.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b103d85ca6e209388771bfb7aa6b68a7aeec4afbf6f0a0264bfbf50360e5212e" -dependencies = [ - "strum", - "strum_macros", - "unicode-width", -] - -[[package]] -name = "constant_time_eq" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "245097e9a4535ee1e3e3931fcfcd55a796a44c643e8596ff6566d68f09b87bbc" - -[[package]] -name = "cpufeatures" -version = "0.2.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59a6001667ab124aebae2a495118e11d30984c3a653e99d86d58971708cf5e4b" -dependencies = [ - "libc", -] - -[[package]] -name = "crc32fast" -version = "1.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b540bd8bc810d3885c6ea91e2018302f68baba2129ab3e88f32389ee9370880d" -dependencies = [ - "cfg-if", -] - -[[package]] -name = "crypto-common" -version = "0.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57952ca27b5e3606ff4dd79b0020231aaf9d6aa76dc05fd30137538c50bd3ce8" -dependencies = [ - "generic-array", - "typenum", -] - -[[package]] -name = "csv" -version = "1.1.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22813a6dc45b335f9bade10bf7271dc477e81113e89eb251a0bc2a8a81c536e1" -dependencies = [ - "bstr", - "csv-core", - "itoa 0.4.8", - "ryu", - "serde", -] - -[[package]] -name = "csv-core" -version = "0.1.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b2466559f260f48ad25fe6317b3c8dac77b5bdb5763ac7d9d6103530663bc90" -dependencies = [ - "memchr", -] - -[[package]] -name = "datafusion" -version = "7.0.0" -dependencies = [ - "ahash", - "arrow", - "async-trait", - "chrono", - "datafusion-common", - "datafusion-data-access", - "datafusion-expr", - "datafusion-physical-expr", - "datafusion-row", - "futures", - "hashbrown 0.12.1", - "lazy_static", - "log", - "num_cpus", - "ordered-float 3.0.0", - "parking_lot", - "parquet", - "paste", - "pin-project-lite", - "rand", - "smallvec", - "sqlparser", - "tempfile", - "tokio", - "tokio-stream", - "uuid", -] - -[[package]] -name = "datafusion-cli" -version = "7.0.0" -dependencies = [ - "arrow", - "clap", - "datafusion", - "dirs", - "env_logger", - "mimalloc", - "rustyline", - "tokio", -] - -[[package]] -name = "datafusion-common" -version = "7.0.0" -dependencies = [ - "arrow", - "ordered-float 3.0.0", - "parquet", - "sqlparser", -] - -[[package]] -name = "datafusion-data-access" -version = "7.0.0" -dependencies = [ - "async-trait", - "chrono", - "futures", - "glob", - "parking_lot", - "tempfile", - "tokio", -] - -[[package]] -name = "datafusion-expr" -version = "7.0.0" -dependencies = [ - "ahash", - "arrow", - "datafusion-common", - "sqlparser", -] - -[[package]] -name = "datafusion-physical-expr" -version = "7.0.0" -dependencies = [ - "ahash", - "arrow", - "blake2", - "blake3", - "chrono", - "datafusion-common", - "datafusion-expr", - "datafusion-row", - "hashbrown 0.12.1", - "lazy_static", - "md-5", - "ordered-float 3.0.0", - "paste", - "rand", - "regex", - "sha2", - "unicode-segmentation", -] - -[[package]] -name = "datafusion-proto" -version = "7.0.0" -dependencies = [ - "datafusion", - "prost", - "tonic-build", -] - -[[package]] -name = "datafusion-row" -version = "7.0.0" -dependencies = [ - "arrow", - "datafusion-common", - "paste", - "rand", -] - -[[package]] -name = "digest" -version = "0.10.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2fb860ca6fafa5552fb6d0e816a69c8e49f0908bf524e30a90d97c85892d506" -dependencies = [ - "block-buffer", - "crypto-common", - "subtle", -] - -[[package]] -name = "dirs" -version = "4.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca3aa72a6f96ea37bbc5aa912f6788242832f75369bdfdadcb0e38423f100059" -dependencies = [ - "dirs-sys", -] - -[[package]] -name = "dirs-next" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b98cf8ebf19c3d1b223e151f99a4f9f0690dca41414773390fc824184ac833e1" -dependencies = [ - "cfg-if", - "dirs-sys-next", -] - -[[package]] -name = "dirs-sys" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b1d1d91c932ef41c0f2663aa8b0ca0342d444d842c06914aa0a7e352d0bada6" -dependencies = [ - "libc", - "redox_users", - "winapi", -] - -[[package]] -name = "dirs-sys-next" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ebda144c4fe02d1f7ea1a7d9641b6fc6b580adcfa024ae48797ecdeb6825b4d" -dependencies = [ - "libc", - "redox_users", - "winapi", -] - -[[package]] -name = "either" -version = "1.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e78d4f1cc4ae33bbfc157ed5d5a5ef3bc29227303d595861deb238fcec4e9457" - -[[package]] -name = "endian-type" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c34f04666d835ff5d62e058c3995147c06f42fe86ff053337632bca83e42702d" - -[[package]] -name = "env_logger" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b2cf0344971ee6c64c31be0d530793fba457d322dfec2810c453d0ef228f9c3" -dependencies = [ - "atty", - "humantime", - "log", - "regex", - "termcolor", -] - -[[package]] -name = "errno" -version = "0.2.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f639046355ee4f37944e44f60642c6f3a7efa3cf6b78c78a0d989a8ce6c396a1" -dependencies = [ - "errno-dragonfly", - "libc", - "winapi", -] - -[[package]] -name = "errno-dragonfly" -version = "0.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf" -dependencies = [ - "cc", - "libc", -] - -[[package]] -name = "error-code" -version = "2.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64f18991e7bf11e7ffee451b5318b5c1a73c52d0d0ada6e5a3017c8c1ced6a21" -dependencies = [ - "libc", - "str-buf", -] - -[[package]] -name = "fastrand" -version = "1.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3fcf0cee53519c866c09b5de1f6c56ff9d647101f81c1964fa632e148896cdf" -dependencies = [ - "instant", -] - -[[package]] -name = "fd-lock" -version = "3.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46e245f4c8ec30c6415c56cb132c07e69e74f1942f6b4a4061da748b49f486ca" -dependencies = [ - "cfg-if", - "rustix", - "windows-sys 0.30.0", -] - -[[package]] -name = "fixedbitset" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "279fb028e20b3c4c320317955b77c5e0c9701f05a1d309905d6fc702cdc5053e" - -[[package]] -name = "flatbuffers" -version = "2.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "86b428b715fdbdd1c364b84573b5fdc0f84f8e423661b9f398735278bc7f2b6a" -dependencies = [ - "bitflags", - "smallvec", - "thiserror", -] - -[[package]] -name = "flate2" -version = "1.0.23" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b39522e96686d38f4bc984b9198e3a0613264abaebaff2c5c918bfa6b6da09af" -dependencies = [ - "cfg-if", - "crc32fast", - "libc", - "miniz_oxide", -] - -[[package]] -name = "fnv" -version = "1.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" - -[[package]] -name = "futures" -version = "0.3.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f73fe65f54d1e12b726f517d3e2135ca3125a437b6d998caf1962961f7172d9e" -dependencies = [ - "futures-channel", - "futures-core", - "futures-executor", - "futures-io", - "futures-sink", - "futures-task", - "futures-util", -] - -[[package]] -name = "futures-channel" -version = "0.3.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c3083ce4b914124575708913bca19bfe887522d6e2e6d0952943f5eac4a74010" -dependencies = [ - "futures-core", - "futures-sink", -] - -[[package]] -name = "futures-core" -version = "0.3.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c09fd04b7e4073ac7156a9539b57a484a8ea920f79c7c675d05d289ab6110d3" - -[[package]] -name = "futures-executor" -version = "0.3.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9420b90cfa29e327d0429f19be13e7ddb68fa1cccb09d65e5706b8c7a749b8a6" -dependencies = [ - "futures-core", - "futures-task", - "futures-util", -] - -[[package]] -name = "futures-io" -version = "0.3.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc4045962a5a5e935ee2fdedaa4e08284547402885ab326734432bed5d12966b" - -[[package]] -name = "futures-macro" -version = "0.3.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "33c1e13800337f4d4d7a316bf45a567dbcb6ffe087f16424852d97e97a91f512" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "futures-sink" -version = "0.3.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "21163e139fa306126e6eedaf49ecdb4588f939600f0b1e770f4205ee4b7fa868" - -[[package]] -name = "futures-task" -version = "0.3.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57c66a976bf5909d801bbef33416c41372779507e7a6b3a5e25e4749c58f776a" - -[[package]] -name = "futures-util" -version = "0.3.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8b7abd5d659d9b90c8cba917f6ec750a74e2dc23902ef9cd4cc8c8b22e6036a" -dependencies = [ - "futures-channel", - "futures-core", - "futures-io", - "futures-macro", - "futures-sink", - "futures-task", - "memchr", - "pin-project-lite", - "pin-utils", - "slab", -] - -[[package]] -name = "generic-array" -version = "0.14.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd48d33ec7f05fbfa152300fdad764757cbded343c1aa1cff2fbaf4134851803" -dependencies = [ - "typenum", - "version_check", -] - -[[package]] -name = "getrandom" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9be70c98951c83b8d2f8f60d7065fa6d5146873094452a1008da8c2f1e4205ad" -dependencies = [ - "cfg-if", - "libc", - "wasi 0.10.2+wasi-snapshot-preview1", -] - -[[package]] -name = "glob" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9b919933a397b79c37e33b77bb2aa3dc8eb6e165ad809e58ff75bc7db2e34574" - -[[package]] -name = "h2" -version = "0.3.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "37a82c6d637fc9515a4694bbf1cb2457b79d81ce52b3108bdeea58b07dd34a57" -dependencies = [ - "bytes", - "fnv", - "futures-core", - "futures-sink", - "futures-util", - "http", - "indexmap", - "slab", - "tokio", - "tokio-util", - "tracing", -] - -[[package]] -name = "half" -version = "1.8.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7" - -[[package]] -name = "hashbrown" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ab5ef0d4909ef3724cc8cce6ccc8572c5c817592e9285f5464f8e86f8bd3726e" - -[[package]] -name = "hashbrown" -version = "0.12.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db0d4cf898abf0081f964436dc980e96670a0f36863e4b83aaacdb65c9d7ccc3" -dependencies = [ - "ahash", -] - -[[package]] -name = "heck" -version = "0.3.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d621efb26863f0e9924c6ac577e8275e5e6b77455db64ffa6c65c904e9e132c" -dependencies = [ - "unicode-segmentation", -] - -[[package]] -name = "heck" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2540771e65fc8cb83cd6e8a237f70c319bd5c29f78ed1084ba5d50eeac86f7f9" - -[[package]] -name = "hermit-abi" -version = "0.1.19" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62b467343b94ba476dcb2500d242dadbb39557df889310ac77c5d99100aaac33" -dependencies = [ - "libc", -] - -[[package]] -name = "hex" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" - -[[package]] -name = "http" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff8670570af52249509a86f5e3e18a08c60b177071826898fde8997cf5f6bfbb" -dependencies = [ - "bytes", - "fnv", - "itoa 1.0.1", -] - -[[package]] -name = "http-body" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ff4f84919677303da5f147645dbea6b1881f368d03ac84e1dc09031ebd7b2c6" -dependencies = [ - "bytes", - "http", - "pin-project-lite", -] - -[[package]] -name = "http-range-header" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bfe8eed0a9285ef776bb792479ea3834e8b94e13d615c2f66d03dd50a435a29" - -[[package]] -name = "httparse" -version = "1.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "496ce29bb5a52785b44e0f7ca2847ae0bb839c9bd28f69acac9b99d461c0c04c" - -[[package]] -name = "httpdate" -version = "1.0.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c4a1e36c821dbe04574f602848a19f742f4fb3c98d40449f11bcad18d6b17421" - -[[package]] -name = "humantime" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4" - -[[package]] -name = "hyper" -version = "0.14.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b26ae0a80afebe130861d90abf98e3814a4f28a4c6ffeb5ab8ebb2be311e0ef2" -dependencies = [ - "bytes", - "futures-channel", - "futures-core", - "futures-util", - "h2", - "http", - "http-body", - "httparse", - "httpdate", - "itoa 1.0.1", - "pin-project-lite", - "socket2", - "tokio", - "tower-service", - "tracing", - "want", -] - -[[package]] -name = "hyper-timeout" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbb958482e8c7be4bc3cf272a766a2b0bf1a6755e7a6ae777f017a31d11b13b1" -dependencies = [ - "hyper", - "pin-project-lite", - "tokio", - "tokio-io-timeout", -] - -[[package]] -name = "indexmap" -version = "1.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f647032dfaa1f8b6dc29bd3edb7bbef4861b8b8007ebb118d6db284fd59f6ee" -dependencies = [ - "autocfg", - "hashbrown 0.11.2", -] - -[[package]] -name = "instant" -version = "0.1.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7a5bbe824c507c5da5956355e86a746d82e0e1464f65d862cc5e71da70e94b2c" -dependencies = [ - "cfg-if", -] - -[[package]] -name = "integer-encoding" -version = "1.1.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48dc51180a9b377fd75814d0cc02199c20f8e99433d6762f650d39cdbbd3b56f" - -[[package]] -name = "io-lifetimes" -version = "0.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9448015e586b611e5d322f6703812bbca2f1e709d5773ecd38ddb4e3bb649504" - -[[package]] -name = "itertools" -version = "0.10.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9a9d19fa1e79b6215ff29b9d6880b706147f16e9b1dbb1e4e5947b5b02bc5e3" -dependencies = [ - "either", -] - -[[package]] -name = "itoa" -version = "0.4.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b71991ff56294aa922b450139ee08b3bfc70982c6b2c7562771375cf73542dd4" - -[[package]] -name = "itoa" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1aab8fc367588b89dcee83ab0fd66b72b50b72fa1904d7095045ace2b0c81c35" - -[[package]] -name = "jobserver" -version = "0.1.24" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af25a77299a7f711a01975c35a6a424eb6862092cc2d6c72c4ed6cbc56dfc1fa" -dependencies = [ - "libc", -] - -[[package]] -name = "lazy_static" -version = "1.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" - -[[package]] -name = "lexical-core" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "92912c4af2e7d9075be3e5e3122c4d7263855fa6cce34fbece4dd08e5884624d" -dependencies = [ - "lexical-parse-float", - "lexical-parse-integer", - "lexical-util", - "lexical-write-float", - "lexical-write-integer", -] - -[[package]] -name = "lexical-parse-float" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f518eed87c3be6debe6d26b855c97358d8a11bf05acec137e5f53080f5ad2dd8" -dependencies = [ - "lexical-parse-integer", - "lexical-util", - "static_assertions", -] - -[[package]] -name = "lexical-parse-integer" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "afc852ec67c6538bbb2b9911116a385b24510e879a69ab516e6a151b15a79168" -dependencies = [ - "lexical-util", - "static_assertions", -] - -[[package]] -name = "lexical-util" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c72a9d52c5c4e62fa2cdc2cb6c694a39ae1382d9c2a17a466f18e272a0930eb1" -dependencies = [ - "static_assertions", -] - -[[package]] -name = "lexical-write-float" -version = "0.8.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8a89ec1d062e481210c309b672f73a0567b7855f21e7d2fae636df44d12e97f9" -dependencies = [ - "lexical-util", - "lexical-write-integer", - "static_assertions", -] - -[[package]] -name = "lexical-write-integer" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "094060bd2a7c2ff3a16d5304a6ae82727cb3cc9d1c70f813cc73f744c319337e" -dependencies = [ - "lexical-util", - "static_assertions", -] - -[[package]] -name = "libc" -version = "0.2.125" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5916d2ae698f6de9bfb891ad7a8d65c09d232dc58cc4ac433c7da3b2fd84bc2b" - -[[package]] -name = "libloading" -version = "0.7.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "efbc0f03f9a775e9f6aed295c6a1ba2253c5757a9e03d55c6caa46a681abcddd" -dependencies = [ - "cfg-if", - "winapi", -] - -[[package]] -name = "libmimalloc-sys" -version = "0.1.25" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11ca136052550448f55df7898c6dbe651c6b574fe38a0d9ea687a9f8088a2e2c" -dependencies = [ - "cc", -] - -[[package]] -name = "linux-raw-sys" -version = "0.0.46" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4d2456c373231a208ad294c33dc5bff30051eafd954cd4caae83a712b12854d" - -[[package]] -name = "lock_api" -version = "0.4.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "327fa5b6a6940e4699ec49a9beae1ea4845c6bab9314e4f84ac68742139d8c53" -dependencies = [ - "autocfg", - "scopeguard", -] - -[[package]] -name = "log" -version = "0.4.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "abb12e687cfb44aa40f41fc3978ef76448f9b6038cad6aef4259d3c095a2382e" -dependencies = [ - "cfg-if", -] - -[[package]] -name = "lz4" -version = "1.23.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4edcb94251b1c375c459e5abe9fb0168c1c826c3370172684844f8f3f8d1a885" -dependencies = [ - "libc", - "lz4-sys", -] - -[[package]] -name = "lz4-sys" -version = "1.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7be8908e2ed6f31c02db8a9fa962f03e36c53fbfde437363eae3306b85d7e17" -dependencies = [ - "cc", - "libc", -] - -[[package]] -name = "matchit" -version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73cbba799671b762df5a175adf59ce145165747bb891505c43d09aefbbf38beb" - -[[package]] -name = "md-5" -version = "0.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "658646b21e0b72f7866c7038ab086d3d5e1cd6271f060fd37defb241949d0582" -dependencies = [ - "digest", -] - -[[package]] -name = "memchr" -version = "2.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" - -[[package]] -name = "memoffset" -version = "0.6.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5aa361d4faea93603064a027415f07bd8e1d5c88c9fbf68bf56a285428fd79ce" -dependencies = [ - "autocfg", -] - -[[package]] -name = "mimalloc" -version = "0.1.29" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f64ad83c969af2e732e907564deb0d0ed393cec4af80776f77dd77a1a427698" -dependencies = [ - "libmimalloc-sys", -] - -[[package]] -name = "mime" -version = "0.3.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2a60c7ce501c71e03a9c9c0d35b861413ae925bd979cc7a4e30d060069aaac8d" - -[[package]] -name = "miniz_oxide" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2b29bd4bc3f33391105ebee3589c19197c4271e3e5a9ec9bfe8127eeff8f082" -dependencies = [ - "adler", -] - -[[package]] -name = "mio" -version = "0.8.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52da4364ffb0e4fe33a9841a98a3f3014fb964045ce4f7a45a398243c8d6b0c9" -dependencies = [ - "libc", - "log", - "miow", - "ntapi", - "wasi 0.11.0+wasi-snapshot-preview1", - "winapi", -] - -[[package]] -name = "miow" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9f1c5b025cda876f66ef43a113f91ebc9f4ccef34843000e0adf6ebbab84e21" -dependencies = [ - "winapi", -] - -[[package]] -name = "multimap" -version = "0.8.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5ce46fe64a9d73be07dcbe690a38ce1b293be448fd8ce1e6c1b8062c9f72c6a" - -[[package]] -name = "multiversion" -version = "0.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "025c962a3dd3cc5e0e520aa9c612201d127dcdf28616974961a649dca64f5373" -dependencies = [ - "multiversion-macros", -] - -[[package]] -name = "multiversion-macros" -version = "0.6.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a8a3e2bde382ebf960c1f3e79689fa5941625fe9bf694a1cb64af3e85faff3af" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "nibble_vec" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77a5d83df9f36fe23f0c3648c6bbb8b0298bb5f1939c8f2704431371f4b84d43" -dependencies = [ - "smallvec", -] - -[[package]] -name = "nix" -version = "0.23.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9f866317acbd3a240710c63f065ffb1e4fd466259045ccb504130b7f668f35c6" -dependencies = [ - "bitflags", - "cc", - "cfg-if", - "libc", - "memoffset", -] - -[[package]] -name = "ntapi" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c28774a7fd2fbb4f0babd8237ce554b73af68021b5f695a3cebd6c59bac0980f" -dependencies = [ - "winapi", -] - -[[package]] -name = "num" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43db66d1170d347f9a065114077f7dccb00c1b9478c89384490a3425279a4606" -dependencies = [ - "num-bigint", - "num-complex", - "num-integer", - "num-iter", - "num-rational", - "num-traits", -] - -[[package]] -name = "num-bigint" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f93ab6289c7b344a8a9f60f88d80aa20032336fe78da341afc91c8a2341fc75f" -dependencies = [ - "autocfg", - "num-integer", - "num-traits", -] - -[[package]] -name = "num-complex" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97fbc387afefefd5e9e39493299f3069e14a140dd34dc19b4c1c1a8fddb6a790" -dependencies = [ - "num-traits", -] - -[[package]] -name = "num-integer" -version = "0.1.45" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" -dependencies = [ - "autocfg", - "num-traits", -] - -[[package]] -name = "num-iter" -version = "0.1.43" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d03e6c028c5dc5cac6e2dec0efda81fc887605bb3d884578bb6d6bf7514e252" -dependencies = [ - "autocfg", - "num-integer", - "num-traits", -] - -[[package]] -name = "num-rational" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d41702bd167c2df5520b384281bc111a4b5efcf7fbc4c9c222c815b07e0a6a6a" -dependencies = [ - "autocfg", - "num-bigint", - "num-integer", - "num-traits", -] - -[[package]] -name = "num-traits" -version = "0.2.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "578ede34cf02f8924ab9447f50c28075b4d3e5b269972345e7e0372b38c6cdcd" -dependencies = [ - "autocfg", -] - -[[package]] -name = "num_cpus" -version = "1.13.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "19e64526ebdee182341572e50e9ad03965aa510cd94427a4549448f285e957a1" -dependencies = [ - "hermit-abi", - "libc", -] - -[[package]] -name = "once_cell" -version = "1.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87f3e037eac156d1775da914196f0f37741a274155e34a0b7e427c35d2a2ecb9" - -[[package]] -name = "ordered-float" -version = "1.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3305af35278dd29f46fcdd139e0b1fbfae2153f0e5928b39b035542dd31e37b7" -dependencies = [ - "num-traits", -] - -[[package]] -name = "ordered-float" -version = "3.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96bcbab4bfea7a59c2c0fe47211a1ac4e3e96bea6eb446d704f310bc5c732ae2" -dependencies = [ - "num-traits", -] - -[[package]] -name = "os_str_bytes" -version = "6.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8e22443d1643a904602595ba1cd8f7d896afe56d26712531c5ff73a15b2fbf64" - -[[package]] -name = "parking_lot" -version = "0.12.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "87f5ec2493a61ac0506c0f4199f99070cbe83857b0337006a30f3e6719b8ef58" -dependencies = [ - "lock_api", - "parking_lot_core", -] - -[[package]] -name = "parking_lot_core" -version = "0.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09a279cbf25cb0757810394fbc1e359949b59e348145c643a939a525692e6929" -dependencies = [ - "cfg-if", - "libc", - "redox_syscall", - "smallvec", - "windows-sys 0.36.1", -] - -[[package]] -name = "parquet" -version = "13.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c6d737baed48775e87a69aa262f1fa2f1d6bd074dedbe9cac244b9aabf2a0b4" -dependencies = [ - "arrow", - "base64", - "brotli", - "byteorder", - "chrono", - "flate2", - "lz4", - "num", - "num-bigint", - "parquet-format", - "rand", - "snap", - "thrift", - "zstd", -] - -[[package]] -name = "parquet-format" -version = "4.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f0c06cdcd5460967c485f9c40a821746f5955ad81990533c7fae95dbd9bc0b5" -dependencies = [ - "thrift", -] - -[[package]] -name = "parse_arg" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14248cc8eced350e20122a291613de29e4fa129ba2731818c4cdbb44fccd3e55" - -[[package]] -name = "paste" -version = "1.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c520e05135d6e763148b6426a837e239041653ba7becd2e538c076c738025fc" - -[[package]] -name = "percent-encoding" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4fd5641d01c8f18a23da7b6fe29298ff4b55afcccdf78973b24cf3175fee32e" - -[[package]] -name = "petgraph" -version = "0.6.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a13a2fa9d0b63e5f22328828741e523766fff0ee9e779316902290dff3f824f" -dependencies = [ - "fixedbitset", - "indexmap", -] - -[[package]] -name = "pin-project" -version = "1.0.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "58ad3879ad3baf4e44784bc6a718a8698867bb991f8ce24d1bcbe2cfb4c3a75e" -dependencies = [ - "pin-project-internal", -] - -[[package]] -name = "pin-project-internal" -version = "1.0.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "744b6f092ba29c3650faf274db506afd39944f48420f6c86b17cfe0ee1cb36bb" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "pin-project-lite" -version = "0.2.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e0a7ae3ac2f1173085d398531c705756c94a4c56843785df85a60c1a0afac116" - -[[package]] -name = "pin-utils" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" - -[[package]] -name = "ppv-lite86" -version = "0.2.16" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb9f9e6e233e5c4a35559a617bf40a4ec447db2e84c20b55a6f83167b7e57872" - -[[package]] -name = "prettyplease" -version = "0.1.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9e07e3a46d0771a8a06b5f4441527802830b43e679ba12f44960f48dd4c6803" -dependencies = [ - "proc-macro2", - "syn", -] - -[[package]] -name = "proc-macro-error" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "da25490ff9892aab3fcf7c36f08cfb902dd3e71ca0f9f9517bea02a73a5ce38c" -dependencies = [ - "proc-macro-error-attr", - "proc-macro2", - "quote", - "syn", - "version_check", -] - -[[package]] -name = "proc-macro-error-attr" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1be40180e52ecc98ad80b184934baf3d0d29f979574e439af5a55274b35f869" -dependencies = [ - "proc-macro2", - "quote", - "version_check", -] - -[[package]] -name = "proc-macro2" -version = "1.0.37" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec757218438d5fda206afc041538b2f6d889286160d649a86a24d37e1235afd1" -dependencies = [ - "unicode-xid", -] - -[[package]] -name = "prost" -version = "0.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a07b0857a71a8cb765763950499cae2413c3f9cede1133478c43600d9e146890" -dependencies = [ - "bytes", - "prost-derive", -] - -[[package]] -name = "prost-build" -version = "0.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "120fbe7988713f39d780a58cf1a7ef0d7ef66c6d87e5aa3438940c05357929f4" -dependencies = [ - "bytes", - "cfg-if", - "cmake", - "heck 0.4.0", - "itertools", - "lazy_static", - "log", - "multimap", - "petgraph", - "prost", - "prost-types", - "regex", - "tempfile", - "which", -] - -[[package]] -name = "prost-derive" -version = "0.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b670f45da57fb8542ebdbb6105a925fe571b67f9e7ed9f47a06a84e72b4e7cc" -dependencies = [ - "anyhow", - "itertools", - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "prost-types" -version = "0.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d0a014229361011dc8e69c8a1ec6c2e8d0f2af7c91e3ea3f5b2170298461e68" -dependencies = [ - "bytes", - "prost", -] - -[[package]] -name = "quote" -version = "1.0.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1feb54ed693b93a84e14094943b84b7c4eae204c512b7ccb95ab0c66d278ad1" -dependencies = [ - "proc-macro2", -] - -[[package]] -name = "radix_trie" -version = "0.2.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c069c179fcdc6a2fe24d8d18305cf085fdbd4f922c041943e203685d6a1c58fd" -dependencies = [ - "endian-type", - "nibble_vec", -] - -[[package]] -name = "rand" -version = "0.8.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" -dependencies = [ - "libc", - "rand_chacha", - "rand_core", -] - -[[package]] -name = "rand_chacha" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" -dependencies = [ - "ppv-lite86", - "rand_core", -] - -[[package]] -name = "rand_core" -version = "0.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d34f1408f55294453790c48b2f1ebbb1c5b4b7563eb1f418bcfcfdbb06ebb4e7" -dependencies = [ - "getrandom", -] - -[[package]] -name = "redox_syscall" -version = "0.2.13" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62f25bc4c7e55e0b0b7a1d43fb893f4fa1361d0abe38b9ce4f323c2adfe6ef42" -dependencies = [ - "bitflags", -] - -[[package]] -name = "redox_users" -version = "0.4.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b033d837a7cf162d7993aded9304e30a83213c648b6e389db233191f891e5c2b" -dependencies = [ - "getrandom", - "redox_syscall", - "thiserror", -] - -[[package]] -name = "regex" -version = "1.5.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a11647b6b25ff05a515cb92c365cec08801e83423a235b51e231e1808747286" -dependencies = [ - "aho-corasick", - "memchr", - "regex-syntax", -] - -[[package]] -name = "regex-automata" -version = "0.1.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" - -[[package]] -name = "regex-syntax" -version = "0.6.25" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f497285884f3fcff424ffc933e56d7cbca511def0c9831a7f9b5f6153e3cc89b" - -[[package]] -name = "remove_dir_all" -version = "0.5.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3acd125665422973a33ac9d3dd2df85edad0f4ae9b00dafb1a05e43a9f5ef8e7" -dependencies = [ - "winapi", -] - -[[package]] -name = "rustc_version" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfa0f585226d2e68097d4f95d113b15b83a82e819ab25717ec0590d9584ef366" -dependencies = [ - "semver", -] - -[[package]] -name = "rustix" -version = "0.34.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3e74b3f02f2b6eb33790923756784614f456de79d821d6b2670dc7d5fbea807" -dependencies = [ - "bitflags", - "errno", - "io-lifetimes", - "libc", - "linux-raw-sys", - "winapi", -] - -[[package]] -name = "rustversion" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2cc38e8fa666e2de3c4aba7edeb5ffc5246c1c2ed0e3d17e560aeeba736b23f" - -[[package]] -name = "rustyline" -version = "9.1.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db7826789c0e25614b03e5a54a0717a86f9ff6e6e5247f92b369472869320039" -dependencies = [ - "bitflags", - "cfg-if", - "clipboard-win", - "dirs-next", - "fd-lock", - "libc", - "log", - "memchr", - "nix", - "radix_trie", - "scopeguard", - "smallvec", - "unicode-segmentation", - "unicode-width", - "utf8parse", - "winapi", -] - -[[package]] -name = "ryu" -version = "1.0.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73b4b750c782965c211b42f022f59af1fbceabdd026623714f104152f1ec149f" - -[[package]] -name = "same-file" -version = "1.0.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" -dependencies = [ - "winapi-util", -] - -[[package]] -name = "scopeguard" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" - -[[package]] -name = "semver" -version = "1.0.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cb243bdfdb5936c8dc3c45762a19d12ab4550cdc753bc247637d4ec35a040fd" - -[[package]] -name = "serde" -version = "1.0.137" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61ea8d54c77f8315140a05f4c7237403bf38b72704d031543aa1d16abbf517d1" -dependencies = [ - "serde_derive", -] - -[[package]] -name = "serde_derive" -version = "1.0.137" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1f26faba0c3959972377d3b2d306ee9f71faee9714294e41bb777f83f88578be" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "serde_json" -version = "1.0.80" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f972498cf015f7c0746cac89ebe1d6ef10c293b94175a243a2d9442c163d9944" -dependencies = [ - "indexmap", - "itoa 1.0.1", - "ryu", - "serde", -] - -[[package]] -name = "sha2" -version = "0.10.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "55deaec60f81eefe3cce0dc50bda92d6d8e88f2a27df7c5033b42afeb1ed2676" -dependencies = [ - "cfg-if", - "cpufeatures", - "digest", -] - -[[package]] -name = "slab" -version = "0.4.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eb703cfe953bccee95685111adeedb76fabe4e97549a58d16f03ea7b9367bb32" - -[[package]] -name = "smallvec" -version = "1.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2dd574626839106c320a323308629dcb1acfc96e32a8cba364ddc61ac23ee83" - -[[package]] -name = "snap" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45456094d1983e2ee2a18fdfebce3189fa451699d0502cb8e3b49dba5ba41451" - -[[package]] -name = "socket2" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "66d72b759436ae32898a2af0a14218dbf55efde3feeb170eb623637db85ee1e0" -dependencies = [ - "libc", - "winapi", -] - -[[package]] -name = "sqlparser" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddc2739f3a9bfc68e2f7b7695589f6cb0181c88af73ceaee0c84215cd2a2ae28" -dependencies = [ - "log", -] - -[[package]] -name = "static_assertions" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" - -[[package]] -name = "str-buf" -version = "1.0.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d44a3643b4ff9caf57abcee9c2c621d6c03d9135e0d8b589bd9afb5992cb176a" - -[[package]] -name = "strsim" -version = "0.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" - -[[package]] -name = "strum" -version = "0.23.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cae14b91c7d11c9a851d3fbc80a963198998c2a64eec840477fa92d8ce9b70bb" - -[[package]] -name = "strum_macros" -version = "0.23.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5bb0dc7ee9c15cea6199cde9a127fa16a4c5819af85395457ad72d68edc85a38" -dependencies = [ - "heck 0.3.3", - "proc-macro2", - "quote", - "rustversion", - "syn", -] - -[[package]] -name = "subtle" -version = "2.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6bdef32e8150c2a081110b42772ffe7d7c9032b606bc226c8260fd97e0976601" - -[[package]] -name = "syn" -version = "1.0.92" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ff7c592601f11445996a06f8ad0c27f094a58857c2f89e97974ab9235b92c52" -dependencies = [ - "proc-macro2", - "quote", - "unicode-xid", -] - -[[package]] -name = "sync_wrapper" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "20518fe4a4c9acf048008599e464deb21beeae3d3578418951a189c235a7a9a8" - -[[package]] -name = "tempfile" -version = "3.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5cdb1ef4eaeeaddc8fbd371e5017057064af0911902ef36b39801f67cc6d79e4" -dependencies = [ - "cfg-if", - "fastrand", - "libc", - "redox_syscall", - "remove_dir_all", - "winapi", -] - -[[package]] -name = "termcolor" -version = "1.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bab24d30b911b2376f3a13cc2cd443142f0c81dda04c118693e35b3835757755" -dependencies = [ - "winapi-util", -] - -[[package]] -name = "textwrap" -version = "0.15.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1141d4d61095b28419e22cb0bbf02755f5e54e0526f97f1e3d1d160e60885fb" - -[[package]] -name = "thiserror" -version = "1.0.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bd829fe32373d27f76265620b5309d0340cb8550f523c1dda251d6298069069a" -dependencies = [ - "thiserror-impl", -] - -[[package]] -name = "thiserror-impl" -version = "1.0.31" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0396bc89e626244658bef819e22d0cc459e795a5ebe878e6ec336d1674a8d79a" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "threadpool" -version = "1.8.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d050e60b33d41c19108b32cea32164033a9013fe3b46cbd4457559bfbf77afaa" -dependencies = [ - "num_cpus", -] - -[[package]] -name = "thrift" -version = "0.13.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c6d965454947cc7266d22716ebfd07b18d84ebaf35eec558586bbb2a8cb6b5b" -dependencies = [ - "byteorder", - "integer-encoding", - "log", - "ordered-float 1.1.1", - "threadpool", -] - -[[package]] -name = "tokio" -version = "1.18.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dce653fb475565de9f6fb0614b28bca8df2c430c0cf84bcd9c843f15de5414cc" -dependencies = [ - "bytes", - "libc", - "memchr", - "mio", - "num_cpus", - "once_cell", - "parking_lot", - "pin-project-lite", - "socket2", - "tokio-macros", - "winapi", -] - -[[package]] -name = "tokio-io-timeout" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30b74022ada614a1b4834de765f9bb43877f910cc8ce4be40e89042c9223a8bf" -dependencies = [ - "pin-project-lite", - "tokio", -] - -[[package]] -name = "tokio-macros" -version = "1.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b557f72f448c511a979e2564e55d74e6c4432fc96ff4f6241bc6bded342643b7" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "tokio-stream" -version = "0.1.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50145484efff8818b5ccd256697f36863f587da82cf8b409c53adf1e840798e3" -dependencies = [ - "futures-core", - "pin-project-lite", - "tokio", -] - -[[package]] -name = "tokio-util" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0edfdeb067411dba2044da6d1cb2df793dd35add7888d73c16e3381ded401764" -dependencies = [ - "bytes", - "futures-core", - "futures-sink", - "pin-project-lite", - "tokio", - "tracing", -] - -[[package]] -name = "tonic" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "30fb54bf1e446f44d870d260d99957e7d11fb9d0a0f5bd1a662ad1411cc103f9" -dependencies = [ - "async-stream", - "async-trait", - "axum", - "base64", - "bytes", - "futures-core", - "futures-util", - "h2", - "http", - "http-body", - "hyper", - "hyper-timeout", - "percent-encoding", - "pin-project", - "prost", - "prost-derive", - "tokio", - "tokio-stream", - "tokio-util", - "tower", - "tower-layer", - "tower-service", - "tracing", - "tracing-futures", -] - -[[package]] -name = "tonic-build" -version = "0.7.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c03447cdc9eaf8feffb6412dcb27baf2db11669a6c4789f29da799aabfb99547" -dependencies = [ - "prettyplease", - "proc-macro2", - "prost-build", - "quote", - "syn", -] - -[[package]] -name = "tower" -version = "0.4.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a89fd63ad6adf737582df5db40d286574513c69a11dac5214dc3b5603d6713e" -dependencies = [ - "futures-core", - "futures-util", - "indexmap", - "pin-project", - "pin-project-lite", - "rand", - "slab", - "tokio", - "tokio-util", - "tower-layer", - "tower-service", - "tracing", -] - -[[package]] -name = "tower-http" -version = "0.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e980386f06883cf4d0578d6c9178c81f68b45d77d00f2c2c1bc034b3439c2c56" -dependencies = [ - "bitflags", - "bytes", - "futures-core", - "futures-util", - "http", - "http-body", - "http-range-header", - "pin-project-lite", - "tower", - "tower-layer", - "tower-service", -] - -[[package]] -name = "tower-layer" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "343bc9466d3fe6b0f960ef45960509f84480bf4fd96f92901afe7ff3df9d3a62" - -[[package]] -name = "tower-service" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "360dfd1d6d30e05fda32ace2c8c70e9c0a9da713275777f5a4dbb8a1893930c6" - -[[package]] -name = "tracing" -version = "0.1.34" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d0ecdcb44a79f0fe9844f0c4f33a342cbcbb5117de8001e6ba0dc2351327d09" -dependencies = [ - "cfg-if", - "log", - "pin-project-lite", - "tracing-attributes", - "tracing-core", -] - -[[package]] -name = "tracing-attributes" -version = "0.1.21" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cc6b8ad3567499f98a1db7a752b07a7c8c7c7c34c332ec00effb2b0027974b7c" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "tracing-core" -version = "0.1.26" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f54c8ca710e81886d498c2fd3331b56c93aa248d49de2222ad2742247c60072f" -dependencies = [ - "lazy_static", -] - -[[package]] -name = "tracing-futures" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" -dependencies = [ - "pin-project", - "tracing", -] - -[[package]] -name = "try-lock" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "59547bce71d9c38b83d9c0e92b6066c4253371f15005def0c30d9657f50c7642" - -[[package]] -name = "typenum" -version = "1.15.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dcf81ac59edc17cc8697ff311e8f5ef2d99fcbd9817b34cec66f90b6c3dfd987" - -[[package]] -name = "unicode-segmentation" -version = "1.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e8820f5d777f6224dc4be3632222971ac30164d4a258d595640799554ebfd99" - -[[package]] -name = "unicode-width" -version = "0.1.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3ed742d4ea2bd1176e236172c8429aaf54486e7ac098db29ffe6529e0ce50973" - -[[package]] -name = "unicode-xid" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "957e51f3646910546462e67d5f7599b9e4fb8acdd304b087a6494730f9eebf04" - -[[package]] -name = "utf8parse" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "936e4b492acfd135421d8dca4b1aa80a7bfc26e702ef3af710e0752684df5372" - -[[package]] -name = "uuid" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cfcd319456c4d6ea10087ed423473267e1a071f3bc0aa89f80d60997843c6f0" -dependencies = [ - "getrandom", -] - -[[package]] -name = "version_check" -version = "0.9.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" - -[[package]] -name = "walkdir" -version = "2.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "808cf2735cd4b6866113f648b791c6adc5714537bc222d9347bb203386ffda56" -dependencies = [ - "same-file", - "winapi", - "winapi-util", -] - -[[package]] -name = "want" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ce8a968cb1cd110d136ff8b819a556d6fb6d919363c61534f6860c7eb172ba0" -dependencies = [ - "log", - "try-lock", -] - -[[package]] -name = "wasi" -version = "0.10.2+wasi-snapshot-preview1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd6fbd9a79829dd1ad0cc20627bf1ed606756a7f77edff7b66b7064f9cb327c6" - -[[package]] -name = "wasi" -version = "0.11.0+wasi-snapshot-preview1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" - -[[package]] -name = "which" -version = "4.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c4fb54e6113b6a8772ee41c3404fb0301ac79604489467e0a9ce1f3e97c24ae" -dependencies = [ - "either", - "lazy_static", - "libc", -] - -[[package]] -name = "winapi" -version = "0.3.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" -dependencies = [ - "winapi-i686-pc-windows-gnu", - "winapi-x86_64-pc-windows-gnu", -] - -[[package]] -name = "winapi-i686-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" - -[[package]] -name = "winapi-util" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" -dependencies = [ - "winapi", -] - -[[package]] -name = "winapi-x86_64-pc-windows-gnu" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" - -[[package]] -name = "windows-sys" -version = "0.30.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "030b7ff91626e57a05ca64a07c481973cbb2db774e4852c9c7ca342408c6a99a" -dependencies = [ - "windows_aarch64_msvc 0.30.0", - "windows_i686_gnu 0.30.0", - "windows_i686_msvc 0.30.0", - "windows_x86_64_gnu 0.30.0", - "windows_x86_64_msvc 0.30.0", -] - -[[package]] -name = "windows-sys" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ea04155a16a59f9eab786fe12a4a450e75cdb175f9e0d80da1e17db09f55b8d2" -dependencies = [ - "windows_aarch64_msvc 0.36.1", - "windows_i686_gnu 0.36.1", - "windows_i686_msvc 0.36.1", - "windows_x86_64_gnu 0.36.1", - "windows_x86_64_msvc 0.36.1", -] - -[[package]] -name = "windows_aarch64_msvc" -version = "0.30.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29277a4435d642f775f63c7d1faeb927adba532886ce0287bd985bffb16b6bca" - -[[package]] -name = "windows_aarch64_msvc" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9bb8c3fd39ade2d67e9874ac4f3db21f0d710bee00fe7cab16949ec184eeaa47" - -[[package]] -name = "windows_i686_gnu" -version = "0.30.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1145e1989da93956c68d1864f32fb97c8f561a8f89a5125f6a2b7ea75524e4b8" - -[[package]] -name = "windows_i686_gnu" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "180e6ccf01daf4c426b846dfc66db1fc518f074baa793aa7d9b9aaeffad6a3b6" - -[[package]] -name = "windows_i686_msvc" -version = "0.30.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4a09e3a0d4753b73019db171c1339cd4362c8c44baf1bcea336235e955954a6" - -[[package]] -name = "windows_i686_msvc" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e2e7917148b2812d1eeafaeb22a97e4813dfa60a3f8f78ebe204bcc88f12f024" - -[[package]] -name = "windows_x86_64_gnu" -version = "0.30.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ca64fcb0220d58db4c119e050e7af03c69e6f4f415ef69ec1773d9aab422d5a" - -[[package]] -name = "windows_x86_64_gnu" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4dcd171b8776c41b97521e5da127a2d86ad280114807d0b2ab1e462bc764d9e1" - -[[package]] -name = "windows_x86_64_msvc" -version = "0.30.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08cabc9f0066848fef4bc6a1c1668e6efce38b661d2aeec75d18d8617eebb5f1" - -[[package]] -name = "windows_x86_64_msvc" -version = "0.36.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c811ca4a8c853ef420abd8592ba53ddbbac90410fab6903b3e79972a631f7680" - -[[package]] -name = "zstd" -version = "0.11.1+zstd.1.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77a16b8414fde0414e90c612eba70985577451c4c504b99885ebed24762cb81a" -dependencies = [ - "zstd-safe", -] - -[[package]] -name = "zstd-safe" -version = "5.0.1+zstd.1.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c12659121420dd6365c5c3de4901f97145b79651fb1d25814020ed2ed0585ae" -dependencies = [ - "libc", - "zstd-sys", -] - -[[package]] -name = "zstd-sys" -version = "2.0.1+zstd.1.5.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fd07cbbc53846d9145dbffdf6dd09a7a0aa52be46741825f5c97bdd4f73f12b" -dependencies = [ - "cc", - "libc", -] diff --git a/ballista-cli/Cargo.toml b/ballista-cli/Cargo.toml deleted file mode 100644 index 3638a6227a02..000000000000 --- a/ballista-cli/Cargo.toml +++ /dev/null @@ -1,41 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -[package] -name = "ballista-cli" -description = "Command Line Client for Ballista distributed query engine." -version = "0.7.0" -authors = ["Apache Arrow "] -edition = "2021" -keywords = [ "ballista", "cli", ] -license = "Apache-2.0" -homepage = "https://github.com/apache/arrow-datafusion" -repository = "https://github.com/apache/arrow-datafusion" -rust-version = "1.59" -readme = "README.md" - -[dependencies] -arrow = { version = "14.0.0" } -ballista = { path = "../ballista/rust/client", version = "0.7.0" } -clap = { version = "3", features = ["derive", "cargo"] } -datafusion = { path = "../datafusion/core", version = "8.0.0" } -datafusion-cli = { path = "../datafusion-cli", version = "8.0.0" } -dirs = "4.0.0" -env_logger = "0.9" -mimalloc = { version = "0.1", default-features = false } -rustyline = "9.0" -tokio = { version = "1.0", features = ["macros", "rt", "rt-multi-thread", "sync", "parking_lot"] } diff --git a/ballista-cli/Dockerfile b/ballista-cli/Dockerfile deleted file mode 100644 index 7a35d185b715..000000000000 --- a/ballista-cli/Dockerfile +++ /dev/null @@ -1,38 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -FROM rust:1.59 as builder - -COPY ./datafusion /usr/src/datafusion - -COPY ./ballista-cli /usr/src/ballista-cli - -COPY ./ballista /usr/src/ballista - -WORKDIR /usr/src/ballista-cli - -RUN rustup component add rustfmt - -RUN cargo build --release - -FROM debian:bullseye-slim - -COPY --from=builder /usr/src/ballista-cli/target/release/ballista-cli /usr/local/bin - -ENTRYPOINT ["ballista-cli"] - -CMD ["--data-path", "/data"] diff --git a/ballista-cli/README.md b/ballista-cli/README.md deleted file mode 100644 index 733329e0cc1d..000000000000 --- a/ballista-cli/README.md +++ /dev/null @@ -1,94 +0,0 @@ - - -# Ballista Command-line Interface - -[Ballista](df) is a distributed query execution framework, written in Rust, that uses Apache Arrow as its in-memory format. - -The Ballista CLI allows SQL queries to be executed by an in-process DataFusion context, or by a distributed -Ballista context. - -```ignore -USAGE: - ballista-cli [FLAGS] [OPTIONS] - -FLAGS: - -h, --help Prints help information - -q, --quiet Reduce printing other than the results and work quietly - -V, --version Prints version information - -OPTIONS: - -c, --batch-size The batch size of each query, or use DataFusion default - -p, --data-path Path to your data, default to current directory - -f, --file ... Execute commands from file(s), then exit - --format Output format [default: table] [possible values: csv, tsv, table, json, ndjson] - --host Ballista scheduler host - --port Ballista scheduler port -``` - -## Example - -Create a CSV file to query. - -```bash,ignore -$ echo "1,2" > data.csv -``` - -```sql,ignore -$ ballista-cli - -Ballista CLI v0.6.0 - -> CREATE EXTERNAL TABLE foo (a INT, b INT) STORED AS CSV LOCATION 'data.csv'; -0 rows in set. Query took 0.001 seconds. - -> SELECT * FROM foo; -+---+---+ -| a | b | -+---+---+ -| 1 | 2 | -+---+---+ -1 row in set. Query took 0.017 seconds. -``` - -## Ballista-Cli - -Build the `ballista-cli` without the feature of ballista. - -```bash -cd arrow-datafusion/ballista-cli -cargo build -``` - -## Ballista - -If you want to execute the SQL in ballista by `ballista-cli`, you must build/compile the `ballista-cli` with features of "ballista" first. - -```bash -cd arrow-datafusion/ballista-cli -cargo build -``` - -The Ballista CLI can connect to a Ballista scheduler for query execution. - -```bash -ballista-cli --host localhost --port 50050 -``` - -[df]: https://crates.io/crates/datafusion diff --git a/ballista-cli/src/command.rs b/ballista-cli/src/command.rs deleted file mode 100644 index 0fd43a3071e5..000000000000 --- a/ballista-cli/src/command.rs +++ /dev/null @@ -1,222 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Command within CLI - -use crate::context::Context; -use crate::functions::{display_all_functions, Function}; -use crate::print_format::PrintFormat; -use crate::print_options::PrintOptions; -use clap::ArgEnum; -use datafusion::arrow::array::{ArrayRef, StringArray}; -use datafusion::arrow::datatypes::{DataType, Field, Schema}; -use datafusion::arrow::record_batch::RecordBatch; -use datafusion::error::{DataFusionError, Result}; -use std::str::FromStr; -use std::sync::Arc; -use std::time::Instant; - -/// Command -#[derive(Debug)] -pub enum Command { - Quit, - Help, - ListTables, - DescribeTable(String), - ListFunctions, - SearchFunctions(String), - QuietMode(Option), - OutputFormat(Option), -} - -pub enum OutputFormat { - ChangeFormat(String), -} - -impl Command { - pub async fn execute( - &self, - ctx: &mut Context, - print_options: &mut PrintOptions, - ) -> Result<()> { - let now = Instant::now(); - match self { - Self::Help => print_options - .print_batches(&[all_commands_info()], now) - .map_err(|e| DataFusionError::Execution(e.to_string())), - Self::ListTables => { - let df = ctx.sql("SHOW TABLES").await?; - let batches = df.collect().await?; - print_options - .print_batches(&batches, now) - .map_err(|e| DataFusionError::Execution(e.to_string())) - } - Self::DescribeTable(name) => { - let df = ctx.sql(&format!("SHOW COLUMNS FROM {}", name)).await?; - let batches = df.collect().await?; - print_options - .print_batches(&batches, now) - .map_err(|e| DataFusionError::Execution(e.to_string())) - } - Self::QuietMode(quiet) => { - if let Some(quiet) = quiet { - print_options.quiet = *quiet; - println!( - "Quiet mode set to {}", - if print_options.quiet { "true" } else { "false" } - ); - } else { - println!( - "Quiet mode is {}", - if print_options.quiet { "true" } else { "false" } - ); - } - Ok(()) - } - Self::Quit => Err(DataFusionError::Execution( - "Unexpected quit, this should be handled outside".into(), - )), - Self::ListFunctions => display_all_functions(), - Self::SearchFunctions(function) => { - if let Ok(func) = function.parse::() { - let details = func.function_details()?; - println!("{}", details); - Ok(()) - } else { - let msg = format!("{} is not a supported function", function); - Err(DataFusionError::Execution(msg)) - } - } - Self::OutputFormat(_) => Err(DataFusionError::Execution( - "Unexpected change output format, this should be handled outside".into(), - )), - } - } - - fn get_name_and_description(&self) -> (&'static str, &'static str) { - match self { - Self::Quit => ("\\q", "quit datafusion-cli"), - Self::ListTables => ("\\d", "list tables"), - Self::DescribeTable(_) => ("\\d name", "describe table"), - Self::Help => ("\\?", "help"), - Self::ListFunctions => ("\\h", "function list"), - Self::SearchFunctions(_) => ("\\h function", "search function"), - Self::QuietMode(_) => ("\\quiet (true|false)?", "print or set quiet mode"), - Self::OutputFormat(_) => { - ("\\pset [NAME [VALUE]]", "set table output option\n(format)") - } - } - } -} - -const ALL_COMMANDS: [Command; 8] = [ - Command::ListTables, - Command::DescribeTable(String::new()), - Command::Quit, - Command::Help, - Command::ListFunctions, - Command::SearchFunctions(String::new()), - Command::QuietMode(None), - Command::OutputFormat(None), -]; - -fn all_commands_info() -> RecordBatch { - let schema = Arc::new(Schema::new(vec![ - Field::new("Command", DataType::Utf8, false), - Field::new("Description", DataType::Utf8, false), - ])); - let (names, description): (Vec<&str>, Vec<&str>) = ALL_COMMANDS - .into_iter() - .map(|c| c.get_name_and_description()) - .unzip(); - RecordBatch::try_new( - schema, - [names, description] - .into_iter() - .map(|i| Arc::new(StringArray::from(i)) as ArrayRef) - .collect::>(), - ) - .expect("This should not fail") -} - -impl FromStr for Command { - type Err = (); - - fn from_str(s: &str) -> std::result::Result { - let (c, arg) = if let Some((a, b)) = s.split_once(' ') { - (a, Some(b)) - } else { - (s, None) - }; - Ok(match (c, arg) { - ("q", None) => Self::Quit, - ("d", None) => Self::ListTables, - ("d", Some(name)) => Self::DescribeTable(name.into()), - ("?", None) => Self::Help, - ("h", None) => Self::ListFunctions, - ("h", Some(function)) => Self::SearchFunctions(function.into()), - ("quiet", Some("true" | "t" | "yes" | "y" | "on")) => { - Self::QuietMode(Some(true)) - } - ("quiet", Some("false" | "f" | "no" | "n" | "off")) => { - Self::QuietMode(Some(false)) - } - ("quiet", None) => Self::QuietMode(None), - ("pset", Some(subcommand)) => { - Self::OutputFormat(Some(subcommand.to_string())) - } - ("pset", None) => Self::OutputFormat(None), - _ => return Err(()), - }) - } -} - -impl FromStr for OutputFormat { - type Err = (); - - fn from_str(s: &str) -> std::result::Result { - let (c, arg) = if let Some((a, b)) = s.split_once(' ') { - (a, Some(b)) - } else { - (s, None) - }; - Ok(match (c, arg) { - ("format", Some(format)) => Self::ChangeFormat(format.to_string()), - _ => return Err(()), - }) - } -} - -impl OutputFormat { - pub async fn execute(&self, print_options: &mut PrintOptions) -> Result<()> { - match self { - Self::ChangeFormat(format) => { - if let Ok(format) = format.parse::() { - print_options.format = format; - println!("Output format is {:?}.", print_options.format); - Ok(()) - } else { - Err(DataFusionError::Execution(format!( - "{:?} is not a valid format type [possible values: {:?}]", - format, - PrintFormat::value_variants() - ))) - } - } - } - } -} diff --git a/ballista-cli/src/context.rs b/ballista-cli/src/context.rs deleted file mode 100644 index c96b0e76624a..000000000000 --- a/ballista-cli/src/context.rs +++ /dev/null @@ -1,90 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Context (remote or local) - -use datafusion::dataframe::DataFrame; -use datafusion::error::{DataFusionError, Result}; -use datafusion::execution::context::{SessionConfig, SessionContext}; -use std::sync::Arc; - -/// The CLI supports using a local DataFusion context or a distributed BallistaContext -pub enum Context { - /// In-process execution with DataFusion - Local(SessionContext), - /// Distributed execution with Ballista (if available) - Remote(BallistaContext), -} - -impl Context { - /// create a new remote context with given host and port - pub async fn new_remote(host: &str, port: u16) -> Result { - Ok(Context::Remote(BallistaContext::try_new(host, port).await?)) - } - - /// create a local context using the given config - pub fn new_local(config: &SessionConfig) -> Context { - Context::Local(SessionContext::with_config(config.clone())) - } - - /// execute an SQL statement against the context - pub async fn sql(&mut self, sql: &str) -> Result> { - match self { - Context::Local(datafusion) => datafusion.sql(sql).await, - Context::Remote(ballista) => ballista.sql(sql).await, - } - } -} - -// implement wrappers around the BallistaContext to support running without ballista - -#[cfg(feature = "ballista")] -pub struct BallistaContext(ballista::context::BallistaContext); -#[cfg(feature = "ballista")] -impl BallistaContext { - pub async fn try_new(host: &str, port: u16) -> Result { - use ballista::context::BallistaContext; - use ballista::prelude::BallistaConfig; - let builder = - BallistaConfig::builder().set("ballista.with_information_schema", "true"); - let config = builder - .build() - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; - let remote_ctx = BallistaContext::remote(host, port, &config) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; - Ok(Self(remote_ctx)) - } - pub async fn sql(&mut self, sql: &str) -> Result> { - self.0.sql(sql).await - } -} - -#[cfg(not(feature = "ballista"))] -pub struct BallistaContext(); -#[cfg(not(feature = "ballista"))] -impl BallistaContext { - pub async fn try_new(_host: &str, _port: u16) -> Result { - Err(DataFusionError::NotImplemented( - "Remote execution not supported. Compile with feature 'ballista' to enable" - .to_string(), - )) - } - pub async fn sql(&mut self, _sql: &str) -> Result> { - unreachable!() - } -} diff --git a/ballista-cli/src/exec.rs b/ballista-cli/src/exec.rs deleted file mode 100644 index dc3c73e5ab43..000000000000 --- a/ballista-cli/src/exec.rs +++ /dev/null @@ -1,170 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Execution functions - -use crate::{ - command::{Command, OutputFormat}, - context::Context, - helper::CliHelper, - print_options::PrintOptions, -}; -use datafusion::error::Result; -use rustyline::error::ReadlineError; -use rustyline::Editor; -use std::fs::File; -use std::io::prelude::*; -use std::io::BufReader; -use std::time::Instant; - -/// run and execute SQL statements and commands from a file, against a context with the given print options -pub async fn exec_from_lines( - ctx: &mut Context, - reader: &mut BufReader, - print_options: &PrintOptions, -) { - let mut query = "".to_owned(); - - for line in reader.lines() { - match line { - Ok(line) if line.starts_with("--") => { - continue; - } - Ok(line) => { - let line = line.trim_end(); - query.push_str(line); - if line.ends_with(';') { - match exec_and_print(ctx, print_options, query).await { - Ok(_) => {} - Err(err) => println!("{:?}", err), - } - query = "".to_owned(); - } else { - query.push('\n'); - } - } - _ => { - break; - } - } - } - - // run the left over query if the last statement doesn't contain ‘;’ - if !query.is_empty() { - match exec_and_print(ctx, print_options, query).await { - Ok(_) => {} - Err(err) => println!("{:?}", err), - } - } -} - -pub async fn exec_from_files( - files: Vec, - ctx: &mut Context, - print_options: &PrintOptions, -) { - let files = files - .into_iter() - .map(|file_path| File::open(file_path).unwrap()) - .collect::>(); - for file in files { - let mut reader = BufReader::new(file); - exec_from_lines(ctx, &mut reader, print_options).await; - } -} - -/// run and execute SQL statements and commands against a context with the given print options -pub async fn exec_from_repl(ctx: &mut Context, print_options: &mut PrintOptions) { - let mut rl = Editor::::new(); - rl.set_helper(Some(CliHelper::default())); - rl.load_history(".history").ok(); - - let mut print_options = print_options.clone(); - - loop { - match rl.readline("❯ ") { - Ok(line) if line.starts_with('\\') => { - rl.add_history_entry(line.trim_end()); - let command = line.split_whitespace().collect::>().join(" "); - if let Ok(cmd) = &command[1..].parse::() { - match cmd { - Command::Quit => break, - Command::OutputFormat(subcommand) => { - if let Some(subcommand) = subcommand { - if let Ok(command) = subcommand.parse::() { - if let Err(e) = - command.execute(&mut print_options).await - { - eprintln!("{}", e) - } - } else { - eprintln!( - "'\\{}' is not a valid command", - &line[1..] - ); - } - } else { - println!("Output format is {:?}.", print_options.format); - } - } - _ => { - if let Err(e) = cmd.execute(ctx, &mut print_options).await { - eprintln!("{}", e) - } - } - } - } else { - eprintln!("'\\{}' is not a valid command", &line[1..]); - } - } - Ok(line) => { - rl.add_history_entry(line.trim_end()); - match exec_and_print(ctx, &print_options, line).await { - Ok(_) => {} - Err(err) => eprintln!("{:?}", err), - } - } - Err(ReadlineError::Interrupted) => { - println!("^C"); - continue; - } - Err(ReadlineError::Eof) => { - println!("\\q"); - break; - } - Err(err) => { - eprintln!("Unknown error happened {:?}", err); - break; - } - } - } - - rl.save_history(".history").ok(); -} - -async fn exec_and_print( - ctx: &mut Context, - print_options: &PrintOptions, - sql: String, -) -> Result<()> { - let now = Instant::now(); - let df = ctx.sql(&sql).await?; - let results = df.collect().await?; - print_options.print_batches(&results, now)?; - - Ok(()) -} diff --git a/ballista-cli/src/lib.rs b/ballista-cli/src/lib.rs deleted file mode 100644 index a51aad013682..000000000000 --- a/ballista-cli/src/lib.rs +++ /dev/null @@ -1,25 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -#![doc = include_str!("../README.md")] -pub const BALLISTA_CLI_VERSION: &str = env!("CARGO_PKG_VERSION"); - -pub mod command; -pub mod context; -pub mod exec; - -pub use datafusion_cli::{helper, print_format, print_options, functions}; diff --git a/ballista-cli/src/main.rs b/ballista-cli/src/main.rs deleted file mode 100644 index 69868091de16..000000000000 --- a/ballista-cli/src/main.rs +++ /dev/null @@ -1,165 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use ballista_cli::{ - context::Context, exec, print_format::PrintFormat, print_options::PrintOptions, - BALLISTA_CLI_VERSION, -}; -use clap::Parser; -use datafusion::error::Result; -use datafusion::execution::context::SessionConfig; -use mimalloc::MiMalloc; -use std::env; -use std::path::Path; - -#[global_allocator] -static GLOBAL: MiMalloc = MiMalloc; - -#[derive(Debug, Parser, PartialEq)] -#[clap(author, version, about, long_about= None)] -struct Args { - #[clap( - short = 'p', - long, - help = "Path to your data, default to current directory", - validator(is_valid_data_dir) - )] - data_path: Option, - - #[clap( - short = 'c', - long, - help = "The batch size of each query, or use DataFusion default", - validator(is_valid_batch_size) - )] - batch_size: Option, - - #[clap( - short, - long, - multiple_values = true, - help = "Execute commands from file(s), then exit", - validator(is_valid_file) - )] - file: Vec, - - #[clap( - short = 'r', - long, - multiple_values = true, - help = "Run the provided files on startup instead of ~/.datafusionrc", - validator(is_valid_file), - conflicts_with = "file" - )] - rc: Option>, - - #[clap(long, arg_enum, default_value_t = PrintFormat::Table)] - format: PrintFormat, - - #[clap(long, help = "Ballista scheduler host")] - host: Option, - - #[clap(long, help = "Ballista scheduler port")] - port: Option, - - #[clap( - short, - long, - help = "Reduce printing other than the results and work quietly" - )] - quiet: bool, -} - -#[tokio::main] -pub async fn main() -> Result<()> { - env_logger::init(); - let args = Args::parse(); - - if !args.quiet { - println!("Ballista CLI v{}", BALLISTA_CLI_VERSION); - } - - if let Some(ref path) = args.data_path { - let p = Path::new(path); - env::set_current_dir(&p).unwrap(); - }; - - let mut session_config = SessionConfig::new().with_information_schema(true); - - if let Some(batch_size) = args.batch_size { - session_config = session_config.with_batch_size(batch_size); - }; - - let mut ctx: Context = match (args.host, args.port) { - (Some(ref h), Some(p)) => Context::new_remote(h, p).await?, - _ => Context::new_local(&session_config), - }; - - let mut print_options = PrintOptions { - format: args.format, - quiet: args.quiet, - }; - - let files = args.file; - let rc = match args.rc { - Some(file) => file, - None => { - let mut files = Vec::new(); - let home = dirs::home_dir(); - if let Some(p) = home { - let home_rc = p.join(".datafusionrc"); - if home_rc.exists() { - files.push(home_rc.into_os_string().into_string().unwrap()); - } - } - files - } - }; - if !files.is_empty() { - exec::exec_from_files(files, &mut ctx, &print_options).await - } else { - if !rc.is_empty() { - exec::exec_from_files(rc, &mut ctx, &print_options).await - } - exec::exec_from_repl(&mut ctx, &mut print_options).await; - } - - Ok(()) -} - -fn is_valid_file(dir: &str) -> std::result::Result<(), String> { - if Path::new(dir).is_file() { - Ok(()) - } else { - Err(format!("Invalid file '{}'", dir)) - } -} - -fn is_valid_data_dir(dir: &str) -> std::result::Result<(), String> { - if Path::new(dir).is_dir() { - Ok(()) - } else { - Err(format!("Invalid data directory '{}'", dir)) - } -} - -fn is_valid_batch_size(size: &str) -> std::result::Result<(), String> { - match size.parse::() { - Ok(size) if size > 0 => Ok(()), - _ => Err(format!("Invalid batch size '{}'", size)), - } -} diff --git a/ballista-examples/Cargo.toml b/ballista-examples/Cargo.toml deleted file mode 100644 index f119401b062e..000000000000 --- a/ballista-examples/Cargo.toml +++ /dev/null @@ -1,44 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -[package] -name = "ballista-examples" -description = "Ballista usage examples" -version = "0.6.0" -homepage = "https://github.com/apache/arrow-datafusion" -repository = "https://github.com/apache/arrow-datafusion" -authors = ["Apache Arrow "] -license = "Apache-2.0" -keywords = [ "arrow", "distributed", "query", "sql" ] -edition = "2021" -publish = false -rust-version = "1.59" - - -[[example]] -name = "test_sql" -path = "examples/test_sql.rs" -required-features = ["ballista/standalone"] - -[dependencies] -ballista = { path = "../ballista/rust/client", version = "0.7.0" } -datafusion = { path = "../datafusion/core" } -futures = "0.3" -num_cpus = "1.13.0" -prost = "0.10" -tokio = { version = "1.0", features = ["macros", "rt", "rt-multi-thread", "sync", "parking_lot"] } -tonic = "0.7" diff --git a/ballista-examples/README.md b/ballista-examples/README.md deleted file mode 100644 index 8dd83fc9ebe8..000000000000 --- a/ballista-examples/README.md +++ /dev/null @@ -1,55 +0,0 @@ - - -# Ballista Examples - -This directory contains examples for executing distributed queries with Ballista. - -For background information on the Ballista architecture, refer to -the [Ballista README](../ballista/README.md). - -## Start a standalone cluster - -From the root of the arrow-datafusion project, build release binaries. - -```bash -cargo build --release -``` - -Start a Ballista scheduler process in a new terminal session. - -```bash -RUST_LOG=info ./target/release/ballista-scheduler -``` - -Start one or more Ballista executor processes in new terminal sessions. When starting more than one -executor, a unique port number must be specified for each executor. - -```bash -RUST_LOG=info ./target/release/ballista-executor -c 4 -``` - -## Running the examples - -The examples can be run using the `cargo run --bin` syntax. - -```bash -cargo run --release --bin ballista-dataframe -``` - diff --git a/ballista-examples/examples/test_sql.rs b/ballista-examples/examples/test_sql.rs deleted file mode 100644 index 136f4f4abb88..000000000000 --- a/ballista-examples/examples/test_sql.rs +++ /dev/null @@ -1,46 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use ballista::prelude::{BallistaConfig, BallistaContext, Result}; -use datafusion::prelude::CsvReadOptions; - -/// This example show the udf plugin is work -/// -#[tokio::main] -async fn main() -> Result<()> { - let config = BallistaConfig::builder() - .set("ballista.shuffle.partitions", "1") - .build()?; - - let ctx = BallistaContext::standalone(&config, 10).await.unwrap(); - - let testdata = datafusion::test_util::arrow_test_data(); - - // register csv file with the execution context - ctx.register_csv( - "aggregate_test_100", - &format!("{}/csv/aggregate_test_100.csv", testdata), - CsvReadOptions::new(), - ) - .await?; - - // test udf - let df = ctx.sql("select count(1) from aggregate_test_100").await?; - - df.show().await?; - Ok(()) -} diff --git a/ballista-examples/src/bin/ballista-dataframe.rs b/ballista-examples/src/bin/ballista-dataframe.rs deleted file mode 100644 index a819950f4267..000000000000 --- a/ballista-examples/src/bin/ballista-dataframe.rs +++ /dev/null @@ -1,45 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use ballista::prelude::*; -use datafusion::prelude::{col, lit, ParquetReadOptions}; - -/// This example demonstrates executing a simple query against an Arrow data source (Parquet) and -/// fetching results, using the DataFrame trait -#[tokio::main] -async fn main() -> Result<()> { - let config = BallistaConfig::builder() - .set("ballista.shuffle.partitions", "4") - .build()?; - let ctx = BallistaContext::remote("localhost", 50050, &config).await?; - - let testdata = datafusion::test_util::parquet_test_data(); - - let filename = &format!("{}/alltypes_plain.parquet", testdata); - - // define the query using the DataFrame trait - let df = ctx - .read_parquet(filename, ParquetReadOptions::default()) - .await? - .select_columns(&["id", "bool_col", "timestamp_col"])? - .filter(col("id").gt(lit(1)))?; - - // print the results - df.show().await?; - - Ok(()) -} diff --git a/ballista-examples/src/bin/ballista-sql.rs b/ballista-examples/src/bin/ballista-sql.rs deleted file mode 100644 index b4a1260f6f8a..000000000000 --- a/ballista-examples/src/bin/ballista-sql.rs +++ /dev/null @@ -1,54 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use ballista::prelude::*; -use datafusion::prelude::CsvReadOptions; - -/// This example demonstrates executing a simple query against an Arrow data source (CSV) and -/// fetching results, using SQL -#[tokio::main] -async fn main() -> Result<()> { - let config = BallistaConfig::builder() - .set("ballista.shuffle.partitions", "4") - .build()?; - let ctx = BallistaContext::remote("localhost", 50050, &config).await?; - - let testdata = datafusion::test_util::arrow_test_data(); - - // register csv file with the execution context - ctx.register_csv( - "aggregate_test_100", - &format!("{}/csv/aggregate_test_100.csv", testdata), - CsvReadOptions::new(), - ) - .await?; - - // execute the query - let df = ctx - .sql( - "SELECT c1, MIN(c12), MAX(c12) \ - FROM aggregate_test_100 \ - WHERE c11 > 0.1 AND c11 < 0.9 \ - GROUP BY c1", - ) - .await?; - - // print the results - df.show().await?; - - Ok(()) -} diff --git a/ballista/CHANGELOG.md b/ballista/CHANGELOG.md deleted file mode 100644 index 07ce062a4f6e..000000000000 --- a/ballista/CHANGELOG.md +++ /dev/null @@ -1,549 +0,0 @@ - - -# Changelog - -## [ballista-0.7.0](https://github.com/apache/arrow-datafusion/tree/ballista-0.7.0) (2022-05-12) - -[Full Changelog](https://github.com/apache/arrow-datafusion/compare/7.1.0-rc1...ballista-0.7.0) - -**Breaking changes:** - -- Make `ExecutionPlan::execute` Sync [\#2434](https://github.com/apache/arrow-datafusion/pull/2434) ([tustvold](https://github.com/tustvold)) -- Add `Expr::Exists` to represent EXISTS subquery expression [\#2339](https://github.com/apache/arrow-datafusion/pull/2339) ([andygrove](https://github.com/andygrove)) -- Remove dependency from `LogicalPlan::TableScan` to `ExecutionPlan` [\#2284](https://github.com/apache/arrow-datafusion/pull/2284) ([andygrove](https://github.com/andygrove)) -- Move logical expression type-coercion code from `physical-expr` crate to `expr` crate [\#2257](https://github.com/apache/arrow-datafusion/pull/2257) ([andygrove](https://github.com/andygrove)) -- feat: 2061 create external table ddl table partition cols [\#2099](https://github.com/apache/arrow-datafusion/pull/2099) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([jychen7](https://github.com/jychen7)) -- Reorganize the project folders [\#2081](https://github.com/apache/arrow-datafusion/pull/2081) ([yahoNanJing](https://github.com/yahoNanJing)) -- Support more ScalarFunction in Ballista [\#2008](https://github.com/apache/arrow-datafusion/pull/2008) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- Merge dataframe and dataframe imp [\#1998](https://github.com/apache/arrow-datafusion/pull/1998) ([vchag](https://github.com/vchag)) -- Rename `ExecutionContext` to `SessionContext`, `ExecutionContextState` to `SessionState`, add `TaskContext` to support multi-tenancy configurations - Part 1 [\#1987](https://github.com/apache/arrow-datafusion/pull/1987) ([mingmwang](https://github.com/mingmwang)) -- Add Coalesce function [\#1969](https://github.com/apache/arrow-datafusion/pull/1969) ([msathis](https://github.com/msathis)) -- Add Create Schema functionality in SQL [\#1959](https://github.com/apache/arrow-datafusion/pull/1959) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([matthewmturner](https://github.com/matthewmturner)) -- remove sync constraint of SendableRecordBatchStream [\#1884](https://github.com/apache/arrow-datafusion/pull/1884) ([doki23](https://github.com/doki23)) - -**Implemented enhancements:** - -- Add `CREATE VIEW` [\#2279](https://github.com/apache/arrow-datafusion/pull/2279) ([matthewmturner](https://github.com/matthewmturner)) -- \[Ballista\] Support Union in ballista. [\#2098](https://github.com/apache/arrow-datafusion/pull/2098) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- Add missing aggr\_expr to PhysicalExprNode for Ballista. [\#1989](https://github.com/apache/arrow-datafusion/pull/1989) ([Ted-Jiang](https://github.com/Ted-Jiang)) - -**Fixed bugs:** - -- Ballista integration tests no longer work [\#2440](https://github.com/apache/arrow-datafusion/issues/2440) -- Ballista crates cannot be released from DafaFusion 7.0.0 source release [\#1980](https://github.com/apache/arrow-datafusion/issues/1980) -- protobuf OctetLength should be deserialized as octet\_length, not length [\#1834](https://github.com/apache/arrow-datafusion/pull/1834) ([carols10cents](https://github.com/carols10cents)) - -**Documentation updates:** - -- MINOR: Make crate READMEs consistent [\#2437](https://github.com/apache/arrow-datafusion/pull/2437) ([andygrove](https://github.com/andygrove)) -- docs: Update the Ballista dev env instructions [\#2419](https://github.com/apache/arrow-datafusion/pull/2419) ([haoxins](https://github.com/haoxins)) -- Revise document of installing ballista pinned to specified version [\#2034](https://github.com/apache/arrow-datafusion/pull/2034) ([WinkerDu](https://github.com/WinkerDu)) -- Fix typos \(Datafusion -\> DataFusion\) [\#1993](https://github.com/apache/arrow-datafusion/pull/1993) ([andygrove](https://github.com/andygrove)) - -**Performance improvements:** - -- Introduce StageManager for managing tasks stage by stage [\#1983](https://github.com/apache/arrow-datafusion/pull/1983) ([yahoNanJing](https://github.com/yahoNanJing)) - -**Closed issues:** - -- Make expected result string in unit tests more readable [\#2412](https://github.com/apache/arrow-datafusion/issues/2412) -- remove duplicated `fn aggregate()` in aggregate expression tests [\#2399](https://github.com/apache/arrow-datafusion/issues/2399) -- split `distinct_expression.rs` into `count_distinct.rs` and `array_agg_distinct.rs` [\#2385](https://github.com/apache/arrow-datafusion/issues/2385) -- move sql tests in `context.rs` to corresponding test files in `datafustion/core/tests/sql` [\#2328](https://github.com/apache/arrow-datafusion/issues/2328) -- Date32/Date64 as join keys for merge join [\#2314](https://github.com/apache/arrow-datafusion/issues/2314) -- Error precision and scale for decimal coercion in logic comparison [\#2232](https://github.com/apache/arrow-datafusion/issues/2232) -- Support Multiple row layout [\#2188](https://github.com/apache/arrow-datafusion/issues/2188) -- Discussion: Is Ballista a standalone system or framework [\#1916](https://github.com/apache/arrow-datafusion/issues/1916) - -**Merged pull requests:** - -- MINOR: Enable multi-statement benchmark queries [\#2507](https://github.com/apache/arrow-datafusion/pull/2507) ([andygrove](https://github.com/andygrove)) -- Persist session configs in scheduler [\#2501](https://github.com/apache/arrow-datafusion/pull/2501) ([thinkharderdev](https://github.com/thinkharderdev)) -- Update to `sqlparser` `0.17.0` [\#2500](https://github.com/apache/arrow-datafusion/pull/2500) ([alamb](https://github.com/alamb)) -- Limit cpu cores used when generating changelog [\#2494](https://github.com/apache/arrow-datafusion/pull/2494) ([andygrove](https://github.com/andygrove)) -- MINOR: Parameterize changelog script [\#2484](https://github.com/apache/arrow-datafusion/pull/2484) ([jychen7](https://github.com/jychen7)) -- Fix stage key extraction [\#2472](https://github.com/apache/arrow-datafusion/pull/2472) ([thinkharderdev](https://github.com/thinkharderdev)) -- Add support for list\_dir\(\) on local fs [\#2467](https://github.com/apache/arrow-datafusion/pull/2467) ([wjones127](https://github.com/wjones127)) -- minor: update versions and paths in changelog scripts [\#2429](https://github.com/apache/arrow-datafusion/pull/2429) ([andygrove](https://github.com/andygrove)) -- Fix Ballista executing during plan [\#2428](https://github.com/apache/arrow-datafusion/pull/2428) ([tustvold](https://github.com/tustvold)) -- Re-organize and rename aggregates physical plan [\#2388](https://github.com/apache/arrow-datafusion/pull/2388) ([yjshen](https://github.com/yjshen)) -- Upgrade to arrow 13 [\#2382](https://github.com/apache/arrow-datafusion/pull/2382) ([alamb](https://github.com/alamb)) -- Grouped Aggregate in row format [\#2375](https://github.com/apache/arrow-datafusion/pull/2375) ([yjshen](https://github.com/yjshen)) -- Stop optimizing queries twice [\#2369](https://github.com/apache/arrow-datafusion/pull/2369) ([andygrove](https://github.com/andygrove)) -- Bump follow-redirects from 1.13.2 to 1.14.9 in /ballista/ui/scheduler [\#2325](https://github.com/apache/arrow-datafusion/pull/2325) ([dependabot[bot]](https://github.com/apps/dependabot)) -- Move FileType enum from sql module to logical\_plan module [\#2290](https://github.com/apache/arrow-datafusion/pull/2290) ([andygrove](https://github.com/andygrove)) -- Add BatchPartitioner \(\#2285\) [\#2287](https://github.com/apache/arrow-datafusion/pull/2287) ([tustvold](https://github.com/tustvold)) -- Update uuid requirement from 0.8 to 1.0 [\#2280](https://github.com/apache/arrow-datafusion/pull/2280) ([dependabot[bot]](https://github.com/apps/dependabot)) -- Bump async from 2.6.3 to 2.6.4 in /ballista/ui/scheduler [\#2277](https://github.com/apache/arrow-datafusion/pull/2277) ([dependabot[bot]](https://github.com/apps/dependabot)) -- Bump minimist from 1.2.5 to 1.2.6 in /ballista/ui/scheduler [\#2276](https://github.com/apache/arrow-datafusion/pull/2276) ([dependabot[bot]](https://github.com/apps/dependabot)) -- Bump url-parse from 1.5.1 to 1.5.10 in /ballista/ui/scheduler [\#2275](https://github.com/apache/arrow-datafusion/pull/2275) ([dependabot[bot]](https://github.com/apps/dependabot)) -- Bump nanoid from 3.1.20 to 3.3.3 in /ballista/ui/scheduler [\#2274](https://github.com/apache/arrow-datafusion/pull/2274) ([dependabot[bot]](https://github.com/apps/dependabot)) -- Update to Arrow 12.0.0, update tonic and prost [\#2253](https://github.com/apache/arrow-datafusion/pull/2253) ([alamb](https://github.com/alamb)) -- Add ExecutorMetricsCollector interface [\#2234](https://github.com/apache/arrow-datafusion/pull/2234) ([thinkharderdev](https://github.com/thinkharderdev)) -- minor: add editor config file [\#2224](https://github.com/apache/arrow-datafusion/pull/2224) ([jackwener](https://github.com/jackwener)) -- \[Ballista\] Enable ApproxPercentileWithWeight in Ballista and fill UT [\#2192](https://github.com/apache/arrow-datafusion/pull/2192) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- make nightly clippy happy [\#2186](https://github.com/apache/arrow-datafusion/pull/2186) ([xudong963](https://github.com/xudong963)) -- \[Ballista\]Make PhysicalAggregateExprNode has repeated PhysicalExprNode [\#2184](https://github.com/apache/arrow-datafusion/pull/2184) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- Add LogicalPlan::SubqueryAlias [\#2172](https://github.com/apache/arrow-datafusion/pull/2172) ([andygrove](https://github.com/andygrove)) -- Implement fast path of with\_new\_children\(\) in ExecutionPlan [\#2168](https://github.com/apache/arrow-datafusion/pull/2168) ([mingmwang](https://github.com/mingmwang)) -- \[MINOR\] ignore suspicious slow test in Ballista [\#2167](https://github.com/apache/arrow-datafusion/pull/2167) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- enable explain for ballista [\#2163](https://github.com/apache/arrow-datafusion/pull/2163) ([doki23](https://github.com/doki23)) -- Add delimiter for create external table [\#2162](https://github.com/apache/arrow-datafusion/pull/2162) ([matthewmturner](https://github.com/matthewmturner)) -- Update sqlparser requirement from 0.15 to 0.16 [\#2152](https://github.com/apache/arrow-datafusion/pull/2152) ([dependabot[bot]](https://github.com/apps/dependabot)) -- Add IF NOT EXISTS to `CREATE TABLE` and `CREATE EXTERNAL TABLE` [\#2143](https://github.com/apache/arrow-datafusion/pull/2143) ([matthewmturner](https://github.com/matthewmturner)) -- Update quarterly roadmap for Q2 [\#2133](https://github.com/apache/arrow-datafusion/pull/2133) ([matthewmturner](https://github.com/matthewmturner)) -- \[Ballista\] Add ballista plugin manager and UDF plugin [\#2131](https://github.com/apache/arrow-datafusion/pull/2131) ([gaojun2048](https://github.com/gaojun2048)) -- Serialize scalar UDFs in physical plan [\#2130](https://github.com/apache/arrow-datafusion/pull/2130) ([thinkharderdev](https://github.com/thinkharderdev)) -- doc: update release schedule [\#2110](https://github.com/apache/arrow-datafusion/pull/2110) ([jychen7](https://github.com/jychen7)) -- Reduce repetition in Decimal binary kernels, upgrade to arrow 11.1 [\#2107](https://github.com/apache/arrow-datafusion/pull/2107) ([alamb](https://github.com/alamb)) -- update zlib version to 1.2.12 [\#2106](https://github.com/apache/arrow-datafusion/pull/2106) ([waitingkuo](https://github.com/waitingkuo)) -- Add CREATE DATABASE command to SQL [\#2094](https://github.com/apache/arrow-datafusion/pull/2094) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([matthewmturner](https://github.com/matthewmturner)) -- Refactor SessionContext, BallistaContext to support multi-tenancy configurations - Part 3 [\#2091](https://github.com/apache/arrow-datafusion/pull/2091) ([mingmwang](https://github.com/mingmwang)) -- Remove dependency of common for the storage crate [\#2076](https://github.com/apache/arrow-datafusion/pull/2076) ([yahoNanJing](https://github.com/yahoNanJing)) -- [MINOR] fix doc in `EXTRACT\(field FROM source\) [\#2074](https://github.com/apache/arrow-datafusion/pull/2074) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- \[Bug\]\[Datafusion\] fix TaskContext session\_config bug [\#2070](https://github.com/apache/arrow-datafusion/pull/2070) ([gaojun2048](https://github.com/gaojun2048)) -- Short-circuit evaluation for `CaseWhen` [\#2068](https://github.com/apache/arrow-datafusion/pull/2068) ([yjshen](https://github.com/yjshen)) -- split datafusion-object-store module [\#2065](https://github.com/apache/arrow-datafusion/pull/2065) ([yahoNanJing](https://github.com/yahoNanJing)) -- Change log level for noisy logs [\#2060](https://github.com/apache/arrow-datafusion/pull/2060) ([thinkharderdev](https://github.com/thinkharderdev)) -- Update to arrow/parquet 11.0 [\#2048](https://github.com/apache/arrow-datafusion/pull/2048) ([alamb](https://github.com/alamb)) -- minor: format comments \(`//` to `// `\) [\#2047](https://github.com/apache/arrow-datafusion/pull/2047) ([jackwener](https://github.com/jackwener)) -- use cargo-tomlfmt to check Cargo.toml formatting in CI [\#2033](https://github.com/apache/arrow-datafusion/pull/2033) ([WinkerDu](https://github.com/WinkerDu)) -- Refactor SessionContext, SessionState and SessionConfig to support multi-tenancy configurations - Part 2 [\#2029](https://github.com/apache/arrow-datafusion/pull/2029) ([mingmwang](https://github.com/mingmwang)) -- Simplify prerequisites for running examples [\#2028](https://github.com/apache/arrow-datafusion/pull/2028) ([doki23](https://github.com/doki23)) -- Use SessionContext to parse Expr protobuf [\#2024](https://github.com/apache/arrow-datafusion/pull/2024) ([thinkharderdev](https://github.com/thinkharderdev)) -- Fix stuck issue for the load testing of Push-based task scheduling [\#2006](https://github.com/apache/arrow-datafusion/pull/2006) ([yahoNanJing](https://github.com/yahoNanJing)) -- Fixing a typo in documentation [\#1997](https://github.com/apache/arrow-datafusion/pull/1997) ([psvri](https://github.com/psvri)) -- Fix minor clippy issue [\#1995](https://github.com/apache/arrow-datafusion/pull/1995) ([alamb](https://github.com/alamb)) -- Make it possible to only scan part of a parquet file in a partition [\#1990](https://github.com/apache/arrow-datafusion/pull/1990) ([yjshen](https://github.com/yjshen)) -- Update Dockerfile to fix integration tests [\#1982](https://github.com/apache/arrow-datafusion/pull/1982) ([andygrove](https://github.com/andygrove)) -- Update sqlparser requirement from 0.14 to 0.15 [\#1966](https://github.com/apache/arrow-datafusion/pull/1966) ([dependabot[bot]](https://github.com/apps/dependabot)) -- fix logical conflict with protobuf [\#1958](https://github.com/apache/arrow-datafusion/pull/1958) ([alamb](https://github.com/alamb)) -- Update to arrow 10.0.0, pyo3 0.16 [\#1957](https://github.com/apache/arrow-datafusion/pull/1957) ([alamb](https://github.com/alamb)) -- update jit-related dependencies [\#1953](https://github.com/apache/arrow-datafusion/pull/1953) ([xudong963](https://github.com/xudong963)) -- Allow different types of query variables \(`@@var`\) rather than just string [\#1943](https://github.com/apache/arrow-datafusion/pull/1943) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([maxburke](https://github.com/maxburke)) -- Pruning serialization [\#1941](https://github.com/apache/arrow-datafusion/pull/1941) ([thinkharderdev](https://github.com/thinkharderdev)) -- Fix select from EmptyExec always return 0 row after optimizer passes [\#1938](https://github.com/apache/arrow-datafusion/pull/1938) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- Introduce Ballista query stage scheduler [\#1935](https://github.com/apache/arrow-datafusion/pull/1935) ([yahoNanJing](https://github.com/yahoNanJing)) -- Add db benchmark script [\#1928](https://github.com/apache/arrow-datafusion/pull/1928) ([matthewmturner](https://github.com/matthewmturner)) -- fix a typo [\#1919](https://github.com/apache/arrow-datafusion/pull/1919) ([vchag](https://github.com/vchag)) -- \[MINOR\] Update copyright year in Docs [\#1918](https://github.com/apache/arrow-datafusion/pull/1918) ([alamb](https://github.com/alamb)) -- add metadata to DFSchema, close \#1806. [\#1914](https://github.com/apache/arrow-datafusion/pull/1914) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([jiacai2050](https://github.com/jiacai2050)) -- Refactor scheduler state mod [\#1913](https://github.com/apache/arrow-datafusion/pull/1913) ([yahoNanJing](https://github.com/yahoNanJing)) -- Refactor the event channel [\#1912](https://github.com/apache/arrow-datafusion/pull/1912) ([yahoNanJing](https://github.com/yahoNanJing)) -- Refactor scheduler server [\#1911](https://github.com/apache/arrow-datafusion/pull/1911) ([yahoNanJing](https://github.com/yahoNanJing)) -- Clippy fix on nightly [\#1907](https://github.com/apache/arrow-datafusion/pull/1907) ([yjshen](https://github.com/yjshen)) -- Updated Rust version to 1.59 in all the files [\#1903](https://github.com/apache/arrow-datafusion/pull/1903) ([NaincyKumariKnoldus](https://github.com/NaincyKumariKnoldus)) -- Remove uneeded Mutex in Ballista Client [\#1898](https://github.com/apache/arrow-datafusion/pull/1898) ([alamb](https://github.com/alamb)) -- Create a `datafusion-proto` crate for datafusion protobuf serialization [\#1887](https://github.com/apache/arrow-datafusion/pull/1887) ([carols10cents](https://github.com/carols10cents)) -- Fix clippy lints [\#1885](https://github.com/apache/arrow-datafusion/pull/1885) ([HaoYang670](https://github.com/HaoYang670)) -- Separate cpu-bound \(query-execution\) and IO-bound\(heartbeat\) to … [\#1883](https://github.com/apache/arrow-datafusion/pull/1883) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- \[Minor\] Clean up DecimalArray API Usage [\#1869](https://github.com/apache/arrow-datafusion/pull/1869) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([alamb](https://github.com/alamb)) -- Changes after went through "Datafusion as a library section" [\#1868](https://github.com/apache/arrow-datafusion/pull/1868) ([nonontb](https://github.com/nonontb)) -- Remove allow unused imports from ballista-core, then fix all warnings [\#1853](https://github.com/apache/arrow-datafusion/pull/1853) ([carols10cents](https://github.com/carols10cents)) -- Update to arrow 9.1.0 [\#1851](https://github.com/apache/arrow-datafusion/pull/1851) ([alamb](https://github.com/alamb)) -- move some tests out of context and into sql [\#1846](https://github.com/apache/arrow-datafusion/pull/1846) ([alamb](https://github.com/alamb)) -- Fix compiling ballista in standalone mode, add build to CI [\#1839](https://github.com/apache/arrow-datafusion/pull/1839) ([alamb](https://github.com/alamb)) -- Update documentation example for change in API [\#1812](https://github.com/apache/arrow-datafusion/pull/1812) ([alamb](https://github.com/alamb)) -- Refactor scheduler state with different management policy for volatile and stable states [\#1810](https://github.com/apache/arrow-datafusion/pull/1810) ([yahoNanJing](https://github.com/yahoNanJing)) -- DataFusion + Conbench Integration [\#1791](https://github.com/apache/arrow-datafusion/pull/1791) ([dianaclarke](https://github.com/dianaclarke)) -- Enable periodic cleanup of work\_dir directories in ballista executor [\#1783](https://github.com/apache/arrow-datafusion/pull/1783) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- Use`eq_dyn`, `neq_dyn`, `lt_dyn`, `lt_eq_dyn`, `gt_dyn`, `gt_eq_dyn` kernels from arrow [\#1475](https://github.com/apache/arrow-datafusion/pull/1475) ([alamb](https://github.com/alamb)) - -## [7.1.0-rc1](https://github.com/apache/arrow-datafusion/tree/7.1.0-rc1) (2022-04-10) - -[Full Changelog](https://github.com/apache/arrow-datafusion/compare/7.0.0-rc2...7.1.0-rc1) - -**Implemented enhancements:** - -- Support substring with three arguments: \(str, from, for\) for DataFrame API and Ballista [\#2092](https://github.com/apache/arrow-datafusion/issues/2092) -- UnionAll support for Ballista [\#2032](https://github.com/apache/arrow-datafusion/issues/2032) -- Separate cpu-bound and IO-bound work in ballista-executor by using diff tokio runtime. [\#1770](https://github.com/apache/arrow-datafusion/issues/1770) -- \[Ballista\] Introduce DAGScheduler for better managing the stage-based task scheduling [\#1704](https://github.com/apache/arrow-datafusion/issues/1704) -- \[Ballista\] Support to better manage cluster state, like alive executors, executor available task slots, etc [\#1703](https://github.com/apache/arrow-datafusion/issues/1703) - -**Closed issues:** - -- Optimize memory usage pattern to avoid "double memory" behavior [\#2149](https://github.com/apache/arrow-datafusion/issues/2149) -- Document approx\_percentile\_cont\_with\_weight in users guide [\#2078](https://github.com/apache/arrow-datafusion/issues/2078) -- \[follow up\]cleaning up statements.remove\(0\) [\#1986](https://github.com/apache/arrow-datafusion/issues/1986) -- Formatting error on documentation for Python [\#1873](https://github.com/apache/arrow-datafusion/issues/1873) -- Remove duplicate tests from `test_const_evaluator_scalar_functions` [\#1727](https://github.com/apache/arrow-datafusion/issues/1727) -- Question: Is the Ballista project providing value to the overall DataFusion project? [\#1273](https://github.com/apache/arrow-datafusion/issues/1273) - -## [7.0.0-rc2](https://github.com/apache/arrow-datafusion/tree/7.0.0-rc2) (2022-02-14) - -[Full Changelog](https://github.com/apache/arrow-datafusion/compare/7.0.0...7.0.0-rc2) - -## [7.0.0](https://github.com/apache/arrow-datafusion/tree/7.0.0) (2022-02-14) - -[Full Changelog](https://github.com/apache/arrow-datafusion/compare/6.0.0-rc0...7.0.0) - -**Breaking changes:** - -- Update `ExecutionPlan` to know about sortedness and repartitioning optimizer pass respect the invariants [\#1776](https://github.com/apache/arrow-datafusion/pull/1776) ([alamb](https://github.com/alamb)) -- Update to `arrow 8.0.0` [\#1673](https://github.com/apache/arrow-datafusion/pull/1673) ([alamb](https://github.com/alamb)) - -**Implemented enhancements:** - -- Task assignment between Scheduler and Executors [\#1221](https://github.com/apache/arrow-datafusion/issues/1221) -- Add `approx_median()` aggregate function [\#1729](https://github.com/apache/arrow-datafusion/pull/1729) ([realno](https://github.com/realno)) -- \[Ballista\] Add Decimal128, Date64, TimestampSecond, TimestampMillisecond, Interv… [\#1659](https://github.com/apache/arrow-datafusion/pull/1659) ([gaojun2048](https://github.com/gaojun2048)) -- Add `corr` aggregate function [\#1561](https://github.com/apache/arrow-datafusion/pull/1561) ([realno](https://github.com/realno)) -- Add `covar`, `covar_pop` and `covar_samp` aggregate functions [\#1551](https://github.com/apache/arrow-datafusion/pull/1551) ([realno](https://github.com/realno)) -- Add `approx_quantile()` aggregation function [\#1539](https://github.com/apache/arrow-datafusion/pull/1539) ([domodwyer](https://github.com/domodwyer)) -- Initial MemoryManager and DiskManager APIs for query execution + External Sort implementation [\#1526](https://github.com/apache/arrow-datafusion/pull/1526) ([yjshen](https://github.com/yjshen)) -- Add `stddev` and `variance` [\#1525](https://github.com/apache/arrow-datafusion/pull/1525) ([realno](https://github.com/realno)) -- Add `rem` operation for Expr [\#1467](https://github.com/apache/arrow-datafusion/pull/1467) ([liukun4515](https://github.com/liukun4515)) -- Implement `array_agg` aggregate function [\#1300](https://github.com/apache/arrow-datafusion/pull/1300) ([viirya](https://github.com/viirya)) - -**Fixed bugs:** - -- Ballista context::tests::test\_standalone\_mode test fails [\#1020](https://github.com/apache/arrow-datafusion/issues/1020) -- \[Ballista\] Fix scheduler state mod bug [\#1655](https://github.com/apache/arrow-datafusion/pull/1655) ([gaojun2048](https://github.com/gaojun2048)) -- Pass local address host so we do not get mismatch between IPv4 and IP… [\#1466](https://github.com/apache/arrow-datafusion/pull/1466) ([thinkharderdev](https://github.com/thinkharderdev)) -- Add Timezone to Scalar::Time\* types, and better timezone awareness to Datafusion's time types [\#1455](https://github.com/apache/arrow-datafusion/pull/1455) ([maxburke](https://github.com/maxburke)) - -**Documentation updates:** - -- Add dependencies to ballista example documentation [\#1346](https://github.com/apache/arrow-datafusion/pull/1346) ([jgoday](https://github.com/jgoday)) -- \[MINOR\] Fix some typos. [\#1310](https://github.com/apache/arrow-datafusion/pull/1310) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- fix some clippy warnings from nightly channel [\#1277](https://github.com/apache/arrow-datafusion/pull/1277) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([Jimexist](https://github.com/Jimexist)) - -**Performance improvements:** - -- Introduce push-based task scheduling for Ballista [\#1560](https://github.com/apache/arrow-datafusion/pull/1560) ([yahoNanJing](https://github.com/yahoNanJing)) - -**Closed issues:** - -- Track memory usage in Non Limited Operators [\#1569](https://github.com/apache/arrow-datafusion/issues/1569) -- \[Question\] Why does ballista store tables in the client instead of in the SchedulerServer [\#1473](https://github.com/apache/arrow-datafusion/issues/1473) -- Why use the expr types before coercion to get the result type? [\#1358](https://github.com/apache/arrow-datafusion/issues/1358) -- A problem about the projection\_push\_down optimizer gathers valid columns [\#1312](https://github.com/apache/arrow-datafusion/issues/1312) -- apply constant folding to `LogicalPlan::Values` [\#1170](https://github.com/apache/arrow-datafusion/issues/1170) -- reduce usage of `IntoIterator` in logical plan builder window fn [\#372](https://github.com/apache/arrow-datafusion/issues/372) - -**Merged pull requests:** - -- Fix verification scripts for 7.0.0 release [\#1830](https://github.com/apache/arrow-datafusion/pull/1830) ([alamb](https://github.com/alamb)) -- update README for ballista [\#1817](https://github.com/apache/arrow-datafusion/pull/1817) ([liukun4515](https://github.com/liukun4515)) -- Fix logical conflict [\#1801](https://github.com/apache/arrow-datafusion/pull/1801) ([alamb](https://github.com/alamb)) -- Improve the error message and UX of tpch benchmark program [\#1800](https://github.com/apache/arrow-datafusion/pull/1800) ([alamb](https://github.com/alamb)) -- Update to sqlparser 0.14 [\#1796](https://github.com/apache/arrow-datafusion/pull/1796) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([alamb](https://github.com/alamb)) -- Update datafusion versions [\#1793](https://github.com/apache/arrow-datafusion/pull/1793) ([matthewmturner](https://github.com/matthewmturner)) -- Update datafusion to use arrow 9.0.0 [\#1775](https://github.com/apache/arrow-datafusion/pull/1775) ([alamb](https://github.com/alamb)) -- Update parking\_lot requirement from 0.11 to 0.12 [\#1735](https://github.com/apache/arrow-datafusion/pull/1735) ([dependabot[bot]](https://github.com/apps/dependabot)) -- substitute `parking_lot::Mutex` for `std::sync::Mutex` [\#1720](https://github.com/apache/arrow-datafusion/pull/1720) ([xudong963](https://github.com/xudong963)) -- Create ListingTableConfig which includes file format and schema inference [\#1715](https://github.com/apache/arrow-datafusion/pull/1715) ([matthewmturner](https://github.com/matthewmturner)) -- Support `create_physical_expr` and `ExecutionContextState` or `DefaultPhysicalPlanner` for faster speed [\#1700](https://github.com/apache/arrow-datafusion/pull/1700) ([alamb](https://github.com/alamb)) -- Use NamedTempFile rather than `String` in DiskManager [\#1680](https://github.com/apache/arrow-datafusion/pull/1680) ([alamb](https://github.com/alamb)) -- Abstract over logical and physical plan representations in Ballista [\#1677](https://github.com/apache/arrow-datafusion/pull/1677) ([thinkharderdev](https://github.com/thinkharderdev)) -- upgrade clap to version 3 [\#1672](https://github.com/apache/arrow-datafusion/pull/1672) ([Jimexist](https://github.com/Jimexist)) -- Improve configuration and resource use of `MemoryManager` and `DiskManager` [\#1668](https://github.com/apache/arrow-datafusion/pull/1668) ([alamb](https://github.com/alamb)) -- Make `MemoryManager` and `MemoryStream` public [\#1664](https://github.com/apache/arrow-datafusion/pull/1664) ([yjshen](https://github.com/yjshen)) -- Consolidate Schema and RecordBatch projection [\#1638](https://github.com/apache/arrow-datafusion/pull/1638) ([alamb](https://github.com/alamb)) -- Update hashbrown requirement from 0.11 to 0.12 [\#1631](https://github.com/apache/arrow-datafusion/pull/1631) ([dependabot[bot]](https://github.com/apps/dependabot)) -- Update etcd-client requirement from 0.7 to 0.8 [\#1626](https://github.com/apache/arrow-datafusion/pull/1626) ([dependabot[bot]](https://github.com/apps/dependabot)) -- update nightly version [\#1597](https://github.com/apache/arrow-datafusion/pull/1597) ([Jimexist](https://github.com/Jimexist)) -- Add support show tables and show columns for ballista [\#1593](https://github.com/apache/arrow-datafusion/pull/1593) ([gaojun2048](https://github.com/gaojun2048)) -- minor: improve the benchmark readme [\#1567](https://github.com/apache/arrow-datafusion/pull/1567) ([xudong963](https://github.com/xudong963)) -- Consolidate `batch_size` configuration in `ExecutionConfig`, `RuntimeConfig` and `PhysicalPlanConfig` [\#1562](https://github.com/apache/arrow-datafusion/pull/1562) ([yjshen](https://github.com/yjshen)) -- Update to rust 1.58 [\#1557](https://github.com/apache/arrow-datafusion/pull/1557) ([xudong963](https://github.com/xudong963)) -- support mathematics operation for decimal data type [\#1554](https://github.com/apache/arrow-datafusion/pull/1554) ([liukun4515](https://github.com/liukun4515)) -- Make call SchedulerServer::new once in ballista-scheduler process [\#1537](https://github.com/apache/arrow-datafusion/pull/1537) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- Add load test command in tpch.rs. [\#1530](https://github.com/apache/arrow-datafusion/pull/1530) ([Ted-Jiang](https://github.com/Ted-Jiang)) -- Remove one copy of ballista datatype serialization code [\#1524](https://github.com/apache/arrow-datafusion/pull/1524) ([alamb](https://github.com/alamb)) -- Update to arrow-7.0.0 [\#1523](https://github.com/apache/arrow-datafusion/pull/1523) ([alamb](https://github.com/alamb)) -- Workaround build failure: Pin quote to 1.0.10 [\#1499](https://github.com/apache/arrow-datafusion/pull/1499) ([alamb](https://github.com/alamb)) -- add rfcs for datafusion [\#1490](https://github.com/apache/arrow-datafusion/pull/1490) ([xudong963](https://github.com/xudong963)) -- support comparison for decimal data type and refactor the binary coercion rule [\#1483](https://github.com/apache/arrow-datafusion/pull/1483) ([liukun4515](https://github.com/liukun4515)) -- Update arrow-rs to 6.4.0 and replace boolean comparison in datafusion with arrow compute kernel [\#1446](https://github.com/apache/arrow-datafusion/pull/1446) ([xudong963](https://github.com/xudong963)) -- support cast/try\_cast for decimal: signed numeric to decimal [\#1442](https://github.com/apache/arrow-datafusion/pull/1442) ([liukun4515](https://github.com/liukun4515)) -- use 0.13 sql parser [\#1435](https://github.com/apache/arrow-datafusion/pull/1435) ([Jimexist](https://github.com/Jimexist)) -- Clarify communication on bi-weekly sync [\#1427](https://github.com/apache/arrow-datafusion/pull/1427) ([alamb](https://github.com/alamb)) -- Minimize features [\#1399](https://github.com/apache/arrow-datafusion/pull/1399) ([carols10cents](https://github.com/carols10cents)) -- Update rust vesion to 1.57 [\#1395](https://github.com/apache/arrow-datafusion/pull/1395) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([xudong963](https://github.com/xudong963)) -- Add coercion rules for AggregateFunctions [\#1387](https://github.com/apache/arrow-datafusion/pull/1387) ([liukun4515](https://github.com/liukun4515)) -- upgrade the arrow-rs version [\#1385](https://github.com/apache/arrow-datafusion/pull/1385) ([liukun4515](https://github.com/liukun4515)) -- Extract logical plan: rename the plan name \(follow up\) [\#1354](https://github.com/apache/arrow-datafusion/pull/1354) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([liukun4515](https://github.com/liukun4515)) -- upgrade arrow-rs to 6.2.0 [\#1334](https://github.com/apache/arrow-datafusion/pull/1334) ([liukun4515](https://github.com/liukun4515)) -- Update release instructions [\#1331](https://github.com/apache/arrow-datafusion/pull/1331) ([alamb](https://github.com/alamb)) -- Extract Aggregate, Sort, and Join to struct from AggregatePlan [\#1326](https://github.com/apache/arrow-datafusion/pull/1326) ([matthewmturner](https://github.com/matthewmturner)) -- Extract `EmptyRelation`, `Limit`, `Values` from `LogicalPlan` [\#1325](https://github.com/apache/arrow-datafusion/pull/1325) ([liukun4515](https://github.com/liukun4515)) -- Extract CrossJoin, Repartition, Union in LogicalPlan [\#1322](https://github.com/apache/arrow-datafusion/pull/1322) ([liukun4515](https://github.com/liukun4515)) -- Extract Explain, Analyze, Extension in LogicalPlan as independent struct [\#1317](https://github.com/apache/arrow-datafusion/pull/1317) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([xudong963](https://github.com/xudong963)) -- Extract CreateMemoryTable, DropTable, CreateExternalTable in LogicalPlan as independent struct [\#1311](https://github.com/apache/arrow-datafusion/pull/1311) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([liukun4515](https://github.com/liukun4515)) -- Extract Projection, Filter, Window in LogicalPlan as independent struct [\#1309](https://github.com/apache/arrow-datafusion/pull/1309) ([ic4y](https://github.com/ic4y)) -- Add PSQL comparison tests for except, intersect [\#1292](https://github.com/apache/arrow-datafusion/pull/1292) ([mrob95](https://github.com/mrob95)) -- Extract logical plans in LogicalPlan as independent struct: TableScan [\#1290](https://github.com/apache/arrow-datafusion/pull/1290) ([xudong963](https://github.com/xudong963)) - -## [6.0.0-rc0](https://github.com/apache/arrow-datafusion/tree/6.0.0-rc0) (2021-11-14) - -[Full Changelog](https://github.com/apache/arrow-datafusion/compare/6.0.0...6.0.0-rc0) - -## [6.0.0](https://github.com/apache/arrow-datafusion/tree/6.0.0) (2021-11-14) - -[Full Changelog](https://github.com/apache/arrow-datafusion/compare/ballista-0.6.0...6.0.0) - - -## [ballista-0.6.0](https://github.com/apache/arrow-datafusion/tree/ballista-0.6.0) (2021-11-13) - -[Full Changelog](https://github.com/apache/arrow-datafusion/compare/ballista-0.5.0...ballista-0.6.0) - -**Breaking changes:** - -- File partitioning for ListingTable [\#1141](https://github.com/apache/arrow-datafusion/pull/1141) ([rdettai](https://github.com/rdettai)) -- Register tables in BallistaContext using TableProviders instead of Dataframe [\#1028](https://github.com/apache/arrow-datafusion/pull/1028) ([rdettai](https://github.com/rdettai)) -- Make TableProvider.scan\(\) and PhysicalPlanner::create\_physical\_plan\(\) async [\#1013](https://github.com/apache/arrow-datafusion/pull/1013) ([rdettai](https://github.com/rdettai)) -- Reorganize table providers by table format [\#1010](https://github.com/apache/arrow-datafusion/pull/1010) ([rdettai](https://github.com/rdettai)) -- Move CBOs and Statistics to physical plan [\#965](https://github.com/apache/arrow-datafusion/pull/965) ([rdettai](https://github.com/rdettai)) -- Update to sqlparser v 0.10.0 [\#934](https://github.com/apache/arrow-datafusion/pull/934) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([alamb](https://github.com/alamb)) -- FilePartition and PartitionedFile for scanning flexibility [\#932](https://github.com/apache/arrow-datafusion/pull/932) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([yjshen](https://github.com/yjshen)) -- Improve SQLMetric APIs, port existing metrics [\#908](https://github.com/apache/arrow-datafusion/pull/908) ([alamb](https://github.com/alamb)) -- Add support for EXPLAIN ANALYZE [\#858](https://github.com/apache/arrow-datafusion/pull/858) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([alamb](https://github.com/alamb)) -- Rename concurrency to target\_partitions [\#706](https://github.com/apache/arrow-datafusion/pull/706) ([andygrove](https://github.com/andygrove)) - -**Implemented enhancements:** - -- Update datafusion-cli to support Ballista, or implement new ballista-cli [\#886](https://github.com/apache/arrow-datafusion/issues/886) -- Prepare Ballista crates for publishing [\#509](https://github.com/apache/arrow-datafusion/issues/509) -- Add drop table support [\#1266](https://github.com/apache/arrow-datafusion/pull/1266) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([viirya](https://github.com/viirya)) -- use arrow 6.1.0 [\#1255](https://github.com/apache/arrow-datafusion/pull/1255) ([Jimexist](https://github.com/Jimexist)) -- Add support for `create table as` via MemTable [\#1243](https://github.com/apache/arrow-datafusion/pull/1243) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([Dandandan](https://github.com/Dandandan)) -- add values list expression [\#1165](https://github.com/apache/arrow-datafusion/pull/1165) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([Jimexist](https://github.com/Jimexist)) -- Multiple files per partitions for CSV Avro Json [\#1138](https://github.com/apache/arrow-datafusion/pull/1138) ([rdettai](https://github.com/rdettai)) -- Implement INTERSECT & INTERSECT DISTINCT [\#1135](https://github.com/apache/arrow-datafusion/pull/1135) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([xudong963](https://github.com/xudong963)) -- Simplify file struct abstractions [\#1120](https://github.com/apache/arrow-datafusion/pull/1120) ([rdettai](https://github.com/rdettai)) -- Implement `is [not] distinct from` [\#1117](https://github.com/apache/arrow-datafusion/pull/1117) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([Dandandan](https://github.com/Dandandan)) -- add digest\(utf8, method\) function and refactor all current hash digest functions [\#1090](https://github.com/apache/arrow-datafusion/pull/1090) ([Jimexist](https://github.com/Jimexist)) -- \[crypto\] add `blake3` algorithm to `digest` function [\#1086](https://github.com/apache/arrow-datafusion/pull/1086) ([Jimexist](https://github.com/Jimexist)) -- \[crypto\] add blake2b and blake2s functions [\#1081](https://github.com/apache/arrow-datafusion/pull/1081) ([Jimexist](https://github.com/Jimexist)) -- Update sqlparser-rs to 0.11 [\#1052](https://github.com/apache/arrow-datafusion/pull/1052) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([alamb](https://github.com/alamb)) -- remove hard coded partition count in ballista logicalplan deserialization [\#1044](https://github.com/apache/arrow-datafusion/pull/1044) ([xudong963](https://github.com/xudong963)) -- Indexed field access for List [\#1006](https://github.com/apache/arrow-datafusion/pull/1006) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([Igosuki](https://github.com/Igosuki)) -- Update DataFusion to arrow 6.0 [\#984](https://github.com/apache/arrow-datafusion/pull/984) ([alamb](https://github.com/alamb)) -- Implement Display for Expr, improve operator display [\#971](https://github.com/apache/arrow-datafusion/pull/971) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([matthewmturner](https://github.com/matthewmturner)) -- ObjectStore API to read from remote storage systems [\#950](https://github.com/apache/arrow-datafusion/pull/950) ([yjshen](https://github.com/yjshen)) -- fixes \#933 replace placeholder fmt\_as fr ExecutionPlan impls [\#939](https://github.com/apache/arrow-datafusion/pull/939) ([tiphaineruy](https://github.com/tiphaineruy)) -- Support `NotLike` in Ballista [\#916](https://github.com/apache/arrow-datafusion/pull/916) ([Dandandan](https://github.com/Dandandan)) -- Avro Table Provider [\#910](https://github.com/apache/arrow-datafusion/pull/910) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([Igosuki](https://github.com/Igosuki)) -- Add BaselineMetrics, Timestamp metrics, add for `CoalescePartitionsExec`, rename output\_time -\> elapsed\_compute [\#909](https://github.com/apache/arrow-datafusion/pull/909) ([alamb](https://github.com/alamb)) -- \[Ballista\] Add executor last seen info to the ui [\#895](https://github.com/apache/arrow-datafusion/pull/895) ([msathis](https://github.com/msathis)) -- add cross join support to ballista [\#891](https://github.com/apache/arrow-datafusion/pull/891) ([houqp](https://github.com/houqp)) -- Add Ballista support to DataFusion CLI [\#889](https://github.com/apache/arrow-datafusion/pull/889) ([andygrove](https://github.com/andygrove)) -- Add support for PostgreSQL regex match [\#870](https://github.com/apache/arrow-datafusion/pull/870) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([b41sh](https://github.com/b41sh)) - -**Fixed bugs:** - -- Test execution\_plans::shuffle\_writer::tests::test Fail [\#1040](https://github.com/apache/arrow-datafusion/issues/1040) -- Integration test fails to build docker images [\#918](https://github.com/apache/arrow-datafusion/issues/918) -- Ballista: Remove hard-coded concurrency from logical plan serde code [\#708](https://github.com/apache/arrow-datafusion/issues/708) -- How can I make ballista distributed compute work? [\#327](https://github.com/apache/arrow-datafusion/issues/327) -- fix subquery alias [\#1067](https://github.com/apache/arrow-datafusion/pull/1067) [[sql](https://github.com/apache/arrow-datafusion/labels/sql)] ([xudong963](https://github.com/xudong963)) -- Fix compilation for ballista in stand-alone mode [\#1008](https://github.com/apache/arrow-datafusion/pull/1008) ([Igosuki](https://github.com/Igosuki)) - -**Documentation updates:** - -- Add Ballista roadmap [\#1166](https://github.com/apache/arrow-datafusion/pull/1166) ([andygrove](https://github.com/andygrove)) -- Adds note on compatible rust version [\#1097](https://github.com/apache/arrow-datafusion/pull/1097) ([1nF0rmed](https://github.com/1nF0rmed)) -- implement `approx_distinct` function using HyperLogLog [\#1087](https://github.com/apache/arrow-datafusion/pull/1087) ([Jimexist](https://github.com/Jimexist)) -- Improve User Guide [\#954](https://github.com/apache/arrow-datafusion/pull/954) ([andygrove](https://github.com/andygrove)) -- Update plan\_query\_stages doc [\#951](https://github.com/apache/arrow-datafusion/pull/951) ([rdettai](https://github.com/rdettai)) -- \[DataFusion\] - Add show and show\_limit function for DataFrame [\#923](https://github.com/apache/arrow-datafusion/pull/923) ([francis-du](https://github.com/francis-du)) -- update docs related to protoc and optional syntax [\#902](https://github.com/apache/arrow-datafusion/pull/902) ([Jimexist](https://github.com/Jimexist)) -- Improve Ballista crate README content [\#878](https://github.com/apache/arrow-datafusion/pull/878) ([andygrove](https://github.com/andygrove)) - -**Performance improvements:** - -- optimize build profile for datafusion python binding, cli and ballista [\#1137](https://github.com/apache/arrow-datafusion/pull/1137) ([houqp](https://github.com/houqp)) - -**Closed issues:** - -- InList expr with NULL literals do not work [\#1190](https://github.com/apache/arrow-datafusion/issues/1190) -- update the homepage README to include values, `approx_distinct`, etc. [\#1171](https://github.com/apache/arrow-datafusion/issues/1171) -- \[Python\]: Inconsistencies with Python package name [\#1011](https://github.com/apache/arrow-datafusion/issues/1011) -- Wanting to contribute to project where to start? [\#983](https://github.com/apache/arrow-datafusion/issues/983) -- delete redundant code [\#973](https://github.com/apache/arrow-datafusion/issues/973) -- How to build DataFusion python wheel [\#853](https://github.com/apache/arrow-datafusion/issues/853) -- Produce a design for a metrics framework [\#21](https://github.com/apache/arrow-datafusion/issues/21) - -**Merged pull requests:** - -- \[nit\] simplify ballista executor `CollectExec` impl codes [\#1140](https://github.com/apache/arrow-datafusion/pull/1140) ([panarch](https://github.com/panarch)) - - -For older versions, see [apache/arrow/CHANGELOG.md](https://github.com/apache/arrow/blob/master/CHANGELOG.md) - -## [ballista-0.5.0](https://github.com/apache/arrow-datafusion/tree/ballista-0.5.0) (2021-08-10) - -[Full Changelog](https://github.com/apache/arrow-datafusion/compare/4.0.0...ballista-0.5.0) - -**Breaking changes:** - -- \[ballista\] support date\_part and date\_turnc ser/de, pass tpch 7 [\#840](https://github.com/apache/arrow-datafusion/pull/840) ([houqp](https://github.com/houqp)) -- Box ScalarValue:Lists, reduce size by half size [\#788](https://github.com/apache/arrow-datafusion/pull/788) ([alamb](https://github.com/alamb)) -- Support DataFrame.collect for Ballista DataFrames [\#785](https://github.com/apache/arrow-datafusion/pull/785) ([andygrove](https://github.com/andygrove)) -- JOIN conditions are order dependent [\#778](https://github.com/apache/arrow-datafusion/pull/778) ([seddonm1](https://github.com/seddonm1)) -- UnresolvedShuffleExec should represent a single shuffle [\#727](https://github.com/apache/arrow-datafusion/pull/727) ([andygrove](https://github.com/andygrove)) -- Ballista: Make shuffle partitions configurable in benchmarks [\#702](https://github.com/apache/arrow-datafusion/pull/702) ([andygrove](https://github.com/andygrove)) -- Rename MergeExec to CoalescePartitionsExec [\#635](https://github.com/apache/arrow-datafusion/pull/635) ([andygrove](https://github.com/andygrove)) -- Ballista: Rename QueryStageExec to ShuffleWriterExec [\#633](https://github.com/apache/arrow-datafusion/pull/633) ([andygrove](https://github.com/andygrove)) -- fix 593, reduce cloning by taking ownership in logical planner's `from` fn [\#610](https://github.com/apache/arrow-datafusion/pull/610) ([Jimexist](https://github.com/Jimexist)) -- fix join column handling logic for `On` and `Using` constraints [\#605](https://github.com/apache/arrow-datafusion/pull/605) ([houqp](https://github.com/houqp)) -- Move ballista standalone mode to client [\#589](https://github.com/apache/arrow-datafusion/pull/589) ([edrevo](https://github.com/edrevo)) -- Ballista: Implement map-side shuffle [\#543](https://github.com/apache/arrow-datafusion/pull/543) ([andygrove](https://github.com/andygrove)) -- ShuffleReaderExec now supports multiple locations per partition [\#541](https://github.com/apache/arrow-datafusion/pull/541) ([andygrove](https://github.com/andygrove)) -- Make external hostname in executor optional [\#232](https://github.com/apache/arrow-datafusion/pull/232) ([edrevo](https://github.com/edrevo)) -- Remove namespace from executors [\#75](https://github.com/apache/arrow-datafusion/pull/75) ([edrevo](https://github.com/edrevo)) -- Support qualified columns in queries [\#55](https://github.com/apache/arrow-datafusion/pull/55) ([houqp](https://github.com/houqp)) -- Read CSV format text from stdin or memory [\#54](https://github.com/apache/arrow-datafusion/pull/54) ([heymind](https://github.com/heymind)) -- Remove Ballista DataFrame [\#48](https://github.com/apache/arrow-datafusion/pull/48) ([andygrove](https://github.com/andygrove)) -- Use atomics for SQLMetric implementation, remove unused name field [\#25](https://github.com/apache/arrow-datafusion/pull/25) ([returnString](https://github.com/returnString)) - -**Implemented enhancements:** - -- Add crate documentation for Ballista crates [\#830](https://github.com/apache/arrow-datafusion/issues/830) -- Support DataFrame.collect for Ballista DataFrames [\#787](https://github.com/apache/arrow-datafusion/issues/787) -- Ballista: Prep for supporting shuffle correctly, part one [\#736](https://github.com/apache/arrow-datafusion/issues/736) -- Ballista: Implement physical plan serde for ShuffleWriterExec [\#710](https://github.com/apache/arrow-datafusion/issues/710) -- Ballista: Finish implementing shuffle mechanism [\#707](https://github.com/apache/arrow-datafusion/issues/707) -- Rename QueryStageExec to ShuffleWriterExec [\#542](https://github.com/apache/arrow-datafusion/issues/542) -- Ballista ShuffleReaderExec should be able to read from multiple locations per partition [\#540](https://github.com/apache/arrow-datafusion/issues/540) -- \[Ballista\] Use deployments in k8s user guide [\#473](https://github.com/apache/arrow-datafusion/issues/473) -- Ballista refactor QueryStageExec in preparation for map-side shuffle [\#458](https://github.com/apache/arrow-datafusion/issues/458) -- Ballista: Implement map-side of shuffle [\#456](https://github.com/apache/arrow-datafusion/issues/456) -- Refactor Ballista to separate Flight logic from execution logic [\#449](https://github.com/apache/arrow-datafusion/issues/449) -- Use published versions of arrow rather than github shas [\#393](https://github.com/apache/arrow-datafusion/issues/393) -- BallistaContext::collect\(\) logging is too noisy [\#352](https://github.com/apache/arrow-datafusion/issues/352) -- Update Ballista to use new physical plan formatter utility [\#343](https://github.com/apache/arrow-datafusion/issues/343) -- Add Ballista Getting Started documentation [\#329](https://github.com/apache/arrow-datafusion/issues/329) -- Remove references to ballistacompute Docker Hub repo [\#325](https://github.com/apache/arrow-datafusion/issues/325) -- Implement scalable distributed joins [\#63](https://github.com/apache/arrow-datafusion/issues/63) -- Remove hard-coded Ballista version from scripts [\#32](https://github.com/apache/arrow-datafusion/issues/32) -- Implement streaming versions of Dataframe.collect methods [\#789](https://github.com/apache/arrow-datafusion/pull/789) ([andygrove](https://github.com/andygrove)) -- Ballista shuffle is finally working as intended, providing scalable distributed joins [\#750](https://github.com/apache/arrow-datafusion/pull/750) ([andygrove](https://github.com/andygrove)) -- Update to use arrow 5.0 [\#721](https://github.com/apache/arrow-datafusion/pull/721) ([alamb](https://github.com/alamb)) -- Implement serde for ShuffleWriterExec [\#712](https://github.com/apache/arrow-datafusion/pull/712) ([andygrove](https://github.com/andygrove)) -- dedup using join column in wildcard expansion [\#678](https://github.com/apache/arrow-datafusion/pull/678) ([houqp](https://github.com/houqp)) -- Implement metrics for shuffle read and write [\#676](https://github.com/apache/arrow-datafusion/pull/676) ([andygrove](https://github.com/andygrove)) -- Remove hard-coded PartitionMode from Ballista serde [\#637](https://github.com/apache/arrow-datafusion/pull/637) ([andygrove](https://github.com/andygrove)) -- Ballista: Implement scalable distributed joins [\#634](https://github.com/apache/arrow-datafusion/pull/634) ([andygrove](https://github.com/andygrove)) -- Add Keda autoscaling for ballista in k8s [\#586](https://github.com/apache/arrow-datafusion/pull/586) ([edrevo](https://github.com/edrevo)) -- Add some resiliency to lost executors [\#568](https://github.com/apache/arrow-datafusion/pull/568) ([edrevo](https://github.com/edrevo)) -- Add `partition by` constructs in window functions and modify logical planning [\#501](https://github.com/apache/arrow-datafusion/pull/501) ([Jimexist](https://github.com/Jimexist)) -- Support anti join [\#482](https://github.com/apache/arrow-datafusion/pull/482) ([Dandandan](https://github.com/Dandandan)) -- add `order by` construct in window function and logical plans [\#463](https://github.com/apache/arrow-datafusion/pull/463) ([Jimexist](https://github.com/Jimexist)) -- Refactor Ballista executor so that FlightService delegates to an Executor struct [\#450](https://github.com/apache/arrow-datafusion/pull/450) ([andygrove](https://github.com/andygrove)) -- implement lead and lag built-in window function [\#429](https://github.com/apache/arrow-datafusion/pull/429) ([Jimexist](https://github.com/Jimexist)) -- Implement fmt\_as for ShuffleReaderExec [\#400](https://github.com/apache/arrow-datafusion/pull/400) ([andygrove](https://github.com/andygrove)) -- Add window expression part 1 - logical and physical planning, structure, to/from proto, and explain, for empty over clause only [\#334](https://github.com/apache/arrow-datafusion/pull/334) ([Jimexist](https://github.com/Jimexist)) -- \[breaking change\] fix 265, log should be log10, and add ln [\#271](https://github.com/apache/arrow-datafusion/pull/271) ([Jimexist](https://github.com/Jimexist)) -- Allow table providers to indicate their type for catalog metadata [\#205](https://github.com/apache/arrow-datafusion/pull/205) ([returnString](https://github.com/returnString)) -- Add query 19 to TPC-H regression tests [\#59](https://github.com/apache/arrow-datafusion/pull/59) ([Dandandan](https://github.com/Dandandan)) -- Use arrow eq kernels in CaseWhen expression evaluation [\#52](https://github.com/apache/arrow-datafusion/pull/52) ([Dandandan](https://github.com/Dandandan)) -- Add option param for standalone mode [\#42](https://github.com/apache/arrow-datafusion/pull/42) ([djKooks](https://github.com/djKooks)) -- \[DataFusion\] Optimize hash join inner workings, null handling fix [\#24](https://github.com/apache/arrow-datafusion/pull/24) ([Dandandan](https://github.com/Dandandan)) -- \[Ballista\] Docker files for ui [\#22](https://github.com/apache/arrow-datafusion/pull/22) ([msathis](https://github.com/msathis)) - -**Fixed bugs:** - -- Ballista: TPC-H q3 @ SF=1000 never completes [\#835](https://github.com/apache/arrow-datafusion/issues/835) -- Ballista does not support MIN/MAX aggregate functions [\#832](https://github.com/apache/arrow-datafusion/issues/832) -- Ballista docker images fail to build [\#828](https://github.com/apache/arrow-datafusion/issues/828) -- Ballista: UnresolvedShuffleExec should only have a single stage\_id [\#726](https://github.com/apache/arrow-datafusion/issues/726) -- Ballista integration tests are failing [\#623](https://github.com/apache/arrow-datafusion/issues/623) -- Integration test build failure due to arrow-rs using unstable feature [\#596](https://github.com/apache/arrow-datafusion/issues/596) -- `cargo build` cannot build the project [\#531](https://github.com/apache/arrow-datafusion/issues/531) -- ShuffleReaderExec does not get formatted correctly in displayable physical plan [\#399](https://github.com/apache/arrow-datafusion/issues/399) -- Implement serde for MIN and MAX [\#833](https://github.com/apache/arrow-datafusion/pull/833) ([andygrove](https://github.com/andygrove)) -- Ballista: Prep for fixing shuffle mechansim, part 1 [\#738](https://github.com/apache/arrow-datafusion/pull/738) ([andygrove](https://github.com/andygrove)) -- Ballista: Shuffle write bug fix [\#714](https://github.com/apache/arrow-datafusion/pull/714) ([andygrove](https://github.com/andygrove)) -- honor table name for csv/parquet scan in ballista plan serde [\#629](https://github.com/apache/arrow-datafusion/pull/629) ([houqp](https://github.com/houqp)) -- MINOR: Fix integration tests by adding datafusion-cli module to docker image [\#322](https://github.com/apache/arrow-datafusion/pull/322) ([andygrove](https://github.com/andygrove)) - -**Documentation updates:** - -- Add minimal crate documentation for Ballista crates [\#831](https://github.com/apache/arrow-datafusion/pull/831) ([andygrove](https://github.com/andygrove)) -- Add Ballista examples [\#775](https://github.com/apache/arrow-datafusion/pull/775) ([andygrove](https://github.com/andygrove)) -- Update ballista.proto link in architecture doc [\#502](https://github.com/apache/arrow-datafusion/pull/502) ([terrycorley](https://github.com/terrycorley)) -- Update k8s user guide to use deployments [\#474](https://github.com/apache/arrow-datafusion/pull/474) ([edrevo](https://github.com/edrevo)) -- use prettier to format md files [\#367](https://github.com/apache/arrow-datafusion/pull/367) ([Jimexist](https://github.com/Jimexist)) -- Make it easier for developers to find Ballista documentation [\#330](https://github.com/apache/arrow-datafusion/pull/330) ([andygrove](https://github.com/andygrove)) -- Instructions for cross-compiling Ballista to the Raspberry Pi [\#263](https://github.com/apache/arrow-datafusion/pull/263) ([andygrove](https://github.com/andygrove)) -- Add install guide in README [\#236](https://github.com/apache/arrow-datafusion/pull/236) ([djKooks](https://github.com/djKooks)) - -**Performance improvements:** - -- Ballista: Avoid sleeping between polling for tasks [\#698](https://github.com/apache/arrow-datafusion/pull/698) ([Dandandan](https://github.com/Dandandan)) -- Make BallistaContext::collect streaming [\#535](https://github.com/apache/arrow-datafusion/pull/535) ([edrevo](https://github.com/edrevo)) - -**Closed issues:** - -- Confirm git tagging strategy for releases [\#770](https://github.com/apache/arrow-datafusion/issues/770) -- arrow::util::pretty::pretty\_format\_batches missing [\#769](https://github.com/apache/arrow-datafusion/issues/769) -- move the `assert_batches_eq!` macros to a non part of datafusion [\#745](https://github.com/apache/arrow-datafusion/issues/745) -- fix an issue where aliases are not respected in generating downstream schemas in window expr [\#592](https://github.com/apache/arrow-datafusion/issues/592) -- make the planner to print more succinct and useful information in window function explain clause [\#526](https://github.com/apache/arrow-datafusion/issues/526) -- move window frame module to be in `logical_plan` [\#517](https://github.com/apache/arrow-datafusion/issues/517) -- use a more rust idiomatic way of handling nth\_value [\#448](https://github.com/apache/arrow-datafusion/issues/448) -- Make Ballista not depend on arrow directly [\#446](https://github.com/apache/arrow-datafusion/issues/446) -- create a test with more than one partition for window functions [\#435](https://github.com/apache/arrow-datafusion/issues/435) -- Implement hash-partitioned hash aggregate [\#27](https://github.com/apache/arrow-datafusion/issues/27) -- Consider using GitHub pages for DataFusion/Ballista documentation [\#18](https://github.com/apache/arrow-datafusion/issues/18) -- Add Ballista to default cargo workspace [\#17](https://github.com/apache/arrow-datafusion/issues/17) -- Update "repository" in Cargo.toml [\#16](https://github.com/apache/arrow-datafusion/issues/16) -- Consolidate TPC-H benchmarks [\#6](https://github.com/apache/arrow-datafusion/issues/6) -- \[Ballista\] Fix integration test script [\#4](https://github.com/apache/arrow-datafusion/issues/4) -- Ballista should not have separate DataFrame implementation [\#2](https://github.com/apache/arrow-datafusion/issues/2) - -**Merged pull requests:** - -- Change datatype of tpch keys from Int32 to UInt64 to support sf=1000 [\#836](https://github.com/apache/arrow-datafusion/pull/836) ([andygrove](https://github.com/andygrove)) -- Add ballista-examples to docker build [\#829](https://github.com/apache/arrow-datafusion/pull/829) ([andygrove](https://github.com/andygrove)) -- Update dependencies: prost to 0.8 and tonic to 0.5 [\#818](https://github.com/apache/arrow-datafusion/pull/818) ([alamb](https://github.com/alamb)) -- Move `hash_array` into hash\_utils.rs [\#807](https://github.com/apache/arrow-datafusion/pull/807) ([alamb](https://github.com/alamb)) -- Fix: Update clippy lints for Rust 1.54 [\#794](https://github.com/apache/arrow-datafusion/pull/794) ([alamb](https://github.com/alamb)) -- MINOR: Remove unused Ballista query execution code path [\#732](https://github.com/apache/arrow-datafusion/pull/732) ([andygrove](https://github.com/andygrove)) -- \[fix\] benchmark run with compose [\#666](https://github.com/apache/arrow-datafusion/pull/666) ([rdettai](https://github.com/rdettai)) -- bring back dev scripts for ballista [\#648](https://github.com/apache/arrow-datafusion/pull/648) ([Jimexist](https://github.com/Jimexist)) -- Remove unnecessary mutex [\#639](https://github.com/apache/arrow-datafusion/pull/639) ([edrevo](https://github.com/edrevo)) -- round trip TPCH queries in tests [\#630](https://github.com/apache/arrow-datafusion/pull/630) ([houqp](https://github.com/houqp)) -- Fix build [\#627](https://github.com/apache/arrow-datafusion/pull/627) ([andygrove](https://github.com/andygrove)) -- in ballista also check for UI prettier changes [\#578](https://github.com/apache/arrow-datafusion/pull/578) ([Jimexist](https://github.com/Jimexist)) -- turn on clippy rule for needless borrow [\#545](https://github.com/apache/arrow-datafusion/pull/545) ([Jimexist](https://github.com/Jimexist)) -- reuse datafusion physical planner in ballista building from protobuf [\#532](https://github.com/apache/arrow-datafusion/pull/532) ([Jimexist](https://github.com/Jimexist)) -- update cargo.toml in python crate and fix unit test due to hash joins [\#483](https://github.com/apache/arrow-datafusion/pull/483) ([Jimexist](https://github.com/Jimexist)) -- make `VOLUME` declaration in tpch datagen docker absolute [\#466](https://github.com/apache/arrow-datafusion/pull/466) ([crepererum](https://github.com/crepererum)) -- Refactor QueryStageExec in preparation for implementing map-side shuffle [\#459](https://github.com/apache/arrow-datafusion/pull/459) ([andygrove](https://github.com/andygrove)) -- Simplified usage of `use arrow` in ballista. [\#447](https://github.com/apache/arrow-datafusion/pull/447) ([jorgecarleitao](https://github.com/jorgecarleitao)) -- Benchmark subcommand to distinguish between DataFusion and Ballista [\#402](https://github.com/apache/arrow-datafusion/pull/402) ([jgoday](https://github.com/jgoday)) -- \#352: BallistaContext::collect\(\) logging is too noisy [\#394](https://github.com/apache/arrow-datafusion/pull/394) ([jgoday](https://github.com/jgoday)) -- cleanup function return type fn [\#350](https://github.com/apache/arrow-datafusion/pull/350) ([Jimexist](https://github.com/Jimexist)) -- Update Ballista to use new physical plan formatter utility [\#344](https://github.com/apache/arrow-datafusion/pull/344) ([andygrove](https://github.com/andygrove)) -- Update arrow dependencies again [\#341](https://github.com/apache/arrow-datafusion/pull/341) ([alamb](https://github.com/alamb)) -- Remove references to Ballista Docker images published to ballistacompute Docker Hub repo [\#326](https://github.com/apache/arrow-datafusion/pull/326) ([andygrove](https://github.com/andygrove)) -- Update arrow-rs deps [\#317](https://github.com/apache/arrow-datafusion/pull/317) ([alamb](https://github.com/alamb)) -- Update arrow deps [\#269](https://github.com/apache/arrow-datafusion/pull/269) ([alamb](https://github.com/alamb)) -- Enable redundant\_field\_names clippy lint [\#261](https://github.com/apache/arrow-datafusion/pull/261) ([Dandandan](https://github.com/Dandandan)) -- Update arrow-rs deps \(to fix build due to flatbuffers update\) [\#224](https://github.com/apache/arrow-datafusion/pull/224) ([alamb](https://github.com/alamb)) -- update arrow-rs deps to latest master [\#216](https://github.com/apache/arrow-datafusion/pull/216) ([alamb](https://github.com/alamb)) - - - -\* *This Changelog was automatically generated by [github_changelog_generator](https://github.com/github-changelog-generator/github-changelog-generator)* diff --git a/ballista/README.md b/ballista/README.md deleted file mode 100644 index 1fc3bdbf064c..000000000000 --- a/ballista/README.md +++ /dev/null @@ -1,71 +0,0 @@ - - -# Ballista: Distributed Compute with Apache Arrow and DataFusion - -Ballista is a distributed compute platform primarily implemented in Rust, and powered by Apache Arrow and -DataFusion. It is built on an architecture that allows other programming languages (such as Python, C++, and -Java) to be supported as first-class citizens without paying a penalty for serialization costs. - -The foundational technologies in Ballista are: - -- [Apache Arrow](https://arrow.apache.org/) memory model and compute kernels for efficient processing of data. -- [Apache Arrow Flight Protocol](https://arrow.apache.org/blog/2019/10/13/introducing-arrow-flight/) for efficient - data transfer between processes. -- [Google Protocol Buffers](https://developers.google.com/protocol-buffers) for serializing query plans. -- [Docker](https://www.docker.com/) for packaging up executors along with user-defined code. - -Ballista can be deployed as a standalone cluster and also supports [Kubernetes](https://kubernetes.io/). In either -case, the scheduler can be configured to use [etcd](https://etcd.io/) as a backing store to (eventually) provide -redundancy in the case of a scheduler failing. - -# Getting Started - -Refer to the core [Ballista crate README](rust/client/README.md) for the Getting Started guide. - -## Distributed Scheduler Overview - -Ballista uses the DataFusion query execution framework to create a physical plan and then transforms it into a -distributed physical plan by breaking the query down into stages whenever the partitioning scheme changes. - -Specifically, any `RepartitionExec` operator is replaced with an `UnresolvedShuffleExec` and the child operator -of the repartition operator is wrapped in a `ShuffleWriterExec` operator and scheduled for execution. - -Each executor polls the scheduler for the next task to run. Tasks are currently always `ShuffleWriterExec` operators -and each task represents one _input_ partition that will be executed. The resulting batches are repartitioned -according to the shuffle partitioning scheme and each _output_ partition is streamed to disk in Arrow IPC format. - -The scheduler will replace `UnresolvedShuffleExec` operators with `ShuffleReaderExec` operators once all shuffle -tasks have completed. The `ShuffleReaderExec` operator connects to other executors as required using the Flight -interface, and streams the shuffle IPC files. - -# How does this compare to Apache Spark? - -Ballista implements a similar design to Apache Spark, but there are some key differences. - -- The choice of Rust as the main execution language means that memory usage is deterministic and avoids the overhead of - GC pauses. -- Ballista is designed from the ground up to use columnar data, enabling a number of efficiencies such as vectorized - processing (SIMD and GPU) and efficient compression. Although Spark does have some columnar support, it is still - largely row-based today. -- The combination of Rust and Arrow provides excellent memory efficiency and memory usage can be 5x - 10x lower than - Apache Spark in some cases, which means that more processing can fit on a single node, reducing the overhead of - distributed compute. -- The use of Apache Arrow as the memory model and network protocol means that data can be exchanged between executors - in any programming language with minimal serialization overhead. diff --git a/ballista/docs/README.md b/ballista/docs/README.md deleted file mode 100644 index 38d3db5dff39..000000000000 --- a/ballista/docs/README.md +++ /dev/null @@ -1,34 +0,0 @@ - - -# Ballista Developer Documentation - -This directory contains documentation for developers that are contributing to Ballista. If you are looking for -end-user documentation for a published release, please start with the -[DataFusion User Guide](../../docs/user-guide) instead. - -## Architecture & Design - -- Read the [Architecture Overview](architecture.md) to get an understanding of the scheduler and executor - processes and how distributed query execution works. - -## Build, Test, Release - -- Setting up a [development environment](dev-env.md). -- [Integration Testing](integration-testing.md) diff --git a/ballista/docs/architecture.md b/ballista/docs/architecture.md deleted file mode 100644 index 2868d52b943e..000000000000 --- a/ballista/docs/architecture.md +++ /dev/null @@ -1,75 +0,0 @@ - - -# Ballista Architecture - -## Overview - -Ballista allows queries to be executed in a distributed cluster. A cluster consists of one or -more scheduler processes and one or more executor processes. See the following sections in this document for more -details about these components. - -The scheduler accepts logical query plans and translates them into physical query plans using DataFusion and then -runs a secondary planning/optimization process to translate the physical query plan into a distributed physical -query plan. - -This process breaks a query down into a number of query stages that can be executed independently. There are -dependencies between query stages and these dependencies form a directionally-acyclic graph (DAG) because a query -stage cannot start until its child query stages have completed. - -Each query stage has one or more partitions that can be processed in parallel by the available -executors in the cluster. This is the basic unit of scalability in Ballista. - -The following diagram shows the flow of requests and responses between the client, scheduler, and executor -processes. - -![Query Execution Flow](images/query-execution.png) - -## Scheduler Process - -The scheduler process implements a gRPC interface (defined in -[ballista.proto](../rust/core/proto/ballista.proto)). The interface provides the following methods: - -| Method | Description | -| -------------------- | -------------------------------------------------------------------- | -| ExecuteQuery | Submit a logical query plan or SQL query for execution | -| GetExecutorsMetadata | Retrieves a list of executors that have registered with a scheduler | -| GetFileMetadata | Retrieve metadata about files available in the cluster file system | -| GetJobStatus | Get the status of a submitted query | -| RegisterExecutor | Executors call this method to register themselves with the scheduler | - -The scheduler can run in standalone mode, or can be run in clustered mode using etcd as backing store for state. - -## Executor Process - -The executor process implements the Apache Arrow Flight gRPC interface and is responsible for: - -- Executing query stages and persisting the results to disk in Apache Arrow IPC Format -- Making query stage results available as Flights so that they can be retrieved by other executors as well as by - clients - -## Rust Client - -The Rust client provides a DataFrame API that is a thin wrapper around the DataFusion DataFrame and provides -the means for a client to build a query plan for execution. - -The client executes the query plan by submitting an `ExecuteLogicalPlan` request to the scheduler and then calls -`GetJobStatus` to check for completion. On completion, the client receives a list of locations for the Flights -containing the results for the query and will then connect to the appropriate executor processes to retrieve -those results. diff --git a/ballista/docs/dev-env.md b/ballista/docs/dev-env.md deleted file mode 100644 index 1c8605a06e5a..000000000000 --- a/ballista/docs/dev-env.md +++ /dev/null @@ -1,51 +0,0 @@ - - -# Setting up a Rust development environment - -You will need a standard Rust development environment. The easiest way to achieve this is by using rustup: https://rustup.rs/ - -## Install OpenSSL - -Follow instructions for [setting up OpenSSL](https://docs.rs/openssl/latest/openssl/). For Ubuntu users, the following -command works. - -```bash -sudo apt-get install pkg-config libssl-dev -``` - -For MacOS users, the following command works. - -```bash -brew install openssl@1.1 -``` - -## Install CMake - -You'll need cmake in order to compile some of ballista's dependencies. Ubuntu users can use the following command: - -```bash -sudo apt-get install cmake -``` - -MacOS users can use the following command: - -```bash -brew install cmake -``` diff --git a/ballista/docs/images/query-execution.png b/ballista/docs/images/query-execution.png deleted file mode 100644 index b35240282bcd..000000000000 Binary files a/ballista/docs/images/query-execution.png and /dev/null differ diff --git a/ballista/docs/integration-testing.md b/ballista/docs/integration-testing.md deleted file mode 100644 index 5c5b9ee9b656..000000000000 --- a/ballista/docs/integration-testing.md +++ /dev/null @@ -1,29 +0,0 @@ - - -# Integration Testing - -We use the [DataFusion Benchmarks](https://github.com/apache/arrow-datafusion/tree/master/benchmarks) for integration -testing. - -The integration tests can be executed by running the following command from the root of the DataFusion repository. - -```bash -./dev/integration-tests.sh -``` diff --git a/ballista/rust/.dockerignore b/ballista/rust/.dockerignore deleted file mode 100644 index 96f99a522ad0..000000000000 --- a/ballista/rust/.dockerignore +++ /dev/null @@ -1,23 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -# Turn .dockerignore to .dockerallow by excluding everything and explicitly -# allowing specific files and directories. This enables us to quickly add -# dependency files to the docker content without scanning the whole directory. -# This setup requires to all of our docker containers have arrow's source -# as a mounted directory. -target \ No newline at end of file diff --git a/ballista/rust/.gitignore b/ballista/rust/.gitignore deleted file mode 100644 index 97eec1640469..000000000000 --- a/ballista/rust/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -target -temp \ No newline at end of file diff --git a/ballista/rust/client/Cargo.toml b/ballista/rust/client/Cargo.toml deleted file mode 100644 index 828044450b8e..000000000000 --- a/ballista/rust/client/Cargo.toml +++ /dev/null @@ -1,45 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -[package] -name = "ballista" -description = "Ballista Distributed Compute" -license = "Apache-2.0" -version = "0.7.0" -homepage = "https://github.com/apache/arrow-datafusion" -repository = "https://github.com/apache/arrow-datafusion" -readme = "README.md" -authors = ["Apache Arrow "] -edition = "2021" -rust-version = "1.59" - -[dependencies] -ballista-core = { path = "../core", version = "0.7.0" } -ballista-executor = { path = "../executor", version = "0.7.0", optional = true } -ballista-scheduler = { path = "../scheduler", version = "0.7.0", optional = true } - -datafusion = { path = "../../../datafusion/core", version = "8.0.0" } -futures = "0.3" -log = "0.4" -parking_lot = "0.12" -sqlparser = "0.17" -tempfile = "3" -tokio = "1.0" - -[features] -default = [] -standalone = ["ballista-executor", "ballista-scheduler"] diff --git a/ballista/rust/client/README.md b/ballista/rust/client/README.md deleted file mode 100644 index ecf364f4f749..000000000000 --- a/ballista/rust/client/README.md +++ /dev/null @@ -1,133 +0,0 @@ - - -# Ballista: Distributed Scheduler for Apache Arrow DataFusion - -Ballista is a distributed compute platform primarily implemented in Rust, and powered by Apache Arrow and -DataFusion. It is built on an architecture that allows other programming languages (such as Python, C++, and -Java) to be supported as first-class citizens without paying a penalty for serialization costs. - -The foundational technologies in Ballista are: - -- [Apache Arrow](https://arrow.apache.org/) memory model and compute kernels for efficient processing of data. -- [Apache Arrow Flight Protocol](https://arrow.apache.org/blog/2019/10/13/introducing-arrow-flight/) for efficient - data transfer between processes. -- [Google Protocol Buffers](https://developers.google.com/protocol-buffers) for serializing query plans. -- [Docker](https://www.docker.com/) for packaging up executors along with user-defined code. - -Ballista can be deployed as a standalone cluster and also supports [Kubernetes](https://kubernetes.io/). In either -case, the scheduler can be configured to use [etcd](https://etcd.io/) as a backing store to (eventually) provide -redundancy in the case of a scheduler failing. - -## Rust Version Compatbility - -This crate is tested with the latest stable version of Rust. We do not currrently test against other, older versions of the Rust compiler. - -## Starting a cluster - -There are numerous ways to start a Ballista cluster, including support for Docker and -Kubernetes. For full documentation, refer to the -[DataFusion User Guide](https://arrow.apache.org/datafusion/user-guide/introduction.html) - -A simple way to start a local cluster for testing purposes is to use cargo to install -the scheduler and executor crates. - -```bash -cargo install --locked ballista-scheduler -cargo install --locked ballista-executor -``` - -With these crates installed, it is now possible to start a scheduler process. - -```bash -RUST_LOG=info ballista-scheduler -``` - -The scheduler will bind to port 50050 by default. - -Next, start an executor processes in a new terminal session with the specified concurrency -level. - -```bash -RUST_LOG=info ballista-executor -c 4 -``` - -The executor will bind to port 50051 by default. Additional executors can be started by -manually specifying a bind port. For example: - -```bash -RUST_LOG=info ballista-executor --bind-port 50052 -c 4 -``` - -## Executing a query - -Ballista provides a `BallistaContext` as a starting point for creating queries. DataFrames can be created -by invoking the `read_csv`, `read_parquet`, and `sql` methods. - -To build a simple ballista example, add the following dependencies to your `Cargo.toml` file: - -```toml -[dependencies] -ballista = "0.6" -datafusion = "7.0" -tokio = "1.0" -``` - -The following example runs a simple aggregate SQL query against a CSV file from the -[New York Taxi and Limousine Commission](https://www1.nyc.gov/site/tlc/about/tlc-trip-record-data.page) -data set. - -```rust,no_run -use ballista::prelude::*; -use datafusion::arrow::util::pretty; -use datafusion::prelude::CsvReadOptions; - -#[tokio::main] -async fn main() -> Result<()> { - // create configuration - let config = BallistaConfig::builder() - .set("ballista.shuffle.partitions", "4") - .build()?; - - // connect to Ballista scheduler - let ctx = BallistaContext::remote("localhost", 50050, &config).await?; - - // register csv file with the execution context - ctx.register_csv( - "tripdata", - "/path/to/yellow_tripdata_2020-01.csv", - CsvReadOptions::new(), - ).await?; - - // execute the query - let df = ctx.sql( - "SELECT passenger_count, MIN(fare_amount), MAX(fare_amount), AVG(fare_amount), SUM(fare_amount) - FROM tripdata - GROUP BY passenger_count - ORDER BY passenger_count", - ).await?; - - // collect the results and print them to stdout - let results = df.collect().await?; - pretty::print_batches(&results)?; - Ok(()) -} -``` - -More [examples](https://github.com/apache/arrow-datafusion/tree/master/ballista-examples) can be found in the arrow-datafusion repository. diff --git a/ballista/rust/client/src/columnar_batch.rs b/ballista/rust/client/src/columnar_batch.rs deleted file mode 100644 index 3431f5612883..000000000000 --- a/ballista/rust/client/src/columnar_batch.rs +++ /dev/null @@ -1,163 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::{collections::HashMap, sync::Arc}; - -use ballista_core::error::{ballista_error, Result}; - -use datafusion::arrow::{ - array::ArrayRef, - datatypes::{DataType, Schema}, - record_batch::RecordBatch, -}; -use datafusion::scalar::ScalarValue; - -pub type MaybeColumnarBatch = Result>; - -/// Batch of columnar data. -#[derive(Debug, Clone)] -pub struct ColumnarBatch { - schema: Arc, - columns: HashMap, -} - -impl ColumnarBatch { - pub fn from_arrow(batch: &RecordBatch) -> Self { - let columns = batch - .columns() - .iter() - .enumerate() - .map(|(i, array)| { - ( - batch.schema().field(i).name().clone(), - ColumnarValue::Columnar(array.clone()), - ) - }) - .collect(); - - Self { - schema: batch.schema(), - columns, - } - } - - pub fn from_values(values: &[ColumnarValue], schema: &Schema) -> Self { - let columns = schema - .fields() - .iter() - .enumerate() - .map(|(i, f)| (f.name().clone(), values[i].clone())) - .collect(); - - Self { - schema: Arc::new(schema.clone()), - columns, - } - } - - pub fn to_arrow(&self) -> Result { - let arrays = self - .schema - .fields() - .iter() - .map(|c| { - match self.column(c.name())? { - ColumnarValue::Columnar(array) => Ok(array.clone()), - ColumnarValue::Scalar(_, _) => { - // note that this can be implemented easily if needed - Err(ballista_error("Cannot convert scalar value to Arrow array")) - } - } - }) - .collect::>>()?; - - Ok(RecordBatch::try_new(self.schema.clone(), arrays)?) - } - - pub fn schema(&self) -> Arc { - self.schema.clone() - } - - pub fn num_columns(&self) -> usize { - self.columns.len() - } - - pub fn num_rows(&self) -> usize { - self.columns[self.schema.field(0).name()].len() - } - - pub fn column(&self, name: &str) -> Result<&ColumnarValue> { - Ok(&self.columns[name]) - } - - pub fn memory_size(&self) -> usize { - self.columns.values().map(|c| c.memory_size()).sum() - } -} - -/// A columnar value can either be a scalar value or an Arrow array. -#[derive(Debug, Clone)] -pub enum ColumnarValue { - Scalar(ScalarValue, usize), - Columnar(ArrayRef), -} - -impl ColumnarValue { - pub fn len(&self) -> usize { - match self { - ColumnarValue::Scalar(_, n) => *n, - ColumnarValue::Columnar(array) => array.len(), - } - } - - pub fn is_empty(&self) -> bool { - self.len() == 0 - } - - pub fn data_type(&self) -> &DataType { - match self { - ColumnarValue::Columnar(array) => array.data_type(), - ColumnarValue::Scalar(value, _) => match value { - ScalarValue::UInt8(_) => &DataType::UInt8, - ScalarValue::UInt16(_) => &DataType::UInt16, - ScalarValue::UInt32(_) => &DataType::UInt32, - ScalarValue::UInt64(_) => &DataType::UInt64, - ScalarValue::Int8(_) => &DataType::Int8, - ScalarValue::Int16(_) => &DataType::Int16, - ScalarValue::Int32(_) => &DataType::Int32, - ScalarValue::Int64(_) => &DataType::Int64, - ScalarValue::Float32(_) => &DataType::Float32, - ScalarValue::Float64(_) => &DataType::Float64, - _ => unimplemented!(), - }, - } - } - - pub fn to_arrow(&self) -> ArrayRef { - match self { - ColumnarValue::Columnar(array) => array.clone(), - ColumnarValue::Scalar(value, n) => value.to_array_of_size(*n), - } - } - - pub fn memory_size(&self) -> usize { - match self { - ColumnarValue::Columnar(array) => array.get_array_memory_size(), - _ => 0, - } - } -} diff --git a/ballista/rust/client/src/context.rs b/ballista/rust/client/src/context.rs deleted file mode 100644 index a413c5289e5c..000000000000 --- a/ballista/rust/client/src/context.rs +++ /dev/null @@ -1,947 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Distributed execution context. - -use log::info; -use parking_lot::Mutex; -use sqlparser::ast::Statement; -use std::collections::HashMap; -use std::fs; -use std::path::PathBuf; -use std::sync::Arc; - -use ballista_core::config::BallistaConfig; -use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient; -use ballista_core::serde::protobuf::{ExecuteQueryParams, KeyValuePair, LogicalPlanNode}; -use ballista_core::utils::create_df_ctx_with_ballista_query_planner; - -use datafusion::catalog::TableReference; -use datafusion::dataframe::DataFrame; -use datafusion::datasource::TableProvider; -use datafusion::error::{DataFusionError, Result}; -use datafusion::logical_plan::{ - source_as_provider, CreateExternalTable, FileType, LogicalPlan, TableScan, -}; -use datafusion::prelude::{ - AvroReadOptions, CsvReadOptions, ParquetReadOptions, SessionConfig, SessionContext, -}; -use datafusion::sql::parser::{DFParser, Statement as DFStatement}; - -struct BallistaContextState { - /// Ballista configuration - config: BallistaConfig, - /// Scheduler host - scheduler_host: String, - /// Scheduler port - scheduler_port: u16, - /// Tables that have been registered with this context - tables: HashMap>, -} - -impl BallistaContextState { - pub fn new( - scheduler_host: String, - scheduler_port: u16, - config: &BallistaConfig, - ) -> Self { - Self { - config: config.clone(), - scheduler_host, - scheduler_port, - tables: HashMap::new(), - } - } - - pub fn config(&self) -> &BallistaConfig { - &self.config - } -} - -pub struct BallistaContext { - state: Arc>, - context: Arc, -} - -impl BallistaContext { - /// Create a context for executing queries against a remote Ballista scheduler instance - pub async fn remote( - host: &str, - port: u16, - config: &BallistaConfig, - ) -> ballista_core::error::Result { - let state = BallistaContextState::new(host.to_owned(), port, config); - - let scheduler_url = - format!("http://{}:{}", &state.scheduler_host, state.scheduler_port); - info!( - "Connecting to Ballista scheduler at {}", - scheduler_url.clone() - ); - let mut scheduler = SchedulerGrpcClient::connect(scheduler_url.clone()) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; - - let remote_session_id = scheduler - .execute_query(ExecuteQueryParams { - query: None, - settings: config - .settings() - .iter() - .map(|(k, v)| KeyValuePair { - key: k.to_owned(), - value: v.to_owned(), - }) - .collect::>(), - optional_session_id: None, - }) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))? - .into_inner() - .session_id; - - info!( - "Server side SessionContext created with session id: {}", - remote_session_id - ); - - let ctx = { - create_df_ctx_with_ballista_query_planner::( - scheduler_url, - remote_session_id, - state.config(), - ) - }; - - Ok(Self { - state: Arc::new(Mutex::new(state)), - context: Arc::new(ctx), - }) - } - - #[cfg(feature = "standalone")] - pub async fn standalone( - config: &BallistaConfig, - concurrent_tasks: usize, - ) -> ballista_core::error::Result { - use ballista_core::serde::protobuf::PhysicalPlanNode; - use ballista_core::serde::BallistaCodec; - - log::info!("Running in local mode. Scheduler will be run in-proc"); - - let addr = ballista_scheduler::standalone::new_standalone_scheduler().await?; - let scheduler_url = format!("http://localhost:{}", addr.port()); - let mut scheduler = loop { - match SchedulerGrpcClient::connect(scheduler_url.clone()).await { - Err(_) => { - tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; - log::info!("Attempting to connect to in-proc scheduler..."); - } - Ok(scheduler) => break scheduler, - } - }; - - let remote_session_id = scheduler - .execute_query(ExecuteQueryParams { - query: None, - settings: config - .settings() - .iter() - .map(|(k, v)| KeyValuePair { - key: k.to_owned(), - value: v.to_owned(), - }) - .collect::>(), - optional_session_id: None, - }) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))? - .into_inner() - .session_id; - - info!( - "Server side SessionContext created with session id: {}", - remote_session_id - ); - - let ctx = { - create_df_ctx_with_ballista_query_planner::( - scheduler_url, - remote_session_id, - config, - ) - }; - - let default_codec: BallistaCodec = - BallistaCodec::default(); - - ballista_executor::new_standalone_executor( - scheduler, - concurrent_tasks, - default_codec, - ) - .await?; - - let state = - BallistaContextState::new("localhost".to_string(), addr.port(), config); - - Ok(Self { - state: Arc::new(Mutex::new(state)), - context: Arc::new(ctx), - }) - } - - /// Create a DataFrame representing an Avro table scan - /// TODO fetch schema from scheduler instead of resolving locally - pub async fn read_avro( - &self, - path: &str, - options: AvroReadOptions<'_>, - ) -> Result> { - // convert to absolute path because the executor likely has a different working directory - let path = PathBuf::from(path); - let path = fs::canonicalize(&path)?; - - let ctx = self.context.clone(); - let df = ctx.read_avro(path.to_str().unwrap(), options).await?; - Ok(df) - } - - /// Create a DataFrame representing a Parquet table scan - /// TODO fetch schema from scheduler instead of resolving locally - pub async fn read_parquet( - &self, - path: &str, - options: ParquetReadOptions<'_>, - ) -> Result> { - // convert to absolute path because the executor likely has a different working directory - let path = PathBuf::from(path); - let path = fs::canonicalize(&path)?; - - let ctx = self.context.clone(); - let df = ctx.read_parquet(path.to_str().unwrap(), options).await?; - Ok(df) - } - - /// Create a DataFrame representing a CSV table scan - /// TODO fetch schema from scheduler instead of resolving locally - pub async fn read_csv( - &self, - path: &str, - options: CsvReadOptions<'_>, - ) -> Result> { - // convert to absolute path because the executor likely has a different working directory - let path = PathBuf::from(path); - let path = fs::canonicalize(&path)?; - - let ctx = self.context.clone(); - let df = ctx.read_csv(path.to_str().unwrap(), options).await?; - Ok(df) - } - - /// Register a DataFrame as a table that can be referenced from a SQL query - pub fn register_table( - &self, - name: &str, - table: Arc, - ) -> Result<()> { - let mut state = self.state.lock(); - state.tables.insert(name.to_owned(), table); - Ok(()) - } - - pub async fn register_csv( - &self, - name: &str, - path: &str, - options: CsvReadOptions<'_>, - ) -> Result<()> { - match self.read_csv(path, options).await?.to_logical_plan()? { - LogicalPlan::TableScan(TableScan { source, .. }) => { - self.register_table(name, source_as_provider(&source)?) - } - _ => Err(DataFusionError::Internal("Expected tables scan".to_owned())), - } - } - - pub async fn register_parquet( - &self, - name: &str, - path: &str, - options: ParquetReadOptions<'_>, - ) -> Result<()> { - match self.read_parquet(path, options).await?.to_logical_plan()? { - LogicalPlan::TableScan(TableScan { source, .. }) => { - self.register_table(name, source_as_provider(&source)?) - } - _ => Err(DataFusionError::Internal("Expected tables scan".to_owned())), - } - } - - pub async fn register_avro( - &self, - name: &str, - path: &str, - options: AvroReadOptions<'_>, - ) -> Result<()> { - match self.read_avro(path, options).await?.to_logical_plan()? { - LogicalPlan::TableScan(TableScan { source, .. }) => { - self.register_table(name, source_as_provider(&source)?) - } - _ => Err(DataFusionError::Internal("Expected tables scan".to_owned())), - } - } - - /// is a 'show *' sql - pub async fn is_show_statement(&self, sql: &str) -> Result { - let mut is_show_variable: bool = false; - let statements = DFParser::parse_sql(sql)?; - - if statements.len() != 1 { - return Err(DataFusionError::NotImplemented( - "The context currently only supports a single SQL statement".to_string(), - )); - } - - if let DFStatement::Statement(s) = &statements[0] { - let st: &Statement = s; - match st { - Statement::ShowVariable { .. } => { - is_show_variable = true; - } - Statement::ShowColumns { .. } => { - is_show_variable = true; - } - _ => { - is_show_variable = false; - } - } - }; - - Ok(is_show_variable) - } - - /// Create a DataFrame from a SQL statement. - /// - /// This method is `async` because queries of type `CREATE EXTERNAL TABLE` - /// might require the schema to be inferred. - pub async fn sql(&self, sql: &str) -> Result> { - let mut ctx = self.context.clone(); - - let is_show = self.is_show_statement(sql).await?; - // the show tables、 show columns sql can not run at scheduler because the tables is store at client - if is_show { - let state = self.state.lock(); - ctx = Arc::new(SessionContext::with_config( - SessionConfig::new().with_information_schema( - state.config.default_with_information_schema(), - ), - )); - } - - // register tables with DataFusion context - { - let state = self.state.lock(); - for (name, prov) in &state.tables { - // ctx is shared between queries, check table exists or not before register - let table_ref = TableReference::Bare { table: name }; - if !ctx.table_exist(table_ref)? { - ctx.register_table( - TableReference::Bare { table: name }, - Arc::clone(prov), - )?; - } - } - } - - let plan = ctx.create_logical_plan(sql)?; - - match plan { - LogicalPlan::CreateExternalTable(CreateExternalTable { - ref schema, - ref name, - ref location, - ref file_type, - ref has_header, - ref delimiter, - ref table_partition_cols, - ref if_not_exists, - }) => { - let table_exists = ctx.table_exist(name.as_str())?; - - match (if_not_exists, table_exists) { - (_, false) => match file_type { - FileType::CSV => { - self.register_csv( - name, - location, - CsvReadOptions::new() - .schema(&schema.as_ref().to_owned().into()) - .has_header(*has_header) - .delimiter(*delimiter as u8) - .table_partition_cols(table_partition_cols.to_vec()), - ) - .await?; - Ok(Arc::new(DataFrame::new(ctx.state.clone(), &plan))) - } - FileType::Parquet => { - self.register_parquet( - name, - location, - ParquetReadOptions::default() - .table_partition_cols(table_partition_cols.to_vec()), - ) - .await?; - Ok(Arc::new(DataFrame::new(ctx.state.clone(), &plan))) - } - FileType::Avro => { - self.register_avro( - name, - location, - AvroReadOptions::default() - .table_partition_cols(table_partition_cols.to_vec()), - ) - .await?; - Ok(Arc::new(DataFrame::new(ctx.state.clone(), &plan))) - } - _ => Err(DataFusionError::NotImplemented(format!( - "Unsupported file type {:?}.", - file_type - ))), - }, - (true, true) => { - Ok(Arc::new(DataFrame::new(ctx.state.clone(), &plan))) - } - (false, true) => Err(DataFusionError::Execution(format!( - "Table '{:?}' already exists", - name - ))), - } - } - _ => ctx.sql(sql).await, - } - } -} - -#[cfg(test)] -mod tests { - - #[tokio::test] - #[cfg(feature = "standalone")] - async fn test_standalone_mode() { - use super::*; - let context = BallistaContext::standalone(&BallistaConfig::new().unwrap(), 1) - .await - .unwrap(); - let df = context.sql("SELECT 1;").await.unwrap(); - df.collect().await.unwrap(); - } - - #[tokio::test] - #[cfg(feature = "standalone")] - async fn test_ballista_show_tables() { - use super::*; - use std::fs::File; - use std::io::Write; - use tempfile::TempDir; - let context = BallistaContext::standalone(&BallistaConfig::new().unwrap(), 1) - .await - .unwrap(); - - let data = "Jorge,2018-12-13T12:12:10.011Z\n\ - Andrew,2018-11-13T17:11:10.011Z"; - - let tmp_dir = TempDir::new().unwrap(); - let file_path = tmp_dir.path().join("timestamps.csv"); - - // scope to ensure the file is closed and written - { - File::create(&file_path) - .expect("creating temp file") - .write_all(data.as_bytes()) - .expect("writing data"); - } - - let sql = format!( - "CREATE EXTERNAL TABLE csv_with_timestamps ( - name VARCHAR, - ts TIMESTAMP - ) - STORED AS CSV - LOCATION '{}' - ", - file_path.to_str().expect("path is utf8") - ); - - context.sql(sql.as_str()).await.unwrap(); - - let df = context.sql("show columns from csv_with_timestamps;").await; - - assert!(df.is_err()); - } - - #[tokio::test] - #[cfg(feature = "standalone")] - async fn test_show_tables_not_with_information_schema() { - use super::*; - use ballista_core::config::{ - BallistaConfigBuilder, BALLISTA_WITH_INFORMATION_SCHEMA, - }; - use std::fs::File; - use std::io::Write; - use tempfile::TempDir; - let config = BallistaConfigBuilder::default() - .set(BALLISTA_WITH_INFORMATION_SCHEMA, "true") - .build() - .unwrap(); - let context = BallistaContext::standalone(&config, 1).await.unwrap(); - - let data = "Jorge,2018-12-13T12:12:10.011Z\n\ - Andrew,2018-11-13T17:11:10.011Z"; - - let tmp_dir = TempDir::new().unwrap(); - let file_path = tmp_dir.path().join("timestamps.csv"); - - // scope to ensure the file is closed and written - { - File::create(&file_path) - .expect("creating temp file") - .write_all(data.as_bytes()) - .expect("writing data"); - } - - let sql = format!( - "CREATE EXTERNAL TABLE csv_with_timestamps ( - name VARCHAR, - ts TIMESTAMP - ) - STORED AS CSV - LOCATION '{}' - ", - file_path.to_str().expect("path is utf8") - ); - - context.sql(sql.as_str()).await.unwrap(); - let df = context.sql("show tables;").await; - assert!(df.is_ok()); - } - - #[tokio::test] - #[cfg(feature = "standalone")] - #[ignore] - // Tracking: https://github.com/apache/arrow-datafusion/issues/1840 - async fn test_task_stuck_when_referenced_task_failed() { - use super::*; - use datafusion::arrow::datatypes::Schema; - use datafusion::arrow::util::pretty; - use datafusion::datasource::file_format::csv::CsvFormat; - use datafusion::datasource::listing::{ - ListingOptions, ListingTable, ListingTableConfig, - }; - - use ballista_core::config::{ - BallistaConfigBuilder, BALLISTA_WITH_INFORMATION_SCHEMA, - }; - let config = BallistaConfigBuilder::default() - .set(BALLISTA_WITH_INFORMATION_SCHEMA, "true") - .build() - .unwrap(); - let context = BallistaContext::standalone(&config, 1).await.unwrap(); - - let testdata = datafusion::test_util::parquet_test_data(); - context - .register_parquet( - "single_nan", - &format!("{}/single_nan.parquet", testdata), - ParquetReadOptions::default(), - ) - .await - .unwrap(); - - { - let mut guard = context.state.lock(); - let csv_table = guard.tables.get("single_nan"); - - if let Some(table_provide) = csv_table { - if let Some(listing_table) = table_provide - .clone() - .as_any() - .downcast_ref::() - { - let x = listing_table.options(); - let error_options = ListingOptions { - file_extension: x.file_extension.clone(), - format: Arc::new(CsvFormat::default()), - table_partition_cols: x.table_partition_cols.clone(), - collect_stat: x.collect_stat, - target_partitions: x.target_partitions, - }; - - let config = ListingTableConfig::new( - listing_table.object_store().clone(), - listing_table.table_path().to_string(), - ) - .with_schema(Arc::new(Schema::new(vec![]))) - .with_listing_options(error_options); - - let error_table = ListingTable::try_new(config).unwrap(); - - // change the table to an error table - guard - .tables - .insert("single_nan".to_string(), Arc::new(error_table)); - } - } - } - - let df = context - .sql("select count(1) from single_nan;") - .await - .unwrap(); - let results = df.collect().await.unwrap(); - pretty::print_batches(&results); - } - - #[tokio::test] - #[cfg(feature = "standalone")] - async fn test_empty_exec_with_one_row() { - use crate::context::BallistaContext; - use ballista_core::config::{ - BallistaConfigBuilder, BALLISTA_WITH_INFORMATION_SCHEMA, - }; - - let config = BallistaConfigBuilder::default() - .set(BALLISTA_WITH_INFORMATION_SCHEMA, "true") - .build() - .unwrap(); - let context = BallistaContext::standalone(&config, 1).await.unwrap(); - - let sql = "select EXTRACT(year FROM to_timestamp('2020-09-08T12:13:14+00:00'));"; - - let df = context.sql(sql).await.unwrap(); - assert!(!df.collect().await.unwrap().is_empty()); - } - - #[tokio::test] - #[cfg(feature = "standalone")] - async fn test_union_and_union_all() { - use super::*; - use ballista_core::config::{ - BallistaConfigBuilder, BALLISTA_WITH_INFORMATION_SCHEMA, - }; - use datafusion::arrow::util::pretty::pretty_format_batches; - let config = BallistaConfigBuilder::default() - .set(BALLISTA_WITH_INFORMATION_SCHEMA, "true") - .build() - .unwrap(); - let context = BallistaContext::standalone(&config, 1).await.unwrap(); - - let df = context - .sql("SELECT 1 as NUMBER union SELECT 1 as NUMBER;") - .await - .unwrap(); - let res1 = df.collect().await.unwrap(); - let expected1 = vec![ - "+--------+", - "| number |", - "+--------+", - "| 1 |", - "+--------+", - ]; - assert_eq!( - expected1, - pretty_format_batches(&*res1) - .unwrap() - .to_string() - .trim() - .lines() - .collect::>() - ); - let expected2 = vec![ - "+--------+", - "| number |", - "+--------+", - "| 1 |", - "| 1 |", - "+--------+", - ]; - let df = context - .sql("SELECT 1 as NUMBER union all SELECT 1 as NUMBER;") - .await - .unwrap(); - let res2 = df.collect().await.unwrap(); - assert_eq!( - expected2, - pretty_format_batches(&*res2) - .unwrap() - .to_string() - .trim() - .lines() - .collect::>() - ); - } - - #[tokio::test] - #[cfg(feature = "standalone")] - async fn test_aggregate_func() { - use crate::context::BallistaContext; - use ballista_core::config::{ - BallistaConfigBuilder, BALLISTA_WITH_INFORMATION_SCHEMA, - }; - use datafusion::arrow; - use datafusion::arrow::util::pretty::pretty_format_batches; - use datafusion::prelude::ParquetReadOptions; - - let config = BallistaConfigBuilder::default() - .set(BALLISTA_WITH_INFORMATION_SCHEMA, "true") - .build() - .unwrap(); - let context = BallistaContext::standalone(&config, 1).await.unwrap(); - - let testdata = datafusion::test_util::parquet_test_data(); - context - .register_parquet( - "test", - &format!("{}/alltypes_plain.parquet", testdata), - ParquetReadOptions::default(), - ) - .await - .unwrap(); - - let df = context.sql("select min(\"id\") from test").await.unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+--------------+", - "| MIN(test.id) |", - "+--------------+", - "| 0 |", - "+--------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context.sql("select max(\"id\") from test").await.unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+--------------+", - "| MAX(test.id) |", - "+--------------+", - "| 7 |", - "+--------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context.sql("select SUM(\"id\") from test").await.unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+--------------+", - "| SUM(test.id) |", - "+--------------+", - "| 28 |", - "+--------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context.sql("select AVG(\"id\") from test").await.unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+--------------+", - "| AVG(test.id) |", - "+--------------+", - "| 3.5 |", - "+--------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context.sql("select COUNT(\"id\") from test").await.unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+----------------+", - "| COUNT(test.id) |", - "+----------------+", - "| 8 |", - "+----------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select approx_distinct(\"id\") from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+-------------------------+", - "| APPROXDISTINCT(test.id) |", - "+-------------------------+", - "| 8 |", - "+-------------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select ARRAY_AGG(\"id\") from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+--------------------------+", - "| ARRAYAGG(test.id) |", - "+--------------------------+", - "| [4, 5, 6, 7, 2, 3, 0, 1] |", - "+--------------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context.sql("select VAR(\"id\") from test").await.unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+-------------------+", - "| VARIANCE(test.id) |", - "+-------------------+", - "| 6.000000000000001 |", - "+-------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select VAR_POP(\"id\") from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+----------------------+", - "| VARIANCEPOP(test.id) |", - "+----------------------+", - "| 5.250000000000001 |", - "+----------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select VAR_SAMP(\"id\") from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+-------------------+", - "| VARIANCE(test.id) |", - "+-------------------+", - "| 6.000000000000001 |", - "+-------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select STDDEV(\"id\") from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+--------------------+", - "| STDDEV(test.id) |", - "+--------------------+", - "| 2.4494897427831783 |", - "+--------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select STDDEV_SAMP(\"id\") from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+--------------------+", - "| STDDEV(test.id) |", - "+--------------------+", - "| 2.4494897427831783 |", - "+--------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select COVAR(id, tinyint_col) from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+--------------------------------------+", - "| COVARIANCE(test.id,test.tinyint_col) |", - "+--------------------------------------+", - "| 0.28571428571428586 |", - "+--------------------------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select CORR(id, tinyint_col) from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+---------------------------------------+", - "| CORRELATION(test.id,test.tinyint_col) |", - "+---------------------------------------+", - "| 0.21821789023599245 |", - "+---------------------------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select approx_percentile_cont_with_weight(\"id\", 2, 0.5) from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+---------------------------------------------------------------+", - "| APPROXPERCENTILECONTWITHWEIGHT(test.id,Int64(2),Float64(0.5)) |", - "+---------------------------------------------------------------+", - "| 1 |", - "+---------------------------------------------------------------+", - ]; - assert_result_eq(expected, &*res); - - let df = context - .sql("select approx_percentile_cont(\"double_col\", 0.5) from test") - .await - .unwrap(); - let res = df.collect().await.unwrap(); - let expected = vec![ - "+----------------------------------------------------+", - "| APPROXPERCENTILECONT(test.double_col,Float64(0.5)) |", - "+----------------------------------------------------+", - "| 7.574999999999999 |", - "+----------------------------------------------------+", - ]; - - assert_result_eq(expected, &*res); - - fn assert_result_eq( - expected: Vec<&str>, - results: &[arrow::record_batch::RecordBatch], - ) { - assert_eq!( - expected, - pretty_format_batches(results) - .unwrap() - .to_string() - .trim() - .lines() - .collect::>() - ); - } - } -} diff --git a/ballista/rust/client/src/lib.rs b/ballista/rust/client/src/lib.rs deleted file mode 100644 index 125278dcca40..000000000000 --- a/ballista/rust/client/src/lib.rs +++ /dev/null @@ -1,22 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -#![doc = include_str!("../README.md")] - -pub mod columnar_batch; -pub mod context; -pub mod prelude; diff --git a/ballista/rust/client/src/prelude.rs b/ballista/rust/client/src/prelude.rs deleted file mode 100644 index 15558a357daa..000000000000 --- a/ballista/rust/client/src/prelude.rs +++ /dev/null @@ -1,26 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Ballista Prelude (common imports) - -pub use crate::context::BallistaContext; -pub use ballista_core::config::BallistaConfig; -pub use ballista_core::config::BALLISTA_DEFAULT_BATCH_SIZE; -pub use ballista_core::config::BALLISTA_DEFAULT_SHUFFLE_PARTITIONS; -pub use ballista_core::error::{BallistaError, Result}; - -pub use futures::StreamExt; diff --git a/ballista/rust/core/Cargo.toml b/ballista/rust/core/Cargo.toml deleted file mode 100644 index 3d43aceb4f0f..000000000000 --- a/ballista/rust/core/Cargo.toml +++ /dev/null @@ -1,67 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -[package] -name = "ballista-core" -description = "Ballista Distributed Compute" -license = "Apache-2.0" -version = "0.7.0" -homepage = "https://github.com/apache/arrow-datafusion" -repository = "https://github.com/apache/arrow-datafusion" -readme = "README.md" -authors = ["Apache Arrow "] -edition = "2018" -build = "build.rs" - -[features] -# Used for testing ONLY: causes all values to hash to the same value (test for collisions) -force_hash_collisions = ["datafusion/force_hash_collisions"] -simd = ["datafusion/simd"] - -[dependencies] -ahash = { version = "0.7", default-features = false } - -arrow-flight = { version = "14.0.0" } -async-trait = "0.1.41" -chrono = { version = "0.4", default-features = false } -clap = { version = "3", features = ["derive", "cargo"] } -datafusion = { path = "../../../datafusion/core", version = "8.0.0" } -datafusion-proto = { path = "../../../datafusion/proto", version = "8.0.0" } -futures = "0.3" -hashbrown = "0.12" - -libloading = "0.7.3" -log = "0.4" -once_cell = "1.9.0" - -parking_lot = "0.12" -parse_arg = "0.1.3" -prost = "0.10" -prost-types = "0.10" -serde = { version = "1", features = ["derive"] } -sqlparser = "0.17" -tokio = "1.0" -tonic = "0.7" -uuid = { version = "1.0", features = ["v4"] } -walkdir = "2.3.2" - -[dev-dependencies] -tempfile = "3" - -[build-dependencies] -rustc_version = "0.4.0" -tonic-build = { version = "0.7" } diff --git a/ballista/rust/core/README.md b/ballista/rust/core/README.md deleted file mode 100644 index 2b4c9fbfd3e4..000000000000 --- a/ballista/rust/core/README.md +++ /dev/null @@ -1,23 +0,0 @@ - - -# Ballista Core Library - -This crate contains the Ballista core library which is used as a dependency by the `ballista-client`, -`ballista-scheduler`, and `ballista-executor` crates. diff --git a/ballista/rust/core/build.rs b/ballista/rust/core/build.rs deleted file mode 100644 index c2acde108a2b..000000000000 --- a/ballista/rust/core/build.rs +++ /dev/null @@ -1,30 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -fn main() -> Result<(), String> { - // for use in docker build where file changes can be wonky - println!("cargo:rerun-if-env-changed=FORCE_REBUILD"); - - println!("cargo:rerun-if-changed=proto/ballista.proto"); - let version = rustc_version::version().unwrap(); - println!("cargo:rustc-env=RUSTC_VERSION={}", version); - println!("cargo:rerun-if-changed=proto/datafusion.proto"); - tonic_build::configure() - .extern_path(".datafusion", "::datafusion_proto::protobuf") - .compile(&["proto/ballista.proto"], &["proto"]) - .map_err(|e| format!("protobuf compilation failed: {}", e)) -} diff --git a/ballista/rust/core/proto/ballista.proto b/ballista/rust/core/proto/ballista.proto deleted file mode 100644 index 1ceb412f1474..000000000000 --- a/ballista/rust/core/proto/ballista.proto +++ /dev/null @@ -1,918 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - *

- * http://www.apache.org/licenses/LICENSE-2.0 - *

- * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -syntax = "proto3"; - -package ballista.protobuf; - -option java_multiple_files = true; -option java_package = "org.ballistacompute.protobuf"; -option java_outer_classname = "BallistaProto"; - -import "datafusion.proto"; - -/////////////////////////////////////////////////////////////////////////////////////////////////// -// Ballista Logical Plan -/////////////////////////////////////////////////////////////////////////////////////////////////// - -// LogicalPlan is a nested type -message LogicalPlanNode { - oneof LogicalPlanType { - ListingTableScanNode listing_scan = 1; - ProjectionNode projection = 3; - SelectionNode selection = 4; - LimitNode limit = 5; - AggregateNode aggregate = 6; - JoinNode join = 7; - SortNode sort = 8; - RepartitionNode repartition = 9; - EmptyRelationNode empty_relation = 10; - CreateExternalTableNode create_external_table = 11; - ExplainNode explain = 12; - WindowNode window = 13; - AnalyzeNode analyze = 14; - CrossJoinNode cross_join = 15; - ValuesNode values = 16; - LogicalExtensionNode extension = 17; - CreateCatalogSchemaNode create_catalog_schema = 18; - UnionNode union = 19; - CreateCatalogNode create_catalog = 20; - SubqueryAliasNode subquery_alias = 21; - CreateViewNode create_view = 22; - OffsetNode offset = 23; - } -} - -message LogicalExtensionNode { - bytes node = 1; - repeated LogicalPlanNode inputs = 2; -} - -message ProjectionColumns { - repeated string columns = 1; -} - -message Statistics { - int64 num_rows = 1; - int64 total_byte_size = 2; - repeated ColumnStats column_stats = 3; - bool is_exact = 4; -} - -message FileRange { - int64 start = 1; - int64 end = 2; -} - -message PartitionedFile { - string path = 1; - uint64 size = 2; - uint64 last_modified_ns = 3; - repeated datafusion.ScalarValue partition_values = 4; - FileRange range = 5; -} - -message CsvFormat { - bool has_header = 1; - string delimiter = 2; -} - -message ParquetFormat { - bool enable_pruning = 1; -} - -message AvroFormat {} - -message ListingTableScanNode { - string table_name = 1; - string path = 2; - string file_extension = 3; - ProjectionColumns projection = 4; - datafusion.Schema schema = 5; - repeated datafusion.LogicalExprNode filters = 6; - repeated string table_partition_cols = 7; - bool collect_stat = 8; - uint32 target_partitions = 9; - oneof FileFormatType { - CsvFormat csv = 10; - ParquetFormat parquet = 11; - AvroFormat avro = 12; - } -} - -message ProjectionNode { - LogicalPlanNode input = 1; - repeated datafusion.LogicalExprNode expr = 2; - oneof optional_alias { - string alias = 3; - } -} - -message SelectionNode { - LogicalPlanNode input = 1; - datafusion.LogicalExprNode expr = 2; -} - -message SortNode { - LogicalPlanNode input = 1; - repeated datafusion.LogicalExprNode expr = 2; -} - -message RepartitionNode { - LogicalPlanNode input = 1; - oneof partition_method { - uint64 round_robin = 2; - HashRepartition hash = 3; - } -} - -message HashRepartition { - repeated datafusion.LogicalExprNode hash_expr = 1; - uint64 partition_count = 2; -} - -message EmptyRelationNode { - bool produce_one_row = 1; -} - -message CreateExternalTableNode { - string name = 1; - string location = 2; - FileType file_type = 3; - bool has_header = 4; - datafusion.DfSchema schema = 5; - repeated string table_partition_cols = 6; - bool if_not_exists = 7; - string delimiter = 8; -} - -message CreateCatalogSchemaNode { - string schema_name = 1; - bool if_not_exists = 2; - datafusion.DfSchema schema = 3; -} - -message CreateCatalogNode { - string catalog_name = 1; - bool if_not_exists = 2; - datafusion.DfSchema schema = 3; -} - -message CreateViewNode { - string name = 1; - LogicalPlanNode input = 2; - bool or_replace = 3; -} - -// a node containing data for defining values list. unlike in SQL where it's two dimensional, here -// the list is flattened, and with the field n_cols it can be parsed and partitioned into rows -message ValuesNode { - uint64 n_cols = 1; - repeated datafusion.LogicalExprNode values_list = 2; -} - -enum FileType { - NdJson = 0; - Parquet = 1; - CSV = 2; - Avro = 3; -} - -message AnalyzeNode { - LogicalPlanNode input = 1; - bool verbose = 2; -} - -message ExplainNode { - LogicalPlanNode input = 1; - bool verbose = 2; -} - -message AggregateNode { - LogicalPlanNode input = 1; - repeated datafusion.LogicalExprNode group_expr = 2; - repeated datafusion.LogicalExprNode aggr_expr = 3; -} - -message WindowNode { - LogicalPlanNode input = 1; - repeated datafusion.LogicalExprNode window_expr = 2; -} - -enum JoinType { - INNER = 0; - LEFT = 1; - RIGHT = 2; - FULL = 3; - SEMI = 4; - ANTI = 5; -} - -enum JoinConstraint { - ON = 0; - USING = 1; -} - -message JoinNode { - LogicalPlanNode left = 1; - LogicalPlanNode right = 2; - JoinType join_type = 3; - JoinConstraint join_constraint = 4; - repeated datafusion.Column left_join_column = 5; - repeated datafusion.Column right_join_column = 6; - bool null_equals_null = 7; -} - -message UnionNode { - repeated LogicalPlanNode inputs = 1; -} - -message CrossJoinNode { - LogicalPlanNode left = 1; - LogicalPlanNode right = 2; -} - -message LimitNode { - LogicalPlanNode input = 1; - uint32 limit = 2; -} - -message OffsetNode { - LogicalPlanNode input = 1; - uint32 offset = 2; -} - -message SelectionExecNode { - datafusion.LogicalExprNode expr = 1; -} - -message SubqueryAliasNode { - LogicalPlanNode input = 1; - string alias = 2; -} - -/////////////////////////////////////////////////////////////////////////////////////////////////// -// Ballista Physical Plan -/////////////////////////////////////////////////////////////////////////////////////////////////// - -// PhysicalPlanNode is a nested type -message PhysicalPlanNode { - oneof PhysicalPlanType { - ParquetScanExecNode parquet_scan = 1; - CsvScanExecNode csv_scan = 2; - EmptyExecNode empty = 3; - ProjectionExecNode projection = 4; - GlobalLimitExecNode global_limit = 6; - LocalLimitExecNode local_limit = 7; - AggregateExecNode aggregate = 8; - HashJoinExecNode hash_join = 9; - ShuffleReaderExecNode shuffle_reader = 10; - SortExecNode sort = 11; - CoalesceBatchesExecNode coalesce_batches = 12; - FilterExecNode filter = 13; - CoalescePartitionsExecNode merge = 14; - UnresolvedShuffleExecNode unresolved = 15; - RepartitionExecNode repartition = 16; - WindowAggExecNode window = 17; - ShuffleWriterExecNode shuffle_writer = 18; - CrossJoinExecNode cross_join = 19; - AvroScanExecNode avro_scan = 20; - PhysicalExtensionNode extension = 21; - UnionExecNode union = 22; - ExplainExecNode explain = 23; - } -} - -message PhysicalExtensionNode { - bytes node = 1; - repeated PhysicalPlanNode inputs = 2; -} - -// physical expressions -message PhysicalExprNode { - oneof ExprType { - // column references - PhysicalColumn column = 1; - - datafusion.ScalarValue literal = 2; - - // binary expressions - PhysicalBinaryExprNode binary_expr = 3; - - // aggregate expressions - PhysicalAggregateExprNode aggregate_expr = 4; - - // null checks - PhysicalIsNull is_null_expr = 5; - PhysicalIsNotNull is_not_null_expr = 6; - PhysicalNot not_expr = 7; - - PhysicalCaseNode case_ = 8; - PhysicalCastNode cast = 9; - PhysicalSortExprNode sort = 10; - PhysicalNegativeNode negative = 11; - PhysicalInListNode in_list = 12; - PhysicalScalarFunctionNode scalar_function = 13; - PhysicalTryCastNode try_cast = 14; - - // window expressions - PhysicalWindowExprNode window_expr = 15; - - PhysicalScalarUdfNode scalar_udf = 16; - } -} - -message PhysicalScalarUdfNode { - string name = 1; - repeated PhysicalExprNode args = 2; - datafusion.ArrowType return_type = 4; -} - -message PhysicalAggregateExprNode { - datafusion.AggregateFunction aggr_function = 1; - repeated PhysicalExprNode expr = 2; -} - -message PhysicalWindowExprNode { - oneof window_function { - datafusion.AggregateFunction aggr_function = 1; - datafusion.BuiltInWindowFunction built_in_function = 2; - // udaf = 3 - } - PhysicalExprNode expr = 4; -} - -message PhysicalIsNull { - PhysicalExprNode expr = 1; -} - -message PhysicalIsNotNull { - PhysicalExprNode expr = 1; -} - -message PhysicalNot { - PhysicalExprNode expr = 1; -} - -message PhysicalAliasNode { - PhysicalExprNode expr = 1; - string alias = 2; -} - -message PhysicalBinaryExprNode { - PhysicalExprNode l = 1; - PhysicalExprNode r = 2; - string op = 3; -} - -message PhysicalSortExprNode { - PhysicalExprNode expr = 1; - bool asc = 2; - bool nulls_first = 3; -} - -message PhysicalWhenThen { - PhysicalExprNode when_expr = 1; - PhysicalExprNode then_expr = 2; -} - -message PhysicalInListNode { - PhysicalExprNode expr = 1; - repeated PhysicalExprNode list = 2; - bool negated = 3; -} - -message PhysicalCaseNode { - PhysicalExprNode expr = 1; - repeated PhysicalWhenThen when_then_expr = 2; - PhysicalExprNode else_expr = 3; -} - -message PhysicalScalarFunctionNode { - string name = 1; - datafusion.ScalarFunction fun = 2; - repeated PhysicalExprNode args = 3; - datafusion.ArrowType return_type = 4; -} - -message PhysicalTryCastNode { - PhysicalExprNode expr = 1; - datafusion.ArrowType arrow_type = 2; -} - -message PhysicalCastNode { - PhysicalExprNode expr = 1; - datafusion.ArrowType arrow_type = 2; -} - -message PhysicalNegativeNode { - PhysicalExprNode expr = 1; -} - -message UnresolvedShuffleExecNode { - uint32 stage_id = 1; - datafusion.Schema schema = 2; - uint32 input_partition_count = 3; - uint32 output_partition_count = 4; -} - -message FilterExecNode { - PhysicalPlanNode input = 1; - PhysicalExprNode expr = 2; -} - -message FileGroup { - repeated PartitionedFile files = 1; -} - -message ScanLimit { - // wrap into a message to make it optional - uint32 limit = 1; -} - -message FileScanExecConf { - repeated FileGroup file_groups = 1; - datafusion.Schema schema = 2; - repeated uint32 projection = 4; - ScanLimit limit = 5; - Statistics statistics = 6; - repeated string table_partition_cols = 7; -} - -message ParquetScanExecNode { - FileScanExecConf base_conf = 1; - datafusion.LogicalExprNode pruning_predicate = 2; -} - -message CsvScanExecNode { - FileScanExecConf base_conf = 1; - bool has_header = 2; - string delimiter = 3; -} - -message AvroScanExecNode { - FileScanExecConf base_conf = 1; -} - -enum PartitionMode { - COLLECT_LEFT = 0; - PARTITIONED = 1; -} - -message HashJoinExecNode { - PhysicalPlanNode left = 1; - PhysicalPlanNode right = 2; - repeated JoinOn on = 3; - JoinType join_type = 4; - PartitionMode partition_mode = 6; - bool null_equals_null = 7; -} - -message UnionExecNode { - repeated PhysicalPlanNode inputs = 1; -} - -message ExplainExecNode { - datafusion.Schema schema = 1; - repeated datafusion.StringifiedPlan stringified_plans = 2; - bool verbose = 3; -} - -message CrossJoinExecNode { - PhysicalPlanNode left = 1; - PhysicalPlanNode right = 2; -} - -message PhysicalColumn { - string name = 1; - uint32 index = 2; -} - -message JoinOn { - PhysicalColumn left = 1; - PhysicalColumn right = 2; -} - -message EmptyExecNode { - bool produce_one_row = 1; - datafusion.Schema schema = 2; -} - -message ProjectionExecNode { - PhysicalPlanNode input = 1; - repeated PhysicalExprNode expr = 2; - repeated string expr_name = 3; -} - -enum AggregateMode { - PARTIAL = 0; - FINAL = 1; - FINAL_PARTITIONED = 2; -} - -message WindowAggExecNode { - PhysicalPlanNode input = 1; - repeated PhysicalExprNode window_expr = 2; - repeated string window_expr_name = 3; - datafusion.Schema input_schema = 4; -} - -message AggregateExecNode { - repeated PhysicalExprNode group_expr = 1; - repeated PhysicalExprNode aggr_expr = 2; - AggregateMode mode = 3; - PhysicalPlanNode input = 4; - repeated string group_expr_name = 5; - repeated string aggr_expr_name = 6; - // we need the input schema to the partial aggregate to pass to the final aggregate - datafusion.Schema input_schema = 7; -} - -message ShuffleWriterExecNode { - //TODO it seems redundant to provide job and stage id here since we also have them - // in the TaskDefinition that wraps this plan - string job_id = 1; - uint32 stage_id = 2; - PhysicalPlanNode input = 3; - PhysicalHashRepartition output_partitioning = 4; -} - -message ShuffleReaderExecNode { - repeated ShuffleReaderPartition partition = 1; - datafusion.Schema schema = 2; -} - -message ShuffleReaderPartition { - // each partition of a shuffle read can read data from multiple locations - repeated PartitionLocation location = 1; -} - -message GlobalLimitExecNode { - PhysicalPlanNode input = 1; - uint32 limit = 2; -} - -message LocalLimitExecNode { - PhysicalPlanNode input = 1; - uint32 limit = 2; -} - -message SortExecNode { - PhysicalPlanNode input = 1; - repeated PhysicalExprNode expr = 2; -} - -message CoalesceBatchesExecNode { - PhysicalPlanNode input = 1; - uint32 target_batch_size = 2; -} - -message CoalescePartitionsExecNode { - PhysicalPlanNode input = 1; -} - -message PhysicalHashRepartition { - repeated PhysicalExprNode hash_expr = 1; - uint64 partition_count = 2; -} - -message RepartitionExecNode{ - PhysicalPlanNode input = 1; - oneof partition_method { - uint64 round_robin = 2; - PhysicalHashRepartition hash = 3; - uint64 unknown = 4; - } -} - -/////////////////////////////////////////////////////////////////////////////////////////////////// -// Ballista Scheduling -/////////////////////////////////////////////////////////////////////////////////////////////////// - -message KeyValuePair { - string key = 1; - string value = 2; -} - -message Action { - - oneof ActionType { - // Fetch a partition from an executor - FetchPartition fetch_partition = 3; - } - - // configuration settings - repeated KeyValuePair settings = 100; -} - -message ExecutePartition { - string job_id = 1; - uint32 stage_id = 2; - repeated uint32 partition_id = 3; - PhysicalPlanNode plan = 4; - // The task could need to read partitions from other executors - repeated PartitionLocation partition_location = 5; - // Output partition for shuffle writer - PhysicalHashRepartition output_partitioning = 6; -} - -message FetchPartition { - string job_id = 1; - uint32 stage_id = 2; - uint32 partition_id = 3; - string path = 4; -} - -// Mapping from partition id to executor id -message PartitionLocation { - PartitionId partition_id = 1; - ExecutorMetadata executor_meta = 2; - PartitionStats partition_stats = 3; - string path = 4; -} - -// Unique identifier for a materialized partition of data -message PartitionId { - string job_id = 1; - uint32 stage_id = 2; - uint32 partition_id = 4; -} - -message PartitionStats { - int64 num_rows = 1; - int64 num_batches = 2; - int64 num_bytes = 3; - repeated ColumnStats column_stats = 4; -} - -message ColumnStats { - datafusion.ScalarValue min_value = 1; - datafusion.ScalarValue max_value = 2; - uint32 null_count = 3; - uint32 distinct_count = 4; -} - -// Used by scheduler -message ExecutorMetadata { - string id = 1; - string host = 2; - uint32 port = 3; - uint32 grpc_port = 4; - ExecutorSpecification specification = 5; -} - -// Used by grpc -message ExecutorRegistration { - string id = 1; - // "optional" keyword is stable in protoc 3.15 but prost is still on 3.14 (see https://github.com/tokio-rs/prost/issues/430 and https://github.com/tokio-rs/prost/pull/455) - // this syntax is ugly but is binary compatible with the "optional" keyword (see https://stackoverflow.com/questions/42622015/how-to-define-an-optional-field-in-protobuf-3) - oneof optional_host { - string host = 2; - } - uint32 port = 3; - uint32 grpc_port = 4; - ExecutorSpecification specification = 5; -} - -message ExecutorHeartbeat { - string executor_id = 1; - // Unix epoch-based timestamp in seconds - uint64 timestamp = 2; - ExecutorState state = 3; -} - -message ExecutorState { - repeated ExecutorMetric metrics = 1; -} - -message ExecutorMetric { - // TODO add more metrics - oneof metric { - uint64 available_memory = 1; - } -} - -message ExecutorSpecification { - repeated ExecutorResource resources = 1; -} - -message ExecutorResource { - // TODO add more resources - oneof resource { - uint32 task_slots = 1; - } -} - -message ExecutorData { - string executor_id = 1; - repeated ExecutorResourcePair resources = 2; -} - -message ExecutorResourcePair { - ExecutorResource total = 1; - ExecutorResource available = 2; -} - -message RunningTask { - string executor_id = 1; -} - -message FailedTask { - string error = 1; -} - -message CompletedTask { - string executor_id = 1; - // 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; -} - -message ShuffleWritePartition { - uint64 partition_id = 1; - string path = 2; - uint64 num_batches = 3; - uint64 num_rows = 4; - uint64 num_bytes = 5; -} - -message TaskStatus { - PartitionId task_id = 1; - oneof status { - RunningTask running = 2; - FailedTask failed = 3; - CompletedTask completed = 4; - } -} - -message PollWorkParams { - ExecutorRegistration metadata = 1; - bool can_accept_task = 2; - // All tasks must be reported until they reach the failed or completed state - repeated TaskStatus task_status = 3; -} - -message TaskDefinition { - PartitionId task_id = 1; - bytes plan = 2; - // Output partition for shuffle writer - PhysicalHashRepartition output_partitioning = 3; - string session_id = 4; - repeated KeyValuePair props = 5; -} - -message JobSessionConfig { - string session_id = 1; - repeated KeyValuePair configs = 2; -} - -message PollWorkResult { - TaskDefinition task = 1; -} - -message RegisterExecutorParams { - ExecutorRegistration metadata = 1; -} - -message RegisterExecutorResult { - bool success = 1; -} - -message HeartBeatParams { - string executor_id = 1; - ExecutorState state = 2; -} - -message HeartBeatResult { - // TODO it's from Spark for BlockManager - bool reregister = 1; -} - -message StopExecutorParams { -} - -message StopExecutorResult { -} - -message UpdateTaskStatusParams { - string executor_id = 1; - // All tasks must be reported until they reach the failed or completed state - repeated TaskStatus task_status = 2; -} - -message UpdateTaskStatusResult { - bool success = 1; -} - -message ExecuteQueryParams { - oneof query { - bytes logical_plan = 1; - string sql = 2; - } - oneof optional_session_id { - string session_id = 3; - } - repeated KeyValuePair settings = 4; -} - -message ExecuteSqlParams { - string sql = 1; -} - -message ExecuteQueryResult { - string job_id = 1; - string session_id = 2; -} - -message GetJobStatusParams { - string job_id = 1; -} - -message CompletedJob { - repeated PartitionLocation partition_location = 1; -} - -message QueuedJob {} - -// TODO: add progress report -message RunningJob {} - -message FailedJob { - string error = 1; -} - -message JobStatus { - oneof status { - QueuedJob queued = 1; - RunningJob running = 2; - FailedJob failed = 3; - CompletedJob completed = 4; - } -} - -message GetJobStatusResult { - JobStatus status = 1; -} - -message GetFileMetadataParams { - string path = 1; - FileType file_type = 2; -} - -message GetFileMetadataResult { - datafusion.Schema schema = 1; -} - -message FilePartitionMetadata { - repeated string filename = 1; -} - -message LaunchTaskParams { - // Allow to launch a task set to an executor at once - repeated TaskDefinition task = 1; -} - -message LaunchTaskResult { - bool success = 1; - // TODO when part of the task set are scheduled successfully -} - -service SchedulerGrpc { - // Executors must poll the scheduler for heartbeat and to receive tasks - rpc PollWork (PollWorkParams) returns (PollWorkResult) {} - - rpc RegisterExecutor(RegisterExecutorParams) returns (RegisterExecutorResult) {} - - // Push-based task scheduler will only leverage this interface - // rather than the PollWork interface to report executor states - rpc HeartBeatFromExecutor (HeartBeatParams) returns (HeartBeatResult) {} - - rpc UpdateTaskStatus (UpdateTaskStatusParams) returns (UpdateTaskStatusResult) {} - - rpc GetFileMetadata (GetFileMetadataParams) returns (GetFileMetadataResult) {} - - rpc ExecuteQuery (ExecuteQueryParams) returns (ExecuteQueryResult) {} - - rpc GetJobStatus (GetJobStatusParams) returns (GetJobStatusResult) {} -} - -service ExecutorGrpc { - rpc LaunchTask (LaunchTaskParams) returns (LaunchTaskResult) {} - - rpc StopExecutor (StopExecutorParams) returns (StopExecutorResult) {} -} diff --git a/ballista/rust/core/proto/datafusion.proto b/ballista/rust/core/proto/datafusion.proto deleted file mode 100644 index 9999abbf2cf0..000000000000 --- a/ballista/rust/core/proto/datafusion.proto +++ /dev/null @@ -1,542 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - *

- * http://www.apache.org/licenses/LICENSE-2.0 - *

- * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -syntax = "proto3"; - -package datafusion; - -option java_multiple_files = true; -option java_package = "org.datafusioncompute.protobuf"; -option java_outer_classname = "DatafusionProto"; - -message ColumnRelation { - string relation = 1; -} - -message Column { - string name = 1; - ColumnRelation relation = 2; -} - -message DfField{ - Field field = 1; - ColumnRelation qualifier = 2; -} - -message DfSchema { - repeated DfField columns = 1; -} - -// logical expressions -message LogicalExprNode { - oneof ExprType { - // column references - Column column = 1; - - // alias - AliasNode alias = 2; - - ScalarValue literal = 3; - - // binary expressions - BinaryExprNode binary_expr = 4; - - // aggregate expressions - AggregateExprNode aggregate_expr = 5; - - // null checks - IsNull is_null_expr = 6; - IsNotNull is_not_null_expr = 7; - Not not_expr = 8; - - BetweenNode between = 9; - CaseNode case_ = 10; - CastNode cast = 11; - SortExprNode sort = 12; - NegativeNode negative = 13; - InListNode in_list = 14; - bool wildcard = 15; - ScalarFunctionNode scalar_function = 16; - TryCastNode try_cast = 17; - - // window expressions - WindowExprNode window_expr = 18; - } -} - -message IsNull { - LogicalExprNode expr = 1; -} - -message IsNotNull { - LogicalExprNode expr = 1; -} - -message Not { - LogicalExprNode expr = 1; -} - -message AliasNode { - LogicalExprNode expr = 1; - string alias = 2; -} - -message BinaryExprNode { - LogicalExprNode l = 1; - LogicalExprNode r = 2; - string op = 3; -} - -message NegativeNode { - LogicalExprNode expr = 1; -} - -message InListNode { - LogicalExprNode expr = 1; - repeated LogicalExprNode list = 2; - bool negated = 3; -} - -enum ScalarFunction { - Abs=0; - Acos=1; - Asin=2; - Atan=3; - Ascii=4; - Ceil=5; - Cos=6; - Digest=7; - Exp=8; - Floor=9; - Ln=10; - Log=11; - Log10=12; - Log2=13; - Round=14; - Signum=15; - Sin=16; - Sqrt=17; - Tan=18; - Trunc=19; - Array=20; - RegexpMatch=21; - BitLength=22; - Btrim=23; - CharacterLength=24; - Chr=25; - Concat=26; - ConcatWithSeparator=27; - DatePart=28; - DateTrunc=29; - InitCap=30; - Left=31; - Lpad=32; - Lower=33; - Ltrim=34; - MD5=35; - NullIf=36; - OctetLength=37; - Random=38; - RegexpReplace=39; - Repeat=40; - Replace=41; - Reverse=42; - Right=43; - Rpad=44; - Rtrim=45; - SHA224=46; - SHA256=47; - SHA384=48; - SHA512=49; - SplitPart=50; - StartsWith=51; - Strpos=52; - Substr=53; - ToHex=54; - ToTimestamp=55; - ToTimestampMillis=56; - ToTimestampMicros=57; - ToTimestampSeconds=58; - Now=59; - Translate=60; - Trim=61; - Upper=62; - Coalesce=63; -} - -message ScalarFunctionNode { - ScalarFunction fun = 1; - repeated LogicalExprNode args = 2; -} - -enum AggregateFunction { - MIN = 0; - MAX = 1; - SUM = 2; - AVG = 3; - COUNT = 4; - APPROX_DISTINCT = 5; - ARRAY_AGG = 6; - VARIANCE=7; - VARIANCE_POP=8; - COVARIANCE=9; - COVARIANCE_POP=10; - STDDEV=11; - STDDEV_POP=12; - CORRELATION=13; - APPROX_PERCENTILE_CONT = 14; - APPROX_MEDIAN=15; - APPROX_PERCENTILE_CONT_WITH_WEIGHT = 16; -} - -message AggregateExprNode { - AggregateFunction aggr_function = 1; - repeated LogicalExprNode expr = 2; -} - -enum BuiltInWindowFunction { - ROW_NUMBER = 0; - RANK = 1; - DENSE_RANK = 2; - PERCENT_RANK = 3; - CUME_DIST = 4; - NTILE = 5; - LAG = 6; - LEAD = 7; - FIRST_VALUE = 8; - LAST_VALUE = 9; - NTH_VALUE = 10; -} - -message WindowExprNode { - oneof window_function { - AggregateFunction aggr_function = 1; - BuiltInWindowFunction built_in_function = 2; - // udaf = 3 - } - LogicalExprNode expr = 4; - repeated LogicalExprNode partition_by = 5; - repeated LogicalExprNode order_by = 6; - // repeated LogicalExprNode filter = 7; - oneof window_frame { - WindowFrame frame = 8; - } -} - -message BetweenNode { - LogicalExprNode expr = 1; - bool negated = 2; - LogicalExprNode low = 3; - LogicalExprNode high = 4; -} - -message CaseNode { - LogicalExprNode expr = 1; - repeated WhenThen when_then_expr = 2; - LogicalExprNode else_expr = 3; -} - -message WhenThen { - LogicalExprNode when_expr = 1; - LogicalExprNode then_expr = 2; -} - -message CastNode { - LogicalExprNode expr = 1; - ArrowType arrow_type = 2; -} - -message TryCastNode { - LogicalExprNode expr = 1; - ArrowType arrow_type = 2; -} - -message SortExprNode { - LogicalExprNode expr = 1; - bool asc = 2; - bool nulls_first = 3; -} - -enum WindowFrameUnits { - ROWS = 0; - RANGE = 1; - GROUPS = 2; -} - -message WindowFrame { - WindowFrameUnits window_frame_units = 1; - WindowFrameBound start_bound = 2; - // "optional" keyword is stable in protoc 3.15 but prost is still on 3.14 (see https://github.com/tokio-rs/prost/issues/430 and https://github.com/tokio-rs/prost/pull/455) - // this syntax is ugly but is binary compatible with the "optional" keyword (see https://stackoverflow.com/questions/42622015/how-to-define-an-optional-field-in-protobuf-3) - oneof end_bound { - WindowFrameBound bound = 3; - } -} - -enum WindowFrameBoundType { - CURRENT_ROW = 0; - PRECEDING = 1; - FOLLOWING = 2; -} - -message WindowFrameBound { - WindowFrameBoundType window_frame_bound_type = 1; - // "optional" keyword is stable in protoc 3.15 but prost is still on 3.14 (see https://github.com/tokio-rs/prost/issues/430 and https://github.com/tokio-rs/prost/pull/455) - // this syntax is ugly but is binary compatible with the "optional" keyword (see https://stackoverflow.com/questions/42622015/how-to-define-an-optional-field-in-protobuf-3) - oneof bound_value { - uint64 value = 2; - } -} - -/////////////////////////////////////////////////////////////////////////////////////////////////// -// Arrow Data Types -/////////////////////////////////////////////////////////////////////////////////////////////////// - -message Schema { - repeated datafusion.Field columns = 1; -} - -message Field { - // name of the field - string name = 1; - ArrowType arrow_type = 2; - bool nullable = 3; - // for complex data types like structs, unions - repeated Field children = 4; -} - -message FixedSizeBinary{ - int32 length = 1; -} - -message Timestamp{ - TimeUnit time_unit = 1; - string timezone = 2; -} - -enum DateUnit{ - Day = 0; - DateMillisecond = 1; -} - -enum TimeUnit{ - Second = 0; - TimeMillisecond = 1; - Microsecond = 2; - Nanosecond = 3; -} - -enum IntervalUnit{ - YearMonth = 0; - DayTime = 1; - MonthDayNano = 2; -} - -message Decimal{ - uint64 whole = 1; - uint64 fractional = 2; -} - -message List{ - Field field_type = 1; -} - -message FixedSizeList{ - Field field_type = 1; - int32 list_size = 2; -} - -message Dictionary{ - ArrowType key = 1; - ArrowType value = 2; -} - -message Struct{ - repeated Field sub_field_types = 1; -} - -enum UnionMode{ - sparse = 0; - dense = 1; -} - -message Union{ - repeated Field union_types = 1; - UnionMode union_mode = 2; -} - -message ScalarListValue{ - ScalarType datatype = 1; - repeated ScalarValue values = 2; -} - -message ScalarValue{ - oneof value { - bool bool_value = 1; - string utf8_value = 2; - string large_utf8_value = 3; - int32 int8_value = 4; - int32 int16_value = 5; - int32 int32_value = 6; - int64 int64_value = 7; - uint32 uint8_value = 8; - uint32 uint16_value = 9; - uint32 uint32_value = 10; - uint64 uint64_value = 11; - float float32_value = 12; - double float64_value = 13; - //Literal Date32 value always has a unit of day - int32 date_32_value = 14; - int64 time_microsecond_value = 15; - int64 time_nanosecond_value = 16; - ScalarListValue list_value = 17; - ScalarType null_list_value = 18; - - PrimitiveScalarType null_value = 19; - Decimal128 decimal128_value = 20; - int64 date_64_value = 21; - int64 time_second_value = 22; - int64 time_millisecond_value = 23; - int32 interval_yearmonth_value = 24; - int64 interval_daytime_value = 25; - } -} - -message Decimal128{ - bytes value = 1; - int64 p = 2; - int64 s = 3; -} - -// Contains all valid datafusion scalar type except for -// List -enum PrimitiveScalarType{ - - BOOL = 0; // arrow::Type::BOOL - UINT8 = 1; // arrow::Type::UINT8 - INT8 = 2; // arrow::Type::INT8 - UINT16 = 3; // represents arrow::Type fields in src/arrow/type.h - INT16 = 4; - UINT32 = 5; - INT32 = 6; - UINT64 = 7; - INT64 = 8; - FLOAT32 = 9; - FLOAT64 = 10; - UTF8 = 11; - LARGE_UTF8 = 12; - DATE32 = 13; - TIME_MICROSECOND = 14; - TIME_NANOSECOND = 15; - NULL = 16; - - DECIMAL128 = 17; - DATE64 = 20; - TIME_SECOND = 21; - TIME_MILLISECOND = 22; - INTERVAL_YEARMONTH = 23; - INTERVAL_DAYTIME = 24; -} - -message ScalarType{ - oneof datatype{ - PrimitiveScalarType scalar = 1; - ScalarListType list = 2; - } -} - -message ScalarListType{ - repeated string field_names = 3; - PrimitiveScalarType deepest_type = 2; -} - -// Broke out into multiple message types so that type -// metadata did not need to be in separate message -//All types that are of the empty message types contain no additional metadata -// about the type -message ArrowType{ - oneof arrow_type_enum{ - EmptyMessage NONE = 1; // arrow::Type::NA - EmptyMessage BOOL = 2; // arrow::Type::BOOL - EmptyMessage UINT8 = 3; // arrow::Type::UINT8 - EmptyMessage INT8 = 4; // arrow::Type::INT8 - EmptyMessage UINT16 =5; // represents arrow::Type fields in src/arrow/type.h - EmptyMessage INT16 = 6; - EmptyMessage UINT32 =7; - EmptyMessage INT32 = 8; - EmptyMessage UINT64 =9; - EmptyMessage INT64 =10 ; - EmptyMessage FLOAT16 =11 ; - EmptyMessage FLOAT32 =12 ; - EmptyMessage FLOAT64 =13 ; - EmptyMessage UTF8 =14 ; - EmptyMessage LARGE_UTF8 = 32; - EmptyMessage BINARY =15 ; - int32 FIXED_SIZE_BINARY =16 ; - EmptyMessage LARGE_BINARY = 31; - EmptyMessage DATE32 =17 ; - EmptyMessage DATE64 =18 ; - TimeUnit DURATION = 19; - Timestamp TIMESTAMP =20 ; - TimeUnit TIME32 =21 ; - TimeUnit TIME64 =22 ; - IntervalUnit INTERVAL =23 ; - Decimal DECIMAL =24 ; - List LIST =25; - List LARGE_LIST = 26; - FixedSizeList FIXED_SIZE_LIST = 27; - Struct STRUCT =28; - Union UNION =29; - Dictionary DICTIONARY =30; - } -} - -//Useful for representing an empty enum variant in rust -// E.G. enum example{One, Two(i32)} -// maps to -// message example{ -// oneof{ -// EmptyMessage One = 1; -// i32 Two = 2; -// } -//} -message EmptyMessage{} - -message OptimizedLogicalPlanType { - string optimizer_name = 1; -} - -message OptimizedPhysicalPlanType { - string optimizer_name = 1; -} - -message PlanType { - oneof plan_type_enum { - EmptyMessage InitialLogicalPlan = 1; - OptimizedLogicalPlanType OptimizedLogicalPlan = 2; - EmptyMessage FinalLogicalPlan = 3; - EmptyMessage InitialPhysicalPlan = 4; - OptimizedPhysicalPlanType OptimizedPhysicalPlan = 5; - EmptyMessage FinalPhysicalPlan = 6; - } -} - -message StringifiedPlan { - PlanType plan_type = 1; - string plan = 2; -} diff --git a/ballista/rust/core/src/client.rs b/ballista/rust/core/src/client.rs deleted file mode 100644 index a5c4a062be8b..000000000000 --- a/ballista/rust/core/src/client.rs +++ /dev/null @@ -1,178 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Client API for sending requests to executors. - -use std::collections::HashMap; -use std::sync::Arc; - -use std::{ - convert::{TryFrom, TryInto}, - task::{Context, Poll}, -}; - -use crate::error::{ballista_error, BallistaError, Result}; -use crate::serde::protobuf::{self}; -use crate::serde::scheduler::Action; - -use arrow_flight::utils::flight_data_to_arrow_batch; -use arrow_flight::Ticket; -use arrow_flight::{flight_service_client::FlightServiceClient, FlightData}; -use datafusion::arrow::array::ArrayRef; -use datafusion::arrow::{ - datatypes::{Schema, SchemaRef}, - error::{ArrowError, Result as ArrowResult}, - record_batch::RecordBatch, -}; - -use datafusion::physical_plan::{RecordBatchStream, SendableRecordBatchStream}; -use futures::{Stream, StreamExt}; -use log::debug; -use prost::Message; -use tonic::Streaming; - -/// Client for interacting with Ballista executors. -#[derive(Clone)] -pub struct BallistaClient { - flight_client: FlightServiceClient, -} - -impl BallistaClient { - /// Create a new BallistaClient to connect to the executor listening on the specified - /// host and port - pub async fn try_new(host: &str, port: u16) -> Result { - let addr = format!("http://{}:{}", host, port); - debug!("BallistaClient connecting to {}", addr); - let flight_client = - FlightServiceClient::connect(addr.clone()) - .await - .map_err(|e| { - BallistaError::General(format!( - "Error connecting to Ballista scheduler or executor at {}: {:?}", - addr, e - )) - })?; - debug!("BallistaClient connected OK"); - - Ok(Self { flight_client }) - } - - /// Fetch a partition from an executor - pub async fn fetch_partition( - &mut self, - job_id: &str, - stage_id: usize, - partition_id: usize, - path: &str, - ) -> Result { - let action = Action::FetchPartition { - job_id: job_id.to_string(), - stage_id, - partition_id, - path: path.to_owned(), - }; - self.execute_action(&action).await - } - - /// Execute an action and retrieve the results - pub async fn execute_action( - &mut self, - action: &Action, - ) -> Result { - let serialized_action: protobuf::Action = action.to_owned().try_into()?; - - let mut buf: Vec = Vec::with_capacity(serialized_action.encoded_len()); - - serialized_action - .encode(&mut buf) - .map_err(|e| BallistaError::General(format!("{:?}", e)))?; - - let request = tonic::Request::new(Ticket { ticket: buf }); - - let mut stream = self - .flight_client - .do_get(request) - .await - .map_err(|e| BallistaError::General(format!("{:?}", e)))? - .into_inner(); - - // the schema should be the first message returned, else client should error - match stream - .message() - .await - .map_err(|e| BallistaError::General(format!("{:?}", e)))? - { - Some(flight_data) => { - // convert FlightData to a stream - let schema = Arc::new(Schema::try_from(&flight_data)?); - - // all the remaining stream messages should be dictionary and record batches - Ok(Box::pin(FlightDataStream::new(stream, schema))) - } - None => Err(ballista_error( - "Did not receive schema batch from flight server", - )), - } - } -} - -struct FlightDataStream { - stream: Streaming, - schema: SchemaRef, - dictionaries_by_id: HashMap, -} - -impl FlightDataStream { - pub fn new(stream: Streaming, schema: SchemaRef) -> Self { - Self { - stream, - schema, - dictionaries_by_id: HashMap::new(), - } - } -} - -impl Stream for FlightDataStream { - type Item = ArrowResult; - - fn poll_next( - mut self: std::pin::Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - self.stream.poll_next_unpin(cx).map(|x| match x { - Some(flight_data_chunk_result) => { - let converted_chunk = flight_data_chunk_result - .map_err(|e| ArrowError::from_external_error(Box::new(e))) - .and_then(|flight_data_chunk| { - flight_data_to_arrow_batch( - &flight_data_chunk, - self.schema.clone(), - &self.dictionaries_by_id, - ) - }); - Some(converted_chunk) - } - None => None, - }) - } -} - -impl RecordBatchStream for FlightDataStream { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} diff --git a/ballista/rust/core/src/config.rs b/ballista/rust/core/src/config.rs deleted file mode 100644 index 1ff02c16d069..000000000000 --- a/ballista/rust/core/src/config.rs +++ /dev/null @@ -1,323 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. -// - -//! Ballista configuration - -use clap::ArgEnum; -use core::fmt; -use std::collections::HashMap; -use std::result; - -use crate::error::{BallistaError, Result}; - -use datafusion::arrow::datatypes::DataType; - -pub const BALLISTA_DEFAULT_SHUFFLE_PARTITIONS: &str = "ballista.shuffle.partitions"; -pub const BALLISTA_DEFAULT_BATCH_SIZE: &str = "ballista.batch.size"; -pub const BALLISTA_REPARTITION_JOINS: &str = "ballista.repartition.joins"; -pub const BALLISTA_REPARTITION_AGGREGATIONS: &str = "ballista.repartition.aggregations"; -pub const BALLISTA_REPARTITION_WINDOWS: &str = "ballista.repartition.windows"; -pub const BALLISTA_PARQUET_PRUNING: &str = "ballista.parquet.pruning"; -pub const BALLISTA_WITH_INFORMATION_SCHEMA: &str = "ballista.with_information_schema"; -/// give a plugin files dir, and then the dynamic library files in this dir will be load when scheduler state init. -pub const BALLISTA_PLUGIN_DIR: &str = "ballista.plugin_dir"; - -pub type ParseResult = result::Result; - -/// Configuration option meta-data -#[derive(Debug, Clone)] -pub struct ConfigEntry { - name: String, - _description: String, - _data_type: DataType, - default_value: Option, -} - -impl ConfigEntry { - fn new( - name: String, - _description: String, - _data_type: DataType, - default_value: Option, - ) -> Self { - Self { - name, - _description, - _data_type, - default_value, - } - } -} - -/// Ballista configuration builder -pub struct BallistaConfigBuilder { - settings: HashMap, -} - -impl Default for BallistaConfigBuilder { - /// Create a new config builder - fn default() -> Self { - Self { - settings: HashMap::new(), - } - } -} - -impl BallistaConfigBuilder { - /// Create a new config with an additional setting - pub fn set(&self, k: &str, v: &str) -> Self { - let mut settings = self.settings.clone(); - settings.insert(k.to_owned(), v.to_owned()); - Self { settings } - } - - pub fn build(&self) -> Result { - BallistaConfig::with_settings(self.settings.clone()) - } -} - -/// Ballista configuration -#[derive(Debug, Clone)] -pub struct BallistaConfig { - /// Settings stored in map for easy serde - settings: HashMap, -} - -impl BallistaConfig { - /// Create a default configuration - pub fn new() -> Result { - Self::with_settings(HashMap::new()) - } - - /// Create a configuration builder - pub fn builder() -> BallistaConfigBuilder { - BallistaConfigBuilder::default() - } - - /// Create a new configuration based on key-value pairs - pub fn with_settings(settings: HashMap) -> Result { - let supported_entries = BallistaConfig::valid_entries(); - for (name, entry) in &supported_entries { - if let Some(v) = settings.get(name) { - // validate that we can parse the user-supplied value - let _ = Self::parse_value(v.as_str(), entry._data_type.clone()).map_err(|e| BallistaError::General(format!("Failed to parse user-supplied value '{}' for configuration setting '{}': {}", name, v, e)))?; - } else if let Some(v) = entry.default_value.clone() { - let _ = Self::parse_value(v.as_str(), entry._data_type.clone()).map_err(|e| BallistaError::General(format!("Failed to parse default value '{}' for configuration setting '{}': {}", name, v, e)))?; - } else { - return Err(BallistaError::General(format!( - "No value specified for mandatory configuration setting '{}'", - name - ))); - } - } - - Ok(Self { settings }) - } - - pub fn parse_value(val: &str, data_type: DataType) -> ParseResult<()> { - match data_type { - DataType::UInt16 => { - val.to_string() - .parse::() - .map_err(|e| format!("{:?}", e))?; - } - DataType::Boolean => { - val.to_string() - .parse::() - .map_err(|e| format!("{:?}", e))?; - } - DataType::Utf8 => { - val.to_string(); - } - _ => { - return Err(format!("not support data type: {}", data_type)); - } - } - - Ok(()) - } - - /// All available configuration options - pub fn valid_entries() -> HashMap { - let entries = vec![ - ConfigEntry::new(BALLISTA_DEFAULT_SHUFFLE_PARTITIONS.to_string(), - "Sets the default number of partitions to create when repartitioning query stages".to_string(), - DataType::UInt16, Some("2".to_string())), - ConfigEntry::new(BALLISTA_DEFAULT_BATCH_SIZE.to_string(), - "Sets the default batch size".to_string(), - DataType::UInt16, Some("8192".to_string())), - ConfigEntry::new(BALLISTA_REPARTITION_JOINS.to_string(), - "Configuration for repartition joins".to_string(), - DataType::Boolean, Some("true".to_string())), - ConfigEntry::new(BALLISTA_REPARTITION_AGGREGATIONS.to_string(), - "Configuration for repartition aggregations".to_string(), - DataType::Boolean,Some("true".to_string())), - ConfigEntry::new(BALLISTA_REPARTITION_WINDOWS.to_string(), - "Configuration for repartition windows".to_string(), - DataType::Boolean,Some("true".to_string())), - ConfigEntry::new(BALLISTA_PARQUET_PRUNING.to_string(), - "Configuration for parquet prune".to_string(), - DataType::Boolean,Some("true".to_string())), - ConfigEntry::new(BALLISTA_WITH_INFORMATION_SCHEMA.to_string(), - "Sets whether enable information_schema".to_string(), - DataType::Boolean,Some("false".to_string())), - ConfigEntry::new(BALLISTA_PLUGIN_DIR.to_string(), - "Sets the plugin dir".to_string(), - DataType::Utf8,Some("".to_string())), - ]; - entries - .iter() - .map(|e| (e.name.clone(), e.clone())) - .collect::>() - } - - pub fn settings(&self) -> &HashMap { - &self.settings - } - - pub fn default_shuffle_partitions(&self) -> usize { - self.get_usize_setting(BALLISTA_DEFAULT_SHUFFLE_PARTITIONS) - } - - pub fn default_plugin_dir(&self) -> String { - self.get_string_setting(BALLISTA_PLUGIN_DIR) - } - - pub fn default_batch_size(&self) -> usize { - self.get_usize_setting(BALLISTA_DEFAULT_BATCH_SIZE) - } - - pub fn repartition_joins(&self) -> bool { - self.get_bool_setting(BALLISTA_REPARTITION_JOINS) - } - - pub fn repartition_aggregations(&self) -> bool { - self.get_bool_setting(BALLISTA_REPARTITION_AGGREGATIONS) - } - - pub fn repartition_windows(&self) -> bool { - self.get_bool_setting(BALLISTA_REPARTITION_WINDOWS) - } - - pub fn parquet_pruning(&self) -> bool { - self.get_bool_setting(BALLISTA_PARQUET_PRUNING) - } - - pub fn default_with_information_schema(&self) -> bool { - self.get_bool_setting(BALLISTA_WITH_INFORMATION_SCHEMA) - } - - fn get_usize_setting(&self, key: &str) -> usize { - if let Some(v) = self.settings.get(key) { - // infallible because we validate all configs in the constructor - v.parse().unwrap() - } else { - let entries = Self::valid_entries(); - // infallible because we validate all configs in the constructor - let v = entries.get(key).unwrap().default_value.as_ref().unwrap(); - v.parse().unwrap() - } - } - - fn get_bool_setting(&self, key: &str) -> bool { - if let Some(v) = self.settings.get(key) { - // infallible because we validate all configs in the constructor - v.parse::().unwrap() - } else { - let entries = Self::valid_entries(); - // infallible because we validate all configs in the constructor - let v = entries.get(key).unwrap().default_value.as_ref().unwrap(); - v.parse::().unwrap() - } - } - fn get_string_setting(&self, key: &str) -> String { - if let Some(v) = self.settings.get(key) { - // infallible because we validate all configs in the constructor - v.to_string() - } else { - let entries = Self::valid_entries(); - // infallible because we validate all configs in the constructor - let v = entries.get(key).unwrap().default_value.as_ref().unwrap(); - v.to_string() - } - } -} - -// an enum used to configure the scheduler policy -// needs to be visible to code generated by configure_me -#[derive(Clone, ArgEnum, Copy, Debug, serde::Deserialize)] -pub enum TaskSchedulingPolicy { - PullStaged, - PushStaged, -} - -impl std::str::FromStr for TaskSchedulingPolicy { - type Err = String; - - fn from_str(s: &str) -> std::result::Result { - ArgEnum::from_str(s, true) - } -} - -impl parse_arg::ParseArgFromStr for TaskSchedulingPolicy { - fn describe_type(mut writer: W) -> fmt::Result { - write!(writer, "The scheduler policy for the scheduler") - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn default_config() -> Result<()> { - let config = BallistaConfig::new()?; - assert_eq!(2, config.default_shuffle_partitions()); - assert!(!config.default_with_information_schema()); - assert_eq!("", config.default_plugin_dir().as_str()); - Ok(()) - } - - #[test] - fn custom_config() -> Result<()> { - let config = BallistaConfig::builder() - .set(BALLISTA_DEFAULT_SHUFFLE_PARTITIONS, "123") - .set(BALLISTA_WITH_INFORMATION_SCHEMA, "true") - .build()?; - assert_eq!(123, config.default_shuffle_partitions()); - assert!(config.default_with_information_schema()); - Ok(()) - } - - #[test] - fn custom_config_invalid() -> Result<()> { - let config = BallistaConfig::builder() - .set(BALLISTA_DEFAULT_SHUFFLE_PARTITIONS, "true") - .set(BALLISTA_PLUGIN_DIR, "test_dir") - .build(); - assert!(config.is_err()); - assert_eq!("General(\"Failed to parse user-supplied value 'ballista.shuffle.partitions' for configuration setting 'true': ParseIntError { kind: InvalidDigit }\")", format!("{:?}", config.unwrap_err())); - - let config = BallistaConfig::builder() - .set(BALLISTA_WITH_INFORMATION_SCHEMA, "123") - .build(); - assert!(config.is_err()); - assert_eq!("General(\"Failed to parse user-supplied value 'ballista.with_information_schema' for configuration setting '123': ParseBoolError\")", format!("{:?}", config.unwrap_err())); - Ok(()) - } -} diff --git a/ballista/rust/core/src/error.rs b/ballista/rust/core/src/error.rs deleted file mode 100644 index ba6cdc74d698..000000000000 --- a/ballista/rust/core/src/error.rs +++ /dev/null @@ -1,185 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Ballista error types - -use std::{ - error::Error, - fmt::{Display, Formatter}, - io, result, -}; - -use datafusion::arrow::error::ArrowError; -use datafusion::error::DataFusionError; -use sqlparser::parser; - -pub type Result = result::Result; - -/// Ballista error -#[derive(Debug)] -pub enum BallistaError { - NotImplemented(String), - General(String), - Internal(String), - ArrowError(ArrowError), - DataFusionError(DataFusionError), - SqlError(parser::ParserError), - IoError(io::Error), - // ReqwestError(reqwest::Error), - // HttpError(http::Error), - // KubeAPIError(kube::error::Error), - // KubeAPIRequestError(k8s_openapi::RequestError), - // KubeAPIResponseError(k8s_openapi::ResponseError), - TonicError(tonic::transport::Error), - GrpcError(tonic::Status), - TokioError(tokio::task::JoinError), -} - -#[allow(clippy::from_over_into)] -impl Into> for BallistaError { - fn into(self) -> Result { - Err(self) - } -} - -pub fn ballista_error(message: &str) -> BallistaError { - BallistaError::General(message.to_owned()) -} - -impl From for BallistaError { - fn from(e: String) -> Self { - BallistaError::General(e) - } -} - -impl From for BallistaError { - fn from(e: ArrowError) -> Self { - BallistaError::ArrowError(e) - } -} - -impl From for BallistaError { - fn from(e: parser::ParserError) -> Self { - BallistaError::SqlError(e) - } -} - -impl From for BallistaError { - fn from(e: DataFusionError) -> Self { - BallistaError::DataFusionError(e) - } -} - -impl From for BallistaError { - fn from(e: io::Error) -> Self { - BallistaError::IoError(e) - } -} - -// impl From for BallistaError { -// fn from(e: reqwest::Error) -> Self { -// BallistaError::ReqwestError(e) -// } -// } -// -// impl From for BallistaError { -// fn from(e: http::Error) -> Self { -// BallistaError::HttpError(e) -// } -// } - -// impl From for BallistaError { -// fn from(e: kube::error::Error) -> Self { -// BallistaError::KubeAPIError(e) -// } -// } - -// impl From for BallistaError { -// fn from(e: k8s_openapi::RequestError) -> Self { -// BallistaError::KubeAPIRequestError(e) -// } -// } - -// impl From for BallistaError { -// fn from(e: k8s_openapi::ResponseError) -> Self { -// BallistaError::KubeAPIResponseError(e) -// } -// } - -impl From for BallistaError { - fn from(e: tonic::transport::Error) -> Self { - BallistaError::TonicError(e) - } -} - -impl From for BallistaError { - fn from(e: tonic::Status) -> Self { - BallistaError::GrpcError(e) - } -} - -impl From for BallistaError { - fn from(e: tokio::task::JoinError) -> Self { - BallistaError::TokioError(e) - } -} - -impl From for BallistaError { - fn from(e: datafusion_proto::from_proto::Error) -> Self { - BallistaError::General(e.to_string()) - } -} - -impl From for BallistaError { - fn from(e: datafusion_proto::to_proto::Error) -> Self { - BallistaError::General(e.to_string()) - } -} - -impl Display for BallistaError { - fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { - match self { - BallistaError::NotImplemented(ref desc) => { - write!(f, "Not implemented: {}", desc) - } - BallistaError::General(ref desc) => write!(f, "General error: {}", desc), - BallistaError::ArrowError(ref desc) => write!(f, "Arrow error: {}", desc), - BallistaError::DataFusionError(ref desc) => { - write!(f, "DataFusion error: {:?}", desc) - } - BallistaError::SqlError(ref desc) => write!(f, "SQL error: {:?}", desc), - BallistaError::IoError(ref desc) => write!(f, "IO error: {}", desc), - // BallistaError::ReqwestError(ref desc) => write!(f, "Reqwest error: {}", desc), - // BallistaError::HttpError(ref desc) => write!(f, "HTTP error: {}", desc), - // BallistaError::KubeAPIError(ref desc) => write!(f, "Kube API error: {}", desc), - // BallistaError::KubeAPIRequestError(ref desc) => { - // write!(f, "KubeAPI request error: {}", desc) - // } - // BallistaError::KubeAPIResponseError(ref desc) => { - // write!(f, "KubeAPI response error: {}", desc) - // } - BallistaError::TonicError(desc) => write!(f, "Tonic error: {}", desc), - BallistaError::GrpcError(desc) => write!(f, "Grpc error: {}", desc), - BallistaError::Internal(desc) => { - write!(f, "Internal Ballista error: {}", desc) - } - BallistaError::TokioError(desc) => write!(f, "Tokio join error: {}", desc), - } - } -} - -impl Error for BallistaError {} diff --git a/ballista/rust/core/src/event_loop.rs b/ballista/rust/core/src/event_loop.rs deleted file mode 100644 index 217b7fc30f5b..000000000000 --- a/ballista/rust/core/src/event_loop.rs +++ /dev/null @@ -1,141 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; - -use async_trait::async_trait; -use log::{error, info}; -use tokio::sync::mpsc; - -use crate::error::{BallistaError, Result}; - -#[async_trait] -pub trait EventAction: Send + Sync { - fn on_start(&self); - - fn on_stop(&self); - - async fn on_receive(&self, event: E) -> Result>; - - fn on_error(&self, error: BallistaError); -} - -#[derive(Clone)] -pub struct EventLoop { - pub name: String, - pub buffer_size: usize, - stopped: Arc, - action: Arc>, - tx_event: Option>, -} - -impl EventLoop { - pub fn new( - name: String, - buffer_size: usize, - action: Arc>, - ) -> Self { - Self { - name, - buffer_size, - stopped: Arc::new(AtomicBool::new(false)), - action, - tx_event: None, - } - } - - fn run(&self, mut rx_event: mpsc::Receiver) { - assert!( - self.tx_event.is_some(), - "The event sender should be initialized first!" - ); - let tx_event = self.tx_event.as_ref().unwrap().clone(); - let name = self.name.clone(); - let stopped = self.stopped.clone(); - let action = self.action.clone(); - tokio::spawn(async move { - info!("Starting the event loop {}", name); - while !stopped.load(Ordering::SeqCst) { - if let Some(event) = rx_event.recv().await { - match action.on_receive(event).await { - Ok(Some(event)) => { - if let Err(e) = tx_event.send(event).await { - let msg = format!("Fail to send event due to {}", e); - error!("{}", msg); - action.on_error(BallistaError::General(msg)); - } - } - Err(e) => { - error!("Fail to process event due to {}", e); - action.on_error(e); - } - _ => {} - } - } else { - info!("Event Channel closed, shutting down"); - break; - } - } - info!("The event loop {} has been stopped", name); - }); - } - - pub fn start(&mut self) -> Result<()> { - if self.stopped.load(Ordering::SeqCst) { - return Err(BallistaError::General(format!( - "{} has already been stopped", - self.name - ))); - } - self.action.on_start(); - - let (tx_event, rx_event) = mpsc::channel::(self.buffer_size); - self.tx_event = Some(tx_event); - self.run(rx_event); - - Ok(()) - } - - pub fn stop(&self) { - if !self.stopped.swap(true, Ordering::SeqCst) { - self.action.on_stop(); - } else { - // Keep quiet to allow calling `stop` multiple times. - } - } - - pub fn get_sender(&self) -> Result> { - Ok(EventSender { - tx_event: self.tx_event.as_ref().cloned().ok_or_else(|| { - BallistaError::General("Event sender not exist!!!".to_string()) - })?, - }) - } -} - -pub struct EventSender { - tx_event: mpsc::Sender, -} - -impl EventSender { - pub async fn post_event(&self, event: E) -> Result<()> { - self.tx_event.send(event).await.map_err(|e| { - BallistaError::General(format!("Fail to send event due to {}", e)) - }) - } -} diff --git a/ballista/rust/core/src/execution_plans/distributed_query.rs b/ballista/rust/core/src/execution_plans/distributed_query.rs deleted file mode 100644 index 0b20acb47969..000000000000 --- a/ballista/rust/core/src/execution_plans/distributed_query.rs +++ /dev/null @@ -1,329 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::any::Any; - -use std::fmt::Debug; -use std::marker::PhantomData; - -use std::sync::Arc; -use std::time::Duration; - -use crate::client::BallistaClient; -use crate::config::BallistaConfig; -use crate::serde::protobuf::{ - execute_query_params::Query, job_status, scheduler_grpc_client::SchedulerGrpcClient, - ExecuteQueryParams, GetJobStatusParams, GetJobStatusResult, KeyValuePair, - PartitionLocation, -}; - -use datafusion::arrow::datatypes::SchemaRef; -use datafusion::error::{DataFusionError, Result}; -use datafusion::logical_plan::LogicalPlan; -use datafusion::physical_plan::expressions::PhysicalSortExpr; -use datafusion::physical_plan::{ - DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics, -}; - -use crate::serde::protobuf::execute_query_params::OptionalSessionId; -use crate::serde::{AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec}; -use datafusion::arrow::error::{ArrowError, Result as ArrowResult}; -use datafusion::arrow::record_batch::RecordBatch; -use datafusion::execution::context::TaskContext; -use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use futures::{Stream, StreamExt, TryFutureExt, TryStreamExt}; -use log::{error, info}; - -/// This operator sends a logical plan to a Ballista scheduler for execution and -/// polls the scheduler until the query is complete and then fetches the resulting -/// batches directly from the executors that hold the results from the final -/// query stage. -#[derive(Debug, Clone)] -pub struct DistributedQueryExec { - /// Ballista scheduler URL - scheduler_url: String, - /// Ballista configuration - config: BallistaConfig, - /// Logical plan to execute - plan: LogicalPlan, - /// Codec for LogicalPlan extensions - extension_codec: Arc, - /// Phantom data for serializable plan message - plan_repr: PhantomData, - /// Session id - session_id: String, -} - -impl DistributedQueryExec { - pub fn new( - scheduler_url: String, - config: BallistaConfig, - plan: LogicalPlan, - session_id: String, - ) -> Self { - Self { - scheduler_url, - config, - plan, - extension_codec: Arc::new(DefaultLogicalExtensionCodec {}), - plan_repr: PhantomData, - session_id, - } - } - - pub fn with_extension( - scheduler_url: String, - config: BallistaConfig, - plan: LogicalPlan, - extension_codec: Arc, - session_id: String, - ) -> Self { - Self { - scheduler_url, - config, - plan, - extension_codec, - plan_repr: PhantomData, - session_id, - } - } - - pub fn with_repr( - scheduler_url: String, - config: BallistaConfig, - plan: LogicalPlan, - extension_codec: Arc, - plan_repr: PhantomData, - session_id: String, - ) -> Self { - Self { - scheduler_url, - config, - plan, - extension_codec, - plan_repr, - session_id, - } - } -} - -impl ExecutionPlan for DistributedQueryExec { - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - self.plan.schema().as_ref().clone().into() - } - - fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(1) - } - - fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { - None - } - - fn relies_on_input_order(&self) -> bool { - false - } - - fn children(&self) -> Vec> { - vec![] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> datafusion::error::Result> { - Ok(Arc::new(DistributedQueryExec { - scheduler_url: self.scheduler_url.clone(), - config: self.config.clone(), - plan: self.plan.clone(), - extension_codec: self.extension_codec.clone(), - plan_repr: self.plan_repr, - session_id: self.session_id.clone(), - })) - } - - fn execute( - &self, - partition: usize, - _context: Arc, - ) -> Result { - assert_eq!(0, partition); - - let mut buf: Vec = vec![]; - let plan_message = - T::try_from_logical_plan(&self.plan, self.extension_codec.as_ref()).map_err( - |e| { - DataFusionError::Internal(format!( - "failed to serialize logical plan: {:?}", - e - )) - }, - )?; - plan_message.try_encode(&mut buf).map_err(|e| { - DataFusionError::Execution(format!("failed to encode logical plan: {:?}", e)) - })?; - - let query = ExecuteQueryParams { - query: Some(Query::LogicalPlan(buf)), - settings: self - .config - .settings() - .iter() - .map(|(k, v)| KeyValuePair { - key: k.to_owned(), - value: v.to_owned(), - }) - .collect::>(), - optional_session_id: Some(OptionalSessionId::SessionId( - self.session_id.clone(), - )), - }; - - let stream = futures::stream::once( - execute_query(self.scheduler_url.clone(), self.session_id.clone(), query) - .map_err(|e| ArrowError::ExternalError(Box::new(e))), - ) - .try_flatten(); - - let schema = self.schema(); - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) - } - - fn fmt_as( - &self, - t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - match t { - DisplayFormatType::Default => { - write!( - f, - "DistributedQueryExec: scheduler_url={}", - self.scheduler_url - ) - } - } - } - - fn statistics(&self) -> Statistics { - // This execution plan sends the logical plan to the scheduler without - // performing the node by node conversion to a full physical plan. - // This implies that we cannot infer the statistics at this stage. - Statistics::default() - } -} - -async fn execute_query( - scheduler_url: String, - session_id: String, - query: ExecuteQueryParams, -) -> Result> + Send> { - info!("Connecting to Ballista scheduler at {}", scheduler_url); - // TODO reuse the scheduler to avoid connecting to the Ballista scheduler again and again - - let mut scheduler = SchedulerGrpcClient::connect(scheduler_url.clone()) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; - - let query_result = scheduler - .execute_query(query) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))? - .into_inner(); - - assert_eq!( - session_id, query_result.session_id, - "Session id inconsistent between Client and Server side in DistributedQueryExec." - ); - - let job_id = query_result.job_id; - let mut prev_status: Option = None; - - loop { - let GetJobStatusResult { status } = scheduler - .get_job_status(GetJobStatusParams { - job_id: job_id.clone(), - }) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))? - .into_inner(); - let status = status.and_then(|s| s.status).ok_or_else(|| { - DataFusionError::Internal("Received empty status message".to_owned()) - })?; - let wait_future = tokio::time::sleep(Duration::from_millis(100)); - let has_status_change = prev_status.map(|x| x != status).unwrap_or(true); - match status { - job_status::Status::Queued(_) => { - if has_status_change { - info!("Job {} still queued...", job_id); - } - wait_future.await; - prev_status = Some(status); - } - job_status::Status::Running(_) => { - if has_status_change { - info!("Job {} is running...", job_id); - } - wait_future.await; - prev_status = Some(status); - } - job_status::Status::Failed(err) => { - let msg = format!("Job {} failed: {}", job_id, err.error); - error!("{}", msg); - break Err(DataFusionError::Execution(msg)); - } - job_status::Status::Completed(completed) => { - let streams = completed.partition_location.into_iter().map(|p| { - let f = fetch_partition(p) - .map_err(|e| ArrowError::ExternalError(Box::new(e))); - - futures::stream::once(f).try_flatten() - }); - - break Ok(futures::stream::iter(streams).flatten()); - } - }; - } -} - -async fn fetch_partition( - location: PartitionLocation, -) -> Result { - let metadata = location.executor_meta.ok_or_else(|| { - DataFusionError::Internal("Received empty executor metadata".to_owned()) - })?; - let partition_id = location.partition_id.ok_or_else(|| { - DataFusionError::Internal("Received empty partition id".to_owned()) - })?; - let mut ballista_client = - BallistaClient::try_new(metadata.host.as_str(), metadata.port as u16) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; - ballista_client - .fetch_partition( - &partition_id.job_id, - partition_id.stage_id as usize, - partition_id.partition_id as usize, - &location.path, - ) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e))) -} diff --git a/ballista/rust/core/src/execution_plans/mod.rs b/ballista/rust/core/src/execution_plans/mod.rs deleted file mode 100644 index 7a5e105c6c4a..000000000000 --- a/ballista/rust/core/src/execution_plans/mod.rs +++ /dev/null @@ -1,29 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! This module contains execution plans that are needed to distribute DataFusion's execution plans into -//! several Ballista executors. - -mod distributed_query; -mod shuffle_reader; -mod shuffle_writer; -mod unresolved_shuffle; - -pub use distributed_query::DistributedQueryExec; -pub use shuffle_reader::ShuffleReaderExec; -pub use shuffle_writer::ShuffleWriterExec; -pub use unresolved_shuffle::UnresolvedShuffleExec; diff --git a/ballista/rust/core/src/execution_plans/shuffle_reader.rs b/ballista/rust/core/src/execution_plans/shuffle_reader.rs deleted file mode 100644 index 3046a2276f89..000000000000 --- a/ballista/rust/core/src/execution_plans/shuffle_reader.rs +++ /dev/null @@ -1,294 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::any::Any; -use std::sync::Arc; - -use crate::client::BallistaClient; -use crate::serde::scheduler::{PartitionLocation, PartitionStats}; - -use datafusion::arrow::datatypes::SchemaRef; - -use datafusion::error::{DataFusionError, Result}; -use datafusion::physical_plan::expressions::PhysicalSortExpr; -use datafusion::physical_plan::metrics::{ - ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, -}; -use datafusion::physical_plan::{ - DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics, -}; -use futures::{StreamExt, TryStreamExt}; - -use datafusion::arrow::error::ArrowError; -use datafusion::execution::context::TaskContext; -use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use log::info; - -/// ShuffleReaderExec reads partitions that have already been materialized by a ShuffleWriterExec -/// being executed by an executor -#[derive(Debug, Clone)] -pub struct ShuffleReaderExec { - /// Each partition of a shuffle can read data from multiple locations - pub(crate) partition: Vec>, - pub(crate) schema: SchemaRef, - /// Execution metrics - metrics: ExecutionPlanMetricsSet, -} - -impl ShuffleReaderExec { - /// Create a new ShuffleReaderExec - pub fn try_new( - partition: Vec>, - schema: SchemaRef, - ) -> Result { - Ok(Self { - partition, - schema, - metrics: ExecutionPlanMetricsSet::new(), - }) - } -} - -impl ExecutionPlan for ShuffleReaderExec { - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - - fn output_partitioning(&self) -> Partitioning { - // TODO partitioning may be known and could be populated here - // see https://github.com/apache/arrow-datafusion/issues/758 - Partitioning::UnknownPartitioning(self.partition.len()) - } - - fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { - None - } - - fn relies_on_input_order(&self) -> bool { - false - } - - fn children(&self) -> Vec> { - vec![] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - Err(DataFusionError::Plan( - "Ballista ShuffleReaderExec does not support with_new_children()".to_owned(), - )) - } - - fn execute( - &self, - partition: usize, - _context: Arc, - ) -> Result { - info!("ShuffleReaderExec::execute({})", partition); - - let fetch_time = - MetricBuilder::new(&self.metrics).subset_time("fetch_time", partition); - - let locations = self.partition[partition].clone(); - let stream = locations.into_iter().map(move |p| { - let fetch_time = fetch_time.clone(); - futures::stream::once(async move { - let timer = fetch_time.timer(); - let r = fetch_partition(&p).await; - timer.done(); - - r.map_err(|e| ArrowError::ExternalError(Box::new(e))) - }) - .try_flatten() - }); - - let result = RecordBatchStreamAdapter::new( - Arc::new(self.schema.as_ref().clone()), - futures::stream::iter(stream).flatten(), - ); - Ok(Box::pin(result)) - } - - fn fmt_as( - &self, - t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - match t { - DisplayFormatType::Default => { - let loc_str = self - .partition - .iter() - .enumerate() - .map(|(partition_id, locations)| { - format!( - "[partition={} paths={}]", - partition_id, - locations - .iter() - .map(|l| l.path.clone()) - .collect::>() - .join(",") - ) - }) - .collect::>() - .join(", "); - write!( - f, - "ShuffleReaderExec: partition_locations({})={}", - self.partition.len(), - loc_str - ) - } - } - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) - } - - fn statistics(&self) -> Statistics { - stats_for_partitions( - self.partition - .iter() - .flatten() - .map(|loc| loc.partition_stats), - ) - } -} - -fn stats_for_partitions( - partition_stats: impl Iterator, -) -> Statistics { - // TODO stats: add column statistics to PartitionStats - partition_stats.fold( - Statistics { - is_exact: true, - num_rows: Some(0), - total_byte_size: Some(0), - column_statistics: None, - }, - |mut acc, part| { - // if any statistic is unkown it makes the entire statistic unkown - acc.num_rows = acc.num_rows.zip(part.num_rows).map(|(a, b)| a + b as usize); - acc.total_byte_size = acc - .total_byte_size - .zip(part.num_bytes) - .map(|(a, b)| a + b as usize); - acc - }, - ) -} - -async fn fetch_partition( - location: &PartitionLocation, -) -> Result { - let metadata = &location.executor_meta; - let partition_id = &location.partition_id; - let mut ballista_client = - BallistaClient::try_new(metadata.host.as_str(), metadata.port as u16) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; - ballista_client - .fetch_partition( - &partition_id.job_id, - partition_id.stage_id as usize, - partition_id.partition_id as usize, - &location.path, - ) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e))) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn test_stats_for_partitions_empty() { - let result = stats_for_partitions(std::iter::empty()); - - let exptected = Statistics { - is_exact: true, - num_rows: Some(0), - total_byte_size: Some(0), - column_statistics: None, - }; - - assert_eq!(result, exptected); - } - - #[tokio::test] - async fn test_stats_for_partitions_full() { - let part_stats = vec![ - PartitionStats { - num_rows: Some(10), - num_bytes: Some(84), - num_batches: Some(1), - }, - PartitionStats { - num_rows: Some(4), - num_bytes: Some(65), - num_batches: None, - }, - ]; - - let result = stats_for_partitions(part_stats.into_iter()); - - let exptected = Statistics { - is_exact: true, - num_rows: Some(14), - total_byte_size: Some(149), - column_statistics: None, - }; - - assert_eq!(result, exptected); - } - - #[tokio::test] - async fn test_stats_for_partitions_missing() { - let part_stats = vec![ - PartitionStats { - num_rows: Some(10), - num_bytes: Some(84), - num_batches: Some(1), - }, - PartitionStats { - num_rows: None, - num_bytes: None, - num_batches: None, - }, - ]; - - let result = stats_for_partitions(part_stats.into_iter()); - - let exptected = Statistics { - is_exact: true, - num_rows: None, - total_byte_size: None, - column_statistics: None, - }; - - assert_eq!(result, exptected); - } -} diff --git a/ballista/rust/core/src/execution_plans/shuffle_writer.rs b/ballista/rust/core/src/execution_plans/shuffle_writer.rs deleted file mode 100644 index 45a102185c44..000000000000 --- a/ballista/rust/core/src/execution_plans/shuffle_writer.rs +++ /dev/null @@ -1,559 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! ShuffleWriterExec represents a section of a query plan that has consistent partitioning and -//! can be executed as one unit with each partition being executed in parallel. The output of each -//! partition is re-partitioned and streamed to disk in Arrow IPC format. Future stages of the query -//! will use the ShuffleReaderExec to read these results. - -use datafusion::physical_plan::expressions::PhysicalSortExpr; - -use std::any::Any; -use std::future::Future; -use std::iter::Iterator; -use std::path::PathBuf; -use std::sync::Arc; -use std::time::Instant; - -use crate::utils; - -use crate::serde::protobuf::ShuffleWritePartition; -use crate::serde::scheduler::PartitionStats; -use datafusion::arrow::array::{ - ArrayBuilder, ArrayRef, StringBuilder, StructBuilder, UInt32Builder, UInt64Builder, -}; -use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; - -use datafusion::arrow::record_batch::RecordBatch; -use datafusion::error::{DataFusionError, Result}; -use datafusion::physical_plan::common::IPCWriter; -use datafusion::physical_plan::memory::MemoryStream; -use datafusion::physical_plan::metrics::{ - self, ExecutionPlanMetricsSet, MetricBuilder, MetricsSet, -}; - -use datafusion::physical_plan::{ - DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics, -}; -use futures::{StreamExt, TryFutureExt, TryStreamExt}; - -use datafusion::arrow::error::ArrowError; -use datafusion::execution::context::TaskContext; -use datafusion::physical_plan::repartition::BatchPartitioner; -use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use log::{debug, info}; - -/// ShuffleWriterExec represents a section of a query plan that has consistent partitioning and -/// can be executed as one unit with each partition being executed in parallel. The output of each -/// partition is re-partitioned and streamed to disk in Arrow IPC format. Future stages of the query -/// will use the ShuffleReaderExec to read these results. -#[derive(Debug, Clone)] -pub struct ShuffleWriterExec { - /// Unique ID for the job (query) that this stage is a part of - job_id: String, - /// Unique query stage ID within the job - stage_id: usize, - /// Physical execution plan for this query stage - plan: Arc, - /// Path to write output streams to - work_dir: String, - /// Optional shuffle output partitioning - shuffle_output_partitioning: Option, - /// Execution metrics - metrics: ExecutionPlanMetricsSet, -} - -#[derive(Debug, Clone)] -struct ShuffleWriteMetrics { - /// Time spend writing batches to shuffle files - write_time: metrics::Time, - repart_time: metrics::Time, - input_rows: metrics::Count, - output_rows: metrics::Count, -} - -impl ShuffleWriteMetrics { - fn new(partition: usize, metrics: &ExecutionPlanMetricsSet) -> Self { - let write_time = MetricBuilder::new(metrics).subset_time("write_time", partition); - let repart_time = - MetricBuilder::new(metrics).subset_time("repart_time", partition); - - let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition); - - let output_rows = MetricBuilder::new(metrics).output_rows(partition); - - Self { - write_time, - repart_time, - input_rows, - output_rows, - } - } -} - -impl ShuffleWriterExec { - /// Create a new shuffle writer - pub fn try_new( - job_id: String, - stage_id: usize, - plan: Arc, - work_dir: String, - shuffle_output_partitioning: Option, - ) -> Result { - Ok(Self { - job_id, - stage_id, - plan, - work_dir, - shuffle_output_partitioning, - metrics: ExecutionPlanMetricsSet::new(), - }) - } - - /// Get the Job ID for this query stage - pub fn job_id(&self) -> &str { - &self.job_id - } - - /// Get the Stage ID for this query stage - pub fn stage_id(&self) -> usize { - self.stage_id - } - - /// Get the true output partitioning - pub fn shuffle_output_partitioning(&self) -> Option<&Partitioning> { - self.shuffle_output_partitioning.as_ref() - } - - pub fn execute_shuffle_write( - &self, - input_partition: usize, - context: Arc, - ) -> impl Future>> { - let mut path = PathBuf::from(&self.work_dir); - path.push(&self.job_id); - path.push(&format!("{}", self.stage_id)); - - let write_metrics = ShuffleWriteMetrics::new(input_partition, &self.metrics); - let output_partitioning = self.shuffle_output_partitioning.clone(); - let plan = self.plan.clone(); - - async move { - let now = Instant::now(); - let mut stream = plan.execute(input_partition, context)?; - - match output_partitioning { - None => { - let timer = write_metrics.write_time.timer(); - path.push(&format!("{}", input_partition)); - std::fs::create_dir_all(&path)?; - path.push("data.arrow"); - let path = path.to_str().unwrap(); - info!("Writing results to {}", path); - - // stream results to disk - let stats = utils::write_stream_to_disk( - &mut stream, - path, - &write_metrics.write_time, - ) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; - - write_metrics - .input_rows - .add(stats.num_rows.unwrap_or(0) as usize); - write_metrics - .output_rows - .add(stats.num_rows.unwrap_or(0) as usize); - timer.done(); - - info!( - "Executed partition {} in {} seconds. Statistics: {}", - input_partition, - now.elapsed().as_secs(), - stats - ); - - Ok(vec![ShuffleWritePartition { - partition_id: input_partition as u64, - path: path.to_owned(), - num_batches: stats.num_batches.unwrap_or(0), - num_rows: stats.num_rows.unwrap_or(0), - num_bytes: stats.num_bytes.unwrap_or(0), - }]) - } - - Some(Partitioning::Hash(exprs, num_output_partitions)) => { - // we won't necessary produce output for every possible partition, so we - // create writers on demand - let mut writers: Vec> = vec![]; - for _ in 0..num_output_partitions { - writers.push(None); - } - - let mut partitioner = BatchPartitioner::try_new( - Partitioning::Hash(exprs, num_output_partitions), - write_metrics.repart_time.clone(), - )?; - - while let Some(result) = stream.next().await { - let input_batch = result?; - - write_metrics.input_rows.add(input_batch.num_rows()); - - partitioner.partition( - input_batch, - |output_partition, output_batch| { - // write non-empty batch out - - // TODO optimize so we don't write or fetch empty partitions - // if output_batch.num_rows() > 0 { - let timer = write_metrics.write_time.timer(); - match &mut writers[output_partition] { - Some(w) => { - w.write(&output_batch)?; - } - None => { - let mut path = path.clone(); - path.push(&format!("{}", output_partition)); - std::fs::create_dir_all(&path)?; - - path.push(format!( - "data-{}.arrow", - input_partition - )); - info!("Writing results to {:?}", path); - - let mut writer = IPCWriter::new( - &path, - stream.schema().as_ref(), - )?; - - writer.write(&output_batch)?; - writers[output_partition] = Some(writer); - } - } - write_metrics.output_rows.add(output_batch.num_rows()); - timer.done(); - Ok(()) - }, - )?; - } - - let mut part_locs = vec![]; - - for (i, w) in writers.iter_mut().enumerate() { - match w { - Some(w) => { - w.finish()?; - info!( - "Finished writing shuffle partition {} at {:?}. Batches: {}. Rows: {}. Bytes: {}.", - i, - w.path(), - w.num_batches, - w.num_rows, - w.num_bytes - ); - - part_locs.push(ShuffleWritePartition { - partition_id: i as u64, - path: w.path().to_string_lossy().to_string(), - num_batches: w.num_batches, - num_rows: w.num_rows, - num_bytes: w.num_bytes, - }); - } - None => {} - } - } - Ok(part_locs) - } - - _ => Err(DataFusionError::Execution( - "Invalid shuffle partitioning scheme".to_owned(), - )), - } - } - } -} - -impl ExecutionPlan for ShuffleWriterExec { - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - self.plan.schema() - } - - fn output_partitioning(&self) -> Partitioning { - // This operator needs to be executed once for each *input* partition and there - // isn't really a mechanism yet in DataFusion to support this use case so we report - // the input partitioning as the output partitioning here. The executor reports - // output partition meta data back to the scheduler. - self.plan.output_partitioning() - } - - fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { - None - } - - fn relies_on_input_order(&self) -> bool { - false - } - - fn children(&self) -> Vec> { - vec![self.plan.clone()] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - Ok(Arc::new(ShuffleWriterExec::try_new( - self.job_id.clone(), - self.stage_id, - children[0].clone(), - self.work_dir.clone(), - self.shuffle_output_partitioning.clone(), - )?)) - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - let schema = result_schema(); - - let schema_captured = schema.clone(); - let fut_stream = self - .execute_shuffle_write(partition, context) - .and_then(|part_loc| async move { - // build metadata result batch - let num_writers = part_loc.len(); - let mut partition_builder = UInt32Builder::new(num_writers); - let mut path_builder = StringBuilder::new(num_writers); - let mut num_rows_builder = UInt64Builder::new(num_writers); - let mut num_batches_builder = UInt64Builder::new(num_writers); - let mut num_bytes_builder = UInt64Builder::new(num_writers); - - for loc in &part_loc { - path_builder.append_value(loc.path.clone())?; - partition_builder.append_value(loc.partition_id as u32)?; - num_rows_builder.append_value(loc.num_rows)?; - num_batches_builder.append_value(loc.num_batches)?; - num_bytes_builder.append_value(loc.num_bytes)?; - } - - // build arrays - let partition_num: ArrayRef = Arc::new(partition_builder.finish()); - let path: ArrayRef = Arc::new(path_builder.finish()); - let field_builders: Vec> = vec![ - Box::new(num_rows_builder), - Box::new(num_batches_builder), - Box::new(num_bytes_builder), - ]; - let mut stats_builder = StructBuilder::new( - PartitionStats::default().arrow_struct_fields(), - field_builders, - ); - for _ in 0..num_writers { - stats_builder.append(true)?; - } - let stats = Arc::new(stats_builder.finish()); - - // build result batch containing metadata - let batch = RecordBatch::try_new( - schema_captured.clone(), - vec![partition_num, path, stats], - )?; - - debug!("RESULTS METADATA:\n{:?}", batch); - - MemoryStream::try_new(vec![batch], schema_captured, None) - }) - .map_err(|e| ArrowError::ExternalError(Box::new(e))); - - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::once(fut_stream).try_flatten(), - ))) - } - - fn metrics(&self) -> Option { - Some(self.metrics.clone_inner()) - } - - fn fmt_as( - &self, - t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - match t { - DisplayFormatType::Default => { - write!( - f, - "ShuffleWriterExec: {:?}", - self.shuffle_output_partitioning - ) - } - } - } - - fn statistics(&self) -> Statistics { - self.plan.statistics() - } -} - -fn result_schema() -> SchemaRef { - let stats = PartitionStats::default(); - Arc::new(Schema::new(vec![ - Field::new("partition", DataType::UInt32, false), - Field::new("path", DataType::Utf8, false), - stats.arrow_struct_repr(), - ])) -} - -#[cfg(test)] -mod tests { - use super::*; - use datafusion::arrow::array::{StringArray, StructArray, UInt32Array, UInt64Array}; - use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; - use datafusion::physical_plan::expressions::Column; - - use datafusion::physical_plan::memory::MemoryExec; - use datafusion::prelude::SessionContext; - use tempfile::TempDir; - - #[tokio::test] - // number of rows in each partition is a function of the hash output, so don't test here - #[cfg(not(feature = "force_hash_collisions"))] - async fn test() -> Result<()> { - let session_ctx = SessionContext::new(); - let task_ctx = session_ctx.task_ctx(); - - let input_plan = Arc::new(CoalescePartitionsExec::new(create_input_plan()?)); - let work_dir = TempDir::new()?; - let query_stage = ShuffleWriterExec::try_new( - "jobOne".to_owned(), - 1, - input_plan, - work_dir.into_path().to_str().unwrap().to_owned(), - Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)), - )?; - let mut stream = query_stage.execute(0, task_ctx)?; - let batches = utils::collect_stream(&mut stream) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; - assert_eq!(1, batches.len()); - let batch = &batches[0]; - assert_eq!(3, batch.num_columns()); - assert_eq!(2, batch.num_rows()); - let path = batch.columns()[1] - .as_any() - .downcast_ref::() - .unwrap(); - - let file0 = path.value(0); - assert!( - file0.ends_with("/jobOne/1/0/data-0.arrow") - || file0.ends_with("\\jobOne\\1\\0\\data-0.arrow") - ); - let file1 = path.value(1); - assert!( - file1.ends_with("/jobOne/1/1/data-0.arrow") - || file1.ends_with("\\jobOne\\1\\1\\data-0.arrow") - ); - - let stats = batch.columns()[2] - .as_any() - .downcast_ref::() - .unwrap(); - - let num_rows = stats - .column_by_name("num_rows") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(4, num_rows.value(0)); - assert_eq!(4, num_rows.value(1)); - - Ok(()) - } - - #[tokio::test] - // number of rows in each partition is a function of the hash output, so don't test here - #[cfg(not(feature = "force_hash_collisions"))] - async fn test_partitioned() -> Result<()> { - let session_ctx = SessionContext::new(); - let task_ctx = session_ctx.task_ctx(); - - let input_plan = create_input_plan()?; - let work_dir = TempDir::new()?; - let query_stage = ShuffleWriterExec::try_new( - "jobOne".to_owned(), - 1, - input_plan, - work_dir.into_path().to_str().unwrap().to_owned(), - Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 2)), - )?; - let mut stream = query_stage.execute(0, task_ctx)?; - let batches = utils::collect_stream(&mut stream) - .await - .map_err(|e| DataFusionError::Execution(format!("{:?}", e)))?; - assert_eq!(1, batches.len()); - let batch = &batches[0]; - assert_eq!(3, batch.num_columns()); - assert_eq!(2, batch.num_rows()); - let stats = batch.columns()[2] - .as_any() - .downcast_ref::() - .unwrap(); - let num_rows = stats - .column_by_name("num_rows") - .unwrap() - .as_any() - .downcast_ref::() - .unwrap(); - assert_eq!(2, num_rows.value(0)); - assert_eq!(2, num_rows.value(1)); - - Ok(()) - } - - fn create_input_plan() -> Result> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::UInt32, true), - Field::new("b", DataType::Utf8, true), - ])); - - // define data. - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(UInt32Array::from(vec![Some(1), Some(2)])), - Arc::new(StringArray::from(vec![Some("hello"), Some("world")])), - ], - )?; - let partition = vec![batch.clone(), batch]; - let partitions = vec![partition.clone(), partition]; - Ok(Arc::new(MemoryExec::try_new(&partitions, schema, None)?)) - } -} diff --git a/ballista/rust/core/src/execution_plans/unresolved_shuffle.rs b/ballista/rust/core/src/execution_plans/unresolved_shuffle.rs deleted file mode 100644 index 15d403fb6a45..000000000000 --- a/ballista/rust/core/src/execution_plans/unresolved_shuffle.rs +++ /dev/null @@ -1,129 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::any::Any; -use std::sync::Arc; - -use datafusion::arrow::datatypes::SchemaRef; -use datafusion::error::{DataFusionError, Result}; -use datafusion::execution::context::TaskContext; -use datafusion::physical_plan::expressions::PhysicalSortExpr; -use datafusion::physical_plan::{ - DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics, -}; - -/// UnresolvedShuffleExec represents a dependency on the results of a ShuffleWriterExec node which hasn't computed yet. -/// -/// An ExecutionPlan that contains an UnresolvedShuffleExec isn't ready for execution. The presence of this ExecutionPlan -/// is used as a signal so the scheduler knows it can't start computation until the dependent shuffle has completed. -#[derive(Debug, Clone)] -pub struct UnresolvedShuffleExec { - // The query stage ids which needs to be computed - pub stage_id: usize, - - // The schema this node will have once it is replaced with a ShuffleReaderExec - pub schema: SchemaRef, - - // The number of shuffle writer partition tasks that will produce the partitions - pub input_partition_count: usize, - - // The partition count this node will have once it is replaced with a ShuffleReaderExec - pub output_partition_count: usize, -} - -impl UnresolvedShuffleExec { - /// Create a new UnresolvedShuffleExec - pub fn new( - stage_id: usize, - schema: SchemaRef, - input_partition_count: usize, - output_partition_count: usize, - ) -> Self { - Self { - stage_id, - schema, - input_partition_count, - output_partition_count, - } - } -} - -impl ExecutionPlan for UnresolvedShuffleExec { - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - self.schema.clone() - } - - fn output_partitioning(&self) -> Partitioning { - // TODO the output partition is known and should be populated here! - // see https://github.com/apache/arrow-datafusion/issues/758 - Partitioning::UnknownPartitioning(self.output_partition_count) - } - - fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { - None - } - - fn relies_on_input_order(&self) -> bool { - false - } - - fn children(&self) -> Vec> { - vec![] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - Err(DataFusionError::Plan( - "Ballista UnresolvedShuffleExec does not support with_new_children()" - .to_owned(), - )) - } - - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> Result { - Err(DataFusionError::Plan( - "Ballista UnresolvedShuffleExec does not support execution".to_owned(), - )) - } - - fn fmt_as( - &self, - t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - match t { - DisplayFormatType::Default => { - write!(f, "UnresolvedShuffleExec") - } - } - } - - fn statistics(&self) -> Statistics { - // The full statistics are computed in the `ShuffleReaderExec` node - // that replaces this one once the previous stage is completed. - Statistics::default() - } -} diff --git a/ballista/rust/core/src/lib.rs b/ballista/rust/core/src/lib.rs deleted file mode 100644 index 34f4699e115a..000000000000 --- a/ballista/rust/core/src/lib.rs +++ /dev/null @@ -1,35 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -#![doc = include_str!("../README.md")] -pub const BALLISTA_VERSION: &str = env!("CARGO_PKG_VERSION"); - -pub fn print_version() { - println!("Ballista version: {}", BALLISTA_VERSION) -} - -pub mod client; -pub mod config; -pub mod error; -pub mod event_loop; -pub mod execution_plans; -/// some plugins -pub mod plugin; -pub mod utils; - -#[macro_use] -pub mod serde; diff --git a/ballista/rust/core/src/plugin/mod.rs b/ballista/rust/core/src/plugin/mod.rs deleted file mode 100644 index a7012af479bc..000000000000 --- a/ballista/rust/core/src/plugin/mod.rs +++ /dev/null @@ -1,127 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use crate::error::Result; -use crate::plugin::udf::UDFPluginManager; -use libloading::Library; -use std::any::Any; -use std::env; -use std::sync::Arc; - -/// plugin manager -pub mod plugin_manager; -/// udf plugin -pub mod udf; - -/// CARGO_PKG_VERSION -pub static CORE_VERSION: &str = env!("CARGO_PKG_VERSION"); -/// RUSTC_VERSION -pub static RUSTC_VERSION: &str = env!("RUSTC_VERSION"); - -/// Top plugin trait -pub trait Plugin { - /// Returns the plugin as [`Any`](std::any::Any) so that it can be - /// downcast to a specific implementation. - fn as_any(&self) -> &dyn Any; -} - -/// The enum of Plugin -#[derive(PartialEq, std::cmp::Eq, std::hash::Hash, Copy, Clone)] -pub enum PluginEnum { - /// UDF/UDAF plugin - UDF, -} - -impl PluginEnum { - /// new a struct which impl the PluginRegistrar trait - pub fn init_plugin_manager(&self) -> Box { - match self { - PluginEnum::UDF => Box::new(UDFPluginManager::default()), - } - } -} - -/// Every plugin need a PluginDeclaration -#[derive(Copy, Clone)] -pub struct PluginDeclaration { - /// Rust doesn’t have a stable ABI, meaning different compiler versions can generate incompatible code. - /// For these reasons, the UDF plug-in must be compiled using the same version of rustc as datafusion. - pub rustc_version: &'static str, - - /// core version of the plugin. The plugin's core_version need same as plugin manager. - pub core_version: &'static str, - - /// One of PluginEnum - pub plugin_type: unsafe extern "C" fn() -> PluginEnum, -} - -/// Plugin Registrar , Every plugin need implement this trait -pub trait PluginRegistrar: Send + Sync + 'static { - /// # Safety - /// load plugin from library - unsafe fn load(&mut self, library: Arc) -> Result<()>; - - /// Returns the plugin as [`Any`](std::any::Any) so that it can be - /// downcast to a specific implementation. - fn as_any(&self) -> &dyn Any; -} - -/// Declare a plugin's PluginDeclaration. -/// -/// # Notes -/// -/// This works by automatically generating an `extern "C"` function named `get_plugin_type` with a -/// pre-defined signature and symbol name. And then generating a PluginDeclaration. -/// Therefore you will only be able to declare one plugin per library. -#[macro_export] -macro_rules! declare_plugin { - ($plugin_type:expr) => { - #[no_mangle] - pub extern "C" fn get_plugin_type() -> $crate::plugin::PluginEnum { - $plugin_type - } - - #[no_mangle] - pub static plugin_declaration: $crate::plugin::PluginDeclaration = - $crate::plugin::PluginDeclaration { - rustc_version: $crate::plugin::RUSTC_VERSION, - core_version: $crate::plugin::CORE_VERSION, - plugin_type: get_plugin_type, - }; - }; -} - -/// get the plugin dir -pub fn plugin_dir() -> String { - let current_exe_dir = match env::current_exe() { - Ok(exe_path) => exe_path.display().to_string(), - Err(_e) => "".to_string(), - }; - - // If current_exe_dir contain `deps` the root dir is the parent dir - // eg: /Users/xxx/workspace/rust/rust_plugin_sty/target/debug/deps/plugins_app-067452b3ff2af70e - // the plugin dir is /Users/xxx/workspace/rust/rust_plugin_sty/target/debug/deps - // else eg: /Users/xxx/workspace/rust/rust_plugin_sty/target/debug/plugins_app - // the plugin dir is /Users/xxx/workspace/rust/rust_plugin_sty/target/debug/ - if current_exe_dir.contains("/deps/") { - let i = current_exe_dir.find("/deps/").unwrap(); - String::from(¤t_exe_dir.as_str()[..i + 6]) - } else { - let i = current_exe_dir.rfind('/').unwrap(); - String::from(¤t_exe_dir.as_str()[..i]) - } -} diff --git a/ballista/rust/core/src/plugin/plugin_manager.rs b/ballista/rust/core/src/plugin/plugin_manager.rs deleted file mode 100644 index e238383b4620..000000000000 --- a/ballista/rust/core/src/plugin/plugin_manager.rs +++ /dev/null @@ -1,150 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. -use crate::error::{BallistaError, Result}; -use libloading::Library; -use log::info; -use std::collections::HashMap; -use std::io; -use std::sync::{Arc, Mutex}; -use walkdir::{DirEntry, WalkDir}; - -use crate::plugin::{ - PluginDeclaration, PluginEnum, PluginRegistrar, CORE_VERSION, RUSTC_VERSION, -}; -use once_cell::sync::OnceCell; - -/// To prevent the library from being loaded multiple times, we use once_cell defines a Arc> -/// Because datafusion is a library, not a service, users may not need to load all plug-ins in the process. -/// So fn global_plugin_manager return Arc>. In this way, users can load the required library through the load method of GlobalPluginManager when needed -static INSTANCE: OnceCell>> = OnceCell::new(); - -/// global_plugin_manager -pub fn global_plugin_manager( - plugin_path: &str, -) -> &'static Arc> { - INSTANCE.get_or_init(move || unsafe { - let mut gpm = GlobalPluginManager::default(); - gpm.load(plugin_path).unwrap(); - Arc::new(Mutex::new(gpm)) - }) -} - -#[derive(Default)] -/// manager all plugin_type's plugin_manager -pub struct GlobalPluginManager { - /// every plugin need a plugin registrar - pub plugin_managers: HashMap>, - - /// loaded plugin files - pub plugin_files: Vec, -} - -impl GlobalPluginManager { - /// # Safety - /// find plugin file from `plugin_path` and load it . - unsafe fn load(&mut self, plugin_path: &str) -> Result<()> { - if "".eq(plugin_path) { - return Ok(()); - } - // find library file from udaf_plugin_path - info!("load plugin from dir:{}", plugin_path); - - let plugin_files = self.get_all_plugin_files(plugin_path)?; - - for plugin_file in plugin_files { - let library = Library::new(plugin_file.path()).map_err(|e| { - BallistaError::IoError(io::Error::new( - io::ErrorKind::Other, - format!("load library error: {}", e), - )) - })?; - - let library = Arc::new(library); - - let dec = library.get::<*mut PluginDeclaration>(b"plugin_declaration\0"); - if dec.is_err() { - info!( - "not found plugin_declaration in the library: {}", - plugin_file.path().to_str().unwrap() - ); - continue; - } - - let dec = dec.unwrap().read(); - - // ersion checks to prevent accidental ABI incompatibilities - if dec.rustc_version != RUSTC_VERSION || dec.core_version != CORE_VERSION { - return Err(BallistaError::IoError(io::Error::new( - io::ErrorKind::Other, - "Version mismatch", - ))); - } - - let plugin_enum = (dec.plugin_type)(); - let curr_plugin_manager = match self.plugin_managers.get_mut(&plugin_enum) { - None => { - let plugin_manager = plugin_enum.init_plugin_manager(); - self.plugin_managers.insert(plugin_enum, plugin_manager); - self.plugin_managers.get_mut(&plugin_enum).unwrap() - } - Some(manager) => manager, - }; - curr_plugin_manager.load(library)?; - self.plugin_files - .push(plugin_file.path().to_str().unwrap().to_string()); - } - - Ok(()) - } - - /// get all plugin file in the dir - fn get_all_plugin_files(&self, plugin_path: &str) -> io::Result> { - let mut plugin_files = Vec::new(); - for entry in WalkDir::new(plugin_path).into_iter().filter_map(|e| { - let item = e.unwrap(); - // every file only load once - if self - .plugin_files - .contains(&item.path().to_str().unwrap().to_string()) - { - return None; - } - - let file_type = item.file_type(); - if !file_type.is_file() { - return None; - } - - if let Some(path) = item.path().extension() { - if let Some(suffix) = path.to_str() { - if suffix == "dylib" || suffix == "so" || suffix == "dll" { - info!( - "load plugin from library file:{}", - item.path().to_str().unwrap() - ); - return Some(item); - } - } - } - - None - }) { - plugin_files.push(entry); - } - Ok(plugin_files) - } -} diff --git a/ballista/rust/core/src/plugin/udf.rs b/ballista/rust/core/src/plugin/udf.rs deleted file mode 100644 index ea82742fb868..000000000000 --- a/ballista/rust/core/src/plugin/udf.rs +++ /dev/null @@ -1,152 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. -use crate::error::{BallistaError, Result}; -use crate::plugin::plugin_manager::global_plugin_manager; -use crate::plugin::{Plugin, PluginEnum, PluginRegistrar}; -use datafusion::physical_plan::udaf::AggregateUDF; -use datafusion::physical_plan::udf::ScalarUDF; -use libloading::{Library, Symbol}; -use std::any::Any; -use std::collections::HashMap; -use std::io; -use std::sync::Arc; - -/// UDF plugin trait -pub trait UDFPlugin: Plugin { - /// get a ScalarUDF by name - fn get_scalar_udf_by_name(&self, fun_name: &str) -> Result; - - /// return all udf names in the plugin - fn udf_names(&self) -> Result>; - - /// get a aggregate udf by name - fn get_aggregate_udf_by_name(&self, fun_name: &str) -> Result; - - /// return all udaf names - fn udaf_names(&self) -> Result>; -} - -/// UDFPluginManager -#[derive(Default, Clone)] -pub struct UDFPluginManager { - /// scalar udfs - pub scalar_udfs: HashMap>, - - /// aggregate udfs - pub aggregate_udfs: HashMap>, - - /// All libraries load from the plugin dir. - pub libraries: Vec>, -} - -impl PluginRegistrar for UDFPluginManager { - unsafe fn load(&mut self, library: Arc) -> Result<()> { - type PluginRegister = unsafe fn() -> Box; - let register_fun: Symbol = - library.get(b"registrar_udf_plugin\0").map_err(|e| { - BallistaError::IoError(io::Error::new( - io::ErrorKind::Other, - format!("not found fn registrar_udf_plugin in the library: {}", e), - )) - })?; - - let udf_plugin: Box = register_fun(); - udf_plugin - .udf_names() - .unwrap() - .iter() - .try_for_each(|udf_name| { - if self.scalar_udfs.contains_key(udf_name) { - Err(BallistaError::IoError(io::Error::new( - io::ErrorKind::Other, - format!("udf name: {} already exists", udf_name), - ))) - } else { - let scalar_udf = udf_plugin.get_scalar_udf_by_name(udf_name)?; - self.scalar_udfs - .insert(udf_name.to_string(), Arc::new(scalar_udf)); - Ok(()) - } - })?; - - udf_plugin - .udaf_names() - .unwrap() - .iter() - .try_for_each(|udaf_name| { - if self.aggregate_udfs.contains_key(udaf_name) { - Err(BallistaError::IoError(io::Error::new( - io::ErrorKind::Other, - format!("udaf name: {} already exists", udaf_name), - ))) - } else { - let aggregate_udf = - udf_plugin.get_aggregate_udf_by_name(udaf_name)?; - self.aggregate_udfs - .insert(udaf_name.to_string(), Arc::new(aggregate_udf)); - Ok(()) - } - })?; - self.libraries.push(library); - Ok(()) - } - - fn as_any(&self) -> &dyn Any { - self - } -} - -/// Declare a udf plugin registrar callback -/// -/// # Notes -/// -/// This works by automatically generating an `extern "C"` function named `registrar_udf_plugin` with a -/// pre-defined signature and symbol name. -/// Therefore you will only be able to declare one plugin per library. -#[macro_export] -macro_rules! declare_udf_plugin { - ($curr_plugin_type:ty, $constructor:path) => { - #[no_mangle] - pub extern "C" fn registrar_udf_plugin() -> Box { - // make sure the constructor is the correct type. - let constructor: fn() -> $curr_plugin_type = $constructor; - let object = constructor(); - Box::new(object) - } - - $crate::declare_plugin!($crate::plugin::PluginEnum::UDF); - }; -} - -/// get a Option of Immutable UDFPluginManager -pub fn get_udf_plugin_manager(path: &str) -> Option { - let udf_plugin_manager_opt = { - let gpm = global_plugin_manager(path).lock().unwrap(); - let plugin_registrar_opt = gpm.plugin_managers.get(&PluginEnum::UDF); - if let Some(plugin_registrar) = plugin_registrar_opt { - if let Some(udf_plugin_manager) = - plugin_registrar.as_any().downcast_ref::() - { - return Some(udf_plugin_manager.clone()); - } else { - return None; - } - } - None - }; - udf_plugin_manager_opt -} diff --git a/ballista/rust/core/src/serde/logical_plan/from_proto.rs b/ballista/rust/core/src/serde/logical_plan/from_proto.rs deleted file mode 100644 index a17c646caaa5..000000000000 --- a/ballista/rust/core/src/serde/logical_plan/from_proto.rs +++ /dev/null @@ -1,52 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Serde code to convert from protocol buffers to Rust data structures. - -use crate::error::BallistaError; -use crate::serde::protobuf; -use std::convert::TryFrom; - -impl TryFrom for protobuf::FileType { - type Error = BallistaError; - fn try_from(value: i32) -> Result { - use protobuf::FileType; - match value { - _x if _x == FileType::NdJson as i32 => Ok(FileType::NdJson), - _x if _x == FileType::Parquet as i32 => Ok(FileType::Parquet), - _x if _x == FileType::Csv as i32 => Ok(FileType::Csv), - _x if _x == FileType::Avro as i32 => Ok(FileType::Avro), - invalid => Err(BallistaError::General(format!( - "Attempted to convert invalid i32 to protobuf::Filetype: {}", - invalid - ))), - } - } -} - -#[allow(clippy::from_over_into)] -impl Into for protobuf::FileType { - fn into(self) -> datafusion::logical_plan::FileType { - use datafusion::logical_plan::FileType; - match self { - protobuf::FileType::NdJson => FileType::NdJson, - protobuf::FileType::Parquet => FileType::Parquet, - protobuf::FileType::Csv => FileType::CSV, - protobuf::FileType::Avro => FileType::Avro, - } - } -} diff --git a/ballista/rust/core/src/serde/logical_plan/mod.rs b/ballista/rust/core/src/serde/logical_plan/mod.rs deleted file mode 100644 index f088f2f1fbb7..000000000000 --- a/ballista/rust/core/src/serde/logical_plan/mod.rs +++ /dev/null @@ -1,1413 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use crate::error::BallistaError; -use crate::serde::protobuf::LogicalExtensionNode; -use crate::serde::{ - byte_to_string, proto_error, protobuf, str_to_byte, AsLogicalPlan, - LogicalExtensionCodec, -}; -use crate::{convert_required, into_logical_plan}; -use datafusion::arrow::datatypes::Schema; -use datafusion::datasource::file_format::avro::AvroFormat; -use datafusion::datasource::file_format::csv::CsvFormat; -use datafusion::datasource::file_format::parquet::ParquetFormat; -use datafusion::datasource::file_format::FileFormat; -use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig}; -use datafusion::logical_plan::plan::{ - Aggregate, EmptyRelation, Filter, Join, Projection, Sort, SubqueryAlias, Window, -}; -use datafusion::logical_plan::{ - provider_as_source, source_as_provider, Column, CreateCatalog, CreateCatalogSchema, - CreateExternalTable, CreateView, CrossJoin, Expr, JoinConstraint, Limit, LogicalPlan, - LogicalPlanBuilder, Offset, Repartition, TableScan, Values, -}; -use datafusion::prelude::SessionContext; - -use datafusion_proto::from_proto::parse_expr; -use prost::bytes::BufMut; -use prost::Message; -use protobuf::listing_table_scan_node::FileFormatType; -use protobuf::logical_plan_node::LogicalPlanType; -use protobuf::LogicalPlanNode; -use std::convert::TryInto; -use std::sync::Arc; - -pub mod from_proto; - -impl AsLogicalPlan for LogicalPlanNode { - fn try_decode(buf: &[u8]) -> Result - where - Self: Sized, - { - LogicalPlanNode::decode(buf).map_err(|e| { - BallistaError::Internal(format!("failed to decode logical plan: {:?}", e)) - }) - } - - fn try_encode(&self, buf: &mut B) -> Result<(), BallistaError> - where - B: BufMut, - Self: Sized, - { - self.encode(buf).map_err(|e| { - BallistaError::Internal(format!("failed to encode logical plan: {:?}", e)) - }) - } - - fn try_into_logical_plan( - &self, - ctx: &SessionContext, - extension_codec: &dyn LogicalExtensionCodec, - ) -> Result { - let plan = self.logical_plan_type.as_ref().ok_or_else(|| { - proto_error(format!( - "logical_plan::from_proto() Unsupported logical plan '{:?}'", - self - )) - })?; - match plan { - LogicalPlanType::Values(values) => { - let n_cols = values.n_cols as usize; - let values: Vec> = if values.values_list.is_empty() { - Ok(Vec::new()) - } else if values.values_list.len() % n_cols != 0 { - Err(BallistaError::General(format!( - "Invalid values list length, expect {} to be divisible by {}", - values.values_list.len(), - n_cols - ))) - } else { - values - .values_list - .chunks_exact(n_cols) - .map(|r| { - r.iter().map(|expr| parse_expr(expr, ctx)).collect::, - datafusion_proto::from_proto::Error, - >>( - ) - }) - .collect::, _>>() - .map_err(|e| e.into()) - }?; - LogicalPlanBuilder::values(values)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Projection(projection) => { - let input: LogicalPlan = - into_logical_plan!(projection.input, ctx, extension_codec)?; - let x: Vec = projection - .expr - .iter() - .map(|expr| parse_expr(expr, ctx)) - .collect::, _>>()?; - LogicalPlanBuilder::from(input) - .project_with_alias( - x, - projection.optional_alias.as_ref().map(|a| match a { - protobuf::projection_node::OptionalAlias::Alias(alias) => { - alias.clone() - } - }), - )? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Selection(selection) => { - let input: LogicalPlan = - into_logical_plan!(selection.input, ctx, extension_codec)?; - let expr: Expr = selection - .expr - .as_ref() - .map(|expr| parse_expr(expr, ctx)) - .transpose()? - .ok_or_else(|| { - BallistaError::General("expression required".to_string()) - })?; - // .try_into()?; - LogicalPlanBuilder::from(input) - .filter(expr)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Window(window) => { - let input: LogicalPlan = - into_logical_plan!(window.input, ctx, extension_codec)?; - let window_expr = window - .window_expr - .iter() - .map(|expr| parse_expr(expr, ctx)) - .collect::, _>>()?; - LogicalPlanBuilder::from(input) - .window(window_expr)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Aggregate(aggregate) => { - let input: LogicalPlan = - into_logical_plan!(aggregate.input, ctx, extension_codec)?; - let group_expr = aggregate - .group_expr - .iter() - .map(|expr| parse_expr(expr, ctx)) - .collect::, _>>()?; - let aggr_expr = aggregate - .aggr_expr - .iter() - .map(|expr| parse_expr(expr, ctx)) - .collect::, _>>()?; - LogicalPlanBuilder::from(input) - .aggregate(group_expr, aggr_expr)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::ListingScan(scan) => { - let schema: Schema = convert_required!(scan.schema)?; - - let mut projection = None; - if let Some(columns) = &scan.projection { - let column_indices = columns - .columns - .iter() - .map(|name| schema.index_of(name)) - .collect::, _>>()?; - projection = Some(column_indices); - } - - let filters = scan - .filters - .iter() - .map(|expr| parse_expr(expr, ctx)) - .collect::, _>>()?; - - let file_format: Arc = - match scan.file_format_type.as_ref().ok_or_else(|| { - proto_error(format!( - "logical_plan::from_proto() Unsupported file format '{:?}'", - self - )) - })? { - &FileFormatType::Parquet(protobuf::ParquetFormat { - enable_pruning, - }) => Arc::new( - ParquetFormat::default().with_enable_pruning(enable_pruning), - ), - FileFormatType::Csv(protobuf::CsvFormat { - has_header, - delimiter, - }) => Arc::new( - CsvFormat::default() - .with_has_header(*has_header) - .with_delimiter(str_to_byte(delimiter)?), - ), - FileFormatType::Avro(..) => Arc::new(AvroFormat::default()), - }; - - let options = ListingOptions { - file_extension: scan.file_extension.clone(), - format: file_format, - table_partition_cols: scan.table_partition_cols.clone(), - collect_stat: scan.collect_stat, - target_partitions: scan.target_partitions as usize, - }; - - let object_store = ctx - .runtime_env() - .object_store(scan.path.as_str()) - .map_err(|e| { - BallistaError::NotImplemented(format!( - "No object store is registered for path {}: {:?}", - scan.path, e - )) - })? - .0; - - println!( - "Found object store {:?} for path {}", - object_store, - scan.path.as_str() - ); - - let config = ListingTableConfig::new(object_store, scan.path.as_str()) - .with_listing_options(options) - .with_schema(Arc::new(schema)); - - let provider = ListingTable::try_new(config)?; - - LogicalPlanBuilder::scan_with_filters( - &scan.table_name, - provider_as_source(Arc::new(provider)), - projection, - filters, - )? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Sort(sort) => { - let input: LogicalPlan = - into_logical_plan!(sort.input, ctx, extension_codec)?; - let sort_expr: Vec = sort - .expr - .iter() - .map(|expr| parse_expr(expr, ctx)) - .collect::, _>>()?; - LogicalPlanBuilder::from(input) - .sort(sort_expr)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Repartition(repartition) => { - use datafusion::logical_plan::Partitioning; - let input: LogicalPlan = - into_logical_plan!(repartition.input, ctx, extension_codec)?; - use protobuf::repartition_node::PartitionMethod; - let pb_partition_method = repartition.partition_method.clone().ok_or_else(|| { - BallistaError::General(String::from( - "Protobuf deserialization error, RepartitionNode was missing required field 'partition_method'", - )) - })?; - - let partitioning_scheme = match pb_partition_method { - PartitionMethod::Hash(protobuf::HashRepartition { - hash_expr: pb_hash_expr, - partition_count, - }) => Partitioning::Hash( - pb_hash_expr - .iter() - .map(|expr| parse_expr(expr, ctx)) - .collect::, _>>()?, - partition_count as usize, - ), - PartitionMethod::RoundRobin(partition_count) => { - Partitioning::RoundRobinBatch(partition_count as usize) - } - }; - - LogicalPlanBuilder::from(input) - .repartition(partitioning_scheme)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::EmptyRelation(empty_relation) => { - LogicalPlanBuilder::empty(empty_relation.produce_one_row) - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::CreateExternalTable(create_extern_table) => { - let pb_schema = (create_extern_table.schema.clone()).ok_or_else(|| { - BallistaError::General(String::from( - "Protobuf deserialization error, CreateExternalTableNode was missing required field schema.", - )) - })?; - - let pb_file_type: protobuf::FileType = - create_extern_table.file_type.try_into()?; - - Ok(LogicalPlan::CreateExternalTable(CreateExternalTable { - schema: pb_schema.try_into()?, - name: create_extern_table.name.clone(), - location: create_extern_table.location.clone(), - file_type: pb_file_type.into(), - has_header: create_extern_table.has_header, - delimiter: create_extern_table.delimiter.chars().next().ok_or_else(|| { - BallistaError::General(String::from("Protobuf deserialization error, unable to parse CSV delimiter")) - })?, - table_partition_cols: create_extern_table - .table_partition_cols - .clone(), - if_not_exists: create_extern_table.if_not_exists, - })) - } - LogicalPlanType::CreateView(create_view) => { - let plan = create_view - .input.clone().ok_or_else(|| BallistaError::General(String::from( - "Protobuf deserialization error, CreateViewNode has invalid LogicalPlan input.", - )))? - .try_into_logical_plan(ctx, extension_codec)?; - - Ok(LogicalPlan::CreateView(CreateView { - name: create_view.name.clone(), - input: Arc::new(plan), - or_replace: create_view.or_replace, - })) - } - LogicalPlanType::CreateCatalogSchema(create_catalog_schema) => { - let pb_schema = (create_catalog_schema.schema.clone()).ok_or_else(|| { - BallistaError::General(String::from( - "Protobuf deserialization error, CreateCatalogSchemaNode was missing required field schema.", - )) - })?; - - Ok(LogicalPlan::CreateCatalogSchema(CreateCatalogSchema { - schema_name: create_catalog_schema.schema_name.clone(), - if_not_exists: create_catalog_schema.if_not_exists, - schema: pb_schema.try_into()?, - })) - } - LogicalPlanType::CreateCatalog(create_catalog) => { - let pb_schema = (create_catalog.schema.clone()).ok_or_else(|| { - BallistaError::General(String::from( - "Protobuf deserialization error, CreateCatalogNode was missing required field schema.", - )) - })?; - - Ok(LogicalPlan::CreateCatalog(CreateCatalog { - catalog_name: create_catalog.catalog_name.clone(), - if_not_exists: create_catalog.if_not_exists, - schema: pb_schema.try_into()?, - })) - } - LogicalPlanType::Analyze(analyze) => { - let input: LogicalPlan = - into_logical_plan!(analyze.input, ctx, extension_codec)?; - LogicalPlanBuilder::from(input) - .explain(analyze.verbose, true)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Explain(explain) => { - let input: LogicalPlan = - into_logical_plan!(explain.input, ctx, extension_codec)?; - LogicalPlanBuilder::from(input) - .explain(explain.verbose, false)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::SubqueryAlias(aliased_relation) => { - let input: LogicalPlan = - into_logical_plan!(aliased_relation.input, ctx, extension_codec)?; - LogicalPlanBuilder::from(input) - .alias(&aliased_relation.alias)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Limit(limit) => { - let input: LogicalPlan = - into_logical_plan!(limit.input, ctx, extension_codec)?; - LogicalPlanBuilder::from(input) - .limit(limit.limit as usize)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Offset(offset) => { - let input: LogicalPlan = - into_logical_plan!(offset.input, ctx, extension_codec)?; - LogicalPlanBuilder::from(input) - .offset(offset.offset as usize)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Join(join) => { - let left_keys: Vec = - join.left_join_column.iter().map(|i| i.into()).collect(); - let right_keys: Vec = - join.right_join_column.iter().map(|i| i.into()).collect(); - let join_type = - protobuf::JoinType::from_i32(join.join_type).ok_or_else(|| { - proto_error(format!( - "Received a JoinNode message with unknown JoinType {}", - join.join_type - )) - })?; - let join_constraint = protobuf::JoinConstraint::from_i32( - join.join_constraint, - ) - .ok_or_else(|| { - proto_error(format!( - "Received a JoinNode message with unknown JoinConstraint {}", - join.join_constraint - )) - })?; - - let builder = LogicalPlanBuilder::from(into_logical_plan!( - join.left, - ctx, - extension_codec - )?); - let builder = match join_constraint.into() { - JoinConstraint::On => builder.join( - &into_logical_plan!(join.right, ctx, extension_codec)?, - join_type.into(), - (left_keys, right_keys), - )?, - JoinConstraint::Using => builder.join_using( - &into_logical_plan!(join.right, ctx, extension_codec)?, - join_type.into(), - left_keys, - )?, - }; - - builder.build().map_err(|e| e.into()) - } - LogicalPlanType::Union(union) => { - let mut input_plans: Vec = union - .inputs - .iter() - .map(|i| i.try_into_logical_plan(ctx, extension_codec)) - .collect::>()?; - - if input_plans.len() < 2 { - return Err( BallistaError::General(String::from( - "Protobuf deserialization error, Union was require at least two input.", - ))); - } - - let mut builder = LogicalPlanBuilder::from(input_plans.pop().unwrap()); - for plan in input_plans { - builder = builder.union(plan)?; - } - builder.build().map_err(|e| e.into()) - } - LogicalPlanType::CrossJoin(crossjoin) => { - let left = into_logical_plan!(crossjoin.left, ctx, extension_codec)?; - let right = into_logical_plan!(crossjoin.right, ctx, extension_codec)?; - - LogicalPlanBuilder::from(left) - .cross_join(&right)? - .build() - .map_err(|e| e.into()) - } - LogicalPlanType::Extension(LogicalExtensionNode { node, inputs }) => { - let input_plans: Vec = inputs - .iter() - .map(|i| i.try_into_logical_plan(ctx, extension_codec)) - .collect::>()?; - - let extension_node = - extension_codec.try_decode(node, &input_plans, ctx)?; - Ok(LogicalPlan::Extension(extension_node)) - } - } - } - - fn try_from_logical_plan( - plan: &LogicalPlan, - extension_codec: &dyn LogicalExtensionCodec, - ) -> Result - where - Self: Sized, - { - match plan { - LogicalPlan::Values(Values { values, .. }) => { - let n_cols = if values.is_empty() { - 0 - } else { - values[0].len() - } as u64; - let values_list = values - .iter() - .flatten() - .map(|v| v.try_into()) - .collect::, _>>()?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Values( - protobuf::ValuesNode { - n_cols, - values_list, - }, - )), - }) - } - LogicalPlan::TableScan(TableScan { - table_name, - source, - filters, - projection, - .. - }) => { - let source = source_as_provider(source)?; - let schema = source.schema(); - let source = source.as_any(); - - let projection = match projection { - None => None, - Some(columns) => { - let column_names = columns - .iter() - .map(|i| schema.field(*i).name().to_owned()) - .collect(); - Some(protobuf::ProjectionColumns { - columns: column_names, - }) - } - }; - let schema: datafusion_proto::protobuf::Schema = schema.as_ref().into(); - - let filters: Vec = filters - .iter() - .map(|filter| filter.try_into()) - .collect::, _>>()?; - - if let Some(listing_table) = source.downcast_ref::() { - let any = listing_table.options().format.as_any(); - let file_format_type = if let Some(parquet) = - any.downcast_ref::() - { - FileFormatType::Parquet(protobuf::ParquetFormat { - enable_pruning: parquet.enable_pruning(), - }) - } else if let Some(csv) = any.downcast_ref::() { - FileFormatType::Csv(protobuf::CsvFormat { - delimiter: byte_to_string(csv.delimiter())?, - has_header: csv.has_header(), - }) - } else if any.is::() { - FileFormatType::Avro(protobuf::AvroFormat {}) - } else { - return Err(proto_error(format!( - "Error converting file format, {:?} is invalid as a datafusion foramt.", - listing_table.options().format - ))); - }; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::ListingScan( - protobuf::ListingTableScanNode { - file_format_type: Some(file_format_type), - table_name: table_name.to_owned(), - collect_stat: listing_table.options().collect_stat, - file_extension: listing_table - .options() - .file_extension - .clone(), - table_partition_cols: listing_table - .options() - .table_partition_cols - .clone(), - path: listing_table.table_path().to_owned(), - schema: Some(schema), - projection, - filters, - target_partitions: listing_table - .options() - .target_partitions - as u32, - }, - )), - }) - } else { - Err(BallistaError::General(format!( - "logical plan to_proto unsupported table provider {:?}", - source - ))) - } - } - LogicalPlan::Projection(Projection { - expr, input, alias, .. - }) => Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Projection(Box::new( - protobuf::ProjectionNode { - input: Some(Box::new( - protobuf::LogicalPlanNode::try_from_logical_plan( - input.as_ref(), - extension_codec, - )?, - )), - expr: expr.iter().map(|expr| expr.try_into()).collect::, - datafusion_proto::to_proto::Error, - >>( - )?, - optional_alias: alias - .clone() - .map(protobuf::projection_node::OptionalAlias::Alias), - }, - ))), - }), - LogicalPlan::Filter(Filter { predicate, input }) => { - let input: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - input.as_ref(), - extension_codec, - )?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Selection(Box::new( - protobuf::SelectionNode { - input: Some(Box::new(input)), - expr: Some(predicate.try_into()?), - }, - ))), - }) - } - LogicalPlan::Window(Window { - input, window_expr, .. - }) => { - let input: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - input.as_ref(), - extension_codec, - )?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Window(Box::new( - protobuf::WindowNode { - input: Some(Box::new(input)), - window_expr: window_expr - .iter() - .map(|expr| expr.try_into()) - .collect::, _>>()?, - }, - ))), - }) - } - LogicalPlan::Aggregate(Aggregate { - group_expr, - aggr_expr, - input, - .. - }) => { - let input: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - input.as_ref(), - extension_codec, - )?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Aggregate(Box::new( - protobuf::AggregateNode { - input: Some(Box::new(input)), - group_expr: group_expr - .iter() - .map(|expr| expr.try_into()) - .collect::, _>>()?, - aggr_expr: aggr_expr - .iter() - .map(|expr| expr.try_into()) - .collect::, _>>()?, - }, - ))), - }) - } - LogicalPlan::Join(Join { - left, - right, - on, - join_type, - join_constraint, - null_equals_null, - .. - }) => { - let left: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - left.as_ref(), - extension_codec, - )?; - let right: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - right.as_ref(), - extension_codec, - )?; - let (left_join_column, right_join_column) = - on.iter().map(|(l, r)| (l.into(), r.into())).unzip(); - let join_type: protobuf::JoinType = join_type.to_owned().into(); - let join_constraint: protobuf::JoinConstraint = - join_constraint.to_owned().into(); - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Join(Box::new( - protobuf::JoinNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - join_type: join_type.into(), - join_constraint: join_constraint.into(), - left_join_column, - right_join_column, - null_equals_null: *null_equals_null, - }, - ))), - }) - } - LogicalPlan::Subquery(_) => { - // note that the ballista and datafusion proto files need refactoring to allow - // LogicalExprNode to reference a LogicalPlanNode - // see https://github.com/apache/arrow-datafusion/issues/2338 - Err(BallistaError::NotImplemented( - "Ballista does not support subqueries".to_string(), - )) - } - LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) => { - let input: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - input.as_ref(), - extension_codec, - )?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::SubqueryAlias(Box::new( - protobuf::SubqueryAliasNode { - input: Some(Box::new(input)), - alias: alias.clone(), - }, - ))), - }) - } - LogicalPlan::Limit(Limit { input, n }) => { - let input: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - input.as_ref(), - extension_codec, - )?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Limit(Box::new( - protobuf::LimitNode { - input: Some(Box::new(input)), - limit: *n as u32, - }, - ))), - }) - } - LogicalPlan::Offset(Offset { input, offset }) => { - let input: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - input.as_ref(), - extension_codec, - )?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Offset(Box::new( - protobuf::OffsetNode { - input: Some(Box::new(input)), - offset: *offset as u32, - }, - ))), - }) - } - LogicalPlan::Sort(Sort { input, expr }) => { - let input: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - input.as_ref(), - extension_codec, - )?; - let selection_expr: Vec = - expr.iter() - .map(|expr| expr.try_into()) - .collect::, datafusion_proto::to_proto::Error>>()?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Sort(Box::new( - protobuf::SortNode { - input: Some(Box::new(input)), - expr: selection_expr, - }, - ))), - }) - } - LogicalPlan::Repartition(Repartition { - input, - partitioning_scheme, - }) => { - use datafusion::logical_plan::Partitioning; - let input: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - input.as_ref(), - extension_codec, - )?; - - // Assumed common usize field was batch size - // Used u64 to avoid any nastyness involving large values, most data clusters are probably uniformly 64 bits any ways - use protobuf::repartition_node::PartitionMethod; - - let pb_partition_method = - match partitioning_scheme { - Partitioning::Hash(exprs, partition_count) => { - PartitionMethod::Hash(protobuf::HashRepartition { - hash_expr: exprs - .iter() - .map(|expr| expr.try_into()) - .collect::, - datafusion_proto::to_proto::Error, - >>()?, - partition_count: *partition_count as u64, - }) - } - Partitioning::RoundRobinBatch(partition_count) => { - PartitionMethod::RoundRobin(*partition_count as u64) - } - }; - - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Repartition(Box::new( - protobuf::RepartitionNode { - input: Some(Box::new(input)), - partition_method: Some(pb_partition_method), - }, - ))), - }) - } - LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row, .. - }) => Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::EmptyRelation( - protobuf::EmptyRelationNode { - produce_one_row: *produce_one_row, - }, - )), - }), - LogicalPlan::CreateExternalTable(CreateExternalTable { - name, - location, - file_type, - has_header, - delimiter, - schema: df_schema, - table_partition_cols, - if_not_exists, - }) => { - use datafusion::logical_plan::FileType; - - let pb_file_type: protobuf::FileType = match file_type { - FileType::NdJson => protobuf::FileType::NdJson, - FileType::Parquet => protobuf::FileType::Parquet, - FileType::CSV => protobuf::FileType::Csv, - FileType::Avro => protobuf::FileType::Avro, - }; - - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::CreateExternalTable( - protobuf::CreateExternalTableNode { - name: name.clone(), - location: location.clone(), - file_type: pb_file_type as i32, - has_header: *has_header, - schema: Some(df_schema.into()), - table_partition_cols: table_partition_cols.clone(), - if_not_exists: *if_not_exists, - delimiter: String::from(*delimiter), - }, - )), - }) - } - LogicalPlan::CreateView(CreateView { - name, - input, - or_replace, - }) => Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::CreateView(Box::new( - protobuf::CreateViewNode { - name: name.clone(), - input: Some(Box::new(LogicalPlanNode::try_from_logical_plan( - input, - extension_codec, - )?)), - or_replace: *or_replace, - }, - ))), - }), - LogicalPlan::CreateCatalogSchema(CreateCatalogSchema { - schema_name, - if_not_exists, - schema: df_schema, - }) => Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::CreateCatalogSchema( - protobuf::CreateCatalogSchemaNode { - schema_name: schema_name.clone(), - if_not_exists: *if_not_exists, - schema: Some(df_schema.into()), - }, - )), - }), - LogicalPlan::CreateCatalog(CreateCatalog { - catalog_name, - if_not_exists, - schema: df_schema, - }) => Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::CreateCatalog( - protobuf::CreateCatalogNode { - catalog_name: catalog_name.clone(), - if_not_exists: *if_not_exists, - schema: Some(df_schema.into()), - }, - )), - }), - LogicalPlan::Analyze(a) => { - let input = protobuf::LogicalPlanNode::try_from_logical_plan( - a.input.as_ref(), - extension_codec, - )?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Analyze(Box::new( - protobuf::AnalyzeNode { - input: Some(Box::new(input)), - verbose: a.verbose, - }, - ))), - }) - } - LogicalPlan::Explain(a) => { - let input = protobuf::LogicalPlanNode::try_from_logical_plan( - a.plan.as_ref(), - extension_codec, - )?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Explain(Box::new( - protobuf::ExplainNode { - input: Some(Box::new(input)), - verbose: a.verbose, - }, - ))), - }) - } - LogicalPlan::Union(union) => { - let inputs: Vec = union - .inputs - .iter() - .map(|i| { - protobuf::LogicalPlanNode::try_from_logical_plan( - i, - extension_codec, - ) - }) - .collect::>()?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Union( - protobuf::UnionNode { inputs }, - )), - }) - } - LogicalPlan::CrossJoin(CrossJoin { left, right, .. }) => { - let left = protobuf::LogicalPlanNode::try_from_logical_plan( - left.as_ref(), - extension_codec, - )?; - let right = protobuf::LogicalPlanNode::try_from_logical_plan( - right.as_ref(), - extension_codec, - )?; - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::CrossJoin(Box::new( - protobuf::CrossJoinNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - }, - ))), - }) - } - LogicalPlan::Extension(extension) => { - let mut buf: Vec = vec![]; - extension_codec.try_encode(extension, &mut buf)?; - - let inputs: Vec = extension - .node - .inputs() - .iter() - .map(|i| { - protobuf::LogicalPlanNode::try_from_logical_plan( - i, - extension_codec, - ) - }) - .collect::>()?; - - Ok(protobuf::LogicalPlanNode { - logical_plan_type: Some(LogicalPlanType::Extension( - LogicalExtensionNode { node: buf, inputs }, - )), - }) - } - LogicalPlan::CreateMemoryTable(_) => Err(proto_error( - "Error converting CreateMemoryTable. Not yet supported in Ballista", - )), - LogicalPlan::DropTable(_) => Err(proto_error( - "Error converting DropTable. Not yet supported in Ballista", - )), - } - } -} - -#[macro_export] -macro_rules! into_logical_plan { - ($PB:expr, $CTX:expr, $CODEC:expr) => {{ - if let Some(field) = $PB.as_ref() { - field.as_ref().try_into_logical_plan($CTX, $CODEC) - } else { - Err(proto_error("Missing required field in protobuf")) - } - }}; -} - -#[cfg(test)] -mod roundtrip_tests { - - use super::super::{super::error::Result, protobuf}; - use crate::serde::{AsLogicalPlan, BallistaCodec}; - use async_trait::async_trait; - use core::panic; - use datafusion::common::DFSchemaRef; - use datafusion::logical_plan::source_as_provider; - use datafusion::{ - arrow::datatypes::{DataType, Field, Schema}, - datafusion_data_access::{ - self, - object_store::{FileMetaStream, ListEntryStream, ObjectReader, ObjectStore}, - SizedFile, - }, - datasource::listing::ListingTable, - logical_plan::{ - col, CreateExternalTable, Expr, FileType, LogicalPlan, LogicalPlanBuilder, - Repartition, ToDFSchema, - }, - prelude::*, - }; - use std::io; - use std::sync::Arc; - - #[derive(Debug)] - struct TestObjectStore {} - - #[async_trait] - impl ObjectStore for TestObjectStore { - async fn list_file( - &self, - _prefix: &str, - ) -> datafusion_data_access::Result { - Err(io::Error::new( - io::ErrorKind::Unsupported, - "this is only a test object store".to_string(), - )) - } - - async fn list_dir( - &self, - _prefix: &str, - _delimiter: Option, - ) -> datafusion_data_access::Result { - Err(io::Error::new( - io::ErrorKind::Unsupported, - "this is only a test object store".to_string(), - )) - } - - fn file_reader( - &self, - _file: SizedFile, - ) -> datafusion_data_access::Result> { - Err(io::Error::new( - io::ErrorKind::Unsupported, - "this is only a test object store".to_string(), - )) - } - } - - // Given a identity of a LogicalPlan converts it to protobuf and back, using debug formatting to test equality. - macro_rules! roundtrip_test { - ($initial_struct:ident, $proto_type:ty, $struct_type:ty) => { - let proto: $proto_type = (&$initial_struct).try_into()?; - - let round_trip: $struct_type = (&proto).try_into()?; - - assert_eq!( - format!("{:?}", $initial_struct), - format!("{:?}", round_trip) - ); - }; - ($initial_struct:ident, $struct_type:ty) => { - roundtrip_test!($initial_struct, protobuf::LogicalPlanNode, $struct_type); - }; - ($initial_struct:ident) => { - let ctx = SessionContext::new(); - let codec: BallistaCodec< - protobuf::LogicalPlanNode, - protobuf::PhysicalPlanNode, - > = BallistaCodec::default(); - let proto: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - &$initial_struct, - codec.logical_extension_codec(), - ) - .expect("from logical plan"); - let round_trip: LogicalPlan = proto - .try_into_logical_plan(&ctx, codec.logical_extension_codec()) - .expect("to logical plan"); - - assert_eq!( - format!("{:?}", $initial_struct), - format!("{:?}", round_trip) - ); - }; - ($initial_struct:ident, $ctx:ident) => { - let codec: BallistaCodec< - protobuf::LogicalPlanNode, - protobuf::PhysicalPlanNode, - > = BallistaCodec::default(); - let proto: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan(&$initial_struct) - .expect("from logical plan"); - let round_trip: LogicalPlan = proto - .try_into_logical_plan(&$ctx, codec.logical_extension_codec()) - .expect("to logical plan"); - - assert_eq!( - format!("{:?}", $initial_struct), - format!("{:?}", round_trip) - ); - }; - } - - #[tokio::test] - async fn roundtrip_repartition() -> Result<()> { - use datafusion::logical_plan::Partitioning; - - let test_partition_counts = [usize::MIN, usize::MAX, 43256]; - - let test_expr: Vec = - vec![col("c1") + col("c2"), Expr::Literal((4.0).into())]; - - let plan = std::sync::Arc::new( - test_scan_csv("employee.csv", Some(vec![3, 4])) - .await? - .sort(vec![col("salary")])? - .build()?, - ); - - for partition_count in test_partition_counts.iter() { - let rr_repartition = Partitioning::RoundRobinBatch(*partition_count); - - let roundtrip_plan = LogicalPlan::Repartition(Repartition { - input: plan.clone(), - partitioning_scheme: rr_repartition, - }); - - roundtrip_test!(roundtrip_plan); - - let h_repartition = Partitioning::Hash(test_expr.clone(), *partition_count); - - let roundtrip_plan = LogicalPlan::Repartition(Repartition { - input: plan.clone(), - partitioning_scheme: h_repartition, - }); - - roundtrip_test!(roundtrip_plan); - - let no_expr_hrepartition = Partitioning::Hash(Vec::new(), *partition_count); - - let roundtrip_plan = LogicalPlan::Repartition(Repartition { - input: plan.clone(), - partitioning_scheme: no_expr_hrepartition, - }); - - roundtrip_test!(roundtrip_plan); - } - - Ok(()) - } - - #[test] - fn roundtrip_create_external_table() -> Result<()> { - let schema = test_schema(); - - let df_schema_ref = schema.to_dfschema_ref()?; - - let filetypes: [FileType; 4] = [ - FileType::NdJson, - FileType::Parquet, - FileType::CSV, - FileType::Avro, - ]; - - for file in filetypes.iter() { - let create_table_node = - LogicalPlan::CreateExternalTable(CreateExternalTable { - schema: df_schema_ref.clone(), - name: String::from("TestName"), - location: String::from("employee.csv"), - file_type: *file, - has_header: true, - delimiter: ',', - table_partition_cols: vec![], - if_not_exists: false, - }); - - roundtrip_test!(create_table_node); - } - - Ok(()) - } - - #[tokio::test] - async fn roundtrip_analyze() -> Result<()> { - let verbose_plan = test_scan_csv("employee.csv", Some(vec![3, 4])) - .await? - .sort(vec![col("salary")])? - .explain(true, true)? - .build()?; - - let plan = test_scan_csv("employee.csv", Some(vec![3, 4])) - .await? - .sort(vec![col("salary")])? - .explain(false, true)? - .build()?; - - roundtrip_test!(plan); - - roundtrip_test!(verbose_plan); - - Ok(()) - } - - #[tokio::test] - async fn roundtrip_explain() -> Result<()> { - let verbose_plan = test_scan_csv("employee.csv", Some(vec![3, 4])) - .await? - .sort(vec![col("salary")])? - .explain(true, false)? - .build()?; - - let plan = test_scan_csv("employee.csv", Some(vec![3, 4])) - .await? - .sort(vec![col("salary")])? - .explain(false, false)? - .build()?; - - roundtrip_test!(plan); - - roundtrip_test!(verbose_plan); - - Ok(()) - } - - #[tokio::test] - async fn roundtrip_join() -> Result<()> { - let scan_plan = test_scan_csv("employee1", Some(vec![0, 3, 4])) - .await? - .build()?; - - let plan = test_scan_csv("employee2", Some(vec![0, 3, 4])) - .await? - .join(&scan_plan, JoinType::Inner, (vec!["id"], vec!["id"]))? - .build()?; - - roundtrip_test!(plan); - Ok(()) - } - - #[tokio::test] - async fn roundtrip_sort() -> Result<()> { - let plan = test_scan_csv("employee.csv", Some(vec![3, 4])) - .await? - .sort(vec![col("salary")])? - .build()?; - roundtrip_test!(plan); - - Ok(()) - } - - #[tokio::test] - async fn roundtrip_empty_relation() -> Result<()> { - let plan_false = LogicalPlanBuilder::empty(false).build()?; - - roundtrip_test!(plan_false); - - let plan_true = LogicalPlanBuilder::empty(true).build()?; - - roundtrip_test!(plan_true); - - Ok(()) - } - - #[tokio::test] - async fn roundtrip_logical_plan() -> Result<()> { - let plan = test_scan_csv("employee.csv", Some(vec![3, 4])) - .await? - .aggregate(vec![col("state")], vec![max(col("salary"))])? - .build()?; - - roundtrip_test!(plan); - - Ok(()) - } - - #[ignore] // see https://github.com/apache/arrow-datafusion/issues/2546 - #[tokio::test] - async fn roundtrip_logical_plan_custom_ctx() -> Result<()> { - let ctx = SessionContext::new(); - let codec: BallistaCodec = - BallistaCodec::default(); - let custom_object_store = Arc::new(TestObjectStore {}); - ctx.runtime_env() - .register_object_store("test", custom_object_store.clone()); - - let (os, uri) = ctx.runtime_env().object_store("test://foo.csv")?; - assert_eq!("TestObjectStore", &format!("{:?}", os)); - assert_eq!("foo.csv", uri); - - let schema = test_schema(); - let plan = ctx - .read_csv( - "test://employee.csv", - CsvReadOptions::new().schema(&schema).has_header(true), - ) - .await? - .to_logical_plan()?; - - let proto: protobuf::LogicalPlanNode = - protobuf::LogicalPlanNode::try_from_logical_plan( - &plan, - codec.logical_extension_codec(), - ) - .expect("from logical plan"); - let round_trip: LogicalPlan = proto - .try_into_logical_plan(&ctx, codec.logical_extension_codec()) - .expect("to logical plan"); - - assert_eq!(format!("{:?}", plan), format!("{:?}", round_trip)); - - let round_trip_store = match round_trip { - LogicalPlan::TableScan(scan) => { - let source = source_as_provider(&scan.source)?; - match source.as_ref().as_any().downcast_ref::() { - Some(listing_table) => { - format!("{:?}", listing_table.object_store()) - } - _ => panic!("expected a ListingTable"), - } - } - _ => panic!("expected a TableScan"), - }; - - assert_eq!(round_trip_store, format!("{:?}", custom_object_store)); - - Ok(()) - } - - fn test_schema() -> Schema { - Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("first_name", DataType::Utf8, false), - Field::new("last_name", DataType::Utf8, false), - Field::new("state", DataType::Utf8, false), - Field::new("salary", DataType::Int32, false), - ]) - } - - async fn test_scan_csv( - table_name: &str, - projection: Option>, - ) -> Result { - let schema = test_schema(); - let ctx = SessionContext::new(); - let options = CsvReadOptions::new().schema(&schema); - let df = ctx.read_csv(table_name, options).await?; - let plan = match df.to_logical_plan()? { - LogicalPlan::TableScan(ref scan) => { - let mut scan = scan.clone(); - scan.projection = projection; - let mut projected_schema = scan.projected_schema.as_ref().clone(); - projected_schema = projected_schema.replace_qualifier(table_name); - scan.projected_schema = DFSchemaRef::new(projected_schema); - LogicalPlan::TableScan(scan) - } - _ => unimplemented!(), - }; - Ok(LogicalPlanBuilder::from(plan)) - } -} diff --git a/ballista/rust/core/src/serde/mod.rs b/ballista/rust/core/src/serde/mod.rs deleted file mode 100644 index 236e66c7dc41..000000000000 --- a/ballista/rust/core/src/serde/mod.rs +++ /dev/null @@ -1,745 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! This crate contains code generated from the Ballista Protocol Buffer Definition as well -//! as convenience code for interacting with the generated code. - -use prost::bytes::BufMut; -use std::fmt::Debug; -use std::marker::PhantomData; -use std::sync::Arc; -use std::{convert::TryInto, io::Cursor}; - -use datafusion::logical_plan::{ - FunctionRegistry, JoinConstraint, JoinType, LogicalPlan, Operator, -}; - -use crate::{error::BallistaError, serde::scheduler::Action as BallistaAction}; - -use datafusion::execution::runtime_env::RuntimeEnv; -use datafusion::logical_plan::plan::Extension; -use datafusion::physical_plan::ExecutionPlan; -use datafusion::prelude::SessionContext; -use prost::Message; - -// include the generated protobuf source as a submodule -#[allow(clippy::all)] -pub mod protobuf { - include!(concat!(env!("OUT_DIR"), "/ballista.protobuf.rs")); -} - -pub mod logical_plan; -pub mod physical_plan; -pub mod scheduler; - -pub fn decode_protobuf(bytes: &[u8]) -> Result { - let mut buf = Cursor::new(bytes); - - protobuf::Action::decode(&mut buf) - .map_err(|e| BallistaError::Internal(format!("{:?}", e))) - .and_then(|node| node.try_into()) -} - -pub(crate) fn proto_error>(message: S) -> BallistaError { - BallistaError::General(message.into()) -} - -pub trait AsLogicalPlan: Debug + Send + Sync + Clone { - fn try_decode(buf: &[u8]) -> Result - where - Self: Sized; - - fn try_encode(&self, buf: &mut B) -> Result<(), BallistaError> - where - B: BufMut, - Self: Sized; - - fn try_into_logical_plan( - &self, - ctx: &SessionContext, - extension_codec: &dyn LogicalExtensionCodec, - ) -> Result; - - fn try_from_logical_plan( - plan: &LogicalPlan, - extension_codec: &dyn LogicalExtensionCodec, - ) -> Result - where - Self: Sized; -} - -pub trait LogicalExtensionCodec: Debug + Send + Sync { - fn try_decode( - &self, - buf: &[u8], - inputs: &[LogicalPlan], - ctx: &SessionContext, - ) -> Result; - - fn try_encode( - &self, - node: &Extension, - buf: &mut Vec, - ) -> Result<(), BallistaError>; -} - -#[derive(Debug, Clone)] -pub struct DefaultLogicalExtensionCodec {} - -impl LogicalExtensionCodec for DefaultLogicalExtensionCodec { - fn try_decode( - &self, - _buf: &[u8], - _inputs: &[LogicalPlan], - _ctx: &SessionContext, - ) -> Result { - Err(BallistaError::NotImplemented( - "LogicalExtensionCodec is not provided".to_string(), - )) - } - - fn try_encode( - &self, - _node: &Extension, - _buf: &mut Vec, - ) -> Result<(), BallistaError> { - Err(BallistaError::NotImplemented( - "LogicalExtensionCodec is not provided".to_string(), - )) - } -} - -pub trait AsExecutionPlan: Debug + Send + Sync + Clone { - fn try_decode(buf: &[u8]) -> Result - where - Self: Sized; - - fn try_encode(&self, buf: &mut B) -> Result<(), BallistaError> - where - B: BufMut, - Self: Sized; - - fn try_into_physical_plan( - &self, - registry: &dyn FunctionRegistry, - runtime: &RuntimeEnv, - extension_codec: &dyn PhysicalExtensionCodec, - ) -> Result, BallistaError>; - - fn try_from_physical_plan( - plan: Arc, - extension_codec: &dyn PhysicalExtensionCodec, - ) -> Result - where - Self: Sized; -} - -pub trait PhysicalExtensionCodec: Debug + Send + Sync { - fn try_decode( - &self, - buf: &[u8], - inputs: &[Arc], - registry: &dyn FunctionRegistry, - ) -> Result, BallistaError>; - - fn try_encode( - &self, - node: Arc, - buf: &mut Vec, - ) -> Result<(), BallistaError>; -} - -#[derive(Debug, Clone)] -pub struct DefaultPhysicalExtensionCodec {} - -impl PhysicalExtensionCodec for DefaultPhysicalExtensionCodec { - fn try_decode( - &self, - _buf: &[u8], - _inputs: &[Arc], - _registry: &dyn FunctionRegistry, - ) -> Result, BallistaError> { - Err(BallistaError::NotImplemented( - "PhysicalExtensionCodec is not provided".to_string(), - )) - } - - fn try_encode( - &self, - _node: Arc, - _buf: &mut Vec, - ) -> Result<(), BallistaError> { - Err(BallistaError::NotImplemented( - "PhysicalExtensionCodec is not provided".to_string(), - )) - } -} - -#[derive(Clone, Debug)] -pub struct BallistaCodec { - logical_extension_codec: Arc, - physical_extension_codec: Arc, - logical_plan_repr: PhantomData, - physical_plan_repr: PhantomData, -} - -impl Default - for BallistaCodec -{ - fn default() -> Self { - Self { - logical_extension_codec: Arc::new(DefaultLogicalExtensionCodec {}), - physical_extension_codec: Arc::new(DefaultPhysicalExtensionCodec {}), - logical_plan_repr: PhantomData, - physical_plan_repr: PhantomData, - } - } -} - -impl BallistaCodec { - pub fn new( - logical_extension_codec: Arc, - physical_extension_codec: Arc, - ) -> Self { - Self { - logical_extension_codec, - physical_extension_codec, - logical_plan_repr: PhantomData, - physical_plan_repr: PhantomData, - } - } - - pub fn logical_extension_codec(&self) -> &dyn LogicalExtensionCodec { - self.logical_extension_codec.as_ref() - } - - pub fn physical_extension_codec(&self) -> &dyn PhysicalExtensionCodec { - self.physical_extension_codec.as_ref() - } -} - -#[macro_export] -macro_rules! convert_required { - ($PB:expr) => {{ - if let Some(field) = $PB.as_ref() { - Ok(field.try_into()?) - } else { - Err(proto_error("Missing required field in protobuf")) - } - }}; -} - -#[macro_export] -macro_rules! into_required { - ($PB:expr) => {{ - if let Some(field) = $PB.as_ref() { - Ok(field.into()) - } else { - Err(proto_error("Missing required field in protobuf")) - } - }}; -} - -#[macro_export] -macro_rules! convert_box_required { - ($PB:expr) => {{ - if let Some(field) = $PB.as_ref() { - field.as_ref().try_into() - } else { - Err(proto_error("Missing required field in protobuf")) - } - }}; -} - -pub(crate) fn from_proto_binary_op(op: &str) -> Result { - match op { - "And" => Ok(Operator::And), - "Or" => Ok(Operator::Or), - "Eq" => Ok(Operator::Eq), - "NotEq" => Ok(Operator::NotEq), - "LtEq" => Ok(Operator::LtEq), - "Lt" => Ok(Operator::Lt), - "Gt" => Ok(Operator::Gt), - "GtEq" => Ok(Operator::GtEq), - "Plus" => Ok(Operator::Plus), - "Minus" => Ok(Operator::Minus), - "Multiply" => Ok(Operator::Multiply), - "Divide" => Ok(Operator::Divide), - "Modulo" => Ok(Operator::Modulo), - "Like" => Ok(Operator::Like), - "NotLike" => Ok(Operator::NotLike), - other => Err(proto_error(format!( - "Unsupported binary operator '{:?}'", - other - ))), - } -} - -impl From for JoinType { - fn from(t: protobuf::JoinType) -> Self { - match t { - protobuf::JoinType::Inner => JoinType::Inner, - protobuf::JoinType::Left => JoinType::Left, - protobuf::JoinType::Right => JoinType::Right, - protobuf::JoinType::Full => JoinType::Full, - protobuf::JoinType::Semi => JoinType::Semi, - protobuf::JoinType::Anti => JoinType::Anti, - } - } -} - -impl From for protobuf::JoinType { - fn from(t: JoinType) -> Self { - match t { - JoinType::Inner => protobuf::JoinType::Inner, - JoinType::Left => protobuf::JoinType::Left, - JoinType::Right => protobuf::JoinType::Right, - JoinType::Full => protobuf::JoinType::Full, - JoinType::Semi => protobuf::JoinType::Semi, - JoinType::Anti => protobuf::JoinType::Anti, - } - } -} - -impl From for JoinConstraint { - fn from(t: protobuf::JoinConstraint) -> Self { - match t { - protobuf::JoinConstraint::On => JoinConstraint::On, - protobuf::JoinConstraint::Using => JoinConstraint::Using, - } - } -} - -impl From for protobuf::JoinConstraint { - fn from(t: JoinConstraint) -> Self { - match t { - JoinConstraint::On => protobuf::JoinConstraint::On, - JoinConstraint::Using => protobuf::JoinConstraint::Using, - } - } -} - -fn byte_to_string(b: u8) -> Result { - let b = &[b]; - let b = std::str::from_utf8(b) - .map_err(|_| BallistaError::General("Invalid CSV delimiter".to_owned()))?; - Ok(b.to_owned()) -} - -fn str_to_byte(s: &str) -> Result { - if s.len() != 1 { - return Err(BallistaError::General("Invalid CSV delimiter".to_owned())); - } - Ok(s.as_bytes()[0]) -} - -#[cfg(test)] -mod tests { - use async_trait::async_trait; - use datafusion::arrow::datatypes::SchemaRef; - use datafusion::error::DataFusionError; - use datafusion::execution::context::{QueryPlanner, SessionState, TaskContext}; - use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv}; - use datafusion::logical_plan::plan::Extension; - use datafusion::logical_plan::{ - col, DFSchemaRef, Expr, FunctionRegistry, LogicalPlan, UserDefinedLogicalNode, - }; - use datafusion::physical_plan::expressions::PhysicalSortExpr; - use datafusion::physical_plan::planner::{DefaultPhysicalPlanner, ExtensionPlanner}; - use datafusion::physical_plan::{ - DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PhysicalPlanner, - SendableRecordBatchStream, Statistics, - }; - use datafusion::prelude::{CsvReadOptions, SessionConfig, SessionContext}; - use prost::Message; - use std::any::Any; - - use datafusion_proto::from_proto::parse_expr; - use std::convert::TryInto; - use std::fmt; - use std::fmt::{Debug, Formatter}; - use std::ops::Deref; - use std::sync::Arc; - - pub mod proto { - #[derive(Clone, PartialEq, ::prost::Message)] - pub struct TopKPlanProto { - #[prost(uint64, tag = "1")] - pub k: u64, - - #[prost(message, optional, tag = "2")] - pub expr: ::core::option::Option, - } - - #[derive(Clone, PartialEq, ::prost::Message)] - pub struct TopKExecProto { - #[prost(uint64, tag = "1")] - pub k: u64, - } - } - - use crate::error::BallistaError; - use crate::serde::protobuf::{LogicalPlanNode, PhysicalPlanNode}; - use crate::serde::{ - AsExecutionPlan, AsLogicalPlan, LogicalExtensionCodec, PhysicalExtensionCodec, - }; - use proto::{TopKExecProto, TopKPlanProto}; - - struct TopKPlanNode { - k: usize, - input: LogicalPlan, - /// The sort expression (this example only supports a single sort - /// expr) - expr: Expr, - } - - impl TopKPlanNode { - pub fn new(k: usize, input: LogicalPlan, expr: Expr) -> Self { - Self { k, input, expr } - } - } - - impl Debug for TopKPlanNode { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - self.fmt_for_explain(f) - } - } - - impl UserDefinedLogicalNode for TopKPlanNode { - fn as_any(&self) -> &dyn Any { - self - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![&self.input] - } - - /// Schema for TopK is the same as the input - fn schema(&self) -> &DFSchemaRef { - self.input.schema() - } - - fn expressions(&self) -> Vec { - vec![self.expr.clone()] - } - - /// For example: `TopK: k=10` - fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!(f, "TopK: k={}", self.k) - } - - fn from_template( - &self, - exprs: &[Expr], - inputs: &[LogicalPlan], - ) -> Arc { - assert_eq!(inputs.len(), 1, "input size inconsistent"); - assert_eq!(exprs.len(), 1, "expression size inconsistent"); - Arc::new(TopKPlanNode { - k: self.k, - input: inputs[0].clone(), - expr: exprs[0].clone(), - }) - } - } - - struct TopKExec { - input: Arc, - /// The maxium number of values - k: usize, - } - - impl TopKExec { - pub fn new(k: usize, input: Arc) -> Self { - Self { input, k } - } - } - - impl Debug for TopKExec { - fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - write!(f, "TopKExec") - } - } - - impl ExecutionPlan for TopKExec { - /// Return a reference to Any that can be used for downcasting - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - self.input.schema() - } - - fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(1) - } - - fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { - None - } - - fn required_child_distribution(&self) -> Distribution { - Distribution::SinglePartition - } - - fn children(&self) -> Vec> { - vec![self.input.clone()] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> datafusion::error::Result> { - Ok(Arc::new(TopKExec { - input: children[0].clone(), - k: self.k, - })) - } - - /// Execute one partition and return an iterator over RecordBatch - fn execute( - &self, - _partition: usize, - _context: Arc, - ) -> datafusion::error::Result { - Err(DataFusionError::NotImplemented( - "not implemented".to_string(), - )) - } - - fn fmt_as( - &self, - t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - match t { - DisplayFormatType::Default => { - write!(f, "TopKExec: k={}", self.k) - } - } - } - - fn statistics(&self) -> Statistics { - // to improve the optimizability of this plan - // better statistics inference could be provided - Statistics::default() - } - } - - struct TopKPlanner {} - - impl ExtensionPlanner for TopKPlanner { - /// Create a physical plan for an extension node - fn plan_extension( - &self, - _planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - _session_state: &SessionState, - ) -> datafusion::error::Result>> { - Ok( - if let Some(topk_node) = node.as_any().downcast_ref::() { - assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs"); - assert_eq!(physical_inputs.len(), 1, "Inconsistent number of inputs"); - // figure out input name - Some(Arc::new(TopKExec { - input: physical_inputs[0].clone(), - k: topk_node.k, - })) - } else { - None - }, - ) - } - } - - struct TopKQueryPlanner {} - - #[async_trait] - impl QueryPlanner for TopKQueryPlanner { - /// Given a `LogicalPlan` created from above, create an - /// `ExecutionPlan` suitable for execution - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &SessionState, - ) -> datafusion::error::Result> { - // Teach the default physical planner how to plan TopK nodes. - let physical_planner = - DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( - TopKPlanner {}, - )]); - // Delegate most work of physical planning to the default physical planner - physical_planner - .create_physical_plan(logical_plan, session_state) - .await - } - } - - #[derive(Debug)] - pub struct TopKExtensionCodec {} - - impl LogicalExtensionCodec for TopKExtensionCodec { - fn try_decode( - &self, - buf: &[u8], - inputs: &[LogicalPlan], - ctx: &SessionContext, - ) -> Result { - if let Some((input, _)) = inputs.split_first() { - let proto = TopKPlanProto::decode(buf).map_err(|e| { - BallistaError::Internal(format!( - "failed to decode logical plan: {:?}", - e - )) - })?; - - if let Some(expr) = proto.expr.as_ref() { - let node = TopKPlanNode::new( - proto.k as usize, - input.clone(), - parse_expr(expr, ctx)?, - ); - - Ok(Extension { - node: Arc::new(node), - }) - } else { - Err(BallistaError::from("invalid plan, no expr".to_string())) - } - } else { - Err(BallistaError::from("invalid plan, no input".to_string())) - } - } - - fn try_encode( - &self, - node: &Extension, - buf: &mut Vec, - ) -> Result<(), BallistaError> { - if let Some(exec) = node.node.as_any().downcast_ref::() { - let proto = TopKPlanProto { - k: exec.k as u64, - expr: Some((&exec.expr).try_into()?), - }; - - proto.encode(buf).map_err(|e| { - BallistaError::Internal(format!( - "failed to encode logical plan: {:?}", - e - )) - })?; - - Ok(()) - } else { - Err(BallistaError::from("unsupported plan type".to_string())) - } - } - } - - impl PhysicalExtensionCodec for TopKExtensionCodec { - fn try_decode( - &self, - buf: &[u8], - inputs: &[Arc], - _registry: &dyn FunctionRegistry, - ) -> Result, BallistaError> { - if let Some((input, _)) = inputs.split_first() { - let proto = TopKExecProto::decode(buf).map_err(|e| { - BallistaError::Internal(format!( - "failed to decode execution plan: {:?}", - e - )) - })?; - Ok(Arc::new(TopKExec::new(proto.k as usize, input.clone()))) - } else { - Err(BallistaError::from("invalid plan, no input".to_string())) - } - } - - fn try_encode( - &self, - node: Arc, - buf: &mut Vec, - ) -> Result<(), BallistaError> { - if let Some(exec) = node.as_any().downcast_ref::() { - let proto = TopKExecProto { k: exec.k as u64 }; - - proto.encode(buf).map_err(|e| { - BallistaError::Internal(format!( - "failed to encode execution plan: {:?}", - e - )) - })?; - - Ok(()) - } else { - Err(BallistaError::from("unsupported plan type".to_string())) - } - } - } - - #[tokio::test] - async fn test_extension_plan() -> crate::error::Result<()> { - let runtime = Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap()); - let session_state = - SessionState::with_config_rt(SessionConfig::new(), runtime.clone()) - .with_query_planner(Arc::new(TopKQueryPlanner {})); - - let ctx = SessionContext::with_state(session_state); - - let scan = ctx - .read_csv( - "../../../datafusion/core/tests/customer.csv", - CsvReadOptions::default(), - ) - .await? - .to_logical_plan()?; - - let topk_plan = LogicalPlan::Extension(Extension { - node: Arc::new(TopKPlanNode::new(3, scan, col("revenue"))), - }); - - let topk_exec = ctx.create_physical_plan(&topk_plan).await?; - - let extension_codec = TopKExtensionCodec {}; - - let proto = LogicalPlanNode::try_from_logical_plan(&topk_plan, &extension_codec)?; - let logical_round_trip = proto.try_into_logical_plan(&ctx, &extension_codec)?; - - assert_eq!( - format!("{:?}", topk_plan), - format!("{:?}", logical_round_trip) - ); - - let proto = PhysicalPlanNode::try_from_physical_plan( - topk_exec.clone(), - &extension_codec, - )?; - let physical_round_trip = - proto.try_into_physical_plan(&ctx, runtime.deref(), &extension_codec)?; - - assert_eq!( - format!("{:?}", topk_exec), - format!("{:?}", physical_round_trip) - ); - - Ok(()) - } -} diff --git a/ballista/rust/core/src/serde/physical_plan/from_proto.rs b/ballista/rust/core/src/serde/physical_plan/from_proto.rs deleted file mode 100644 index 6cd2dc18c8c8..000000000000 --- a/ballista/rust/core/src/serde/physical_plan/from_proto.rs +++ /dev/null @@ -1,396 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Serde code to convert from protocol buffers to Rust data structures. - -use std::convert::{TryFrom, TryInto}; -use std::ops::Deref; -use std::sync::Arc; - -use crate::error::BallistaError; - -use crate::convert_required; -use crate::serde::{from_proto_binary_op, proto_error, protobuf}; -use chrono::{TimeZone, Utc}; - -use datafusion::datafusion_data_access::{ - object_store::local::LocalFileSystem, FileMeta, SizedFile, -}; -use datafusion::datasource::listing::{FileRange, PartitionedFile}; -use datafusion::execution::context::ExecutionProps; -use datafusion::logical_plan::FunctionRegistry; - -use datafusion::physical_plan::file_format::FileScanConfig; - -use datafusion::logical_expr::window_function::WindowFunction; - -use datafusion::physical_plan::{ - expressions::{ - BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, - Literal, NegativeExpr, NotExpr, TryCastExpr, DEFAULT_DATAFUSION_CAST_OPTIONS, - }, - functions::{self, ScalarFunctionExpr}, - Partitioning, -}; -use datafusion::physical_plan::{ColumnStatistics, PhysicalExpr, Statistics}; - -use protobuf::physical_expr_node::ExprType; - -impl From<&protobuf::PhysicalColumn> for Column { - fn from(c: &protobuf::PhysicalColumn) -> Column { - Column::new(&c.name, c.index as usize) - } -} - -pub(crate) fn parse_physical_expr( - proto: &protobuf::PhysicalExprNode, - registry: &dyn FunctionRegistry, -) -> Result, BallistaError> { - let expr_type = proto - .expr_type - .as_ref() - .ok_or_else(|| proto_error("Unexpected empty physical expression"))?; - - let pexpr: Arc = match expr_type { - ExprType::Column(c) => { - let pcol: Column = c.into(); - Arc::new(pcol) - } - ExprType::Literal(scalar) => { - Arc::new(Literal::new(convert_required!(scalar.value)?)) - } - ExprType::BinaryExpr(binary_expr) => Arc::new(BinaryExpr::new( - parse_required_physical_box_expr(&binary_expr.l, registry, "left")?, - from_proto_binary_op(&binary_expr.op)?, - parse_required_physical_box_expr(&binary_expr.r, registry, "right")?, - )), - ExprType::AggregateExpr(_) => { - return Err(BallistaError::General( - "Cannot convert aggregate expr node to physical expression".to_owned(), - )); - } - ExprType::WindowExpr(_) => { - return Err(BallistaError::General( - "Cannot convert window expr node to physical expression".to_owned(), - )); - } - ExprType::Sort(_) => { - return Err(BallistaError::General( - "Cannot convert sort expr node to physical expression".to_owned(), - )); - } - ExprType::IsNullExpr(e) => Arc::new(IsNullExpr::new( - parse_required_physical_box_expr(&e.expr, registry, "expr")?, - )), - ExprType::IsNotNullExpr(e) => Arc::new(IsNotNullExpr::new( - parse_required_physical_box_expr(&e.expr, registry, "expr")?, - )), - ExprType::NotExpr(e) => Arc::new(NotExpr::new(parse_required_physical_box_expr( - &e.expr, registry, "expr", - )?)), - ExprType::Negative(e) => Arc::new(NegativeExpr::new( - parse_required_physical_box_expr(&e.expr, registry, "expr")?, - )), - ExprType::InList(e) => Arc::new(InListExpr::new( - parse_required_physical_box_expr(&e.expr, registry, "expr")?, - e.list - .iter() - .map(|x| parse_physical_expr(x, registry)) - .collect::, _>>()?, - e.negated, - )), - ExprType::Case(e) => Arc::new(CaseExpr::try_new( - e.expr - .as_ref() - .map(|e| parse_physical_expr(e.as_ref(), registry)) - .transpose()?, - e.when_then_expr - .iter() - .map(|e| { - Ok(( - parse_required_physical_expr( - &e.when_expr, - registry, - "when_expr", - )?, - parse_required_physical_expr( - &e.then_expr, - registry, - "then_expr", - )?, - )) - }) - .collect::, BallistaError>>()? - .as_slice(), - e.else_expr - .as_ref() - .map(|e| parse_physical_expr(e.as_ref(), registry)) - .transpose()?, - )?), - ExprType::Cast(e) => Arc::new(CastExpr::new( - parse_required_physical_box_expr(&e.expr, registry, "expr")?, - convert_required!(e.arrow_type)?, - DEFAULT_DATAFUSION_CAST_OPTIONS, - )), - ExprType::TryCast(e) => Arc::new(TryCastExpr::new( - parse_required_physical_box_expr(&e.expr, registry, "expr")?, - convert_required!(e.arrow_type)?, - )), - ExprType::ScalarFunction(e) => { - let scalar_function = datafusion_proto::protobuf::ScalarFunction::from_i32( - e.fun, - ) - .ok_or_else(|| { - proto_error(format!("Received an unknown scalar function: {}", e.fun,)) - })?; - - let args = e - .args - .iter() - .map(|x| parse_physical_expr(x, registry)) - .collect::, _>>()?; - - // TODO Do not create new the ExecutionProps - let execution_props = ExecutionProps::new(); - - let fun_expr = functions::create_physical_fun( - &(&scalar_function).into(), - &execution_props, - )?; - - Arc::new(ScalarFunctionExpr::new( - &e.name, - fun_expr, - args, - &convert_required!(e.return_type)?, - )) - } - ExprType::ScalarUdf(e) => { - let scalar_fun = registry.udf(e.name.as_str())?.deref().clone().fun; - - let args = e - .args - .iter() - .map(|x| parse_physical_expr(x, registry)) - .collect::, _>>()?; - - Arc::new(ScalarFunctionExpr::new( - e.name.as_str(), - scalar_fun, - args, - &convert_required!(e.return_type)?, - )) - } - }; - - Ok(pexpr) -} - -fn parse_required_physical_box_expr( - expr: &Option>, - registry: &dyn FunctionRegistry, - field: &str, -) -> Result, BallistaError> { - expr.as_ref() - .map(|e| parse_physical_expr(e.as_ref(), registry)) - .transpose()? - .ok_or_else(|| { - BallistaError::General(format!("Missing required field {:?}", field)) - }) -} - -fn parse_required_physical_expr( - expr: &Option, - registry: &dyn FunctionRegistry, - field: &str, -) -> Result, BallistaError> { - expr.as_ref() - .map(|e| parse_physical_expr(e, registry)) - .transpose()? - .ok_or_else(|| { - BallistaError::General(format!("Missing required field {:?}", field)) - }) -} - -impl TryFrom<&protobuf::physical_window_expr_node::WindowFunction> for WindowFunction { - type Error = BallistaError; - - fn try_from( - expr: &protobuf::physical_window_expr_node::WindowFunction, - ) -> Result { - match expr { - protobuf::physical_window_expr_node::WindowFunction::AggrFunction(n) => { - let f = datafusion_proto::protobuf::AggregateFunction::from_i32(*n) - .ok_or_else(|| { - proto_error(format!( - "Received an unknown window aggregate function: {}", - n - )) - })?; - - Ok(WindowFunction::AggregateFunction(f.into())) - } - protobuf::physical_window_expr_node::WindowFunction::BuiltInFunction(n) => { - let f = datafusion_proto::protobuf::BuiltInWindowFunction::from_i32(*n) - .ok_or_else(|| { - proto_error(format!( - "Received an unknown window builtin function: {}", - n - )) - })?; - - Ok(WindowFunction::BuiltInWindowFunction(f.into())) - } - } - } -} - -pub fn parse_protobuf_hash_partitioning( - partitioning: Option<&protobuf::PhysicalHashRepartition>, - registry: &dyn FunctionRegistry, -) -> Result, BallistaError> { - match partitioning { - Some(hash_part) => { - let expr = hash_part - .hash_expr - .iter() - .map(|e| parse_physical_expr(e, registry)) - .collect::>, _>>()?; - - Ok(Some(Partitioning::Hash( - expr, - hash_part.partition_count.try_into().unwrap(), - ))) - } - None => Ok(None), - } -} - -impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile { - type Error = BallistaError; - - fn try_from(val: &protobuf::PartitionedFile) -> Result { - Ok(PartitionedFile { - file_meta: FileMeta { - sized_file: SizedFile { - path: val.path.clone(), - size: val.size, - }, - last_modified: if val.last_modified_ns == 0 { - None - } else { - Some(Utc.timestamp_nanos(val.last_modified_ns as i64)) - }, - }, - partition_values: val - .partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - range: val.range.as_ref().map(|v| v.try_into()).transpose()?, - }) - } -} - -impl TryFrom<&protobuf::FileRange> for FileRange { - type Error = BallistaError; - - fn try_from(value: &protobuf::FileRange) -> Result { - Ok(FileRange { - start: value.start, - end: value.end, - }) - } -} - -impl TryFrom<&protobuf::FileGroup> for Vec { - type Error = BallistaError; - - fn try_from(val: &protobuf::FileGroup) -> Result { - val.files - .iter() - .map(|f| f.try_into()) - .collect::, _>>() - } -} - -impl From<&protobuf::ColumnStats> for ColumnStatistics { - fn from(cs: &protobuf::ColumnStats) -> ColumnStatistics { - ColumnStatistics { - null_count: Some(cs.null_count as usize), - max_value: cs.max_value.as_ref().map(|m| m.try_into().unwrap()), - min_value: cs.min_value.as_ref().map(|m| m.try_into().unwrap()), - distinct_count: Some(cs.distinct_count as usize), - } - } -} - -impl TryInto for &protobuf::Statistics { - type Error = BallistaError; - - fn try_into(self) -> Result { - let column_statistics = self - .column_stats - .iter() - .map(|s| s.into()) - .collect::>(); - Ok(Statistics { - num_rows: Some(self.num_rows as usize), - total_byte_size: Some(self.total_byte_size as usize), - // No column statistic (None) is encoded with empty array - column_statistics: if column_statistics.is_empty() { - None - } else { - Some(column_statistics) - }, - is_exact: self.is_exact, - }) - } -} - -impl TryInto for &protobuf::FileScanExecConf { - type Error = BallistaError; - - fn try_into(self) -> Result { - let schema = Arc::new(convert_required!(self.schema)?); - let projection = self - .projection - .iter() - .map(|i| *i as usize) - .collect::>(); - let projection = if projection.is_empty() { - None - } else { - Some(projection) - }; - let statistics = convert_required!(self.statistics)?; - - Ok(FileScanConfig { - object_store: Arc::new(LocalFileSystem {}), - file_schema: schema, - file_groups: self - .file_groups - .iter() - .map(|f| f.try_into()) - .collect::, _>>()?, - statistics, - projection, - limit: self.limit.as_ref().map(|sl| sl.limit as usize), - table_partition_cols: vec![], - }) - } -} diff --git a/ballista/rust/core/src/serde/physical_plan/mod.rs b/ballista/rust/core/src/serde/physical_plan/mod.rs deleted file mode 100644 index f5b495b67f1d..000000000000 --- a/ballista/rust/core/src/serde/physical_plan/mod.rs +++ /dev/null @@ -1,1396 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::convert::TryInto; -use std::sync::Arc; - -use prost::bytes::BufMut; -use prost::Message; - -use datafusion::arrow::compute::SortOptions; -use datafusion::arrow::datatypes::SchemaRef; -use datafusion::datafusion_data_access::object_store::local::LocalFileSystem; -use datafusion::datasource::listing::PartitionedFile; -use datafusion::execution::runtime_env::RuntimeEnv; -use datafusion::logical_plan::window_frames::WindowFrame; -use datafusion::logical_plan::FunctionRegistry; -use datafusion::physical_plan::aggregates::AggregateExec; -use datafusion::physical_plan::aggregates::{create_aggregate_expr, AggregateMode}; -use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; -use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion::physical_plan::cross_join::CrossJoinExec; -use datafusion::physical_plan::empty::EmptyExec; -use datafusion::physical_plan::explain::ExplainExec; -use datafusion::physical_plan::expressions::{Column, PhysicalSortExpr}; -use datafusion::physical_plan::file_format::{ - AvroExec, CsvExec, FileScanConfig, ParquetExec, -}; -use datafusion::physical_plan::filter::FilterExec; -use datafusion::physical_plan::hash_join::{HashJoinExec, PartitionMode}; -use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; -use datafusion::physical_plan::projection::ProjectionExec; -use datafusion::physical_plan::repartition::RepartitionExec; -use datafusion::physical_plan::sorts::sort::SortExec; -use datafusion::physical_plan::union::UnionExec; -use datafusion::physical_plan::windows::{create_window_expr, WindowAggExec}; -use datafusion::physical_plan::{ - AggregateExpr, ExecutionPlan, Partitioning, PhysicalExpr, WindowExpr, -}; -use datafusion_proto::from_proto::parse_expr; - -use crate::error::BallistaError; -use crate::execution_plans::{ - ShuffleReaderExec, ShuffleWriterExec, UnresolvedShuffleExec, -}; -use crate::serde::physical_plan::from_proto::{ - parse_physical_expr, parse_protobuf_hash_partitioning, -}; -use crate::serde::protobuf::physical_expr_node::ExprType; -use crate::serde::protobuf::physical_plan_node::PhysicalPlanType; -use crate::serde::protobuf::repartition_exec_node::PartitionMethod; -use crate::serde::protobuf::{PhysicalExtensionNode, PhysicalPlanNode}; -use crate::serde::scheduler::PartitionLocation; -use crate::serde::{ - byte_to_string, proto_error, protobuf, str_to_byte, AsExecutionPlan, - PhysicalExtensionCodec, -}; -use crate::{convert_required, into_physical_plan, into_required}; - -pub mod from_proto; -pub mod to_proto; - -impl AsExecutionPlan for PhysicalPlanNode { - fn try_decode(buf: &[u8]) -> Result - where - Self: Sized, - { - PhysicalPlanNode::decode(buf).map_err(|e| { - BallistaError::Internal(format!("failed to decode physical plan: {:?}", e)) - }) - } - - fn try_encode(&self, buf: &mut B) -> Result<(), BallistaError> - where - B: BufMut, - Self: Sized, - { - self.encode(buf).map_err(|e| { - BallistaError::Internal(format!("failed to encode physical plan: {:?}", e)) - }) - } - - fn try_into_physical_plan( - &self, - registry: &dyn FunctionRegistry, - runtime: &RuntimeEnv, - extension_codec: &dyn PhysicalExtensionCodec, - ) -> Result, BallistaError> { - let plan = self.physical_plan_type.as_ref().ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unsupported physical plan '{:?}'", - self - )) - })?; - match plan { - PhysicalPlanType::Explain(explain) => Ok(Arc::new(ExplainExec::new( - Arc::new(explain.schema.as_ref().unwrap().try_into()?), - explain - .stringified_plans - .iter() - .map(|plan| plan.into()) - .collect(), - explain.verbose, - ))), - PhysicalPlanType::Projection(projection) => { - let input: Arc = into_physical_plan!( - projection.input, - registry, - runtime, - extension_codec - )?; - let exprs = projection - .expr - .iter() - .zip(projection.expr_name.iter()) - .map(|(expr, name)| Ok((parse_physical_expr(expr,registry)?, name.to_string()))) - .collect::, String)>, BallistaError>>( - )?; - Ok(Arc::new(ProjectionExec::try_new(exprs, input)?)) - } - PhysicalPlanType::Filter(filter) => { - let input: Arc = into_physical_plan!( - filter.input, - registry, - runtime, - extension_codec - )?; - let predicate = filter - .expr - .as_ref() - .map(|expr| parse_physical_expr(expr, registry)) - .transpose()? - .ok_or_else(|| { - BallistaError::General( - "filter (FilterExecNode) in PhysicalPlanNode is missing." - .to_owned(), - ) - })?; - Ok(Arc::new(FilterExec::try_new(predicate, input)?)) - } - PhysicalPlanType::CsvScan(scan) => Ok(Arc::new(CsvExec::new( - decode_scan_config(scan.base_conf.as_ref().unwrap(), runtime)?, - scan.has_header, - str_to_byte(&scan.delimiter)?, - ))), - PhysicalPlanType::ParquetScan(scan) => { - let predicate = scan - .pruning_predicate - .as_ref() - .map(|expr| parse_expr(expr, registry)) - .transpose()?; - Ok(Arc::new(ParquetExec::new( - decode_scan_config(scan.base_conf.as_ref().unwrap(), runtime)?, - predicate, - ))) - } - PhysicalPlanType::AvroScan(scan) => Ok(Arc::new(AvroExec::new( - decode_scan_config(scan.base_conf.as_ref().unwrap(), runtime)?, - ))), - PhysicalPlanType::CoalesceBatches(coalesce_batches) => { - let input: Arc = into_physical_plan!( - coalesce_batches.input, - registry, - runtime, - extension_codec - )?; - Ok(Arc::new(CoalesceBatchesExec::new( - input, - coalesce_batches.target_batch_size as usize, - ))) - } - PhysicalPlanType::Merge(merge) => { - let input: Arc = - into_physical_plan!(merge.input, registry, runtime, extension_codec)?; - Ok(Arc::new(CoalescePartitionsExec::new(input))) - } - PhysicalPlanType::Repartition(repart) => { - let input: Arc = into_physical_plan!( - repart.input, - registry, - runtime, - extension_codec - )?; - match repart.partition_method { - Some(PartitionMethod::Hash(ref hash_part)) => { - let expr = hash_part - .hash_expr - .iter() - .map(|e| parse_physical_expr(e, registry)) - .collect::>, _>>()?; - - Ok(Arc::new(RepartitionExec::try_new( - input, - Partitioning::Hash( - expr, - hash_part.partition_count.try_into().unwrap(), - ), - )?)) - } - Some(PartitionMethod::RoundRobin(partition_count)) => { - Ok(Arc::new(RepartitionExec::try_new( - input, - Partitioning::RoundRobinBatch( - partition_count.try_into().unwrap(), - ), - )?)) - } - Some(PartitionMethod::Unknown(partition_count)) => { - Ok(Arc::new(RepartitionExec::try_new( - input, - Partitioning::UnknownPartitioning( - partition_count.try_into().unwrap(), - ), - )?)) - } - _ => Err(BallistaError::General( - "Invalid partitioning scheme".to_owned(), - )), - } - } - PhysicalPlanType::GlobalLimit(limit) => { - let input: Arc = - into_physical_plan!(limit.input, registry, runtime, extension_codec)?; - Ok(Arc::new(GlobalLimitExec::new(input, limit.limit as usize))) - } - PhysicalPlanType::LocalLimit(limit) => { - let input: Arc = - into_physical_plan!(limit.input, registry, runtime, extension_codec)?; - Ok(Arc::new(LocalLimitExec::new(input, limit.limit as usize))) - } - PhysicalPlanType::Window(window_agg) => { - let input: Arc = into_physical_plan!( - window_agg.input, - registry, - runtime, - extension_codec - )?; - let input_schema = window_agg - .input_schema - .as_ref() - .ok_or_else(|| { - BallistaError::General( - "input_schema in WindowAggrNode is missing.".to_owned(), - ) - })? - .clone(); - let physical_schema: SchemaRef = - SchemaRef::new((&input_schema).try_into()?); - - let physical_window_expr: Vec> = window_agg - .window_expr - .iter() - .zip(window_agg.window_expr_name.iter()) - .map(|(expr, name)| { - let expr_type = expr.expr_type.as_ref().ok_or_else(|| { - proto_error("Unexpected empty window physical expression") - })?; - - match expr_type { - ExprType::WindowExpr(window_node) => { - let window_node_expr = window_node - .expr - .as_ref() - .map(|e| parse_physical_expr(e.as_ref(), registry)) - .transpose()? - .ok_or_else(|| { - proto_error( - "missing window_node expr expression" - .to_string(), - ) - })?; - - Ok(create_window_expr( - &convert_required!(window_node.window_function)?, - name.to_owned(), - &[window_node_expr], - &[], - &[], - Some(WindowFrame::default()), - &physical_schema, - )?) - } - _ => Err(BallistaError::General( - "Invalid expression for WindowAggrExec".to_string(), - )), - } - }) - .collect::, _>>()?; - - Ok(Arc::new(WindowAggExec::try_new( - physical_window_expr, - input, - Arc::new((&input_schema).try_into()?), - )?)) - } - PhysicalPlanType::Aggregate(hash_agg) => { - let input: Arc = into_physical_plan!( - hash_agg.input, - registry, - runtime, - extension_codec - )?; - let mode = protobuf::AggregateMode::from_i32(hash_agg.mode).ok_or_else( - || { - proto_error(format!( - "Received a AggregateNode message with unknown AggregateMode {}", - hash_agg.mode - )) - }, - )?; - let agg_mode: AggregateMode = match mode { - protobuf::AggregateMode::Partial => AggregateMode::Partial, - protobuf::AggregateMode::Final => AggregateMode::Final, - protobuf::AggregateMode::FinalPartitioned => { - AggregateMode::FinalPartitioned - } - }; - let group = hash_agg - .group_expr - .iter() - .zip(hash_agg.group_expr_name.iter()) - .map(|(expr, name)| { - parse_physical_expr(expr, registry) - .map(|expr| (expr, name.to_string())) - }) - .collect::, _>>()?; - - let input_schema = hash_agg - .input_schema - .as_ref() - .ok_or_else(|| { - BallistaError::General( - "input_schema in AggregateNode is missing.".to_owned(), - ) - })? - .clone(); - let physical_schema: SchemaRef = - SchemaRef::new((&input_schema).try_into()?); - - let physical_aggr_expr: Vec> = hash_agg - .aggr_expr - .iter() - .zip(hash_agg.aggr_expr_name.iter()) - .map(|(expr, name)| { - let expr_type = expr.expr_type.as_ref().ok_or_else(|| { - proto_error("Unexpected empty aggregate physical expression") - })?; - - match expr_type { - ExprType::AggregateExpr(agg_node) => { - let aggr_function = - datafusion_proto::protobuf::AggregateFunction::from_i32( - agg_node.aggr_function, - ) - .ok_or_else( - || { - proto_error(format!( - "Received an unknown aggregate function: {}", - agg_node.aggr_function - )) - }, - )?; - - let input_phy_expr: Vec> = agg_node.expr.iter() - .map(|e| parse_physical_expr(e, registry).unwrap()).collect(); - - Ok(create_aggregate_expr( - &aggr_function.into(), - false, - input_phy_expr.as_slice(), - &physical_schema, - name.to_string(), - )?) - } - _ => Err(BallistaError::General( - "Invalid aggregate expression for AggregateExec" - .to_string(), - )), - } - }) - .collect::, _>>()?; - - Ok(Arc::new(AggregateExec::try_new( - agg_mode, - group, - physical_aggr_expr, - input, - Arc::new((&input_schema).try_into()?), - )?)) - } - PhysicalPlanType::HashJoin(hashjoin) => { - let left: Arc = into_physical_plan!( - hashjoin.left, - registry, - runtime, - extension_codec - )?; - let right: Arc = into_physical_plan!( - hashjoin.right, - registry, - runtime, - extension_codec - )?; - let on: Vec<(Column, Column)> = hashjoin - .on - .iter() - .map(|col| { - let left = into_required!(col.left)?; - let right = into_required!(col.right)?; - Ok((left, right)) - }) - .collect::>()?; - let join_type = protobuf::JoinType::from_i32(hashjoin.join_type) - .ok_or_else(|| { - proto_error(format!( - "Received a HashJoinNode message with unknown JoinType {}", - hashjoin.join_type - )) - })?; - - let partition_mode = - protobuf::PartitionMode::from_i32(hashjoin.partition_mode) - .ok_or_else(|| { - proto_error(format!( - "Received a HashJoinNode message with unknown PartitionMode {}", - hashjoin.partition_mode - )) - })?; - let partition_mode = match partition_mode { - protobuf::PartitionMode::CollectLeft => PartitionMode::CollectLeft, - protobuf::PartitionMode::Partitioned => PartitionMode::Partitioned, - }; - Ok(Arc::new(HashJoinExec::try_new( - left, - right, - on, - &join_type.into(), - partition_mode, - &hashjoin.null_equals_null, - )?)) - } - PhysicalPlanType::Union(union) => { - let mut inputs: Vec> = vec![]; - for input in &union.inputs { - inputs.push(input.try_into_physical_plan( - registry, - runtime, - extension_codec, - )?); - } - Ok(Arc::new(UnionExec::new(inputs))) - } - PhysicalPlanType::CrossJoin(crossjoin) => { - let left: Arc = into_physical_plan!( - crossjoin.left, - registry, - runtime, - extension_codec - )?; - let right: Arc = into_physical_plan!( - crossjoin.right, - registry, - runtime, - extension_codec - )?; - Ok(Arc::new(CrossJoinExec::try_new(left, right)?)) - } - PhysicalPlanType::ShuffleWriter(shuffle_writer) => { - let input: Arc = into_physical_plan!( - shuffle_writer.input, - registry, - runtime, - extension_codec - )?; - - let output_partitioning = parse_protobuf_hash_partitioning( - shuffle_writer.output_partitioning.as_ref(), - registry, - )?; - - Ok(Arc::new(ShuffleWriterExec::try_new( - shuffle_writer.job_id.clone(), - shuffle_writer.stage_id as usize, - input, - "".to_string(), // this is intentional but hacky - the executor will fill this in - output_partitioning, - )?)) - } - PhysicalPlanType::ShuffleReader(shuffle_reader) => { - let schema = Arc::new(convert_required!(shuffle_reader.schema)?); - let partition_location: Vec> = shuffle_reader - .partition - .iter() - .map(|p| { - p.location - .iter() - .map(|l| l.clone().try_into()) - .collect::, _>>() - }) - .collect::, BallistaError>>()?; - let shuffle_reader = - ShuffleReaderExec::try_new(partition_location, schema)?; - Ok(Arc::new(shuffle_reader)) - } - PhysicalPlanType::Empty(empty) => { - let schema = Arc::new(convert_required!(empty.schema)?); - Ok(Arc::new(EmptyExec::new(empty.produce_one_row, schema))) - } - PhysicalPlanType::Sort(sort) => { - let input: Arc = - into_physical_plan!(sort.input, registry, runtime, extension_codec)?; - let exprs = sort - .expr - .iter() - .map(|expr| { - let expr = expr.expr_type.as_ref().ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected expr {:?}", - self - )) - })?; - if let protobuf::physical_expr_node::ExprType::Sort(sort_expr) = expr { - let expr = sort_expr - .expr - .as_ref() - .ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected sort expr {:?}", - self - )) - })? - .as_ref(); - Ok(PhysicalSortExpr { - expr: parse_physical_expr(expr,registry)?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - } else { - Err(BallistaError::General(format!( - "physical_plan::from_proto() {:?}", - self - ))) - } - }) - .collect::, _>>()?; - Ok(Arc::new(SortExec::try_new(exprs, input)?)) - } - PhysicalPlanType::Unresolved(unresolved_shuffle) => { - let schema = Arc::new(convert_required!(unresolved_shuffle.schema)?); - Ok(Arc::new(UnresolvedShuffleExec { - stage_id: unresolved_shuffle.stage_id as usize, - schema, - input_partition_count: unresolved_shuffle.input_partition_count - as usize, - output_partition_count: unresolved_shuffle.output_partition_count - as usize, - })) - } - PhysicalPlanType::Extension(extension) => { - let inputs: Vec> = extension - .inputs - .iter() - .map(|i| i.try_into_physical_plan(registry, runtime, extension_codec)) - .collect::>()?; - - let extension_node = extension_codec.try_decode( - extension.node.as_slice(), - &inputs, - registry, - )?; - - Ok(extension_node) - } - } - } - - fn try_from_physical_plan( - plan: Arc, - extension_codec: &dyn PhysicalExtensionCodec, - ) -> Result - where - Self: Sized, - { - let plan_clone = plan.clone(); - let plan = plan.as_any(); - - if let Some(exec) = plan.downcast_ref::() { - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Explain( - protobuf::ExplainExecNode { - schema: Some(exec.schema().as_ref().into()), - stringified_plans: exec - .stringified_plans() - .iter() - .map(|plan| plan.into()) - .collect(), - verbose: exec.verbose(), - }, - )), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.input().to_owned(), - extension_codec, - )?; - let expr = exec - .expr() - .iter() - .map(|expr| expr.0.clone().try_into()) - .collect::, BallistaError>>()?; - let expr_name = exec.expr().iter().map(|expr| expr.1.clone()).collect(); - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Projection(Box::new( - protobuf::ProjectionExecNode { - input: Some(Box::new(input)), - expr, - expr_name, - }, - ))), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.input().to_owned(), - extension_codec, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Filter(Box::new( - protobuf::FilterExecNode { - input: Some(Box::new(input)), - expr: Some(exec.predicate().clone().try_into()?), - }, - ))), - }) - } else if let Some(limit) = plan.downcast_ref::() { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - limit.input().to_owned(), - extension_codec, - )?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::GlobalLimit(Box::new( - protobuf::GlobalLimitExecNode { - input: Some(Box::new(input)), - limit: limit.limit() as u32, - }, - ))), - }) - } else if let Some(limit) = plan.downcast_ref::() { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - limit.input().to_owned(), - extension_codec, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::LocalLimit(Box::new( - protobuf::LocalLimitExecNode { - input: Some(Box::new(input)), - limit: limit.limit() as u32, - }, - ))), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.left().to_owned(), - extension_codec, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.right().to_owned(), - extension_codec, - )?; - let on: Vec = exec - .on() - .iter() - .map(|tuple| protobuf::JoinOn { - left: Some(protobuf::PhysicalColumn { - name: tuple.0.name().to_string(), - index: tuple.0.index() as u32, - }), - right: Some(protobuf::PhysicalColumn { - name: tuple.1.name().to_string(), - index: tuple.1.index() as u32, - }), - }) - .collect(); - let join_type: protobuf::JoinType = exec.join_type().to_owned().into(); - - let partition_mode = match exec.partition_mode() { - PartitionMode::CollectLeft => protobuf::PartitionMode::CollectLeft, - PartitionMode::Partitioned => protobuf::PartitionMode::Partitioned, - }; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::HashJoin(Box::new( - protobuf::HashJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - on, - join_type: join_type.into(), - partition_mode: partition_mode.into(), - null_equals_null: *exec.null_equals_null(), - }, - ))), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.left().to_owned(), - extension_codec, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.right().to_owned(), - extension_codec, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CrossJoin(Box::new( - protobuf::CrossJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - }, - ))), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let groups = exec - .group_expr() - .iter() - .map(|expr| expr.0.to_owned().try_into()) - .collect::, BallistaError>>()?; - let group_names = exec - .group_expr() - .iter() - .map(|expr| expr.1.to_owned()) - .collect(); - let agg = exec - .aggr_expr() - .iter() - .map(|expr| expr.to_owned().try_into()) - .collect::, BallistaError>>()?; - let agg_names = exec - .aggr_expr() - .iter() - .map(|expr| match expr.field() { - Ok(field) => Ok(field.name().clone()), - Err(e) => Err(BallistaError::DataFusionError(e)), - }) - .collect::>()?; - - let agg_mode = match exec.mode() { - AggregateMode::Partial => protobuf::AggregateMode::Partial, - AggregateMode::Final => protobuf::AggregateMode::Final, - AggregateMode::FinalPartitioned => { - protobuf::AggregateMode::FinalPartitioned - } - }; - let input_schema = exec.input_schema(); - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.input().to_owned(), - extension_codec, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Aggregate(Box::new( - protobuf::AggregateExecNode { - group_expr: groups, - group_expr_name: group_names, - aggr_expr: agg, - aggr_expr_name: agg_names, - mode: agg_mode as i32, - input: Some(Box::new(input)), - input_schema: Some(input_schema.as_ref().into()), - }, - ))), - }) - } else if let Some(empty) = plan.downcast_ref::() { - let schema = empty.schema().as_ref().into(); - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Empty( - protobuf::EmptyExecNode { - produce_one_row: empty.produce_one_row(), - schema: Some(schema), - }, - )), - }) - } else if let Some(coalesce_batches) = plan.downcast_ref::() - { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - coalesce_batches.input().to_owned(), - extension_codec, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CoalesceBatches(Box::new( - protobuf::CoalesceBatchesExecNode { - input: Some(Box::new(input)), - target_batch_size: coalesce_batches.target_batch_size() as u32, - }, - ))), - }) - } else if let Some(exec) = plan.downcast_ref::() { - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CsvScan( - protobuf::CsvScanExecNode { - base_conf: Some(exec.base_config().try_into()?), - has_header: exec.has_header(), - delimiter: byte_to_string(exec.delimiter())?, - }, - )), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let pruning_expr = exec - .pruning_predicate() - .map(|pred| pred.logical_expr().try_into()) - .transpose()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ParquetScan( - protobuf::ParquetScanExecNode { - base_conf: Some(exec.base_config().try_into()?), - pruning_predicate: pruning_expr, - }, - )), - }) - } else if let Some(exec) = plan.downcast_ref::() { - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::AvroScan( - protobuf::AvroScanExecNode { - base_conf: Some(exec.base_config().try_into()?), - }, - )), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let mut partition = vec![]; - for location in &exec.partition { - partition.push(protobuf::ShuffleReaderPartition { - location: location - .iter() - .map(|l| l.clone().try_into()) - .collect::, _>>()?, - }); - } - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ShuffleReader( - protobuf::ShuffleReaderExecNode { - partition, - schema: Some(exec.schema().as_ref().into()), - }, - )), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.input().to_owned(), - extension_codec, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Merge(Box::new( - protobuf::CoalescePartitionsExecNode { - input: Some(Box::new(input)), - }, - ))), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.input().to_owned(), - extension_codec, - )?; - - let pb_partition_method = match exec.partitioning() { - Partitioning::Hash(exprs, partition_count) => { - PartitionMethod::Hash(protobuf::PhysicalHashRepartition { - hash_expr: exprs - .iter() - .map(|expr| expr.clone().try_into()) - .collect::, BallistaError>>()?, - partition_count: *partition_count as u64, - }) - } - Partitioning::RoundRobinBatch(partition_count) => { - PartitionMethod::RoundRobin(*partition_count as u64) - } - Partitioning::UnknownPartitioning(partition_count) => { - PartitionMethod::Unknown(*partition_count as u64) - } - }; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Repartition(Box::new( - protobuf::RepartitionExecNode { - input: Some(Box::new(input)), - partition_method: Some(pb_partition_method), - }, - ))), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.input().to_owned(), - extension_codec, - )?; - let expr = exec - .expr() - .iter() - .map(|expr| { - let sort_expr = Box::new(protobuf::PhysicalSortExprNode { - expr: Some(Box::new(expr.expr.to_owned().try_into()?)), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }); - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::Sort( - sort_expr, - )), - }) - }) - .collect::, BallistaError>>()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Sort(Box::new( - protobuf::SortExecNode { - input: Some(Box::new(input)), - expr, - }, - ))), - }) - } else if let Some(exec) = plan.downcast_ref::() { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan( - exec.children()[0].to_owned(), - extension_codec, - )?; - // note that we use shuffle_output_partitioning() rather than output_partitioning() - // to get the true output partitioning - let output_partitioning = match exec.shuffle_output_partitioning() { - Some(Partitioning::Hash(exprs, partition_count)) => { - Some(protobuf::PhysicalHashRepartition { - hash_expr: exprs - .iter() - .map(|expr| expr.clone().try_into()) - .collect::, BallistaError>>()?, - partition_count: *partition_count as u64, - }) - } - None => None, - other => { - return Err(BallistaError::General(format!( - "physical_plan::to_proto() invalid partitioning for ShuffleWriterExec: {:?}", - other - ))) - } - }; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ShuffleWriter(Box::new( - protobuf::ShuffleWriterExecNode { - job_id: exec.job_id().to_string(), - stage_id: exec.stage_id() as u32, - input: Some(Box::new(input)), - output_partitioning, - }, - ))), - }) - } else if let Some(exec) = plan.downcast_ref::() { - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Unresolved( - protobuf::UnresolvedShuffleExecNode { - stage_id: exec.stage_id as u32, - schema: Some(exec.schema().as_ref().into()), - input_partition_count: exec.input_partition_count as u32, - output_partition_count: exec.output_partition_count as u32, - }, - )), - }) - } else if let Some(union) = plan.downcast_ref::() { - let mut inputs: Vec = vec![]; - for input in union.inputs() { - inputs.push(protobuf::PhysicalPlanNode::try_from_physical_plan( - input.to_owned(), - extension_codec, - )?); - } - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Union( - protobuf::UnionExecNode { inputs }, - )), - }) - } else { - let mut buf: Vec = vec![]; - extension_codec.try_encode(plan_clone.clone(), &mut buf)?; - - let inputs: Vec = plan_clone - .children() - .into_iter() - .map(|i| PhysicalPlanNode::try_from_physical_plan(i, extension_codec)) - .collect::>()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Extension( - PhysicalExtensionNode { node: buf, inputs }, - )), - }) - } - } -} - -fn decode_scan_config( - proto: &protobuf::FileScanExecConf, - runtime: &RuntimeEnv, -) -> Result { - let schema = Arc::new(convert_required!(proto.schema)?); - let projection = proto - .projection - .iter() - .map(|i| *i as usize) - .collect::>(); - let projection = if projection.is_empty() { - None - } else { - Some(projection) - }; - let statistics = convert_required!(proto.statistics)?; - - let file_groups: Vec> = proto - .file_groups - .iter() - .map(|f| f.try_into()) - .collect::, _>>()?; - - let object_store = if let Some(file) = file_groups.get(0).and_then(|h| h.get(0)) { - runtime.object_store(file.file_meta.path())?.0 - } else { - Arc::new(LocalFileSystem {}) - }; - - Ok(FileScanConfig { - object_store, - file_schema: schema, - file_groups, - statistics, - projection, - limit: proto.limit.as_ref().map(|sl| sl.limit as usize), - table_partition_cols: vec![], - }) -} - -#[macro_export] -macro_rules! into_physical_plan { - ($PB:expr, $REG:expr, $RUNTIME:expr, $CODEC:expr) => {{ - if let Some(field) = $PB.as_ref() { - field - .as_ref() - .try_into_physical_plan($REG, $RUNTIME, $CODEC) - } else { - Err(proto_error("Missing required field in protobuf")) - } - }}; -} - -#[cfg(test)] -mod roundtrip_tests { - use std::ops::Deref; - use std::sync::Arc; - - use datafusion::arrow::array::ArrayRef; - use datafusion::execution::context::ExecutionProps; - use datafusion::logical_expr::{BuiltinScalarFunction, Volatility}; - use datafusion::logical_plan::create_udf; - use datafusion::physical_plan::functions; - use datafusion::physical_plan::functions::{ - make_scalar_function, ScalarFunctionExpr, - }; - use datafusion::physical_plan::projection::ProjectionExec; - use datafusion::{ - arrow::{ - compute::kernels::sort::SortOptions, - datatypes::{DataType, Field, Schema}, - }, - datafusion_data_access::object_store::local::LocalFileSystem, - datasource::listing::PartitionedFile, - logical_plan::{JoinType, Operator}, - physical_plan::{ - aggregates::{AggregateExec, AggregateMode}, - empty::EmptyExec, - expressions::{binary, col, lit, InListExpr, NotExpr}, - expressions::{Avg, Column, PhysicalSortExpr}, - file_format::{FileScanConfig, ParquetExec}, - filter::FilterExec, - hash_join::{HashJoinExec, PartitionMode}, - limit::{GlobalLimitExec, LocalLimitExec}, - sorts::sort::SortExec, - AggregateExpr, ExecutionPlan, Partitioning, PhysicalExpr, Statistics, - }, - prelude::SessionContext, - scalar::ScalarValue, - }; - - use crate::execution_plans::ShuffleWriterExec; - use crate::serde::protobuf::{LogicalPlanNode, PhysicalPlanNode}; - use crate::serde::{AsExecutionPlan, BallistaCodec}; - - use super::super::super::error::Result; - use super::super::protobuf; - - fn roundtrip_test(exec_plan: Arc) -> Result<()> { - let ctx = SessionContext::new(); - let codec: BallistaCodec = - BallistaCodec::default(); - let proto: protobuf::PhysicalPlanNode = - protobuf::PhysicalPlanNode::try_from_physical_plan( - exec_plan.clone(), - codec.physical_extension_codec(), - ) - .expect("to proto"); - let runtime = ctx.runtime_env(); - let result_exec_plan: Arc = proto - .try_into_physical_plan( - &ctx, - runtime.deref(), - codec.physical_extension_codec(), - ) - .expect("from proto"); - assert_eq!( - format!("{:?}", exec_plan), - format!("{:?}", result_exec_plan) - ); - Ok(()) - } - - fn roundtrip_test_with_context( - exec_plan: Arc, - ctx: SessionContext, - ) -> Result<()> { - let codec: BallistaCodec = - BallistaCodec::default(); - let proto: protobuf::PhysicalPlanNode = - protobuf::PhysicalPlanNode::try_from_physical_plan( - exec_plan.clone(), - codec.physical_extension_codec(), - ) - .expect("to proto"); - let runtime = ctx.runtime_env(); - let result_exec_plan: Arc = proto - .try_into_physical_plan( - &ctx, - runtime.deref(), - codec.physical_extension_codec(), - ) - .expect("from proto"); - assert_eq!( - format!("{:?}", exec_plan), - format!("{:?}", result_exec_plan) - ); - Ok(()) - } - - #[test] - fn roundtrip_empty() -> Result<()> { - roundtrip_test(Arc::new(EmptyExec::new(false, Arc::new(Schema::empty())))) - } - - #[test] - fn roundtrip_local_limit() -> Result<()> { - roundtrip_test(Arc::new(LocalLimitExec::new( - Arc::new(EmptyExec::new(false, Arc::new(Schema::empty()))), - 25, - ))) - } - - #[test] - fn roundtrip_global_limit() -> Result<()> { - roundtrip_test(Arc::new(GlobalLimitExec::new( - Arc::new(EmptyExec::new(false, Arc::new(Schema::empty()))), - 25, - ))) - } - - #[test] - fn roundtrip_hash_join() -> Result<()> { - let field_a = Field::new("col", DataType::Int64, false); - let schema_left = Schema::new(vec![field_a.clone()]); - let schema_right = Schema::new(vec![field_a]); - let on = vec![( - Column::new("col", schema_left.index_of("col")?), - Column::new("col", schema_right.index_of("col")?), - )]; - - let schema_left = Arc::new(schema_left); - let schema_right = Arc::new(schema_right); - for join_type in &[ - JoinType::Inner, - JoinType::Left, - JoinType::Right, - JoinType::Full, - JoinType::Anti, - JoinType::Semi, - ] { - for partition_mode in - &[PartitionMode::Partitioned, PartitionMode::CollectLeft] - { - roundtrip_test(Arc::new(HashJoinExec::try_new( - Arc::new(EmptyExec::new(false, schema_left.clone())), - Arc::new(EmptyExec::new(false, schema_right.clone())), - on.clone(), - join_type, - *partition_mode, - &false, - )?))?; - } - } - Ok(()) - } - - #[test] - fn rountrip_aggregate() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let groups: Vec<(Arc, String)> = - vec![(col("a", &schema)?, "unused".to_string())]; - - let aggregates: Vec> = vec![Arc::new(Avg::new( - col("b", &schema)?, - "AVG(b)".to_string(), - DataType::Float64, - ))]; - - roundtrip_test(Arc::new(AggregateExec::try_new( - AggregateMode::Final, - groups.clone(), - aggregates.clone(), - Arc::new(EmptyExec::new(false, schema.clone())), - schema, - )?)) - } - - #[test] - fn roundtrip_filter_with_not_and_in_list() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Int64, false); - let field_c = Field::new("c", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b, field_c])); - let not = Arc::new(NotExpr::new(col("a", &schema)?)); - let in_list = Arc::new(InListExpr::new( - col("b", &schema)?, - vec![ - lit(ScalarValue::Int64(Some(1))), - lit(ScalarValue::Int64(Some(2))), - ], - false, - )); - let and = binary(not, Operator::And, in_list, &schema)?; - roundtrip_test(Arc::new(FilterExec::try_new( - and, - Arc::new(EmptyExec::new(false, schema.clone())), - )?)) - } - - #[test] - fn roundtrip_sort() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let sort_exprs = vec![ - PhysicalSortExpr { - expr: col("a", &schema)?, - options: SortOptions { - descending: true, - nulls_first: false, - }, - }, - PhysicalSortExpr { - expr: col("b", &schema)?, - options: SortOptions { - descending: false, - nulls_first: true, - }, - }, - ]; - roundtrip_test(Arc::new(SortExec::try_new( - sort_exprs, - Arc::new(EmptyExec::new(false, schema)), - )?)) - } - - #[test] - fn roundtrip_shuffle_writer() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - roundtrip_test(Arc::new(ShuffleWriterExec::try_new( - "job123".to_string(), - 123, - Arc::new(EmptyExec::new(false, schema)), - "".to_string(), - Some(Partitioning::Hash(vec![Arc::new(Column::new("a", 0))], 4)), - )?)) - } - - #[test] - fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { - let scan_config = FileScanConfig { - object_store: Arc::new(LocalFileSystem {}), - file_schema: Arc::new(Schema::new(vec![Field::new( - "col", - DataType::Utf8, - false, - )])), - file_groups: vec![vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )]], - statistics: Statistics { - num_rows: Some(100), - total_byte_size: Some(1024), - column_statistics: None, - is_exact: false, - }, - projection: None, - limit: None, - table_partition_cols: vec![], - }; - - let predicate = datafusion::prelude::col("col").eq(datafusion::prelude::lit("1")); - roundtrip_test(Arc::new(ParquetExec::new(scan_config, Some(predicate)))) - } - - #[test] - fn roundtrip_builtin_scalar_function() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let input = Arc::new(EmptyExec::new(false, schema.clone())); - - let execution_props = ExecutionProps::new(); - - let fun_expr = functions::create_physical_fun( - &BuiltinScalarFunction::Abs, - &execution_props, - )?; - - let expr = ScalarFunctionExpr::new( - "abs", - fun_expr, - vec![col("a", &schema)?], - &DataType::Int64, - ); - - let project = - ProjectionExec::try_new(vec![(Arc::new(expr), "a".to_string())], input)?; - - roundtrip_test(Arc::new(project)) - } - - #[test] - fn roundtrip_scalar_udf() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let input = Arc::new(EmptyExec::new(false, schema.clone())); - - let fn_impl = |args: &[ArrayRef]| Ok(Arc::new(args[0].clone()) as ArrayRef); - - let scalar_fn = make_scalar_function(fn_impl); - - let udf = create_udf( - "dummy", - vec![DataType::Int64], - Arc::new(DataType::Int64), - Volatility::Immutable, - scalar_fn.clone(), - ); - - let expr = ScalarFunctionExpr::new( - "dummy", - scalar_fn, - vec![col("a", &schema)?], - &DataType::Int64, - ); - - let project = - ProjectionExec::try_new(vec![(Arc::new(expr), "a".to_string())], input)?; - - let mut ctx = SessionContext::new(); - - ctx.register_udf(udf); - - roundtrip_test_with_context(Arc::new(project), ctx) - } -} diff --git a/ballista/rust/core/src/serde/physical_plan/to_proto.rs b/ballista/rust/core/src/serde/physical_plan/to_proto.rs deleted file mode 100644 index 1b96a8556c0b..000000000000 --- a/ballista/rust/core/src/serde/physical_plan/to_proto.rs +++ /dev/null @@ -1,446 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License.language governing permissions and -// limitations under the License. - -//! Serde code to convert Arrow schemas and DataFusion logical plans to Ballista protocol -//! buffer format, allowing DataFusion physical plans to be serialized and transmitted between -//! processes. - -use std::{ - convert::{TryFrom, TryInto}, - str::FromStr, - sync::Arc, -}; - -use datafusion::physical_plan::expressions::{CastExpr, TryCastExpr}; -use datafusion::physical_plan::ColumnStatistics; -use datafusion::physical_plan::{ - expressions::{ - CaseExpr, InListExpr, IsNotNullExpr, IsNullExpr, NegativeExpr, NotExpr, - }, - Statistics, -}; - -use datafusion::datasource::listing::{FileRange, PartitionedFile}; -use datafusion::physical_plan::file_format::FileScanConfig; - -use datafusion::physical_plan::expressions::{Count, Literal}; - -use datafusion::physical_plan::expressions::{Avg, BinaryExpr, Column, Max, Min, Sum}; -use datafusion::physical_plan::{AggregateExpr, PhysicalExpr}; - -use crate::serde::{protobuf, BallistaError}; - -use datafusion::logical_expr::BuiltinScalarFunction; -use datafusion::physical_plan::functions::ScalarFunctionExpr; - -impl TryInto for Arc { - type Error = BallistaError; - - fn try_into(self) -> Result { - use datafusion::physical_plan::expressions; - use datafusion_proto::protobuf::AggregateFunction; - let aggr_function = if self.as_any().downcast_ref::().is_some() { - Ok(AggregateFunction::Avg.into()) - } else if self.as_any().downcast_ref::().is_some() { - Ok(AggregateFunction::Sum.into()) - } else if self.as_any().downcast_ref::().is_some() { - Ok(AggregateFunction::Count.into()) - } else if self.as_any().downcast_ref::().is_some() { - Ok(AggregateFunction::Min.into()) - } else if self.as_any().downcast_ref::().is_some() { - Ok(AggregateFunction::Max.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::ApproxDistinct.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::ArrayAgg.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::Variance.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::VariancePop.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::Covariance.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::CovariancePop.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::Stddev.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::StddevPop.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::Correlation.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::ApproxPercentileCont.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::ApproxPercentileContWithWeight.into()) - } else if self - .as_any() - .downcast_ref::() - .is_some() - { - Ok(AggregateFunction::ApproxMedian.into()) - } else { - Err(BallistaError::NotImplemented(format!( - "Aggregate function not supported: {:?}", - self - ))) - }?; - let expressions: Vec = self - .expressions() - .iter() - .map(|e| e.clone().try_into()) - .collect::, BallistaError>>()?; - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::AggregateExpr( - protobuf::PhysicalAggregateExprNode { - aggr_function, - expr: expressions, - }, - )), - }) - } -} - -impl TryFrom> for protobuf::PhysicalExprNode { - type Error = BallistaError; - - fn try_from(value: Arc) -> Result { - let expr = value.as_any(); - - if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::Column( - protobuf::PhysicalColumn { - name: expr.name().to_string(), - index: expr.index() as u32, - }, - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - let binary_expr = Box::new(protobuf::PhysicalBinaryExprNode { - l: Some(Box::new(expr.left().to_owned().try_into()?)), - r: Some(Box::new(expr.right().to_owned().try_into()?)), - op: format!("{:?}", expr.op()), - }); - - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::BinaryExpr( - binary_expr, - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some( - protobuf::physical_expr_node::ExprType::Case( - Box::new( - protobuf::PhysicalCaseNode { - expr: expr - .expr() - .as_ref() - .map(|exp| exp.clone().try_into().map(Box::new)) - .transpose()?, - when_then_expr: expr - .when_then_expr() - .iter() - .map(|(when_expr, then_expr)| { - try_parse_when_then_expr(when_expr, then_expr) - }) - .collect::, - Self::Error, - >>()?, - else_expr: expr - .else_expr() - .map(|a| a.clone().try_into().map(Box::new)) - .transpose()?, - }, - ), - ), - ), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::NotExpr( - Box::new(protobuf::PhysicalNot { - expr: Some(Box::new(expr.arg().to_owned().try_into()?)), - }), - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::IsNullExpr( - Box::new(protobuf::PhysicalIsNull { - expr: Some(Box::new(expr.arg().to_owned().try_into()?)), - }), - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr( - Box::new(protobuf::PhysicalIsNotNull { - expr: Some(Box::new(expr.arg().to_owned().try_into()?)), - }), - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some( - protobuf::physical_expr_node::ExprType::InList( - Box::new( - protobuf::PhysicalInListNode { - expr: Some(Box::new(expr.expr().to_owned().try_into()?)), - list: expr - .list() - .iter() - .map(|a| a.clone().try_into()) - .collect::, - Self::Error, - >>()?, - negated: expr.negated(), - }, - ), - ), - ), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::Negative( - Box::new(protobuf::PhysicalNegativeNode { - expr: Some(Box::new(expr.arg().to_owned().try_into()?)), - }), - )), - }) - } else if let Some(lit) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::Literal( - lit.value().try_into()?, - )), - }) - } else if let Some(cast) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::Cast(Box::new( - protobuf::PhysicalCastNode { - expr: Some(Box::new(cast.expr().clone().try_into()?)), - arrow_type: Some(cast.cast_type().into()), - }, - ))), - }) - } else if let Some(cast) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::TryCast( - Box::new(protobuf::PhysicalTryCastNode { - expr: Some(Box::new(cast.expr().clone().try_into()?)), - arrow_type: Some(cast.cast_type().into()), - }), - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - let args: Vec = expr - .args() - .iter() - .map(|e| e.to_owned().try_into()) - .collect::, _>>()?; - if let Ok(fun) = BuiltinScalarFunction::from_str(expr.name()) { - let fun: datafusion_proto::protobuf::ScalarFunction = - (&fun).try_into()?; - - Ok(protobuf::PhysicalExprNode { - expr_type: Some( - protobuf::physical_expr_node::ExprType::ScalarFunction( - protobuf::PhysicalScalarFunctionNode { - name: expr.name().to_string(), - fun: fun.into(), - args, - return_type: Some(expr.return_type().into()), - }, - ), - ), - }) - } else { - Ok(protobuf::PhysicalExprNode { - expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarUdf( - protobuf::PhysicalScalarUdfNode { - name: expr.name().to_string(), - args, - return_type: Some(expr.return_type().into()), - }, - )), - }) - } - } else { - Err(BallistaError::General(format!( - "physical_plan::to_proto() unsupported expression {:?}", - value - ))) - } - } -} - -fn try_parse_when_then_expr( - when_expr: &Arc, - then_expr: &Arc, -) -> Result { - Ok(protobuf::PhysicalWhenThen { - when_expr: Some(when_expr.clone().try_into()?), - then_expr: Some(then_expr.clone().try_into()?), - }) -} - -impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile { - type Error = BallistaError; - - fn try_from(pf: &PartitionedFile) -> Result { - Ok(protobuf::PartitionedFile { - path: pf.file_meta.path().to_owned(), - size: pf.file_meta.size(), - last_modified_ns: pf - .file_meta - .last_modified - .map(|ts| ts.timestamp_nanos() as u64) - .unwrap_or(0), - partition_values: pf - .partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - range: pf.range.as_ref().map(|r| r.try_into()).transpose()?, - }) - } -} - -impl TryFrom<&FileRange> for protobuf::FileRange { - type Error = BallistaError; - - fn try_from(value: &FileRange) -> Result { - Ok(protobuf::FileRange { - start: value.start, - end: value.end, - }) - } -} - -impl TryFrom<&[PartitionedFile]> for protobuf::FileGroup { - type Error = BallistaError; - - fn try_from(gr: &[PartitionedFile]) -> Result { - Ok(protobuf::FileGroup { - files: gr - .iter() - .map(|f| f.try_into()) - .collect::, _>>()?, - }) - } -} - -impl From<&ColumnStatistics> for protobuf::ColumnStats { - fn from(cs: &ColumnStatistics) -> protobuf::ColumnStats { - protobuf::ColumnStats { - min_value: cs.min_value.as_ref().map(|m| m.try_into().unwrap()), - max_value: cs.max_value.as_ref().map(|m| m.try_into().unwrap()), - null_count: cs.null_count.map(|n| n as u32).unwrap_or(0), - distinct_count: cs.distinct_count.map(|n| n as u32).unwrap_or(0), - } - } -} - -impl From<&Statistics> for protobuf::Statistics { - fn from(s: &Statistics) -> protobuf::Statistics { - let none_value = -1_i64; - let column_stats = match &s.column_statistics { - None => vec![], - Some(column_stats) => column_stats.iter().map(|s| s.into()).collect(), - }; - protobuf::Statistics { - num_rows: s.num_rows.map(|n| n as i64).unwrap_or(none_value), - total_byte_size: s.total_byte_size.map(|n| n as i64).unwrap_or(none_value), - column_stats, - is_exact: s.is_exact, - } - } -} - -impl TryFrom<&FileScanConfig> for protobuf::FileScanExecConf { - type Error = BallistaError; - fn try_from( - conf: &FileScanConfig, - ) -> Result { - let file_groups = conf - .file_groups - .iter() - .map(|p| p.as_slice().try_into()) - .collect::, _>>()?; - - Ok(protobuf::FileScanExecConf { - file_groups, - statistics: Some((&conf.statistics).into()), - limit: conf.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }), - projection: conf - .projection - .as_ref() - .unwrap_or(&vec![]) - .iter() - .map(|n| *n as u32) - .collect(), - schema: Some(conf.file_schema.as_ref().into()), - table_partition_cols: conf.table_partition_cols.to_vec(), - }) - } -} diff --git a/ballista/rust/core/src/serde/scheduler/from_proto.rs b/ballista/rust/core/src/serde/scheduler/from_proto.rs deleted file mode 100644 index b401f1fdf789..000000000000 --- a/ballista/rust/core/src/serde/scheduler/from_proto.rs +++ /dev/null @@ -1,106 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::convert::TryInto; - -use crate::error::BallistaError; -use crate::serde::protobuf; -use crate::serde::protobuf::action::ActionType; -use crate::serde::scheduler::{Action, PartitionId, PartitionLocation, PartitionStats}; - -impl TryInto for protobuf::Action { - type Error = BallistaError; - - fn try_into(self) -> Result { - match self.action_type { - Some(ActionType::FetchPartition(fetch)) => Ok(Action::FetchPartition { - job_id: fetch.job_id, - stage_id: fetch.stage_id as usize, - partition_id: fetch.partition_id as usize, - path: fetch.path, - }), - _ => Err(BallistaError::General( - "scheduler::from_proto(Action) invalid or missing action".to_owned(), - )), - } - } -} - -impl TryInto for protobuf::PartitionId { - type Error = BallistaError; - - fn try_into(self) -> Result { - Ok(PartitionId::new( - &self.job_id, - self.stage_id as usize, - self.partition_id as usize, - )) - } -} - -#[allow(clippy::from_over_into)] -impl Into for protobuf::PartitionStats { - fn into(self) -> PartitionStats { - PartitionStats::new( - foo(self.num_rows), - foo(self.num_batches), - foo(self.num_bytes), - ) - } -} - -fn foo(n: i64) -> Option { - if n < 0 { - None - } else { - Some(n as u64) - } -} - -impl TryInto for protobuf::PartitionLocation { - type Error = BallistaError; - - fn try_into(self) -> Result { - Ok(PartitionLocation { - partition_id: self - .partition_id - .ok_or_else(|| { - BallistaError::General( - "partition_id in PartitionLocation is missing.".to_owned(), - ) - })? - .try_into()?, - executor_meta: self - .executor_meta - .ok_or_else(|| { - BallistaError::General( - "executor_meta in PartitionLocation is missing".to_owned(), - ) - })? - .into(), - partition_stats: self - .partition_stats - .ok_or_else(|| { - BallistaError::General( - "partition_stats in PartitionLocation is missing".to_owned(), - ) - })? - .into(), - path: self.path, - }) - } -} diff --git a/ballista/rust/core/src/serde/scheduler/mod.rs b/ballista/rust/core/src/serde/scheduler/mod.rs deleted file mode 100644 index 369a87d2255b..000000000000 --- a/ballista/rust/core/src/serde/scheduler/mod.rs +++ /dev/null @@ -1,423 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::{collections::HashMap, fmt, sync::Arc}; - -use datafusion::arrow::array::{ - ArrayBuilder, StructArray, StructBuilder, UInt64Array, UInt64Builder, -}; -use datafusion::arrow::datatypes::{DataType, Field}; - -use datafusion::physical_plan::ExecutionPlan; -use datafusion::physical_plan::Partitioning; -use serde::Serialize; - -use super::protobuf; -use crate::error::BallistaError; - -pub mod from_proto; -pub mod to_proto; - -/// Action that can be sent to an executor -#[derive(Debug, Clone)] -pub enum Action { - /// Collect a shuffle partition - FetchPartition { - job_id: String, - stage_id: usize, - partition_id: usize, - path: String, - }, -} - -/// Unique identifier for the output partition of an operator. -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] -pub struct PartitionId { - pub job_id: String, - pub stage_id: usize, - pub partition_id: usize, -} - -impl PartitionId { - pub fn new(job_id: &str, stage_id: usize, partition_id: usize) -> Self { - Self { - job_id: job_id.to_string(), - stage_id, - partition_id, - } - } -} - -#[derive(Debug, Clone)] -pub struct PartitionLocation { - pub partition_id: PartitionId, - pub executor_meta: ExecutorMetadata, - pub partition_stats: PartitionStats, - pub path: String, -} - -/// Meta-data for an executor, used when fetching shuffle partitions from other executors -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] -pub struct ExecutorMetadata { - pub id: String, - pub host: String, - pub port: u16, - pub grpc_port: u16, - pub specification: ExecutorSpecification, -} - -#[allow(clippy::from_over_into)] -impl Into for ExecutorMetadata { - fn into(self) -> protobuf::ExecutorMetadata { - protobuf::ExecutorMetadata { - id: self.id, - host: self.host, - port: self.port as u32, - grpc_port: self.grpc_port as u32, - specification: Some(self.specification.into()), - } - } -} - -impl From for ExecutorMetadata { - fn from(meta: protobuf::ExecutorMetadata) -> Self { - Self { - id: meta.id, - host: meta.host, - port: meta.port as u16, - grpc_port: meta.grpc_port as u16, - specification: meta.specification.unwrap().into(), - } - } -} - -/// Specification of an executor, indicting executor resources, like total task slots -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] -pub struct ExecutorSpecification { - pub task_slots: u32, -} - -#[allow(clippy::from_over_into)] -impl Into for ExecutorSpecification { - fn into(self) -> protobuf::ExecutorSpecification { - protobuf::ExecutorSpecification { - resources: vec![protobuf::executor_resource::Resource::TaskSlots( - self.task_slots, - )] - .into_iter() - .map(|r| protobuf::ExecutorResource { resource: Some(r) }) - .collect(), - } - } -} - -impl From for ExecutorSpecification { - fn from(input: protobuf::ExecutorSpecification) -> Self { - let mut ret = Self { task_slots: 0 }; - for resource in input.resources { - if let Some(protobuf::executor_resource::Resource::TaskSlots(task_slots)) = - resource.resource - { - ret.task_slots = task_slots - } - } - ret - } -} - -/// From Spark, available resources for an executor, like available task slots -#[derive(Debug, Clone, Serialize)] -pub struct ExecutorData { - pub executor_id: String, - pub total_task_slots: u32, - pub available_task_slots: u32, -} - -pub struct ExecutorDataChange { - pub executor_id: String, - pub task_slots: i32, -} - -struct ExecutorResourcePair { - total: protobuf::executor_resource::Resource, - available: protobuf::executor_resource::Resource, -} - -#[allow(clippy::from_over_into)] -impl Into for ExecutorData { - fn into(self) -> protobuf::ExecutorData { - protobuf::ExecutorData { - executor_id: self.executor_id, - resources: vec![ExecutorResourcePair { - total: protobuf::executor_resource::Resource::TaskSlots( - self.total_task_slots, - ), - available: protobuf::executor_resource::Resource::TaskSlots( - self.available_task_slots, - ), - }] - .into_iter() - .map(|r| protobuf::ExecutorResourcePair { - total: Some(protobuf::ExecutorResource { - resource: Some(r.total), - }), - available: Some(protobuf::ExecutorResource { - resource: Some(r.available), - }), - }) - .collect(), - } - } -} - -impl From for ExecutorData { - fn from(input: protobuf::ExecutorData) -> Self { - let mut ret = Self { - executor_id: input.executor_id, - total_task_slots: 0, - available_task_slots: 0, - }; - for resource in input.resources { - if let Some(task_slots) = resource.total { - if let Some(protobuf::executor_resource::Resource::TaskSlots( - task_slots, - )) = task_slots.resource - { - ret.total_task_slots = task_slots - } - }; - if let Some(task_slots) = resource.available { - if let Some(protobuf::executor_resource::Resource::TaskSlots( - task_slots, - )) = task_slots.resource - { - ret.available_task_slots = task_slots - } - }; - } - ret - } -} - -/// The internal state of an executor, like cpu usage, memory usage, etc -#[derive(Debug, Clone, Copy, Serialize)] -pub struct ExecutorState { - // in bytes - pub available_memory_size: u64, -} - -#[allow(clippy::from_over_into)] -impl Into for ExecutorState { - fn into(self) -> protobuf::ExecutorState { - protobuf::ExecutorState { - metrics: vec![protobuf::executor_metric::Metric::AvailableMemory( - self.available_memory_size, - )] - .into_iter() - .map(|m| protobuf::ExecutorMetric { metric: Some(m) }) - .collect(), - } - } -} - -impl From for ExecutorState { - fn from(input: protobuf::ExecutorState) -> Self { - let mut ret = Self { - available_memory_size: u64::MAX, - }; - for metric in input.metrics { - if let Some(protobuf::executor_metric::Metric::AvailableMemory( - available_memory_size, - )) = metric.metric - { - ret.available_memory_size = available_memory_size - } - } - ret - } -} - -/// Summary of executed partition -#[derive(Debug, Copy, Clone, Default)] -pub struct PartitionStats { - pub(crate) num_rows: Option, - pub(crate) num_batches: Option, - pub(crate) num_bytes: Option, -} - -impl fmt::Display for PartitionStats { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - write!( - f, - "numBatches={:?}, numRows={:?}, numBytes={:?}", - self.num_batches, self.num_rows, self.num_bytes - ) - } -} - -impl PartitionStats { - pub fn new( - num_rows: Option, - num_batches: Option, - num_bytes: Option, - ) -> Self { - Self { - num_rows, - num_batches, - num_bytes, - } - } - - pub fn arrow_struct_repr(self) -> Field { - Field::new( - "partition_stats", - DataType::Struct(self.arrow_struct_fields()), - false, - ) - } - - pub fn arrow_struct_fields(self) -> Vec { - vec![ - Field::new("num_rows", DataType::UInt64, false), - Field::new("num_batches", DataType::UInt64, false), - Field::new("num_bytes", DataType::UInt64, false), - ] - } - - pub fn to_arrow_arrayref(self) -> Result, BallistaError> { - let mut field_builders = Vec::new(); - - let mut num_rows_builder = UInt64Builder::new(1); - match self.num_rows { - Some(n) => num_rows_builder.append_value(n)?, - None => num_rows_builder.append_null()?, - } - field_builders.push(Box::new(num_rows_builder) as Box); - - let mut num_batches_builder = UInt64Builder::new(1); - match self.num_batches { - Some(n) => num_batches_builder.append_value(n)?, - None => num_batches_builder.append_null()?, - } - field_builders.push(Box::new(num_batches_builder) as Box); - - let mut num_bytes_builder = UInt64Builder::new(1); - match self.num_bytes { - Some(n) => num_bytes_builder.append_value(n)?, - None => num_bytes_builder.append_null()?, - } - field_builders.push(Box::new(num_bytes_builder) as Box); - - let mut struct_builder = - StructBuilder::new(self.arrow_struct_fields(), field_builders); - struct_builder.append(true)?; - Ok(Arc::new(struct_builder.finish())) - } - - pub fn from_arrow_struct_array(struct_array: &StructArray) -> PartitionStats { - let num_rows = struct_array - .column_by_name("num_rows") - .expect("from_arrow_struct_array expected a field num_rows") - .as_any() - .downcast_ref::() - .expect("from_arrow_struct_array expected num_rows to be a UInt64Array"); - let num_batches = struct_array - .column_by_name("num_batches") - .expect("from_arrow_struct_array expected a field num_batches") - .as_any() - .downcast_ref::() - .expect("from_arrow_struct_array expected num_batches to be a UInt64Array"); - let num_bytes = struct_array - .column_by_name("num_bytes") - .expect("from_arrow_struct_array expected a field num_bytes") - .as_any() - .downcast_ref::() - .expect("from_arrow_struct_array expected num_bytes to be a UInt64Array"); - PartitionStats { - num_rows: Some(num_rows.value(0).to_owned()), - num_batches: Some(num_batches.value(0).to_owned()), - num_bytes: Some(num_bytes.value(0).to_owned()), - } - } -} - -/// Task that can be sent to an executor to execute one stage of a query and write -/// results out to disk -#[derive(Debug, Clone)] -pub struct ExecutePartition { - /// Unique ID representing this query execution - pub job_id: String, - /// Unique ID representing this query stage within the overall query - pub stage_id: usize, - /// The partitions to execute. The same plan could be sent to multiple executors and each - /// executor will execute a range of partitions per QueryStageTask - pub partition_id: Vec, - /// The physical plan for this query stage - pub plan: Arc, - /// Location of shuffle partitions that this query stage may depend on - pub shuffle_locations: HashMap, - /// Output partitioning for shuffle writes - pub output_partitioning: Option, -} - -impl ExecutePartition { - pub fn new( - job_id: String, - stage_id: usize, - partition_id: Vec, - plan: Arc, - shuffle_locations: HashMap, - output_partitioning: Option, - ) -> Self { - Self { - job_id, - stage_id, - partition_id, - plan, - shuffle_locations, - output_partitioning, - } - } - - pub fn key(&self) -> String { - format!("{}.{}.{:?}", self.job_id, self.stage_id, self.partition_id) - } -} - -#[derive(Debug)] -pub struct ExecutePartitionResult { - /// Path containing results for this partition - path: String, - stats: PartitionStats, -} - -impl ExecutePartitionResult { - pub fn new(path: &str, stats: PartitionStats) -> Self { - Self { - path: path.to_owned(), - stats, - } - } - - pub fn path(&self) -> &str { - &self.path - } - - pub fn statistics(&self) -> &PartitionStats { - &self.stats - } -} diff --git a/ballista/rust/core/src/serde/scheduler/to_proto.rs b/ballista/rust/core/src/serde/scheduler/to_proto.rs deleted file mode 100644 index 4c1c5d161892..000000000000 --- a/ballista/rust/core/src/serde/scheduler/to_proto.rs +++ /dev/null @@ -1,107 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::convert::TryInto; - -use crate::error::BallistaError; -use crate::serde::protobuf; -use crate::serde::protobuf::action::ActionType; -use crate::serde::scheduler::{Action, PartitionId, PartitionLocation, PartitionStats}; -use datafusion::physical_plan::Partitioning; - -impl TryInto for Action { - type Error = BallistaError; - - fn try_into(self) -> Result { - match self { - Action::FetchPartition { - job_id, - stage_id, - partition_id, - path, - } => Ok(protobuf::Action { - action_type: Some(ActionType::FetchPartition(protobuf::FetchPartition { - job_id, - stage_id: stage_id as u32, - partition_id: partition_id as u32, - path, - })), - settings: vec![], - }), - } - } -} - -#[allow(clippy::from_over_into)] -impl Into for PartitionId { - fn into(self) -> protobuf::PartitionId { - protobuf::PartitionId { - job_id: self.job_id, - stage_id: self.stage_id as u32, - partition_id: self.partition_id as u32, - } - } -} - -impl TryInto for PartitionLocation { - type Error = BallistaError; - - fn try_into(self) -> Result { - Ok(protobuf::PartitionLocation { - partition_id: Some(self.partition_id.into()), - executor_meta: Some(self.executor_meta.into()), - partition_stats: Some(self.partition_stats.into()), - path: self.path, - }) - } -} - -#[allow(clippy::from_over_into)] -impl Into for PartitionStats { - fn into(self) -> protobuf::PartitionStats { - let none_value = -1_i64; - protobuf::PartitionStats { - num_rows: self.num_rows.map(|n| n as i64).unwrap_or(none_value), - num_batches: self.num_batches.map(|n| n as i64).unwrap_or(none_value), - num_bytes: self.num_bytes.map(|n| n as i64).unwrap_or(none_value), - column_stats: vec![], - } - } -} - -pub fn hash_partitioning_to_proto( - output_partitioning: Option<&Partitioning>, -) -> Result, BallistaError> { - match output_partitioning { - Some(Partitioning::Hash(exprs, partition_count)) => { - Ok(Some(protobuf::PhysicalHashRepartition { - hash_expr: exprs - .iter() - .map(|expr| expr.clone().try_into()) - .collect::, BallistaError>>()?, - partition_count: *partition_count as u64, - })) - } - None => Ok(None), - other => { - return Err(BallistaError::General(format!( - "scheduler::to_proto() invalid partitioning for ExecutePartition: {:?}", - other - ))) - } - } -} diff --git a/ballista/rust/core/src/utils.rs b/ballista/rust/core/src/utils.rs deleted file mode 100644 index 1418aecb31a2..000000000000 --- a/ballista/rust/core/src/utils.rs +++ /dev/null @@ -1,315 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::io::{BufWriter, Write}; -use std::marker::PhantomData; - -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::Arc; -use std::{fs::File, pin::Pin}; - -use crate::error::{BallistaError, Result}; -use crate::execution_plans::{ - DistributedQueryExec, ShuffleWriterExec, UnresolvedShuffleExec, -}; -use crate::serde::scheduler::PartitionStats; - -use crate::config::BallistaConfig; -use crate::serde::{AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec}; -use async_trait::async_trait; -use datafusion::arrow::datatypes::Schema; -use datafusion::arrow::{ipc::writer::FileWriter, record_batch::RecordBatch}; -use datafusion::error::DataFusionError; -use datafusion::execution::context::{ - QueryPlanner, SessionConfig, SessionContext, SessionState, -}; -use datafusion::logical_plan::LogicalPlan; - -use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; -use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion::physical_plan::common::batch_byte_size; -use datafusion::physical_plan::empty::EmptyExec; - -use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv}; -use datafusion::physical_plan::aggregates::AggregateExec; -use datafusion::physical_plan::file_format::{CsvExec, ParquetExec}; -use datafusion::physical_plan::filter::FilterExec; -use datafusion::physical_plan::hash_join::HashJoinExec; -use datafusion::physical_plan::projection::ProjectionExec; -use datafusion::physical_plan::sorts::sort::SortExec; -use datafusion::physical_plan::{metrics, ExecutionPlan, RecordBatchStream}; -use futures::StreamExt; - -/// Stream data to disk in Arrow IPC format - -pub async fn write_stream_to_disk( - stream: &mut Pin>, - path: &str, - disk_write_metric: &metrics::Time, -) -> Result { - let file = File::create(&path).map_err(|e| { - BallistaError::General(format!( - "Failed to create partition file at {}: {:?}", - path, e - )) - })?; - - let mut num_rows = 0; - let mut num_batches = 0; - let mut num_bytes = 0; - let mut writer = FileWriter::try_new(file, stream.schema().as_ref())?; - - while let Some(result) = stream.next().await { - let batch = result?; - - let batch_size_bytes: usize = batch_byte_size(&batch); - num_batches += 1; - num_rows += batch.num_rows(); - num_bytes += batch_size_bytes; - - let timer = disk_write_metric.timer(); - writer.write(&batch)?; - timer.done(); - } - let timer = disk_write_metric.timer(); - writer.finish()?; - timer.done(); - Ok(PartitionStats::new( - Some(num_rows as u64), - Some(num_batches), - Some(num_bytes as u64), - )) -} - -pub async fn collect_stream( - stream: &mut Pin>, -) -> Result> { - let mut batches = vec![]; - while let Some(batch) = stream.next().await { - batches.push(batch?); - } - Ok(batches) -} - -pub fn produce_diagram(filename: &str, stages: &[Arc]) -> Result<()> { - let write_file = File::create(filename)?; - let mut w = BufWriter::new(&write_file); - writeln!(w, "digraph G {{")?; - - // draw stages and entities - for stage in stages { - writeln!(w, "\tsubgraph cluster{} {{", stage.stage_id())?; - writeln!(w, "\t\tlabel = \"Stage {}\";", stage.stage_id())?; - let mut id = AtomicUsize::new(0); - build_exec_plan_diagram( - &mut w, - stage.children()[0].as_ref(), - stage.stage_id(), - &mut id, - true, - )?; - writeln!(w, "\t}}")?; - } - - // draw relationships - for stage in stages { - let mut id = AtomicUsize::new(0); - build_exec_plan_diagram( - &mut w, - stage.children()[0].as_ref(), - stage.stage_id(), - &mut id, - false, - )?; - } - - write!(w, "}}")?; - Ok(()) -} - -fn build_exec_plan_diagram( - w: &mut BufWriter<&File>, - plan: &dyn ExecutionPlan, - stage_id: usize, - id: &mut AtomicUsize, - draw_entity: bool, -) -> Result { - let operator_str = if plan.as_any().downcast_ref::().is_some() { - "AggregateExec" - } else if plan.as_any().downcast_ref::().is_some() { - "SortExec" - } else if plan.as_any().downcast_ref::().is_some() { - "ProjectionExec" - } else if plan.as_any().downcast_ref::().is_some() { - "HashJoinExec" - } else if plan.as_any().downcast_ref::().is_some() { - "ParquetExec" - } else if plan.as_any().downcast_ref::().is_some() { - "CsvExec" - } else if plan.as_any().downcast_ref::().is_some() { - "FilterExec" - } else if plan.as_any().downcast_ref::().is_some() { - "ShuffleWriterExec" - } else if plan - .as_any() - .downcast_ref::() - .is_some() - { - "UnresolvedShuffleExec" - } else if plan - .as_any() - .downcast_ref::() - .is_some() - { - "CoalesceBatchesExec" - } else if plan - .as_any() - .downcast_ref::() - .is_some() - { - "CoalescePartitionsExec" - } else { - println!("Unknown: {:?}", plan); - "Unknown" - }; - - let node_id = id.load(Ordering::SeqCst); - id.store(node_id + 1, Ordering::SeqCst); - - if draw_entity { - writeln!( - w, - "\t\tstage_{}_exec_{} [shape=box, label=\"{}\"];", - stage_id, node_id, operator_str - )?; - } - for child in plan.children() { - if let Some(shuffle) = child.as_any().downcast_ref::() { - if !draw_entity { - writeln!( - w, - "\tstage_{}_exec_1 -> stage_{}_exec_{};", - shuffle.stage_id, stage_id, node_id - )?; - } - } else { - // relationships within same entity - let child_id = - build_exec_plan_diagram(w, child.as_ref(), stage_id, id, draw_entity)?; - if draw_entity { - writeln!( - w, - "\t\tstage_{}_exec_{} -> stage_{}_exec_{};", - stage_id, child_id, stage_id, node_id - )?; - } - } - } - Ok(node_id) -} - -/// Create a client DataFusion context that uses the BallistaQueryPlanner to send logical plans -/// to a Ballista scheduler -pub fn create_df_ctx_with_ballista_query_planner( - scheduler_url: String, - session_id: String, - config: &BallistaConfig, -) -> SessionContext { - let planner: Arc> = - Arc::new(BallistaQueryPlanner::new(scheduler_url, config.clone())); - - let session_config = SessionConfig::new() - .with_target_partitions(config.default_shuffle_partitions()) - .with_information_schema(true); - let mut session_state = SessionState::with_config_rt( - session_config, - Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap()), - ) - .with_query_planner(planner); - session_state.session_id = session_id; - // the SessionContext created here is the client side context, but the session_id is from server side. - SessionContext::with_state(session_state) -} - -pub struct BallistaQueryPlanner { - scheduler_url: String, - config: BallistaConfig, - extension_codec: Arc, - plan_repr: PhantomData, -} - -impl BallistaQueryPlanner { - pub fn new(scheduler_url: String, config: BallistaConfig) -> Self { - Self { - scheduler_url, - config, - extension_codec: Arc::new(DefaultLogicalExtensionCodec {}), - plan_repr: PhantomData, - } - } - - pub fn with_extension( - scheduler_url: String, - config: BallistaConfig, - extension_codec: Arc, - ) -> Self { - Self { - scheduler_url, - config, - extension_codec, - plan_repr: PhantomData, - } - } - - pub fn with_repr( - scheduler_url: String, - config: BallistaConfig, - extension_codec: Arc, - plan_repr: PhantomData, - ) -> Self { - Self { - scheduler_url, - config, - extension_codec, - plan_repr, - } - } -} - -#[async_trait] -impl QueryPlanner for BallistaQueryPlanner { - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &SessionState, - ) -> std::result::Result, DataFusionError> { - match logical_plan { - LogicalPlan::CreateExternalTable(_) => { - // table state is managed locally in the BallistaContext, not in the scheduler - Ok(Arc::new(EmptyExec::new(false, Arc::new(Schema::empty())))) - } - _ => Ok(Arc::new(DistributedQueryExec::with_repr( - self.scheduler_url.clone(), - self.config.clone(), - logical_plan.clone(), - self.extension_codec.clone(), - self.plan_repr, - session_state.session_id.clone(), - ))), - } - } -} diff --git a/ballista/rust/executor/Cargo.toml b/ballista/rust/executor/Cargo.toml deleted file mode 100644 index 959a55078dab..000000000000 --- a/ballista/rust/executor/Cargo.toml +++ /dev/null @@ -1,63 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -[package] -name = "ballista-executor" -description = "Ballista Distributed Compute - Executor" -license = "Apache-2.0" -version = "0.7.0" -homepage = "https://github.com/apache/arrow-datafusion" -repository = "https://github.com/apache/arrow-datafusion" -readme = "README.md" -authors = ["Apache Arrow "] -edition = "2018" - -[package.metadata.configure_me.bin] -executor = "executor_config_spec.toml" - -[features] -snmalloc = ["snmalloc-rs"] - -[dependencies] -anyhow = "1" -arrow = { version = "14.0.0" } -arrow-flight = { version = "14.0.0" } -async-trait = "0.1.41" -ballista-core = { path = "../core", version = "0.7.0" } -chrono = { version = "0.4", default-features = false } -configure_me = "0.4.0" -datafusion = { path = "../../../datafusion/core", version = "8.0.0" } -env_logger = "0.9" -futures = "0.3" -hyper = "0.14.4" -log = "0.4" -parking_lot = "0.12" -snmalloc-rs = { version = "0.2", optional = true } -tempfile = "3" -tokio = { version = "1.0", features = ["macros", "rt", "rt-multi-thread", "parking_lot"] } -tokio-stream = { version = "0.1", features = ["net"] } -tonic = "0.7" -uuid = { version = "1.0", features = ["v4"] } - -[dev-dependencies] - -[build-dependencies] -configure_me_codegen = "0.4.0" - -# use libc on unix like platforms to set worker priority in DedicatedExecutor -[target."cfg(unix)".dependencies.libc] -version = "0.2" diff --git a/ballista/rust/executor/README.md b/ballista/rust/executor/README.md deleted file mode 100644 index 91f4c3266c48..000000000000 --- a/ballista/rust/executor/README.md +++ /dev/null @@ -1,22 +0,0 @@ - - -# Ballista Executor Process - -This crate contains the Ballista executor process. diff --git a/ballista/rust/executor/build.rs b/ballista/rust/executor/build.rs deleted file mode 100644 index 1c9e32b0b894..000000000000 --- a/ballista/rust/executor/build.rs +++ /dev/null @@ -1,24 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -extern crate configure_me_codegen; - -fn main() -> Result<(), String> { - println!("cargo:rerun-if-changed=executor_config_spec.toml"); - configure_me_codegen::build_script_auto() - .map_err(|e| format!("configure_me code generation failed: {}", e)) -} diff --git a/ballista/rust/executor/executor_config_spec.toml b/ballista/rust/executor/executor_config_spec.toml deleted file mode 100644 index 86e712bda284..000000000000 --- a/ballista/rust/executor/executor_config_spec.toml +++ /dev/null @@ -1,104 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -[general] -name = "Ballista Executor" -env_prefix = "BALLISTA_EXECUTOR" -conf_file_param = "config_file" - -[[switch]] -name = "version" -doc = "Print version of this executable" - -[[param]] -name = "scheduler_host" -type = "String" -default = "std::string::String::from(\"localhost\")" -doc = "Scheduler host" - -[[param]] -name = "scheduler_port" -type = "u16" -default = "50050" -doc = "scheduler port" - -[[param]] -name = "bind_host" -type = "String" -default = "std::string::String::from(\"0.0.0.0\")" -doc = "Local IP address to bind to." - -[[param]] -name = "external_host" -type = "String" -doc = "Host name or IP address to register with scheduler so that other executors can connect to this executor. If none is provided, the scheduler will use the connecting IP address to communicate with the executor." - -[[param]] -abbr = "p" -name = "bind_port" -type = "u16" -default = "50051" -doc = "bind port" - -[[param]] -name = "bind_grpc_port" -type = "u16" -default = "50052" -doc = "bind grpc service port" - -[[param]] -name = "work_dir" -type = "String" -doc = "Directory for temporary IPC files" - -[[param]] -abbr = "c" -name = "concurrent_tasks" -type = "usize" -default = "4" -doc = "Max concurrent tasks." - -[[param]] -abbr = "s" -name = "task_scheduling_policy" -type = "ballista_core::config::TaskSchedulingPolicy" -doc = "The task scheduing policy for the scheduler, see TaskSchedulingPolicy::variants() for options. Default: PullStaged" -default = "ballista_core::config::TaskSchedulingPolicy::PullStaged" - -[[param]] -name = "executor_cleanup_enable" -type = "bool" -doc = "Enable periodic cleanup of work_dir directories." -default = "false" - -[[param]] -name = "executor_cleanup_interval" -type = "u64" -doc = "Controls the interval in seconds , which the worker cleans up old job dirs on the local machine." -default = "1800" - -[[param]] -name = "executor_cleanup_ttl" -type = "u64" -doc = "The number of seconds to retain job directories on each worker 604800 (7 days, 7 * 24 * 3600), In other words, after job done, how long the resulting data is retained" -default = "604800" - -[[param]] -name = "plugin_dir" -type = "String" -doc = "plugin dir" -default = "std::string::String::from(\"\")" \ No newline at end of file diff --git a/ballista/rust/executor/src/collect.rs b/ballista/rust/executor/src/collect.rs deleted file mode 100644 index 54e97550a68f..000000000000 --- a/ballista/rust/executor/src/collect.rs +++ /dev/null @@ -1,135 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! The CollectExec operator retrieves results from the cluster and returns them as a single -//! vector of [RecordBatch]. - -use std::sync::Arc; -use std::task::{Context, Poll}; -use std::{any::Any, pin::Pin}; - -use datafusion::arrow::{ - datatypes::SchemaRef, error::Result as ArrowResult, record_batch::RecordBatch, -}; -use datafusion::error::DataFusionError; -use datafusion::execution::context::TaskContext; -use datafusion::physical_plan::expressions::PhysicalSortExpr; -use datafusion::physical_plan::{ - DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream, Statistics, -}; -use datafusion::{error::Result, physical_plan::RecordBatchStream}; -use futures::stream::SelectAll; -use futures::Stream; - -/// The CollectExec operator retrieves results from the cluster and returns them as a single -/// vector of [RecordBatch]. -#[derive(Debug, Clone)] -pub struct CollectExec { - plan: Arc, -} - -impl CollectExec { - pub fn new(plan: Arc) -> Self { - Self { plan } - } -} - -impl ExecutionPlan for CollectExec { - fn as_any(&self) -> &dyn Any { - self - } - - fn schema(&self) -> SchemaRef { - self.plan.schema() - } - - fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(1) - } - - fn output_ordering(&self) -> Option<&[PhysicalSortExpr]> { - None - } - - fn children(&self) -> Vec> { - vec![self.plan.clone()] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - unimplemented!() - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - assert_eq!(0, partition); - let num_partitions = self.plan.output_partitioning().partition_count(); - - let streams = (0..num_partitions) - .map(|i| self.plan.execute(i, context.clone())) - .collect::>>() - .map_err(|e| DataFusionError::Execution(format!("BallistaError: {:?}", e)))?; - - Ok(Box::pin(MergedRecordBatchStream { - schema: self.schema(), - select_all: Box::pin(futures::stream::select_all(streams)), - })) - } - - fn fmt_as( - &self, - t: DisplayFormatType, - f: &mut std::fmt::Formatter, - ) -> std::fmt::Result { - match t { - DisplayFormatType::Default => { - write!(f, "CollectExec") - } - } - } - - fn statistics(&self) -> Statistics { - self.plan.statistics() - } -} - -struct MergedRecordBatchStream { - schema: SchemaRef, - select_all: Pin>>, -} - -impl Stream for MergedRecordBatchStream { - type Item = ArrowResult; - - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - self.select_all.as_mut().poll_next(cx) - } -} - -impl RecordBatchStream for MergedRecordBatchStream { - fn schema(&self) -> SchemaRef { - self.schema.clone() - } -} diff --git a/ballista/rust/executor/src/cpu_bound_executor.rs b/ballista/rust/executor/src/cpu_bound_executor.rs deleted file mode 100644 index d6c53bc42280..000000000000 --- a/ballista/rust/executor/src/cpu_bound_executor.rs +++ /dev/null @@ -1,378 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//Inspire by https://thenewstack.io/using-rustlangs-async-tokio-runtime-for-cpu-bound-tasks/ - -//! This module contains a dedicated thread pool for running "cpu -//! intensive" workloads as query plans - -use log::warn; -use parking_lot::Mutex; -use std::{pin::Pin, sync::Arc}; -use tokio::sync::oneshot::Receiver; - -use futures::Future; - -/// The type of thing that the dedicated executor runs -type Task = Pin + Send>>; - -/// Runs futures (and any `tasks` that are `tokio::task::spawned` by -/// them) on a separate tokio runtime, like separate CPU-bound (execute a datafusion plan) tasks -/// from IO-bound tasks(heartbeats). Get more from the above blog. -#[derive(Clone)] -pub struct DedicatedExecutor { - state: Arc>, -} - -/// Runs futures (and any `tasks` that are `tokio::task::spawned` by -/// them) on a separate tokio Executor -struct State { - /// The number of threads in this pool - num_threads: usize, - - /// The name of the threads for this executor - thread_name: String, - - /// Channel for requests -- the dedicated executor takes requests - /// from here and runs them. - requests: Option>, - - /// The thread that is doing the work - thread: Option>, -} - -/// The default worker priority (value passed to `libc::setpriority`); -const WORKER_PRIORITY: i32 = 10; - -impl std::fmt::Debug for DedicatedExecutor { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - let state = self.state.lock(); - - let mut d = f.debug_struct("DedicatedExecutor"); - - d.field("num_threads", &state.num_threads) - .field("thread_name", &state.thread_name); - - if state.requests.is_some() { - d.field("requests", &"Some(...)") - } else { - d.field("requests", &"None") - }; - - if state.thread.is_some() { - d.field("thread", &"Some(...)") - } else { - d.field("thread", &"None") - }; - - d.finish() - } -} - -impl DedicatedExecutor { - /// https://stackoverflow.com/questions/62536566 - /// Creates a new `DedicatedExecutor` with a dedicated tokio - /// runtime that is separate from the `[tokio::main]` threadpool. - /// - /// The worker thread priority is set to low so that such tasks do - /// not starve other more important tasks (such as answering health checks) - /// - pub fn new(thread_name: impl Into, num_threads: usize) -> Self { - let thread_name = thread_name.into(); - let name_copy = thread_name.to_string(); - - let (tx, rx) = std::sync::mpsc::channel(); - - // Cannot create a separated tokio runtime in another tokio runtime, - // So use std::thread to spawn a thread - let thread = std::thread::spawn(move || { - let runtime = tokio::runtime::Builder::new_multi_thread() - .enable_all() - .thread_name(&name_copy) - .worker_threads(num_threads) - .on_thread_start(move || set_current_thread_priority(WORKER_PRIORITY)) - .build() - .expect("Creating tokio runtime"); - - // By entering the context, all calls to `tokio::spawn` go - // to this executor - let _guard = runtime.enter(); - - while let Ok(request) = rx.recv() { - // TODO feedback request status - tokio::task::spawn(request); - } - }); - - let state = State { - num_threads, - thread_name, - requests: Some(tx), - thread: Some(thread), - }; - - Self { - state: Arc::new(Mutex::new(state)), - } - } - - /// Runs the specified Future (and any tasks it spawns) on the - /// `DedicatedExecutor`. - /// - /// Currently all tasks are added to the tokio executor - /// immediately and compete for the threadpool's resources. - pub fn spawn(&self, task: T) -> Receiver - where - T: Future + Send + 'static, - T::Output: Send + 'static, - { - let (tx, rx) = tokio::sync::oneshot::channel(); - - // create a execution plan to spawn - let job = Box::pin(async move { - let task_output = task.await; - if tx.send(task_output).is_err() { - warn!("Spawned task output ignored: receiver dropped"); - } - }); - - let mut state = self.state.lock(); - - if let Some(requests) = &mut state.requests { - // would fail if someone has started shutdown - requests.send(job).ok(); - } else { - warn!("tried to schedule task on an executor that was shutdown"); - } - - rx - } - - /// signals shutdown of this executor and any Clones - #[allow(dead_code)] - pub fn shutdown(&self) { - // hang up the channel which will cause the dedicated thread - // to quit - let mut state = self.state.lock(); - // remaining job will still running - state.requests = None; - } - - /// Stops all subsequent task executions, and waits for the worker - /// thread to complete. Note this will shutdown all clones of this - /// `DedicatedExecutor` as well. - /// - /// Only the first one to `join` will actually wait for the - /// executing thread to complete. All other calls to join will - /// complete immediately. - #[allow(dead_code)] - pub fn join(&self) { - self.shutdown(); - - // take the thread out when mutex is held - let thread = { - let mut state = self.state.lock(); - state.thread.take() - }; - - // wait for completion while not holding the mutex to avoid - // deadlocks - if let Some(thread) = thread { - thread.join().ok(); - } - } -} - -#[cfg(unix)] -fn set_current_thread_priority(prio: i32) { - unsafe { libc::setpriority(0, 0, prio) }; -} - -#[cfg(not(unix))] -fn set_current_thread_priority(_prio: i32) { - warn!("Setting worker thread priority not supported on this platform"); -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::{Arc, Barrier}; - - #[cfg(unix)] - fn get_current_thread_priority() -> i32 { - // on linux setpriority sets the current thread's priority - // (as opposed to the current process). - unsafe { libc::getpriority(0, 0) } - } - - #[cfg(not(unix))] - fn get_current_thread_priority() -> i32 { - WORKER_PRIORITY - } - - #[tokio::test] - async fn basic_test_in_diff_thread() { - let barrier = Arc::new(Barrier::new(2)); - - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 1); - let dedicated_task = exec.spawn(do_work(42, Arc::clone(&barrier))); - - // Note the dedicated task will never complete if it runs on - // the main tokio thread - //#[tokio::test] will only create one thread, if we running use tokio spwan - // after call do_work with barrier.wait() the only thread will be blocked and never finished - barrier.wait(); - - // should be able to get the result - assert_eq!(dedicated_task.await.unwrap(), 42); - } - - #[tokio::test] - async fn basic_clone() { - let barrier = Arc::new(Barrier::new(2)); - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 1); - let dedicated_task = exec.clone().spawn(do_work(42, Arc::clone(&barrier))); - barrier.wait(); - assert_eq!(dedicated_task.await.unwrap(), 42); - } - - #[tokio::test] - async fn multi_task() { - let barrier = Arc::new(Barrier::new(3)); - - // make an executor with two threads - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 2); - let dedicated_task1 = exec.spawn(do_work(11, Arc::clone(&barrier))); - let dedicated_task2 = exec.spawn(do_work(42, Arc::clone(&barrier))); - - // block main thread until completion of other two tasks - barrier.wait(); - - // should be able to get the result - assert_eq!(dedicated_task1.await.unwrap(), 11); - assert_eq!(dedicated_task2.await.unwrap(), 42); - - exec.join(); - } - - #[tokio::test] - async fn worker_priority() { - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 2); - - let dedicated_task = exec.spawn(async move { get_current_thread_priority() }); - - assert_eq!(dedicated_task.await.unwrap(), WORKER_PRIORITY); - } - - #[tokio::test] - async fn tokio_spawn() { - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 2); - - // spawn a task that spawns to other tasks and ensure they run on the dedicated - // executor - let dedicated_task = exec.spawn(async move { - // spawn separate tasks - let t1 = tokio::task::spawn(async { - assert_eq!( - std::thread::current().name(), - Some("Test DedicatedExecutor") - ); - 25usize - }); - t1.await.unwrap() - }); - - assert_eq!(dedicated_task.await.unwrap(), 25); - } - - #[tokio::test] - async fn panic_on_executor() { - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 1); - let dedicated_task = exec.spawn(async move { - panic!("At the disco, on the dedicated task scheduler"); - }); - - // should not be able to get the result - dedicated_task.await.unwrap_err(); - } - - #[tokio::test] - #[ignore] - // related https://github.com/apache/arrow-datafusion/issues/2140 - async fn executor_shutdown_while_task_running() { - let barrier = Arc::new(Barrier::new(2)); - - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 1); - let dedicated_task = exec.spawn(do_work(42, Arc::clone(&barrier))); - - exec.shutdown(); - // block main thread until completion of the outstanding task - barrier.wait(); - - // task should complete successfully - assert_eq!(dedicated_task.await.unwrap(), 42); - } - - #[tokio::test] - async fn executor_submit_task_after_shutdown() { - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 1); - - // Simulate trying to submit tasks once executor has shutdown - exec.shutdown(); - let dedicated_task = exec.spawn(async { 11 }); - - // task should complete, but return an error - dedicated_task.await.unwrap_err(); - } - - #[tokio::test] - async fn executor_submit_task_after_clone_shutdown() { - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 1); - - // shutdown the clone (but not the exec) - exec.clone().join(); - - // Simulate trying to submit tasks once executor has shutdown - let dedicated_task = exec.spawn(async { 11 }); - - // task should complete, but return an error - dedicated_task.await.unwrap_err(); - } - - #[tokio::test] - async fn executor_join() { - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 1); - // test it doesn't hang - exec.join() - } - - #[tokio::test] - #[allow(clippy::redundant_clone)] - async fn executor_clone_join() { - let exec = DedicatedExecutor::new("Test DedicatedExecutor", 1); - // test not hang - exec.clone().join(); - exec.clone().join(); - exec.join(); - } - - /// Wait for the barrier and then return `result` - async fn do_work(result: usize, barrier: Arc) -> usize { - barrier.wait(); - result - } -} diff --git a/ballista/rust/executor/src/execution_loop.rs b/ballista/rust/executor/src/execution_loop.rs deleted file mode 100644 index 06128f8dbc53..000000000000 --- a/ballista/rust/executor/src/execution_loop.rs +++ /dev/null @@ -1,213 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::HashMap; -use std::ops::Deref; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::mpsc::{Receiver, Sender, TryRecvError}; -use std::{sync::Arc, time::Duration}; - -use datafusion::physical_plan::ExecutionPlan; -use log::{debug, error, info, warn}; -use tonic::transport::Channel; - -use ballista_core::serde::protobuf::{ - scheduler_grpc_client::SchedulerGrpcClient, PollWorkParams, PollWorkResult, - TaskDefinition, TaskStatus, -}; - -use crate::as_task_status; -use crate::executor::Executor; -use ballista_core::error::BallistaError; -use ballista_core::serde::physical_plan::from_proto::parse_protobuf_hash_partitioning; -use ballista_core::serde::scheduler::ExecutorSpecification; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec}; -use datafusion::execution::context::TaskContext; - -pub async fn poll_loop( - mut scheduler: SchedulerGrpcClient, - executor: Arc, - codec: BallistaCodec, -) { - let executor_specification: ExecutorSpecification = executor - .metadata - .specification - .as_ref() - .unwrap() - .clone() - .into(); - let available_tasks_slots = - Arc::new(AtomicUsize::new(executor_specification.task_slots as usize)); - let (task_status_sender, mut task_status_receiver) = - std::sync::mpsc::channel::(); - - loop { - debug!("Starting registration loop with scheduler"); - - let task_status: Vec = - sample_tasks_status(&mut task_status_receiver).await; - - // Keeps track of whether we received task in last iteration - // to avoid going in sleep mode between polling - let mut active_job = false; - - let poll_work_result: anyhow::Result< - tonic::Response, - tonic::Status, - > = scheduler - .poll_work(PollWorkParams { - metadata: Some(executor.metadata.clone()), - can_accept_task: available_tasks_slots.load(Ordering::SeqCst) > 0, - task_status, - }) - .await; - - let task_status_sender = task_status_sender.clone(); - - match poll_work_result { - Ok(result) => { - if let Some(task) = result.into_inner().task { - match run_received_tasks( - executor.clone(), - available_tasks_slots.clone(), - task_status_sender, - task, - &codec, - ) - .await - { - Ok(_) => { - active_job = true; - } - Err(e) => { - warn!("Failed to run task: {:?}", e); - active_job = false; - } - } - } else { - active_job = false; - } - } - Err(error) => { - warn!("Executor registration failed. If this continues to happen the executor might be marked as dead by the scheduler. Error: {}", error); - } - } - if !active_job { - tokio::time::sleep(Duration::from_millis(100)).await; - } - } -} - -async fn run_received_tasks( - executor: Arc, - available_tasks_slots: Arc, - task_status_sender: Sender, - task: TaskDefinition, - codec: &BallistaCodec, -) -> Result<(), BallistaError> { - let task_id = task.task_id.unwrap(); - let task_id_log = format!( - "{}/{}/{}", - task_id.job_id, task_id.stage_id, task_id.partition_id - ); - info!("Received task {}", task_id_log); - available_tasks_slots.fetch_sub(1, Ordering::SeqCst); - - let runtime = executor.runtime.clone(); - let session_id = task.session_id; - let mut task_props = HashMap::new(); - for kv_pair in task.props { - task_props.insert(kv_pair.key, kv_pair.value); - } - - let mut task_scalar_functions = HashMap::new(); - let mut task_aggregate_functions = HashMap::new(); - // TODO combine the functions from Executor's functions and TaskDefintion's function resources - for scalar_func in executor.scalar_functions.clone() { - task_scalar_functions.insert(scalar_func.0.clone(), scalar_func.1); - } - for agg_func in executor.aggregate_functions.clone() { - task_aggregate_functions.insert(agg_func.0, agg_func.1); - } - let task_context = Arc::new(TaskContext::new( - task_id_log.clone(), - session_id, - task_props, - task_scalar_functions, - task_aggregate_functions, - runtime.clone(), - )); - - let plan: Arc = - U::try_decode(task.plan.as_slice()).and_then(|proto| { - proto.try_into_physical_plan( - task_context.deref(), - runtime.deref(), - codec.physical_extension_codec(), - ) - })?; - - let shuffle_output_partitioning = parse_protobuf_hash_partitioning( - task.output_partitioning.as_ref(), - task_context.as_ref(), - )?; - - tokio::spawn(async move { - let execution_result = executor - .execute_shuffle_write( - task_id.job_id.clone(), - task_id.stage_id as usize, - task_id.partition_id as usize, - plan, - task_context, - shuffle_output_partitioning, - ) - .await; - info!("Done with task {}", task_id_log); - debug!("Statistics: {:?}", execution_result); - available_tasks_slots.fetch_add(1, Ordering::SeqCst); - let _ = task_status_sender.send(as_task_status( - execution_result, - executor.metadata.id.clone(), - task_id, - )); - }); - - Ok(()) -} - -async fn sample_tasks_status( - task_status_receiver: &mut Receiver, -) -> Vec { - let mut task_status: Vec = vec![]; - - loop { - match task_status_receiver.try_recv() { - anyhow::Result::Ok(status) => { - task_status.push(status); - } - Err(TryRecvError::Empty) => { - break; - } - Err(TryRecvError::Disconnected) => { - error!("Task statuses channel disconnected"); - } - } - } - - task_status -} diff --git a/ballista/rust/executor/src/executor.rs b/ballista/rust/executor/src/executor.rs deleted file mode 100644 index fa092137abd6..000000000000 --- a/ballista/rust/executor/src/executor.rs +++ /dev/null @@ -1,119 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Ballista executor logic - -use std::collections::HashMap; -use std::sync::Arc; - -use crate::metrics::ExecutorMetricsCollector; -use ballista_core::error::BallistaError; -use ballista_core::execution_plans::ShuffleWriterExec; -use ballista_core::serde::protobuf; -use ballista_core::serde::protobuf::ExecutorRegistration; -use datafusion::error::DataFusionError; -use datafusion::execution::context::TaskContext; -use datafusion::execution::runtime_env::RuntimeEnv; - -use datafusion::physical_plan::udaf::AggregateUDF; -use datafusion::physical_plan::udf::ScalarUDF; -use datafusion::physical_plan::{ExecutionPlan, Partitioning}; - -/// Ballista executor -pub struct Executor { - /// Metadata - pub metadata: ExecutorRegistration, - - /// Directory for storing partial results - pub work_dir: String, - - /// Scalar functions that are registered in the Executor - pub scalar_functions: HashMap>, - - /// Aggregate functions registered in the Executor - pub aggregate_functions: HashMap>, - - /// Runtime environment for Executor - pub runtime: Arc, - - /// Collector for runtime execution metrics - pub metrics_collector: Arc, -} - -impl Executor { - /// Create a new executor instance - pub fn new( - metadata: ExecutorRegistration, - work_dir: &str, - runtime: Arc, - metrics_collector: Arc, - ) -> Self { - Self { - metadata, - work_dir: work_dir.to_owned(), - // TODO add logic to dynamically load UDF/UDAFs libs from files - scalar_functions: HashMap::new(), - aggregate_functions: HashMap::new(), - runtime, - metrics_collector, - } - } -} - -impl Executor { - /// Execute one partition of a query stage and persist the result to disk in IPC format. On - /// success, return a RecordBatch containing metadata about the results, including path - /// and statistics. - pub async fn execute_shuffle_write( - &self, - job_id: String, - stage_id: usize, - part: usize, - plan: Arc, - task_ctx: Arc, - _shuffle_output_partitioning: Option, - ) -> Result, BallistaError> { - let exec = if let Some(shuffle_writer) = - plan.as_any().downcast_ref::() - { - // recreate the shuffle writer with the correct working directory - ShuffleWriterExec::try_new( - job_id.clone(), - stage_id, - plan.children()[0].clone(), - self.work_dir.clone(), - shuffle_writer.shuffle_output_partitioning().cloned(), - ) - } else { - Err(DataFusionError::Internal( - "Plan passed to execute_shuffle_write is not a ShuffleWriterExec" - .to_string(), - )) - }?; - - let partitions = exec.execute_shuffle_write(part, task_ctx).await?; - - self.metrics_collector - .record_stage(&job_id, stage_id, part, exec); - - Ok(partitions) - } - - pub fn work_dir(&self) -> &str { - &self.work_dir - } -} diff --git a/ballista/rust/executor/src/executor_server.rs b/ballista/rust/executor/src/executor_server.rs deleted file mode 100644 index 11a5c75527cc..000000000000 --- a/ballista/rust/executor/src/executor_server.rs +++ /dev/null @@ -1,353 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::HashMap; -use std::ops::Deref; -use std::sync::Arc; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use tokio::sync::mpsc; - -use log::{debug, error, info}; -use tonic::transport::{Channel, Server}; -use tonic::{Request, Response, Status}; - -use ballista_core::error::BallistaError; -use ballista_core::serde::physical_plan::from_proto::parse_protobuf_hash_partitioning; -use ballista_core::serde::protobuf::executor_grpc_server::{ - ExecutorGrpc, ExecutorGrpcServer, -}; -use ballista_core::serde::protobuf::executor_registration::OptionalHost; -use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient; -use ballista_core::serde::protobuf::{ - HeartBeatParams, LaunchTaskParams, LaunchTaskResult, RegisterExecutorParams, - StopExecutorParams, StopExecutorResult, TaskDefinition, UpdateTaskStatusParams, -}; -use ballista_core::serde::scheduler::ExecutorState; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec}; -use datafusion::execution::context::TaskContext; -use datafusion::physical_plan::ExecutionPlan; - -use crate::as_task_status; -use crate::cpu_bound_executor::DedicatedExecutor; -use crate::executor::Executor; - -pub async fn startup( - mut scheduler: SchedulerGrpcClient, - executor: Arc, - codec: BallistaCodec, -) { - // TODO make the buffer size configurable - let (tx_task, rx_task) = mpsc::channel::(1000); - - let executor_server = ExecutorServer::new( - scheduler.clone(), - executor.clone(), - ExecutorEnv { tx_task }, - codec, - ); - - // 1. Start executor grpc service - { - let executor_meta = executor.metadata.clone(); - let addr = format!( - "{}:{}", - executor_meta - .optional_host - .map(|h| match h { - OptionalHost::Host(host) => host, - }) - .unwrap_or_else(|| String::from("127.0.0.1")), - executor_meta.grpc_port - ); - let addr = addr.parse().unwrap(); - info!("Setup executor grpc service for {:?}", addr); - - let server = ExecutorGrpcServer::new(executor_server.clone()); - let grpc_server_future = Server::builder().add_service(server).serve(addr); - tokio::spawn(async move { grpc_server_future.await }); - } - - let executor_server = Arc::new(executor_server); - - // 2. Do executor registration - match register_executor(&mut scheduler, executor.clone()).await { - Ok(_) => { - info!("Executor registration succeed"); - } - Err(error) => { - panic!("Executor registration failed due to: {}", error); - } - }; - - // 3. Start Heartbeater - { - let heartbeater = Heartbeater::new(executor_server.clone()); - heartbeater.start().await; - } - - // 4. Start TaskRunnerPool - { - let task_runner_pool = TaskRunnerPool::new(executor_server.clone()); - task_runner_pool.start(rx_task).await; - } -} - -#[allow(clippy::clone_on_copy)] -async fn register_executor( - scheduler: &mut SchedulerGrpcClient, - executor: Arc, -) -> Result<(), BallistaError> { - let result = scheduler - .register_executor(RegisterExecutorParams { - metadata: Some(executor.metadata.clone()), - }) - .await?; - if result.into_inner().success { - Ok(()) - } else { - Err(BallistaError::General( - "Executor registration failed!!!".to_owned(), - )) - } -} - -#[derive(Clone)] -pub struct ExecutorServer { - _start_time: u128, - executor: Arc, - scheduler: SchedulerGrpcClient, - executor_env: ExecutorEnv, - codec: BallistaCodec, -} - -#[derive(Clone)] -struct ExecutorEnv { - tx_task: mpsc::Sender, -} - -unsafe impl Sync for ExecutorEnv {} - -impl ExecutorServer { - fn new( - scheduler: SchedulerGrpcClient, - executor: Arc, - executor_env: ExecutorEnv, - codec: BallistaCodec, - ) -> Self { - Self { - _start_time: SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap() - .as_millis(), - executor, - scheduler, - executor_env, - codec, - } - } - - async fn heartbeat(&self) { - // TODO Error handling - self.scheduler - .clone() - .heart_beat_from_executor(HeartBeatParams { - executor_id: self.executor.metadata.id.clone(), - state: Some(self.get_executor_state().into()), - }) - .await - .unwrap(); - } - - async fn run_task(&self, task: TaskDefinition) -> Result<(), BallistaError> { - let task_id = task.task_id.unwrap(); - let task_id_log = format!( - "{}/{}/{}", - task_id.job_id, task_id.stage_id, task_id.partition_id - ); - info!("Start to run task {}", task_id_log); - - let runtime = self.executor.runtime.clone(); - let session_id = task.session_id; - let mut task_props = HashMap::new(); - for kv_pair in task.props { - task_props.insert(kv_pair.key, kv_pair.value); - } - - let mut task_scalar_functions = HashMap::new(); - let mut task_aggregate_functions = HashMap::new(); - // TODO combine the functions from Executor's functions and TaskDefintion's function resources - for scalar_func in self.executor.scalar_functions.clone() { - task_scalar_functions.insert(scalar_func.0, scalar_func.1); - } - for agg_func in self.executor.aggregate_functions.clone() { - task_aggregate_functions.insert(agg_func.0, agg_func.1); - } - let task_context = Arc::new(TaskContext::new( - task_id_log.clone(), - session_id, - task_props, - task_scalar_functions, - task_aggregate_functions, - runtime.clone(), - )); - - let encoded_plan = &task.plan.as_slice(); - - let plan: Arc = - U::try_decode(encoded_plan).and_then(|proto| { - proto.try_into_physical_plan( - task_context.deref(), - runtime.deref(), - self.codec.physical_extension_codec(), - ) - })?; - - let shuffle_output_partitioning = parse_protobuf_hash_partitioning( - task.output_partitioning.as_ref(), - task_context.as_ref(), - )?; - - let execution_result = self - .executor - .execute_shuffle_write( - task_id.job_id.clone(), - task_id.stage_id as usize, - task_id.partition_id as usize, - plan, - task_context, - shuffle_output_partitioning, - ) - .await; - info!("Done with task {}", task_id_log); - debug!("Statistics: {:?}", execution_result); - - let executor_id = &self.executor.metadata.id; - // TODO use another channel to update the status of a task set - self.scheduler - .clone() - .update_task_status(UpdateTaskStatusParams { - executor_id: executor_id.clone(), - task_status: vec![as_task_status( - execution_result, - executor_id.clone(), - task_id, - )], - }) - .await?; - - Ok(()) - } - - // TODO with real state - fn get_executor_state(&self) -> ExecutorState { - ExecutorState { - available_memory_size: u64::MAX, - } - } -} - -struct Heartbeater { - executor_server: Arc>, -} - -impl Heartbeater { - fn new(executor_server: Arc>) -> Self { - Self { executor_server } - } - - async fn start(&self) { - let executor_server = self.executor_server.clone(); - tokio::spawn(async move { - info!("Starting heartbeater to send heartbeat the scheduler periodically"); - loop { - executor_server.heartbeat().await; - tokio::time::sleep(Duration::from_millis(60000)).await; - } - }); - } -} - -struct TaskRunnerPool { - executor_server: Arc>, -} - -impl TaskRunnerPool { - fn new(executor_server: Arc>) -> Self { - Self { executor_server } - } - - async fn start(&self, mut rx_task: mpsc::Receiver) { - let executor_server = self.executor_server.clone(); - tokio::spawn(async move { - info!("Starting the task runner pool"); - // Use a dedicated executor for CPU bound tasks so that the main tokio - // executor can still answer requests even when under load - // TODO make it configurable - let dedicated_executor = DedicatedExecutor::new("task_runner", 4); - loop { - if let Some(task) = rx_task.recv().await { - if let Some(task_id) = &task.task_id { - let task_id_log = format!( - "{}/{}/{}", - task_id.job_id, task_id.stage_id, task_id.partition_id - ); - info!("Received task {:?}", &task_id_log); - - let server = executor_server.clone(); - dedicated_executor.spawn(async move { - server.run_task(task).await.unwrap_or_else(|e| { - error!( - "Fail to run the task {:?} due to {:?}", - task_id_log, e - ); - }); - }); - } else { - error!("There's no task id in the task definition {:?}", task); - } - } else { - info!("Channel is closed and will exit the loop"); - return; - } - } - }); - } -} - -#[tonic::async_trait] -impl ExecutorGrpc - for ExecutorServer -{ - async fn launch_task( - &self, - request: Request, - ) -> Result, Status> { - let tasks = request.into_inner().task; - let task_sender = self.executor_env.tx_task.clone(); - for task in tasks { - task_sender.send(task).await.unwrap(); - } - Ok(Response::new(LaunchTaskResult { success: true })) - } - - async fn stop_executor( - &self, - _request: Request, - ) -> Result, Status> { - todo!() - } -} diff --git a/ballista/rust/executor/src/flight_service.rs b/ballista/rust/executor/src/flight_service.rs deleted file mode 100644 index 31fd8b002485..000000000000 --- a/ballista/rust/executor/src/flight_service.rs +++ /dev/null @@ -1,245 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Implementation of the Apache Arrow Flight protocol that wraps an executor. - -use std::fs::File; -use std::pin::Pin; -use std::sync::Arc; - -use crate::executor::Executor; -use arrow_flight::SchemaAsIpc; -use ballista_core::error::BallistaError; -use ballista_core::serde::decode_protobuf; -use ballista_core::serde::scheduler::Action as BallistaAction; - -use arrow_flight::{ - flight_service_server::FlightService, Action, ActionType, Criteria, Empty, - FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, - PutResult, SchemaResult, Ticket, -}; -use datafusion::arrow::{ - error::ArrowError, ipc::reader::FileReader, ipc::writer::IpcWriteOptions, - record_batch::RecordBatch, -}; -use futures::{Stream, StreamExt}; -use log::{info, warn}; -use std::io::{Read, Seek}; -use tokio::sync::mpsc::channel; -use tokio::{ - sync::mpsc::{Receiver, Sender}, - task, -}; -use tokio_stream::wrappers::ReceiverStream; -use tonic::{Request, Response, Status, Streaming}; - -type FlightDataSender = Sender>; -type FlightDataReceiver = Receiver>; - -/// Service implementing the Apache Arrow Flight Protocol -#[derive(Clone)] -pub struct BallistaFlightService { - /// Executor - _executor: Arc, -} - -impl BallistaFlightService { - pub fn new(_executor: Arc) -> Self { - Self { _executor } - } -} - -type BoxedFlightStream = - Pin> + Send + Sync + 'static>>; - -#[tonic::async_trait] -impl FlightService for BallistaFlightService { - type DoActionStream = BoxedFlightStream; - type DoExchangeStream = BoxedFlightStream; - type DoGetStream = BoxedFlightStream; - type DoPutStream = BoxedFlightStream; - type HandshakeStream = BoxedFlightStream; - type ListActionsStream = BoxedFlightStream; - type ListFlightsStream = BoxedFlightStream; - - async fn do_get( - &self, - request: Request, - ) -> Result, Status> { - let ticket = request.into_inner(); - - let action = - decode_protobuf(&ticket.ticket).map_err(|e| from_ballista_err(&e))?; - - match &action { - BallistaAction::FetchPartition { path, .. } => { - info!("FetchPartition reading {}", &path); - let file = File::open(&path) - .map_err(|e| { - BallistaError::General(format!( - "Failed to open partition file at {}: {:?}", - path, e - )) - }) - .map_err(|e| from_ballista_err(&e))?; - let reader = - FileReader::try_new(file, None).map_err(|e| from_arrow_err(&e))?; - - let (tx, rx): (FlightDataSender, FlightDataReceiver) = channel(2); - - // Arrow IPC reader does not implement Sync + Send so we need to use a channel - // to communicate - task::spawn(async move { - if let Err(e) = stream_flight_data(reader, tx).await { - warn!("Error streaming results: {:?}", e); - } - }); - - Ok(Response::new( - Box::pin(ReceiverStream::new(rx)) as Self::DoGetStream - )) - } - } - } - - async fn get_schema( - &self, - _request: Request, - ) -> Result, Status> { - Err(Status::unimplemented("get_schema")) - } - - async fn get_flight_info( - &self, - _request: Request, - ) -> Result, Status> { - Err(Status::unimplemented("get_flight_info")) - } - - async fn handshake( - &self, - _request: Request>, - ) -> Result, Status> { - Err(Status::unimplemented("handshake")) - } - - async fn list_flights( - &self, - _request: Request, - ) -> Result, Status> { - Err(Status::unimplemented("list_flights")) - } - - async fn do_put( - &self, - request: Request>, - ) -> Result, Status> { - let mut request = request.into_inner(); - - while let Some(data) = request.next().await { - let _data = data?; - } - - Err(Status::unimplemented("do_put")) - } - - async fn do_action( - &self, - request: Request, - ) -> Result, Status> { - let action = request.into_inner(); - - let _action = - decode_protobuf(&action.body.to_vec()).map_err(|e| from_ballista_err(&e))?; - - Err(Status::unimplemented("do_action")) - } - - async fn list_actions( - &self, - _request: Request, - ) -> Result, Status> { - Err(Status::unimplemented("list_actions")) - } - - async fn do_exchange( - &self, - _request: Request>, - ) -> Result, Status> { - Err(Status::unimplemented("do_exchange")) - } -} - -/// Convert a single RecordBatch into an iterator of FlightData (containing -/// dictionaries and batches) -fn create_flight_iter( - batch: &RecordBatch, - options: &IpcWriteOptions, -) -> Box>> { - let (flight_dictionaries, flight_batch) = - arrow_flight::utils::flight_data_from_arrow_batch(batch, options); - Box::new( - flight_dictionaries - .into_iter() - .chain(std::iter::once(flight_batch)) - .map(Ok), - ) -} - -async fn stream_flight_data( - reader: FileReader, - tx: FlightDataSender, -) -> Result<(), Status> -where - T: Read + Seek, -{ - let options = arrow::ipc::writer::IpcWriteOptions::default(); - let schema_flight_data = SchemaAsIpc::new(reader.schema().as_ref(), &options).into(); - send_response(&tx, Ok(schema_flight_data)).await?; - - let mut row_count = 0; - for batch in reader { - if let Ok(x) = &batch { - row_count += x.num_rows(); - } - let batch_flight_data: Vec<_> = batch - .map(|b| create_flight_iter(&b, &options).collect()) - .map_err(|e| from_arrow_err(&e))?; - for batch in batch_flight_data.into_iter() { - send_response(&tx, batch).await?; - } - } - info!("FetchPartition streamed {} rows", row_count); - Ok(()) -} - -async fn send_response( - tx: &FlightDataSender, - data: Result, -) -> Result<(), Status> { - tx.send(data) - .await - .map_err(|e| Status::internal(format!("{:?}", e))) -} - -fn from_arrow_err(e: &ArrowError) -> Status { - Status::internal(format!("ArrowError: {:?}", e)) -} - -fn from_ballista_err(e: &ballista_core::error::BallistaError) -> Status { - Status::internal(format!("Ballista Error: {:?}", e)) -} diff --git a/ballista/rust/executor/src/lib.rs b/ballista/rust/executor/src/lib.rs deleted file mode 100644 index 4d145b269efd..000000000000 --- a/ballista/rust/executor/src/lib.rs +++ /dev/null @@ -1,68 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -#![doc = include_str!("../README.md")] - -pub mod collect; -pub mod execution_loop; -pub mod executor; -pub mod executor_server; -pub mod flight_service; -pub mod metrics; - -mod cpu_bound_executor; -mod standalone; - -pub use standalone::new_standalone_executor; - -use log::info; - -use ballista_core::serde::protobuf::{ - task_status, CompletedTask, FailedTask, PartitionId, ShuffleWritePartition, - TaskStatus, -}; - -pub fn as_task_status( - execution_result: ballista_core::error::Result>, - executor_id: String, - task_id: PartitionId, -) -> TaskStatus { - match execution_result { - Ok(partitions) => { - info!("Task {:?} finished", task_id); - - TaskStatus { - task_id: Some(task_id), - status: Some(task_status::Status::Completed(CompletedTask { - executor_id, - partitions, - })), - } - } - Err(e) => { - let error_msg = e.to_string(); - info!("Task {:?} failed: {}", task_id, error_msg); - - TaskStatus { - task_id: Some(task_id), - status: Some(task_status::Status::Failed(FailedTask { - error: format!("Task failed due to Tokio error: {}", error_msg), - })), - } - } - } -} diff --git a/ballista/rust/executor/src/main.rs b/ballista/rust/executor/src/main.rs deleted file mode 100644 index 825ddd4d8fa5..000000000000 --- a/ballista/rust/executor/src/main.rs +++ /dev/null @@ -1,295 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Ballista Rust executor binary. - -use chrono::{DateTime, Duration, Utc}; -use std::sync::Arc; -use std::time::Duration as Core_Duration; - -use anyhow::{Context, Result}; -use arrow_flight::flight_service_server::FlightServiceServer; -use ballista_executor::{execution_loop, executor_server}; -use log::{error, info}; -use tempfile::TempDir; -use tokio::fs::ReadDir; -use tokio::{fs, time}; -use tonic::transport::Server; -use uuid::Uuid; - -use ballista_core::config::TaskSchedulingPolicy; -use ballista_core::error::BallistaError; -use ballista_core::serde::protobuf::{ - executor_registration, scheduler_grpc_client::SchedulerGrpcClient, - ExecutorRegistration, LogicalPlanNode, PhysicalPlanNode, -}; -use ballista_core::serde::scheduler::ExecutorSpecification; -use ballista_core::serde::BallistaCodec; -use ballista_core::{print_version, BALLISTA_VERSION}; -use ballista_executor::executor::Executor; -use ballista_executor::flight_service::BallistaFlightService; -use ballista_executor::metrics::LoggingMetricsCollector; -use config::prelude::*; -use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv}; - -#[macro_use] -extern crate configure_me; - -#[allow(clippy::all, warnings)] -mod config { - // Ideally we would use the include_config macro from configure_me, but then we cannot use - // #[allow(clippy::all)] to silence clippy warnings from the generated code - include!(concat!(env!("OUT_DIR"), "/executor_configure_me_config.rs")); -} - -#[cfg(feature = "snmalloc")] -#[global_allocator] -static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; - -#[tokio::main] -async fn main() -> Result<()> { - env_logger::init(); - - // parse command-line arguments - let (opt, _remaining_args) = - Config::including_optional_config_files(&["/etc/ballista/executor.toml"]) - .unwrap_or_exit(); - - if opt.version { - print_version(); - std::process::exit(0); - } - - let external_host = opt.external_host; - let bind_host = opt.bind_host; - let port = opt.bind_port; - let grpc_port = opt.bind_grpc_port; - - let addr = format!("{}:{}", bind_host, port); - let addr = addr - .parse() - .with_context(|| format!("Could not parse address: {}", addr))?; - - let scheduler_host = opt.scheduler_host; - let scheduler_port = opt.scheduler_port; - let scheduler_url = format!("http://{}:{}", scheduler_host, scheduler_port); - - let work_dir = opt.work_dir.unwrap_or( - TempDir::new()? - .into_path() - .into_os_string() - .into_string() - .unwrap(), - ); - info!("Running with config:"); - info!("work_dir: {}", work_dir); - info!("concurrent_tasks: {}", opt.concurrent_tasks); - - let executor_meta = ExecutorRegistration { - id: Uuid::new_v4().to_string(), // assign this executor a unique ID - optional_host: external_host - .clone() - .map(executor_registration::OptionalHost::Host), - port: port as u32, - grpc_port: grpc_port as u32, - specification: Some( - ExecutorSpecification { - task_slots: opt.concurrent_tasks as u32, - } - .into(), - ), - }; - - let config = RuntimeConfig::new().with_temp_file_path(work_dir.clone()); - let runtime = Arc::new(RuntimeEnv::new(config).map_err(|_| { - BallistaError::Internal("Failed to init Executor RuntimeEnv".to_owned()) - })?); - - let metrics_collector = Arc::new(LoggingMetricsCollector::default()); - - let executor = Arc::new(Executor::new( - executor_meta, - &work_dir, - runtime, - metrics_collector, - )); - - let scheduler = SchedulerGrpcClient::connect(scheduler_url) - .await - .context("Could not connect to scheduler")?; - - let default_codec: BallistaCodec = - BallistaCodec::default(); - - let scheduler_policy = opt.task_scheduling_policy; - let cleanup_ttl = opt.executor_cleanup_ttl; - - if opt.executor_cleanup_enable { - let mut interval_time = - time::interval(Core_Duration::from_secs(opt.executor_cleanup_interval)); - tokio::spawn(async move { - loop { - interval_time.tick().await; - if let Err(e) = - clean_shuffle_data_loop(&work_dir, cleanup_ttl as i64).await - { - error!("Ballista executor fail to clean_shuffle_data {:?}", e) - } - } - }); - } - - match scheduler_policy { - TaskSchedulingPolicy::PushStaged => { - tokio::spawn(executor_server::startup( - scheduler, - executor.clone(), - default_codec, - )); - } - _ => { - tokio::spawn(execution_loop::poll_loop( - scheduler, - executor.clone(), - default_codec, - )); - } - } - - // Arrow flight service - { - let service = BallistaFlightService::new(executor.clone()); - let server = FlightServiceServer::new(service); - info!( - "Ballista v{} Rust Executor listening on {:?}", - BALLISTA_VERSION, addr - ); - let server_future = - tokio::spawn(Server::builder().add_service(server).serve(addr)); - server_future - .await - .context("Tokio error")? - .context("Could not start executor server")?; - } - - Ok(()) -} - -/// This function will scheduled periodically for cleanup executor. -/// Will only clean the dir under work_dir not include file -async fn clean_shuffle_data_loop(work_dir: &str, seconds: i64) -> Result<()> { - let mut dir = fs::read_dir(work_dir).await?; - let mut to_deleted = Vec::new(); - let mut need_delete_dir; - while let Some(child) = dir.next_entry().await? { - if let Ok(metadata) = child.metadata().await { - // only delete the job dir - if metadata.is_dir() { - let dir = fs::read_dir(child.path()).await?; - match check_modified_time_in_dirs(vec![dir], seconds).await { - Ok(x) => match x { - true => { - need_delete_dir = child.path().into_os_string(); - to_deleted.push(need_delete_dir) - } - false => {} - }, - Err(e) => { - error!("Fail in clean_shuffle_data_loop {:?}", e) - } - } - } - } else { - error!("Can not get metadata from file: {:?}", child) - } - } - info!( - "The work_dir {:?} that have not been modified for {:?} seconds will be deleted", - &to_deleted, seconds - ); - for del in to_deleted { - fs::remove_dir_all(del).await?; - } - Ok(()) -} - -/// Determines if a directory all files are older than cutoff seconds. -async fn check_modified_time_in_dirs( - mut vec: Vec, - ttl_seconds: i64, -) -> Result { - let cutoff = Utc::now() - Duration::seconds(ttl_seconds); - - while !vec.is_empty() { - let mut dir = vec.pop().unwrap(); - while let Some(child) = dir.next_entry().await? { - let meta = child.metadata().await?; - if meta.is_dir() { - let dir = fs::read_dir(child.path()).await?; - // check in next loop - vec.push(dir); - } else { - let modified_time: DateTime = - meta.modified().map(chrono::DateTime::from)?; - if modified_time > cutoff { - // if one file has been modified in ttl we won't delete the whole dir - return Ok(false); - } - } - } - } - Ok(true) -} - -#[cfg(test)] -mod tests { - use crate::clean_shuffle_data_loop; - use std::fs; - use std::fs::File; - use std::io::Write; - use std::time::Duration; - use tempfile::TempDir; - - #[tokio::test] - async fn test_executor_clean_up() { - let work_dir = TempDir::new().unwrap().into_path(); - let job_dir = work_dir.as_path().join("job_id"); - let file_path = job_dir.as_path().join("tmp.csv"); - let data = "Jorge,2018-12-13T12:12:10.011Z\n\ - Andrew,2018-11-13T17:11:10.011Z"; - fs::create_dir(job_dir).unwrap(); - File::create(&file_path) - .expect("creating temp file") - .write_all(data.as_bytes()) - .expect("writing data"); - - let work_dir_clone = work_dir.clone(); - - let count1 = fs::read_dir(work_dir.clone()).unwrap().count(); - assert_eq!(count1, 1); - let mut handles = vec![]; - handles.push(tokio::spawn(async move { - tokio::time::sleep(Duration::from_secs(2)).await; - clean_shuffle_data_loop(work_dir_clone.to_str().unwrap(), 1) - .await - .unwrap(); - })); - futures::future::join_all(handles).await; - let count2 = fs::read_dir(work_dir.clone()).unwrap().count(); - assert_eq!(count2, 0); - } -} diff --git a/ballista/rust/executor/src/metrics/mod.rs b/ballista/rust/executor/src/metrics/mod.rs deleted file mode 100644 index 2c7e1d504141..000000000000 --- a/ballista/rust/executor/src/metrics/mod.rs +++ /dev/null @@ -1,58 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use ballista_core::execution_plans::ShuffleWriterExec; -use datafusion::physical_plan::display::DisplayableExecutionPlan; - -/// `ExecutorMetricsCollector` records metrics for `ShuffleWriteExec` -/// after they are executed. -/// -/// After each stage completes, `ShuffleWriteExec::record_stage` will be -/// called. -pub trait ExecutorMetricsCollector: Send + Sync { - /// Record metrics for stage after it is executed - fn record_stage( - &self, - job_id: &str, - stage_id: usize, - partition: usize, - plan: ShuffleWriterExec, - ); -} - -/// Implementation of `ExecutorMetricsCollector` which logs the completed -/// plan to stdout. -#[derive(Default)] -pub struct LoggingMetricsCollector {} - -impl ExecutorMetricsCollector for LoggingMetricsCollector { - fn record_stage( - &self, - job_id: &str, - stage_id: usize, - partition: usize, - plan: ShuffleWriterExec, - ) { - println!( - "=== [{}/{}/{}] Physical plan with metrics ===\n{}\n", - job_id, - stage_id, - partition, - DisplayableExecutionPlan::with_metrics(&plan).indent() - ); - } -} diff --git a/ballista/rust/executor/src/standalone.rs b/ballista/rust/executor/src/standalone.rs deleted file mode 100644 index d68052af8c7a..000000000000 --- a/ballista/rust/executor/src/standalone.rs +++ /dev/null @@ -1,95 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::Arc; - -use arrow_flight::flight_service_server::FlightServiceServer; - -use ballista_core::serde::scheduler::ExecutorSpecification; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec}; -use ballista_core::{ - error::Result, - serde::protobuf::executor_registration::OptionalHost, - serde::protobuf::{scheduler_grpc_client::SchedulerGrpcClient, ExecutorRegistration}, - BALLISTA_VERSION, -}; -use datafusion::execution::runtime_env::{RuntimeConfig, RuntimeEnv}; -use log::info; -use tempfile::TempDir; -use tokio::net::TcpListener; -use tonic::transport::{Channel, Server}; -use uuid::Uuid; - -use crate::metrics::LoggingMetricsCollector; -use crate::{execution_loop, executor::Executor, flight_service::BallistaFlightService}; - -pub async fn new_standalone_executor< - T: 'static + AsLogicalPlan, - U: 'static + AsExecutionPlan, ->( - scheduler: SchedulerGrpcClient, - concurrent_tasks: usize, - codec: BallistaCodec, -) -> Result<()> { - // Let the OS assign a random, free port - let listener = TcpListener::bind("localhost:0").await?; - let addr = listener.local_addr()?; - info!( - "Ballista v{} Rust Executor listening on {:?}", - BALLISTA_VERSION, addr - ); - - let executor_meta = ExecutorRegistration { - id: Uuid::new_v4().to_string(), // assign this executor a unique ID - optional_host: Some(OptionalHost::Host("localhost".to_string())), - port: addr.port() as u32, - // TODO Make it configurable - grpc_port: 50020, - specification: Some( - ExecutorSpecification { - task_slots: concurrent_tasks as u32, - } - .into(), - ), - }; - let work_dir = TempDir::new()? - .into_path() - .into_os_string() - .into_string() - .unwrap(); - info!("work_dir: {}", work_dir); - - let config = RuntimeConfig::new().with_temp_file_path(work_dir.clone()); - - let executor = Arc::new(Executor::new( - executor_meta, - &work_dir, - Arc::new(RuntimeEnv::new(config).unwrap()), - Arc::new(LoggingMetricsCollector::default()), - )); - - let service = BallistaFlightService::new(executor.clone()); - let server = FlightServiceServer::new(service); - tokio::spawn( - Server::builder().add_service(server).serve_with_incoming( - tokio_stream::wrappers::TcpListenerStream::new(listener), - ), - ); - - tokio::spawn(execution_loop::poll_loop(scheduler, executor, codec)); - Ok(()) -} diff --git a/ballista/rust/scheduler/Cargo.toml b/ballista/rust/scheduler/Cargo.toml deleted file mode 100644 index 50509f1eea8d..000000000000 --- a/ballista/rust/scheduler/Cargo.toml +++ /dev/null @@ -1,70 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -[package] -name = "ballista-scheduler" -description = "Ballista Distributed Compute - Scheduler" -license = "Apache-2.0" -version = "0.7.0" -homepage = "https://github.com/apache/arrow-datafusion" -repository = "https://github.com/apache/arrow-datafusion" -readme = "README.md" -authors = ["Apache Arrow "] -edition = "2018" - -[package.metadata.configure_me.bin] -scheduler = "scheduler_config_spec.toml" - -[features] -default = ["etcd", "sled"] -etcd = ["etcd-client"] -sled = ["sled_package", "tokio-stream"] - -[dependencies] -anyhow = "1" -async-recursion = "1.0.0" -async-trait = "0.1.41" -ballista-core = { path = "../core", version = "0.7.0" } -clap = { version = "3", features = ["derive", "cargo"] } -configure_me = "0.4.0" -datafusion = { path = "../../../datafusion/core", version = "8.0.0" } -env_logger = "0.9" -etcd-client = { version = "0.9", optional = true } -futures = "0.3" -http = "0.2" -http-body = "0.4" -hyper = "0.14.4" -log = "0.4" -parking_lot = "0.12" -parse_arg = "0.1.3" -prost = "0.10" -rand = "0.8" -serde = { version = "1", features = ["derive"] } -sled_package = { package = "sled", version = "0.34", optional = true } -tokio = { version = "1.0", features = ["full"] } -tokio-stream = { version = "0.1", features = ["net"], optional = true } -tonic = "0.7" -tower = { version = "0.4" } -warp = "0.3" - -[dev-dependencies] -ballista-core = { path = "../core", version = "0.7.0" } -uuid = { version = "1.0", features = ["v4"] } - -[build-dependencies] -configure_me_codegen = "0.4.1" -tonic-build = { version = "0.7" } diff --git a/ballista/rust/scheduler/README.md b/ballista/rust/scheduler/README.md deleted file mode 100644 index 382cce8cb7bb..000000000000 --- a/ballista/rust/scheduler/README.md +++ /dev/null @@ -1,22 +0,0 @@ - - -# Ballista Scheduler Process - -This crate contains the Ballista scheduler process. diff --git a/ballista/rust/scheduler/build.rs b/ballista/rust/scheduler/build.rs deleted file mode 100644 index e90bd495a9e4..000000000000 --- a/ballista/rust/scheduler/build.rs +++ /dev/null @@ -1,29 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -extern crate configure_me_codegen; - -fn main() -> Result<(), String> { - println!("cargo:rerun-if-changed=scheduler_config_spec.toml"); - configure_me_codegen::build_script_auto() - .map_err(|e| format!("configure_me code generation failed: {}", e))?; - - println!("cargo:rerun-if-changed=proto/keda.proto"); - tonic_build::configure() - .compile(&["proto/keda.proto"], &["proto"]) - .map_err(|e| format!("protobuf compilation failed: {}", e)) -} diff --git a/ballista/rust/scheduler/proto/keda.proto b/ballista/rust/scheduler/proto/keda.proto deleted file mode 100644 index 051dd438f41a..000000000000 --- a/ballista/rust/scheduler/proto/keda.proto +++ /dev/null @@ -1,63 +0,0 @@ -/* - Copyright 2020 The KEDA Authors. - - and others that have contributed code to the public domain. - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at. - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. -*/ -// This file comes from https://github.com/kedacore/keda/blob/main/pkg/scalers/externalscaler/externalscaler.proto -syntax = "proto3"; - -package externalscaler; -option go_package = ".;externalscaler"; - -service ExternalScaler { - rpc IsActive(ScaledObjectRef) returns (IsActiveResponse) {} - // Commented out since we aren't supporting the streaming scaler interface at the moment - // rpc StreamIsActive(ScaledObjectRef) returns (stream IsActiveResponse) {} - rpc GetMetricSpec(ScaledObjectRef) returns (GetMetricSpecResponse) {} - rpc GetMetrics(GetMetricsRequest) returns (GetMetricsResponse) {} -} - -message ScaledObjectRef { - string name = 1; - string namespace = 2; - map scalerMetadata = 3; -} - -message IsActiveResponse { - bool result = 1; -} - -message GetMetricSpecResponse { - repeated MetricSpec metricSpecs = 1; -} - -message MetricSpec { - string metricName = 1; - int64 targetSize = 2; -} - -message GetMetricsRequest { - ScaledObjectRef scaledObjectRef = 1; - string metricName = 2; -} - -message GetMetricsResponse { - repeated MetricValue metricValues = 1; -} - -message MetricValue { - string metricName = 1; - int64 metricValue = 2; -} \ No newline at end of file diff --git a/ballista/rust/scheduler/scheduler_config_spec.toml b/ballista/rust/scheduler/scheduler_config_spec.toml deleted file mode 100644 index cca96edfa954..000000000000 --- a/ballista/rust/scheduler/scheduler_config_spec.toml +++ /dev/null @@ -1,73 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -[general] -name = "Ballista Scheduler" -env_prefix = "BALLISTA_SCHEDULER" -conf_file_param = "config_file" - -[[switch]] -name = "version" -doc = "Print version of this executable" - -[[param]] -abbr = "b" -name = "config_backend" -type = "ballista_scheduler::state::backend::StateBackend" -doc = "The configuration backend for the scheduler, see StateBackend::variants() for options. Default: Standalone" -default = "ballista_scheduler::state::backend::StateBackend::Standalone" - -[[param]] -abbr = "n" -name = "namespace" -type = "String" -doc = "Namespace for the ballista cluster that this executor will join. Default: ballista" -default = "std::string::String::from(\"ballista\")" - -[[param]] -abbr = "e" -name = "etcd_urls" -type = "String" -doc = "etcd urls for use when discovery mode is `etcd`. Default: localhost:2379" -default = "std::string::String::from(\"localhost:2379\")" - -[[param]] -abbr = "h" -name = "bind_host" -type = "String" -default = "std::string::String::from(\"0.0.0.0\")" -doc = "Local host name or IP address to bind to. Default: 0.0.0.0" - -[[param]] -abbr = "p" -name = "bind_port" -type = "u16" -default = "50050" -doc = "bind port. Default: 50050" - -[[param]] -abbr = "s" -name = "scheduler_policy" -type = "ballista_core::config::TaskSchedulingPolicy" -doc = "The scheduing policy for the scheduler, see TaskSchedulingPolicy::variants() for options. Default: PullStaged" -default = "ballista_core::config::TaskSchedulingPolicy::PullStaged" - -[[param]] -name = "plugin_dir" -type = "String" -doc = "plugin dir" -default = "std::string::String::from(\"\")" \ No newline at end of file diff --git a/ballista/rust/scheduler/src/api/handlers.rs b/ballista/rust/scheduler/src/api/handlers.rs deleted file mode 100644 index 72f17445a52f..000000000000 --- a/ballista/rust/scheduler/src/api/handlers.rs +++ /dev/null @@ -1,56 +0,0 @@ -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use crate::scheduler_server::SchedulerServer; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan}; -use ballista_core::BALLISTA_VERSION; -use warp::Rejection; - -#[derive(Debug, serde::Serialize)] -struct StateResponse { - executors: Vec, - started: u128, - version: &'static str, -} - -#[derive(Debug, serde::Serialize)] -pub struct ExecutorMetaResponse { - pub id: String, - pub host: String, - pub port: u16, - pub last_seen: u128, -} - -pub(crate) async fn scheduler_state( - data_server: SchedulerServer, -) -> Result { - // TODO: Display last seen information in UI - let executors: Vec = data_server - .state - .get_executors_metadata() - .await - .unwrap_or_default() - .into_iter() - .map(|(metadata, duration)| ExecutorMetaResponse { - id: metadata.id, - host: metadata.host, - port: metadata.port, - last_seen: duration.as_millis(), - }) - .collect(); - let response = StateResponse { - executors, - started: data_server.start_time, - version: BALLISTA_VERSION, - }; - Ok(warp::reply::json(&response)) -} diff --git a/ballista/rust/scheduler/src/api/mod.rs b/ballista/rust/scheduler/src/api/mod.rs deleted file mode 100644 index 98b7046398c5..000000000000 --- a/ballista/rust/scheduler/src/api/mod.rs +++ /dev/null @@ -1,91 +0,0 @@ -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -mod handlers; - -use crate::scheduler_server::SchedulerServer; -use anyhow::Result; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan}; -use std::{ - pin::Pin, - task::{Context as TaskContext, Poll}, -}; -use warp::filters::BoxedFilter; -use warp::{Buf, Filter, Reply}; - -pub enum EitherBody { - Left(A), - Right(B), -} - -pub type Error = Box; -pub type HttpBody = dyn http_body::Body + 'static; - -impl http_body::Body for EitherBody -where - A: http_body::Body + Send + Unpin, - B: http_body::Body + Send + Unpin, - A::Error: Into, - B::Error: Into, -{ - type Data = A::Data; - type Error = Error; - - fn poll_data( - self: Pin<&mut Self>, - cx: &mut TaskContext<'_>, - ) -> Poll>> { - match self.get_mut() { - EitherBody::Left(b) => Pin::new(b).poll_data(cx).map(map_option_err), - EitherBody::Right(b) => Pin::new(b).poll_data(cx).map(map_option_err), - } - } - - fn poll_trailers( - self: Pin<&mut Self>, - cx: &mut TaskContext<'_>, - ) -> Poll, Self::Error>> { - match self.get_mut() { - EitherBody::Left(b) => Pin::new(b).poll_trailers(cx).map_err(Into::into), - EitherBody::Right(b) => Pin::new(b).poll_trailers(cx).map_err(Into::into), - } - } - - fn is_end_stream(&self) -> bool { - match self { - EitherBody::Left(b) => b.is_end_stream(), - EitherBody::Right(b) => b.is_end_stream(), - } - } -} - -fn map_option_err>( - err: Option>, -) -> Option> { - err.map(|e| e.map_err(Into::into)) -} - -fn with_data_server( - db: SchedulerServer, -) -> impl Filter,), Error = std::convert::Infallible> + Clone -{ - warp::any().map(move || db.clone()) -} - -pub fn get_routes( - scheduler_server: SchedulerServer, -) -> BoxedFilter<(impl Reply,)> { - let routes = warp::path("state") - .and(with_data_server(scheduler_server)) - .and_then(handlers::scheduler_state); - routes.boxed() -} diff --git a/ballista/rust/scheduler/src/lib.rs b/ballista/rust/scheduler/src/lib.rs deleted file mode 100644 index ea39ef02efd4..000000000000 --- a/ballista/rust/scheduler/src/lib.rs +++ /dev/null @@ -1,28 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -#![doc = include_str ! ("../README.md")] - -pub mod api; -pub mod planner; -pub mod scheduler_server; -#[cfg(feature = "sled")] -pub mod standalone; -pub mod state; - -#[cfg(test)] -pub mod test_utils; diff --git a/ballista/rust/scheduler/src/main.rs b/ballista/rust/scheduler/src/main.rs deleted file mode 100644 index 4304821823e0..000000000000 --- a/ballista/rust/scheduler/src/main.rs +++ /dev/null @@ -1,199 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Ballista Rust scheduler binary. - -use anyhow::{Context, Result}; -use ballista_scheduler::scheduler_server::externalscaler::external_scaler_server::ExternalScalerServer; -use futures::future::{self, Either, TryFutureExt}; -use hyper::{server::conn::AddrStream, service::make_service_fn, Server}; -use std::convert::Infallible; -use std::{net::SocketAddr, sync::Arc}; -use tonic::transport::server::Connected; -use tonic::transport::Server as TonicServer; -use tower::Service; - -use ballista_core::BALLISTA_VERSION; -use ballista_core::{ - print_version, - serde::protobuf::{ - scheduler_grpc_server::SchedulerGrpcServer, LogicalPlanNode, PhysicalPlanNode, - }, -}; -use ballista_scheduler::api::{get_routes, EitherBody, Error}; -#[cfg(feature = "etcd")] -use ballista_scheduler::state::backend::etcd::EtcdClient; -#[cfg(feature = "sled")] -use ballista_scheduler::state::backend::standalone::StandaloneClient; - -use ballista_scheduler::scheduler_server::SchedulerServer; -use ballista_scheduler::state::backend::{StateBackend, StateBackendClient}; - -use ballista_core::config::TaskSchedulingPolicy; -use ballista_core::serde::BallistaCodec; -use log::info; - -#[macro_use] -extern crate configure_me; - -#[allow(clippy::all, warnings)] -mod config { - // Ideally we would use the include_config macro from configure_me, but then we cannot use - // #[allow(clippy::all)] to silence clippy warnings from the generated code - include!(concat!( - env!("OUT_DIR"), - "/scheduler_configure_me_config.rs" - )); -} - -use config::prelude::*; -use datafusion::execution::context::default_session_builder; - -async fn start_server( - config_backend: Arc, - namespace: String, - addr: SocketAddr, - policy: TaskSchedulingPolicy, -) -> Result<()> { - info!( - "Ballista v{} Scheduler listening on {:?}", - BALLISTA_VERSION, addr - ); - // Should only call SchedulerServer::new() once in the process - info!( - "Starting Scheduler grpc server with task scheduling policy of {:?}", - policy - ); - let mut scheduler_server: SchedulerServer = - match policy { - TaskSchedulingPolicy::PushStaged => SchedulerServer::new_with_policy( - config_backend.clone(), - namespace.clone(), - policy, - BallistaCodec::default(), - default_session_builder, - ), - _ => SchedulerServer::new( - config_backend.clone(), - namespace.clone(), - BallistaCodec::default(), - ), - }; - - scheduler_server.init().await?; - - Server::bind(&addr) - .serve(make_service_fn(move |request: &AddrStream| { - let scheduler_grpc_server = - SchedulerGrpcServer::new(scheduler_server.clone()); - - let keda_scaler = ExternalScalerServer::new(scheduler_server.clone()); - - let mut tonic = TonicServer::builder() - .add_service(scheduler_grpc_server) - .add_service(keda_scaler) - .into_service(); - let mut warp = warp::service(get_routes(scheduler_server.clone())); - - let connect_info = request.connect_info(); - future::ok::<_, Infallible>(tower::service_fn( - move |req: hyper::Request| { - // Set the connect info from hyper to tonic - let (mut parts, body) = req.into_parts(); - parts.extensions.insert(connect_info.clone()); - let req = http::Request::from_parts(parts, body); - - let header = req.headers().get(hyper::header::ACCEPT); - if header.is_some() && header.unwrap().eq("application/json") { - return Either::Left( - warp.call(req) - .map_ok(|res| res.map(EitherBody::Left)) - .map_err(Error::from), - ); - } - Either::Right( - tonic - .call(req) - .map_ok(|res| res.map(EitherBody::Right)) - .map_err(Error::from), - ) - }, - )) - })) - .await - .context("Could not start grpc server") -} - -#[tokio::main] -async fn main() -> Result<()> { - env_logger::init(); - - // parse options - let (opt, _remaining_args) = - Config::including_optional_config_files(&["/etc/ballista/scheduler.toml"]) - .unwrap_or_exit(); - - if opt.version { - print_version(); - std::process::exit(0); - } - - let namespace = opt.namespace; - let bind_host = opt.bind_host; - let port = opt.bind_port; - - let addr = format!("{}:{}", bind_host, port); - let addr = addr.parse()?; - - let client: Arc = match opt.config_backend { - #[cfg(not(any(feature = "sled", feature = "etcd")))] - _ => std::compile_error!( - "To build the scheduler enable at least one config backend feature (`etcd` or `sled`)" - ), - #[cfg(feature = "etcd")] - StateBackend::Etcd => { - let etcd = etcd_client::Client::connect(&[opt.etcd_urls], None) - .await - .context("Could not connect to etcd")?; - Arc::new(EtcdClient::new(etcd)) - } - #[cfg(not(feature = "etcd"))] - StateBackend::Etcd => { - unimplemented!( - "build the scheduler with the `etcd` feature to use the etcd config backend" - ) - } - #[cfg(feature = "sled")] - StateBackend::Standalone => { - // TODO: Use a real file and make path is configurable - Arc::new( - StandaloneClient::try_new_temporary() - .context("Could not create standalone config backend")?, - ) - } - #[cfg(not(feature = "sled"))] - StateBackend::Standalone => { - unimplemented!( - "build the scheduler with the `sled` feature to use the standalone config backend" - ) - } - }; - - let policy: TaskSchedulingPolicy = opt.scheduler_policy; - start_server(client, namespace, addr, policy).await?; - Ok(()) -} diff --git a/ballista/rust/scheduler/src/planner.rs b/ballista/rust/scheduler/src/planner.rs deleted file mode 100644 index af1ce1c612f6..000000000000 --- a/ballista/rust/scheduler/src/planner.rs +++ /dev/null @@ -1,619 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Distributed query execution -//! -//! This code is EXPERIMENTAL and still under development - -use std::collections::HashMap; -use std::sync::Arc; - -use ballista_core::error::{BallistaError, Result}; -use ballista_core::{ - execution_plans::{ShuffleReaderExec, ShuffleWriterExec, UnresolvedShuffleExec}, - serde::scheduler::PartitionLocation, -}; -use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion::physical_plan::repartition::RepartitionExec; -use datafusion::physical_plan::windows::WindowAggExec; -use datafusion::physical_plan::{ - with_new_children_if_necessary, ExecutionPlan, Partitioning, -}; -use futures::future::BoxFuture; -use futures::FutureExt; -use log::info; - -type PartialQueryStageResult = (Arc, Vec>); - -pub struct DistributedPlanner { - next_stage_id: usize, -} - -impl DistributedPlanner { - pub fn new() -> Self { - Self { next_stage_id: 0 } - } -} - -impl Default for DistributedPlanner { - fn default() -> Self { - Self::new() - } -} - -impl DistributedPlanner { - /// Returns a vector of ExecutionPlans, where the root node is a [ShuffleWriterExec]. - /// Plans that depend on the input of other plans will have leaf nodes of type [UnresolvedShuffleExec]. - /// A [ShuffleWriterExec] is created whenever the partitioning changes. - pub async fn plan_query_stages<'a>( - &'a mut self, - job_id: &'a str, - execution_plan: Arc, - ) -> Result>> { - info!("planning query stages"); - let (new_plan, mut stages) = self - .plan_query_stages_internal(job_id, execution_plan) - .await?; - stages.push(create_shuffle_writer( - job_id, - self.next_stage_id(), - new_plan, - None, - )?); - Ok(stages) - } - - /// Returns a potentially modified version of the input execution_plan along with the resulting query stages. - /// This function is needed because the input execution_plan might need to be modified, but it might not hold a - /// complete query stage (its parent might also belong to the same stage) - fn plan_query_stages_internal<'a>( - &'a mut self, - job_id: &'a str, - execution_plan: Arc, - ) -> BoxFuture<'a, Result> { - async move { - // recurse down and replace children - if execution_plan.children().is_empty() { - return Ok((execution_plan, vec![])); - } - - let mut stages = vec![]; - let mut children = vec![]; - for child in execution_plan.children() { - let (new_child, mut child_stages) = self - .plan_query_stages_internal(job_id, child.clone()) - .await?; - children.push(new_child); - stages.append(&mut child_stages); - } - - if let Some(_coalesce) = execution_plan - .as_any() - .downcast_ref::() - { - let shuffle_writer = create_shuffle_writer( - job_id, - self.next_stage_id(), - children[0].clone(), - None, - )?; - let unresolved_shuffle = Arc::new(UnresolvedShuffleExec::new( - shuffle_writer.stage_id(), - shuffle_writer.schema(), - shuffle_writer.output_partitioning().partition_count(), - shuffle_writer - .shuffle_output_partitioning() - .map(|p| p.partition_count()) - .unwrap_or_else(|| { - shuffle_writer.output_partitioning().partition_count() - }), - )); - stages.push(shuffle_writer); - Ok(( - with_new_children_if_necessary( - execution_plan, - vec![unresolved_shuffle], - )?, - stages, - )) - } else if let Some(repart) = - execution_plan.as_any().downcast_ref::() - { - match repart.output_partitioning() { - Partitioning::Hash(_, _) => { - let shuffle_writer = create_shuffle_writer( - job_id, - self.next_stage_id(), - children[0].clone(), - Some(repart.partitioning().to_owned()), - )?; - let unresolved_shuffle = Arc::new(UnresolvedShuffleExec::new( - shuffle_writer.stage_id(), - shuffle_writer.schema(), - shuffle_writer.output_partitioning().partition_count(), - shuffle_writer - .shuffle_output_partitioning() - .map(|p| p.partition_count()) - .unwrap_or_else(|| { - shuffle_writer.output_partitioning().partition_count() - }), - )); - stages.push(shuffle_writer); - Ok((unresolved_shuffle, stages)) - } - _ => { - // remove any non-hash repartition from the distributed plan - Ok((children[0].clone(), stages)) - } - } - } else if let Some(window) = - execution_plan.as_any().downcast_ref::() - { - Err(BallistaError::NotImplemented(format!( - "WindowAggExec with window {:?}", - window - ))) - } else { - Ok(( - with_new_children_if_necessary(execution_plan, children)?, - stages, - )) - } - } - .boxed() - } - - /// Generate a new stage ID - fn next_stage_id(&mut self) -> usize { - self.next_stage_id += 1; - self.next_stage_id - } -} - -/// Returns the unresolved shuffles in the execution plan -pub fn find_unresolved_shuffles( - plan: &Arc, -) -> Result> { - if let Some(unresolved_shuffle) = - plan.as_any().downcast_ref::() - { - Ok(vec![unresolved_shuffle.clone()]) - } else { - Ok(plan - .children() - .iter() - .map(find_unresolved_shuffles) - .collect::>>()? - .into_iter() - .flatten() - .collect()) - } -} - -pub fn remove_unresolved_shuffles( - stage: Arc, - partition_locations: &HashMap>>, -) -> Result> { - let mut new_children: Vec> = vec![]; - for child in stage.children() { - if let Some(unresolved_shuffle) = - child.as_any().downcast_ref::() - { - let mut relevant_locations = vec![]; - let p = partition_locations - .get(&unresolved_shuffle.stage_id) - .ok_or_else(|| { - BallistaError::General( - "Missing partition location. Could not remove unresolved shuffles" - .to_owned(), - ) - })? - .clone(); - - for i in 0..unresolved_shuffle.output_partition_count { - if let Some(x) = p.get(&i) { - relevant_locations.push(x.to_owned()); - } else { - relevant_locations.push(vec![]); - } - } - info!( - "Creating shuffle reader: {}", - relevant_locations - .iter() - .map(|c| c - .iter() - .map(|l| l.path.clone()) - .collect::>() - .join(", ")) - .collect::>() - .join("\n") - ); - new_children.push(Arc::new(ShuffleReaderExec::try_new( - relevant_locations, - unresolved_shuffle.schema().clone(), - )?)) - } else { - new_children.push(remove_unresolved_shuffles(child, partition_locations)?); - } - } - Ok(with_new_children_if_necessary(stage, new_children)?) -} - -fn create_shuffle_writer( - job_id: &str, - stage_id: usize, - plan: Arc, - partitioning: Option, -) -> Result> { - Ok(Arc::new(ShuffleWriterExec::try_new( - job_id.to_owned(), - stage_id, - plan, - "".to_owned(), // executor will decide on the work_dir path - partitioning, - )?)) -} - -#[cfg(test)] -mod test { - use crate::planner::DistributedPlanner; - use crate::test_utils::datafusion_test_context; - use ballista_core::error::BallistaError; - use ballista_core::execution_plans::UnresolvedShuffleExec; - use ballista_core::serde::{protobuf, AsExecutionPlan, BallistaCodec}; - use datafusion::physical_plan::aggregates::{AggregateExec, AggregateMode}; - use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; - use datafusion::physical_plan::hash_join::HashJoinExec; - use datafusion::physical_plan::sorts::sort::SortExec; - use datafusion::physical_plan::{ - coalesce_partitions::CoalescePartitionsExec, projection::ProjectionExec, - }; - use datafusion::physical_plan::{displayable, ExecutionPlan}; - use datafusion::prelude::SessionContext; - use std::ops::Deref; - - use ballista_core::serde::protobuf::{LogicalPlanNode, PhysicalPlanNode}; - use std::sync::Arc; - use uuid::Uuid; - - macro_rules! downcast_exec { - ($exec: expr, $ty: ty) => { - $exec.as_any().downcast_ref::<$ty>().unwrap() - }; - } - - #[tokio::test] - async fn distributed_aggregate_plan() -> Result<(), BallistaError> { - let ctx = datafusion_test_context("testdata").await?; - - // simplified form of TPC-H query 1 - let df = ctx - .sql( - "select l_returnflag, sum(l_extendedprice * 1) as sum_disc_price - from lineitem - group by l_returnflag - order by l_returnflag", - ) - .await?; - - let plan = df.to_logical_plan()?; - let plan = ctx.optimize(&plan)?; - let plan = ctx.create_physical_plan(&plan).await?; - - let mut planner = DistributedPlanner::new(); - let job_uuid = Uuid::new_v4(); - let stages = planner - .plan_query_stages(&job_uuid.to_string(), plan) - .await?; - for stage in &stages { - println!("{}", displayable(stage.as_ref()).indent()); - } - - /* Expected result: - - ShuffleWriterExec: Some(Hash([Column { name: "l_returnflag", index: 0 }], 2)) - AggregateExec: mode=Partial, gby=[l_returnflag@1 as l_returnflag], aggr=[SUM(l_extendedprice Multiply Int64(1))] - CsvExec: source=Path(testdata/lineitem: [testdata/lineitem/partition0.tbl,testdata/lineitem/partition1.tbl]), has_header=false - - ShuffleWriterExec: None - ProjectionExec: expr=[l_returnflag@0 as l_returnflag, SUM(lineitem.l_extendedprice Multiply Int64(1))@1 as sum_disc_price] - AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag], aggr=[SUM(l_extendedprice Multiply Int64(1))] - CoalesceBatchesExec: target_batch_size=4096 - UnresolvedShuffleExec - - ShuffleWriterExec: None - SortExec: [l_returnflag@0 ASC] - CoalescePartitionsExec - UnresolvedShuffleExec - */ - - assert_eq!(3, stages.len()); - - // verify stage 0 - let stage0 = stages[0].children()[0].clone(); - let partial_hash = downcast_exec!(stage0, AggregateExec); - assert!(*partial_hash.mode() == AggregateMode::Partial); - - // verify stage 1 - let stage1 = stages[1].children()[0].clone(); - let projection = downcast_exec!(stage1, ProjectionExec); - let final_hash = projection.children()[0].clone(); - let final_hash = downcast_exec!(final_hash, AggregateExec); - assert!(*final_hash.mode() == AggregateMode::FinalPartitioned); - let coalesce = final_hash.children()[0].clone(); - let coalesce = downcast_exec!(coalesce, CoalesceBatchesExec); - let unresolved_shuffle = coalesce.children()[0].clone(); - let unresolved_shuffle = - downcast_exec!(unresolved_shuffle, UnresolvedShuffleExec); - assert_eq!(unresolved_shuffle.stage_id, 1); - assert_eq!(unresolved_shuffle.input_partition_count, 2); - assert_eq!(unresolved_shuffle.output_partition_count, 2); - - // verify stage 2 - let stage2 = stages[2].children()[0].clone(); - let sort = downcast_exec!(stage2, SortExec); - let coalesce_partitions = sort.children()[0].clone(); - let coalesce_partitions = - downcast_exec!(coalesce_partitions, CoalescePartitionsExec); - assert_eq!( - coalesce_partitions.output_partitioning().partition_count(), - 1 - ); - let unresolved_shuffle = coalesce_partitions.children()[0].clone(); - let unresolved_shuffle = - downcast_exec!(unresolved_shuffle, UnresolvedShuffleExec); - assert_eq!(unresolved_shuffle.stage_id, 2); - assert_eq!(unresolved_shuffle.input_partition_count, 2); - assert_eq!(unresolved_shuffle.output_partition_count, 2); - - Ok(()) - } - - #[tokio::test] - async fn distributed_join_plan() -> Result<(), BallistaError> { - let ctx = datafusion_test_context("testdata").await?; - - // simplified form of TPC-H query 12 - let df = ctx - .sql( - "select - l_shipmode, - sum(case - when o_orderpriority = '1-URGENT' - or o_orderpriority = '2-HIGH' - then 1 - else 0 - end) as high_line_count, - sum(case - when o_orderpriority <> '1-URGENT' - and o_orderpriority <> '2-HIGH' - then 1 - else 0 - end) as low_line_count -from - lineitem - join - orders - on - l_orderkey = o_orderkey -where - l_shipmode in ('MAIL', 'SHIP') - and l_commitdate < l_receiptdate - and l_shipdate < l_commitdate - and l_receiptdate >= date '1994-01-01' - and l_receiptdate < date '1995-01-01' -group by - l_shipmode -order by - l_shipmode; -", - ) - .await?; - - let plan = df.to_logical_plan()?; - let plan = ctx.optimize(&plan)?; - let plan = ctx.create_physical_plan(&plan).await?; - - let mut planner = DistributedPlanner::new(); - let job_uuid = Uuid::new_v4(); - let stages = planner - .plan_query_stages(&job_uuid.to_string(), plan) - .await?; - for stage in &stages { - println!("{}", displayable(stage.as_ref()).indent()); - } - - /* Expected result: - - ShuffleWriterExec: Some(Hash([Column { name: "l_orderkey", index: 0 }], 2)) - CoalesceBatchesExec: target_batch_size=4096 - FilterExec: l_shipmode@4 IN ([Literal { value: Utf8("MAIL") }, Literal { value: Utf8("SHIP") }]) AND l_commitdate@2 < l_receiptdate@3 AND l_shipdate@1 < l_commitdate@2 AND l_receiptdate@3 >= 8766 AND l_receiptdate@3 < 9131 - CsvExec: source=Path(testdata/lineitem: [testdata/lineitem/partition0.tbl,testdata/lineitem/partition1.tbl]), has_header=false - - ShuffleWriterExec: Some(Hash([Column { name: "o_orderkey", index: 0 }], 2)) - CsvExec: source=Path(testdata/orders: [testdata/orders/orders.tbl]), has_header=false - - ShuffleWriterExec: Some(Hash([Column { name: "l_shipmode", index: 0 }], 2)) - AggregateExec: mode=Partial, gby=[l_shipmode@4 as l_shipmode], aggr=[SUM(CASE WHEN #orders.o_orderpriority Eq Utf8("1-URGENT") Or #orders.o_orderpriority Eq Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END), SUM(CASE WHEN #orders.o_orderpriority NotEq Utf8("1-URGENT") And #orders.o_orderpriority NotEq Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)] - CoalesceBatchesExec: target_batch_size=4096 - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(Column { name: "l_orderkey", index: 0 }, Column { name: "o_orderkey", index: 0 })] - CoalesceBatchesExec: target_batch_size=4096 - UnresolvedShuffleExec - CoalesceBatchesExec: target_batch_size=4096 - UnresolvedShuffleExec - - ShuffleWriterExec: None - ProjectionExec: expr=[l_shipmode@0 as l_shipmode, SUM(CASE WHEN #orders.o_orderpriority Eq Utf8("1-URGENT") Or #orders.o_orderpriority Eq Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)@1 as high_line_count, SUM(CASE WHEN #orders.o_orderpriority NotEq Utf8("1-URGENT") And #orders.o_orderpriority NotEq Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)@2 as low_line_count] - AggregateExec: mode=FinalPartitioned, gby=[l_shipmode@0 as l_shipmode], aggr=[SUM(CASE WHEN #orders.o_orderpriority Eq Utf8("1-URGENT") Or #orders.o_orderpriority Eq Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END), SUM(CASE WHEN #orders.o_orderpriority NotEq Utf8("1-URGENT") And #orders.o_orderpriority NotEq Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)] - CoalesceBatchesExec: target_batch_size=4096 - UnresolvedShuffleExec - - ShuffleWriterExec: None - SortExec: [l_shipmode@0 ASC] - CoalescePartitionsExec - UnresolvedShuffleExec - */ - - assert_eq!(5, stages.len()); - - // verify partitioning for each stage - - // csv "lineitem" (2 files) - assert_eq!( - 2, - stages[0].children()[0] - .output_partitioning() - .partition_count() - ); - assert_eq!( - 2, - stages[0] - .shuffle_output_partitioning() - .unwrap() - .partition_count() - ); - - // csv "orders" (1 file) - assert_eq!( - 1, - stages[1].children()[0] - .output_partitioning() - .partition_count() - ); - assert_eq!( - 2, - stages[1] - .shuffle_output_partitioning() - .unwrap() - .partition_count() - ); - - // join and partial hash aggregate - let input = stages[2].children()[0].clone(); - assert_eq!(2, input.output_partitioning().partition_count()); - assert_eq!( - 2, - stages[2] - .shuffle_output_partitioning() - .unwrap() - .partition_count() - ); - - let hash_agg = downcast_exec!(input, AggregateExec); - - let coalesce_batches = hash_agg.children()[0].clone(); - let coalesce_batches = downcast_exec!(coalesce_batches, CoalesceBatchesExec); - - let join = coalesce_batches.children()[0].clone(); - let join = downcast_exec!(join, HashJoinExec); - - let join_input_1 = join.children()[0].clone(); - // skip CoalesceBatches - let join_input_1 = join_input_1.children()[0].clone(); - let unresolved_shuffle_reader_1 = - downcast_exec!(join_input_1, UnresolvedShuffleExec); - assert_eq!(unresolved_shuffle_reader_1.input_partition_count, 2); // lineitem - assert_eq!(unresolved_shuffle_reader_1.output_partition_count, 2); - - let join_input_2 = join.children()[1].clone(); - // skip CoalesceBatches - let join_input_2 = join_input_2.children()[0].clone(); - let unresolved_shuffle_reader_2 = - downcast_exec!(join_input_2, UnresolvedShuffleExec); - assert_eq!(unresolved_shuffle_reader_2.input_partition_count, 1); // orders - assert_eq!(unresolved_shuffle_reader_2.output_partition_count, 2); - - // final partitioned hash aggregate - assert_eq!( - 2, - stages[3].children()[0] - .output_partitioning() - .partition_count() - ); - assert!(stages[3].shuffle_output_partitioning().is_none()); - - // coalesce partitions and sort - assert_eq!( - 1, - stages[4].children()[0] - .output_partitioning() - .partition_count() - ); - assert!(stages[4].shuffle_output_partitioning().is_none()); - - Ok(()) - } - - #[tokio::test] - async fn roundtrip_serde_aggregate() -> Result<(), BallistaError> { - let ctx = datafusion_test_context("testdata").await?; - - // simplified form of TPC-H query 1 - let df = ctx - .sql( - "select l_returnflag, sum(l_extendedprice * 1) as sum_disc_price - from lineitem - group by l_returnflag - order by l_returnflag", - ) - .await?; - - let plan = df.to_logical_plan()?; - let plan = ctx.optimize(&plan)?; - let plan = ctx.create_physical_plan(&plan).await?; - - let mut planner = DistributedPlanner::new(); - let job_uuid = Uuid::new_v4(); - let stages = planner - .plan_query_stages(&job_uuid.to_string(), plan) - .await?; - - let partial_hash = stages[0].children()[0].clone(); - let partial_hash_serde = roundtrip_operator(partial_hash.clone())?; - - let partial_hash = downcast_exec!(partial_hash, AggregateExec); - let partial_hash_serde = downcast_exec!(partial_hash_serde, AggregateExec); - - assert_eq!( - format!("{:?}", partial_hash), - format!("{:?}", partial_hash_serde) - ); - - Ok(()) - } - - fn roundtrip_operator( - plan: Arc, - ) -> Result, BallistaError> { - let ctx = SessionContext::new(); - let codec: BallistaCodec = - BallistaCodec::default(); - let proto: protobuf::PhysicalPlanNode = - protobuf::PhysicalPlanNode::try_from_physical_plan( - plan.clone(), - codec.physical_extension_codec(), - )?; - let runtime = ctx.runtime_env(); - let result_exec_plan: Arc = (&proto).try_into_physical_plan( - &ctx, - runtime.deref(), - codec.physical_extension_codec(), - )?; - Ok(result_exec_plan) - } -} diff --git a/ballista/rust/scheduler/src/scheduler_server/event.rs b/ballista/rust/scheduler/src/scheduler_server/event.rs deleted file mode 100644 index 9252453e6fe3..000000000000 --- a/ballista/rust/scheduler/src/scheduler_server/event.rs +++ /dev/null @@ -1,33 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use datafusion::physical_plan::ExecutionPlan; -use std::sync::Arc; - -#[derive(Clone)] -pub(crate) enum SchedulerServerEvent { - // number of offer rounds - ReviveOffers(u32), -} - -#[derive(Clone)] -pub enum QueryStageSchedulerEvent { - JobSubmitted(String, Arc), - StageFinished(String, u32), - JobFinished(String), - JobFailed(String, u32, String), -} diff --git a/ballista/rust/scheduler/src/scheduler_server/event_loop.rs b/ballista/rust/scheduler/src/scheduler_server/event_loop.rs deleted file mode 100644 index c343743ca35f..000000000000 --- a/ballista/rust/scheduler/src/scheduler_server/event_loop.rs +++ /dev/null @@ -1,169 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::Arc; -use std::time::Duration; - -use async_trait::async_trait; -use log::{debug, warn}; - -use crate::scheduler_server::event::SchedulerServerEvent; -use ballista_core::error::{BallistaError, Result}; -use ballista_core::event_loop::EventAction; -use ballista_core::serde::protobuf::{LaunchTaskParams, TaskDefinition}; -use ballista_core::serde::scheduler::ExecutorDataChange; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan}; - -use crate::scheduler_server::ExecutorsClient; -use crate::state::task_scheduler::TaskScheduler; -use crate::state::SchedulerState; - -pub(crate) struct SchedulerServerEventAction< - T: 'static + AsLogicalPlan, - U: 'static + AsExecutionPlan, -> { - state: Arc>, - executors_client: ExecutorsClient, -} - -impl - SchedulerServerEventAction -{ - pub fn new( - state: Arc>, - executors_client: ExecutorsClient, - ) -> Self { - Self { - state, - executors_client, - } - } - - #[allow(unused_variables)] - async fn offer_resources(&self, n: u32) -> Result> { - let mut available_executors = - self.state.executor_manager.get_available_executors_data(); - // In case of there's no enough resources, reschedule the tasks of the job - if available_executors.is_empty() { - // TODO Maybe it's better to use an exclusive runtime for this kind task scheduling - warn!("Not enough available executors for task running"); - tokio::time::sleep(Duration::from_millis(100)).await; - return Ok(Some(SchedulerServerEvent::ReviveOffers(1))); - } - - let mut executors_data_change: Vec = available_executors - .iter() - .map(|executor_data| ExecutorDataChange { - executor_id: executor_data.executor_id.clone(), - task_slots: executor_data.available_task_slots as i32, - }) - .collect(); - - let (tasks_assigment, num_tasks) = self - .state - .fetch_schedulable_tasks(&mut available_executors, n) - .await?; - for (data_change, data) in executors_data_change - .iter_mut() - .zip(available_executors.iter()) - { - data_change.task_slots = - data.available_task_slots as i32 - data_change.task_slots; - } - - #[cfg(not(test))] - if num_tasks > 0 { - self.launch_tasks(&executors_data_change, tasks_assigment) - .await?; - } - - Ok(None) - } - - #[allow(dead_code)] - async fn launch_tasks( - &self, - executors: &[ExecutorDataChange], - tasks_assigment: Vec>, - ) -> Result<()> { - for (idx_executor, tasks) in tasks_assigment.into_iter().enumerate() { - if !tasks.is_empty() { - let executor_data_change = &executors[idx_executor]; - debug!( - "Start to launch tasks {:?} to executor {:?}", - tasks - .iter() - .map(|task| { - if let Some(task_id) = task.task_id.as_ref() { - format!( - "{}/{}/{}", - task_id.job_id, - task_id.stage_id, - task_id.partition_id - ) - } else { - "".to_string() - } - }) - .collect::>(), - executor_data_change.executor_id - ); - let mut client = { - let clients = self.executors_client.read().await; - clients - .get(&executor_data_change.executor_id) - .unwrap() - .clone() - }; - // TODO check whether launching task is successful or not - client.launch_task(LaunchTaskParams { task: tasks }).await?; - self.state - .executor_manager - .update_executor_data(executor_data_change); - } else { - // Since the task assignment policy is round robin, - // if find tasks for one executor is empty, just break fast - break; - } - } - - Ok(()) - } -} - -#[async_trait] -impl - EventAction for SchedulerServerEventAction -{ - // TODO - fn on_start(&self) {} - - // TODO - fn on_stop(&self) {} - - async fn on_receive( - &self, - event: SchedulerServerEvent, - ) -> Result> { - match event { - SchedulerServerEvent::ReviveOffers(n) => self.offer_resources(n).await, - } - } - - // TODO - fn on_error(&self, _error: BallistaError) {} -} diff --git a/ballista/rust/scheduler/src/scheduler_server/external_scaler.rs b/ballista/rust/scheduler/src/scheduler_server/external_scaler.rs deleted file mode 100644 index 4b3966df2701..000000000000 --- a/ballista/rust/scheduler/src/scheduler_server/external_scaler.rs +++ /dev/null @@ -1,65 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use crate::scheduler_server::externalscaler::{ - external_scaler_server::ExternalScaler, GetMetricSpecResponse, GetMetricsRequest, - GetMetricsResponse, IsActiveResponse, MetricSpec, MetricValue, ScaledObjectRef, -}; -use crate::scheduler_server::SchedulerServer; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan}; -use log::debug; -use tonic::{Request, Response}; - -const INFLIGHT_TASKS_METRIC_NAME: &str = "inflight_tasks"; - -#[tonic::async_trait] -impl ExternalScaler - for SchedulerServer -{ - async fn is_active( - &self, - _request: Request, - ) -> Result, tonic::Status> { - let result = self.state.stage_manager.has_running_tasks(); - debug!("Are there active tasks? {}", result); - Ok(Response::new(IsActiveResponse { result })) - } - - async fn get_metric_spec( - &self, - _request: Request, - ) -> Result, tonic::Status> { - Ok(Response::new(GetMetricSpecResponse { - metric_specs: vec![MetricSpec { - metric_name: INFLIGHT_TASKS_METRIC_NAME.to_string(), - target_size: 1, - }], - })) - } - - async fn get_metrics( - &self, - _request: Request, - ) -> Result, tonic::Status> { - Ok(Response::new(GetMetricsResponse { - metric_values: vec![MetricValue { - metric_name: INFLIGHT_TASKS_METRIC_NAME.to_string(), - metric_value: 10000000, // A very high number to saturate the HPA - }], - })) - } -} diff --git a/ballista/rust/scheduler/src/scheduler_server/grpc.rs b/ballista/rust/scheduler/src/scheduler_server/grpc.rs deleted file mode 100644 index 10be478653ef..000000000000 --- a/ballista/rust/scheduler/src/scheduler_server/grpc.rs +++ /dev/null @@ -1,634 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use anyhow::Context; -use ballista_core::config::{BallistaConfig, TaskSchedulingPolicy}; -use ballista_core::error::BallistaError; -use ballista_core::serde::protobuf::execute_query_params::{OptionalSessionId, Query}; -use ballista_core::serde::protobuf::executor_grpc_client::ExecutorGrpcClient; -use ballista_core::serde::protobuf::executor_registration::OptionalHost; -use ballista_core::serde::protobuf::scheduler_grpc_server::SchedulerGrpc; -use ballista_core::serde::protobuf::{ - job_status, ExecuteQueryParams, ExecuteQueryResult, ExecutorHeartbeat, FailedJob, - FileType, GetFileMetadataParams, GetFileMetadataResult, GetJobStatusParams, - GetJobStatusResult, HeartBeatParams, HeartBeatResult, JobStatus, PollWorkParams, - PollWorkResult, QueuedJob, RegisterExecutorParams, RegisterExecutorResult, - UpdateTaskStatusParams, UpdateTaskStatusResult, -}; -use ballista_core::serde::scheduler::{ - ExecutorData, ExecutorDataChange, ExecutorMetadata, -}; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan}; -use datafusion::datafusion_data_access::object_store::{ - local::LocalFileSystem, ObjectStore, -}; -use datafusion::datasource::file_format::parquet::ParquetFormat; -use datafusion::datasource::file_format::FileFormat; -use futures::StreamExt; -use log::{debug, error, info, trace, warn}; -use rand::{distributions::Alphanumeric, thread_rng, Rng}; -use std::convert::TryInto; -use std::ops::Deref; -use std::sync::Arc; -use std::time::Instant; -use std::time::{SystemTime, UNIX_EPOCH}; -use tonic::{Request, Response, Status}; - -use crate::scheduler_server::event::QueryStageSchedulerEvent; -use crate::scheduler_server::{ - create_datafusion_context, update_datafusion_context, SchedulerServer, -}; -use crate::state::task_scheduler::TaskScheduler; - -#[tonic::async_trait] -impl SchedulerGrpc - for SchedulerServer -{ - async fn poll_work( - &self, - request: Request, - ) -> std::result::Result, tonic::Status> { - if let TaskSchedulingPolicy::PushStaged = self.policy { - error!("Poll work interface is not supported for push-based task scheduling"); - return Err(tonic::Status::failed_precondition( - "Bad request because poll work is not supported for push-based task scheduling", - )); - } - let remote_addr = request.remote_addr(); - if let PollWorkParams { - metadata: Some(metadata), - can_accept_task, - task_status, - } = request.into_inner() - { - debug!("Received poll_work request for {:?}", metadata); - let metadata = ExecutorMetadata { - id: metadata.id, - host: metadata - .optional_host - .map(|h| match h { - OptionalHost::Host(host) => host, - }) - .unwrap_or_else(|| remote_addr.unwrap().ip().to_string()), - port: metadata.port as u16, - grpc_port: metadata.grpc_port as u16, - specification: metadata.specification.unwrap().into(), - }; - let executor_heartbeat = ExecutorHeartbeat { - executor_id: metadata.id.clone(), - timestamp: SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("Time went backwards") - .as_secs(), - state: None, - }; - // In case that it's the first time to poll work, do registration - if self.state.get_executor_metadata(&metadata.id).is_none() { - self.state - .save_executor_metadata(metadata.clone()) - .await - .map_err(|e| { - let msg = format!("Could not save executor metadata: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - } - self.state - .executor_manager - .save_executor_heartbeat(executor_heartbeat); - self.update_task_status(task_status).await.map_err(|e| { - let msg = format!( - "Fail to update tasks status from executor {:?} due to {:?}", - &metadata.id, e - ); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - let task: Result, Status> = if can_accept_task { - let mut executors_data = vec![ExecutorData { - executor_id: metadata.id.clone(), - total_task_slots: 1, - available_task_slots: 1, - }]; - let (mut tasks, num_tasks) = self - .state - .fetch_schedulable_tasks(&mut executors_data, 1) - .await - .map_err(|e| { - let msg = format!("Error finding next assignable task: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - if num_tasks == 0 { - Ok(None) - } else { - assert_eq!(tasks.len(), 1); - let mut task = tasks.pop().unwrap(); - assert_eq!(task.len(), 1); - let task = task.pop().unwrap(); - Ok(Some(task)) - } - } else { - Ok(None) - }; - Ok(Response::new(PollWorkResult { task: task? })) - } else { - warn!("Received invalid executor poll_work request"); - Err(tonic::Status::invalid_argument( - "Missing metadata in request", - )) - } - } - - async fn register_executor( - &self, - request: Request, - ) -> Result, Status> { - let remote_addr = request.remote_addr(); - if let RegisterExecutorParams { - metadata: Some(metadata), - } = request.into_inner() - { - info!("Received register executor request for {:?}", metadata); - let metadata = ExecutorMetadata { - id: metadata.id, - host: metadata - .optional_host - .map(|h| match h { - OptionalHost::Host(host) => host, - }) - .unwrap_or_else(|| remote_addr.unwrap().ip().to_string()), - port: metadata.port as u16, - grpc_port: metadata.grpc_port as u16, - specification: metadata.specification.unwrap().into(), - }; - // Check whether the executor starts the grpc service - { - let executor_url = - format!("http://{}:{}", metadata.host, metadata.grpc_port); - info!("Connect to executor {:?}", executor_url); - let executor_client = ExecutorGrpcClient::connect(executor_url) - .await - .context("Could not connect to executor") - .map_err(|e| tonic::Status::internal(format!("{:?}", e)))?; - let mut clients = self.executors_client.as_ref().unwrap().write().await; - // TODO check duplicated registration - clients.insert(metadata.id.clone(), executor_client); - info!("Size of executor clients: {:?}", clients.len()); - } - self.state - .save_executor_metadata(metadata.clone()) - .await - .map_err(|e| { - let msg = format!("Could not save executor metadata: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - let executor_data = ExecutorData { - executor_id: metadata.id.clone(), - total_task_slots: metadata.specification.task_slots, - available_task_slots: metadata.specification.task_slots, - }; - self.state - .executor_manager - .save_executor_data(executor_data); - Ok(Response::new(RegisterExecutorResult { success: true })) - } else { - warn!("Received invalid register executor request"); - Err(tonic::Status::invalid_argument( - "Missing metadata in request", - )) - } - } - - async fn heart_beat_from_executor( - &self, - request: Request, - ) -> Result, Status> { - let HeartBeatParams { executor_id, state } = request.into_inner(); - - debug!("Received heart beat request for {:?}", executor_id); - trace!("Related executor state is {:?}", state); - let executor_heartbeat = ExecutorHeartbeat { - executor_id, - timestamp: SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("Time went backwards") - .as_secs(), - state, - }; - self.state - .executor_manager - .save_executor_heartbeat(executor_heartbeat); - Ok(Response::new(HeartBeatResult { reregister: false })) - } - - async fn update_task_status( - &self, - request: Request, - ) -> Result, Status> { - let UpdateTaskStatusParams { - executor_id, - task_status, - } = request.into_inner(); - - debug!( - "Received task status update request for executor {:?}", - executor_id - ); - let num_tasks = task_status.len(); - if let Some(executor_data) = - self.state.executor_manager.get_executor_data(&executor_id) - { - self.state - .executor_manager - .update_executor_data(&ExecutorDataChange { - executor_id: executor_data.executor_id, - task_slots: num_tasks as i32, - }); - } else { - error!("Fail to get executor data for {:?}", &executor_id); - } - - self.update_task_status(task_status).await.map_err(|e| { - let msg = format!( - "Fail to update tasks status from executor {:?} due to {:?}", - &executor_id, e - ); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - - Ok(Response::new(UpdateTaskStatusResult { success: true })) - } - - async fn get_file_metadata( - &self, - request: Request, - ) -> std::result::Result, tonic::Status> { - // TODO support multiple object stores - let obj_store = LocalFileSystem {}; - // TODO shouldn't this take a ListingOption object as input? - - let GetFileMetadataParams { path, file_type } = request.into_inner(); - - let file_type: FileType = file_type.try_into().map_err(|e| { - let msg = format!("Error reading request: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - - let file_format: Arc = match file_type { - FileType::Parquet => Ok(Arc::new(ParquetFormat::default())), - // TODO implement for CSV - _ => Err(tonic::Status::unimplemented( - "get_file_metadata unsupported file type", - )), - }?; - - let file_metas = obj_store.list_file(&path).await.map_err(|e| { - let msg = format!("Error listing files: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - - let obj_readers = file_metas.map(move |f| obj_store.file_reader(f?.sized_file)); - - let schema = file_format - .infer_schema(Box::pin(obj_readers)) - .await - .map_err(|e| { - let msg = format!("Error infering schema: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - - Ok(Response::new(GetFileMetadataResult { - schema: Some(schema.as_ref().into()), - })) - } - - async fn execute_query( - &self, - request: Request, - ) -> std::result::Result, tonic::Status> { - let query_params = request.into_inner(); - if let ExecuteQueryParams { - query: Some(query), - settings, - optional_session_id, - } = query_params - { - // parse config - let mut config_builder = BallistaConfig::builder(); - for kv_pair in &settings { - config_builder = config_builder.set(&kv_pair.key, &kv_pair.value); - } - let config = config_builder.build().map_err(|e| { - let msg = format!("Could not parse configs: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - - let df_session = match optional_session_id { - Some(OptionalSessionId::SessionId(session_id)) => { - let session_ctx = self - .state - .session_registry() - .lookup_session(session_id.as_str()) - .await - .ok_or_else(|| { - Status::invalid_argument(format!( - "SessionContext not found for session ID {}", - session_id - )) - })?; - update_datafusion_context(session_ctx, &config) - } - _ => { - let df_session = - create_datafusion_context(&config, self.session_builder); - self.state - .session_registry() - .register_session(df_session.clone()) - .await; - df_session - } - }; - - let plan = match query { - Query::LogicalPlan(message) => T::try_decode(message.as_slice()) - .and_then(|m| { - m.try_into_logical_plan( - df_session.deref(), - self.codec.logical_extension_codec(), - ) - }) - .map_err(|e| { - let msg = format!("Could not parse logical plan protobuf: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?, - Query::Sql(sql) => df_session - .sql(&sql) - .await - .and_then(|df| df.to_logical_plan()) - .map_err(|e| { - let msg = format!("Error parsing SQL: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?, - }; - debug!("Received plan for execution: {:?}", plan); - - // Generate job id. - // TODO Maybe the format will be changed in the future - let job_id = generate_job_id(); - let session_id = df_session.session_id(); - let state = self.state.clone(); - let query_stage_event_sender = - self.query_stage_event_loop.get_sender().map_err(|e| { - tonic::Status::internal(format!( - "Could not get query stage event sender due to: {}", - e - )) - })?; - - // Save placeholder job metadata - state - .save_job_metadata( - &job_id, - &JobStatus { - status: Some(job_status::Status::Queued(QueuedJob {})), - }, - ) - .await - .map_err(|e| { - tonic::Status::internal(format!("Could not save job metadata: {}", e)) - })?; - - state - .save_job_session(&job_id, &session_id, settings) - .await - .map_err(|e| { - tonic::Status::internal(format!( - "Could not save job session mapping: {}", - e - )) - })?; - - let job_id_spawn = job_id.clone(); - let ctx = df_session.clone(); - tokio::spawn(async move { - if let Err(e) = async { - // create physical plan - let start = Instant::now(); - let plan = async { - let optimized_plan = ctx.optimize(&plan).map_err(|e| { - let msg = - format!("Could not create optimized logical plan: {}", e); - error!("{}", msg); - - BallistaError::General(msg) - })?; - - debug!("Calculated optimized plan: {:?}", optimized_plan); - - ctx.create_physical_plan(&optimized_plan) - .await - .map_err(|e| { - let msg = - format!("Could not create physical plan: {}", e); - error!("{}", msg); - - BallistaError::General(msg) - }) - } - .await?; - info!( - "DataFusion created physical plan in {} milliseconds", - start.elapsed().as_millis() - ); - - query_stage_event_sender - .post_event(QueryStageSchedulerEvent::JobSubmitted( - job_id_spawn.clone(), - plan, - )) - .await?; - - Ok::<(), BallistaError>(()) - } - .await - { - let msg = format!("Job {} failed due to {}", job_id_spawn, e); - warn!("{}", msg); - state - .save_job_metadata( - &job_id_spawn, - &JobStatus { - status: Some(job_status::Status::Failed(FailedJob { - error: msg.to_string(), - })), - }, - ) - .await - .unwrap_or_else(|_| { - panic!( - "Fail to update job status to failed for {}", - job_id_spawn - ) - }); - } - }); - - Ok(Response::new(ExecuteQueryResult { job_id, session_id })) - } else if let ExecuteQueryParams { - query: None, - settings, - optional_session_id: None, - } = query_params - { - // parse config for new session - let mut config_builder = BallistaConfig::builder(); - for kv_pair in &settings { - config_builder = config_builder.set(&kv_pair.key, &kv_pair.value); - } - let config = config_builder.build().map_err(|e| { - let msg = format!("Could not parse configs: {}", e); - error!("{}", msg); - tonic::Status::internal(msg) - })?; - let df_session = create_datafusion_context(&config, self.session_builder); - self.state - .session_registry() - .register_session(df_session.clone()) - .await; - Ok(Response::new(ExecuteQueryResult { - job_id: "NA".to_owned(), - session_id: df_session.session_id(), - })) - } else { - Err(tonic::Status::internal("Error parsing request")) - } - } - - async fn get_job_status( - &self, - request: Request, - ) -> std::result::Result, tonic::Status> { - let job_id = request.into_inner().job_id; - debug!("Received get_job_status request for job {}", job_id); - let job_meta = self.state.get_job_metadata(&job_id).unwrap(); - Ok(Response::new(GetJobStatusResult { - status: Some(job_meta), - })) - } -} - -fn generate_job_id() -> String { - let mut rng = thread_rng(); - std::iter::repeat(()) - .map(|()| rng.sample(Alphanumeric)) - .map(char::from) - .take(7) - .collect() -} - -#[cfg(all(test, feature = "sled"))] -mod test { - use std::sync::Arc; - - use tonic::Request; - - use crate::state::{backend::standalone::StandaloneClient, SchedulerState}; - use ballista_core::error::BallistaError; - use ballista_core::serde::protobuf::{ - executor_registration::OptionalHost, ExecutorRegistration, LogicalPlanNode, - PhysicalPlanNode, PollWorkParams, - }; - use ballista_core::serde::scheduler::ExecutorSpecification; - use ballista_core::serde::BallistaCodec; - use datafusion::execution::context::default_session_builder; - - use super::{SchedulerGrpc, SchedulerServer}; - - #[tokio::test] - async fn test_poll_work() -> Result<(), BallistaError> { - let state_storage = Arc::new(StandaloneClient::try_new_temporary()?); - let namespace = "default"; - let scheduler: SchedulerServer = - SchedulerServer::new( - state_storage.clone(), - namespace.to_owned(), - BallistaCodec::default(), - ); - let exec_meta = ExecutorRegistration { - id: "abc".to_owned(), - optional_host: Some(OptionalHost::Host("".to_owned())), - port: 0, - grpc_port: 0, - specification: Some(ExecutorSpecification { task_slots: 2 }.into()), - }; - let request: Request = Request::new(PollWorkParams { - metadata: Some(exec_meta.clone()), - can_accept_task: false, - task_status: vec![], - }); - let response = scheduler - .poll_work(request) - .await - .expect("Received error response") - .into_inner(); - // no response task since we told the scheduler we didn't want to accept one - assert!(response.task.is_none()); - let state: SchedulerState = - SchedulerState::new( - state_storage.clone(), - namespace.to_string(), - default_session_builder, - BallistaCodec::default(), - ); - state.init().await?; - // executor should be registered - assert_eq!(state.get_executors_metadata().await.unwrap().len(), 1); - - let request: Request = Request::new(PollWorkParams { - metadata: Some(exec_meta.clone()), - can_accept_task: true, - task_status: vec![], - }); - let response = scheduler - .poll_work(request) - .await - .expect("Received error response") - .into_inner(); - // still no response task since there are no tasks in the scheduelr - assert!(response.task.is_none()); - let state: SchedulerState = - SchedulerState::new( - state_storage.clone(), - namespace.to_string(), - default_session_builder, - BallistaCodec::default(), - ); - state.init().await?; - // executor should be registered - assert_eq!(state.get_executors_metadata().await.unwrap().len(), 1); - Ok(()) - } -} diff --git a/ballista/rust/scheduler/src/scheduler_server/mod.rs b/ballista/rust/scheduler/src/scheduler_server/mod.rs deleted file mode 100644 index d61f570cb4c2..000000000000 --- a/ballista/rust/scheduler/src/scheduler_server/mod.rs +++ /dev/null @@ -1,617 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::HashMap; -use std::sync::Arc; -use std::time::{SystemTime, UNIX_EPOCH}; - -use tokio::sync::RwLock; -use tonic::transport::Channel; - -use ballista_core::config::{BallistaConfig, TaskSchedulingPolicy}; -use ballista_core::error::Result; -use ballista_core::event_loop::EventLoop; -use ballista_core::serde::protobuf::executor_grpc_client::ExecutorGrpcClient; -use ballista_core::serde::protobuf::TaskStatus; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec}; -use datafusion::execution::context::{default_session_builder, SessionState}; -use datafusion::prelude::{SessionConfig, SessionContext}; - -use crate::scheduler_server::event::{QueryStageSchedulerEvent, SchedulerServerEvent}; -use crate::scheduler_server::event_loop::SchedulerServerEventAction; -use crate::scheduler_server::query_stage_scheduler::QueryStageScheduler; -use crate::state::backend::StateBackendClient; -use crate::state::SchedulerState; - -// include the generated protobuf source as a submodule -#[allow(clippy::all)] -pub mod externalscaler { - include!(concat!(env!("OUT_DIR"), "/externalscaler.rs")); -} - -pub mod event; -mod event_loop; -mod external_scaler; -mod grpc; -mod query_stage_scheduler; - -type ExecutorsClient = Arc>>>; -pub(crate) type SessionBuilder = fn(SessionConfig) -> SessionState; - -#[derive(Clone)] -pub struct SchedulerServer { - pub(crate) state: Arc>, - pub start_time: u128, - policy: TaskSchedulingPolicy, - executors_client: Option, - event_loop: Option>, - query_stage_event_loop: EventLoop, - codec: BallistaCodec, - /// SessionState Builder - session_builder: SessionBuilder, -} - -impl SchedulerServer { - pub fn new( - config: Arc, - namespace: String, - codec: BallistaCodec, - ) -> Self { - SchedulerServer::new_with_policy( - config, - namespace, - TaskSchedulingPolicy::PullStaged, - codec, - default_session_builder, - ) - } - - pub fn new_with_builder( - config: Arc, - namespace: String, - codec: BallistaCodec, - session_builder: SessionBuilder, - ) -> Self { - SchedulerServer::new_with_policy( - config, - namespace, - TaskSchedulingPolicy::PullStaged, - codec, - session_builder, - ) - } - - pub fn new_with_policy( - config: Arc, - namespace: String, - policy: TaskSchedulingPolicy, - codec: BallistaCodec, - session_builder: SessionBuilder, - ) -> Self { - let state = Arc::new(SchedulerState::new( - config, - namespace, - session_builder, - codec.clone(), - )); - - let (executors_client, event_loop) = - if matches!(policy, TaskSchedulingPolicy::PushStaged) { - let executors_client = Arc::new(RwLock::new(HashMap::new())); - let event_action: Arc> = - Arc::new(SchedulerServerEventAction::new( - state.clone(), - executors_client.clone(), - )); - let event_loop = - EventLoop::new("scheduler".to_owned(), 10000, event_action); - (Some(executors_client), Some(event_loop)) - } else { - (None, None) - }; - let query_stage_scheduler = - Arc::new(QueryStageScheduler::new(state.clone(), None)); - let query_stage_event_loop = - EventLoop::new("query_stage".to_owned(), 10000, query_stage_scheduler); - Self { - state, - start_time: SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap() - .as_millis(), - policy, - executors_client, - event_loop, - query_stage_event_loop, - codec, - session_builder, - } - } - - pub async fn init(&mut self) -> Result<()> { - { - // initialize state - self.state.init().await?; - } - - { - if let Some(event_loop) = self.event_loop.as_mut() { - event_loop.start()?; - - let query_stage_scheduler = Arc::new(QueryStageScheduler::new( - self.state.clone(), - Some(event_loop.get_sender()?), - )); - let query_stage_event_loop = EventLoop::new( - self.query_stage_event_loop.name.clone(), - self.query_stage_event_loop.buffer_size, - query_stage_scheduler, - ); - self.query_stage_event_loop = query_stage_event_loop; - } - - self.query_stage_event_loop.start()?; - } - - Ok(()) - } - - pub(crate) async fn update_task_status( - &self, - tasks_status: Vec, - ) -> Result<()> { - let num_tasks_status = tasks_status.len() as u32; - let stage_events = self.state.stage_manager.update_tasks_status(tasks_status); - if stage_events.is_empty() { - if let Some(event_loop) = self.event_loop.as_ref() { - event_loop - .get_sender()? - .post_event(SchedulerServerEvent::ReviveOffers(num_tasks_status)) - .await?; - } - } else { - for stage_event in stage_events { - self.post_stage_event(stage_event).await?; - } - } - - Ok(()) - } - - async fn post_stage_event(&self, event: QueryStageSchedulerEvent) -> Result<()> { - self.query_stage_event_loop - .get_sender()? - .post_event(event) - .await - } -} - -/// Create a DataFusion session context that is compatible with Ballista Configuration -pub fn create_datafusion_context( - config: &BallistaConfig, - session_builder: SessionBuilder, -) -> Arc { - let config = SessionConfig::new() - .with_target_partitions(config.default_shuffle_partitions()) - .with_batch_size(config.default_batch_size()) - .with_repartition_joins(config.repartition_joins()) - .with_repartition_aggregations(config.repartition_aggregations()) - .with_repartition_windows(config.repartition_windows()) - .with_parquet_pruning(config.parquet_pruning()); - let session_state = session_builder(config); - Arc::new(SessionContext::with_state(session_state)) -} - -/// Update the existing DataFusion session context with Ballista Configuration -pub fn update_datafusion_context( - session_ctx: Arc, - config: &BallistaConfig, -) -> Arc { - { - let mut mut_state = session_ctx.state.write(); - mut_state.config.target_partitions = config.default_shuffle_partitions(); - mut_state.config.batch_size = config.default_batch_size(); - mut_state.config.repartition_joins = config.repartition_joins(); - mut_state.config.repartition_aggregations = config.repartition_aggregations(); - mut_state.config.repartition_windows = config.repartition_windows(); - mut_state.config.parquet_pruning = config.parquet_pruning(); - } - session_ctx -} - -/// A Registry holds all the datafusion session contexts -pub struct SessionContextRegistry { - /// A map from session_id to SessionContext - pub running_sessions: RwLock>>, -} - -impl Default for SessionContextRegistry { - fn default() -> Self { - Self::new() - } -} - -impl SessionContextRegistry { - /// Create the registry that session contexts can registered into. - /// ['LocalFileSystem'] store is registered in by default to support read local files natively. - pub fn new() -> Self { - Self { - running_sessions: RwLock::new(HashMap::new()), - } - } - - /// Adds a new session to this registry. - pub async fn register_session( - &self, - session_ctx: Arc, - ) -> Option> { - let session_id = session_ctx.session_id(); - let mut sessions = self.running_sessions.write().await; - sessions.insert(session_id, session_ctx) - } - - /// Lookup the session context registered - pub async fn lookup_session(&self, session_id: &str) -> Option> { - let sessions = self.running_sessions.read().await; - sessions.get(session_id).cloned() - } - - /// Remove a session from this registry. - pub async fn unregister_session( - &self, - session_id: &str, - ) -> Option> { - let mut sessions = self.running_sessions.write().await; - sessions.remove(session_id) - } -} -#[cfg(all(test, feature = "sled"))] -mod test { - use std::sync::Arc; - use std::time::{Duration, Instant}; - - use ballista_core::config::TaskSchedulingPolicy; - use ballista_core::error::{BallistaError, Result}; - use ballista_core::execution_plans::ShuffleWriterExec; - use ballista_core::serde::protobuf::{ - job_status, task_status, CompletedTask, LogicalPlanNode, PartitionId, - PhysicalPlanNode, TaskStatus, - }; - use ballista_core::serde::scheduler::ExecutorData; - use ballista_core::serde::BallistaCodec; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use datafusion::execution::context::default_session_builder; - use datafusion::logical_plan::{col, sum, LogicalPlan}; - use datafusion::prelude::{SessionConfig, SessionContext}; - use datafusion::test_util::scan_empty; - - use crate::scheduler_server::event::QueryStageSchedulerEvent; - use crate::scheduler_server::SchedulerServer; - use crate::state::backend::standalone::StandaloneClient; - use crate::state::task_scheduler::TaskScheduler; - - #[tokio::test] - async fn test_pull_based_task_scheduling() -> Result<()> { - let now = Instant::now(); - test_task_scheduling(TaskSchedulingPolicy::PullStaged, test_plan(), 4).await?; - println!( - "pull-based task scheduling cost {}ms", - now.elapsed().as_millis() - ); - - Ok(()) - } - - #[tokio::test] - async fn test_push_based_task_scheduling() -> Result<()> { - let now = Instant::now(); - test_task_scheduling(TaskSchedulingPolicy::PushStaged, test_plan(), 4).await?; - println!( - "push-based task scheduling cost {}ms", - now.elapsed().as_millis() - ); - - Ok(()) - } - - async fn test_task_scheduling( - policy: TaskSchedulingPolicy, - plan_of_linear_stages: LogicalPlan, - total_available_task_slots: usize, - ) -> Result<()> { - let scheduler = test_scheduler(policy).await?; - if matches!(policy, TaskSchedulingPolicy::PushStaged) { - let executors = test_executors(total_available_task_slots); - for executor_data in executors { - scheduler - .state - .executor_manager - .save_executor_data(executor_data); - } - } - let config = - SessionConfig::new().with_target_partitions(total_available_task_slots); - let ctx = Arc::new(SessionContext::with_config(config)); - let plan = async { - let optimized_plan = ctx.optimize(&plan_of_linear_stages).map_err(|e| { - BallistaError::General(format!( - "Could not create optimized logical plan: {}", - e - )) - })?; - - ctx.create_physical_plan(&optimized_plan) - .await - .map_err(|e| { - BallistaError::General(format!( - "Could not create physical plan: {}", - e - )) - }) - } - .await?; - - let job_id = "job"; - scheduler - .state - .session_registry() - .register_session(ctx.clone()) - .await; - scheduler - .state - .save_job_session(job_id, ctx.session_id().as_str(), vec![]) - .await?; - { - // verify job submit - scheduler - .post_stage_event(QueryStageSchedulerEvent::JobSubmitted( - job_id.to_owned(), - plan, - )) - .await?; - - let waiting_time_ms = - test_waiting_async(|| scheduler.state.get_job_metadata(job_id).is_some()) - .await; - let job_status = scheduler.state.get_job_metadata(job_id); - assert!( - job_status.is_some(), - "Fail to receive JobSubmitted event within {}ms", - waiting_time_ms - ); - } - - let stage_task_num = test_get_job_stage_task_num(&scheduler, job_id); - let first_stage_id = 1u32; - let final_stage_id = stage_task_num.len() as u32 - 1; - assert!(scheduler - .state - .stage_manager - .is_final_stage(job_id, final_stage_id)); - - if matches!(policy, TaskSchedulingPolicy::PullStaged) { - assert!(!scheduler.state.stage_manager.has_running_tasks()); - assert!(scheduler - .state - .stage_manager - .is_running_stage(job_id, first_stage_id)); - if first_stage_id != final_stage_id { - assert!(scheduler - .state - .stage_manager - .is_pending_stage(job_id, final_stage_id)); - } - } - - // complete stage one by one - for stage_id in first_stage_id..final_stage_id { - let next_stage_id = stage_id + 1; - let num_tasks = stage_task_num[stage_id as usize] as usize; - if matches!(policy, TaskSchedulingPolicy::PullStaged) { - let mut executors = test_executors(total_available_task_slots); - let _fet_tasks = scheduler - .state - .fetch_schedulable_tasks(&mut executors, 1) - .await?; - } - assert!(scheduler.state.stage_manager.has_running_tasks()); - assert!(scheduler - .state - .stage_manager - .is_running_stage(job_id, stage_id)); - assert!(scheduler - .state - .stage_manager - .is_pending_stage(job_id, next_stage_id)); - - test_complete_stage(&scheduler, job_id, 1, num_tasks).await?; - assert!(!scheduler.state.stage_manager.has_running_tasks()); - assert!(!scheduler - .state - .stage_manager - .is_running_stage(job_id, stage_id)); - assert!(scheduler - .state - .stage_manager - .is_completed_stage(job_id, stage_id)); - let waiting_time_ms = test_waiting_async(|| { - !scheduler - .state - .stage_manager - .is_pending_stage(job_id, next_stage_id) - }) - .await; - assert!( - !scheduler - .state - .stage_manager - .is_pending_stage(job_id, next_stage_id), - "Fail to update stage state machine within {}ms", - waiting_time_ms - ); - assert!(scheduler - .state - .stage_manager - .is_running_stage(job_id, next_stage_id)); - } - - // complete the final stage - { - let num_tasks = stage_task_num[final_stage_id as usize] as usize; - if matches!(policy, TaskSchedulingPolicy::PullStaged) { - let mut executors = test_executors(total_available_task_slots); - let _fet_tasks = scheduler - .state - .fetch_schedulable_tasks(&mut executors, 1) - .await?; - } - assert!(scheduler.state.stage_manager.has_running_tasks()); - - test_complete_stage(&scheduler, job_id, final_stage_id, num_tasks).await?; - assert!(!scheduler.state.stage_manager.has_running_tasks()); - assert!(!scheduler - .state - .stage_manager - .is_running_stage(job_id, final_stage_id)); - assert!(scheduler - .state - .stage_manager - .is_completed_stage(job_id, final_stage_id)); - let waiting_time_ms = test_waiting_async(|| { - let job_status = scheduler.state.get_job_metadata(job_id).unwrap(); - matches!(job_status.status, Some(job_status::Status::Completed(_))) - }) - .await; - - let job_status = scheduler.state.get_job_metadata(job_id).unwrap(); - assert!( - matches!(job_status.status, Some(job_status::Status::Completed(_))), - "Fail to update job state machine within {}ms", - waiting_time_ms - ); - } - - Ok(()) - } - - async fn test_waiting_async(cond: F) -> u64 - where - F: Fn() -> bool, - { - let round_waiting_time = 10; - let num_round = 5; - for _i in 0..num_round { - if cond() { - break; - } - tokio::time::sleep(Duration::from_millis(round_waiting_time)).await; - } - - round_waiting_time * num_round - } - - async fn test_complete_stage( - scheduler: &SchedulerServer, - job_id: &str, - stage_id: u32, - num_tasks: usize, - ) -> Result<()> { - let tasks_status: Vec = (0..num_tasks as u32) - .into_iter() - .map(|task_id| TaskStatus { - status: Some(task_status::Status::Completed(CompletedTask { - executor_id: "localhost".to_owned(), - partitions: Vec::new(), - })), - task_id: Some(PartitionId { - job_id: job_id.to_owned(), - stage_id, - partition_id: task_id, - }), - }) - .collect(); - scheduler.update_task_status(tasks_status).await - } - - async fn test_scheduler( - policy: TaskSchedulingPolicy, - ) -> Result> { - let state_storage = Arc::new(StandaloneClient::try_new_temporary()?); - let mut scheduler: SchedulerServer = - SchedulerServer::new_with_policy( - state_storage.clone(), - "default".to_owned(), - policy, - BallistaCodec::default(), - default_session_builder, - ); - scheduler.init().await?; - - Ok(scheduler) - } - - fn test_executors(num_partitions: usize) -> Vec { - let task_slots = (num_partitions as u32 + 1) / 2; - - vec![ - ExecutorData { - executor_id: "localhost1".to_owned(), - total_task_slots: task_slots, - available_task_slots: task_slots, - }, - ExecutorData { - executor_id: "localhost2".to_owned(), - total_task_slots: num_partitions as u32 - task_slots, - available_task_slots: num_partitions as u32 - task_slots, - }, - ] - } - - fn test_get_job_stage_task_num( - scheduler: &SchedulerServer, - job_id: &str, - ) -> Vec { - let mut ret = vec![0, 1]; - let mut stage_id = 1; - while let Some(stage_plan) = scheduler.state.get_stage_plan(job_id, stage_id) { - if let Some(shuffle_writer) = - stage_plan.as_any().downcast_ref::() - { - if let Some(partitions) = shuffle_writer.shuffle_output_partitioning() { - ret.push(partitions.partition_count() as u32) - } - } - stage_id += 1; - } - - ret - } - - fn test_plan() -> LogicalPlan { - let schema = Schema::new(vec![ - Field::new("id", DataType::Utf8, false), - Field::new("gmv", DataType::UInt64, false), - ]); - - scan_empty(None, &schema, Some(vec![0, 1])) - .unwrap() - .aggregate(vec![col("id")], vec![sum(col("gmv"))]) - .unwrap() - .build() - .unwrap() - } -} diff --git a/ballista/rust/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/rust/scheduler/src/scheduler_server/query_stage_scheduler.rs deleted file mode 100644 index 58e9f8f3b814..000000000000 --- a/ballista/rust/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ /dev/null @@ -1,476 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::{HashMap, HashSet}; -use std::sync::Arc; - -use async_recursion::async_recursion; -use async_trait::async_trait; -use log::{debug, error, info, warn}; - -use ballista_core::error::{BallistaError, Result}; - -use ballista_core::event_loop::{EventAction, EventSender}; -use ballista_core::execution_plans::UnresolvedShuffleExec; -use ballista_core::serde::protobuf::{ - job_status, task_status, CompletedJob, CompletedTask, FailedJob, FailedTask, - JobStatus, RunningJob, TaskStatus, -}; -use ballista_core::serde::scheduler::{ExecutorMetadata, PartitionStats}; -use ballista_core::serde::{protobuf, AsExecutionPlan, AsLogicalPlan}; -use datafusion::physical_plan::ExecutionPlan; - -use crate::planner::{ - find_unresolved_shuffles, remove_unresolved_shuffles, DistributedPlanner, -}; -use crate::scheduler_server::event::{QueryStageSchedulerEvent, SchedulerServerEvent}; -use crate::state::SchedulerState; - -pub(crate) struct QueryStageScheduler< - T: 'static + AsLogicalPlan, - U: 'static + AsExecutionPlan, -> { - state: Arc>, - event_sender: Option>, -} - -impl QueryStageScheduler { - pub(crate) fn new( - state: Arc>, - event_sender: Option>, - ) -> Self { - Self { - state, - event_sender, - } - } - - async fn generate_stages( - &self, - job_id: &str, - plan: Arc, - ) -> Result<()> { - let mut planner = DistributedPlanner::new(); - // The last one is the final stage - let stages = planner.plan_query_stages(job_id, plan).await.map_err(|e| { - let msg = format!("Could not plan query stages: {}", e); - error!("{}", msg); - BallistaError::General(msg) - })?; - - let mut stages_dependency: HashMap> = HashMap::new(); - // save stages into state - for shuffle_writer in stages.iter() { - let stage_id = shuffle_writer.stage_id(); - let stage_plan: Arc = shuffle_writer.clone(); - self.state - .save_stage_plan(job_id, stage_id, stage_plan.clone()) - .await - .map_err(|e| { - let msg = format!("Could not save stage plan: {}", e); - error!("{}", msg); - BallistaError::General(msg) - })?; - - for child in find_unresolved_shuffles(&stage_plan)? { - stages_dependency - .entry(child.stage_id as u32) - .or_insert_with(HashSet::new) - .insert(stage_id as u32); - } - } - - self.state - .stage_manager - .add_stages_dependency(job_id, stages_dependency); - - let final_stage_id = stages.last().as_ref().unwrap().stage_id(); - self.state - .stage_manager - .add_final_stage(job_id, final_stage_id as u32); - self.submit_stage(job_id, final_stage_id).await?; - - Ok(()) - } - - async fn submit_pending_stages(&self, job_id: &str, stage_id: usize) -> Result<()> { - if let Some(parent_stages) = self - .state - .stage_manager - .get_parent_stages(job_id, stage_id as u32) - { - self.state - .stage_manager - .remove_pending_stage(job_id, &parent_stages); - for parent_stage in parent_stages { - self.submit_stage(job_id, parent_stage as usize).await?; - } - } - - Ok(()) - } - - #[async_recursion] - async fn submit_stage(&self, job_id: &str, stage_id: usize) -> Result<()> { - { - if self - .state - .stage_manager - .is_running_stage(job_id, stage_id as u32) - { - debug!("stage {}/{} has already been submitted", job_id, stage_id); - return Ok(()); - } - if self - .state - .stage_manager - .is_pending_stage(job_id, stage_id as u32) - { - debug!( - "stage {}/{} has already been added to the pending list", - job_id, stage_id - ); - return Ok(()); - } - } - if let Some(stage_plan) = self.state.get_stage_plan(job_id, stage_id) { - if let Some(incomplete_unresolved_shuffles) = self - .try_resolve_stage(job_id, stage_id, stage_plan.clone()) - .await? - { - assert!( - !incomplete_unresolved_shuffles.is_empty(), - "there are no incomplete unresolved shuffles" - ); - for incomplete_unresolved_shuffle in incomplete_unresolved_shuffles { - self.submit_stage(job_id, incomplete_unresolved_shuffle.stage_id) - .await?; - } - self.state - .stage_manager - .add_pending_stage(job_id, stage_id as u32); - } else { - self.state.stage_manager.add_running_stage( - job_id, - stage_id as u32, - stage_plan.output_partitioning().partition_count() as u32, - ); - } - } else { - return Err(BallistaError::General(format!( - "Fail to find stage plan for {}/{}", - job_id, stage_id - ))); - } - Ok(()) - } - - /// Try to resolve a stage if all of the unresolved shuffles are completed. - /// Return the unresolved shuffles which are incomplete - async fn try_resolve_stage( - &self, - job_id: &str, - stage_id: usize, - stage_plan: Arc, - ) -> Result>> { - // Find all of the unresolved shuffles - let unresolved_shuffles = find_unresolved_shuffles(&stage_plan)?; - - // If no dependent shuffles - if unresolved_shuffles.is_empty() { - return Ok(None); - } - - // Find all of the incomplete unresolved shuffles - let (incomplete_unresolved_shuffles, unresolved_shuffles): ( - Vec, - Vec, - ) = unresolved_shuffles.into_iter().partition(|s| { - !self - .state - .stage_manager - .is_completed_stage(job_id, s.stage_id as u32) - }); - - if !incomplete_unresolved_shuffles.is_empty() { - return Ok(Some(incomplete_unresolved_shuffles)); - } - - // All of the unresolved shuffles are completed, update the stage plan - { - let mut partition_locations: HashMap< - usize, // input stage id - HashMap< - usize, // task id of this stage - Vec, // shuffle partitions - >, - > = HashMap::new(); - for unresolved_shuffle in unresolved_shuffles.iter() { - let input_stage_id = unresolved_shuffle.stage_id; - let stage_shuffle_partition_locations = partition_locations - .entry(input_stage_id) - .or_insert_with(HashMap::new); - if let Some(input_stage_tasks) = self - .state - .stage_manager - .get_stage_tasks(job_id, input_stage_id as u32) - { - // each input partition can produce multiple output partitions - for (shuffle_input_partition_id, task_status) in - input_stage_tasks.iter().enumerate() - { - match &task_status.status { - Some(task_status::Status::Completed(CompletedTask { - executor_id, - partitions, - })) => { - debug!( - "Task for unresolved shuffle input partition {} completed and produced these shuffle partitions:\n\t{}", - shuffle_input_partition_id, - partitions.iter().map(|p| format!("{}={}", p.partition_id, &p.path)).collect::>().join("\n\t") - ); - - for shuffle_write_partition in partitions { - let temp = stage_shuffle_partition_locations - .entry( - shuffle_write_partition.partition_id as usize, - ) - .or_insert(Vec::new()); - let executor_meta = self - .state - .get_executor_metadata(executor_id) - .ok_or_else(|| { - BallistaError::General(format!( - "Fail to find executor metadata for {}", - &executor_id - )) - })?; - let partition_location = - ballista_core::serde::scheduler::PartitionLocation { - partition_id: - ballista_core::serde::scheduler::PartitionId { - job_id: job_id.to_owned(), - stage_id: unresolved_shuffle.stage_id, - partition_id: shuffle_write_partition - .partition_id - as usize, - }, - executor_meta, - partition_stats: PartitionStats::new( - Some(shuffle_write_partition.num_rows), - Some(shuffle_write_partition.num_batches), - Some(shuffle_write_partition.num_bytes), - ), - path: shuffle_write_partition.path.clone(), - }; - debug!( - "Scheduler storing stage {} output partition {} path: {}", - unresolved_shuffle.stage_id, - partition_location.partition_id.partition_id, - partition_location.path - ); - temp.push(partition_location); - } - } - _ => { - debug!( - "Stage {} input partition {} has not completed yet", - unresolved_shuffle.stage_id, - shuffle_input_partition_id - ); - // TODO task error handling - } - } - } - } else { - return Err(BallistaError::General(format!( - "Fail to find completed stage for {}/{}", - job_id, stage_id - ))); - } - } - - let plan = remove_unresolved_shuffles(stage_plan, &partition_locations)?; - self.state.save_stage_plan(job_id, stage_id, plan).await?; - } - - Ok(None) - } -} - -#[async_trait] -impl - EventAction for QueryStageScheduler -{ - // TODO - fn on_start(&self) {} - - // TODO - fn on_stop(&self) {} - - async fn on_receive( - &self, - event: QueryStageSchedulerEvent, - ) -> Result> { - match event { - QueryStageSchedulerEvent::JobSubmitted(job_id, plan) => { - info!("Job {} submitted", job_id); - match self.generate_stages(&job_id, plan).await { - Err(e) => { - let msg = format!("Job {} failed due to {}", job_id, e); - warn!("{}", msg); - self.state - .save_job_metadata( - &job_id, - &JobStatus { - status: Some(job_status::Status::Failed(FailedJob { - error: msg.to_string(), - })), - }, - ) - .await?; - } - Ok(()) => { - if let Err(e) = self - .state - .save_job_metadata( - &job_id, - &JobStatus { - status: Some(job_status::Status::Running( - RunningJob {}, - )), - }, - ) - .await - { - warn!( - "Could not update job {} status to running: {}", - job_id, e - ); - } - } - } - } - QueryStageSchedulerEvent::StageFinished(job_id, stage_id) => { - info!("Job stage {}/{} finished", job_id, stage_id); - self.submit_pending_stages(&job_id, stage_id as usize) - .await?; - } - QueryStageSchedulerEvent::JobFinished(job_id) => { - info!("Job {} finished", job_id); - let tasks_for_complete_final_stage = self - .state - .stage_manager - .get_tasks_for_complete_final_stage(&job_id)?; - let executors: HashMap = self - .state - .get_executors_metadata() - .await? - .into_iter() - .map(|(meta, _)| (meta.id.to_string(), meta)) - .collect(); - let job_status = get_job_status_from_tasks( - &tasks_for_complete_final_stage, - &executors, - ); - self.state.save_job_metadata(&job_id, &job_status).await?; - } - QueryStageSchedulerEvent::JobFailed(job_id, stage_id, fail_message) => { - error!( - "Job stage {}/{} failed due to {}", - &job_id, stage_id, fail_message - ); - let job_status = JobStatus { - status: Some(job_status::Status::Failed(FailedJob { - error: fail_message, - })), - }; - self.state.save_job_metadata(&job_id, &job_status).await?; - } - } - - if let Some(event_sender) = self.event_sender.as_ref() { - // The stage event must triggerred with releasing some resources. Therefore, revive offers for the scheduler - event_sender - .post_event(SchedulerServerEvent::ReviveOffers(1)) - .await?; - }; - Ok(None) - } - - // TODO - fn on_error(&self, _error: BallistaError) {} -} - -fn get_job_status_from_tasks( - tasks: &[Arc], - executors: &HashMap, -) -> JobStatus { - let mut job_status = tasks - .iter() - .map(|task| match &task.status { - Some(task_status::Status::Completed(CompletedTask { - executor_id, - partitions, - })) => Ok((task, executor_id, partitions)), - _ => Err(BallistaError::General("Task not completed".to_string())), - }) - .collect::>>() - .ok() - .map(|info| { - let mut partition_location = vec![]; - for (status, executor_id, partitions) in info { - let input_partition_id = status.task_id.as_ref().unwrap(); // TODO unwrap - let executor_meta = executors.get(executor_id).map(|e| e.clone().into()); - for shuffle_write_partition in partitions { - let shuffle_input_partition_id = Some(protobuf::PartitionId { - job_id: input_partition_id.job_id.clone(), - stage_id: input_partition_id.stage_id, - partition_id: input_partition_id.partition_id, - }); - partition_location.push(protobuf::PartitionLocation { - partition_id: shuffle_input_partition_id.clone(), - executor_meta: executor_meta.clone(), - partition_stats: Some(protobuf::PartitionStats { - num_batches: shuffle_write_partition.num_batches as i64, - num_rows: shuffle_write_partition.num_rows as i64, - num_bytes: shuffle_write_partition.num_bytes as i64, - column_stats: vec![], - }), - path: shuffle_write_partition.path.clone(), - }); - } - } - job_status::Status::Completed(CompletedJob { partition_location }) - }); - - if job_status.is_none() { - // Update other statuses - for task in tasks.iter() { - if let Some(task_status::Status::Failed(FailedTask { error })) = &task.status - { - let error = error.clone(); - job_status = Some(job_status::Status::Failed(FailedJob { error })); - break; - } - } - } - - JobStatus { - status: Some(job_status.unwrap_or(job_status::Status::Running(RunningJob {}))), - } -} diff --git a/ballista/rust/scheduler/src/standalone.rs b/ballista/rust/scheduler/src/standalone.rs deleted file mode 100644 index f333587e495d..000000000000 --- a/ballista/rust/scheduler/src/standalone.rs +++ /dev/null @@ -1,58 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use ballista_core::serde::protobuf::{LogicalPlanNode, PhysicalPlanNode}; -use ballista_core::serde::BallistaCodec; -use ballista_core::{ - error::Result, serde::protobuf::scheduler_grpc_server::SchedulerGrpcServer, - BALLISTA_VERSION, -}; -use log::info; -use std::{net::SocketAddr, sync::Arc}; -use tokio::net::TcpListener; -use tonic::transport::Server; - -use crate::{ - scheduler_server::SchedulerServer, state::backend::standalone::StandaloneClient, -}; - -pub async fn new_standalone_scheduler() -> Result { - let client = StandaloneClient::try_new_temporary()?; - - let mut scheduler_server: SchedulerServer = - SchedulerServer::new( - Arc::new(client), - "ballista".to_string(), - BallistaCodec::default(), - ); - scheduler_server.init().await?; - let server = SchedulerGrpcServer::new(scheduler_server.clone()); - // Let the OS assign a random, free port - let listener = TcpListener::bind("localhost:0").await?; - let addr = listener.local_addr()?; - info!( - "Ballista v{} Rust Scheduler listening on {:?}", - BALLISTA_VERSION, addr - ); - tokio::spawn( - Server::builder().add_service(server).serve_with_incoming( - tokio_stream::wrappers::TcpListenerStream::new(listener), - ), - ); - - Ok(addr) -} diff --git a/ballista/rust/scheduler/src/state/backend/etcd.rs b/ballista/rust/scheduler/src/state/backend/etcd.rs deleted file mode 100644 index fa85e54d5f85..000000000000 --- a/ballista/rust/scheduler/src/state/backend/etcd.rs +++ /dev/null @@ -1,187 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Etcd config backend. - -use std::task::Poll; - -use ballista_core::error::{ballista_error, Result}; - -use etcd_client::{GetOptions, LockResponse, WatchOptions, WatchStream, Watcher}; -use futures::{Stream, StreamExt}; -use log::warn; - -use crate::state::backend::{Lock, StateBackendClient, Watch, WatchEvent}; - -/// A [`StateBackendClient`] implementation that uses etcd to save cluster configuration. -#[derive(Clone)] -pub struct EtcdClient { - etcd: etcd_client::Client, -} - -impl EtcdClient { - pub fn new(etcd: etcd_client::Client) -> Self { - Self { etcd } - } -} - -#[tonic::async_trait] -impl StateBackendClient for EtcdClient { - async fn get(&self, key: &str) -> Result> { - Ok(self - .etcd - .clone() - .get(key, None) - .await - .map_err(|e| ballista_error(&format!("etcd error {:?}", e)))? - .kvs() - .get(0) - .map(|kv| kv.value().to_owned()) - .unwrap_or_default()) - } - - async fn get_from_prefix(&self, prefix: &str) -> Result)>> { - Ok(self - .etcd - .clone() - .get(prefix, Some(GetOptions::new().with_prefix())) - .await - .map_err(|e| ballista_error(&format!("etcd error {:?}", e)))? - .kvs() - .iter() - .map(|kv| (kv.key_str().unwrap().to_owned(), kv.value().to_owned())) - .collect()) - } - - async fn put(&self, key: String, value: Vec) -> Result<()> { - let mut etcd = self.etcd.clone(); - etcd.put(key.clone(), value.clone(), None) - .await - .map_err(|e| { - warn!("etcd put failed: {}", e); - ballista_error("etcd put failed") - }) - .map(|_| ()) - } - - async fn lock(&self) -> Result> { - let mut etcd = self.etcd.clone(); - // TODO: make this a namespaced-lock - let lock = etcd - .lock("/ballista_global_lock", None) - .await - .map_err(|e| { - warn!("etcd lock failed: {}", e); - ballista_error("etcd lock failed") - })?; - Ok(Box::new(EtcdLockGuard { etcd, lock })) - } - - async fn watch(&self, prefix: String) -> Result> { - let mut etcd = self.etcd.clone(); - let options = WatchOptions::new().with_prefix(); - let (watcher, stream) = etcd.watch(prefix, Some(options)).await.map_err(|e| { - warn!("etcd watch failed: {}", e); - ballista_error("etcd watch failed") - })?; - Ok(Box::new(EtcdWatch { - watcher, - stream, - buffered_events: Vec::new(), - })) - } -} - -struct EtcdWatch { - watcher: Watcher, - stream: WatchStream, - buffered_events: Vec, -} - -#[tonic::async_trait] -impl Watch for EtcdWatch { - async fn cancel(&mut self) -> Result<()> { - self.watcher.cancel().await.map_err(|e| { - warn!("etcd watch cancel failed: {}", e); - ballista_error("etcd watch cancel failed") - }) - } -} - -impl Stream for EtcdWatch { - type Item = WatchEvent; - - fn poll_next( - self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> Poll> { - let self_mut = self.get_mut(); - if let Some(event) = self_mut.buffered_events.pop() { - Poll::Ready(Some(event)) - } else { - loop { - match self_mut.stream.poll_next_unpin(cx) { - Poll::Ready(Some(Err(e))) => { - warn!("Error when watching etcd prefix: {}", e); - continue; - } - Poll::Ready(Some(Ok(v))) => { - self_mut.buffered_events.extend(v.events().iter().map(|ev| { - match ev.event_type() { - etcd_client::EventType::Put => { - let kv = ev.kv().unwrap(); - WatchEvent::Put( - kv.key_str().unwrap().to_string(), - kv.value().to_owned(), - ) - } - etcd_client::EventType::Delete => { - let kv = ev.kv().unwrap(); - WatchEvent::Delete(kv.key_str().unwrap().to_string()) - } - } - })); - if let Some(event) = self_mut.buffered_events.pop() { - return Poll::Ready(Some(event)); - } else { - continue; - } - } - Poll::Ready(None) => return Poll::Ready(None), - Poll::Pending => return Poll::Pending, - } - } - } - } - - fn size_hint(&self) -> (usize, Option) { - self.stream.size_hint() - } -} - -struct EtcdLockGuard { - etcd: etcd_client::Client, - lock: LockResponse, -} - -// Cannot use Drop because we need this to be async -#[tonic::async_trait] -impl Lock for EtcdLockGuard { - async fn unlock(&mut self) { - self.etcd.unlock(self.lock.key()).await.unwrap(); - } -} diff --git a/ballista/rust/scheduler/src/state/backend/mod.rs b/ballista/rust/scheduler/src/state/backend/mod.rs deleted file mode 100644 index 15f244b693a6..000000000000 --- a/ballista/rust/scheduler/src/state/backend/mod.rs +++ /dev/null @@ -1,94 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use ballista_core::error::Result; -use clap::ArgEnum; -use futures::Stream; -use std::fmt; -use tokio::sync::OwnedMutexGuard; - -#[cfg(feature = "etcd")] -pub mod etcd; -#[cfg(feature = "sled")] -pub mod standalone; - -// an enum used to configure the backend -// needs to be visible to code generated by configure_me -#[derive(Debug, Clone, ArgEnum, serde::Deserialize)] -pub enum StateBackend { - Etcd, - Standalone, -} - -impl std::str::FromStr for StateBackend { - type Err = String; - - fn from_str(s: &str) -> std::result::Result { - ArgEnum::from_str(s, true) - } -} - -impl parse_arg::ParseArgFromStr for StateBackend { - fn describe_type(mut writer: W) -> fmt::Result { - write!(writer, "The configuration backend for the scheduler") - } -} - -/// A trait that contains the necessary methods to save and retrieve the state and configuration of a cluster. -#[tonic::async_trait] -pub trait StateBackendClient: Send + Sync { - /// Retrieve the data associated with a specific key. - /// - /// An empty vec is returned if the key does not exist. - async fn get(&self, key: &str) -> Result>; - - /// Retrieve all data associated with a specific key. - async fn get_from_prefix(&self, prefix: &str) -> Result)>>; - - /// Saves the value into the provided key, overriding any previous data that might have been associated to that key. - async fn put(&self, key: String, value: Vec) -> Result<()>; - - async fn lock(&self) -> Result>; - - /// Watch all events that happen on a specific prefix. - async fn watch(&self, prefix: String) -> Result>; -} - -/// A Watch is a cancelable stream of put or delete events in the [StateBackendClient] -#[tonic::async_trait] -pub trait Watch: Stream + Send + Unpin { - async fn cancel(&mut self) -> Result<()>; -} - -#[derive(Debug, PartialEq)] -pub enum WatchEvent { - /// Contains the inserted or updated key and the new value - Put(String, Vec), - - /// Contains the deleted key - Delete(String), -} - -#[tonic::async_trait] -pub trait Lock: Send + Sync { - async fn unlock(&mut self); -} - -#[tonic::async_trait] -impl Lock for OwnedMutexGuard { - async fn unlock(&mut self) {} -} diff --git a/ballista/rust/scheduler/src/state/backend/standalone.rs b/ballista/rust/scheduler/src/state/backend/standalone.rs deleted file mode 100644 index 5bb4e384132f..000000000000 --- a/ballista/rust/scheduler/src/state/backend/standalone.rs +++ /dev/null @@ -1,216 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::{sync::Arc, task::Poll}; - -use ballista_core::error::{ballista_error, BallistaError, Result}; - -use futures::{FutureExt, Stream}; -use log::warn; -use sled_package as sled; -use tokio::sync::Mutex; - -use crate::state::backend::{Lock, StateBackendClient, Watch, WatchEvent}; - -/// A [`StateBackendClient`] implementation that uses file-based storage to save cluster configuration. -#[derive(Clone)] -pub struct StandaloneClient { - db: sled::Db, - lock: Arc>, -} - -impl StandaloneClient { - /// Creates a StandaloneClient that saves data to the specified file. - pub fn try_new>(path: P) -> Result { - Ok(Self { - db: sled::open(path).map_err(sled_to_ballista_error)?, - lock: Arc::new(Mutex::new(())), - }) - } - - /// Creates a StandaloneClient that saves data to a temp file. - pub fn try_new_temporary() -> Result { - Ok(Self { - db: sled::Config::new() - .temporary(true) - .open() - .map_err(sled_to_ballista_error)?, - lock: Arc::new(Mutex::new(())), - }) - } -} - -fn sled_to_ballista_error(e: sled::Error) -> BallistaError { - match e { - sled::Error::Io(io) => BallistaError::IoError(io), - _ => BallistaError::General(format!("{}", e)), - } -} - -#[tonic::async_trait] -impl StateBackendClient for StandaloneClient { - async fn get(&self, key: &str) -> Result> { - Ok(self - .db - .get(key) - .map_err(|e| ballista_error(&format!("sled error {:?}", e)))? - .map(|v| v.to_vec()) - .unwrap_or_default()) - } - - async fn get_from_prefix(&self, prefix: &str) -> Result)>> { - Ok(self - .db - .scan_prefix(prefix) - .map(|v| { - v.map(|(key, value)| { - ( - std::str::from_utf8(&key).unwrap().to_owned(), - value.to_vec(), - ) - }) - }) - .collect::, _>>() - .map_err(|e| ballista_error(&format!("sled error {:?}", e)))?) - } - - async fn put(&self, key: String, value: Vec) -> Result<()> { - self.db - .insert(key, value) - .map_err(|e| { - warn!("sled insert failed: {}", e); - ballista_error("sled insert failed") - }) - .map(|_| ()) - } - - async fn lock(&self) -> Result> { - Ok(Box::new(self.lock.clone().lock_owned().await)) - } - - async fn watch(&self, prefix: String) -> Result> { - Ok(Box::new(SledWatch { - subscriber: self.db.watch_prefix(prefix), - })) - } -} - -struct SledWatch { - subscriber: sled::Subscriber, -} - -#[tonic::async_trait] -impl Watch for SledWatch { - async fn cancel(&mut self) -> Result<()> { - Ok(()) - } -} - -impl Stream for SledWatch { - type Item = WatchEvent; - - fn poll_next( - self: std::pin::Pin<&mut Self>, - cx: &mut std::task::Context<'_>, - ) -> std::task::Poll> { - match self.get_mut().subscriber.poll_unpin(cx) { - Poll::Pending => Poll::Pending, - Poll::Ready(None) => Poll::Ready(None), - Poll::Ready(Some(sled::Event::Insert { key, value })) => { - let key = std::str::from_utf8(&key).unwrap().to_owned(); - Poll::Ready(Some(WatchEvent::Put(key, value.to_vec()))) - } - Poll::Ready(Some(sled::Event::Remove { key })) => { - let key = std::str::from_utf8(&key).unwrap().to_owned(); - Poll::Ready(Some(WatchEvent::Delete(key))) - } - } - } - - fn size_hint(&self) -> (usize, Option) { - self.subscriber.size_hint() - } -} - -#[cfg(test)] -mod tests { - use super::{StandaloneClient, StateBackendClient, Watch, WatchEvent}; - - use futures::StreamExt; - use std::result::Result; - - fn create_instance() -> Result> { - Ok(StandaloneClient::try_new_temporary()?) - } - - #[tokio::test] - async fn put_read() -> Result<(), Box> { - let client = create_instance()?; - let key = "key"; - let value = "value".as_bytes(); - client.put(key.to_owned(), value.to_vec()).await?; - assert_eq!(client.get(key).await?, value); - Ok(()) - } - - #[tokio::test] - async fn read_empty() -> Result<(), Box> { - let client = create_instance()?; - let key = "key"; - let empty: &[u8] = &[]; - assert_eq!(client.get(key).await?, empty); - Ok(()) - } - - #[tokio::test] - async fn read_prefix() -> Result<(), Box> { - let client = create_instance()?; - let key = "key"; - let value = "value".as_bytes(); - client.put(format!("{}/1", key), value.to_vec()).await?; - client.put(format!("{}/2", key), value.to_vec()).await?; - assert_eq!( - client.get_from_prefix(key).await?, - vec![ - ("key/1".to_owned(), value.to_vec()), - ("key/2".to_owned(), value.to_vec()) - ] - ); - Ok(()) - } - - #[tokio::test] - async fn read_watch() -> Result<(), Box> { - let client = create_instance()?; - let key = "key"; - let value = "value".as_bytes(); - let mut watch: Box = client.watch(key.to_owned()).await?; - client.put(key.to_owned(), value.to_vec()).await?; - assert_eq!( - watch.next().await, - Some(WatchEvent::Put(key.to_owned(), value.to_owned())) - ); - let value2 = "value2".as_bytes(); - client.put(key.to_owned(), value2.to_vec()).await?; - assert_eq!( - watch.next().await, - Some(WatchEvent::Put(key.to_owned(), value2.to_owned())) - ); - watch.cancel().await?; - Ok(()) - } -} diff --git a/ballista/rust/scheduler/src/state/executor_manager.rs b/ballista/rust/scheduler/src/state/executor_manager.rs deleted file mode 100644 index 40821beabb51..000000000000 --- a/ballista/rust/scheduler/src/state/executor_manager.rs +++ /dev/null @@ -1,145 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -use ballista_core::serde::protobuf::ExecutorHeartbeat; -use ballista_core::serde::scheduler::{ExecutorData, ExecutorDataChange}; -use log::{error, info, warn}; -use parking_lot::RwLock; -use std::collections::{HashMap, HashSet}; -use std::sync::Arc; - -#[derive(Clone)] -pub(crate) struct ExecutorManager { - executors_heartbeat: Arc>>, - executors_data: Arc>>, -} - -impl ExecutorManager { - pub(crate) fn new() -> Self { - Self { - executors_heartbeat: Arc::new(RwLock::new(HashMap::new())), - executors_data: Arc::new(RwLock::new(HashMap::new())), - } - } - - pub(crate) fn save_executor_heartbeat(&self, heartbeat: ExecutorHeartbeat) { - let mut executors_heartbeat = self.executors_heartbeat.write(); - executors_heartbeat.insert(heartbeat.executor_id.clone(), heartbeat); - } - - pub(crate) fn get_executors_heartbeat(&self) -> Vec { - let executors_heartbeat = self.executors_heartbeat.read(); - executors_heartbeat - .iter() - .map(|(_exec, heartbeat)| heartbeat.clone()) - .collect() - } - - /// last_seen_ts_threshold is in seconds - pub(crate) fn get_alive_executors( - &self, - last_seen_ts_threshold: u64, - ) -> HashSet { - let executors_heartbeat = self.executors_heartbeat.read(); - executors_heartbeat - .iter() - .filter_map(|(exec, heartbeat)| { - (heartbeat.timestamp > last_seen_ts_threshold).then(|| exec.clone()) - }) - .collect() - } - - #[allow(dead_code)] - fn get_alive_executors_within_one_minute(&self) -> HashSet { - let now_epoch_ts = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("Time went backwards"); - let last_seen_threshold = now_epoch_ts - .checked_sub(Duration::from_secs(60)) - .unwrap_or_else(|| Duration::from_secs(0)); - self.get_alive_executors(last_seen_threshold.as_secs()) - } - - pub(crate) fn save_executor_data(&self, executor_data: ExecutorData) { - let mut executors_data = self.executors_data.write(); - executors_data.insert(executor_data.executor_id.clone(), executor_data); - } - - pub(crate) fn update_executor_data(&self, executor_data_change: &ExecutorDataChange) { - let mut executors_data = self.executors_data.write(); - if let Some(executor_data) = - executors_data.get_mut(&executor_data_change.executor_id) - { - let available_task_slots = executor_data.available_task_slots as i32 - + executor_data_change.task_slots; - if available_task_slots < 0 { - error!( - "Available task slots {} for executor {} is less than 0", - available_task_slots, executor_data.executor_id - ); - } else { - info!( - "available_task_slots for executor {} becomes {}", - executor_data.executor_id, available_task_slots - ); - executor_data.available_task_slots = available_task_slots as u32; - } - } else { - warn!( - "Could not find executor data for {}", - executor_data_change.executor_id - ); - } - } - - pub(crate) fn get_executor_data(&self, executor_id: &str) -> Option { - let executors_data = self.executors_data.read(); - executors_data.get(executor_id).cloned() - } - - /// There are two checks: - /// 1. firstly alive - /// 2. secondly available task slots > 0 - #[cfg(not(test))] - #[allow(dead_code)] - pub(crate) fn get_available_executors_data(&self) -> Vec { - let mut res = { - let alive_executors = self.get_alive_executors_within_one_minute(); - let executors_data = self.executors_data.read(); - executors_data - .iter() - .filter_map(|(exec, data)| { - (data.available_task_slots > 0 && alive_executors.contains(exec)) - .then(|| data.clone()) - }) - .collect::>() - }; - res.sort_by(|a, b| Ord::cmp(&b.available_task_slots, &a.available_task_slots)); - res - } - - #[cfg(test)] - #[allow(dead_code)] - pub(crate) fn get_available_executors_data(&self) -> Vec { - let mut res: Vec = - self.executors_data.read().values().cloned().collect(); - res.sort_by(|a, b| Ord::cmp(&b.available_task_slots, &a.available_task_slots)); - res - } -} diff --git a/ballista/rust/scheduler/src/state/mod.rs b/ballista/rust/scheduler/src/state/mod.rs deleted file mode 100644 index 5f301346f32a..000000000000 --- a/ballista/rust/scheduler/src/state/mod.rs +++ /dev/null @@ -1,258 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::HashMap; -use std::sync::Arc; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -use datafusion::physical_plan::ExecutionPlan; - -use ballista_core::error::Result; - -use crate::scheduler_server::{SessionBuilder, SessionContextRegistry}; -use ballista_core::serde::protobuf::{ExecutorHeartbeat, JobStatus, KeyValuePair}; -use ballista_core::serde::scheduler::ExecutorMetadata; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan, BallistaCodec}; - -use crate::state::backend::StateBackendClient; -use crate::state::executor_manager::ExecutorManager; -use crate::state::persistent_state::PersistentSchedulerState; -use crate::state::stage_manager::StageManager; - -pub mod backend; -mod executor_manager; -mod persistent_state; -mod stage_manager; -pub mod task_scheduler; - -#[derive(Clone)] -pub(super) struct SchedulerState -{ - persistent_state: PersistentSchedulerState, - pub executor_manager: ExecutorManager, - pub stage_manager: StageManager, -} - -impl SchedulerState { - pub fn new( - config_client: Arc, - namespace: String, - session_builder: SessionBuilder, - codec: BallistaCodec, - ) -> Self { - Self { - persistent_state: PersistentSchedulerState::new( - config_client, - namespace, - session_builder, - codec, - ), - executor_manager: ExecutorManager::new(), - stage_manager: StageManager::new(), - } - } - - pub async fn init(&self) -> Result<()> { - self.persistent_state.init().await?; - - Ok(()) - } - - pub fn get_codec(&self) -> &BallistaCodec { - &self.persistent_state.codec - } - - pub async fn get_executors_metadata( - &self, - ) -> Result> { - let mut result = vec![]; - - let executors_heartbeat = self - .executor_manager - .get_executors_heartbeat() - .into_iter() - .map(|heartbeat| (heartbeat.executor_id.clone(), heartbeat)) - .collect::>(); - - let executors_metadata = self.persistent_state.get_executors_metadata(); - - let now_epoch_ts = SystemTime::now() - .duration_since(UNIX_EPOCH) - .expect("Time went backwards"); - - for meta in executors_metadata.into_iter() { - // If there's no heartbeat info for an executor, regard its heartbeat timestamp as 0 - // so that it will always be excluded when requesting alive executors - let ts = executors_heartbeat - .get(&meta.id) - .map(|heartbeat| Duration::from_secs(heartbeat.timestamp)) - .unwrap_or_else(|| Duration::from_secs(0)); - let time_since_last_seen = now_epoch_ts - .checked_sub(ts) - .unwrap_or_else(|| Duration::from_secs(0)); - result.push((meta, time_since_last_seen)); - } - Ok(result) - } - - pub fn get_executor_metadata(&self, executor_id: &str) -> Option { - self.persistent_state.get_executor_metadata(executor_id) - } - - pub async fn save_executor_metadata( - &self, - executor_meta: ExecutorMetadata, - ) -> Result<()> { - self.persistent_state - .save_executor_metadata(executor_meta) - .await - } - - pub async fn save_job_session( - &self, - job_id: &str, - session_id: &str, - configs: Vec, - ) -> Result<()> { - self.persistent_state - .save_job_session(job_id, session_id, configs) - .await - } - - pub fn get_session_from_job(&self, job_id: &str) -> Option { - self.persistent_state.get_session_from_job(job_id) - } - - pub async fn save_job_metadata( - &self, - job_id: &str, - status: &JobStatus, - ) -> Result<()> { - self.persistent_state - .save_job_metadata(job_id, status) - .await - } - - pub fn get_job_metadata(&self, job_id: &str) -> Option { - self.persistent_state.get_job_metadata(job_id) - } - - pub async fn save_stage_plan( - &self, - job_id: &str, - stage_id: usize, - plan: Arc, - ) -> Result<()> { - self.persistent_state - .save_stage_plan(job_id, stage_id, plan) - .await - } - - pub fn get_stage_plan( - &self, - job_id: &str, - stage_id: usize, - ) -> Option> { - self.persistent_state.get_stage_plan(job_id, stage_id) - } - - pub fn session_registry(&self) -> Arc { - self.persistent_state.session_registry() - } -} - -#[cfg(all(test, feature = "sled"))] -mod test { - use std::sync::Arc; - - use ballista_core::error::BallistaError; - use ballista_core::serde::protobuf::{ - job_status, JobStatus, LogicalPlanNode, PhysicalPlanNode, QueuedJob, - }; - use ballista_core::serde::scheduler::{ExecutorMetadata, ExecutorSpecification}; - use ballista_core::serde::BallistaCodec; - use datafusion::execution::context::default_session_builder; - - use super::{backend::standalone::StandaloneClient, SchedulerState}; - - #[tokio::test] - async fn executor_metadata() -> Result<(), BallistaError> { - let state: SchedulerState = - SchedulerState::new( - Arc::new(StandaloneClient::try_new_temporary()?), - "test".to_string(), - default_session_builder, - BallistaCodec::default(), - ); - let meta = ExecutorMetadata { - id: "123".to_owned(), - host: "localhost".to_owned(), - port: 123, - grpc_port: 124, - specification: ExecutorSpecification { task_slots: 2 }, - }; - state.save_executor_metadata(meta.clone()).await?; - let result: Vec<_> = state - .get_executors_metadata() - .await? - .into_iter() - .map(|(meta, _)| meta) - .collect(); - assert_eq!(vec![meta], result); - Ok(()) - } - - #[tokio::test] - async fn job_metadata() -> Result<(), BallistaError> { - let state: SchedulerState = - SchedulerState::new( - Arc::new(StandaloneClient::try_new_temporary()?), - "test".to_string(), - default_session_builder, - BallistaCodec::default(), - ); - let meta = JobStatus { - status: Some(job_status::Status::Queued(QueuedJob {})), - }; - state.save_job_metadata("job", &meta).await?; - let result = state.get_job_metadata("job").unwrap(); - assert!(result.status.is_some()); - match result.status.unwrap() { - job_status::Status::Queued(_) => (), - _ => panic!("Unexpected status"), - } - Ok(()) - } - - #[tokio::test] - async fn job_metadata_non_existant() -> Result<(), BallistaError> { - let state: SchedulerState = - SchedulerState::new( - Arc::new(StandaloneClient::try_new_temporary()?), - "test".to_string(), - default_session_builder, - BallistaCodec::default(), - ); - let meta = JobStatus { - status: Some(job_status::Status::Queued(QueuedJob {})), - }; - state.save_job_metadata("job", &meta).await?; - let result = state.get_job_metadata("job2"); - assert!(result.is_none()); - Ok(()) - } -} diff --git a/ballista/rust/scheduler/src/state/persistent_state.rs b/ballista/rust/scheduler/src/state/persistent_state.rs deleted file mode 100644 index 428d523ac423..000000000000 --- a/ballista/rust/scheduler/src/state/persistent_state.rs +++ /dev/null @@ -1,528 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use ballista_core::config::BallistaConfig; -use log::{debug, error}; -use parking_lot::RwLock; -use prost::Message; -use std::any::type_name; -use std::collections::HashMap; -use std::ops::Deref; -use std::sync::Arc; - -use ballista_core::error::{BallistaError, Result}; - -use ballista_core::serde::protobuf::{JobSessionConfig, JobStatus, KeyValuePair}; - -use crate::scheduler_server::{ - create_datafusion_context, SessionBuilder, SessionContextRegistry, -}; -use crate::state::backend::StateBackendClient; -use crate::state::stage_manager::StageKey; -use ballista_core::serde::scheduler::ExecutorMetadata; -use ballista_core::serde::{protobuf, AsExecutionPlan, AsLogicalPlan, BallistaCodec}; -use datafusion::physical_plan::ExecutionPlan; - -#[derive(Clone)] -pub(crate) struct PersistentSchedulerState< - T: 'static + AsLogicalPlan, - U: 'static + AsExecutionPlan, -> { - // for db - config_client: Arc, - namespace: String, - pub(crate) codec: BallistaCodec, - - // for in-memory cache - executors_metadata: Arc>>, - - // TODO add remove logic - jobs: Arc>>, - stages: Arc>>>, - job2session: Arc>>, - - /// DataFusion session contexts that are registered within the Scheduler - session_context_registry: Arc, - - session_builder: SessionBuilder, -} - -impl - PersistentSchedulerState -{ - pub(crate) fn new( - config_client: Arc, - namespace: String, - session_builder: SessionBuilder, - codec: BallistaCodec, - ) -> Self { - Self { - config_client, - namespace, - codec, - executors_metadata: Arc::new(RwLock::new(HashMap::new())), - jobs: Arc::new(RwLock::new(HashMap::new())), - stages: Arc::new(RwLock::new(HashMap::new())), - job2session: Arc::new(RwLock::new(HashMap::new())), - session_context_registry: Arc::new(SessionContextRegistry::default()), - session_builder, - } - } - - /// Load the state stored in storage into memory - pub(crate) async fn init(&self) -> Result<()> { - self.init_executors_metadata_from_storage().await?; - self.init_jobs_from_storage().await?; - self.init_stages_from_storage().await?; - - Ok(()) - } - - async fn init_executors_metadata_from_storage(&self) -> Result<()> { - let entries = self - .config_client - .get_from_prefix(&get_executors_metadata_prefix(&self.namespace)) - .await?; - - let mut executors_metadata = self.executors_metadata.write(); - for (_key, entry) in entries { - let meta: protobuf::ExecutorMetadata = decode_protobuf(&entry)?; - executors_metadata.insert(meta.id.clone(), meta.into()); - } - - Ok(()) - } - - async fn init_jobs_from_storage(&self) -> Result<()> { - let entries = self - .config_client - .get_from_prefix(&get_job_prefix(&self.namespace)) - .await?; - - let mut jobs = self.jobs.write(); - for (key, entry) in entries { - let job: JobStatus = decode_protobuf(&entry)?; - let job_id = extract_job_id_from_job_key(&key) - .map(|job_id| job_id.to_string()) - .unwrap(); - jobs.insert(job_id, job); - } - - Ok(()) - } - - async fn init_stages_from_storage(&self) -> Result<()> { - let entries = self - .config_client - .get_from_prefix(&get_stage_prefix(&self.namespace)) - .await?; - - let mut tmp_stages: HashMap> = HashMap::new(); - { - for (key, entry) in entries { - let (job_id, stage_id) = extract_stage_id_from_stage_key(&key).unwrap(); - let job_session = self - .config_client - .get(&get_job_config_key(&self.namespace, &job_id)) - .await?; - let job_session: JobSessionConfig = decode_protobuf(&job_session)?; - - // Rebuild SessionContext from serialized settings - let mut config_builder = BallistaConfig::builder(); - for kv_pair in &job_session.configs { - config_builder = config_builder.set(&kv_pair.key, &kv_pair.value); - } - let config = config_builder.build().map_err(|e| { - let msg = format!("Could not parse configs: {}", e); - error!("{}", msg); - BallistaError::Internal(format!( - "Error building configs for job ID {}", - job_id - )) - })?; - - let session_ctx = - create_datafusion_context(&config, self.session_builder); - self.session_registry() - .register_session(session_ctx.clone()) - .await; - - let value = U::try_decode(&entry)?; - let runtime = session_ctx.runtime_env(); - let plan = value.try_into_physical_plan( - session_ctx.deref(), - runtime.deref(), - self.codec.physical_extension_codec(), - )?; - - let mut job2_sess = self.job2session.write(); - job2_sess.insert(job_id.clone(), job_session.session_id); - - tmp_stages.insert((job_id, stage_id), plan); - } - } - let mut stages = self.stages.write(); - for tmp_stage in tmp_stages { - stages.insert(tmp_stage.0, tmp_stage.1); - } - Ok(()) - } - - pub(crate) async fn save_executor_metadata( - &self, - executor_meta: ExecutorMetadata, - ) -> Result<()> { - { - // Save in db - let key = get_executor_metadata_key(&self.namespace, &executor_meta.id); - let value = { - let executor_meta: protobuf::ExecutorMetadata = - executor_meta.clone().into(); - encode_protobuf(&executor_meta)? - }; - self.synchronize_save(key, value).await?; - } - - { - // Save in memory - let mut executors_metadata = self.executors_metadata.write(); - executors_metadata.insert(executor_meta.id.clone(), executor_meta); - } - - Ok(()) - } - - pub(crate) fn get_executor_metadata( - &self, - executor_id: &str, - ) -> Option { - let executors_metadata = self.executors_metadata.read(); - executors_metadata.get(executor_id).cloned() - } - - pub(crate) fn get_executors_metadata(&self) -> Vec { - let executors_metadata = self.executors_metadata.read(); - executors_metadata.values().cloned().collect() - } - - pub(crate) async fn save_job_session( - &self, - job_id: &str, - session_id: &str, - configs: Vec, - ) -> Result<()> { - let key = get_job_config_key(&self.namespace, job_id); - let value = encode_protobuf(&protobuf::JobSessionConfig { - session_id: session_id.to_string(), - configs, - })?; - - self.synchronize_save(key, value).await?; - - let mut job2_sess = self.job2session.write(); - job2_sess.insert(job_id.to_string(), session_id.to_string()); - - Ok(()) - } - - pub(crate) fn get_session_from_job(&self, job_id: &str) -> Option { - let job_session = self.job2session.read(); - job_session.get(job_id).cloned() - } - - pub(crate) async fn save_job_metadata( - &self, - job_id: &str, - status: &JobStatus, - ) -> Result<()> { - debug!("Saving job metadata: {:?}", status); - { - // Save in db - let key = get_job_key(&self.namespace, job_id); - let value = encode_protobuf(status)?; - self.synchronize_save(key, value).await?; - } - - { - // Save in memory - let mut jobs = self.jobs.write(); - jobs.insert(job_id.to_string(), status.clone()); - } - - Ok(()) - } - - pub(crate) fn get_job_metadata(&self, job_id: &str) -> Option { - let jobs = self.jobs.read(); - jobs.get(job_id).cloned() - } - - pub(crate) async fn save_stage_plan( - &self, - job_id: &str, - stage_id: usize, - plan: Arc, - ) -> Result<()> { - { - // Save in db - let key = get_stage_plan_key(&self.namespace, job_id, stage_id as u32); - let value = { - let mut buf: Vec = vec![]; - let proto = U::try_from_physical_plan( - plan.clone(), - self.codec.physical_extension_codec(), - )?; - proto.try_encode(&mut buf)?; - - buf - }; - self.synchronize_save(key, value).await?; - } - - { - // Save in memory - let mut stages = self.stages.write(); - stages.insert((job_id.to_string(), stage_id as u32), plan); - } - - Ok(()) - } - - pub(crate) fn get_stage_plan( - &self, - job_id: &str, - stage_id: usize, - ) -> Option> { - let stages = self.stages.read(); - let key = (job_id.to_string(), stage_id as u32); - stages.get(&key).cloned() - } - - async fn synchronize_save(&self, key: String, value: Vec) -> Result<()> { - let mut lock = self.config_client.lock().await?; - self.config_client.put(key, value).await?; - lock.unlock().await; - - Ok(()) - } - - pub fn session_registry(&self) -> Arc { - self.session_context_registry.clone() - } -} - -fn get_executors_metadata_prefix(namespace: &str) -> String { - format!("/ballista/{}/executor_metadata", namespace) -} - -fn get_executor_metadata_key(namespace: &str, id: &str) -> String { - format!("{}/{}", get_executors_metadata_prefix(namespace), id) -} - -fn get_job_prefix(namespace: &str) -> String { - format!("/ballista/{}/jobs", namespace) -} - -fn get_job_key(namespace: &str, id: &str) -> String { - format!("{}/{}", get_job_prefix(namespace), id) -} - -fn get_job_config_key(namespace: &str, id: &str) -> String { - format!("config/{}/{}", get_job_prefix(namespace), id) -} - -fn get_stage_prefix(namespace: &str) -> String { - format!("/ballista/{}/stages", namespace,) -} - -fn get_stage_plan_key(namespace: &str, job_id: &str, stage_id: u32) -> String { - format!("{}/{}/{}", get_stage_prefix(namespace), job_id, stage_id,) -} -fn extract_job_id_from_job_key(job_key: &str) -> Result<&str> { - job_key.split('/').nth(2).ok_or_else(|| { - BallistaError::Internal(format!("Unexpected task key: {}", job_key)) - }) -} - -fn extract_stage_id_from_stage_key(stage_key: &str) -> Result { - let splits: Vec<&str> = stage_key.split('/').collect(); - if splits.len() > 4 { - Ok(( - splits[splits.len() - 2].to_string(), - splits[splits.len() - 1].parse::().map_err(|e| { - BallistaError::Internal(format!( - "Invalid stage ID in stage key: {}, {:?}", - stage_key, e - )) - })?, - )) - } else { - Err(BallistaError::Internal(format!( - "Unexpected stage key: {}", - stage_key - ))) - } -} - -fn decode_protobuf(bytes: &[u8]) -> Result { - T::decode(bytes).map_err(|e| { - BallistaError::Internal(format!( - "Could not deserialize {}: {}", - type_name::(), - e - )) - }) -} - -fn encode_protobuf(msg: &T) -> Result> { - let mut value: Vec = Vec::with_capacity(msg.encoded_len()); - msg.encode(&mut value).map_err(|e| { - BallistaError::Internal(format!( - "Could not serialize {}: {}", - type_name::(), - e - )) - })?; - Ok(value) -} - -#[cfg(test)] -mod test { - use super::extract_stage_id_from_stage_key; - use crate::state::backend::standalone::StandaloneClient; - - use crate::state::persistent_state::PersistentSchedulerState; - - use ballista_core::serde::protobuf::job_status::Status; - use ballista_core::serde::protobuf::{ - JobStatus, LogicalPlanNode, PhysicalPlanNode, QueuedJob, - }; - use ballista_core::serde::BallistaCodec; - use datafusion::execution::context::default_session_builder; - use datafusion::logical_plan::LogicalPlanBuilder; - use datafusion::prelude::SessionContext; - - use std::sync::Arc; - - #[test] - fn test_extract_stage_id_from_stage_key() { - let (job_id, stage_id) = - extract_stage_id_from_stage_key("/ballista/default/stages/2Yoyba8/1") - .expect("extracting stage key"); - - assert_eq!(job_id.as_str(), "2Yoyba8"); - assert_eq!(stage_id, 1); - - let (job_id, stage_id) = - extract_stage_id_from_stage_key("ballista/default/stages/2Yoyba8/1") - .expect("extracting stage key"); - - assert_eq!(job_id.as_str(), "2Yoyba8"); - assert_eq!(stage_id, 1); - - let (job_id, stage_id) = - extract_stage_id_from_stage_key("ballista//stages/2Yoyba8/1") - .expect("extracting stage key"); - - assert_eq!(job_id.as_str(), "2Yoyba8"); - assert_eq!(stage_id, 1); - } - - #[tokio::test] - async fn test_init_from_storage() { - let ctx = SessionContext::new(); - - let plan = LogicalPlanBuilder::empty(true) - .build() - .expect("create empty logical plan"); - let plan = ctx - .create_physical_plan(&plan) - .await - .expect("create physical plan"); - - let expected_plan = format!("{:?}", plan); - - let job_id = "job-id".to_string(); - let session_id = "session-id".to_string(); - - let config_client = Arc::new( - StandaloneClient::try_new_temporary().expect("creating config client"), - ); - - let persistent_state: PersistentSchedulerState< - LogicalPlanNode, - PhysicalPlanNode, - > = PersistentSchedulerState::new( - config_client.clone(), - "default".to_string(), - default_session_builder, - BallistaCodec::default(), - ); - - persistent_state - .save_job_session(&job_id, &session_id, vec![]) - .await - .expect("saving session"); - persistent_state - .save_job_metadata( - &job_id, - &JobStatus { - status: Some(Status::Queued(QueuedJob {})), - }, - ) - .await - .expect("saving job metadata"); - persistent_state - .save_stage_plan(&job_id, 1, plan) - .await - .expect("saving stage plan"); - - assert_eq!( - persistent_state - .get_stage_plan(&job_id, 1) - .map(|plan| format!("{:?}", plan)), - Some(expected_plan.clone()) - ); - assert_eq!( - persistent_state.get_session_from_job(&job_id), - Some("session-id".to_string()) - ); - - let persistent_state: PersistentSchedulerState< - LogicalPlanNode, - PhysicalPlanNode, - > = PersistentSchedulerState::new( - config_client.clone(), - "default".to_string(), - default_session_builder, - BallistaCodec::default(), - ); - - persistent_state.init().await.expect("initializing state"); - - assert_eq!( - persistent_state - .get_stage_plan(&job_id, 1) - .map(|plan| format!("{:?}", plan)), - Some(expected_plan.clone()) - ); - assert_eq!( - persistent_state.get_session_from_job(&job_id), - Some("session-id".to_string()) - ); - } -} diff --git a/ballista/rust/scheduler/src/state/stage_manager.rs b/ballista/rust/scheduler/src/state/stage_manager.rs deleted file mode 100644 index e926c1db444c..000000000000 --- a/ballista/rust/scheduler/src/state/stage_manager.rs +++ /dev/null @@ -1,783 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::{HashMap, HashSet}; -use std::sync::Arc; - -use log::{debug, error, warn}; -use parking_lot::RwLock; -use rand::Rng; - -use crate::scheduler_server::event::QueryStageSchedulerEvent; -use crate::state::task_scheduler::StageScheduler; -use ballista_core::error::{BallistaError, Result}; -use ballista_core::serde::protobuf; -use ballista_core::serde::protobuf::{task_status, FailedTask, TaskStatus}; - -/// job_id + stage_id -pub type StageKey = (String, u32); - -#[derive(Clone)] -pub struct StageManager { - stage_distribution: Arc>, - - // The final stage id for jobs - final_stages: Arc>>, - - // (job_id, stage_id) -> stage set in which each one depends on (job_id, stage_id) - stages_dependency: Arc>>>, - - // job_id -> pending stages - pending_stages: Arc>>>, -} - -impl StageManager { - pub fn new() -> Self { - Self { - stage_distribution: Arc::new(RwLock::new(StageDistribution::new())), - final_stages: Arc::new(RwLock::new(HashMap::new())), - stages_dependency: Arc::new(RwLock::new(HashMap::new())), - pending_stages: Arc::new(RwLock::new(HashMap::new())), - } - } - - pub fn add_final_stage(&self, job_id: &str, stage_id: u32) { - let mut final_stages = self.final_stages.write(); - final_stages.insert(job_id.to_owned(), stage_id); - } - - pub fn is_final_stage(&self, job_id: &str, stage_id: u32) -> bool { - self.get_final_stage_id(job_id) - .map(|final_stage_id| final_stage_id == stage_id) - .unwrap_or(false) - } - - fn get_final_stage_id(&self, job_id: &str) -> Option { - let final_stages = self.final_stages.read(); - final_stages.get(job_id).cloned() - } - - pub fn get_tasks_for_complete_final_stage( - &self, - job_id: &str, - ) -> Result>> { - let final_stage_id = self.get_final_stage_id(job_id).ok_or_else(|| { - BallistaError::General(format!( - "Fail to find final stage id for job {}", - job_id - )) - })?; - - let stage_key = (job_id.to_owned(), final_stage_id); - let stage_distribution = self.stage_distribution.read(); - - if let Some(stage) = stage_distribution.stages_completed.get(&stage_key) { - Ok(stage.tasks.clone()) - } else { - Err(BallistaError::General(format!( - "The final stage id {} has not been completed yet", - final_stage_id - ))) - } - } - - pub fn add_pending_stage(&self, job_id: &str, stage_id: u32) { - let mut pending_stages = self.pending_stages.write(); - pending_stages - .entry(job_id.to_owned()) - .or_insert_with(HashSet::new) - .insert(stage_id); - } - - pub fn is_pending_stage(&self, job_id: &str, stage_id: u32) -> bool { - let pending_stages = self.pending_stages.read(); - if let Some(pending_stages) = pending_stages.get(job_id) { - pending_stages.contains(&stage_id) - } else { - false - } - } - - pub fn remove_pending_stage( - &self, - job_id: &str, - stages_remove: &HashSet, - ) -> bool { - let mut pending_stages = self.pending_stages.write(); - let mut is_stages_empty = false; - let ret = if let Some(stages) = pending_stages.get_mut(job_id) { - let len_before_remove = stages.len(); - for stage_id in stages_remove { - stages.remove(stage_id); - } - is_stages_empty = stages.is_empty(); - stages.len() != len_before_remove - } else { - false - }; - - if is_stages_empty { - pending_stages.remove(job_id); - } - - ret - } - - pub fn add_stages_dependency( - &self, - job_id: &str, - dependencies: HashMap>, - ) { - let mut stages_dependency = self.stages_dependency.write(); - for (stage_id, parent_stages) in dependencies.into_iter() { - stages_dependency.insert((job_id.to_owned(), stage_id), parent_stages); - } - } - - pub fn get_parent_stages(&self, job_id: &str, stage_id: u32) -> Option> { - let stage_key = (job_id.to_owned(), stage_id); - let stages_dependency = self.stages_dependency.read(); - stages_dependency.get(&stage_key).cloned() - } - - pub fn add_running_stage(&self, job_id: &str, stage_id: u32, num_partitions: u32) { - let stage = Stage::new(job_id, stage_id, num_partitions); - - let mut stage_distribution = self.stage_distribution.write(); - stage_distribution - .stages_running - .insert((job_id.to_string(), stage_id), stage); - } - - pub fn is_running_stage(&self, job_id: &str, stage_id: u32) -> bool { - let stage_key = (job_id.to_owned(), stage_id); - let stage_distribution = self.stage_distribution.read(); - stage_distribution.stages_running.get(&stage_key).is_some() - } - - pub fn is_completed_stage(&self, job_id: &str, stage_id: u32) -> bool { - let stage_key = (job_id.to_owned(), stage_id); - let stage_distribution = self.stage_distribution.read(); - stage_distribution - .stages_completed - .get(&stage_key) - .is_some() - } - - pub(crate) fn get_stage_tasks( - &self, - job_id: &str, - stage_id: u32, - ) -> Option>> { - let stage_key = (job_id.to_owned(), stage_id); - let stage_distribution = self.stage_distribution.read(); - if let Some(stage) = stage_distribution.stages_running.get(&stage_key) { - Some(stage.tasks.clone()) - } else { - stage_distribution - .stages_completed - .get(&stage_key) - .map(|task| task.tasks.clone()) - } - } - - pub(crate) fn update_tasks_status( - &self, - tasks_status: Vec, - ) -> Vec { - let mut all_tasks_status: HashMap> = HashMap::new(); - for task_status in tasks_status { - if let Some(task_id) = task_status.task_id.as_ref() { - let stage_tasks_status = all_tasks_status - .entry((task_id.job_id.clone(), task_id.stage_id)) - .or_insert_with(Vec::new); - stage_tasks_status.push(task_status); - } else { - error!("There's no task id when updating status"); - } - } - - let mut ret = vec![]; - let mut stage_distribution = self.stage_distribution.write(); - for (stage_key, stage_tasks_status) in all_tasks_status.into_iter() { - if let Some(stage) = stage_distribution.stages_running.get_mut(&stage_key) { - for task_status in &stage_tasks_status { - stage.update_task_status(task_status); - } - if let Some(fail_message) = stage.get_fail_message() { - ret.push(QueryStageSchedulerEvent::JobFailed( - stage_key.0.clone(), - stage_key.1, - fail_message, - )); - } else if stage.is_completed() { - stage_distribution.complete_stage(stage_key.clone()); - if self.is_final_stage(&stage_key.0, stage_key.1) { - ret.push(QueryStageSchedulerEvent::JobFinished( - stage_key.0.clone(), - )); - } else { - ret.push(QueryStageSchedulerEvent::StageFinished( - stage_key.0.clone(), - stage_key.1, - )); - } - } - } else { - error!("Fail to find stage for {:?}/{}", &stage_key.0, stage_key.1); - } - } - - ret - } - - pub fn fetch_pending_tasks( - &self, - max_num: usize, - cond: F, - ) -> Option<(String, u32, Vec)> - where - F: Fn(&StageKey) -> bool, - { - if let Some(next_stage) = self.fetch_schedulable_stage(cond) { - if let Some(next_tasks) = - self.find_stage_pending_tasks(&next_stage.0, next_stage.1, max_num) - { - Some((next_stage.0.to_owned(), next_stage.1, next_tasks)) - } else { - warn!( - "Fail to find pending tasks for stage {}/{}", - next_stage.0, next_stage.1 - ); - None - } - } else { - None - } - } - - fn find_stage_pending_tasks( - &self, - job_id: &str, - stage_id: u32, - max_num: usize, - ) -> Option> { - let stage_key = (job_id.to_owned(), stage_id); - let stage_distribution = self.stage_distribution.read(); - stage_distribution - .stages_running - .get(&stage_key) - .map(|stage| stage.find_pending_tasks(max_num)) - } - - pub fn has_running_tasks(&self) -> bool { - let stage_distribution = self.stage_distribution.read(); - for stage in stage_distribution.stages_running.values() { - if !stage.get_running_tasks().is_empty() { - return true; - } - } - - false - } -} - -// TODO Currently, it will randomly choose a stage. In the future, we can add more sophisticated stage choose algorithm here, like priority, etc. -impl StageScheduler for StageManager { - fn fetch_schedulable_stage(&self, cond: F) -> Option - where - F: Fn(&StageKey) -> bool, - { - let mut rng = rand::thread_rng(); - let stage_distribution = self.stage_distribution.read(); - let stages_running = &stage_distribution.stages_running; - if stages_running.is_empty() { - debug!("There's no running stages"); - return None; - } - let stages = stages_running - .iter() - .filter(|entry| entry.1.is_schedulable() && cond(entry.0)) - .map(|entry| entry.0) - .collect::>(); - if stages.is_empty() { - None - } else { - let n_th = rng.gen_range(0..stages.len()); - Some(stages[n_th].clone()) - } - } -} - -struct StageDistribution { - // The key is (job_id, stage_id) - stages_running: HashMap, - stages_completed: HashMap, -} - -impl StageDistribution { - fn new() -> Self { - Self { - stages_running: HashMap::new(), - stages_completed: HashMap::new(), - } - } - - fn complete_stage(&mut self, stage_key: StageKey) { - if let Some(stage) = self.stages_running.remove(&stage_key) { - assert!( - stage.is_completed(), - "Stage {}/{} is not completed", - stage_key.0, - stage_key.1 - ); - self.stages_completed.insert(stage_key, stage); - } else { - warn!( - "Fail to find running stage {:?}/{}", - stage_key.0, stage_key.1 - ); - } - } -} - -pub struct Stage { - pub stage_id: u32, - tasks: Vec>, - - tasks_distribution: TaskStatusDistribution, -} - -impl Stage { - fn new(job_id: &str, stage_id: u32, num_partitions: u32) -> Self { - let mut tasks = vec![]; - for partition_id in 0..num_partitions { - let pending_status = Arc::new(TaskStatus { - task_id: Some(protobuf::PartitionId { - job_id: job_id.to_owned(), - stage_id, - partition_id, - }), - status: None, - }); - - tasks.push(pending_status); - } - - Stage { - stage_id, - tasks, - tasks_distribution: TaskStatusDistribution::new(num_partitions as usize), - } - } - - // If error happens for updating some task status, just quietly print the error message - fn update_task_status(&mut self, task: &TaskStatus) { - if let Some(task_id) = &task.task_id { - let task_idx = task_id.partition_id as usize; - if task_idx < self.tasks.len() { - let existing_task_status = self.tasks[task_idx].clone(); - if self.tasks_distribution.update( - task_idx, - &existing_task_status.status, - &task.status, - ) { - self.tasks[task_idx] = Arc::new(task.clone()); - } else { - error!( - "Fail to update status from {:?} to {:?} for task: {:?}/{:?}/{:?}", &existing_task_status.status, &task.status, - &task_id.job_id, &task_id.stage_id, task_idx - ) - } - } else { - error!( - "Fail to find existing task: {:?}/{:?}/{:?}", - &task_id.job_id, &task_id.stage_id, task_idx - ) - } - } else { - error!("Fail to update task status due to no task id"); - } - } - - fn is_schedulable(&self) -> bool { - self.tasks_distribution.is_schedulable() - } - - fn is_completed(&self) -> bool { - self.tasks_distribution.is_completed() - } - - // If return None, means no failed tasks - fn get_fail_message(&self) -> Option { - if self.tasks_distribution.is_failed() { - let task_idx = self.tasks_distribution.sample_failed_index(); - if let Some(task) = self.tasks.get(task_idx) { - if let Some(task_status::Status::Failed(FailedTask { error })) = - &task.status - { - Some(error.clone()) - } else { - warn!("task {:?} is not failed", task); - None - } - } else { - warn!("Could not find error tasks"); - None - } - } else { - None - } - } - - pub fn find_pending_tasks(&self, max_num: usize) -> Vec { - self.tasks_distribution.find_pending_indicators(max_num) - } - - fn get_running_tasks(&self) -> Vec> { - self.tasks_distribution - .running_indicator - .indicator - .iter() - .enumerate() - .filter(|(_i, is_running)| **is_running) - .map(|(i, _is_running)| self.tasks[i].clone()) - .collect() - } -} - -#[derive(Clone)] -struct TaskStatusDistribution { - len: usize, - pending_indicator: TaskStatusIndicator, - running_indicator: TaskStatusIndicator, - failed_indicator: TaskStatusIndicator, - completed_indicator: TaskStatusIndicator, -} - -impl TaskStatusDistribution { - fn new(len: usize) -> Self { - Self { - len, - pending_indicator: TaskStatusIndicator { - indicator: (0..len).map(|_| true).collect::>(), - n_of_true: len, - }, - running_indicator: TaskStatusIndicator { - indicator: (0..len).map(|_| false).collect::>(), - n_of_true: 0, - }, - failed_indicator: TaskStatusIndicator { - indicator: (0..len).map(|_| false).collect::>(), - n_of_true: 0, - }, - completed_indicator: TaskStatusIndicator { - indicator: (0..len).map(|_| false).collect::>(), - n_of_true: 0, - }, - } - } - - fn is_schedulable(&self) -> bool { - self.pending_indicator.n_of_true != 0 - } - - fn is_completed(&self) -> bool { - self.completed_indicator.n_of_true == self.len - } - - fn is_failed(&self) -> bool { - self.failed_indicator.n_of_true != 0 - } - - fn sample_failed_index(&self) -> usize { - for i in 0..self.len { - if self.failed_indicator.indicator[i] { - return i; - } - } - - self.len - } - - fn find_pending_indicators(&self, max_num: usize) -> Vec { - let mut ret = vec![]; - if max_num < 1 { - return ret; - } - - let len = std::cmp::min(max_num, self.len); - for idx in 0..self.len { - if self.pending_indicator.indicator[idx] { - ret.push(idx as u32); - if ret.len() >= len { - break; - } - } - } - - ret - } - - fn update( - &mut self, - idx: usize, - from: &Option, - to: &Option, - ) -> bool { - assert!( - idx < self.len, - "task index {} should be smaller than {}", - idx, - self.len - ); - - match (from, to) { - (Some(from), Some(to)) => match (from, to) { - (task_status::Status::Running(_), task_status::Status::Failed(_)) => { - self.running_indicator.set_false(idx); - self.failed_indicator.set_true(idx); - } - (task_status::Status::Running(_), task_status::Status::Completed(_)) => { - self.running_indicator.set_false(idx); - self.completed_indicator.set_true(idx); - } - _ => { - return false; - } - }, - (None, Some(task_status::Status::Running(_))) => { - self.pending_indicator.set_false(idx); - self.running_indicator.set_true(idx); - } - (Some(from), None) => match from { - task_status::Status::Failed(_) => { - self.failed_indicator.set_false(idx); - self.pending_indicator.set_true(idx); - } - task_status::Status::Completed(_) => { - self.completed_indicator.set_false(idx); - self.pending_indicator.set_true(idx); - } - _ => { - return false; - } - }, - _ => { - return false; - } - } - - true - } -} - -#[derive(Clone)] -struct TaskStatusIndicator { - indicator: Vec, - n_of_true: usize, -} - -impl TaskStatusIndicator { - fn set_false(&mut self, idx: usize) { - self.indicator[idx] = false; - self.n_of_true -= 1; - } - - fn set_true(&mut self, idx: usize) { - self.indicator[idx] = true; - self.n_of_true += 1; - } -} - -#[cfg(test)] -mod test { - use crate::state::stage_manager::StageManager; - use ballista_core::error::Result; - use ballista_core::serde::protobuf::{ - task_status, CompletedTask, FailedTask, PartitionId, RunningTask, TaskStatus, - }; - - #[tokio::test] - async fn test_task_status_state_machine_failed() -> Result<()> { - let stage_manager = StageManager::new(); - - let num_partitions = 3; - let job_id = "job"; - let stage_id = 1u32; - - stage_manager.add_running_stage(job_id, stage_id, num_partitions); - - let task_id = PartitionId { - job_id: job_id.to_owned(), - stage_id, - partition_id: 2, - }; - - { - // Invalid transformation from Pending to Failed - stage_manager.update_tasks_status(vec![TaskStatus { - status: Some(task_status::Status::Failed(FailedTask { - error: "error".to_owned(), - })), - task_id: Some(task_id.clone()), - }]); - let ret = stage_manager.get_stage_tasks(job_id, stage_id); - assert!(ret.is_some()); - assert!(ret - .unwrap() - .get(task_id.partition_id as usize) - .unwrap() - .status - .is_none()); - } - - { - // Valid transformation from Pending to Running to Failed - stage_manager.update_tasks_status(vec![TaskStatus { - status: Some(task_status::Status::Running(RunningTask { - executor_id: "localhost".to_owned(), - })), - task_id: Some(task_id.clone()), - }]); - stage_manager.update_tasks_status(vec![TaskStatus { - status: Some(task_status::Status::Failed(FailedTask { - error: "error".to_owned(), - })), - task_id: Some(task_id.clone()), - }]); - let ret = stage_manager.get_stage_tasks(job_id, stage_id); - assert!(ret.is_some()); - match ret - .unwrap() - .get(task_id.partition_id as usize) - .unwrap() - .status - .as_ref() - .unwrap() - { - task_status::Status::Failed(_) => (), - _ => panic!("Unexpected status"), - } - } - - Ok(()) - } - - #[tokio::test] - async fn test_task_status_state_machine_completed() -> Result<()> { - let stage_manager = StageManager::new(); - - let num_partitions = 3; - let job_id = "job"; - let stage_id = 1u32; - - stage_manager.add_running_stage(job_id, stage_id, num_partitions); - - let task_id = PartitionId { - job_id: job_id.to_owned(), - stage_id, - partition_id: 2, - }; - - // Valid transformation from Pending to Running to Completed to Pending - task_from_pending_to_completed(&stage_manager, &task_id); - let ret = stage_manager.get_stage_tasks(job_id, stage_id); - assert!(ret.is_some()); - match ret - .unwrap() - .get(task_id.partition_id as usize) - .unwrap() - .status - .as_ref() - .unwrap() - { - task_status::Status::Completed(_) => (), - _ => panic!("Unexpected status"), - } - stage_manager.update_tasks_status(vec![TaskStatus { - status: None, - task_id: Some(task_id.clone()), - }]); - let ret = stage_manager.get_stage_tasks(job_id, stage_id); - assert!(ret.is_some()); - assert!(ret - .unwrap() - .get(task_id.partition_id as usize) - .unwrap() - .status - .is_none()); - - Ok(()) - } - - #[tokio::test] - async fn test_stage_state_machine_completed() -> Result<()> { - let stage_manager = StageManager::new(); - - let num_partitions = 3; - let job_id = "job"; - let stage_id = 1u32; - - // Valid transformation from Running to Completed - stage_manager.add_running_stage(job_id, stage_id, num_partitions); - assert!(stage_manager.is_running_stage(job_id, stage_id)); - for partition_id in 0..num_partitions { - task_from_pending_to_completed( - &stage_manager, - &PartitionId { - job_id: job_id.to_owned(), - stage_id, - partition_id, - }, - ); - } - assert!(stage_manager.is_completed_stage(job_id, stage_id)); - - // Valid transformation from Completed to Running - stage_manager.update_tasks_status(vec![TaskStatus { - status: None, - task_id: Some(PartitionId { - job_id: job_id.to_owned(), - stage_id, - partition_id: 0, - }), - }]); - assert!(!stage_manager.is_running_stage(job_id, stage_id)); - - Ok(()) - } - - fn task_from_pending_to_completed( - stage_manager: &StageManager, - task_id: &PartitionId, - ) { - stage_manager.update_tasks_status(vec![TaskStatus { - status: Some(task_status::Status::Running(RunningTask { - executor_id: "localhost".to_owned(), - })), - task_id: Some(task_id.clone()), - }]); - stage_manager.update_tasks_status(vec![TaskStatus { - status: Some(task_status::Status::Completed(CompletedTask { - executor_id: "localhost".to_owned(), - partitions: Vec::new(), - })), - task_id: Some(task_id.clone()), - }]); - } -} diff --git a/ballista/rust/scheduler/src/state/task_scheduler.rs b/ballista/rust/scheduler/src/state/task_scheduler.rs deleted file mode 100644 index 132fbcf1c695..000000000000 --- a/ballista/rust/scheduler/src/state/task_scheduler.rs +++ /dev/null @@ -1,211 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use crate::state::stage_manager::StageKey; -use crate::state::SchedulerState; -use async_trait::async_trait; -use ballista_core::error::BallistaError; -use ballista_core::execution_plans::ShuffleWriterExec; -use ballista_core::serde::protobuf::{ - job_status, task_status, FailedJob, KeyValuePair, RunningTask, TaskDefinition, - TaskStatus, -}; -use ballista_core::serde::scheduler::to_proto::hash_partitioning_to_proto; -use ballista_core::serde::scheduler::{ExecutorData, PartitionId}; -use ballista_core::serde::{AsExecutionPlan, AsLogicalPlan}; -use log::{debug, info}; - -#[async_trait] -pub trait TaskScheduler { - // For each round, it will fetch tasks from one stage - async fn fetch_schedulable_tasks( - &self, - available_executors: &mut [ExecutorData], - n_round: u32, - ) -> Result<(Vec>, usize), BallistaError>; -} - -pub trait StageScheduler { - fn fetch_schedulable_stage(&self, cond: F) -> Option - where - F: Fn(&StageKey) -> bool; -} - -#[async_trait] -impl TaskScheduler - for SchedulerState -{ - async fn fetch_schedulable_tasks( - &self, - available_executors: &mut [ExecutorData], - n_round: u32, - ) -> Result<(Vec>, usize), BallistaError> { - let mut ret: Vec> = - Vec::with_capacity(available_executors.len()); - let mut max_task_num = 0u32; - for executor in available_executors.iter() { - ret.push(Vec::new()); - max_task_num += executor.available_task_slots; - } - - let mut tasks_status = vec![]; - let mut has_resources = true; - for i in 0..n_round { - if !has_resources { - break; - } - let mut num_tasks = 0; - // For each round, it will fetch tasks from one stage - if let Some((job_id, stage_id, tasks)) = - self.stage_manager.fetch_pending_tasks( - max_task_num as usize - tasks_status.len(), - |stage_key| { - // Don't scheduler stages for jobs with error status - if let Some(job_meta) = self.get_job_metadata(&stage_key.0) { - if !matches!( - &job_meta.status, - Some(job_status::Status::Failed(FailedJob { error: _ })) - ) { - true - } else { - info!("Stage {}/{} not to be scheduled due to its job failed", stage_key.0, stage_key.1); - false - } - } else { - false - } - }, - ) - { - let plan = - self.get_stage_plan(&job_id, stage_id as usize) - .ok_or_else(|| { - BallistaError::General(format!( - "Fail to find execution plan for stage {}/{}", - job_id, stage_id - )) - })?; - loop { - debug!("Go inside fetching task loop for stage {}/{}", job_id, stage_id); - - let mut has_tasks = true; - for (idx, executor) in available_executors.iter_mut().enumerate() { - if executor.available_task_slots == 0 { - has_resources = false; - break; - } - - if num_tasks >= tasks.len() { - has_tasks = false; - break; - } - - let task_id = PartitionId { - job_id: job_id.clone(), - stage_id: stage_id as usize, - partition_id: tasks[num_tasks] as usize, - }; - - let task_id = Some(task_id.into()); - let running_task = TaskStatus { - task_id: task_id.clone(), - status: Some(task_status::Status::Running(RunningTask { - executor_id: executor.executor_id.to_owned(), - })), - }; - tasks_status.push(running_task); - - let plan_clone = plan.clone(); - let output_partitioning = if let Some(shuffle_writer) = - plan_clone.as_any().downcast_ref::() - { - shuffle_writer.shuffle_output_partitioning() - } else { - return Err(BallistaError::General(format!( - "Task root plan was not a ShuffleWriterExec: {:?}", - plan_clone - ))); - }; - - let mut buf: Vec = vec![]; - U::try_from_physical_plan( - plan.clone(), - self.get_codec().physical_extension_codec(), - ) - .and_then(|m| m.try_encode(&mut buf)) - .map_err(|e| { - tonic::Status::internal(format!( - "error serializing execution plan: {:?}", - e - )) - })?; - - let session_id = self.get_session_from_job(&job_id).expect("session id does not exist for job"); - let session_props = self - .session_registry() - .lookup_session(&session_id) - .await - .expect("SessionContext does not exist in SessionContextRegistry.") - .copied_config() - .to_props(); - let task_props = session_props - .iter() - .map(|(k, v)| KeyValuePair { - key: k.to_owned(), - value: v.to_owned(), - }) - .collect::>(); - - ret[idx].push(TaskDefinition { - plan: buf, - task_id, - output_partitioning: hash_partitioning_to_proto( - output_partitioning, - ) - .map_err(|_| tonic::Status::internal("TBD".to_string()))?, - session_id, - props: task_props, - }); - executor.available_task_slots -= 1; - num_tasks += 1; - } - if !has_tasks { - break; - } - if !has_resources { - break; - } - } - } - if !has_resources { - info!( - "Not enough resource for task running. Stopped at round {}", - i - ); - break; - } - } - - let total_task_num = tasks_status.len(); - debug!("{} tasks to be scheduled", total_task_num); - - // No need to deal with the stage event, since the task status is changing from pending to running - self.stage_manager.update_tasks_status(tasks_status); - - Ok((ret, total_task_num)) - } -} diff --git a/ballista/rust/scheduler/src/test_utils.rs b/ballista/rust/scheduler/src/test_utils.rs deleted file mode 100644 index 3f8b57670834..000000000000 --- a/ballista/rust/scheduler/src/test_utils.rs +++ /dev/null @@ -1,138 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use ballista_core::error::Result; - -use datafusion::arrow::datatypes::{DataType, Field, Schema}; -use datafusion::execution::context::{SessionConfig, SessionContext}; -use datafusion::prelude::CsvReadOptions; - -pub const TPCH_TABLES: &[&str] = &[ - "part", "supplier", "partsupp", "customer", "orders", "lineitem", "nation", "region", -]; - -pub async fn datafusion_test_context(path: &str) -> Result { - let default_shuffle_partitions = 2; - let config = SessionConfig::new().with_target_partitions(default_shuffle_partitions); - let ctx = SessionContext::with_config(config); - for table in TPCH_TABLES { - let schema = get_tpch_schema(table); - let options = CsvReadOptions::new() - .schema(&schema) - .delimiter(b'|') - .has_header(false) - .file_extension(".tbl"); - let dir = format!("{}/{}", path, table); - ctx.register_csv(table, &dir, options).await?; - } - Ok(ctx) -} - -pub fn get_tpch_schema(table: &str) -> Schema { - // note that the schema intentionally uses signed integers so that any generated Parquet - // files can also be used to benchmark tools that only support signed integers, such as - // Apache Spark - - match table { - "part" => Schema::new(vec![ - Field::new("p_partkey", DataType::Int32, false), - Field::new("p_name", DataType::Utf8, false), - Field::new("p_mfgr", DataType::Utf8, false), - Field::new("p_brand", DataType::Utf8, false), - Field::new("p_type", DataType::Utf8, false), - Field::new("p_size", DataType::Int32, false), - Field::new("p_container", DataType::Utf8, false), - Field::new("p_retailprice", DataType::Float64, false), - Field::new("p_comment", DataType::Utf8, false), - ]), - - "supplier" => Schema::new(vec![ - Field::new("s_suppkey", DataType::Int32, false), - Field::new("s_name", DataType::Utf8, false), - Field::new("s_address", DataType::Utf8, false), - Field::new("s_nationkey", DataType::Int32, false), - Field::new("s_phone", DataType::Utf8, false), - Field::new("s_acctbal", DataType::Float64, false), - Field::new("s_comment", DataType::Utf8, false), - ]), - - "partsupp" => Schema::new(vec![ - Field::new("ps_partkey", DataType::Int32, false), - Field::new("ps_suppkey", DataType::Int32, false), - Field::new("ps_availqty", DataType::Int32, false), - Field::new("ps_supplycost", DataType::Float64, false), - Field::new("ps_comment", DataType::Utf8, false), - ]), - - "customer" => Schema::new(vec![ - Field::new("c_custkey", DataType::Int32, false), - Field::new("c_name", DataType::Utf8, false), - Field::new("c_address", DataType::Utf8, false), - Field::new("c_nationkey", DataType::Int32, false), - Field::new("c_phone", DataType::Utf8, false), - Field::new("c_acctbal", DataType::Float64, false), - Field::new("c_mktsegment", DataType::Utf8, false), - Field::new("c_comment", DataType::Utf8, false), - ]), - - "orders" => Schema::new(vec![ - Field::new("o_orderkey", DataType::Int32, false), - Field::new("o_custkey", DataType::Int32, false), - Field::new("o_orderstatus", DataType::Utf8, false), - Field::new("o_totalprice", DataType::Float64, false), - Field::new("o_orderdate", DataType::Date32, false), - Field::new("o_orderpriority", DataType::Utf8, false), - Field::new("o_clerk", DataType::Utf8, false), - Field::new("o_shippriority", DataType::Int32, false), - Field::new("o_comment", DataType::Utf8, false), - ]), - - "lineitem" => Schema::new(vec![ - Field::new("l_orderkey", DataType::Int32, false), - Field::new("l_partkey", DataType::Int32, false), - Field::new("l_suppkey", DataType::Int32, false), - Field::new("l_linenumber", DataType::Int32, false), - Field::new("l_quantity", DataType::Float64, false), - Field::new("l_extendedprice", DataType::Float64, false), - Field::new("l_discount", DataType::Float64, false), - Field::new("l_tax", DataType::Float64, false), - Field::new("l_returnflag", DataType::Utf8, false), - Field::new("l_linestatus", DataType::Utf8, false), - Field::new("l_shipdate", DataType::Date32, false), - Field::new("l_commitdate", DataType::Date32, false), - Field::new("l_receiptdate", DataType::Date32, false), - Field::new("l_shipinstruct", DataType::Utf8, false), - Field::new("l_shipmode", DataType::Utf8, false), - Field::new("l_comment", DataType::Utf8, false), - ]), - - "nation" => Schema::new(vec![ - Field::new("n_nationkey", DataType::Int32, false), - Field::new("n_name", DataType::Utf8, false), - Field::new("n_regionkey", DataType::Int32, false), - Field::new("n_comment", DataType::Utf8, false), - ]), - - "region" => Schema::new(vec![ - Field::new("r_regionkey", DataType::Int32, false), - Field::new("r_name", DataType::Utf8, false), - Field::new("r_comment", DataType::Utf8, false), - ]), - - _ => unimplemented!(), - } -} diff --git a/ballista/rust/scheduler/testdata/customer/customer.tbl b/ballista/rust/scheduler/testdata/customer/customer.tbl deleted file mode 100644 index afa5a739ab32..000000000000 --- a/ballista/rust/scheduler/testdata/customer/customer.tbl +++ /dev/null @@ -1,10 +0,0 @@ -1|Customer#000000001|IVhzIApeRb ot,c,E|15|25-989-741-2988|711.56|BUILDING|to the even, regular platelets. regular, ironic epitaphs nag e| -2|Customer#000000002|XSTf4,NCwDVaWNe6tEgvwfmRchLXak|13|23-768-687-3665|121.65|AUTOMOBILE|l accounts. blithely ironic theodolites integrate boldly: caref| -3|Customer#000000003|MG9kdTD2WBHm|1|11-719-748-3364|7498.12|AUTOMOBILE| deposits eat slyly ironic, even instructions. express foxes detect slyly. blithely even accounts abov| -4|Customer#000000004|XxVSJsLAGtn|4|14-128-190-5944|2866.83|MACHINERY| requests. final, regular ideas sleep final accou| -5|Customer#000000005|KvpyuHCplrB84WgAiGV6sYpZq7Tj|3|13-750-942-6364|794.47|HOUSEHOLD|n accounts will have to unwind. foxes cajole accor| -6|Customer#000000006|sKZz0CsnMD7mp4Xd0YrBvx,LREYKUWAh yVn|20|30-114-968-4951|7638.57|AUTOMOBILE|tions. even deposits boost according to the slyly bold packages. final accounts cajole requests. furious| -7|Customer#000000007|TcGe5gaZNgVePxU5kRrvXBfkasDTea|18|28-190-982-9759|9561.95|AUTOMOBILE|ainst the ironic, express theodolites. express, even pinto beans among the exp| -8|Customer#000000008|I0B10bB0AymmC, 0PrRYBCP1yGJ8xcBPmWhl5|17|27-147-574-9335|6819.74|BUILDING|among the slyly regular theodolites kindle blithely courts. carefully even theodolites haggle slyly along the ide| -9|Customer#000000009|xKiAFTjUsCuxfeleNqefumTrjS|8|18-338-906-3675|8324.07|FURNITURE|r theodolites according to the requests wake thinly excuses: pending requests haggle furiousl| -10|Customer#000000010|6LrEaV6KR6PLVcgl2ArL Q3rqzLzcT1 v2|5|15-741-346-9870|2753.54|HOUSEHOLD|es regular deposits haggle. fur| diff --git a/ballista/rust/scheduler/testdata/lineitem/partition0.tbl b/ballista/rust/scheduler/testdata/lineitem/partition0.tbl deleted file mode 100644 index b7424c2138b7..000000000000 --- a/ballista/rust/scheduler/testdata/lineitem/partition0.tbl +++ /dev/null @@ -1,10 +0,0 @@ -1|155190|7706|1|17|21168.23|0.04|0.02|N|O|1996-03-13|1996-02-12|1996-03-22|DELIVER IN PERSON|TRUCK|egular courts above the| -1|67310|7311|2|36|45983.16|0.09|0.06|N|O|1996-04-12|1996-02-28|1996-04-20|TAKE BACK RETURN|MAIL|ly final dependencies: slyly bold | -1|63700|3701|3|8|13309.60|0.10|0.02|N|O|1996-01-29|1996-03-05|1996-01-31|TAKE BACK RETURN|REG AIR|riously. regular, express dep| -1|2132|4633|4|28|28955.64|0.09|0.06|N|O|1996-04-21|1996-03-30|1996-05-16|NONE|AIR|lites. fluffily even de| -1|24027|1534|5|24|22824.48|0.10|0.04|N|O|1996-03-30|1996-03-14|1996-04-01|NONE|FOB| pending foxes. slyly re| -1|15635|638|6|32|49620.16|0.07|0.02|N|O|1996-01-30|1996-02-07|1996-02-03|DELIVER IN PERSON|MAIL|arefully slyly ex| -2|106170|1191|1|38|44694.46|0.00|0.05|N|O|1997-01-28|1997-01-14|1997-02-02|TAKE BACK RETURN|RAIL|ven requests. deposits breach a| -3|4297|1798|1|45|54058.05|0.06|0.00|R|F|1994-02-02|1994-01-04|1994-02-23|NONE|AIR|ongside of the furiously brave acco| -3|19036|6540|2|49|46796.47|0.10|0.00|R|F|1993-11-09|1993-12-20|1993-11-24|TAKE BACK RETURN|RAIL| unusual accounts. eve| -3|128449|3474|3|27|39890.88|0.06|0.07|A|F|1994-01-16|1993-11-22|1994-01-23|DELIVER IN PERSON|SHIP|nal foxes wake. | diff --git a/ballista/rust/scheduler/testdata/lineitem/partition1.tbl b/ballista/rust/scheduler/testdata/lineitem/partition1.tbl deleted file mode 100644 index b7424c2138b7..000000000000 --- a/ballista/rust/scheduler/testdata/lineitem/partition1.tbl +++ /dev/null @@ -1,10 +0,0 @@ -1|155190|7706|1|17|21168.23|0.04|0.02|N|O|1996-03-13|1996-02-12|1996-03-22|DELIVER IN PERSON|TRUCK|egular courts above the| -1|67310|7311|2|36|45983.16|0.09|0.06|N|O|1996-04-12|1996-02-28|1996-04-20|TAKE BACK RETURN|MAIL|ly final dependencies: slyly bold | -1|63700|3701|3|8|13309.60|0.10|0.02|N|O|1996-01-29|1996-03-05|1996-01-31|TAKE BACK RETURN|REG AIR|riously. regular, express dep| -1|2132|4633|4|28|28955.64|0.09|0.06|N|O|1996-04-21|1996-03-30|1996-05-16|NONE|AIR|lites. fluffily even de| -1|24027|1534|5|24|22824.48|0.10|0.04|N|O|1996-03-30|1996-03-14|1996-04-01|NONE|FOB| pending foxes. slyly re| -1|15635|638|6|32|49620.16|0.07|0.02|N|O|1996-01-30|1996-02-07|1996-02-03|DELIVER IN PERSON|MAIL|arefully slyly ex| -2|106170|1191|1|38|44694.46|0.00|0.05|N|O|1997-01-28|1997-01-14|1997-02-02|TAKE BACK RETURN|RAIL|ven requests. deposits breach a| -3|4297|1798|1|45|54058.05|0.06|0.00|R|F|1994-02-02|1994-01-04|1994-02-23|NONE|AIR|ongside of the furiously brave acco| -3|19036|6540|2|49|46796.47|0.10|0.00|R|F|1993-11-09|1993-12-20|1993-11-24|TAKE BACK RETURN|RAIL| unusual accounts. eve| -3|128449|3474|3|27|39890.88|0.06|0.07|A|F|1994-01-16|1993-11-22|1994-01-23|DELIVER IN PERSON|SHIP|nal foxes wake. | diff --git a/ballista/rust/scheduler/testdata/nation/nation.tbl b/ballista/rust/scheduler/testdata/nation/nation.tbl deleted file mode 100644 index c31ad6be0fac..000000000000 --- a/ballista/rust/scheduler/testdata/nation/nation.tbl +++ /dev/null @@ -1,10 +0,0 @@ -0|ALGERIA|0| haggle. carefully final deposits detect slyly agai| -1|ARGENTINA|1|al foxes promise slyly according to the regular accounts. bold requests alon| -2|BRAZIL|1|y alongside of the pending deposits. carefully special packages are about the ironic forges. slyly special | -3|CANADA|1|eas hang ironic, silent packages. slyly regular packages are furiously over the tithes. fluffily bold| -4|EGYPT|4|y above the carefully unusual theodolites. final dugouts are quickly across the furiously regular d| -5|ETHIOPIA|0|ven packages wake quickly. regu| -6|FRANCE|3|refully final requests. regular, ironi| -7|GERMANY|3|l platelets. regular accounts x-ray: unusual, regular acco| -8|INDIA|2|ss excuses cajole slyly across the packages. deposits print aroun| -9|INDONESIA|2| slyly express asymptotes. regular deposits haggle slyly. carefully ironic hockey players sleep blithely. carefull| diff --git a/ballista/rust/scheduler/testdata/orders/orders.tbl b/ballista/rust/scheduler/testdata/orders/orders.tbl deleted file mode 100644 index f5fa65b09a7a..000000000000 --- a/ballista/rust/scheduler/testdata/orders/orders.tbl +++ /dev/null @@ -1,10 +0,0 @@ -1|36901|O|173665.47|1996-01-02|5-LOW|Clerk#000000951|0|nstructions sleep furiously among | -2|78002|O|46929.18|1996-12-01|1-URGENT|Clerk#000000880|0| foxes. pending accounts at the pending, silent asymptot| -3|123314|F|193846.25|1993-10-14|5-LOW|Clerk#000000955|0|sly final accounts boost. carefully regular ideas cajole carefully. depos| -4|136777|O|32151.78|1995-10-11|5-LOW|Clerk#000000124|0|sits. slyly regular warthogs cajole. regular, regular theodolites acro| -5|44485|F|144659.20|1994-07-30|5-LOW|Clerk#000000925|0|quickly. bold deposits sleep slyly. packages use slyly| -6|55624|F|58749.59|1992-02-21|4-NOT SPECIFIED|Clerk#000000058|0|ggle. special, final requests are against the furiously specia| -7|39136|O|252004.18|1996-01-10|2-HIGH|Clerk#000000470|0|ly special requests | -32|130057|O|208660.75|1995-07-16|2-HIGH|Clerk#000000616|0|ise blithely bold, regular requests. quickly unusual dep| -33|66958|F|163243.98|1993-10-27|3-MEDIUM|Clerk#000000409|0|uriously. furiously final request| -34|61001|O|58949.67|1998-07-21|3-MEDIUM|Clerk#000000223|0|ly final packages. fluffily final deposits wake blithely ideas. spe| diff --git a/ballista/rust/scheduler/testdata/part/part.tbl b/ballista/rust/scheduler/testdata/part/part.tbl deleted file mode 100644 index 0c6f0e2f3e10..000000000000 --- a/ballista/rust/scheduler/testdata/part/part.tbl +++ /dev/null @@ -1,10 +0,0 @@ -1|goldenrod lavender spring chocolate lace|Manufacturer#1|Brand#13|PROMO BURNISHED COPPER|7|JUMBO PKG|901.00|ly. slyly ironi| -2|blush thistle blue yellow saddle|Manufacturer#1|Brand#13|LARGE BRUSHED BRASS|1|LG CASE|902.00|lar accounts amo| -3|spring green yellow purple cornsilk|Manufacturer#4|Brand#42|STANDARD POLISHED BRASS|21|WRAP CASE|903.00|egular deposits hag| -4|cornflower chocolate smoke green pink|Manufacturer#3|Brand#34|SMALL PLATED BRASS|14|MED DRUM|904.00|p furiously r| -5|forest brown coral puff cream|Manufacturer#3|Brand#32|STANDARD POLISHED TIN|15|SM PKG|905.00| wake carefully | -6|bisque cornflower lawn forest magenta|Manufacturer#2|Brand#24|PROMO PLATED STEEL|4|MED BAG|906.00|sual a| -7|moccasin green thistle khaki floral|Manufacturer#1|Brand#11|SMALL PLATED COPPER|45|SM BAG|907.00|lyly. ex| -8|misty lace thistle snow royal|Manufacturer#4|Brand#44|PROMO BURNISHED TIN|41|LG DRUM|908.00|eposi| -9|thistle dim navajo dark gainsboro|Manufacturer#4|Brand#43|SMALL BURNISHED STEEL|12|WRAP CASE|909.00|ironic foxe| -10|linen pink saddle puff powder|Manufacturer#5|Brand#54|LARGE BURNISHED STEEL|44|LG CAN|910.01|ithely final deposit| diff --git a/ballista/rust/scheduler/testdata/partsupp/partsupp.tbl b/ballista/rust/scheduler/testdata/partsupp/partsupp.tbl deleted file mode 100644 index 45145385a165..000000000000 --- a/ballista/rust/scheduler/testdata/partsupp/partsupp.tbl +++ /dev/null @@ -1,10 +0,0 @@ -1|2|3325|771.64|, even theodolites. regular, final theodolites eat after the carefully pending foxes. furiously regular deposits sleep slyly. carefully bold realms above the ironic dependencies haggle careful| -1|2502|8076|993.49|ven ideas. quickly even packages print. pending multipliers must have to are fluff| -1|5002|3956|337.09|after the fluffily ironic deposits? blithely special dependencies integrate furiously even excuses. blithely silent theodolites could have to haggle pending, express requests; fu| -1|7502|4069|357.84|al, regular dependencies serve carefully after the quickly final pinto beans. furiously even deposits sleep quickly final, silent pinto beans. fluffily reg| -2|3|8895|378.49|nic accounts. final accounts sleep furiously about the ironic, bold packages. regular, regular accounts| -2|2503|4969|915.27|ptotes. quickly pending dependencies integrate furiously. fluffily ironic ideas impress blithely above the express accounts. furiously even epitaphs need to wak| -2|5003|8539|438.37|blithely bold ideas. furiously stealthy packages sleep fluffily. slyly special deposits snooze furiously carefully regular accounts. regular deposits according to the accounts nag carefully slyl| -2|7503|3025|306.39|olites. deposits wake carefully. even, express requests cajole. carefully regular ex| -3|4|4651|920.92|ilent foxes affix furiously quickly unusual requests. even packages across the carefully even theodolites nag above the sp| -3|2504|4093|498.13|ending dependencies haggle fluffily. regular deposits boost quickly carefully regular requests. deposits affix furiously around the pinto beans. ironic, unusual platelets across the p| diff --git a/ballista/rust/scheduler/testdata/region/region.tbl b/ballista/rust/scheduler/testdata/region/region.tbl deleted file mode 100644 index c5ebb63b621f..000000000000 --- a/ballista/rust/scheduler/testdata/region/region.tbl +++ /dev/null @@ -1,5 +0,0 @@ -0|AFRICA|lar deposits. blithely final packages cajole. regular waters are final requests. regular accounts are according to | -1|AMERICA|hs use ironic, even requests. s| -2|ASIA|ges. thinly even pinto beans ca| -3|EUROPE|ly final courts cajole furiously final excuse| -4|MIDDLE EAST|uickly special accounts cajole carefully blithely close requests. carefully final asymptotes haggle furiousl| diff --git a/ballista/rust/scheduler/testdata/supplier/supplier.tbl b/ballista/rust/scheduler/testdata/supplier/supplier.tbl deleted file mode 100644 index d9c0e9f7e201..000000000000 --- a/ballista/rust/scheduler/testdata/supplier/supplier.tbl +++ /dev/null @@ -1,10 +0,0 @@ -1|Supplier#000000001| N kD4on9OM Ipw3,gf0JBoQDd7tgrzrddZ|17|27-918-335-1736|5755.94|each slyly above the careful| -2|Supplier#000000002|89eJ5ksX3ImxJQBvxObC,|5|15-679-861-2259|4032.68| slyly bold instructions. idle dependen| -3|Supplier#000000003|q1,G3Pj6OjIuUYfUoH18BFTKP5aU9bEV3|1|11-383-516-1199|4192.40|blithely silent requests after the express dependencies are sl| -4|Supplier#000000004|Bk7ah4CK8SYQTepEmvMkkgMwg|15|25-843-787-7479|4641.08|riously even requests above the exp| -5|Supplier#000000005|Gcdm2rJRzl5qlTVzc|11|21-151-690-3663|-283.84|. slyly regular pinto bea| -6|Supplier#000000006|tQxuVm7s7CnK|14|24-696-997-4969|1365.79|final accounts. regular dolphins use against the furiously ironic decoys. | -7|Supplier#000000007|s,4TicNGB4uO6PaSqNBUq|23|33-990-965-2201|6820.35|s unwind silently furiously regular courts. final requests are deposits. requests wake quietly blit| -8|Supplier#000000008|9Sq4bBH2FQEmaFOocY45sRTxo6yuoG|17|27-498-742-3860|7627.85|al pinto beans. asymptotes haggl| -9|Supplier#000000009|1KhUgZegwM3ua7dsYmekYBsK|10|20-403-398-8662|5302.37|s. unusual, even requests along the furiously regular pac| -10|Supplier#000000010|Saygah3gYWMp72i PY|24|34-852-489-8585|3891.91|ing waters. regular requests ar| diff --git a/ballista/ui/scheduler/.dockerignore b/ballista/ui/scheduler/.dockerignore deleted file mode 100644 index dd87e2d73f9f..000000000000 --- a/ballista/ui/scheduler/.dockerignore +++ /dev/null @@ -1,2 +0,0 @@ -node_modules -build diff --git a/ballista/ui/scheduler/.gitignore b/ballista/ui/scheduler/.gitignore deleted file mode 100644 index 4d29575de804..000000000000 --- a/ballista/ui/scheduler/.gitignore +++ /dev/null @@ -1,23 +0,0 @@ -# See https://help.github.com/articles/ignoring-files/ for more about ignoring files. - -# dependencies -/node_modules -/.pnp -.pnp.js - -# testing -/coverage - -# production -/build - -# misc -.DS_Store -.env.local -.env.development.local -.env.test.local -.env.production.local - -npm-debug.log* -yarn-debug.log* -yarn-error.log* diff --git a/ballista/ui/scheduler/README.md b/ballista/ui/scheduler/README.md deleted file mode 100644 index 9f93143bb5f1..000000000000 --- a/ballista/ui/scheduler/README.md +++ /dev/null @@ -1,55 +0,0 @@ - - -# Ballista UI - -## Start project from source - -### Run scheduler/executor - -First, run scheduler from project: - -```shell -$ cd rust/scheduler -$ RUST_LOG=info cargo run --release -... - Finished release [optimized] target(s) in 11.92s - Running `/path-to-project/target/release/ballista-scheduler` -``` - -and run executor in new terminal: - -```shell -$ cd rust/executor -$ RUST_LOG=info cargo run --release - Finished release [optimized] target(s) in 0.09s - Running `/path-to-project/target/release/ballista-executor` -``` - -### Run Client project - -```shell -$ cd ui/scheduler -$ yarn -Resolving packages... -$ yarn start -Starting the development server... -``` - -Now access to http://localhost:3000/ diff --git a/ballista/ui/scheduler/index.d.ts b/ballista/ui/scheduler/index.d.ts deleted file mode 100644 index be0f38ba098b..000000000000 --- a/ballista/ui/scheduler/index.d.ts +++ /dev/null @@ -1,18 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -declare module "@chakra-ui/icons"; diff --git a/ballista/ui/scheduler/package.json b/ballista/ui/scheduler/package.json deleted file mode 100644 index 9a5d2006dcce..000000000000 --- a/ballista/ui/scheduler/package.json +++ /dev/null @@ -1,58 +0,0 @@ -{ - "name": "scheduler-ui", - "version": "0.1.0", - "private": true, - "dependencies": { - "@chakra-ui/icons": "^1.0.5", - "@chakra-ui/react": "^1.3.3", - "@emotion/react": "^11.1.5", - "@emotion/styled": "^11.1.5", - "@testing-library/jest-dom": "^5.11.4", - "@testing-library/react": "^11.1.0", - "@testing-library/user-event": "^12.1.10", - "@types/jest": "^26.0.15", - "@types/node": "^12.0.0", - "@types/react": "^17.0.0", - "@types/react-dom": "^17.0.0", - "framer-motion": "^3.7.0", - "react": "^17.0.1", - "react-dom": "^17.0.1", - "react-icons": "^4.2.0", - "react-router-dom": "^5.2.0", - "react-scripts": "4.0.3", - "react-table": "^7.6.3", - "react-timeago": "^5.2.0", - "typescript": "^4.1.2", - "web-vitals": "^1.0.1" - }, - "scripts": { - "start": "react-scripts start", - "build": "react-scripts build", - "test": "react-scripts test", - "eject": "react-scripts eject" - }, - "eslintConfig": { - "extends": [ - "react-app", - "react-app/jest" - ] - }, - "browserslist": { - "production": [ - ">0.2%", - "not dead", - "not op_mini all" - ], - "development": [ - "last 1 chrome version", - "last 1 firefox version", - "last 1 safari version" - ] - }, - "devDependencies": { - "@types/react-table": "^7.0.28", - "@types/react-timeago": "^4.1.2", - "prettier": "^2.3.0" - }, - "proxy": "http://localhost:50050" -} diff --git a/ballista/ui/scheduler/public/favicon.ico b/ballista/ui/scheduler/public/favicon.ico deleted file mode 100644 index a11777cc471a..000000000000 Binary files a/ballista/ui/scheduler/public/favicon.ico and /dev/null differ diff --git a/ballista/ui/scheduler/public/index.html b/ballista/ui/scheduler/public/index.html deleted file mode 100644 index d902333f0342..000000000000 --- a/ballista/ui/scheduler/public/index.html +++ /dev/null @@ -1,62 +0,0 @@ - - - - - - - - - - - - - - - Ballista UI - - - -

- - - diff --git a/ballista/ui/scheduler/public/logo192.png b/ballista/ui/scheduler/public/logo192.png deleted file mode 100644 index fc44b0a3796c..000000000000 Binary files a/ballista/ui/scheduler/public/logo192.png and /dev/null differ diff --git a/ballista/ui/scheduler/public/logo512.png b/ballista/ui/scheduler/public/logo512.png deleted file mode 100644 index a4e47a6545bc..000000000000 Binary files a/ballista/ui/scheduler/public/logo512.png and /dev/null differ diff --git a/ballista/ui/scheduler/public/manifest.json b/ballista/ui/scheduler/public/manifest.json deleted file mode 100644 index 080d6c77ac21..000000000000 --- a/ballista/ui/scheduler/public/manifest.json +++ /dev/null @@ -1,25 +0,0 @@ -{ - "short_name": "React App", - "name": "Create React App Sample", - "icons": [ - { - "src": "favicon.ico", - "sizes": "64x64 32x32 24x24 16x16", - "type": "image/x-icon" - }, - { - "src": "logo192.png", - "type": "image/png", - "sizes": "192x192" - }, - { - "src": "logo512.png", - "type": "image/png", - "sizes": "512x512" - } - ], - "start_url": ".", - "display": "standalone", - "theme_color": "#000000", - "background_color": "#ffffff" -} diff --git a/ballista/ui/scheduler/public/robots.txt b/ballista/ui/scheduler/public/robots.txt deleted file mode 100644 index dc045698d092..000000000000 --- a/ballista/ui/scheduler/public/robots.txt +++ /dev/null @@ -1,20 +0,0 @@ -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. - -# https://www.robotstxt.org/robotstxt.html -User-agent: * -Disallow: diff --git a/ballista/ui/scheduler/react-table-config.d.ts b/ballista/ui/scheduler/react-table-config.d.ts deleted file mode 100644 index 2c9994f91d2a..000000000000 --- a/ballista/ui/scheduler/react-table-config.d.ts +++ /dev/null @@ -1,146 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -import { - UseColumnOrderInstanceProps, - UseColumnOrderState, - UseExpandedHooks, - UseExpandedInstanceProps, - UseExpandedOptions, - UseExpandedRowProps, - UseExpandedState, - UseFiltersColumnOptions, - UseFiltersColumnProps, - UseFiltersInstanceProps, - UseFiltersOptions, - UseFiltersState, - UseGlobalFiltersColumnOptions, - UseGlobalFiltersInstanceProps, - UseGlobalFiltersOptions, - UseGlobalFiltersState, - UseGroupByCellProps, - UseGroupByColumnOptions, - UseGroupByColumnProps, - UseGroupByHooks, - UseGroupByInstanceProps, - UseGroupByOptions, - UseGroupByRowProps, - UseGroupByState, - UsePaginationInstanceProps, - UsePaginationOptions, - UsePaginationState, - UseResizeColumnsColumnOptions, - UseResizeColumnsColumnProps, - UseResizeColumnsOptions, - UseResizeColumnsState, - UseRowSelectHooks, - UseRowSelectInstanceProps, - UseRowSelectOptions, - UseRowSelectRowProps, - UseRowSelectState, - UseRowStateCellProps, - UseRowStateInstanceProps, - UseRowStateOptions, - UseRowStateRowProps, - UseRowStateState, - UseSortByColumnOptions, - UseSortByColumnProps, - UseSortByHooks, - UseSortByInstanceProps, - UseSortByOptions, - UseSortByState, -} from "react-table"; - -declare module "react-table" { - // take this file as-is, or comment out the sections that don't apply to your plugin configuration - - export interface TableOptions< - D extends Record - > extends UseExpandedOptions, - UseFiltersOptions, - UseGlobalFiltersOptions, - UseGroupByOptions, - UsePaginationOptions, - UseResizeColumnsOptions, - UseRowSelectOptions, - UseRowStateOptions, - UseSortByOptions, - // note that having Record here allows you to add anything to the options, this matches the spirit of the - // underlying js library, but might be cleaner if it's replaced by a more specific type that matches your - // feature set, this is a safe default. - Record {} - - export interface Hooks< - D extends Record = Record - > extends UseExpandedHooks, - UseGroupByHooks, - UseRowSelectHooks, - UseSortByHooks {} - - export interface TableInstance< - D extends Record = Record - > extends UseColumnOrderInstanceProps, - UseExpandedInstanceProps, - UseFiltersInstanceProps, - UseGlobalFiltersInstanceProps, - UseGroupByInstanceProps, - UsePaginationInstanceProps, - UseRowSelectInstanceProps, - UseRowStateInstanceProps, - UseSortByInstanceProps {} - - export interface TableState< - D extends Record = Record - > extends UseColumnOrderState, - UseExpandedState, - UseFiltersState, - UseGlobalFiltersState, - UseGroupByState, - UsePaginationState, - UseResizeColumnsState, - UseRowSelectState, - UseRowStateState, - UseSortByState {} - - export interface ColumnInterface< - D extends Record = Record - > extends UseFiltersColumnOptions, - UseGlobalFiltersColumnOptions, - UseGroupByColumnOptions, - UseResizeColumnsColumnOptions, - UseSortByColumnOptions {} - - export interface ColumnInstance< - D extends Record = Record - > extends UseFiltersColumnProps, - UseGroupByColumnProps, - UseResizeColumnsColumnProps, - UseSortByColumnProps {} - - export interface Cell< - D extends Record = Record, - V = any - > extends UseGroupByCellProps, - UseRowStateCellProps {} - - export interface Row< - D extends Record = Record - > extends UseExpandedRowProps, - UseGroupByRowProps, - UseRowSelectRowProps, - UseRowStateRowProps {} -} diff --git a/ballista/ui/scheduler/src/App.css b/ballista/ui/scheduler/src/App.css deleted file mode 100644 index bea95535e9e6..000000000000 --- a/ballista/ui/scheduler/src/App.css +++ /dev/null @@ -1,18 +0,0 @@ -/* - Licensed to the Apache Software Foundation (ASF) under one - or more contributor license agreements. See the NOTICE file - distributed with this work for additional information - regarding copyright ownership. The ASF licenses this file - to you under the Apache License, Version 2.0 (the - "License"); you may not use this file except in compliance - with the License. You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, - software distributed under the License is distributed on an - "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - KIND, either express or implied. See the License for the - specific language governing permissions and limitations - under the License. -*/ diff --git a/ballista/ui/scheduler/src/App.test.tsx b/ballista/ui/scheduler/src/App.test.tsx deleted file mode 100644 index 20dca216eb27..000000000000 --- a/ballista/ui/scheduler/src/App.test.tsx +++ /dev/null @@ -1,26 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -import React from "react"; -import { render, screen } from "@testing-library/react"; -import App from "./App"; - -test("renders learn react link", () => { - render(); - const linkElement = screen.getByText(/learn react/i); - expect(linkElement).toBeInTheDocument(); -}); diff --git a/ballista/ui/scheduler/src/App.tsx b/ballista/ui/scheduler/src/App.tsx deleted file mode 100644 index adb5896a81d1..000000000000 --- a/ballista/ui/scheduler/src/App.tsx +++ /dev/null @@ -1,97 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -import React, { useState, useEffect } from "react"; -import { Box, Grid, VStack } from "@chakra-ui/react"; -import { Header } from "./components/Header"; -import { Summary } from "./components/Summary"; -import { QueriesList, Query, QueryStatus } from "./components/QueriesList"; -import { Footer } from "./components/Footer"; - -import "./App.css"; - -function uuidv4() { - return "xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx".replace(/[xy]/g, function (c) { - var r = (Math.random() * 16) | 0, - v = c === "x" ? r : (r & 0x3) | 0x8; - return v.toString(16); - }); -} - -const getRandomQueries = (num: number): Query[] => { - const nodes: Query[] = []; - - for (let i = 0; i < num; i++) { - nodes.push({ - started: new Date().toISOString(), - query: - "SELECT \n" + - " employee.id,\n" + - " employee.first_name,\n" + - " employee.last_name,\n" + - ' SUM(DATEDIFF("SECOND", call.start_time, call.end_time)) AS call_duration_sum\n' + - "FROM call\n" + - "INNER JOIN employee ON call.employee_id = employee.id\n" + - "GROUP BY\n" + - " employee.id,\n" + - " employee.first_name,\n" + - " employee.last_name\n" + - "ORDER BY\n" + - " employee.id ASC;", - status: QueryStatus.RUNNING, - progress: Math.round(Math.random() * 100), - uuid: uuidv4(), - }); - } - return nodes; -}; - -const queries = getRandomQueries(17); - -const App: React.FunctionComponent = () => { - const [schedulerState, setSchedulerState] = useState(undefined); - - function getSchedulerState() { - return fetch(`/state`, { - method: "POST", - headers: { - Accept: "application/json", - }, - }) - .then((res) => res.json()) - .then((res) => setSchedulerState(res)); - } - - useEffect(() => { - getSchedulerState(); - }, []); - - return ( - - - -
- - -