diff --git a/.github/workflows/extended.yml b/.github/workflows/extended.yml index 19910957a85b9..3f882d7a3a821 100644 --- a/.github/workflows/extended.yml +++ b/.github/workflows/extended.yml @@ -23,59 +23,67 @@ concurrency: # https://docs.github.com/en/actions/writing-workflows/choosing-when-your-workflow-runs/events-that-trigger-workflows#running-your-pull_request-workflow-when-a-pull-request-merges # -# These jobs only run on the `main` branch as they are time consuming +# These jobs are not run as part of PR checks as they are time-consuming # and should not fail often. # # We still run them as they provide important coverage to ensure correctness # in the (very rare) event of a hash failure or sqlite library query failure. on: - # Run on all commits to main push: - branches: - - main jobs: # Check crate compiles and base cargo check passes linux-build-lib: name: linux build test runs-on: ubuntu-latest - container: - image: amd64/rust steps: - uses: actions/checkout@v4 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder with: - rust-version: stable + submodules: true + fetch-depth: 1 + - name: Install Rust + run: | + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + source $HOME/.cargo/env + rustup default stable + - name: Install Protobuf Compiler + run: sudo apt-get install -y protobuf-compiler - name: Prepare cargo build run: | cargo check --profile ci --all-targets cargo clean -# # Run extended tests (with feature 'extended_tests') -# # Disabling as it is running out of disk space -# # see https://github.com/apache/datafusion/issues/14576 -# linux-test-extended: -# name: cargo test 'extended_tests' (amd64) -# needs: linux-build-lib -# runs-on: ubuntu-latest -# container: -# image: amd64/rust -# steps: -# - uses: actions/checkout@v4 -# with: -# submodules: true -# fetch-depth: 1 -# - name: Setup Rust toolchain -# uses: ./.github/actions/setup-builder -# with: -# rust-version: stable -# - name: Run tests (excluding doctests) -# run: cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --workspace --lib --tests --bins --features avro,json,backtrace,extended_tests -# - name: Verify Working Directory Clean -# run: git diff --exit-code -# - name: Cleanup -# run: cargo clean + # Run extended tests (with feature 'extended_tests') + linux-test-extended: + name: cargo test 'extended_tests' (amd64) + needs: linux-build-lib + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + submodules: true + fetch-depth: 1 + - name: Free Disk Space (Ubuntu) + uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be + - name: Install Rust + run: | + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + source $HOME/.cargo/env + rustup default stable + - name: Install Protobuf Compiler + run: sudo apt-get install -y protobuf-compiler + # For debugging, test binaries can be large. + - name: Show available disk space + run: | + df -h + - name: Run tests (excluding doctests) + env: + RUST_BACKTRACE: 1 + run: cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --workspace --lib --tests --bins --features avro,json,backtrace,extended_tests + - name: Verify Working Directory Clean + run: git diff --exit-code + - name: Cleanup + run: cargo clean # Check answers are correct when hash values collide hash-collisions: @@ -95,7 +103,7 @@ jobs: - name: Run tests run: | cd datafusion - cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-sqllogictest --workspace --lib --tests --features=force_hash_collisions,avro,extended_tests + cargo test --profile ci --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-sqllogictest --workspace --lib --tests --features=force_hash_collisions,avro cargo clean sqllogictest-sqlite: diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index a743d0e8fd07a..99aaa7d6f2907 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -389,21 +389,20 @@ jobs: test-datafusion-pyarrow: name: cargo test pyarrow (amd64) needs: linux-build-lib - runs-on: ubuntu-20.04 + runs-on: ubuntu-latest container: - image: amd64/rust:bullseye # Workaround https://github.com/actions/setup-python/issues/721 + image: amd64/rust:bullseye # Use the bullseye tag image which comes with python3.9 steps: - uses: actions/checkout@v4 with: submodules: true fetch-depth: 1 - - uses: actions/setup-python@v5 - with: - python-version: "3.8" - name: Install PyArrow run: | echo "LIBRARY_PATH=$LD_LIBRARY_PATH" >> $GITHUB_ENV - python -m pip install pyarrow + apt-get update + apt-get install python3-pip -y + python3 -m pip install pyarrow - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: diff --git a/Cargo.lock b/Cargo.lock index b7794d731b757..d5edabdf4585a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -674,9 +674,9 @@ dependencies = [ [[package]] name = "aws-sdk-sso" -version = "1.58.0" +version = "1.59.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16ff718c9ee45cc1ebd4774a0e086bb80a6ab752b4902edf1c9f56b86ee1f770" +checksum = "00a35fc7e74f5be45839eb753568535c074a592185dd0a2d406685018d581c43" dependencies = [ "aws-credential-types", "aws-runtime", @@ -696,9 +696,9 @@ dependencies = [ [[package]] name = "aws-sdk-ssooidc" -version = "1.59.0" +version = "1.60.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5183e088715cc135d8d396fdd3bc02f018f0da4c511f53cb8d795b6a31c55809" +checksum = "f8fa655b4f313124ce272cbc38c5fef13793c832279cec750103e5e6b71a54b8" dependencies = [ "aws-credential-types", "aws-runtime", @@ -718,9 +718,9 @@ dependencies = [ [[package]] name = "aws-sdk-sts" -version = "1.59.0" +version = "1.60.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c9f944ef032717596639cea4a2118a3a457268ef51bbb5fde9637e54c465da00" +checksum = "dc1cfe5e16b90421ea031f4c6348b534ef442e76f6bf4a1b2b592c12cc2c6af9" dependencies = [ "aws-credential-types", "aws-runtime", @@ -741,9 +741,9 @@ dependencies = [ [[package]] name = "aws-sigv4" -version = "1.2.8" +version = "1.2.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bc5bbd1e4a2648fd8c5982af03935972c24a2f9846b396de661d351ee3ce837" +checksum = "9bfe75fad52793ce6dec0dc3d4b1f388f038b5eb866c8d4d7f3a8e21b5ea5051" dependencies = [ "aws-credential-types", "aws-smithy-http", @@ -1078,7 +1078,7 @@ dependencies = [ "hyperlocal", "log", "pin-project-lite", - "rustls 0.23.22", + "rustls 0.23.23", "rustls-native-certs 0.8.1", "rustls-pemfile 2.2.0", "rustls-pki-types", @@ -1249,9 +1249,9 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.2.13" +version = "1.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7777341816418c02e033934a09f20dc0ccaf65a5201ef8a450ae0105a573fda" +checksum = "0c3d1b2e905a3a7b00a6141adb0e4c0bb941d11caf55349d863942a1cc44e3c9" dependencies = [ "jobserver", "libc", @@ -1643,9 +1643,9 @@ dependencies = [ [[package]] name = "csv-core" -version = "0.1.11" +version = "0.1.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5efa2b3d7902f4b634a20cae3c9c4e6209dc4779feb6863329607560143efa70" +checksum = "7d02f3b0da4c6504f86e9cd789d8dbafab48c2321be74e9987593de5a894d93d" dependencies = [ "memchr", ] @@ -1918,6 +1918,8 @@ dependencies = [ "datafusion-common-runtime", "datafusion-execution", "datafusion-expr", + "datafusion-physical-expr", + "datafusion-physical-expr-common", "datafusion-physical-plan", "flate2", "futures", @@ -2396,13 +2398,11 @@ dependencies = [ "datafusion-execution", "datafusion-expr", "datafusion-optimizer", - "datafusion-physical-expr", "datafusion-physical-plan", "datafusion-sql", "getrandom 0.2.15", "tokio", "wasm-bindgen", - "wasm-bindgen-futures", "wasm-bindgen-test", ] @@ -2563,9 +2563,9 @@ dependencies = [ [[package]] name = "equivalent" -version = "1.0.1" +version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" [[package]] name = "errno" @@ -2924,9 +2924,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.7" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ccae279728d634d083c00f6099cb58f01cc99c145b84b8be2f6c74618d79922e" +checksum = "5017294ff4bb30944501348f6f8e42e6ad28f42c8bbef7a74029aff064a4e3c2" dependencies = [ "atomic-waker", "bytes", @@ -3134,7 +3134,7 @@ dependencies = [ "bytes", "futures-channel", "futures-util", - "h2 0.4.7", + "h2 0.4.8", "http 1.2.0", "http-body 1.0.1", "httparse", @@ -3187,7 +3187,7 @@ dependencies = [ "http 1.2.0", "hyper 1.6.0", "hyper-util", - "rustls 0.23.22", + "rustls 0.23.23", "rustls-native-certs 0.8.1", "rustls-pki-types", "tokio", @@ -3796,9 +3796,9 @@ dependencies = [ [[package]] name = "miniz_oxide" -version = "0.8.3" +version = "0.8.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8402cab7aefae129c6977bb0ff1b8fd9a04eb5b51efc50a70bea51cda0c7924" +checksum = "b3b1c9bd4fe1f0f8b387f6eb9eb3b4a1aa26185e5750efb9140301703f62cd1b" dependencies = [ "adler2", ] @@ -4494,9 +4494,9 @@ dependencies = [ [[package]] name = "psm" -version = "0.1.24" +version = "0.1.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "200b9ff220857e53e184257720a14553b2f4aa02577d2ed9842d45d4b9654810" +checksum = "f58e5423e24c18cc840e1c98370b3993c6649cd1678b4d24318bcf0a083cbe88" dependencies = [ "cc", ] @@ -4611,7 +4611,7 @@ dependencies = [ "quinn-proto", "quinn-udp", "rustc-hash", - "rustls 0.23.22", + "rustls 0.23.23", "socket2", "thiserror 2.0.11", "tokio", @@ -4629,7 +4629,7 @@ dependencies = [ "rand 0.8.5", "ring", "rustc-hash", - "rustls 0.23.22", + "rustls 0.23.23", "rustls-pki-types", "slab", "thiserror 2.0.11", @@ -4640,9 +4640,9 @@ dependencies = [ [[package]] name = "quinn-udp" -version = "0.5.9" +version = "0.5.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c40286217b4ba3a71d644d752e6a0b71f13f1b6a2c5311acfcbe0c2418ed904" +checksum = "e46f3055866785f6b92bc6164b76be02ca8f2eb4b002c0354b28cf4c119e5944" dependencies = [ "cfg_aliases", "libc", @@ -4695,8 +4695,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" dependencies = [ "rand_chacha 0.9.0", - "rand_core 0.9.0", - "zerocopy 0.8.17", + "rand_core 0.9.1", + "zerocopy 0.8.18", ] [[package]] @@ -4716,7 +4716,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core 0.9.0", + "rand_core 0.9.1", ] [[package]] @@ -4730,12 +4730,12 @@ dependencies = [ [[package]] name = "rand_core" -version = "0.9.0" +version = "0.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b08f3c9802962f7e1b25113931d94f43ed9725bebc59db9d0c3e9a23b67e15ff" +checksum = "a88e0da7a2c97baa202165137c158d0a2e824ac465d13d81046727b34cb247d3" dependencies = [ "getrandom 0.3.1", - "zerocopy 0.8.17", + "zerocopy 0.8.18", ] [[package]] @@ -4896,7 +4896,7 @@ dependencies = [ "bytes", "futures-core", "futures-util", - "h2 0.4.7", + "h2 0.4.8", "http 1.2.0", "http-body 1.0.1", "http-body-util", @@ -4911,7 +4911,7 @@ dependencies = [ "percent-encoding", "pin-project-lite", "quinn", - "rustls 0.23.22", + "rustls 0.23.23", "rustls-native-certs 0.8.1", "rustls-pemfile 2.2.0", "rustls-pki-types", @@ -4934,15 +4934,14 @@ dependencies = [ [[package]] name = "ring" -version = "0.17.8" +version = "0.17.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c17fa4cb658e3583423e915b9f3acc01cceaee1860e33d59ebae66adc3a2dc0d" +checksum = "e75ec5e92c4d8aede845126adc388046234541629e76029599ed35a003c7ed24" dependencies = [ "cc", "cfg-if", "getrandom 0.2.15", "libc", - "spin", "untrusted", "windows-sys 0.52.0", ] @@ -5088,9 +5087,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.22" +version = "0.23.23" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9fb9263ab4eb695e42321db096e3b8fbd715a59b154d5c88d82db2175b681ba7" +checksum = "47796c98c480fce5406ef69d1c76378375492c3b0a0de587be0c1d9feb12f395" dependencies = [ "once_cell", "ring", @@ -5323,9 +5322,9 @@ checksum = "a3f0bf26fd526d2a95683cd0f87bf103b8539e2ca1ef48ce002d67aad59aa0b4" [[package]] name = "serde" -version = "1.0.217" +version = "1.0.218" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "02fc4265df13d6fa1d00ecff087228cc0a2b5f3c0e87e258d8b94a156e984c70" +checksum = "e8dfc9d19bdbf6d17e22319da49161d5d0108e4188e8b680aef6299eed22df60" dependencies = [ "serde_derive", ] @@ -5341,9 +5340,9 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.217" +version = "1.0.218" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a9bf7cf98d04a2b28aead066b7496853d4779c9cc183c440dbac457641e19a0" +checksum = "f09503e191f4e797cb8aac08e9a4a4695c5edf6a2e70e376d961ddd5c969f82b" dependencies = [ "proc-macro2", "quote", @@ -5506,9 +5505,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.13.2" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c5e1a9a646d36c3599cd173a41282daf47c44583ad367b8e6837255952e5c67" +checksum = "7fcf8323ef1faaee30a44a340193b1ac6814fd9b7b4e88e9d4519a3e4abe1cfd" [[package]] name = "snafu" @@ -5565,17 +5564,11 @@ dependencies = [ "windows-sys 0.52.0", ] -[[package]] -name = "spin" -version = "0.9.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" - [[package]] name = "sqllogictest" -version = "0.27.1" +version = "0.27.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07a06aea5e52b0a63b9d8328b46ea2740cdab4cac13def8ef4f2e5288610f9ed" +checksum = "6f1c93848602f92e5925690d4805ccbc1ccdb61bee7d4ae79ad6862b542a539c" dependencies = [ "async-trait", "educe", @@ -5891,9 +5884,9 @@ dependencies = [ [[package]] name = "testcontainers" -version = "0.23.2" +version = "0.23.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "042009c52a4204476bff461ca8ef17bab6f1a91628504a8a36c6fd2c1cde2d5e" +checksum = "59a4f01f39bb10fc2a5ab23eb0d888b1e2bb168c157f61a1b98e6c501c639c74" dependencies = [ "async-trait", "bollard", @@ -6133,7 +6126,7 @@ version = "0.26.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f6d0975eaace0cf0fcadee4e4aaa5da15b5c079146f2cffb67c113be122bf37" dependencies = [ - "rustls 0.23.22", + "rustls 0.23.23", "tokio", ] @@ -6204,7 +6197,7 @@ dependencies = [ "axum", "base64 0.22.1", "bytes", - "h2 0.4.7", + "h2 0.4.8", "http 1.2.0", "http-body 1.0.1", "http-body-util", @@ -6360,9 +6353,9 @@ dependencies = [ [[package]] name = "typenum" -version = "1.17.0" +version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42ff0bf0c66b8238c6f3b578df37d0b7848e55df8577b3f74f92a69acceeb825" +checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f" [[package]] name = "typify" @@ -7094,11 +7087,11 @@ dependencies = [ [[package]] name = "zerocopy" -version = "0.8.17" +version = "0.8.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aa91407dacce3a68c56de03abe2760159582b846c6a4acd2f456618087f12713" +checksum = "79386d31a42a4996e3336b0919ddb90f81112af416270cff95b5f5af22b839c2" dependencies = [ - "zerocopy-derive 0.8.17", + "zerocopy-derive 0.8.18", ] [[package]] @@ -7114,9 +7107,9 @@ dependencies = [ [[package]] name = "zerocopy-derive" -version = "0.8.17" +version = "0.8.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "06718a168365cad3d5ff0bb133aad346959a2074bd4a85c121255a11304a8626" +checksum = "76331675d372f91bf8d17e13afbd5fe639200b73d01f0fc748bb059f9cca2db7" dependencies = [ "proc-macro2", "quote", diff --git a/Cargo.toml b/Cargo.toml index ccf3f02a2fdef..f7d39aeb3003f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -66,7 +66,9 @@ homepage = "https://datafusion.apache.org" license = "Apache-2.0" readme = "README.md" repository = "https://github.com/apache/datafusion" -rust-version = "1.81.0" +# Define Minimum Supported Rust Version (MSRV) +rust-version = "1.82.0" +# Define DataFusion version version = "45.0.0" [workspace.dependencies] @@ -95,6 +97,7 @@ async-trait = "0.1.73" bigdecimal = "0.4.7" bytes = "1.10" chrono = { version = "0.4.38", default-features = false } +criterion = "0.5.1" ctor = "0.2.9" dashmap = "6.0.1" datafusion = { path = "datafusion/core", version = "45.0.0", default-features = false } @@ -143,8 +146,7 @@ pbjson = { version = "0.7.0" } pbjson-types = "0.7" # Should match arrow-flight's version of prost. prost = "0.13.1" -prost-derive = "0.13.1" -rand = "0.8" +rand = "0.8.5" recursive = "0.1.1" regex = "1.8" rstest = "0.24.0" diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index ad8debaf2fa38..860089063c4df 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -43,7 +43,7 @@ futures = { workspace = true } log = { workspace = true } mimalloc = { version = "0.1", optional = true, default-features = false } parquet = { workspace = true, default-features = true } -serde = { version = "1.0.136", features = ["derive"] } +serde = { version = "1.0.218", features = ["derive"] } serde_json = { workspace = true } snmalloc-rs = { version = "0.3", optional = true } structopt = { version = "0.3", default-features = false } diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml index d88f8fccb928a..20cd1c4295e05 100644 --- a/datafusion-cli/Cargo.toml +++ b/datafusion-cli/Cargo.toml @@ -27,6 +27,9 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } diff --git a/datafusion-cli/src/command.rs b/datafusion-cli/src/command.rs index f0eb58a233910..fc7d1a2617cf6 100644 --- a/datafusion-cli/src/command.rs +++ b/datafusion-cli/src/command.rs @@ -62,7 +62,9 @@ impl Command { Self::Help => { let now = Instant::now(); let command_batch = all_commands_info(); - print_options.print_batches(command_batch.schema(), &[command_batch], now) + let schema = command_batch.schema(); + let num_rows = command_batch.num_rows(); + print_options.print_batches(schema, &[command_batch], now, num_rows) } Self::ListTables => { exec_and_print(ctx, print_options, "SHOW TABLES".into()).await diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index a4f154b2de92d..d560dee987f50 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -17,11 +17,6 @@ //! Execution functions -use std::collections::HashMap; -use std::fs::File; -use std::io::prelude::*; -use std::io::BufReader; - use crate::cli_context::CliSessionContext; use crate::helper::split_from_semicolon; use crate::print_format::PrintFormat; @@ -31,6 +26,11 @@ use crate::{ object_storage::get_object_store, print_options::{MaxRows, PrintOptions}, }; +use futures::StreamExt; +use std::collections::HashMap; +use std::fs::File; +use std::io::prelude::*; +use std::io::BufReader; use datafusion::common::instant::Instant; use datafusion::common::{plan_datafusion_err, plan_err}; @@ -39,10 +39,12 @@ use datafusion::datasource::listing::ListingTableUrl; use datafusion::error::{DataFusionError, Result}; use datafusion::logical_expr::{DdlStatement, LogicalPlan}; use datafusion::physical_plan::execution_plan::EmissionType; -use datafusion::physical_plan::{collect, execute_stream, ExecutionPlanProperties}; +use datafusion::physical_plan::{execute_stream, ExecutionPlanProperties}; use datafusion::sql::parser::{DFParser, Statement}; use datafusion::sql::sqlparser::dialect::dialect_from_str; +use datafusion::execution::memory_pool::MemoryConsumer; +use datafusion::physical_plan::spill::get_record_batch_memory_size; use datafusion::sql::sqlparser; use rustyline::error::ReadlineError; use rustyline::Editor; @@ -235,6 +237,10 @@ pub(super) async fn exec_and_print( let df = ctx.execute_logical_plan(plan).await?; let physical_plan = df.create_physical_plan().await?; + // Track memory usage for the query result if it's bounded + let mut reservation = + MemoryConsumer::new("DataFusion-Cli").register(task_ctx.memory_pool()); + if physical_plan.boundedness().is_unbounded() { if physical_plan.pipeline_behavior() == EmissionType::Final { return plan_err!( @@ -247,10 +253,31 @@ pub(super) async fn exec_and_print( let stream = execute_stream(physical_plan, task_ctx.clone())?; print_options.print_stream(stream, now).await?; } else { - // Bounded stream; collected results are printed after all input consumed. + // Bounded stream; collected results size is limited by the maxrows option let schema = physical_plan.schema(); - let results = collect(physical_plan, task_ctx.clone()).await?; - adjusted.into_inner().print_batches(schema, &results, now)?; + let mut stream = execute_stream(physical_plan, task_ctx.clone())?; + let mut results = vec![]; + let mut row_count = 0_usize; + let max_rows = match print_options.maxrows { + MaxRows::Unlimited => usize::MAX, + MaxRows::Limited(n) => n, + }; + while let Some(batch) = stream.next().await { + let batch = batch?; + let curr_num_rows = batch.num_rows(); + // Stop collecting results if the number of rows exceeds the limit + // results batch should include the last batch that exceeds the limit + if row_count < max_rows + curr_num_rows { + // Try to grow the reservation to accommodate the batch in memory + reservation.try_grow(get_record_batch_memory_size(&batch))?; + results.push(batch); + } + row_count += curr_num_rows; + } + adjusted + .into_inner() + .print_batches(schema, &results, now, row_count)?; + reservation.free(); } } diff --git a/datafusion-cli/src/functions.rs b/datafusion-cli/src/functions.rs index b5bcb8243ea94..13d2d5fd3547b 100644 --- a/datafusion-cli/src/functions.rs +++ b/datafusion-cli/src/functions.rs @@ -28,10 +28,10 @@ use arrow::record_batch::RecordBatch; use arrow::util::pretty::pretty_format_batches; use datafusion::catalog::{Session, TableFunctionImpl}; use datafusion::common::{plan_err, Column}; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::TableProvider; use datafusion::error::Result; use datafusion::logical_expr::Expr; -use datafusion::physical_plan::memory::MemorySourceConfig; use datafusion::physical_plan::ExecutionPlan; use datafusion::scalar::ScalarValue; diff --git a/datafusion-cli/src/lib.rs b/datafusion-cli/src/lib.rs index fbfc9242a61d4..34fba6f79304b 100644 --- a/datafusion-cli/src/lib.rs +++ b/datafusion-cli/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![doc = include_str!("../README.md")] pub const DATAFUSION_CLI_VERSION: &str = env!("CARGO_PKG_VERSION"); diff --git a/datafusion-cli/src/print_options.rs b/datafusion-cli/src/print_options.rs index e80cc55663ae2..9557e783e8a7c 100644 --- a/datafusion-cli/src/print_options.rs +++ b/datafusion-cli/src/print_options.rs @@ -102,6 +102,7 @@ impl PrintOptions { schema: SchemaRef, batches: &[RecordBatch], query_start_time: Instant, + row_count: usize, ) -> Result<()> { let stdout = std::io::stdout(); let mut writer = stdout.lock(); @@ -109,7 +110,6 @@ impl PrintOptions { self.format .print_batches(&mut writer, schema, batches, self.maxrows, true)?; - let row_count: usize = batches.iter().map(|b| b.num_rows()).sum(); let formatted_exec_details = get_execution_details_formatted( row_count, if self.format == PrintFormat::Table { diff --git a/datafusion-examples/examples/csv_json_opener.rs b/datafusion-examples/examples/csv_json_opener.rs index 7613578e8c3a3..ef4ff9d51e7f0 100644 --- a/datafusion-examples/examples/csv_json_opener.rs +++ b/datafusion-examples/examples/csv_json_opener.rs @@ -71,7 +71,7 @@ async fn csv_opener() -> Result<()> { .with_batch_size(8192) .with_projection(&scan_config); - let opener = config.create_file_opener(Ok(object_store), &scan_config, 0)?; + let opener = config.create_file_opener(object_store, &scan_config, 0); let mut result = vec![]; let mut stream = diff --git a/datafusion-examples/examples/parquet_exec_visitor.rs b/datafusion-examples/examples/parquet_exec_visitor.rs index 20809a1121c1c..6c9f1a3544309 100644 --- a/datafusion-examples/examples/parquet_exec_visitor.rs +++ b/datafusion-examples/examples/parquet_exec_visitor.rs @@ -20,10 +20,10 @@ use std::sync::Arc; use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ListingOptions, PartitionedFile}; use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; use datafusion::error::DataFusionError; use datafusion::execution::context::SessionContext; use datafusion::physical_plan::metrics::MetricValue; -use datafusion::physical_plan::source::DataSourceExec; use datafusion::physical_plan::{ execute_stream, visit_execution_plan, ExecutionPlan, ExecutionPlanVisitor, }; diff --git a/datafusion-examples/examples/remote_catalog.rs b/datafusion-examples/examples/remote_catalog.rs index f84c6a0302cee..70c0963545e08 100644 --- a/datafusion-examples/examples/remote_catalog.rs +++ b/datafusion-examples/examples/remote_catalog.rs @@ -36,9 +36,9 @@ use datafusion::catalog::TableProvider; use datafusion::catalog::{AsyncSchemaProvider, Session}; use datafusion::common::Result; use datafusion::common::{assert_batches_eq, internal_datafusion_err, plan_err}; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion::execution::SendableRecordBatchStream; use datafusion::logical_expr::{Expr, TableType}; -use datafusion::physical_plan::memory::MemorySourceConfig; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::{DataFrame, SessionContext}; diff --git a/datafusion-examples/examples/simple_udtf.rs b/datafusion-examples/examples/simple_udtf.rs index afba4c390f71c..d2b2d1bf96551 100644 --- a/datafusion-examples/examples/simple_udtf.rs +++ b/datafusion-examples/examples/simple_udtf.rs @@ -23,13 +23,13 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::catalog::Session; use datafusion::catalog::TableFunctionImpl; use datafusion::common::{plan_err, ScalarValue}; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::TableProvider; use datafusion::error::Result; use datafusion::execution::context::ExecutionProps; use datafusion::logical_expr::simplify::SimplifyContext; use datafusion::logical_expr::{Expr, TableType}; use datafusion::optimizer::simplify_expressions::ExprSimplifier; -use datafusion::physical_plan::memory::MemorySourceConfig; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::*; use std::fs::File; diff --git a/datafusion/catalog-listing/Cargo.toml b/datafusion/catalog-listing/Cargo.toml index 68d0ca3a149fd..c6f9ea8bd0557 100644 --- a/datafusion/catalog-listing/Cargo.toml +++ b/datafusion/catalog-listing/Cargo.toml @@ -27,6 +27,9 @@ repository.workspace = true rust-version.workspace = true version.workspace = true +[package.metadata.docs.rs] +all-features = true + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } diff --git a/datafusion/catalog-listing/src/mod.rs b/datafusion/catalog-listing/src/mod.rs index b98790e864551..cb0d86d8666ef 100644 --- a/datafusion/catalog-listing/src/mod.rs +++ b/datafusion/catalog-listing/src/mod.rs @@ -15,4 +15,10 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + pub mod helpers; diff --git a/datafusion/catalog/Cargo.toml b/datafusion/catalog/Cargo.toml index 73ac44a0316e1..8647494111980 100644 --- a/datafusion/catalog/Cargo.toml +++ b/datafusion/catalog/Cargo.toml @@ -27,6 +27,9 @@ repository.workspace = true rust-version.workspace = true version.workspace = true +[package.metadata.docs.rs] +all-features = true + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } diff --git a/datafusion/catalog/src/lib.rs b/datafusion/catalog/src/lib.rs index 0b8d73fabd25c..a339d4916b8d5 100644 --- a/datafusion/catalog/src/lib.rs +++ b/datafusion/catalog/src/lib.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + //! Interfaces and default implementations of catalogs and schemas. //! //! Implementations diff --git a/datafusion/common-runtime/Cargo.toml b/datafusion/common-runtime/Cargo.toml index a21c72cd9f839..6fd9b7ac8fe8d 100644 --- a/datafusion/common-runtime/Cargo.toml +++ b/datafusion/common-runtime/Cargo.toml @@ -28,12 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_common_runtime" -path = "src/lib.rs" [dependencies] log = { workspace = true } diff --git a/datafusion/common-runtime/src/lib.rs b/datafusion/common-runtime/src/lib.rs index 51cb988ea06a3..7bd8dc4cfe364 100644 --- a/datafusion/common-runtime/src/lib.rs +++ b/datafusion/common-runtime/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 3be666ce7974f..fc5c08ee2316a 100644 --- a/datafusion/common/Cargo.toml +++ b/datafusion/common/Cargo.toml @@ -28,12 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_common" -path = "src/lib.rs" [features] avro = ["apache-avro"] diff --git a/datafusion/common/src/lib.rs b/datafusion/common/src/lib.rs index df1ae100f581b..d5b7c22a546c1 100644 --- a/datafusion/common/src/lib.rs +++ b/datafusion/common/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/common/src/spans.rs b/datafusion/common/src/spans.rs index 40ebdeffb6010..5111e264123ce 100644 --- a/datafusion/common/src/spans.rs +++ b/datafusion/common/src/spans.rs @@ -140,7 +140,7 @@ impl Span { /// the column a that comes from SELECT 1 AS a UNION ALL SELECT 2 AS a you'll /// need two spans. #[derive(Debug, Clone)] -// Store teh first [`Span`] on the stack because that is by far the most common +// Store the first [`Span`] on the stack because that is by far the most common // case. More will spill onto the heap. pub struct Spans(pub Vec); diff --git a/datafusion/common/src/table_reference.rs b/datafusion/common/src/table_reference.rs index bb53a30dcb234..9b6f9696c00bb 100644 --- a/datafusion/common/src/table_reference.rs +++ b/datafusion/common/src/table_reference.rs @@ -193,8 +193,7 @@ impl TableReference { match self { TableReference::Bare { table } => **table == *other.table(), TableReference::Partial { schema, table } => { - **table == *other.table() - && other.schema().map_or(true, |s| *s == **schema) + **table == *other.table() && other.schema().is_none_or(|s| *s == **schema) } TableReference::Full { catalog, @@ -202,8 +201,8 @@ impl TableReference { table, } => { **table == *other.table() - && other.schema().map_or(true, |s| *s == **schema) - && other.catalog().map_or(true, |c| *c == **catalog) + && other.schema().is_none_or(|s| *s == **schema) + && other.catalog().is_none_or(|c| *c == **catalog) } } } diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 8a706ca19f4df..1e0f27ccdfc87 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -29,13 +29,12 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true -[lib] -name = "datafusion" -path = "src/lib.rs" - [features] nested_expressions = ["datafusion-functions-nested"] # This feature is deprecated. Use the `nested_expressions` feature instead. @@ -132,7 +131,7 @@ zstd = { version = "0.13", optional = true, default-features = false } [dev-dependencies] async-trait = { workspace = true } -criterion = { version = "0.5", features = ["async_tokio"] } +criterion = { workspace = true, features = ["async_tokio"] } ctor = { workspace = true } dashmap = "6.1.0" datafusion-doc = { workspace = true } @@ -222,3 +221,7 @@ required-features = ["nested_expressions"] [[bench]] harness = false name = "dataframe" + +[[bench]] +harness = false +name = "spm" diff --git a/datafusion/core/benches/physical_plan.rs b/datafusion/core/benches/physical_plan.rs index 53c245ecc2b5c..aae1457ab9e6d 100644 --- a/datafusion/core/benches/physical_plan.rs +++ b/datafusion/core/benches/physical_plan.rs @@ -33,9 +33,9 @@ use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMerge use datafusion::physical_plan::{ collect, expressions::{col, PhysicalSortExpr}, - memory::MemorySourceConfig, }; use datafusion::prelude::SessionContext; +use datafusion_datasource::memory::MemorySourceConfig; use datafusion_physical_expr_common::sort_expr::LexOrdering; // Initialize the operator using the provided record batches and the sort key diff --git a/datafusion/core/benches/sort.rs b/datafusion/core/benches/sort.rs index 4d71d4c56a6da..8f0b3753f67c2 100644 --- a/datafusion/core/benches/sort.rs +++ b/datafusion/core/benches/sort.rs @@ -79,12 +79,13 @@ use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::{ execution::context::TaskContext, physical_plan::{ - coalesce_partitions::CoalescePartitionsExec, memory::MemorySourceConfig, + coalesce_partitions::CoalescePartitionsExec, sorts::sort_preserving_merge::SortPreservingMergeExec, ExecutionPlan, ExecutionPlanProperties, }, prelude::SessionContext, }; +use datafusion_datasource::memory::MemorySourceConfig; use datafusion_physical_expr::{expressions::col, PhysicalSortExpr}; use datafusion_physical_expr_common::sort_expr::LexOrdering; diff --git a/datafusion/physical-plan/benches/spm.rs b/datafusion/core/benches/spm.rs similarity index 98% rename from datafusion/physical-plan/benches/spm.rs rename to datafusion/core/benches/spm.rs index 3a2ecb57394b7..63b06f20cd86a 100644 --- a/datafusion/physical-plan/benches/spm.rs +++ b/datafusion/core/benches/spm.rs @@ -27,7 +27,7 @@ use datafusion_physical_plan::{collect, ExecutionPlan}; use criterion::async_executor::FuturesExecutor; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use datafusion_physical_plan::memory::MemorySourceConfig; +use datafusion_datasource::memory::MemorySourceConfig; fn generate_spm_for_round_robin_tie_breaker( has_same_value: bool, diff --git a/datafusion/core/src/datasource/data_source.rs b/datafusion/core/src/datasource/data_source.rs index d31b68019e308..2db79c5c839da 100644 --- a/datafusion/core/src/datasource/data_source.rs +++ b/datafusion/core/src/datasource/data_source.rs @@ -26,6 +26,8 @@ use crate::datasource::physical_plan::{FileOpener, FileScanConfig}; use arrow::datatypes::SchemaRef; use datafusion_common::Statistics; +use datafusion_datasource::file_groups::FileGroupPartitioner; +use datafusion_physical_expr::LexOrdering; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::DisplayFormatType; @@ -38,10 +40,10 @@ pub trait FileSource: Send + Sync { /// Creates a `dyn FileOpener` based on given parameters fn create_file_opener( &self, - object_store: datafusion_common::Result>, + object_store: Arc, base_config: &FileScanConfig, partition: usize, - ) -> datafusion_common::Result>; + ) -> Arc; /// Any fn as_any(&self) -> &dyn Any; /// Initialize new type with batch size configuration @@ -62,9 +64,33 @@ pub trait FileSource: Send + Sync { fn fmt_extra(&self, _t: DisplayFormatType, _f: &mut Formatter) -> fmt::Result { Ok(()) } - /// Return true if the file format supports repartition + + /// If supported by the [`FileSource`], redistribute files across partitions according to their size. + /// Allows custom file formats to implement their own repartitioning logic. /// - /// If this returns true, the DataSourceExec may repartition the data - /// by breaking up the input files into multiple smaller groups. - fn supports_repartition(&self, config: &FileScanConfig) -> bool; + /// Provides a default repartitioning behavior, see comments on [`FileGroupPartitioner`] for more detail. + fn repartitioned( + &self, + target_partitions: usize, + repartition_file_min_size: usize, + output_ordering: Option, + config: &FileScanConfig, + ) -> datafusion_common::Result> { + if config.file_compression_type.is_compressed() || config.new_lines_in_values { + return Ok(None); + } + + let repartitioned_file_groups_option = FileGroupPartitioner::new() + .with_target_partitions(target_partitions) + .with_repartition_file_min_size(repartition_file_min_size) + .with_preserve_order_within_groups(output_ordering.is_some()) + .repartition_file_groups(&config.file_groups); + + if let Some(repartitioned_file_groups) = repartitioned_file_groups_option { + let mut source = config.clone(); + source.file_groups = repartitioned_file_groups; + return Ok(Some(source)); + } + Ok(None) + } } diff --git a/datafusion/core/src/datasource/empty.rs b/datafusion/core/src/datasource/empty.rs index abda7fa9ec4b6..77686c5eb7c27 100644 --- a/datafusion/core/src/datasource/empty.rs +++ b/datafusion/core/src/datasource/empty.rs @@ -28,7 +28,8 @@ use datafusion_common::project_schema; use crate::datasource::{TableProvider, TableType}; use crate::error::Result; use crate::logical_expr::Expr; -use crate::physical_plan::{empty::EmptyExec, ExecutionPlan}; +use datafusion_physical_plan::empty::EmptyExec; +use datafusion_physical_plan::ExecutionPlan; /// An empty plan that is useful for testing and generating plans /// without mapping them to actual data. diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index 3be8af59ea2a0..819da155a1a24 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -37,7 +37,8 @@ use datafusion_common::{config_err, DataFusionError, Result}; use datafusion_expr::dml::InsertOp; use datafusion_expr::{utils::conjunction, Expr, TableProviderFilterPushDown}; use datafusion_expr::{SortExpr, TableType}; -use datafusion_physical_plan::{empty::EmptyExec, ExecutionPlan, Statistics}; +use datafusion_physical_plan::empty::EmptyExec; +use datafusion_physical_plan::{ExecutionPlan, Statistics}; use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, SchemaRef}; use datafusion_common::{ diff --git a/datafusion/core/src/datasource/memory.rs b/datafusion/core/src/datasource/memory.rs index 94c6e45804e81..b8bec410070c3 100644 --- a/datafusion/core/src/datasource/memory.rs +++ b/datafusion/core/src/datasource/memory.rs @@ -38,11 +38,11 @@ use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_catalog::Session; use datafusion_common::{not_impl_err, plan_err, Constraints, DFSchema, SchemaExt}; +pub use datafusion_datasource::memory::MemorySourceConfig; +pub use datafusion_datasource::source::DataSourceExec; use datafusion_execution::TaskContext; use datafusion_expr::dml::InsertOp; use datafusion_expr::SortExpr; -use datafusion_physical_plan::memory::MemorySourceConfig; -use datafusion_physical_plan::source::DataSourceExec; use async_trait::async_trait; use futures::StreamExt; diff --git a/datafusion/core/src/datasource/mod.rs b/datafusion/core/src/datasource/mod.rs index 12dd9d7cab389..96687913fb425 100644 --- a/datafusion/core/src/datasource/mod.rs +++ b/datafusion/core/src/datasource/mod.rs @@ -36,6 +36,8 @@ mod statistics; pub mod stream; pub mod view; +pub use datafusion_datasource::source; + // backwards compatibility pub use self::default_table_source::{ provider_as_source, source_as_provider, DefaultTableSource, diff --git a/datafusion/core/src/datasource/physical_plan/arrow_file.rs b/datafusion/core/src/datasource/physical_plan/arrow_file.rs index 4a7cdc192cd34..e5523063c782f 100644 --- a/datafusion/core/src/datasource/physical_plan/arrow_file.rs +++ b/datafusion/core/src/datasource/physical_plan/arrow_file.rs @@ -32,12 +32,12 @@ use arrow::datatypes::SchemaRef; use arrow_ipc::reader::FileDecoder; use datafusion_common::config::ConfigOptions; use datafusion_common::{Constraints, Statistics}; +use datafusion_datasource::source::DataSourceExec; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, }; @@ -211,14 +211,14 @@ pub struct ArrowSource { impl FileSource for ArrowSource { fn create_file_opener( &self, - object_store: Result>, + object_store: Arc, base_config: &FileScanConfig, _partition: usize, - ) -> Result> { - Ok(Arc::new(ArrowOpener { - object_store: object_store?, + ) -> Arc { + Arc::new(ArrowOpener { + object_store, projection: base_config.file_column_projection_indices(), - })) + }) } fn as_any(&self) -> &dyn Any { @@ -256,10 +256,6 @@ impl FileSource for ArrowSource { fn file_type(&self) -> &str { "arrow" } - - fn supports_repartition(&self, config: &FileScanConfig) -> bool { - !(config.file_compression_type.is_compressed() || config.new_lines_in_values) - } } /// The struct arrow that implements `[FileOpener]` trait diff --git a/datafusion/core/src/datasource/physical_plan/avro.rs b/datafusion/core/src/datasource/physical_plan/avro.rs index b0a1d8c8c9e22..1674814d76a7e 100644 --- a/datafusion/core/src/datasource/physical_plan/avro.rs +++ b/datafusion/core/src/datasource/physical_plan/avro.rs @@ -29,12 +29,12 @@ use crate::error::Result; use arrow::datatypes::SchemaRef; use datafusion_common::{Constraints, Statistics}; +use datafusion_datasource::source::DataSourceExec; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, }; @@ -194,23 +194,23 @@ impl FileSource for AvroSource { #[cfg(feature = "avro")] fn create_file_opener( &self, - object_store: Result>, + object_store: Arc, _base_config: &FileScanConfig, _partition: usize, - ) -> Result> { - Ok(Arc::new(private::AvroOpener { + ) -> Arc { + Arc::new(private::AvroOpener { config: Arc::new(self.clone()), - object_store: object_store?, - })) + object_store, + }) } #[cfg(not(feature = "avro"))] fn create_file_opener( &self, - _object_store: Result>, + _object_store: Arc, _base_config: &FileScanConfig, _partition: usize, - ) -> Result> { + ) -> Arc { panic!("Avro feature is not enabled in this build") } @@ -255,10 +255,15 @@ impl FileSource for AvroSource { fn file_type(&self) -> &str { "avro" } - fn supports_repartition(&self, config: &FileScanConfig) -> bool { - !(config.file_compression_type.is_compressed() - || config.new_lines_in_values - || self.as_any().downcast_ref::().is_some()) + + fn repartitioned( + &self, + _target_partitions: usize, + _repartition_file_min_size: usize, + _output_ordering: Option, + _config: &FileScanConfig, + ) -> Result> { + Ok(None) } } diff --git a/datafusion/core/src/datasource/physical_plan/csv.rs b/datafusion/core/src/datasource/physical_plan/csv.rs index c0952229b5e04..629d452064f59 100644 --- a/datafusion/core/src/datasource/physical_plan/csv.rs +++ b/datafusion/core/src/datasource/physical_plan/csv.rs @@ -37,13 +37,13 @@ use arrow::csv; use arrow::datatypes::SchemaRef; use datafusion_common::config::ConfigOptions; use datafusion_common::{Constraints, Statistics}; +use datafusion_datasource::source::DataSourceExec; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use datafusion_physical_plan::projection::ProjectionExec; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PlanProperties}; use futures::{StreamExt, TryStreamExt}; @@ -409,7 +409,7 @@ impl ExecutionPlan for CsvExec { /// # }; /// # use datafusion::datasource::physical_plan::CsvSource; /// # use datafusion_execution::object_store::ObjectStoreUrl; -/// # use datafusion_physical_plan::source::DataSourceExec; +/// # use datafusion::datasource::source::DataSourceExec; /// /// # let object_store_url = ObjectStoreUrl::local_filesystem(); /// # let file_schema = Arc::new(Schema::empty()); @@ -564,15 +564,15 @@ impl CsvOpener { impl FileSource for CsvSource { fn create_file_opener( &self, - object_store: Result>, + object_store: Arc, base_config: &FileScanConfig, _partition: usize, - ) -> Result> { - Ok(Arc::new(CsvOpener { + ) -> Arc { + Arc::new(CsvOpener { config: Arc::new(self.clone()), file_compression_type: base_config.file_compression_type, - object_store: object_store?, - })) + object_store, + }) } fn as_any(&self) -> &dyn Any { @@ -618,9 +618,6 @@ impl FileSource for CsvSource { fn fmt_extra(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { write!(f, ", has_header={}", self.has_header) } - fn supports_repartition(&self, config: &FileScanConfig) -> bool { - !(config.file_compression_type.is_compressed() || config.new_lines_in_values) - } } impl FileOpener for CsvOpener { diff --git a/datafusion/core/src/datasource/physical_plan/file_scan_config.rs b/datafusion/core/src/datasource/physical_plan/file_scan_config.rs index 123ecc2f95823..6b74f6be79ebd 100644 --- a/datafusion/core/src/datasource/physical_plan/file_scan_config.rs +++ b/datafusion/core/src/datasource/physical_plan/file_scan_config.rs @@ -19,8 +19,8 @@ //! file sources. use super::{ - get_projected_output_ordering, statistics::MinMaxStatistics, FileGroupPartitioner, - FileGroupsDisplay, FileStream, + get_projected_output_ordering, statistics::MinMaxStatistics, FileGroupsDisplay, + FileStream, }; use crate::datasource::file_format::file_compression_type::FileCompressionType; use crate::datasource::{listing::PartitionedFile, object_store::ObjectStoreUrl}; @@ -36,13 +36,13 @@ use datafusion_physical_expr::{EquivalenceProperties, LexOrdering, Partitioning} use crate::datasource::data_source::FileSource; pub use datafusion_datasource::file_scan_config::*; +use datafusion_datasource::source::{DataSource, DataSourceExec}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_plan::display::{display_orderings, ProjectSchemaDisplay}; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion_physical_plan::projection::{ all_alias_free_columns, new_projections_for_columns, ProjectionExec, }; -use datafusion_physical_plan::source::{DataSource, DataSourceExec}; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan}; /// Convert type to a type suitable for use as a [`ListingTable`] @@ -162,7 +162,7 @@ impl DataSource for FileScanConfig { partition: usize, context: Arc, ) -> Result { - let object_store = context.runtime_env().object_store(&self.object_store_url); + let object_store = context.runtime_env().object_store(&self.object_store_url)?; let source = self .source @@ -170,7 +170,7 @@ impl DataSource for FileScanConfig { .with_schema(Arc::clone(&self.file_schema)) .with_projection(self); - let opener = source.create_file_opener(object_store, self, partition)?; + let opener = source.create_file_opener(object_store, self, partition); let stream = FileStream::new(self, partition, opener, source.metrics())?; Ok(Box::pin(stream)) @@ -203,30 +203,21 @@ impl DataSource for FileScanConfig { self.fmt_file_source(t, f) } - /// Redistribute files across partitions according to their size - /// See comments on [`FileGroupPartitioner`] for more detail. + /// If supported by the underlying [`FileSource`], redistribute files across partitions according to their size. fn repartitioned( &self, target_partitions: usize, repartition_file_min_size: usize, output_ordering: Option, ) -> Result>> { - if !self.source.supports_repartition(self) { - return Ok(None); - } - - let repartitioned_file_groups_option = FileGroupPartitioner::new() - .with_target_partitions(target_partitions) - .with_repartition_file_min_size(repartition_file_min_size) - .with_preserve_order_within_groups(output_ordering.is_some()) - .repartition_file_groups(&self.file_groups); - - if let Some(repartitioned_file_groups) = repartitioned_file_groups_option { - let mut source = self.clone(); - source.file_groups = repartitioned_file_groups; - return Ok(Some(Arc::new(source))); - } - Ok(None) + let source = self.source.repartitioned( + target_partitions, + repartition_file_min_size, + output_ordering, + self, + )?; + + Ok(source.map(|s| Arc::new(s) as _)) } fn output_partitioning(&self) -> Partitioning { diff --git a/datafusion/core/src/datasource/physical_plan/json.rs b/datafusion/core/src/datasource/physical_plan/json.rs index 590b1cb88dcd4..d1ae13b083ab9 100644 --- a/datafusion/core/src/datasource/physical_plan/json.rs +++ b/datafusion/core/src/datasource/physical_plan/json.rs @@ -35,12 +35,12 @@ use crate::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use arrow::json::ReaderBuilder; use arrow::{datatypes::SchemaRef, json}; use datafusion_common::{Constraints, Statistics}; +use datafusion_datasource::source::DataSourceExec; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PlanProperties}; use futures::{StreamExt, TryStreamExt}; @@ -262,18 +262,18 @@ impl JsonSource { impl FileSource for JsonSource { fn create_file_opener( &self, - object_store: Result>, + object_store: Arc, base_config: &FileScanConfig, _partition: usize, - ) -> Result> { - Ok(Arc::new(JsonOpener { + ) -> Arc { + Arc::new(JsonOpener { batch_size: self .batch_size .expect("Batch size must set before creating opener"), projected_schema: base_config.projected_file_schema(), file_compression_type: base_config.file_compression_type, - object_store: object_store?, - })) + object_store, + }) } fn as_any(&self) -> &dyn Any { @@ -313,10 +313,6 @@ impl FileSource for JsonSource { fn file_type(&self) -> &str { "json" } - - fn supports_repartition(&self, config: &FileScanConfig) -> bool { - !(config.file_compression_type.is_compressed() || config.new_lines_in_values) - } } impl FileOpener for JsonOpener { diff --git a/datafusion/core/src/datasource/physical_plan/parquet/mod.rs b/datafusion/core/src/datasource/physical_plan/parquet/mod.rs index 4bd43cd1aacae..2a2d6d7fefdf5 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/mod.rs +++ b/datafusion/core/src/datasource/physical_plan/parquet/mod.rs @@ -50,10 +50,10 @@ pub use access_plan::{ParquetAccessPlan, RowGroupAccess}; use arrow::datatypes::SchemaRef; use datafusion_common::config::ConfigOptions; use datafusion_common::Constraints; +use datafusion_datasource::source::DataSourceExec; use datafusion_physical_expr::{EquivalenceProperties, LexOrdering, PhysicalExpr}; use datafusion_physical_optimizer::pruning::PruningPredicate; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; -use datafusion_physical_plan::source::DataSourceExec; pub use metrics::ParquetFileMetrics; pub use page_filter::PagePruningAccessPlanFilter; pub use reader::{DefaultParquetFileReaderFactory, ParquetFileReaderFactory}; @@ -579,10 +579,10 @@ mod tests { use arrow::record_batch::RecordBatch; use bytes::{BufMut, BytesMut}; use datafusion_common::{assert_contains, ScalarValue}; + use datafusion_datasource::source::DataSourceExec; use datafusion_expr::{col, lit, when, Expr}; use datafusion_physical_expr::planner::logical2physical; use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; - use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use crate::datasource::physical_plan::parquet::source::ParquetSource; diff --git a/datafusion/core/src/datasource/physical_plan/parquet/source.rs b/datafusion/core/src/datasource/physical_plan/parquet/source.rs index 21881112075de..178de8f51ae40 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet/source.rs +++ b/datafusion/core/src/datasource/physical_plan/parquet/source.rs @@ -81,7 +81,7 @@ use object_store::ObjectStore; /// # use datafusion::datasource::listing::PartitionedFile; /// # use datafusion_execution::object_store::ObjectStoreUrl; /// # use datafusion_physical_expr::expressions::lit; -/// # use datafusion_physical_plan::source::DataSourceExec; +/// # use datafusion::datasource::source::DataSourceExec; /// # use datafusion_common::config::TableParquetOptions; /// /// # let file_schema = Arc::new(Schema::empty()); @@ -160,7 +160,7 @@ use object_store::ObjectStore; /// # use arrow::datatypes::Schema; /// # use datafusion::datasource::physical_plan::FileScanConfig; /// # use datafusion::datasource::listing::PartitionedFile; -/// # use datafusion_physical_plan::source::DataSourceExec; +/// # use datafusion::datasource::source::DataSourceExec; /// /// # fn parquet_exec() -> DataSourceExec { unimplemented!() } /// // Split a single DataSourceExec into multiple DataSourceExecs, one for each file @@ -202,7 +202,7 @@ use object_store::ObjectStore; /// # use datafusion::datasource::physical_plan::FileScanConfig; /// # use datafusion::datasource::physical_plan::parquet::source::ParquetSource; /// # use datafusion_execution::object_store::ObjectStoreUrl; -/// # use datafusion_physical_plan::source::DataSourceExec; +/// # use datafusion::datasource::source::DataSourceExec; /// /// # fn schema() -> SchemaRef { /// # Arc::new(Schema::empty()) @@ -463,10 +463,10 @@ impl ParquetSource { impl FileSource for ParquetSource { fn create_file_opener( &self, - object_store: datafusion_common::Result>, + object_store: Arc, base_config: &FileScanConfig, partition: usize, - ) -> datafusion_common::Result> { + ) -> Arc { let projection = base_config .file_column_projection_indices() .unwrap_or_else(|| (0..base_config.file_schema.fields().len()).collect()); @@ -475,17 +475,12 @@ impl FileSource for ParquetSource { .clone() .unwrap_or_else(|| Arc::new(DefaultSchemaAdapterFactory)); - let parquet_file_reader_factory = self - .parquet_file_reader_factory - .as_ref() - .map(|f| Ok(Arc::clone(f))) - .unwrap_or_else(|| { - object_store.map(|store| { - Arc::new(DefaultParquetFileReaderFactory::new(store)) as _ - }) - })?; - - Ok(Arc::new(ParquetOpener { + let parquet_file_reader_factory = + self.parquet_file_reader_factory.clone().unwrap_or_else(|| { + Arc::new(DefaultParquetFileReaderFactory::new(object_store)) as _ + }); + + Arc::new(ParquetOpener { partition_index: partition, projection: Arc::from(projection), batch_size: self @@ -504,7 +499,7 @@ impl FileSource for ParquetSource { enable_page_index: self.enable_page_index(), enable_bloom_filter: self.bloom_filter_on_read(), schema_adapter_factory, - })) + }) } fn as_any(&self) -> &dyn Any { @@ -586,7 +581,4 @@ impl FileSource for ParquetSource { } } } - fn supports_repartition(&self, _config: &FileScanConfig) -> bool { - true - } } diff --git a/datafusion/core/src/execution/session_state_defaults.rs b/datafusion/core/src/execution/session_state_defaults.rs index 92f649781cfd9..33bf01cf35cd8 100644 --- a/datafusion/core/src/execution/session_state_defaults.rs +++ b/datafusion/core/src/execution/session_state_defaults.rs @@ -94,6 +94,8 @@ impl SessionStateDefaults { feature = "unicode_expressions" ))] Arc::new(functions::planner::UserDefinedFunctionPlanner), + Arc::new(functions_aggregate::planner::AggregateFunctionPlanner), + Arc::new(functions_window::planner::WindowFunctionPlanner), ]; expr_planners diff --git a/datafusion/core/src/lib.rs b/datafusion/core/src/lib.rs index f4aa366500ef1..9a0d0157c1ae2 100644 --- a/datafusion/core/src/lib.rs +++ b/datafusion/core/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![cfg_attr(not(test), deny(clippy::clone_on_ref_ptr))] #![warn(missing_docs, clippy::needless_borrow)] diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index d73b7d81536a0..a74cdcc5920bf 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -38,7 +38,6 @@ use crate::logical_expr::{ use crate::physical_expr::{create_physical_expr, create_physical_exprs}; use crate::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; use crate::physical_plan::analyze::AnalyzeExec; -use crate::physical_plan::empty::EmptyExec; use crate::physical_plan::explain::ExplainExec; use crate::physical_plan::expressions::PhysicalSortExpr; use crate::physical_plan::filter::FilterExec; @@ -48,7 +47,6 @@ use crate::physical_plan::joins::{ }; use crate::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use crate::physical_plan::projection::ProjectionExec; -use crate::physical_plan::recursive_query::RecursiveQueryExec; use crate::physical_plan::repartition::RepartitionExec; use crate::physical_plan::sorts::sort::SortExec; use crate::physical_plan::union::UnionExec; @@ -58,6 +56,8 @@ use crate::physical_plan::{ displayable, windows, ExecutionPlan, ExecutionPlanProperties, InputOrderMode, Partitioning, PhysicalExpr, WindowExpr, }; +use datafusion_physical_plan::empty::EmptyExec; +use datafusion_physical_plan::recursive_query::RecursiveQueryExec; use arrow::array::{builder::StringBuilder, RecordBatch}; use arrow::compute::SortOptions; @@ -68,6 +68,7 @@ use datafusion_common::{ exec_err, internal_datafusion_err, internal_err, not_impl_err, plan_err, DFSchema, ScalarValue, }; +use datafusion_datasource::memory::MemorySourceConfig; use datafusion_expr::dml::{CopyTo, InsertOp}; use datafusion_expr::expr::{ physical_name, AggregateFunction, AggregateFunctionParams, Alias, GroupingSet, @@ -84,7 +85,6 @@ use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::LexOrdering; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_plan::execution_plan::InvariantLevel; -use datafusion_physical_plan::memory::MemorySourceConfig; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::unnest::ListUnnest; diff --git a/datafusion/core/src/test/mod.rs b/datafusion/core/src/test/mod.rs index ba85f9afb6da3..c569113a27bd6 100644 --- a/datafusion/core/src/test/mod.rs +++ b/datafusion/core/src/test/mod.rs @@ -42,7 +42,7 @@ use arrow::array::{self, Array, ArrayRef, Decimal128Builder, Int32Array}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::DataFusionError; -use datafusion_physical_plan::source::DataSourceExec; +use datafusion_datasource::source::DataSourceExec; #[cfg(feature = "compression")] use bzip2::write::BzEncoder; diff --git a/datafusion/core/src/test_util/parquet.rs b/datafusion/core/src/test_util/parquet.rs index 0e0090ef028e1..fc98b43051f86 100644 --- a/datafusion/core/src/test_util/parquet.rs +++ b/datafusion/core/src/test_util/parquet.rs @@ -37,7 +37,7 @@ use crate::physical_plan::metrics::MetricsSet; use crate::physical_plan::ExecutionPlan; use crate::prelude::{Expr, SessionConfig, SessionContext}; -use datafusion_physical_plan::source::DataSourceExec; +use datafusion_datasource::source::DataSourceExec; use object_store::path::Path; use object_store::ObjectMeta; use parquet::arrow::ArrowWriter; diff --git a/datafusion/core/tests/dataframe/dataframe_functions.rs b/datafusion/core/tests/dataframe/dataframe_functions.rs index 29c24948fbf09..33f32e8f0f669 100644 --- a/datafusion/core/tests/dataframe/dataframe_functions.rs +++ b/datafusion/core/tests/dataframe/dataframe_functions.rs @@ -22,6 +22,7 @@ use arrow::{ array::{Int32Array, StringArray}, record_batch::RecordBatch, }; +use datafusion_functions_aggregate::count::count_all; use std::sync::Arc; use datafusion::error::Result; @@ -31,7 +32,7 @@ use datafusion::prelude::*; use datafusion::assert_batches_eq; use datafusion_common::{DFSchema, ScalarValue}; use datafusion_expr::expr::Alias; -use datafusion_expr::ExprSchemable; +use datafusion_expr::{table_scan, ExprSchemable, LogicalPlanBuilder}; use datafusion_functions_aggregate::expr_fn::{approx_median, approx_percentile_cont}; use datafusion_functions_nested::map::map; @@ -1123,3 +1124,34 @@ async fn test_fn_map() -> Result<()> { Ok(()) } + +/// Call count wildcard from dataframe API +#[tokio::test] +async fn test_count_wildcard() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + Field::new("c", DataType::UInt32, false), + ]); + + let table_scan = table_scan(Some("test"), &schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(table_scan) + .aggregate(vec![col("b")], vec![count_all()]) + .unwrap() + .project(vec![count_all()]) + .unwrap() + .sort(vec![count_all().sort(true, false)]) + .unwrap() + .build() + .unwrap(); + + let expected = "Sort: count(*) ASC NULLS LAST [count(*):Int64]\ + \n Projection: count(*) [count(*):Int64]\ + \n Aggregate: groupBy=[[test.b]], aggr=[[count(*)]] [b:UInt32, count(*):Int64]\ + \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; + + let formatted_plan = plan.display_indent_schema().to_string(); + assert_eq!(formatted_plan, expected); + + Ok(()) +} diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index d545157607c7d..b05029e8e3b18 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -32,7 +32,8 @@ use arrow::datatypes::{ }; use arrow::error::ArrowError; use arrow::util::pretty::pretty_format_batches; -use datafusion_functions_aggregate::count::count_udaf; +use datafusion_expr::utils::COUNT_STAR_EXPANSION; +use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::{ array_agg, avg, count, count_distinct, max, median, min, sum, }; @@ -72,7 +73,7 @@ use datafusion_expr::expr::{GroupingSet, Sort, WindowFunction}; use datafusion_expr::var_provider::{VarProvider, VarType}; use datafusion_expr::{ cast, col, create_udf, exists, in_subquery, lit, out_ref_col, placeholder, - scalar_subquery, when, wildcard, Expr, ExprFunctionExt, ExprSchemable, LogicalPlan, + scalar_subquery, when, Expr, ExprFunctionExt, ExprSchemable, LogicalPlan, ScalarFunctionImplementation, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, }; @@ -2463,8 +2464,8 @@ async fn test_count_wildcard_on_sort() -> Result<()> { let df_results = ctx .table("t1") .await? - .aggregate(vec![col("b")], vec![count(wildcard())])? - .sort(vec![count(wildcard()).sort(true, false)])? + .aggregate(vec![col("b")], vec![count_all()])? + .sort(vec![count_all().sort(true, false)])? .explain(false, false)? .collect() .await?; @@ -2498,8 +2499,8 @@ async fn test_count_wildcard_on_where_in() -> Result<()> { Arc::new( ctx.table("t2") .await? - .aggregate(vec![], vec![count(wildcard())])? - .select(vec![count(wildcard())])? + .aggregate(vec![], vec![count_all()])? + .select(vec![count_all()])? .into_optimized_plan()?, ), ))? @@ -2532,8 +2533,8 @@ async fn test_count_wildcard_on_where_exist() -> Result<()> { .filter(exists(Arc::new( ctx.table("t2") .await? - .aggregate(vec![], vec![count(wildcard())])? - .select(vec![count(wildcard())])? + .aggregate(vec![], vec![count_all()])? + .select(vec![count_all()])? .into_unoptimized_plan(), // Usually, into_optimized_plan() should be used here, but due to // https://github.com/apache/datafusion/issues/5771, @@ -2568,7 +2569,7 @@ async fn test_count_wildcard_on_window() -> Result<()> { .await? .select(vec![Expr::WindowFunction(WindowFunction::new( WindowFunctionDefinition::AggregateUDF(count_udaf()), - vec![wildcard()], + vec![Expr::Literal(COUNT_STAR_EXPANSION)], )) .order_by(vec![Sort::new(col("a"), false, true)]) .window_frame(WindowFrame::new_bounds( @@ -2599,17 +2600,16 @@ async fn test_count_wildcard_on_aggregate() -> Result<()> { let sql_results = ctx .sql("select count(*) from t1") .await? - .select(vec![col("count(*)")])? .explain(false, false)? .collect() .await?; - // add `.select(vec![count(wildcard())])?` to make sure we can analyze all node instead of just top node. + // add `.select(vec![count_wildcard()])?` to make sure we can analyze all node instead of just top node. let df_results = ctx .table("t1") .await? - .aggregate(vec![], vec![count(wildcard())])? - .select(vec![count(wildcard())])? + .aggregate(vec![], vec![count_all()])? + .select(vec![count_all()])? .explain(false, false)? .collect() .await?; @@ -2646,8 +2646,8 @@ async fn test_count_wildcard_on_where_scalar_subquery() -> Result<()> { ctx.table("t2") .await? .filter(out_ref_col(DataType::UInt32, "t1.a").eq(col("t2.a")))? - .aggregate(vec![], vec![count(wildcard())])? - .select(vec![col(count(wildcard()).to_string())])? + .aggregate(vec![], vec![count_all()])? + .select(vec![col(count_all().to_string())])? .into_unoptimized_plan(), )) .gt(lit(ScalarValue::UInt8(Some(0)))), diff --git a/datafusion/core/tests/execution/logical_plan.rs b/datafusion/core/tests/execution/logical_plan.rs index a17bb5eec8a3c..b30636ddf6a81 100644 --- a/datafusion/core/tests/execution/logical_plan.rs +++ b/datafusion/core/tests/execution/logical_plan.rs @@ -15,6 +15,9 @@ // specific language governing permissions and limitations // under the License. +//! Logical plans need to provide stable semantics, as downstream projects +//! create them and depend on them. Test executable semantics of logical plans. + use arrow::array::Int64Array; use arrow::datatypes::{DataType, Field}; use datafusion::execution::session_state::SessionStateBuilder; @@ -30,9 +33,6 @@ use std::fmt::Debug; use std::ops::Deref; use std::sync::Arc; -///! Logical plans need to provide stable semantics, as downstream projects -///! create them and depend on them. Test executable semantics of logical plans. - #[tokio::test] async fn count_only_nulls() -> Result<()> { // Input: VALUES (NULL), (NULL), (NULL) AS _(col) diff --git a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs index 5e1f263b4c76b..1025a49ea1e38 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs @@ -30,6 +30,8 @@ use arrow::datatypes::{ }; use arrow::util::pretty::pretty_format_batches; use datafusion::common::Result; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; use datafusion::datasource::MemTable; use datafusion::physical_expr::aggregate::AggregateExprBuilder; use datafusion::physical_plan::aggregates::{ @@ -43,8 +45,6 @@ use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::expressions::col; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr_common::sort_expr::LexOrdering; -use datafusion_physical_plan::memory::MemorySourceConfig; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::InputOrderMode; use test_utils::{add_empty_batches, StringBatchGenerator}; diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index 5dd29f90ef838..da93dd5edf291 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -26,6 +26,8 @@ use arrow::datatypes::Schema; use arrow::record_batch::RecordBatch; use arrow::util::pretty::pretty_format_batches; use datafusion::common::JoinSide; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; use datafusion::logical_expr::{JoinType, Operator}; use datafusion::physical_expr::expressions::BinaryExpr; use datafusion::physical_plan::collect; @@ -38,8 +40,6 @@ use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::ScalarValue; use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::PhysicalExprRef; -use datafusion_physical_plan::memory::MemorySourceConfig; -use datafusion_physical_plan::source::DataSourceExec; use itertools::Itertools; use rand::Rng; diff --git a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs index 35fca789ddcb6..92f3755250663 100644 --- a/datafusion/core/tests/fuzz_cases/merge_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/merge_fuzz.rs @@ -24,10 +24,10 @@ use arrow::{ compute::SortOptions, record_batch::RecordBatch, }; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion::physical_plan::{ collect, expressions::{col, PhysicalSortExpr}, - memory::MemorySourceConfig, sorts::sort_preserving_merge::SortPreservingMergeExec, }; use datafusion::prelude::{SessionConfig, SessionContext}; diff --git a/datafusion/core/tests/fuzz_cases/sort_fuzz.rs b/datafusion/core/tests/fuzz_cases/sort_fuzz.rs index 51a5bc87efd97..0b0f0aa2f105a 100644 --- a/datafusion/core/tests/fuzz_cases/sort_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/sort_fuzz.rs @@ -24,6 +24,7 @@ use arrow::{ compute::SortOptions, record_batch::RecordBatch, }; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_plan::expressions::PhysicalSortExpr; use datafusion::physical_plan::sorts::sort::SortExec; @@ -33,7 +34,6 @@ use datafusion_common::cast::as_int32_array; use datafusion_execution::memory_pool::GreedyMemoryPool; use datafusion_physical_expr::expressions::col; use datafusion_physical_expr_common::sort_expr::LexOrdering; -use datafusion_physical_plan::memory::MemorySourceConfig; use rand::Rng; use test_utils::{batches_to_vec, partitions_to_sorted_vec}; diff --git a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs index d23408743f9fa..06b93d41af362 100644 --- a/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/sort_preserving_repartition_fuzz.rs @@ -44,9 +44,9 @@ mod sp_repartition_fuzz_tests { }; use test_utils::add_empty_batches; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; use datafusion_physical_expr_common::sort_expr::LexOrdering; - use datafusion_physical_plan::memory::MemorySourceConfig; - use datafusion_physical_plan::source::DataSourceExec; use itertools::izip; use rand::{rngs::StdRng, seq::SliceRandom, Rng, SeedableRng}; diff --git a/datafusion/core/tests/fuzz_cases/window_fuzz.rs b/datafusion/core/tests/fuzz_cases/window_fuzz.rs index 4a484221a88ad..a7f9e38c9ae31 100644 --- a/datafusion/core/tests/fuzz_cases/window_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/window_fuzz.rs @@ -22,6 +22,8 @@ use arrow::compute::{concat_batches, SortOptions}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use arrow::util::pretty::pretty_format_batches; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; use datafusion::functions_window::row_number::row_number_udwf; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::windows::{ @@ -48,8 +50,6 @@ use datafusion_functions_window::rank::{dense_rank_udwf, rank_udwf}; use datafusion_physical_expr::expressions::{cast, col, lit}; use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr}; use datafusion_physical_expr_common::sort_expr::LexOrdering; -use datafusion_physical_plan::memory::MemorySourceConfig; -use datafusion_physical_plan::source::DataSourceExec; use rand::distributions::Alphanumeric; use rand::rngs::StdRng; diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs index 1789f37535a94..64ab1378340aa 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs @@ -21,12 +21,14 @@ //! This file is organized as: //! - Test runners that spawn individual test processes //! - Test cases that contain the actual validation logic -use std::{process::Command, str}; - use log::info; +use std::sync::Once; +use std::{process::Command, str}; use crate::memory_limit::memory_limit_validation::utils; +static INIT: Once = Once::new(); + // =========================================================================== // Test runners: // Runners are splitted into multiple tests to run in parallel @@ -67,10 +69,35 @@ fn sort_with_mem_limit_2_cols_2_runner() { spawn_test_process("sort_with_mem_limit_2_cols_2"); } +/// `spawn_test_process` might trigger multiple recompilations and the test binary +/// size might grow indefinitely. This initializer ensures recompilation is only done +/// once and the target size is bounded. +/// +/// TODO: This is a hack, can be cleaned up if we have a better way to let multiple +/// test cases run in different processes (instead of different threads by default) +fn init_once() { + INIT.call_once(|| { + let _ = Command::new("cargo") + .arg("test") + .arg("--no-run") + .arg("--package") + .arg("datafusion") + .arg("--test") + .arg("core_integration") + .arg("--features") + .arg("extended_tests") + .env("DATAFUSION_TEST_MEM_LIMIT_VALIDATION", "1") + .output() + .expect("Failed to execute test command"); + }); +} + /// Helper function that executes a test in a separate process with the required environment /// variable set. Memory limit validation tasks need to measure memory resident set /// size (RSS), so they must run in a separate process. fn spawn_test_process(test: &str) { + init_once(); + let test_path = format!( "memory_limit::memory_limit_validation::sort_mem_validation::{}", test diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 669294d38af1f..a1985a1aa4470 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -27,6 +27,8 @@ use arrow::array::{ArrayRef, DictionaryArray, RecordBatch}; use arrow::compute::SortOptions; use arrow::datatypes::{Int32Type, SchemaRef}; use datafusion::assert_batches_eq; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; use datafusion::datasource::{MemTable, TableProvider}; use datafusion::execution::disk_manager::DiskManagerConfig; use datafusion::execution::runtime_env::RuntimeEnvBuilder; @@ -46,8 +48,6 @@ use datafusion_expr::{Expr, TableType}; use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; use datafusion_physical_optimizer::join_selection::JoinSelection; use datafusion_physical_optimizer::PhysicalOptimizerRule; -use datafusion_physical_plan::memory::MemorySourceConfig; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::spill::get_record_batch_memory_size; use test_utils::AccessLogGenerator; diff --git a/datafusion/core/tests/parquet/file_statistics.rs b/datafusion/core/tests/parquet/file_statistics.rs index 82024a731ed3b..ad75cf2607c45 100644 --- a/datafusion/core/tests/parquet/file_statistics.rs +++ b/datafusion/core/tests/parquet/file_statistics.rs @@ -22,6 +22,7 @@ use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, }; +use datafusion::datasource::source::DataSourceExec; use datafusion::datasource::TableProvider; use datafusion::execution::context::SessionState; use datafusion::execution::session_state::SessionStateBuilder; @@ -34,7 +35,6 @@ use datafusion_execution::cache::cache_unit::{ use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::{col, lit, Expr}; -use datafusion_physical_plan::source::DataSourceExec; use datafusion::datasource::physical_plan::FileScanConfig; use tempfile::tempdir; diff --git a/datafusion/core/tests/parquet/page_pruning.rs b/datafusion/core/tests/parquet/page_pruning.rs index ea86bf3685bb1..fe96a2eb5e719 100644 --- a/datafusion/core/tests/parquet/page_pruning.rs +++ b/datafusion/core/tests/parquet/page_pruning.rs @@ -25,6 +25,7 @@ use datafusion::datasource::file_format::FileFormat; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; use datafusion::execution::context::SessionState; use datafusion::physical_plan::metrics::MetricValue; use datafusion::physical_plan::ExecutionPlan; @@ -33,7 +34,6 @@ use datafusion_common::{ScalarValue, ToDFSchema}; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::{col, lit, Expr}; use datafusion_physical_expr::create_physical_expr; -use datafusion_physical_plan::source::DataSourceExec; use futures::StreamExt; use object_store::path::Path; diff --git a/datafusion/core/tests/parquet/utils.rs b/datafusion/core/tests/parquet/utils.rs index dd5541461ff60..8cb50b22cf63d 100644 --- a/datafusion/core/tests/parquet/utils.rs +++ b/datafusion/core/tests/parquet/utils.rs @@ -18,8 +18,8 @@ //! Utilities for parquet tests use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; use datafusion_physical_plan::metrics::MetricsSet; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::{accept, ExecutionPlan, ExecutionPlanVisitor}; /// Find the metrics from the first DataSourceExec encountered in the plan diff --git a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs index 1757c7150bfec..a79d743cb253d 100644 --- a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs +++ b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs @@ -22,6 +22,8 @@ use crate::physical_optimizer::test_utils::TestAggregate; use arrow::array::Int32Array; use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; use datafusion_common::cast::as_int64_array; use datafusion_common::config::ConfigOptions; use datafusion_common::Result; @@ -36,9 +38,7 @@ use datafusion_physical_plan::aggregates::PhysicalGroupBy; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::common; use datafusion_physical_plan::filter::FilterExec; -use datafusion_physical_plan::memory::MemorySourceConfig; use datafusion_physical_plan::projection::ProjectionExec; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::ExecutionPlan; /// Mock data using a MemorySourceConfig which has an exact count statistic diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 50c67f09c7045..66d1380e09c38 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -31,6 +31,7 @@ use datafusion::datasource::file_format::file_compression_type::FileCompressionT use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::object_store::ObjectStoreUrl; use datafusion::datasource::physical_plan::{CsvSource, FileScanConfig, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; use datafusion_common::error::Result; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::ScalarValue; @@ -57,7 +58,6 @@ use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::PlanProperties; diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 89bd97881e3a1..dfba57a584ea1 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -21,7 +21,9 @@ use std::sync::Arc; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::{CsvSource, FileScanConfig}; +use datafusion::datasource::source::DataSourceExec; use datafusion_common::config::ConfigOptions; use datafusion_common::Result; use datafusion_common::{JoinSide, JoinType, ScalarValue}; @@ -46,12 +48,10 @@ use datafusion_physical_plan::joins::{ HashJoinExec, NestedLoopJoinExec, PartitionMode, StreamJoinPartitionMode, SymmetricHashJoinExec, }; -use datafusion_physical_plan::memory::MemorySourceConfig; use datafusion_physical_plan::projection::{update_expr, ProjectionExec}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::streaming::PartitionStream; use datafusion_physical_plan::streaming::StreamingTableExec; use datafusion_physical_plan::union::UnionExec; diff --git a/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs b/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs index d9b569dfa6115..58eb866c590cc 100644 --- a/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs +++ b/datafusion/core/tests/physical_optimizer/replace_with_order_preserving_variants.rs @@ -32,13 +32,13 @@ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::collect; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; -use datafusion_physical_plan::memory::MemorySourceConfig; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::{ displayable, get_plan_string, ExecutionPlan, Partitioning, }; -use datafusion_physical_plan::source::DataSourceExec; +use datafusion::datasource::source::DataSourceExec; use datafusion_common::tree_node::{TransformedResult, TreeNode}; use datafusion_common::Result; use datafusion_expr::{JoinType, Operator}; diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 162f93facc90d..e4d72c112c381 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -26,7 +26,9 @@ use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::utils::expr::COUNT_STAR_EXPANSION; @@ -52,11 +54,9 @@ use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::utils::{JoinFilter, JoinOn}; use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; -use datafusion_physical_plan::memory::MemorySourceConfig; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; -use datafusion_physical_plan::source::DataSourceExec; use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec}; use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::union::UnionExec; diff --git a/datafusion/core/tests/sql/explain_analyze.rs b/datafusion/core/tests/sql/explain_analyze.rs index d4b5ae8b28203..128d1d0aa4b63 100644 --- a/datafusion/core/tests/sql/explain_analyze.rs +++ b/datafusion/core/tests/sql/explain_analyze.rs @@ -780,7 +780,7 @@ async fn explain_logical_plan_only() { let expected = vec![ vec![ "logical_plan", - "Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]]\ + "Aggregate: groupBy=[[]], aggr=[[count(*)]]\ \n SubqueryAlias: t\ \n Projection: \ \n Values: (Utf8(\"a\"), Int64(1), Int64(100)), (Utf8(\"a\"), Int64(2), Int64(150))" diff --git a/datafusion/core/tests/sql/path_partition.rs b/datafusion/core/tests/sql/path_partition.rs index 6345f5e4352fd..1a19bfe9e86f0 100644 --- a/datafusion/core/tests/sql/path_partition.rs +++ b/datafusion/core/tests/sql/path_partition.rs @@ -26,6 +26,7 @@ use std::sync::Arc; use arrow::datatypes::DataType; use datafusion::datasource::listing::ListingTableUrl; use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; use datafusion::{ assert_batches_sorted_eq, datasource::{ @@ -43,7 +44,6 @@ use datafusion_common::ScalarValue; use datafusion_execution::config::SessionConfig; use datafusion_expr::{col, lit, Expr, Operator}; use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; -use datafusion_physical_plan::source::DataSourceExec; use async_trait::async_trait; use bytes::Bytes; diff --git a/datafusion/core/tests/user_defined/user_defined_table_functions.rs b/datafusion/core/tests/user_defined/user_defined_table_functions.rs index 0ec9a5fd76201..618f0590ab3d6 100644 --- a/datafusion/core/tests/user_defined/user_defined_table_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_table_functions.rs @@ -26,6 +26,7 @@ use arrow::csv::ReaderBuilder; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::TableProvider; use datafusion::error::Result; use datafusion::execution::TaskContext; @@ -35,7 +36,6 @@ use datafusion_catalog::Session; use datafusion_catalog::TableFunctionImpl; use datafusion_common::{assert_batches_eq, DFSchema, ScalarValue}; use datafusion_expr::{EmptyRelation, Expr, LogicalPlan, Projection, TableType}; -use datafusion_physical_plan::memory::MemorySourceConfig; use async_trait::async_trait; diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index caf1c60a785d2..a1d7af106f697 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -27,6 +27,9 @@ repository.workspace = true rust-version.workspace = true version.workspace = true +[package.metadata.docs.rs] +all-features = true + [features] compression = ["async-compression", "xz2", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] @@ -49,6 +52,8 @@ datafusion-common = { workspace = true, features = ["object_store"] } datafusion-common-runtime = { workspace = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } +datafusion-physical-expr = { workspace = true } +datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } flate2 = { version = "1.0.24", optional = true } futures = { workspace = true } diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs new file mode 100644 index 0000000000000..efb178ad078ee --- /dev/null +++ b/datafusion/datasource/src/memory.rs @@ -0,0 +1,926 @@ +// 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 plan for reading in-memory batches of data + +use std::any::Any; +use std::fmt; +use std::sync::Arc; + +use crate::source::{DataSource, DataSourceExec}; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::memory::MemoryStream; +use datafusion_physical_plan::projection::{ + all_alias_free_columns, new_projections_for_columns, ProjectionExec, +}; +use datafusion_physical_plan::{ + common, ColumnarValue, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PhysicalExpr, PlanProperties, SendableRecordBatchStream, Statistics, +}; + +use arrow::array::{RecordBatch, RecordBatchOptions}; +use arrow::datatypes::{Schema, SchemaRef}; +use datafusion_common::{ + internal_err, plan_err, project_schema, Constraints, Result, ScalarValue, +}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::equivalence::ProjectionMapping; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::{EquivalenceProperties, LexOrdering}; + +/// Execution plan for reading in-memory batches of data +#[derive(Clone)] +#[deprecated( + since = "46.0.0", + note = "use MemorySourceConfig and DataSourceExec instead" +)] +pub struct MemoryExec { + inner: DataSourceExec, + /// The partitions to query + partitions: Vec>, + /// Optional projection + projection: Option>, + // Sort information: one or more equivalent orderings + sort_information: Vec, + /// if partition sizes should be displayed + show_sizes: bool, +} + +#[allow(unused, deprecated)] +impl fmt::Debug for MemoryExec { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + self.inner.fmt_as(DisplayFormatType::Default, f) + } +} + +#[allow(unused, deprecated)] +impl DisplayAs for MemoryExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + self.inner.fmt_as(t, f) + } +} + +#[allow(unused, deprecated)] +impl ExecutionPlan for MemoryExec { + fn name(&self) -> &'static str { + "MemoryExec" + } + + /// Return a reference to Any that can be used for downcasting + fn as_any(&self) -> &dyn Any { + self + } + + fn properties(&self) -> &PlanProperties { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + // This is a leaf node and has no children + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + // MemoryExec has no children + if children.is_empty() { + Ok(self) + } else { + internal_err!("Children cannot be replaced in {self:?}") + } + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } + + /// We recompute the statistics dynamically from the arrow metadata as it is pretty cheap to do so + fn statistics(&self) -> Result { + self.inner.statistics() + } + + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result>> { + self.inner.try_swapping_with_projection(projection) + } +} + +#[allow(unused, deprecated)] +impl MemoryExec { + /// Create a new execution plan for reading in-memory record batches + /// The provided `schema` should not have the projection applied. + pub fn try_new( + partitions: &[Vec], + schema: SchemaRef, + projection: Option>, + ) -> Result { + let source = MemorySourceConfig::try_new(partitions, schema, projection.clone())?; + let data_source = DataSourceExec::new(Arc::new(source)); + Ok(Self { + inner: data_source, + partitions: partitions.to_vec(), + projection, + sort_information: vec![], + show_sizes: true, + }) + } + + /// Create a new execution plan from a list of constant values (`ValuesExec`) + pub fn try_new_as_values( + schema: SchemaRef, + data: Vec>>, + ) -> Result { + if data.is_empty() { + return plan_err!("Values list cannot be empty"); + } + + let n_row = data.len(); + let n_col = schema.fields().len(); + + // We have this single row batch as a placeholder to satisfy evaluation argument + // and generate a single output row + let placeholder_schema = Arc::new(Schema::empty()); + let placeholder_batch = RecordBatch::try_new_with_options( + Arc::clone(&placeholder_schema), + vec![], + &RecordBatchOptions::new().with_row_count(Some(1)), + )?; + + // Evaluate each column + let arrays = (0..n_col) + .map(|j| { + (0..n_row) + .map(|i| { + let expr = &data[i][j]; + let result = expr.evaluate(&placeholder_batch)?; + + match result { + ColumnarValue::Scalar(scalar) => Ok(scalar), + ColumnarValue::Array(array) if array.len() == 1 => { + ScalarValue::try_from_array(&array, 0) + } + ColumnarValue::Array(_) => { + plan_err!("Cannot have array values in a values list") + } + } + }) + .collect::>>() + .and_then(ScalarValue::iter_to_array) + }) + .collect::>>()?; + + let batch = RecordBatch::try_new_with_options( + Arc::clone(&schema), + arrays, + &RecordBatchOptions::new().with_row_count(Some(n_row)), + )?; + + let partitions = vec![batch]; + Self::try_new_from_batches(Arc::clone(&schema), partitions) + } + + /// Create a new plan using the provided schema and batches. + /// + /// Errors if any of the batches don't match the provided schema, or if no + /// batches are provided. + pub fn try_new_from_batches( + schema: SchemaRef, + batches: Vec, + ) -> Result { + if batches.is_empty() { + return plan_err!("Values list cannot be empty"); + } + + for batch in &batches { + let batch_schema = batch.schema(); + if batch_schema != schema { + return plan_err!( + "Batch has invalid schema. Expected: {}, got: {}", + schema, + batch_schema + ); + } + } + + let partitions = vec![batches]; + let source = MemorySourceConfig { + partitions: partitions.clone(), + schema: Arc::clone(&schema), + projected_schema: Arc::clone(&schema), + projection: None, + sort_information: vec![], + show_sizes: true, + fetch: None, + }; + let data_source = DataSourceExec::new(Arc::new(source)); + Ok(Self { + inner: data_source, + partitions, + projection: None, + sort_information: vec![], + show_sizes: true, + }) + } + + fn memory_source_config(&self) -> MemorySourceConfig { + self.inner + .source() + .as_any() + .downcast_ref::() + .unwrap() + .clone() + } + + pub fn with_constraints(mut self, constraints: Constraints) -> Self { + self.inner = self.inner.with_constraints(constraints); + self + } + + /// Set `show_sizes` to determine whether to display partition sizes + pub fn with_show_sizes(mut self, show_sizes: bool) -> Self { + let mut memory_source = self.memory_source_config(); + memory_source.show_sizes = show_sizes; + self.show_sizes = show_sizes; + self.inner = DataSourceExec::new(Arc::new(memory_source)); + self + } + + /// Ref to constraints + pub fn constraints(&self) -> &Constraints { + self.properties().equivalence_properties().constraints() + } + + /// Ref to partitions + pub fn partitions(&self) -> &[Vec] { + &self.partitions + } + + /// Ref to projection + pub fn projection(&self) -> &Option> { + &self.projection + } + + /// Show sizes + pub fn show_sizes(&self) -> bool { + self.show_sizes + } + + /// Ref to sort information + pub fn sort_information(&self) -> &[LexOrdering] { + &self.sort_information + } + + /// A memory table can be ordered by multiple expressions simultaneously. + /// [`EquivalenceProperties`] keeps track of expressions that describe the + /// global ordering of the schema. These columns are not necessarily same; e.g. + /// ```text + /// ┌-------┐ + /// | a | b | + /// |---|---| + /// | 1 | 9 | + /// | 2 | 8 | + /// | 3 | 7 | + /// | 5 | 5 | + /// └---┴---┘ + /// ``` + /// where both `a ASC` and `b DESC` can describe the table ordering. With + /// [`EquivalenceProperties`], we can keep track of these equivalences + /// and treat `a ASC` and `b DESC` as the same ordering requirement. + /// + /// Note that if there is an internal projection, that projection will be + /// also applied to the given `sort_information`. + pub fn try_with_sort_information( + mut self, + sort_information: Vec, + ) -> Result { + self.sort_information = sort_information.clone(); + let mut memory_source = self.memory_source_config(); + memory_source = memory_source.try_with_sort_information(sort_information)?; + self.inner = DataSourceExec::new(Arc::new(memory_source)); + Ok(self) + } + + /// Arc clone of ref to original schema + pub fn original_schema(&self) -> SchemaRef { + Arc::clone(&self.inner.schema()) + } + + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. + fn compute_properties( + schema: SchemaRef, + orderings: &[LexOrdering], + constraints: Constraints, + partitions: &[Vec], + ) -> PlanProperties { + PlanProperties::new( + EquivalenceProperties::new_with_orderings(schema, orderings) + .with_constraints(constraints), + Partitioning::UnknownPartitioning(partitions.len()), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } +} + +/// Data source configuration for reading in-memory batches of data +#[derive(Clone)] +pub struct MemorySourceConfig { + /// The partitions to query + partitions: Vec>, + /// Schema representing the data before projection + schema: SchemaRef, + /// Schema representing the data after the optional projection is applied + projected_schema: SchemaRef, + /// Optional projection + projection: Option>, + /// Sort information: one or more equivalent orderings + sort_information: Vec, + /// if partition sizes should be displayed + show_sizes: bool, + /// The maximum number of records to read from this plan. If `None`, + /// all records after filtering are returned. + fetch: Option, +} + +impl DataSource for MemorySourceConfig { + fn open( + &self, + partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin( + MemoryStream::try_new( + self.partitions[partition].clone(), + Arc::clone(&self.projected_schema), + self.projection.clone(), + )? + .with_fetch(self.fetch), + )) + } + + fn as_any(&self) -> &dyn Any { + self + } + + fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let partition_sizes: Vec<_> = + self.partitions.iter().map(|b| b.len()).collect(); + + let output_ordering = self + .sort_information + .first() + .map(|output_ordering| { + format!(", output_ordering={}", output_ordering) + }) + .unwrap_or_default(); + + let eq_properties = self.eq_properties(); + let constraints = eq_properties.constraints(); + let constraints = if constraints.is_empty() { + String::new() + } else { + format!(", {}", constraints) + }; + + let limit = self + .fetch + .map_or(String::new(), |limit| format!(", fetch={}", limit)); + if self.show_sizes { + write!( + f, + "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", + partition_sizes.len(), + ) + } else { + write!( + f, + "partitions={}{limit}{output_ordering}{constraints}", + partition_sizes.len(), + ) + } + } + } + } + + fn output_partitioning(&self) -> Partitioning { + Partitioning::UnknownPartitioning(self.partitions.len()) + } + + fn eq_properties(&self) -> EquivalenceProperties { + EquivalenceProperties::new_with_orderings( + Arc::clone(&self.projected_schema), + self.sort_information.as_slice(), + ) + } + + fn statistics(&self) -> Result { + Ok(common::compute_record_batch_statistics( + &self.partitions, + &self.schema, + self.projection.clone(), + )) + } + + fn with_fetch(&self, limit: Option) -> Option> { + let source = self.clone(); + Some(Arc::new(source.with_limit(limit))) + } + + fn fetch(&self) -> Option { + self.fetch + } + + fn try_swapping_with_projection( + &self, + projection: &ProjectionExec, + ) -> Result>> { + // If there is any non-column or alias-carrier expression, Projection should not be removed. + // This process can be moved into MemoryExec, but it would be an overlap of their responsibility. + all_alias_free_columns(projection.expr()) + .then(|| { + let all_projections = (0..self.schema.fields().len()).collect(); + let new_projections = new_projections_for_columns( + projection, + self.projection().as_ref().unwrap_or(&all_projections), + ); + + MemorySourceConfig::try_new_exec( + self.partitions(), + self.original_schema(), + Some(new_projections), + ) + .map(|e| e as _) + }) + .transpose() + } +} + +impl MemorySourceConfig { + /// Create a new `MemorySourceConfig` for reading in-memory record batches + /// The provided `schema` should not have the projection applied. + pub fn try_new( + partitions: &[Vec], + schema: SchemaRef, + projection: Option>, + ) -> Result { + let projected_schema = project_schema(&schema, projection.as_ref())?; + Ok(Self { + partitions: partitions.to_vec(), + schema, + projected_schema, + projection, + sort_information: vec![], + show_sizes: true, + fetch: None, + }) + } + + /// Create a new `DataSourceExec` plan for reading in-memory record batches + /// The provided `schema` should not have the projection applied. + pub fn try_new_exec( + partitions: &[Vec], + schema: SchemaRef, + projection: Option>, + ) -> Result> { + let source = Self::try_new(partitions, schema, projection)?; + Ok(Arc::new(DataSourceExec::new(Arc::new(source)))) + } + + /// Create a new execution plan from a list of constant values (`ValuesExec`) + pub fn try_new_as_values( + schema: SchemaRef, + data: Vec>>, + ) -> Result> { + if data.is_empty() { + return plan_err!("Values list cannot be empty"); + } + + let n_row = data.len(); + let n_col = schema.fields().len(); + + // We have this single row batch as a placeholder to satisfy evaluation argument + // and generate a single output row + let placeholder_schema = Arc::new(Schema::empty()); + let placeholder_batch = RecordBatch::try_new_with_options( + Arc::clone(&placeholder_schema), + vec![], + &RecordBatchOptions::new().with_row_count(Some(1)), + )?; + + // Evaluate each column + let arrays = (0..n_col) + .map(|j| { + (0..n_row) + .map(|i| { + let expr = &data[i][j]; + let result = expr.evaluate(&placeholder_batch)?; + + match result { + ColumnarValue::Scalar(scalar) => Ok(scalar), + ColumnarValue::Array(array) if array.len() == 1 => { + ScalarValue::try_from_array(&array, 0) + } + ColumnarValue::Array(_) => { + plan_err!("Cannot have array values in a values list") + } + } + }) + .collect::>>() + .and_then(ScalarValue::iter_to_array) + }) + .collect::>>()?; + + let batch = RecordBatch::try_new_with_options( + Arc::clone(&schema), + arrays, + &RecordBatchOptions::new().with_row_count(Some(n_row)), + )?; + + let partitions = vec![batch]; + Self::try_new_from_batches(Arc::clone(&schema), partitions) + } + + /// Create a new plan using the provided schema and batches. + /// + /// Errors if any of the batches don't match the provided schema, or if no + /// batches are provided. + pub fn try_new_from_batches( + schema: SchemaRef, + batches: Vec, + ) -> Result> { + if batches.is_empty() { + return plan_err!("Values list cannot be empty"); + } + + for batch in &batches { + let batch_schema = batch.schema(); + if batch_schema != schema { + return plan_err!( + "Batch has invalid schema. Expected: {}, got: {}", + schema, + batch_schema + ); + } + } + + let partitions = vec![batches]; + let source = Self { + partitions, + schema: Arc::clone(&schema), + projected_schema: Arc::clone(&schema), + projection: None, + sort_information: vec![], + show_sizes: true, + fetch: None, + }; + Ok(Arc::new(DataSourceExec::new(Arc::new(source)))) + } + + /// Set the limit of the files + pub fn with_limit(mut self, limit: Option) -> Self { + self.fetch = limit; + self + } + + /// Set `show_sizes` to determine whether to display partition sizes + pub fn with_show_sizes(mut self, show_sizes: bool) -> Self { + self.show_sizes = show_sizes; + self + } + + /// Ref to partitions + pub fn partitions(&self) -> &[Vec] { + &self.partitions + } + + /// Ref to projection + pub fn projection(&self) -> &Option> { + &self.projection + } + + /// Show sizes + pub fn show_sizes(&self) -> bool { + self.show_sizes + } + + /// Ref to sort information + pub fn sort_information(&self) -> &[LexOrdering] { + &self.sort_information + } + + /// A memory table can be ordered by multiple expressions simultaneously. + /// [`EquivalenceProperties`] keeps track of expressions that describe the + /// global ordering of the schema. These columns are not necessarily same; e.g. + /// ```text + /// ┌-------┐ + /// | a | b | + /// |---|---| + /// | 1 | 9 | + /// | 2 | 8 | + /// | 3 | 7 | + /// | 5 | 5 | + /// └---┴---┘ + /// ``` + /// where both `a ASC` and `b DESC` can describe the table ordering. With + /// [`EquivalenceProperties`], we can keep track of these equivalences + /// and treat `a ASC` and `b DESC` as the same ordering requirement. + /// + /// Note that if there is an internal projection, that projection will be + /// also applied to the given `sort_information`. + pub fn try_with_sort_information( + mut self, + mut sort_information: Vec, + ) -> Result { + // All sort expressions must refer to the original schema + let fields = self.schema.fields(); + let ambiguous_column = sort_information + .iter() + .flat_map(|ordering| ordering.clone()) + .flat_map(|expr| collect_columns(&expr.expr)) + .find(|col| { + fields + .get(col.index()) + .map(|field| field.name() != col.name()) + .unwrap_or(true) + }); + if let Some(col) = ambiguous_column { + return internal_err!( + "Column {:?} is not found in the original schema of the MemorySourceConfig", + col + ); + } + + // If there is a projection on the source, we also need to project orderings + if let Some(projection) = &self.projection { + let base_eqp = EquivalenceProperties::new_with_orderings( + self.original_schema(), + &sort_information, + ); + let proj_exprs = projection + .iter() + .map(|idx| { + let base_schema = self.original_schema(); + let name = base_schema.field(*idx).name(); + (Arc::new(Column::new(name, *idx)) as _, name.to_string()) + }) + .collect::>(); + let projection_mapping = + ProjectionMapping::try_new(&proj_exprs, &self.original_schema())?; + sort_information = base_eqp + .project(&projection_mapping, Arc::clone(&self.projected_schema)) + .into_oeq_class() + .into_inner(); + } + + self.sort_information = sort_information; + Ok(self) + } + + /// Arc clone of ref to original schema + pub fn original_schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod memory_source_tests { + use std::sync::Arc; + + use crate::memory::MemorySourceConfig; + use crate::source::DataSourceExec; + use datafusion_physical_plan::ExecutionPlan; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_physical_expr_common::sort_expr::LexOrdering; + + #[test] + fn test_memory_order_eq() -> datafusion_common::Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + Field::new("c", DataType::Int64, false), + ])); + let sort1 = LexOrdering::new(vec![ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions::default(), + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions::default(), + }, + ]); + let sort2 = LexOrdering::new(vec![PhysicalSortExpr { + expr: col("c", &schema)?, + options: SortOptions::default(), + }]); + let mut expected_output_order = LexOrdering::default(); + expected_output_order.extend(sort1.clone()); + expected_output_order.extend(sort2.clone()); + + let sort_information = vec![sort1.clone(), sort2.clone()]; + let mem_exec = Arc::new(DataSourceExec::new(Arc::new( + MemorySourceConfig::try_new(&[vec![]], schema, None)? + .try_with_sort_information(sort_information)?, + ))); + + assert_eq!( + mem_exec.properties().output_ordering().unwrap(), + &expected_output_order + ); + let eq_properties = mem_exec.properties().equivalence_properties(); + assert!(eq_properties.oeq_class().contains(&sort1)); + assert!(eq_properties.oeq_class().contains(&sort2)); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ArrayRef, Int32Array}; + use datafusion_physical_plan::expressions::lit; + use std::collections::HashMap; + + use arrow::datatypes::{DataType, Field}; + use datafusion_common::assert_batches_eq; + use datafusion_common::stats::{ColumnStatistics, Precision}; + use futures::StreamExt; + + // Return a RecordBatch with a single Int32 array with values (0..sz) in a field named "i" + pub fn make_partition(sz: i32) -> RecordBatch { + let seq_start = 0; + let seq_end = sz; + let values = (seq_start..seq_end).collect::>(); + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)])); + let arr = Arc::new(Int32Array::from(values)); + let arr = arr as ArrayRef; + + RecordBatch::try_new(schema, vec![arr]).unwrap() + } + + #[tokio::test] + async fn exec_with_limit() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let batch = make_partition(7); + let schema = batch.schema(); + let batches = vec![batch.clone(), batch]; + + let exec = MemorySourceConfig::try_new_from_batches(schema, batches).unwrap(); + assert_eq!(exec.fetch(), None); + + let exec = exec.with_fetch(Some(4)).unwrap(); + assert_eq!(exec.fetch(), Some(4)); + + let mut it = exec.execute(0, task_ctx)?; + let mut results = vec![]; + while let Some(batch) = it.next().await { + results.push(batch?); + } + + let expected = [ + "+---+", "| i |", "+---+", "| 0 |", "| 1 |", "| 2 |", "| 3 |", "+---+", + ]; + assert_batches_eq!(expected, &results); + Ok(()) + } + + /// Get the schema for the aggregate_test_* csv files + pub fn aggr_test_schema() -> SchemaRef { + let mut f1 = Field::new("c1", DataType::Utf8, false); + f1.set_metadata(HashMap::from_iter(vec![("testing".into(), "test".into())])); + let schema = Schema::new(vec![ + f1, + Field::new("c2", DataType::UInt32, false), + Field::new("c3", DataType::Int8, false), + Field::new("c4", DataType::Int16, false), + Field::new("c5", DataType::Int32, false), + Field::new("c6", DataType::Int64, false), + Field::new("c7", DataType::UInt8, false), + Field::new("c8", DataType::UInt16, false), + Field::new("c9", DataType::UInt32, false), + Field::new("c10", DataType::UInt64, false), + Field::new("c11", DataType::Float32, false), + Field::new("c12", DataType::Float64, false), + Field::new("c13", DataType::Utf8, false), + ]); + + Arc::new(schema) + } + + #[tokio::test] + async fn values_empty_case() -> Result<()> { + let schema = aggr_test_schema(); + let empty = MemorySourceConfig::try_new_as_values(schema, vec![]); + assert!(empty.is_err()); + Ok(()) + } + + #[test] + fn new_exec_with_batches() { + let batch = make_partition(7); + let schema = batch.schema(); + let batches = vec![batch.clone(), batch]; + let _exec = MemorySourceConfig::try_new_from_batches(schema, batches).unwrap(); + } + + #[test] + fn new_exec_with_batches_empty() { + let batch = make_partition(7); + let schema = batch.schema(); + let _ = MemorySourceConfig::try_new_from_batches(schema, Vec::new()).unwrap_err(); + } + + #[test] + fn new_exec_with_batches_invalid_schema() { + let batch = make_partition(7); + let batches = vec![batch.clone(), batch]; + + let invalid_schema = Arc::new(Schema::new(vec![ + Field::new("col0", DataType::UInt32, false), + Field::new("col1", DataType::Utf8, false), + ])); + let _ = MemorySourceConfig::try_new_from_batches(invalid_schema, batches) + .unwrap_err(); + } + + // Test issue: https://github.com/apache/datafusion/issues/8763 + #[test] + fn new_exec_with_non_nullable_schema() { + let schema = Arc::new(Schema::new(vec![Field::new( + "col0", + DataType::UInt32, + false, + )])); + let _ = MemorySourceConfig::try_new_as_values( + Arc::clone(&schema), + vec![vec![lit(1u32)]], + ) + .unwrap(); + // Test that a null value is rejected + let _ = MemorySourceConfig::try_new_as_values( + schema, + vec![vec![lit(ScalarValue::UInt32(None))]], + ) + .unwrap_err(); + } + + #[test] + fn values_stats_with_nulls_only() -> Result<()> { + let data = vec![ + vec![lit(ScalarValue::Null)], + vec![lit(ScalarValue::Null)], + vec![lit(ScalarValue::Null)], + ]; + let rows = data.len(); + let values = MemorySourceConfig::try_new_as_values( + Arc::new(Schema::new(vec![Field::new("col0", DataType::Null, true)])), + data, + )?; + + assert_eq!( + values.statistics()?, + Statistics { + num_rows: Precision::Exact(rows), + total_byte_size: Precision::Exact(8), // not important + column_statistics: vec![ColumnStatistics { + null_count: Precision::Exact(rows), // there are only nulls + distinct_count: Precision::Absent, + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + },], + } + ); + + Ok(()) + } +} diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index c735c3108b3d9..8183d7b53244f 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + //! A table that uses the `ObjectStore` listing capability //! to get the list of files to process. @@ -24,6 +30,8 @@ pub mod file_meta; pub mod file_scan_config; pub mod file_sink_config; pub mod file_stream; +pub mod memory; +pub mod source; pub mod url; pub mod write; use chrono::TimeZone; diff --git a/datafusion/physical-plan/src/source.rs b/datafusion/datasource/src/source.rs similarity index 95% rename from datafusion/physical-plan/src/source.rs rename to datafusion/datasource/src/source.rs index 0c1dfddd26786..1b0d769029726 100644 --- a/datafusion/physical-plan/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -20,10 +20,12 @@ use std::fmt; use std::fmt::{Debug, Formatter}; use std::sync::Arc; -use crate::execution_plan::{Boundedness, EmissionType}; -use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; -use crate::projection::ProjectionExec; -use crate::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, +}; use datafusion_common::config::ConfigOptions; use datafusion_common::{Constraints, Statistics}; diff --git a/datafusion/doc/Cargo.toml b/datafusion/doc/Cargo.toml index c188bcb2a5352..fa316348a6daa 100644 --- a/datafusion/doc/Cargo.toml +++ b/datafusion/doc/Cargo.toml @@ -27,9 +27,11 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_doc" -path = "src/lib.rs" diff --git a/datafusion/doc/src/lib.rs b/datafusion/doc/src/lib.rs index 6940a8ef3ca26..68ed1e2352ca4 100644 --- a/datafusion/doc/src/lib.rs +++ b/datafusion/doc/src/lib.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + #[allow(rustdoc::broken_intra_doc_links)] /// Documentation for use by [`ScalarUDFImpl`](ScalarUDFImpl), /// [`AggregateUDFImpl`](AggregateUDFImpl) and [`WindowUDFImpl`](WindowUDFImpl) functions. diff --git a/datafusion/execution/Cargo.toml b/datafusion/execution/Cargo.toml index bb86868a82146..8f642f3384d2e 100644 --- a/datafusion/execution/Cargo.toml +++ b/datafusion/execution/Cargo.toml @@ -28,12 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_execution" -path = "src/lib.rs" [dependencies] arrow = { workspace = true } diff --git a/datafusion/execution/src/disk_manager.rs b/datafusion/execution/src/disk_manager.rs index 756da7ed5b468..caa62eefe14c7 100644 --- a/datafusion/execution/src/disk_manager.rs +++ b/datafusion/execution/src/disk_manager.rs @@ -250,7 +250,8 @@ mod tests { #[test] fn test_disk_manager_create_spill_folder() { - let config = DiskManagerConfig::new_specified(vec!["DOESNT_EXIST".into()]); + let dir = TempDir::new().unwrap(); + let config = DiskManagerConfig::new_specified(vec![dir.path().to_owned()]); DiskManager::try_new(config) .unwrap() diff --git a/datafusion/execution/src/lib.rs b/datafusion/execution/src/lib.rs index 317bd3203ab1b..a9e3a27f80356 100644 --- a/datafusion/execution/src/lib.rs +++ b/datafusion/execution/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/execution/src/memory_pool/mod.rs b/datafusion/execution/src/memory_pool/mod.rs index 5e49debac7e7f..71d40aeab53c7 100644 --- a/datafusion/execution/src/memory_pool/mod.rs +++ b/datafusion/execution/src/memory_pool/mod.rs @@ -55,7 +55,10 @@ pub use pool::*; /// "large" amounts of memory (proportional to number of input rows), such as /// `GroupByHashExec`. It does NOT track and limit memory used internally by /// other operators such as `DataSourceExec` or the `RecordBatch`es that flow -/// between operators. +/// between operators. Furthermore, operators should not reserve memory for the +/// batches they produce. Instead, if a parent operator needs to hold batches +/// from its children in memory for an extended period, it is the parent +/// operator's responsibility to reserve the necessary memory for those batches. /// /// In order to avoid allocating memory until the OS or the container system /// kills the process, DataFusion `ExecutionPlan`s (operators) that consume diff --git a/datafusion/execution/src/runtime_env.rs b/datafusion/execution/src/runtime_env.rs index 2b08b7ff9e889..95f14f485792a 100644 --- a/datafusion/execution/src/runtime_env.rs +++ b/datafusion/execution/src/runtime_env.rs @@ -27,7 +27,7 @@ use crate::{ }; use crate::cache::cache_manager::{CacheManager, CacheManagerConfig}; -use datafusion_common::{DataFusionError, Result}; +use datafusion_common::Result; use object_store::ObjectStore; use std::path::PathBuf; use std::sync::Arc; @@ -150,9 +150,7 @@ impl RuntimeEnv { /// registry. See [`ObjectStoreRegistry::get_store`] for more /// details. pub fn object_store(&self, url: impl AsRef) -> Result> { - self.object_store_registry - .get_store(url.as_ref()) - .map_err(DataFusionError::from) + self.object_store_registry.get_store(url.as_ref()) } } diff --git a/datafusion/expr-common/Cargo.toml b/datafusion/expr-common/Cargo.toml index abc78a9f084b4..14717dd78135d 100644 --- a/datafusion/expr-common/Cargo.toml +++ b/datafusion/expr-common/Cargo.toml @@ -27,14 +27,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_expr_common" -path = "src/lib.rs" - -[features] [dependencies] arrow = { workspace = true } diff --git a/datafusion/expr-common/src/interval_arithmetic.rs b/datafusion/expr-common/src/interval_arithmetic.rs index 7f20020c34572..9d00b45962bc2 100644 --- a/datafusion/expr-common/src/interval_arithmetic.rs +++ b/datafusion/expr-common/src/interval_arithmetic.rs @@ -17,12 +17,13 @@ //! Interval arithmetic library -use crate::operator::Operator; -use crate::type_coercion::binary::BinaryTypeCoercer; use std::borrow::Borrow; use std::fmt::{self, Display, Formatter}; use std::ops::{AddAssign, SubAssign}; +use crate::operator::Operator; +use crate::type_coercion::binary::{comparison_coercion_numeric, BinaryTypeCoercer}; + use arrow::compute::{cast_with_options, CastOptions}; use arrow::datatypes::{ DataType, IntervalDayTime, IntervalMonthDayNano, IntervalUnit, TimeUnit, @@ -168,7 +169,7 @@ macro_rules! value_transition { /// limits after any operation, they either become unbounded or they are fixed /// to the maximum/minimum value of the datatype, depending on the direction /// of the overflowing endpoint, opting for the safer choice. -/// +/// /// 4. **Floating-point special cases**: /// - `INF` values are converted to `NULL`s while constructing an interval to /// ensure consistency, with other data types. @@ -405,13 +406,18 @@ impl Interval { // There must be no way to create an interval whose endpoints have // different types. - assert!( + debug_assert!( lower_type == upper_type, "Interval bounds have different types: {lower_type} != {upper_type}" ); lower_type } + /// Checks if the interval is unbounded (on either side). + pub fn is_unbounded(&self) -> bool { + self.lower.is_null() || self.upper.is_null() + } + /// Casts this interval to `data_type` using `cast_options`. pub fn cast_to( &self, @@ -645,7 +651,7 @@ impl Interval { let upper = min_of_bounds(&self.upper, &rhs.upper); // New lower and upper bounds must always construct a valid interval. - assert!( + debug_assert!( (lower.is_null() || upper.is_null() || (lower <= upper)), "The intersection of two intervals can not be an invalid interval" ); @@ -653,26 +659,70 @@ impl Interval { Ok(Some(Self { lower, upper })) } - /// Decide if this interval certainly contains, possibly contains, or can't - /// contain a [`ScalarValue`] (`other`) by returning `[true, true]`, - /// `[false, true]` or `[false, false]` respectively. + /// Compute the union of this interval with the given interval. /// /// NOTE: This function only works with intervals of the same data type. /// Attempting to compare intervals of different data types will lead /// to an error. - pub fn contains_value>(&self, other: T) -> Result { + pub fn union>(&self, other: T) -> Result { let rhs = other.borrow(); if self.data_type().ne(&rhs.data_type()) { + return internal_err!( + "Cannot calculate the union of intervals with different data types, lhs:{}, rhs:{}", + self.data_type(), + rhs.data_type() + ); + }; + + let lower = if self.lower.is_null() + || (!rhs.lower.is_null() && self.lower <= rhs.lower) + { + self.lower.clone() + } else { + rhs.lower.clone() + }; + let upper = if self.upper.is_null() + || (!rhs.upper.is_null() && self.upper >= rhs.upper) + { + self.upper.clone() + } else { + rhs.upper.clone() + }; + + // New lower and upper bounds must always construct a valid interval. + debug_assert!( + (lower.is_null() || upper.is_null() || (lower <= upper)), + "The union of two intervals can not be an invalid interval" + ); + + Ok(Self { lower, upper }) + } + + /// Decide if this interval contains a [`ScalarValue`] (`other`) by returning `true` or `false`. + pub fn contains_value>(&self, other: T) -> Result { + let rhs = other.borrow(); + + let (lhs_lower, lhs_upper, rhs) = if self.data_type().eq(&rhs.data_type()) { + (&self.lower, &self.upper, rhs) + } else if let Some(common_type) = + comparison_coercion_numeric(&self.data_type(), &rhs.data_type()) + { + ( + &self.lower.cast_to(&common_type)?, + &self.upper.cast_to(&common_type)?, + &rhs.cast_to(&common_type)?, + ) + } else { return internal_err!( "Data types must be compatible for containment checks, lhs:{}, rhs:{}", self.data_type(), rhs.data_type() ); - } + }; // We only check the upper bound for a `None` value because `None` // values are less than `Some` values according to Rust. - Ok(&self.lower <= rhs && (self.upper.is_null() || rhs <= &self.upper)) + Ok(lhs_lower <= rhs && (lhs_upper.is_null() || rhs <= lhs_upper)) } /// Decide if this interval is a superset of, overlaps with, or @@ -825,6 +875,17 @@ impl Interval { } } + /// Computes the width of this interval; i.e. the difference between its + /// bounds. For unbounded intervals, this function will return a `NULL` + /// `ScalarValue` If the underlying data type doesn't support subtraction, + /// this function will return an error. + pub fn width(&self) -> Result { + let dt = self.data_type(); + let width_dt = + BinaryTypeCoercer::new(&dt, &Operator::Minus, &dt).get_result_type()?; + Ok(sub_bounds::(&width_dt, &self.upper, &self.lower)) + } + /// Returns the cardinality of this interval, which is the number of all /// distinct points inside it. This function returns `None` if: /// - The interval is unbounded from either side, or @@ -874,10 +935,10 @@ impl Interval { /// This method computes the arithmetic negation of the interval, reflecting /// it about the origin of the number line. This operation swaps and negates /// the lower and upper bounds of the interval. - pub fn arithmetic_negate(self) -> Result { + pub fn arithmetic_negate(&self) -> Result { Ok(Self { - lower: self.upper().clone().arithmetic_negate()?, - upper: self.lower().clone().arithmetic_negate()?, + lower: self.upper.arithmetic_negate()?, + upper: self.lower.arithmetic_negate()?, }) } } @@ -1119,11 +1180,11 @@ fn next_value_helper(value: ScalarValue) -> ScalarValue { match value { // f32/f64::NEG_INF/INF and f32/f64::NaN values should not emerge at this point. Float32(Some(val)) => { - assert!(val.is_finite(), "Non-standardized floating point usage"); + debug_assert!(val.is_finite(), "Non-standardized floating point usage"); Float32(Some(if INC { next_up(val) } else { next_down(val) })) } Float64(Some(val)) => { - assert!(val.is_finite(), "Non-standardized floating point usage"); + debug_assert!(val.is_finite(), "Non-standardized floating point usage"); Float64(Some(if INC { next_up(val) } else { next_down(val) })) } Int8(Some(val)) => Int8(Some(increment_decrement::(val))), @@ -1275,7 +1336,7 @@ pub fn satisfy_greater( } else { right.upper.clone() }; - + // No possibility to create an invalid interval: Ok(Some(( Interval::new(new_left_lower, left.upper.clone()), Interval::new(right.lower.clone(), new_right_upper), @@ -1868,6 +1929,7 @@ mod tests { }; use arrow::datatypes::DataType; + use datafusion_common::rounding::{next_down, next_up}; use datafusion_common::{Result, ScalarValue}; #[test] @@ -2532,6 +2594,126 @@ mod tests { Ok(()) } + #[test] + fn union_test() -> Result<()> { + let possible_cases = vec![ + ( + Interval::make(Some(1000_i64), None)?, + Interval::make::(None, None)?, + Interval::make_unbounded(&DataType::Int64)?, + ), + ( + Interval::make(Some(1000_i64), None)?, + Interval::make(None, Some(1000_i64))?, + Interval::make_unbounded(&DataType::Int64)?, + ), + ( + Interval::make(Some(1000_i64), None)?, + Interval::make(None, Some(2000_i64))?, + Interval::make_unbounded(&DataType::Int64)?, + ), + ( + Interval::make(Some(1000_i64), Some(2000_i64))?, + Interval::make(Some(1000_i64), None)?, + Interval::make(Some(1000_i64), None)?, + ), + ( + Interval::make(Some(1000_i64), Some(2000_i64))?, + Interval::make(Some(1000_i64), Some(1500_i64))?, + Interval::make(Some(1000_i64), Some(2000_i64))?, + ), + ( + Interval::make(Some(1000_i64), Some(2000_i64))?, + Interval::make(Some(500_i64), Some(1500_i64))?, + Interval::make(Some(500_i64), Some(2000_i64))?, + ), + ( + Interval::make::(None, None)?, + Interval::make::(None, None)?, + Interval::make::(None, None)?, + ), + ( + Interval::make(Some(1000_i64), None)?, + Interval::make(None, Some(0_i64))?, + Interval::make_unbounded(&DataType::Int64)?, + ), + ( + Interval::make(Some(1000_i64), None)?, + Interval::make(None, Some(999_i64))?, + Interval::make_unbounded(&DataType::Int64)?, + ), + ( + Interval::make(Some(1500_i64), Some(2000_i64))?, + Interval::make(Some(1000_i64), Some(1499_i64))?, + Interval::make(Some(1000_i64), Some(2000_i64))?, + ), + ( + Interval::make(Some(0_i64), Some(1000_i64))?, + Interval::make(Some(2000_i64), Some(3000_i64))?, + Interval::make(Some(0_i64), Some(3000_i64))?, + ), + ( + Interval::make(None, Some(2000_u64))?, + Interval::make(Some(500_u64), None)?, + Interval::make(Some(0_u64), None)?, + ), + ( + Interval::make(Some(0_u64), Some(0_u64))?, + Interval::make(Some(0_u64), None)?, + Interval::make(Some(0_u64), None)?, + ), + ( + Interval::make(Some(1000.0_f32), None)?, + Interval::make(None, Some(1000.0_f32))?, + Interval::make_unbounded(&DataType::Float32)?, + ), + ( + Interval::make(Some(1000.0_f32), Some(1500.0_f32))?, + Interval::make(Some(0.0_f32), Some(1500.0_f32))?, + Interval::make(Some(0.0_f32), Some(1500.0_f32))?, + ), + ( + Interval::try_new( + prev_value(ScalarValue::Float32(Some(1.0))), + prev_value(ScalarValue::Float32(Some(1.0))), + )?, + Interval::make(Some(1.0_f32), Some(1.0_f32))?, + Interval::try_new( + prev_value(ScalarValue::Float32(Some(1.0))), + ScalarValue::Float32(Some(1.0)), + )?, + ), + ( + Interval::try_new( + next_value(ScalarValue::Float32(Some(1.0))), + next_value(ScalarValue::Float32(Some(1.0))), + )?, + Interval::make(Some(1.0_f32), Some(1.0_f32))?, + Interval::try_new( + ScalarValue::Float32(Some(1.0)), + next_value(ScalarValue::Float32(Some(1.0))), + )?, + ), + ( + Interval::make(Some(-1000.0_f64), Some(1500.0_f64))?, + Interval::make(Some(-1500.0_f64), Some(2000.0_f64))?, + Interval::make(Some(-1500.0_f64), Some(2000.0_f64))?, + ), + ( + Interval::make(Some(16.0_f64), Some(32.0_f64))?, + Interval::make(Some(32.0_f64), Some(64.0_f64))?, + Interval::make(Some(16.0_f64), Some(64.0_f64))?, + ), + ]; + for (first, second, expected) in possible_cases { + println!("{}", first); + println!("{}", second); + assert_eq!(first.union(second)?, expected) + } + + Ok(()) + } + #[test] fn test_contains() -> Result<()> { let possible_cases = vec![ @@ -2594,6 +2776,43 @@ mod tests { Ok(()) } + #[test] + fn test_contains_value() -> Result<()> { + let possible_cases = vec![ + ( + Interval::make(Some(0), Some(100))?, + ScalarValue::Int32(Some(50)), + true, + ), + ( + Interval::make(Some(0), Some(100))?, + ScalarValue::Int32(Some(150)), + false, + ), + ( + Interval::make(Some(0), Some(100))?, + ScalarValue::Float64(Some(50.)), + true, + ), + ( + Interval::make(Some(0), Some(100))?, + ScalarValue::Float64(Some(next_down(100.))), + true, + ), + ( + Interval::make(Some(0), Some(100))?, + ScalarValue::Float64(Some(next_up(100.))), + false, + ), + ]; + + for (interval, value, expected) in possible_cases { + assert_eq!(interval.contains_value(value)?, expected) + } + + Ok(()) + } + #[test] fn test_add() -> Result<()> { let cases = vec![ @@ -3208,6 +3427,53 @@ mod tests { Ok(()) } + #[test] + fn test_width_of_intervals() -> Result<()> { + let intervals = [ + ( + Interval::make(Some(0.25_f64), Some(0.50_f64))?, + ScalarValue::from(0.25_f64), + ), + ( + Interval::make(Some(0.5_f64), Some(1.0_f64))?, + ScalarValue::from(0.5_f64), + ), + ( + Interval::make(Some(1.0_f64), Some(2.0_f64))?, + ScalarValue::from(1.0_f64), + ), + ( + Interval::make(Some(32.0_f64), Some(64.0_f64))?, + ScalarValue::from(32.0_f64), + ), + ( + Interval::make(Some(-0.50_f64), Some(-0.25_f64))?, + ScalarValue::from(0.25_f64), + ), + ( + Interval::make(Some(-32.0_f64), Some(-16.0_f64))?, + ScalarValue::from(16.0_f64), + ), + ( + Interval::make(Some(-0.50_f64), Some(0.25_f64))?, + ScalarValue::from(0.75_f64), + ), + ( + Interval::make(Some(-32.0_f64), Some(16.0_f64))?, + ScalarValue::from(48.0_f64), + ), + ( + Interval::make(Some(-32_i64), Some(16_i64))?, + ScalarValue::from(48_i64), + ), + ]; + for (interval, expected) in intervals { + assert_eq!(interval.width()?, expected); + } + + Ok(()) + } + #[test] fn test_cardinality_of_intervals() -> Result<()> { // In IEEE 754 standard for floating-point arithmetic, if we keep the sign and exponent fields same, diff --git a/datafusion/expr-common/src/lib.rs b/datafusion/expr-common/src/lib.rs index 179dd75ace85a..ee40038beb21e 100644 --- a/datafusion/expr-common/src/lib.rs +++ b/datafusion/expr-common/src/lib.rs @@ -19,10 +19,15 @@ //! //! This crate contains types and traits that are used by both Logical and Physical expressions. //! They are kept in their own crate to avoid physical expressions depending on logical expressions. -//! +//! //! //! [DataFusion]: +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] @@ -33,4 +38,5 @@ pub mod interval_arithmetic; pub mod operator; pub mod signature; pub mod sort_properties; +pub mod statistics; pub mod type_coercion; diff --git a/datafusion/expr-common/src/statistics.rs b/datafusion/expr-common/src/statistics.rs new file mode 100644 index 0000000000000..7e0bc88087efb --- /dev/null +++ b/datafusion/expr-common/src/statistics.rs @@ -0,0 +1,1620 @@ +// 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::f64::consts::LN_2; + +use crate::interval_arithmetic::{apply_operator, Interval}; +use crate::operator::Operator; +use crate::type_coercion::binary::binary_numeric_coercion; + +use arrow::array::ArrowNativeTypeOp; +use arrow::datatypes::DataType; +use datafusion_common::rounding::alter_fp_rounding_mode; +use datafusion_common::{internal_err, not_impl_err, Result, ScalarValue}; + +/// This object defines probabilistic distributions that encode uncertain +/// information about a single, scalar value. Currently, we support five core +/// statistical distributions. New variants will be added over time. +/// +/// This object is the lowest-level object in the statistics hierarchy, and it +/// is the main unit of calculus when evaluating expressions in a statistical +/// context. Notions like column and table statistics are built on top of this +/// object and the operations it supports. +#[derive(Clone, Debug, PartialEq)] +pub enum Distribution { + Uniform(UniformDistribution), + Exponential(ExponentialDistribution), + Gaussian(GaussianDistribution), + Bernoulli(BernoulliDistribution), + Generic(GenericDistribution), +} + +use Distribution::{Bernoulli, Exponential, Gaussian, Generic, Uniform}; + +impl Distribution { + /// Constructs a new [`Uniform`] distribution from the given [`Interval`]. + pub fn new_uniform(interval: Interval) -> Result { + UniformDistribution::try_new(interval).map(Uniform) + } + + /// Constructs a new [`Exponential`] distribution from the given rate/offset + /// pair, and validates the given parameters. + pub fn new_exponential( + rate: ScalarValue, + offset: ScalarValue, + positive_tail: bool, + ) -> Result { + ExponentialDistribution::try_new(rate, offset, positive_tail).map(Exponential) + } + + /// Constructs a new [`Gaussian`] distribution from the given mean/variance + /// pair, and validates the given parameters. + pub fn new_gaussian(mean: ScalarValue, variance: ScalarValue) -> Result { + GaussianDistribution::try_new(mean, variance).map(Gaussian) + } + + /// Constructs a new [`Bernoulli`] distribution from the given success + /// probability, and validates the given parameters. + pub fn new_bernoulli(p: ScalarValue) -> Result { + BernoulliDistribution::try_new(p).map(Bernoulli) + } + + /// Constructs a new [`Generic`] distribution from the given mean, median, + /// variance, and range values after validating the given parameters. + pub fn new_generic( + mean: ScalarValue, + median: ScalarValue, + variance: ScalarValue, + range: Interval, + ) -> Result { + GenericDistribution::try_new(mean, median, variance, range).map(Generic) + } + + /// Constructs a new [`Generic`] distribution from the given range. Other + /// parameters (mean, median and variance) are initialized with null values. + pub fn new_from_interval(range: Interval) -> Result { + let null = ScalarValue::try_from(range.data_type())?; + Distribution::new_generic(null.clone(), null.clone(), null, range) + } + + /// Extracts the mean value of this uncertain quantity, depending on its + /// distribution: + /// - A [`Uniform`] distribution's interval determines its mean value, which + /// is the arithmetic average of the interval endpoints. + /// - An [`Exponential`] distribution's mean is calculable by the formula + /// `offset + 1 / λ`, where `λ` is the (non-negative) rate. + /// - A [`Gaussian`] distribution contains the mean explicitly. + /// - A [`Bernoulli`] distribution's mean is equal to its success probability `p`. + /// - A [`Generic`] distribution _may_ have it explicitly, or this information + /// may be absent. + pub fn mean(&self) -> Result { + match &self { + Uniform(u) => u.mean(), + Exponential(e) => e.mean(), + Gaussian(g) => Ok(g.mean().clone()), + Bernoulli(b) => Ok(b.mean().clone()), + Generic(u) => Ok(u.mean().clone()), + } + } + + /// Extracts the median value of this uncertain quantity, depending on its + /// distribution: + /// - A [`Uniform`] distribution's interval determines its median value, which + /// is the arithmetic average of the interval endpoints. + /// - An [`Exponential`] distribution's median is calculable by the formula + /// `offset + ln(2) / λ`, where `λ` is the (non-negative) rate. + /// - A [`Gaussian`] distribution's median is equal to its mean, which is + /// specified explicitly. + /// - A [`Bernoulli`] distribution's median is `1` if `p > 0.5` and `0` + /// otherwise, where `p` is the success probability. + /// - A [`Generic`] distribution _may_ have it explicitly, or this information + /// may be absent. + pub fn median(&self) -> Result { + match &self { + Uniform(u) => u.median(), + Exponential(e) => e.median(), + Gaussian(g) => Ok(g.median().clone()), + Bernoulli(b) => b.median(), + Generic(u) => Ok(u.median().clone()), + } + } + + /// Extracts the variance value of this uncertain quantity, depending on + /// its distribution: + /// - A [`Uniform`] distribution's interval determines its variance value, which + /// is calculable by the formula `(upper - lower) ^ 2 / 12`. + /// - An [`Exponential`] distribution's variance is calculable by the formula + /// `1 / (λ ^ 2)`, where `λ` is the (non-negative) rate. + /// - A [`Gaussian`] distribution's variance is specified explicitly. + /// - A [`Bernoulli`] distribution's median is given by the formula `p * (1 - p)` + /// where `p` is the success probability. + /// - A [`Generic`] distribution _may_ have it explicitly, or this information + /// may be absent. + pub fn variance(&self) -> Result { + match &self { + Uniform(u) => u.variance(), + Exponential(e) => e.variance(), + Gaussian(g) => Ok(g.variance.clone()), + Bernoulli(b) => b.variance(), + Generic(u) => Ok(u.variance.clone()), + } + } + + /// Extracts the range of this uncertain quantity, depending on its + /// distribution: + /// - A [`Uniform`] distribution's range is simply its interval. + /// - An [`Exponential`] distribution's range is `[offset, +∞)`. + /// - A [`Gaussian`] distribution's range is unbounded. + /// - A [`Bernoulli`] distribution's range is [`Interval::UNCERTAIN`], if + /// `p` is neither `0` nor `1`. Otherwise, it is [`Interval::CERTAINLY_FALSE`] + /// and [`Interval::CERTAINLY_TRUE`], respectively. + /// - A [`Generic`] distribution is unbounded by default, but more information + /// may be present. + pub fn range(&self) -> Result { + match &self { + Uniform(u) => Ok(u.range().clone()), + Exponential(e) => e.range(), + Gaussian(g) => g.range(), + Bernoulli(b) => Ok(b.range()), + Generic(u) => Ok(u.range().clone()), + } + } + + /// Returns the data type of the statistical parameters comprising this + /// distribution. + pub fn data_type(&self) -> DataType { + match &self { + Uniform(u) => u.data_type(), + Exponential(e) => e.data_type(), + Gaussian(g) => g.data_type(), + Bernoulli(b) => b.data_type(), + Generic(u) => u.data_type(), + } + } + + pub fn target_type(args: &[&ScalarValue]) -> Result { + let mut arg_types = args + .iter() + .filter(|&&arg| (arg != &ScalarValue::Null)) + .map(|&arg| arg.data_type()); + + let Some(dt) = arg_types.next().map_or_else( + || Some(DataType::Null), + |first| { + arg_types + .try_fold(first, |target, arg| binary_numeric_coercion(&target, &arg)) + }, + ) else { + return internal_err!("Can only evaluate statistics for numeric types"); + }; + Ok(dt) + } +} + +/// Uniform distribution, represented by its range. If the given range extends +/// towards infinity, the distribution will be improper -- which is OK. For a +/// more in-depth discussion, see: +/// +/// +/// +#[derive(Clone, Debug, PartialEq)] +pub struct UniformDistribution { + interval: Interval, +} + +/// Exponential distribution with an optional shift. The probability density +/// function (PDF) is defined as follows: +/// +/// For a positive tail (when `positive_tail` is `true`): +/// +/// `f(x; λ, offset) = λ exp(-λ (x - offset)) for x ≥ offset` +/// +/// For a negative tail (when `positive_tail` is `false`): +/// +/// `f(x; λ, offset) = λ exp(-λ (offset - x)) for x ≤ offset` +/// +/// +/// In both cases, the PDF is `0` outside the specified domain. +/// +/// For more information, see: +/// +/// +#[derive(Clone, Debug, PartialEq)] +pub struct ExponentialDistribution { + rate: ScalarValue, + offset: ScalarValue, + /// Indicates whether the exponential distribution has a positive tail; + /// i.e. it extends towards positive infinity. + positive_tail: bool, +} + +/// Gaussian (normal) distribution, represented by its mean and variance. +/// For a more in-depth discussion, see: +/// +/// +#[derive(Clone, Debug, PartialEq)] +pub struct GaussianDistribution { + mean: ScalarValue, + variance: ScalarValue, +} + +/// Bernoulli distribution with success probability `p`. If `p` has a null value, +/// the success probability is unknown. For a more in-depth discussion, see: +/// +/// +#[derive(Clone, Debug, PartialEq)] +pub struct BernoulliDistribution { + p: ScalarValue, +} + +/// A generic distribution whose functional form is not available, which is +/// approximated via some summary statistics. For a more in-depth discussion, see: +/// +/// +#[derive(Clone, Debug, PartialEq)] +pub struct GenericDistribution { + mean: ScalarValue, + median: ScalarValue, + variance: ScalarValue, + range: Interval, +} + +impl UniformDistribution { + fn try_new(interval: Interval) -> Result { + if interval.data_type().eq(&DataType::Boolean) { + return internal_err!( + "Construction of a boolean `Uniform` distribution is prohibited, create a `Bernoulli` distribution instead." + ); + } + + Ok(Self { interval }) + } + + pub fn data_type(&self) -> DataType { + self.interval.data_type() + } + + /// Computes the mean value of this distribution. In case of improper + /// distributions (i.e. when the range is unbounded), the function returns + /// a `NULL` `ScalarValue`. + pub fn mean(&self) -> Result { + // TODO: Should we ensure that this always returns a real number data type? + let dt = self.data_type(); + let two = ScalarValue::from(2).cast_to(&dt)?; + let result = self + .interval + .lower() + .add_checked(self.interval.upper())? + .div(two); + debug_assert!( + !self.interval.is_unbounded() || result.as_ref().is_ok_and(|r| r.is_null()) + ); + result + } + + pub fn median(&self) -> Result { + self.mean() + } + + /// Computes the variance value of this distribution. In case of improper + /// distributions (i.e. when the range is unbounded), the function returns + /// a `NULL` `ScalarValue`. + pub fn variance(&self) -> Result { + // TODO: Should we ensure that this always returns a real number data type? + let width = self.interval.width()?; + let dt = width.data_type(); + let twelve = ScalarValue::from(12).cast_to(&dt)?; + let result = width.mul_checked(&width)?.div(twelve); + debug_assert!( + !self.interval.is_unbounded() || result.as_ref().is_ok_and(|r| r.is_null()) + ); + result + } + + pub fn range(&self) -> &Interval { + &self.interval + } +} + +impl ExponentialDistribution { + fn try_new( + rate: ScalarValue, + offset: ScalarValue, + positive_tail: bool, + ) -> Result { + let dt = rate.data_type(); + if offset.data_type() != dt { + internal_err!("Rate and offset must have the same data type") + } else if offset.is_null() { + internal_err!("Offset of an `ExponentialDistribution` cannot be null") + } else if rate.is_null() { + internal_err!("Rate of an `ExponentialDistribution` cannot be null") + } else if rate.le(&ScalarValue::new_zero(&dt)?) { + internal_err!("Rate of an `ExponentialDistribution` must be positive") + } else { + Ok(Self { + rate, + offset, + positive_tail, + }) + } + } + + pub fn data_type(&self) -> DataType { + self.rate.data_type() + } + + pub fn rate(&self) -> &ScalarValue { + &self.rate + } + + pub fn offset(&self) -> &ScalarValue { + &self.offset + } + + pub fn positive_tail(&self) -> bool { + self.positive_tail + } + + pub fn mean(&self) -> Result { + // TODO: Should we ensure that this always returns a real number data type? + let one = ScalarValue::new_one(&self.data_type())?; + let tail_mean = one.div(&self.rate)?; + if self.positive_tail { + self.offset.add_checked(tail_mean) + } else { + self.offset.sub_checked(tail_mean) + } + } + + pub fn median(&self) -> Result { + // TODO: Should we ensure that this always returns a real number data type? + let ln_two = ScalarValue::from(LN_2).cast_to(&self.data_type())?; + let tail_median = ln_two.div(&self.rate)?; + if self.positive_tail { + self.offset.add_checked(tail_median) + } else { + self.offset.sub_checked(tail_median) + } + } + + pub fn variance(&self) -> Result { + // TODO: Should we ensure that this always returns a real number data type? + let one = ScalarValue::new_one(&self.data_type())?; + let rate_squared = self.rate.mul_checked(&self.rate)?; + one.div(rate_squared) + } + + pub fn range(&self) -> Result { + let end = ScalarValue::try_from(&self.data_type())?; + if self.positive_tail { + Interval::try_new(self.offset.clone(), end) + } else { + Interval::try_new(end, self.offset.clone()) + } + } +} + +impl GaussianDistribution { + fn try_new(mean: ScalarValue, variance: ScalarValue) -> Result { + let dt = mean.data_type(); + if variance.data_type() != dt { + internal_err!("Mean and variance must have the same data type") + } else if variance.is_null() { + internal_err!("Variance of a `GaussianDistribution` cannot be null") + } else if variance.lt(&ScalarValue::new_zero(&dt)?) { + internal_err!("Variance of a `GaussianDistribution` must be positive") + } else { + Ok(Self { mean, variance }) + } + } + + pub fn data_type(&self) -> DataType { + self.mean.data_type() + } + + pub fn mean(&self) -> &ScalarValue { + &self.mean + } + + pub fn variance(&self) -> &ScalarValue { + &self.variance + } + + pub fn median(&self) -> &ScalarValue { + self.mean() + } + + pub fn range(&self) -> Result { + Interval::make_unbounded(&self.data_type()) + } +} + +impl BernoulliDistribution { + fn try_new(p: ScalarValue) -> Result { + if p.is_null() { + Ok(Self { p }) + } else { + let dt = p.data_type(); + let zero = ScalarValue::new_zero(&dt)?; + let one = ScalarValue::new_one(&dt)?; + if p.ge(&zero) && p.le(&one) { + Ok(Self { p }) + } else { + internal_err!( + "Success probability of a `BernoulliDistribution` must be in [0, 1]" + ) + } + } + } + + pub fn data_type(&self) -> DataType { + self.p.data_type() + } + + pub fn p_value(&self) -> &ScalarValue { + &self.p + } + + pub fn mean(&self) -> &ScalarValue { + &self.p + } + + /// Computes the median value of this distribution. In case of an unknown + /// success probability, the function returns a `NULL` `ScalarValue`. + pub fn median(&self) -> Result { + let dt = self.data_type(); + if self.p.is_null() { + ScalarValue::try_from(&dt) + } else { + let one = ScalarValue::new_one(&dt)?; + if one.sub_checked(&self.p)?.lt(&self.p) { + ScalarValue::new_one(&dt) + } else { + ScalarValue::new_zero(&dt) + } + } + } + + /// Computes the variance value of this distribution. In case of an unknown + /// success probability, the function returns a `NULL` `ScalarValue`. + pub fn variance(&self) -> Result { + let dt = self.data_type(); + let one = ScalarValue::new_one(&dt)?; + let result = one.sub_checked(&self.p)?.mul_checked(&self.p); + debug_assert!(!self.p.is_null() || result.as_ref().is_ok_and(|r| r.is_null())); + result + } + + pub fn range(&self) -> Interval { + let dt = self.data_type(); + // Unwraps are safe as the constructor guarantees that the data type + // supports zero and one values. + if ScalarValue::new_zero(&dt).unwrap().eq(&self.p) { + Interval::CERTAINLY_FALSE + } else if ScalarValue::new_one(&dt).unwrap().eq(&self.p) { + Interval::CERTAINLY_TRUE + } else { + Interval::UNCERTAIN + } + } +} + +impl GenericDistribution { + fn try_new( + mean: ScalarValue, + median: ScalarValue, + variance: ScalarValue, + range: Interval, + ) -> Result { + if range.data_type().eq(&DataType::Boolean) { + return internal_err!( + "Construction of a boolean `Generic` distribution is prohibited, create a `Bernoulli` distribution instead." + ); + } + + let validate_location = |m: &ScalarValue| -> Result { + // Checks whether the given location estimate is within the range. + if m.is_null() { + Ok(true) + } else { + range.contains_value(m) + } + }; + + if !validate_location(&mean)? + || !validate_location(&median)? + || (!variance.is_null() + && variance.lt(&ScalarValue::new_zero(&variance.data_type())?)) + { + internal_err!("Tried to construct an invalid `GenericDistribution` instance") + } else { + Ok(Self { + mean, + median, + variance, + range, + }) + } + } + + pub fn data_type(&self) -> DataType { + self.mean.data_type() + } + + pub fn mean(&self) -> &ScalarValue { + &self.mean + } + + pub fn median(&self) -> &ScalarValue { + &self.median + } + + pub fn variance(&self) -> &ScalarValue { + &self.variance + } + + pub fn range(&self) -> &Interval { + &self.range + } +} + +/// This function takes a logical operator and two Bernoulli distributions, +/// and it returns a new Bernoulli distribution that represents the result of +/// the operation. Currently, only `AND` and `OR` operations are supported. +pub fn combine_bernoullis( + op: &Operator, + left: &BernoulliDistribution, + right: &BernoulliDistribution, +) -> Result { + let left_p = left.p_value(); + let right_p = right.p_value(); + match op { + Operator::And => match (left_p.is_null(), right_p.is_null()) { + (false, false) => { + BernoulliDistribution::try_new(left_p.mul_checked(right_p)?) + } + (false, true) if left_p.eq(&ScalarValue::new_zero(&left_p.data_type())?) => { + Ok(left.clone()) + } + (true, false) + if right_p.eq(&ScalarValue::new_zero(&right_p.data_type())?) => + { + Ok(right.clone()) + } + _ => { + let dt = Distribution::target_type(&[left_p, right_p])?; + BernoulliDistribution::try_new(ScalarValue::try_from(&dt)?) + } + }, + Operator::Or => match (left_p.is_null(), right_p.is_null()) { + (false, false) => { + let sum = left_p.add_checked(right_p)?; + let product = left_p.mul_checked(right_p)?; + let or_success = sum.sub_checked(product)?; + BernoulliDistribution::try_new(or_success) + } + (false, true) if left_p.eq(&ScalarValue::new_one(&left_p.data_type())?) => { + Ok(left.clone()) + } + (true, false) if right_p.eq(&ScalarValue::new_one(&right_p.data_type())?) => { + Ok(right.clone()) + } + _ => { + let dt = Distribution::target_type(&[left_p, right_p])?; + BernoulliDistribution::try_new(ScalarValue::try_from(&dt)?) + } + }, + _ => { + not_impl_err!("Statistical evaluation only supports AND and OR operators") + } + } +} + +/// Applies the given operation to the given Gaussian distributions. Currently, +/// this function handles only addition and subtraction operations. If the +/// result is not a Gaussian random variable, it returns `None`. For details, +/// see: +/// +/// +pub fn combine_gaussians( + op: &Operator, + left: &GaussianDistribution, + right: &GaussianDistribution, +) -> Result> { + match op { + Operator::Plus => GaussianDistribution::try_new( + left.mean().add_checked(right.mean())?, + left.variance().add_checked(right.variance())?, + ) + .map(Some), + Operator::Minus => GaussianDistribution::try_new( + left.mean().sub_checked(right.mean())?, + left.variance().add_checked(right.variance())?, + ) + .map(Some), + _ => Ok(None), + } +} + +/// Creates a new `Bernoulli` distribution by computing the resulting probability. +/// Expects `op` to be a comparison operator, with `left` and `right` having +/// numeric distributions. The resulting distribution has the `Float64` data +/// type. +pub fn create_bernoulli_from_comparison( + op: &Operator, + left: &Distribution, + right: &Distribution, +) -> Result { + match (left, right) { + (Uniform(left), Uniform(right)) => { + match op { + Operator::Eq | Operator::NotEq => { + let (li, ri) = (left.range(), right.range()); + if let Some(intersection) = li.intersect(ri)? { + // If the ranges are not disjoint, calculate the probability + // of equality using cardinalities: + if let (Some(lc), Some(rc), Some(ic)) = ( + li.cardinality(), + ri.cardinality(), + intersection.cardinality(), + ) { + // Avoid overflow by widening the type temporarily: + let pairs = ((lc as u128) * (rc as u128)) as f64; + let p = (ic as f64).div_checked(pairs)?; + // Alternative approach that may be more stable: + // let p = (ic as f64) + // .div_checked(lc as f64)? + // .div_checked(rc as f64)?; + + let mut p_value = ScalarValue::from(p); + if op == &Operator::NotEq { + let one = ScalarValue::from(1.0); + p_value = alter_fp_rounding_mode::( + &one, + &p_value, + |lhs, rhs| lhs.sub_checked(rhs), + )?; + }; + return Distribution::new_bernoulli(p_value); + } + } else if op == &Operator::Eq { + // If the ranges are disjoint, probability of equality is 0. + return Distribution::new_bernoulli(ScalarValue::from(0.0)); + } else { + // If the ranges are disjoint, probability of not-equality is 1. + return Distribution::new_bernoulli(ScalarValue::from(1.0)); + } + } + Operator::Lt | Operator::LtEq | Operator::Gt | Operator::GtEq => { + // TODO: We can handle inequality operators and calculate a + // `p` value instead of falling back to an unknown Bernoulli + // distribution. Note that the strict and non-strict inequalities + // may require slightly different logic in case of real vs. + // integral data types. + } + _ => {} + } + } + (Gaussian(_), Gaussian(_)) => { + // TODO: We can handle Gaussian comparisons and calculate a `p` value + // instead of falling back to an unknown Bernoulli distribution. + } + _ => {} + } + let (li, ri) = (left.range()?, right.range()?); + let range_evaluation = apply_operator(op, &li, &ri)?; + if range_evaluation.eq(&Interval::CERTAINLY_FALSE) { + Distribution::new_bernoulli(ScalarValue::from(0.0)) + } else if range_evaluation.eq(&Interval::CERTAINLY_TRUE) { + Distribution::new_bernoulli(ScalarValue::from(1.0)) + } else if range_evaluation.eq(&Interval::UNCERTAIN) { + Distribution::new_bernoulli(ScalarValue::try_from(&DataType::Float64)?) + } else { + internal_err!("This function must be called with a comparison operator") + } +} + +/// Creates a new [`Generic`] distribution that represents the result of the +/// given binary operation on two unknown quantities represented by their +/// [`Distribution`] objects. The function computes the mean, median and +/// variance if possible. +pub fn new_generic_from_binary_op( + op: &Operator, + left: &Distribution, + right: &Distribution, +) -> Result { + Distribution::new_generic( + compute_mean(op, left, right)?, + compute_median(op, left, right)?, + compute_variance(op, left, right)?, + apply_operator(op, &left.range()?, &right.range()?)?, + ) +} + +/// Computes the mean value for the result of the given binary operation on +/// two unknown quantities represented by their [`Distribution`] objects. +pub fn compute_mean( + op: &Operator, + left: &Distribution, + right: &Distribution, +) -> Result { + let (left_mean, right_mean) = (left.mean()?, right.mean()?); + + match op { + Operator::Plus => return left_mean.add_checked(right_mean), + Operator::Minus => return left_mean.sub_checked(right_mean), + // Note the independence assumption below: + Operator::Multiply => return left_mean.mul_checked(right_mean), + // TODO: We can calculate the mean for division when we support reciprocals, + // or know the distributions of the operands. For details, see: + // + // + // + // + // Fall back to an unknown mean value for division: + Operator::Divide => {} + // Fall back to an unknown mean value for other cases: + _ => {} + } + let target_type = Distribution::target_type(&[&left_mean, &right_mean])?; + ScalarValue::try_from(target_type) +} + +/// Computes the median value for the result of the given binary operation on +/// two unknown quantities represented by its [`Distribution`] objects. Currently, +/// the median is calculable only for addition and subtraction operations on: +/// - [`Uniform`] and [`Uniform`] distributions, and +/// - [`Gaussian`] and [`Gaussian`] distributions. +pub fn compute_median( + op: &Operator, + left: &Distribution, + right: &Distribution, +) -> Result { + match (left, right) { + (Uniform(lu), Uniform(ru)) => { + let (left_median, right_median) = (lu.median()?, ru.median()?); + // Under the independence assumption, the result is a symmetric + // triangular distribution, so we can simply add/subtract the + // median values: + match op { + Operator::Plus => return left_median.add_checked(right_median), + Operator::Minus => return left_median.sub_checked(right_median), + // Fall back to an unknown median value for other cases: + _ => {} + } + } + // Under the independence assumption, the result is another Gaussian + // distribution, so we can simply add/subtract the median values: + (Gaussian(lg), Gaussian(rg)) => match op { + Operator::Plus => return lg.mean().add_checked(rg.mean()), + Operator::Minus => return lg.mean().sub_checked(rg.mean()), + // Fall back to an unknown median value for other cases: + _ => {} + }, + // Fall back to an unknown median value for other cases: + _ => {} + } + + let (left_median, right_median) = (left.median()?, right.median()?); + let target_type = Distribution::target_type(&[&left_median, &right_median])?; + ScalarValue::try_from(target_type) +} + +/// Computes the variance value for the result of the given binary operation on +/// two unknown quantities represented by their [`Distribution`] objects. +pub fn compute_variance( + op: &Operator, + left: &Distribution, + right: &Distribution, +) -> Result { + let (left_variance, right_variance) = (left.variance()?, right.variance()?); + + match op { + // Note the independence assumption below: + Operator::Plus => return left_variance.add_checked(right_variance), + // Note the independence assumption below: + Operator::Minus => return left_variance.add_checked(right_variance), + // Note the independence assumption below: + Operator::Multiply => { + // For more details, along with an explanation of the formula below, see: + // + // + let (left_mean, right_mean) = (left.mean()?, right.mean()?); + let left_mean_sq = left_mean.mul_checked(&left_mean)?; + let right_mean_sq = right_mean.mul_checked(&right_mean)?; + let left_sos = left_variance.add_checked(&left_mean_sq)?; + let right_sos = right_variance.add_checked(&right_mean_sq)?; + let pos = left_mean_sq.mul_checked(right_mean_sq)?; + return left_sos.mul_checked(right_sos)?.sub_checked(pos); + } + // TODO: We can calculate the variance for division when we support reciprocals, + // or know the distributions of the operands. For details, see: + // + // + // + // + // Fall back to an unknown variance value for division: + Operator::Divide => {} + // Fall back to an unknown variance value for other cases: + _ => {} + } + let target_type = Distribution::target_type(&[&left_variance, &right_variance])?; + ScalarValue::try_from(target_type) +} + +#[cfg(test)] +mod tests { + use super::{ + combine_bernoullis, combine_gaussians, compute_mean, compute_median, + compute_variance, create_bernoulli_from_comparison, new_generic_from_binary_op, + BernoulliDistribution, Distribution, GaussianDistribution, UniformDistribution, + }; + use crate::interval_arithmetic::{apply_operator, Interval}; + use crate::operator::Operator; + + use arrow::datatypes::DataType; + use datafusion_common::{HashSet, Result, ScalarValue}; + + #[test] + fn uniform_dist_is_valid_test() -> Result<()> { + assert_eq!( + Distribution::new_uniform(Interval::make_zero(&DataType::Int8)?)?, + Distribution::Uniform(UniformDistribution { + interval: Interval::make_zero(&DataType::Int8)?, + }) + ); + + assert!(Distribution::new_uniform(Interval::UNCERTAIN).is_err()); + Ok(()) + } + + #[test] + fn exponential_dist_is_valid_test() { + // This array collects test cases of the form (distribution, validity). + let exponentials = vec![ + ( + Distribution::new_exponential(ScalarValue::Null, ScalarValue::Null, true), + false, + ), + ( + Distribution::new_exponential( + ScalarValue::from(0_f32), + ScalarValue::from(1_f32), + true, + ), + false, + ), + ( + Distribution::new_exponential( + ScalarValue::from(100_f32), + ScalarValue::from(1_f32), + true, + ), + true, + ), + ( + Distribution::new_exponential( + ScalarValue::from(-100_f32), + ScalarValue::from(1_f32), + true, + ), + false, + ), + ]; + for case in exponentials { + assert_eq!(case.0.is_ok(), case.1); + } + } + + #[test] + fn gaussian_dist_is_valid_test() { + // This array collects test cases of the form (distribution, validity). + let gaussians = vec![ + ( + Distribution::new_gaussian(ScalarValue::Null, ScalarValue::Null), + false, + ), + ( + Distribution::new_gaussian( + ScalarValue::from(0_f32), + ScalarValue::from(0_f32), + ), + true, + ), + ( + Distribution::new_gaussian( + ScalarValue::from(0_f32), + ScalarValue::from(0.5_f32), + ), + true, + ), + ( + Distribution::new_gaussian( + ScalarValue::from(0_f32), + ScalarValue::from(-0.5_f32), + ), + false, + ), + ]; + for case in gaussians { + assert_eq!(case.0.is_ok(), case.1); + } + } + + #[test] + fn bernoulli_dist_is_valid_test() { + // This array collects test cases of the form (distribution, validity). + let bernoullis = vec![ + (Distribution::new_bernoulli(ScalarValue::Null), true), + (Distribution::new_bernoulli(ScalarValue::from(0.)), true), + (Distribution::new_bernoulli(ScalarValue::from(0.25)), true), + (Distribution::new_bernoulli(ScalarValue::from(1.)), true), + (Distribution::new_bernoulli(ScalarValue::from(11.)), false), + (Distribution::new_bernoulli(ScalarValue::from(-11.)), false), + (Distribution::new_bernoulli(ScalarValue::from(0_i64)), true), + (Distribution::new_bernoulli(ScalarValue::from(1_i64)), true), + ( + Distribution::new_bernoulli(ScalarValue::from(11_i64)), + false, + ), + ( + Distribution::new_bernoulli(ScalarValue::from(-11_i64)), + false, + ), + ]; + for case in bernoullis { + assert_eq!(case.0.is_ok(), case.1); + } + } + + #[test] + fn generic_dist_is_valid_test() -> Result<()> { + // This array collects test cases of the form (distribution, validity). + let generic_dists = vec![ + // Using a boolean range to construct a Generic distribution is prohibited. + ( + Distribution::new_generic( + ScalarValue::Null, + ScalarValue::Null, + ScalarValue::Null, + Interval::UNCERTAIN, + ), + false, + ), + ( + Distribution::new_generic( + ScalarValue::Null, + ScalarValue::Null, + ScalarValue::Null, + Interval::make_zero(&DataType::Float32)?, + ), + true, + ), + ( + Distribution::new_generic( + ScalarValue::from(0_f32), + ScalarValue::Float32(None), + ScalarValue::Float32(None), + Interval::make_zero(&DataType::Float32)?, + ), + true, + ), + ( + Distribution::new_generic( + ScalarValue::Float64(None), + ScalarValue::from(0.), + ScalarValue::Float64(None), + Interval::make_zero(&DataType::Float32)?, + ), + true, + ), + ( + Distribution::new_generic( + ScalarValue::from(-10_f32), + ScalarValue::Float32(None), + ScalarValue::Float32(None), + Interval::make_zero(&DataType::Float32)?, + ), + false, + ), + ( + Distribution::new_generic( + ScalarValue::Float32(None), + ScalarValue::from(10_f32), + ScalarValue::Float32(None), + Interval::make_zero(&DataType::Float32)?, + ), + false, + ), + ( + Distribution::new_generic( + ScalarValue::Null, + ScalarValue::Null, + ScalarValue::Null, + Interval::make_zero(&DataType::Float32)?, + ), + true, + ), + ( + Distribution::new_generic( + ScalarValue::from(0), + ScalarValue::from(0), + ScalarValue::Int32(None), + Interval::make_zero(&DataType::Int32)?, + ), + true, + ), + ( + Distribution::new_generic( + ScalarValue::from(0_f32), + ScalarValue::from(0_f32), + ScalarValue::Float32(None), + Interval::make_zero(&DataType::Float32)?, + ), + true, + ), + ( + Distribution::new_generic( + ScalarValue::from(50.), + ScalarValue::from(50.), + ScalarValue::Float64(None), + Interval::make(Some(0.), Some(100.))?, + ), + true, + ), + ( + Distribution::new_generic( + ScalarValue::from(50.), + ScalarValue::from(50.), + ScalarValue::Float64(None), + Interval::make(Some(-100.), Some(0.))?, + ), + false, + ), + ( + Distribution::new_generic( + ScalarValue::Float64(None), + ScalarValue::Float64(None), + ScalarValue::from(1.), + Interval::make_zero(&DataType::Float64)?, + ), + true, + ), + ( + Distribution::new_generic( + ScalarValue::Float64(None), + ScalarValue::Float64(None), + ScalarValue::from(-1.), + Interval::make_zero(&DataType::Float64)?, + ), + false, + ), + ]; + for case in generic_dists { + assert_eq!(case.0.is_ok(), case.1, "{:?}", case.0); + } + + Ok(()) + } + + #[test] + fn mean_extraction_test() -> Result<()> { + // This array collects test cases of the form (distribution, mean value). + let dists = vec![ + ( + Distribution::new_uniform(Interval::make_zero(&DataType::Int64)?), + ScalarValue::from(0_i64), + ), + ( + Distribution::new_uniform(Interval::make_zero(&DataType::Float64)?), + ScalarValue::from(0.), + ), + ( + Distribution::new_uniform(Interval::make(Some(1), Some(100))?), + ScalarValue::from(50), + ), + ( + Distribution::new_uniform(Interval::make(Some(-100), Some(-1))?), + ScalarValue::from(-50), + ), + ( + Distribution::new_uniform(Interval::make(Some(-100), Some(100))?), + ScalarValue::from(0), + ), + ( + Distribution::new_exponential( + ScalarValue::from(2.), + ScalarValue::from(0.), + true, + ), + ScalarValue::from(0.5), + ), + ( + Distribution::new_exponential( + ScalarValue::from(2.), + ScalarValue::from(1.), + true, + ), + ScalarValue::from(1.5), + ), + ( + Distribution::new_gaussian(ScalarValue::from(0.), ScalarValue::from(1.)), + ScalarValue::from(0.), + ), + ( + Distribution::new_gaussian( + ScalarValue::from(-2.), + ScalarValue::from(0.5), + ), + ScalarValue::from(-2.), + ), + ( + Distribution::new_bernoulli(ScalarValue::from(0.5)), + ScalarValue::from(0.5), + ), + ( + Distribution::new_generic( + ScalarValue::from(42.), + ScalarValue::from(42.), + ScalarValue::Float64(None), + Interval::make(Some(25.), Some(50.))?, + ), + ScalarValue::from(42.), + ), + ]; + + for case in dists { + assert_eq!(case.0?.mean()?, case.1); + } + + Ok(()) + } + + #[test] + fn median_extraction_test() -> Result<()> { + // This array collects test cases of the form (distribution, median value). + let dists = vec![ + ( + Distribution::new_uniform(Interval::make_zero(&DataType::Int64)?), + ScalarValue::from(0_i64), + ), + ( + Distribution::new_uniform(Interval::make(Some(25.), Some(75.))?), + ScalarValue::from(50.), + ), + ( + Distribution::new_exponential( + ScalarValue::from(2_f64.ln()), + ScalarValue::from(0.), + true, + ), + ScalarValue::from(1.), + ), + ( + Distribution::new_gaussian(ScalarValue::from(2.), ScalarValue::from(1.)), + ScalarValue::from(2.), + ), + ( + Distribution::new_bernoulli(ScalarValue::from(0.25)), + ScalarValue::from(0.), + ), + ( + Distribution::new_bernoulli(ScalarValue::from(0.75)), + ScalarValue::from(1.), + ), + ( + Distribution::new_gaussian(ScalarValue::from(2.), ScalarValue::from(1.)), + ScalarValue::from(2.), + ), + ( + Distribution::new_generic( + ScalarValue::from(12.), + ScalarValue::from(12.), + ScalarValue::Float64(None), + Interval::make(Some(0.), Some(25.))?, + ), + ScalarValue::from(12.), + ), + ]; + + for case in dists { + assert_eq!(case.0?.median()?, case.1); + } + + Ok(()) + } + + #[test] + fn variance_extraction_test() -> Result<()> { + // This array collects test cases of the form (distribution, variance value). + let dists = vec![ + ( + Distribution::new_uniform(Interval::make(Some(0.), Some(12.))?), + ScalarValue::from(12.), + ), + ( + Distribution::new_exponential( + ScalarValue::from(10.), + ScalarValue::from(0.), + true, + ), + ScalarValue::from(0.01), + ), + ( + Distribution::new_gaussian(ScalarValue::from(0.), ScalarValue::from(1.)), + ScalarValue::from(1.), + ), + ( + Distribution::new_bernoulli(ScalarValue::from(0.5)), + ScalarValue::from(0.25), + ), + ( + Distribution::new_generic( + ScalarValue::Float64(None), + ScalarValue::Float64(None), + ScalarValue::from(0.02), + Interval::make_zero(&DataType::Float64)?, + ), + ScalarValue::from(0.02), + ), + ]; + + for case in dists { + assert_eq!(case.0?.variance()?, case.1); + } + + Ok(()) + } + + #[test] + fn test_calculate_generic_properties_gauss_gauss() -> Result<()> { + let dist_a = + Distribution::new_gaussian(ScalarValue::from(10.), ScalarValue::from(0.0))?; + let dist_b = + Distribution::new_gaussian(ScalarValue::from(20.), ScalarValue::from(0.0))?; + + let test_data = vec![ + // Mean: + ( + compute_mean(&Operator::Plus, &dist_a, &dist_b)?, + ScalarValue::from(30.), + ), + ( + compute_mean(&Operator::Minus, &dist_a, &dist_b)?, + ScalarValue::from(-10.), + ), + // Median: + ( + compute_median(&Operator::Plus, &dist_a, &dist_b)?, + ScalarValue::from(30.), + ), + ( + compute_median(&Operator::Minus, &dist_a, &dist_b)?, + ScalarValue::from(-10.), + ), + ]; + for (actual, expected) in test_data { + assert_eq!(actual, expected); + } + + Ok(()) + } + + #[test] + fn test_combine_bernoullis_and_op() -> Result<()> { + let op = Operator::And; + let left = BernoulliDistribution::try_new(ScalarValue::from(0.5))?; + let right = BernoulliDistribution::try_new(ScalarValue::from(0.4))?; + let left_null = BernoulliDistribution::try_new(ScalarValue::Null)?; + let right_null = BernoulliDistribution::try_new(ScalarValue::Null)?; + + assert_eq!( + combine_bernoullis(&op, &left, &right)?.p_value(), + &ScalarValue::from(0.5 * 0.4) + ); + assert_eq!( + combine_bernoullis(&op, &left_null, &right)?.p_value(), + &ScalarValue::Float64(None) + ); + assert_eq!( + combine_bernoullis(&op, &left, &right_null)?.p_value(), + &ScalarValue::Float64(None) + ); + assert_eq!( + combine_bernoullis(&op, &left_null, &left_null)?.p_value(), + &ScalarValue::Null + ); + + Ok(()) + } + + #[test] + fn test_combine_bernoullis_or_op() -> Result<()> { + let op = Operator::Or; + let left = BernoulliDistribution::try_new(ScalarValue::from(0.6))?; + let right = BernoulliDistribution::try_new(ScalarValue::from(0.4))?; + let left_null = BernoulliDistribution::try_new(ScalarValue::Null)?; + let right_null = BernoulliDistribution::try_new(ScalarValue::Null)?; + + assert_eq!( + combine_bernoullis(&op, &left, &right)?.p_value(), + &ScalarValue::from(0.6 + 0.4 - (0.6 * 0.4)) + ); + assert_eq!( + combine_bernoullis(&op, &left_null, &right)?.p_value(), + &ScalarValue::Float64(None) + ); + assert_eq!( + combine_bernoullis(&op, &left, &right_null)?.p_value(), + &ScalarValue::Float64(None) + ); + assert_eq!( + combine_bernoullis(&op, &left_null, &left_null)?.p_value(), + &ScalarValue::Null + ); + + Ok(()) + } + + #[test] + fn test_combine_bernoullis_unsupported_ops() -> Result<()> { + let mut operator_set = operator_set(); + operator_set.remove(&Operator::And); + operator_set.remove(&Operator::Or); + + let left = BernoulliDistribution::try_new(ScalarValue::from(0.6))?; + let right = BernoulliDistribution::try_new(ScalarValue::from(0.4))?; + for op in operator_set { + assert!( + combine_bernoullis(&op, &left, &right).is_err(), + "Operator {op} should not be supported for Bernoulli distributions" + ); + } + + Ok(()) + } + + #[test] + fn test_combine_gaussians_addition() -> Result<()> { + let op = Operator::Plus; + let left = GaussianDistribution::try_new( + ScalarValue::from(3.0), + ScalarValue::from(2.0), + )?; + let right = GaussianDistribution::try_new( + ScalarValue::from(4.0), + ScalarValue::from(1.0), + )?; + + let result = combine_gaussians(&op, &left, &right)?.unwrap(); + + assert_eq!(result.mean(), &ScalarValue::from(7.0)); // 3.0 + 4.0 + assert_eq!(result.variance(), &ScalarValue::from(3.0)); // 2.0 + 1.0 + Ok(()) + } + + #[test] + fn test_combine_gaussians_subtraction() -> Result<()> { + let op = Operator::Minus; + let left = GaussianDistribution::try_new( + ScalarValue::from(7.0), + ScalarValue::from(2.0), + )?; + let right = GaussianDistribution::try_new( + ScalarValue::from(4.0), + ScalarValue::from(1.0), + )?; + + let result = combine_gaussians(&op, &left, &right)?.unwrap(); + + assert_eq!(result.mean(), &ScalarValue::from(3.0)); // 7.0 - 4.0 + assert_eq!(result.variance(), &ScalarValue::from(3.0)); // 2.0 + 1.0 + + Ok(()) + } + + #[test] + fn test_combine_gaussians_unsupported_ops() -> Result<()> { + let mut operator_set = operator_set(); + operator_set.remove(&Operator::Plus); + operator_set.remove(&Operator::Minus); + + let left = GaussianDistribution::try_new( + ScalarValue::from(7.0), + ScalarValue::from(2.0), + )?; + let right = GaussianDistribution::try_new( + ScalarValue::from(4.0), + ScalarValue::from(1.0), + )?; + for op in operator_set { + assert!( + combine_gaussians(&op, &left, &right)?.is_none(), + "Operator {op} should not be supported for Gaussian distributions" + ); + } + + Ok(()) + } + + // Expected test results were calculated in Wolfram Mathematica, by using: + // + // *METHOD_NAME*[TransformedDistribution[ + // x *op* y, + // {x ~ *DISTRIBUTION_X*[..], y ~ *DISTRIBUTION_Y*[..]} + // ]] + #[test] + fn test_calculate_generic_properties_uniform_uniform() -> Result<()> { + let dist_a = Distribution::new_uniform(Interval::make(Some(0.), Some(12.))?)?; + let dist_b = Distribution::new_uniform(Interval::make(Some(12.), Some(36.))?)?; + + let test_data = vec![ + // Mean: + ( + compute_mean(&Operator::Plus, &dist_a, &dist_b)?, + ScalarValue::from(30.), + ), + ( + compute_mean(&Operator::Minus, &dist_a, &dist_b)?, + ScalarValue::from(-18.), + ), + ( + compute_mean(&Operator::Multiply, &dist_a, &dist_b)?, + ScalarValue::from(144.), + ), + // Median: + ( + compute_median(&Operator::Plus, &dist_a, &dist_b)?, + ScalarValue::from(30.), + ), + ( + compute_median(&Operator::Minus, &dist_a, &dist_b)?, + ScalarValue::from(-18.), + ), + // Variance: + ( + compute_variance(&Operator::Plus, &dist_a, &dist_b)?, + ScalarValue::from(60.), + ), + ( + compute_variance(&Operator::Minus, &dist_a, &dist_b)?, + ScalarValue::from(60.), + ), + ( + compute_variance(&Operator::Multiply, &dist_a, &dist_b)?, + ScalarValue::from(9216.), + ), + ]; + for (actual, expected) in test_data { + assert_eq!(actual, expected); + } + + Ok(()) + } + + /// Test for `Uniform`-`Uniform`, `Uniform`-`Generic`, `Generic`-`Uniform`, + /// `Generic`-`Generic` pairs, where range is always present. + #[test] + fn test_compute_range_where_present() -> Result<()> { + let a = &Interval::make(Some(0.), Some(12.0))?; + let b = &Interval::make(Some(0.), Some(12.0))?; + let mean = ScalarValue::from(6.0); + for (dist_a, dist_b) in [ + ( + Distribution::new_uniform(a.clone())?, + Distribution::new_uniform(b.clone())?, + ), + ( + Distribution::new_generic( + mean.clone(), + mean.clone(), + ScalarValue::Float64(None), + a.clone(), + )?, + Distribution::new_uniform(b.clone())?, + ), + ( + Distribution::new_uniform(a.clone())?, + Distribution::new_generic( + mean.clone(), + mean.clone(), + ScalarValue::Float64(None), + b.clone(), + )?, + ), + ( + Distribution::new_generic( + mean.clone(), + mean.clone(), + ScalarValue::Float64(None), + a.clone(), + )?, + Distribution::new_generic( + mean.clone(), + mean.clone(), + ScalarValue::Float64(None), + b.clone(), + )?, + ), + ] { + use super::Operator::{ + Divide, Eq, Gt, GtEq, Lt, LtEq, Minus, Multiply, NotEq, Plus, + }; + for op in [Plus, Minus, Multiply, Divide] { + assert_eq!( + new_generic_from_binary_op(&op, &dist_a, &dist_b)?.range()?, + apply_operator(&op, a, b)?, + "Failed for {:?} {op} {:?}", + dist_a, + dist_b + ); + } + for op in [Gt, GtEq, Lt, LtEq, Eq, NotEq] { + assert_eq!( + create_bernoulli_from_comparison(&op, &dist_a, &dist_b)?.range()?, + apply_operator(&op, a, b)?, + "Failed for {:?} {op} {:?}", + dist_a, + dist_b + ); + } + } + + Ok(()) + } + + fn operator_set() -> HashSet { + use super::Operator::*; + + let all_ops = vec![ + And, + Or, + Eq, + NotEq, + Gt, + GtEq, + Lt, + LtEq, + Plus, + Minus, + Multiply, + Divide, + Modulo, + IsDistinctFrom, + IsNotDistinctFrom, + RegexMatch, + RegexIMatch, + RegexNotMatch, + RegexNotIMatch, + LikeMatch, + ILikeMatch, + NotLikeMatch, + NotILikeMatch, + BitwiseAnd, + BitwiseOr, + BitwiseXor, + BitwiseShiftRight, + BitwiseShiftLeft, + StringConcat, + AtArrow, + ArrowAt, + ]; + + all_ops.into_iter().collect() + } +} diff --git a/datafusion/expr/Cargo.toml b/datafusion/expr/Cargo.toml index b4f3f7fb680f6..37e1ed1936fb4 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -28,12 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_expr" -path = "src/lib.rs" [features] recursive_protection = ["dep:recursive"] diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index df79b3568ce60..f8baf9c94b3cf 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -2294,7 +2294,6 @@ impl Display for SchemaDisplay<'_> { | Expr::OuterReferenceColumn(..) | Expr::Placeholder(_) | Expr::Wildcard { .. } => write!(f, "{}", self.0), - Expr::AggregateFunction(AggregateFunction { func, params }) => { match func.schema_name(params) { Ok(name) => { diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index 2f04f234eb1db..02684928bac71 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] @@ -46,7 +51,6 @@ pub mod function; pub mod groups_accumulator { pub use datafusion_expr_common::groups_accumulator::*; } - pub mod interval_arithmetic { pub use datafusion_expr_common::interval_arithmetic::*; } @@ -57,6 +61,9 @@ pub mod simplify; pub mod sort_properties { pub use datafusion_expr_common::sort_properties::*; } +pub mod statistics { + pub use datafusion_expr_common::statistics::*; +} pub mod test; pub mod tree_node; pub mod type_coercion; diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs index 04cc26c910cb0..a2ed0592efdb4 100644 --- a/datafusion/expr/src/planner.rs +++ b/datafusion/expr/src/planner.rs @@ -25,9 +25,12 @@ use datafusion_common::{ config::ConfigOptions, file_options::file_type::FileType, not_impl_err, DFSchema, Result, TableReference, }; -use sqlparser::ast; +use sqlparser::ast::{self, NullTreatment}; -use crate::{AggregateUDF, Expr, GetFieldAccess, ScalarUDF, TableSource, WindowUDF}; +use crate::{ + AggregateUDF, Expr, GetFieldAccess, ScalarUDF, SortExpr, TableSource, WindowFrame, + WindowFunctionDefinition, WindowUDF, +}; /// Provides the `SQL` query planner meta-data about tables and /// functions referenced in SQL statements, without a direct dependency on the @@ -138,7 +141,7 @@ pub trait ExprPlanner: Debug + Send + Sync { /// Plan an array literal, such as `[1, 2, 3]` /// - /// Returns origin expression arguments if not possible + /// Returns original expression arguments if not possible fn plan_array_literal( &self, exprs: Vec, @@ -149,14 +152,14 @@ pub trait ExprPlanner: Debug + Send + Sync { /// Plan a `POSITION` expression, such as `POSITION( in )` /// - /// returns origin expression arguments if not possible + /// Returns original expression arguments if not possible fn plan_position(&self, args: Vec) -> Result>> { Ok(PlannerResult::Original(args)) } /// Plan a dictionary literal, such as `{ key: value, ...}` /// - /// Returns origin expression arguments if not possible + /// Returns original expression arguments if not possible fn plan_dictionary_literal( &self, expr: RawDictionaryExpr, @@ -167,14 +170,14 @@ pub trait ExprPlanner: Debug + Send + Sync { /// Plan an extract expression, such as`EXTRACT(month FROM foo)` /// - /// Returns origin expression arguments if not possible + /// Returns original expression arguments if not possible fn plan_extract(&self, args: Vec) -> Result>> { Ok(PlannerResult::Original(args)) } /// Plan an substring expression, such as `SUBSTRING( [FROM ] [FOR ])` /// - /// Returns origin expression arguments if not possible + /// Returns original expression arguments if not possible fn plan_substring(&self, args: Vec) -> Result>> { Ok(PlannerResult::Original(args)) } @@ -195,14 +198,14 @@ pub trait ExprPlanner: Debug + Send + Sync { /// Plans an overlay expression, such as `overlay(str PLACING substr FROM pos [FOR count])` /// - /// Returns origin expression arguments if not possible + /// Returns original expression arguments if not possible fn plan_overlay(&self, args: Vec) -> Result>> { Ok(PlannerResult::Original(args)) } /// Plans a `make_map` expression, such as `make_map(key1, value1, key2, value2, ...)` /// - /// Returns origin expression arguments if not possible + /// Returns original expression arguments if not possible fn plan_make_map(&self, args: Vec) -> Result>> { Ok(PlannerResult::Original(args)) } @@ -230,6 +233,23 @@ pub trait ExprPlanner: Debug + Send + Sync { fn plan_any(&self, expr: RawBinaryExpr) -> Result> { Ok(PlannerResult::Original(expr)) } + + /// Plans aggregate functions, such as `COUNT()` + /// + /// Returns original expression arguments if not possible + fn plan_aggregate( + &self, + expr: RawAggregateExpr, + ) -> Result> { + Ok(PlannerResult::Original(expr)) + } + + /// Plans window functions, such as `COUNT()` + /// + /// Returns original expression arguments if not possible + fn plan_window(&self, expr: RawWindowExpr) -> Result> { + Ok(PlannerResult::Original(expr)) + } } /// An operator with two arguments to plan @@ -266,6 +286,30 @@ pub struct RawDictionaryExpr { pub values: Vec, } +/// This structure is used by `AggregateFunctionPlanner` to plan operators with +/// custom expressions. +#[derive(Debug, Clone)] +pub struct RawAggregateExpr { + pub func: Arc, + pub args: Vec, + pub distinct: bool, + pub filter: Option>, + pub order_by: Option>, + pub null_treatment: Option, +} + +/// This structure is used by `WindowFunctionPlanner` to plan operators with +/// custom expressions. +#[derive(Debug, Clone)] +pub struct RawWindowExpr { + pub func_def: WindowFunctionDefinition, + pub args: Vec, + pub partition_by: Vec, + pub order_by: Vec, + pub window_frame: WindowFrame, + pub null_treatment: Option, +} + /// Result of planning a raw expr with [`ExprPlanner`] #[derive(Debug, Clone)] pub enum PlannerResult { diff --git a/datafusion/expr/src/udaf.rs b/datafusion/expr/src/udaf.rs index 2b9e2bddd1844..ae7196c9b10f9 100644 --- a/datafusion/expr/src/udaf.rs +++ b/datafusion/expr/src/udaf.rs @@ -515,9 +515,9 @@ pub trait AggregateUDFImpl: Debug + Send + Sync { null_treatment, } = params; - let mut schema_name = String::new(); + let mut display_name = String::new(); - schema_name.write_fmt(format_args!( + display_name.write_fmt(format_args!( "{}({}{})", self.name(), if *distinct { "DISTINCT " } else { "" }, @@ -525,17 +525,22 @@ pub trait AggregateUDFImpl: Debug + Send + Sync { ))?; if let Some(nt) = null_treatment { - schema_name.write_fmt(format_args!(" {}", nt))?; + display_name.write_fmt(format_args!(" {}", nt))?; } if let Some(fe) = filter { - schema_name.write_fmt(format_args!(" FILTER (WHERE {fe})"))?; + display_name.write_fmt(format_args!(" FILTER (WHERE {fe})"))?; } - if let Some(order_by) = order_by { - schema_name - .write_fmt(format_args!(" ORDER BY [{}]", expr_vec_fmt!(order_by)))?; + if let Some(ob) = order_by { + display_name.write_fmt(format_args!( + " ORDER BY [{}]", + ob.iter() + .map(|o| format!("{o}")) + .collect::>() + .join(", ") + ))?; } - Ok(schema_name) + Ok(display_name) } /// Returns the user-defined display name of function, given the arguments diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 86c0f9ad637c6..56c1e64554a90 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -832,7 +832,7 @@ pub fn exprlist_len( .enumerate() .filter_map(|(idx, field)| { let (maybe_table_ref, _) = schema.qualified_field(idx); - if maybe_table_ref.map_or(true, |q| q == qualifier) { + if maybe_table_ref.is_none_or(|q| q == qualifier) { Some((maybe_table_ref.cloned(), Arc::clone(field))) } else { None diff --git a/datafusion/ffi/Cargo.toml b/datafusion/ffi/Cargo.toml index 4c396144347c4..97914666688fd 100644 --- a/datafusion/ffi/Cargo.toml +++ b/datafusion/ffi/Cargo.toml @@ -27,12 +27,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_ffi" -path = "src/lib.rs" crate-type = ["cdylib", "rlib"] [dependencies] diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index 8087acfa33c84..00602474d621d 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -219,17 +219,17 @@ impl TryFrom<&FFI_ExecutionPlan> for ForeignExecutionPlan { let properties: PlanProperties = (plan.properties)(plan).try_into()?; let children_rvec = (plan.children)(plan); - let children: Result> = children_rvec + let children = children_rvec .iter() .map(ForeignExecutionPlan::try_from) .map(|child| child.map(|c| Arc::new(c) as Arc)) - .collect(); + .collect::>>()?; Ok(Self { name, plan: plan.clone(), properties, - children: children?, + children, }) } } @@ -281,6 +281,7 @@ impl ExecutionPlan for ForeignExecutionPlan { #[cfg(test)] mod tests { + use arrow::datatypes::{DataType, Field, Schema}; use datafusion::{ physical_plan::{ execution_plan::{Boundedness, EmissionType}, @@ -294,6 +295,7 @@ mod tests { #[derive(Debug)] pub struct EmptyExec { props: PlanProperties, + children: Vec>, } impl EmptyExec { @@ -305,6 +307,7 @@ mod tests { EmissionType::Incremental, Boundedness::Bounded, ), + children: Vec::default(), } } } @@ -333,14 +336,17 @@ mod tests { } fn children(&self) -> Vec<&Arc> { - vec![] + self.children.iter().collect() } fn with_new_children( self: Arc, - _: Vec>, + children: Vec>, ) -> Result> { - unimplemented!() + Ok(Arc::new(EmptyExec { + props: self.props.clone(), + children, + })) } fn execute( @@ -358,7 +364,6 @@ mod tests { #[test] fn test_round_trip_ffi_execution_plan() -> Result<()> { - use arrow::datatypes::{DataType, Field, Schema}; let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); let ctx = SessionContext::new(); @@ -372,6 +377,49 @@ mod tests { assert!(original_name == foreign_plan.name()); + let display = datafusion::physical_plan::display::DisplayableExecutionPlan::new( + &foreign_plan, + ); + + let buf = display.one_line().to_string(); + assert_eq!(buf.trim(), "FFI_ExecutionPlan(number_of_children=0)"); + + Ok(()) + } + + #[test] + fn test_ffi_execution_plan_children() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); + let ctx = SessionContext::new(); + + // Version 1: Adding child to the foreign plan + let child_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let child_local = FFI_ExecutionPlan::new(child_plan, ctx.task_ctx(), None); + let child_foreign = Arc::new(ForeignExecutionPlan::try_from(&child_local)?); + + let parent_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let parent_local = FFI_ExecutionPlan::new(parent_plan, ctx.task_ctx(), None); + let parent_foreign = Arc::new(ForeignExecutionPlan::try_from(&parent_local)?); + + assert_eq!(parent_foreign.children().len(), 0); + assert_eq!(child_foreign.children().len(), 0); + + let parent_foreign = parent_foreign.with_new_children(vec![child_foreign])?; + assert_eq!(parent_foreign.children().len(), 1); + + // Version 2: Adding child to the local plan + let child_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let child_local = FFI_ExecutionPlan::new(child_plan, ctx.task_ctx(), None); + let child_foreign = Arc::new(ForeignExecutionPlan::try_from(&child_local)?); + + let parent_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let parent_plan = parent_plan.with_new_children(vec![child_foreign])?; + let parent_local = FFI_ExecutionPlan::new(parent_plan, ctx.task_ctx(), None); + let parent_foreign = Arc::new(ForeignExecutionPlan::try_from(&parent_local)?); + + assert_eq!(parent_foreign.children().len(), 1); + Ok(()) } } diff --git a/datafusion/ffi/src/insert_op.rs b/datafusion/ffi/src/insert_op.rs index e44262377405b..8e8693076cc0e 100644 --- a/datafusion/ffi/src/insert_op.rs +++ b/datafusion/ffi/src/insert_op.rs @@ -47,3 +47,26 @@ impl From for FFI_InsertOp { } } } + +#[cfg(test)] +mod tests { + use datafusion::logical_expr::dml::InsertOp; + + use super::FFI_InsertOp; + + fn test_round_trip_insert_op(insert_op: InsertOp) { + let ffi_insert_op: FFI_InsertOp = insert_op.into(); + let round_trip: InsertOp = ffi_insert_op.into(); + + assert_eq!(insert_op, round_trip); + } + + /// This test ensures we have not accidentally mapped the FFI + /// enums to the wrong internal enums values. + #[test] + fn test_all_round_trip_insert_ops() { + test_round_trip_insert_op(InsertOp::Append); + test_round_trip_insert_op(InsertOp::Overwrite); + test_round_trip_insert_op(InsertOp::Replace); + } +} diff --git a/datafusion/ffi/src/lib.rs b/datafusion/ffi/src/lib.rs index bbcdd85ff80a5..4eabf91d892a9 100644 --- a/datafusion/ffi/src/lib.rs +++ b/datafusion/ffi/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/ffi/src/table_source.rs b/datafusion/ffi/src/table_source.rs index a59836622ee65..418fdf16a564f 100644 --- a/datafusion/ffi/src/table_source.rs +++ b/datafusion/ffi/src/table_source.rs @@ -85,3 +85,43 @@ impl From for FFI_TableType { } } } + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::error::Result; + + fn round_trip_filter_pushdown(pushdown: TableProviderFilterPushDown) -> Result<()> { + let ffi_pushdown: FFI_TableProviderFilterPushDown = (&pushdown).into(); + let round_trip: TableProviderFilterPushDown = (&ffi_pushdown).into(); + + assert_eq!(pushdown, round_trip); + Ok(()) + } + + #[test] + fn round_trip_all_filter_pushdowns() -> Result<()> { + round_trip_filter_pushdown(TableProviderFilterPushDown::Exact)?; + round_trip_filter_pushdown(TableProviderFilterPushDown::Inexact)?; + round_trip_filter_pushdown(TableProviderFilterPushDown::Unsupported)?; + + Ok(()) + } + + fn round_trip_table_type(table_type: TableType) -> Result<()> { + let ffi_type: FFI_TableType = table_type.into(); + let round_trip_type: TableType = ffi_type.into(); + + assert_eq!(table_type, round_trip_type); + Ok(()) + } + + #[test] + fn test_round_all_trip_table_type() -> Result<()> { + round_trip_table_type(TableType::Base)?; + round_trip_table_type(TableType::Temporary)?; + round_trip_table_type(TableType::View)?; + + Ok(()) + } +} diff --git a/datafusion/ffi/src/volatility.rs b/datafusion/ffi/src/volatility.rs index 8b565b91b76dc..0aaf68a174cfd 100644 --- a/datafusion/ffi/src/volatility.rs +++ b/datafusion/ffi/src/volatility.rs @@ -46,3 +46,24 @@ impl From<&FFI_Volatility> for Volatility { } } } + +#[cfg(test)] +mod tests { + use datafusion::logical_expr::Volatility; + + use super::FFI_Volatility; + + fn test_round_trip_volatility(volatility: Volatility) { + let ffi_volatility: FFI_Volatility = volatility.into(); + let round_trip: Volatility = (&ffi_volatility).into(); + + assert_eq!(volatility, round_trip); + } + + #[test] + fn test_all_round_trip_volatility() { + test_round_trip_volatility(Volatility::Immutable); + test_round_trip_volatility(Volatility::Stable); + test_round_trip_volatility(Volatility::Volatile); + } +} diff --git a/datafusion/functions-aggregate-common/Cargo.toml b/datafusion/functions-aggregate-common/Cargo.toml index cf6eb99e60c62..cf065ca1cb174 100644 --- a/datafusion/functions-aggregate-common/Cargo.toml +++ b/datafusion/functions-aggregate-common/Cargo.toml @@ -27,14 +27,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_functions_aggregate_common" -path = "src/lib.rs" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] ahash = { workspace = true } @@ -44,7 +44,7 @@ datafusion-expr-common = { workspace = true } datafusion-physical-expr-common = { workspace = true } [dev-dependencies] -criterion = "0.5" +criterion = { workspace = true } rand = { workspace = true } [[bench]] diff --git a/datafusion/functions-aggregate-common/src/lib.rs b/datafusion/functions-aggregate-common/src/lib.rs index cc50ff70913b0..6f9dfca30c19a 100644 --- a/datafusion/functions-aggregate-common/src/lib.rs +++ b/datafusion/functions-aggregate-common/src/lib.rs @@ -22,6 +22,11 @@ //! //! [DataFusion]: +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/functions-aggregate/Cargo.toml b/datafusion/functions-aggregate/Cargo.toml index 007e1e76a3be2..ec6e6b633bb81 100644 --- a/datafusion/functions-aggregate/Cargo.toml +++ b/datafusion/functions-aggregate/Cargo.toml @@ -28,14 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_functions_aggregate" -path = "src/lib.rs" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] ahash = { workspace = true } @@ -54,7 +54,7 @@ paste = "1.0.14" [dev-dependencies] arrow = { workspace = true, features = ["test_utils"] } -criterion = "0.5" +criterion = { workspace = true } rand = { workspace = true } [[bench]] diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index cb59042ef4684..c11329d7f5b38 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -17,11 +17,15 @@ use ahash::RandomState; use datafusion_common::stats::Precision; +use datafusion_expr::expr::{ + schema_name_from_exprs, schema_name_from_sorts, AggregateFunctionParams, + WindowFunctionParams, +}; use datafusion_functions_aggregate_common::aggregate::count_distinct::BytesViewDistinctCountAccumulator; use datafusion_macros::user_doc; use datafusion_physical_expr::expressions; use std::collections::HashSet; -use std::fmt::Debug; +use std::fmt::{Debug, Write}; use std::mem::{size_of, size_of_val}; use std::ops::BitAnd; use std::sync::Arc; @@ -47,11 +51,11 @@ use datafusion_common::{ downcast_value, internal_err, not_impl_err, Result, ScalarValue, }; use datafusion_expr::function::StateFieldsArgs; +use datafusion_expr::{expr_vec_fmt, Expr, ReversedUDAF, StatisticsArgs, TypeSignature}; use datafusion_expr::{ function::AccumulatorArgs, utils::format_state_name, Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupsAccumulator, SetMonotonicity, Signature, Volatility, }; -use datafusion_expr::{Expr, ReversedUDAF, StatisticsArgs, TypeSignature}; use datafusion_functions_aggregate_common::aggregate::count_distinct::{ BytesDistinctCountAccumulator, FloatDistinctCountAccumulator, PrimitiveDistinctCountAccumulator, @@ -79,6 +83,11 @@ pub fn count_distinct(expr: Expr) -> Expr { )) } +/// Creates aggregation to count all rows, equivalent to `COUNT(*)`, `COUNT()`, `COUNT(1)` +pub fn count_all() -> Expr { + count(Expr::Literal(COUNT_STAR_EXPANSION)) +} + #[user_doc( doc_section(label = "General Functions"), description = "Returns the number of non-null values in the specified column. To include null values in the total count, use `count(*)`.", @@ -139,6 +148,185 @@ impl AggregateUDFImpl for Count { "count" } + fn schema_name(&self, params: &AggregateFunctionParams) -> Result { + let AggregateFunctionParams { + args, + distinct, + filter, + order_by, + null_treatment, + } = params; + + let mut schema_name = String::new(); + + if is_count_wildcard(args) { + schema_name.write_str("count(*)")?; + } else { + schema_name.write_fmt(format_args!( + "{}({}{})", + self.name(), + if *distinct { "DISTINCT " } else { "" }, + schema_name_from_exprs(args)? + ))?; + } + + if let Some(null_treatment) = null_treatment { + schema_name.write_fmt(format_args!(" {}", null_treatment))?; + } + + if let Some(filter) = filter { + schema_name.write_fmt(format_args!(" FILTER (WHERE {filter})"))?; + }; + + if let Some(order_by) = order_by { + schema_name.write_fmt(format_args!( + " ORDER BY [{}]", + schema_name_from_sorts(order_by)? + ))?; + }; + + Ok(schema_name) + } + + fn window_function_schema_name( + &self, + params: &WindowFunctionParams, + ) -> Result { + let WindowFunctionParams { + args, + partition_by, + order_by, + window_frame, + null_treatment, + } = params; + + let mut schema_name = String::new(); + + if is_count_wildcard(args) { + schema_name.write_str("count(*)")?; + } else { + schema_name.write_fmt(format_args!( + "{}({})", + self.name(), + schema_name_from_exprs(args)? + ))?; + } + + if let Some(null_treatment) = null_treatment { + schema_name.write_fmt(format_args!(" {}", null_treatment))?; + } + + if !partition_by.is_empty() { + schema_name.write_fmt(format_args!( + " PARTITION BY [{}]", + schema_name_from_exprs(partition_by)? + ))?; + } + + if !order_by.is_empty() { + schema_name.write_fmt(format_args!( + " ORDER BY [{}]", + schema_name_from_sorts(order_by)? + ))?; + }; + + schema_name.write_fmt(format_args!(" {window_frame}"))?; + + Ok(schema_name) + } + + fn display_name(&self, params: &AggregateFunctionParams) -> Result { + let AggregateFunctionParams { + args, + distinct, + filter, + order_by, + null_treatment, + } = params; + + let mut display_name = String::new(); + + if is_count_wildcard(args) { + display_name.write_str("count(*)")?; + } else { + display_name.write_fmt(format_args!( + "{}({}{})", + self.name(), + if *distinct { "DISTINCT " } else { "" }, + args.iter() + .map(|arg| format!("{arg}")) + .collect::>() + .join(", ") + ))?; + } + + if let Some(nt) = null_treatment { + display_name.write_fmt(format_args!(" {}", nt))?; + } + if let Some(fe) = filter { + display_name.write_fmt(format_args!(" FILTER (WHERE {fe})"))?; + } + if let Some(ob) = order_by { + display_name.write_fmt(format_args!( + " ORDER BY [{}]", + ob.iter() + .map(|o| format!("{o}")) + .collect::>() + .join(", ") + ))?; + } + + Ok(display_name) + } + + fn window_function_display_name( + &self, + params: &WindowFunctionParams, + ) -> Result { + let WindowFunctionParams { + args, + partition_by, + order_by, + window_frame, + null_treatment, + } = params; + + let mut display_name = String::new(); + + if is_count_wildcard(args) { + display_name.write_str("count(*)")?; + } else { + display_name.write_fmt(format_args!( + "{}({})", + self.name(), + expr_vec_fmt!(args) + ))?; + } + + if let Some(null_treatment) = null_treatment { + display_name.write_fmt(format_args!(" {}", null_treatment))?; + } + + if !partition_by.is_empty() { + display_name.write_fmt(format_args!( + " PARTITION BY [{}]", + expr_vec_fmt!(partition_by) + ))?; + } + + if !order_by.is_empty() { + display_name + .write_fmt(format_args!(" ORDER BY [{}]", expr_vec_fmt!(order_by)))?; + }; + + display_name.write_fmt(format_args!( + " {} BETWEEN {} AND {}", + window_frame.units, window_frame.start_bound, window_frame.end_bound + ))?; + + Ok(display_name) + } + fn signature(&self) -> &Signature { &self.signature } @@ -359,6 +547,15 @@ impl AggregateUDFImpl for Count { } } +fn is_count_wildcard(args: &[Expr]) -> bool { + match args { + [] => true, // count() + // All const should be coerced to int64 or rejected by the signature + [Expr::Literal(ScalarValue::Int64(Some(_)))] => true, // count(1) + _ => false, // More than one argument or non-matching cases + } +} + #[derive(Debug)] struct CountAccumulator { count: i64, diff --git a/datafusion/functions-aggregate/src/lib.rs b/datafusion/functions-aggregate/src/lib.rs index 746873442d9aa..a5c84298e9d5d 100644 --- a/datafusion/functions-aggregate/src/lib.rs +++ b/datafusion/functions-aggregate/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] @@ -59,28 +64,29 @@ pub mod macros; pub mod approx_distinct; +pub mod approx_median; +pub mod approx_percentile_cont; +pub mod approx_percentile_cont_with_weight; pub mod array_agg; +pub mod average; +pub mod bit_and_or_xor; +pub mod bool_and_or; pub mod correlation; pub mod count; pub mod covariance; pub mod first_last; +pub mod grouping; pub mod hyperloglog; pub mod median; pub mod min_max; +pub mod nth_value; pub mod regr; pub mod stddev; +pub mod string_agg; pub mod sum; pub mod variance; -pub mod approx_median; -pub mod approx_percentile_cont; -pub mod approx_percentile_cont_with_weight; -pub mod average; -pub mod bit_and_or_xor; -pub mod bool_and_or; -pub mod grouping; -pub mod nth_value; -pub mod string_agg; +pub mod planner; use crate::approx_percentile_cont::approx_percentile_cont_udaf; use crate::approx_percentile_cont_with_weight::approx_percentile_cont_with_weight_udaf; diff --git a/datafusion/functions-aggregate/src/planner.rs b/datafusion/functions-aggregate/src/planner.rs new file mode 100644 index 0000000000000..1f0a42c4c71bb --- /dev/null +++ b/datafusion/functions-aggregate/src/planner.rs @@ -0,0 +1,63 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! SQL planning extensions like [`AggregateFunctionPlanner`] + +use datafusion_common::Result; +use datafusion_expr::{ + expr::AggregateFunction, + lit, + planner::{ExprPlanner, PlannerResult, RawAggregateExpr}, + utils::COUNT_STAR_EXPANSION, + Expr, +}; + +#[derive(Debug)] +pub struct AggregateFunctionPlanner; + +impl ExprPlanner for AggregateFunctionPlanner { + fn plan_aggregate( + &self, + expr: RawAggregateExpr, + ) -> Result> { + if expr.func.name() == "count" + && (expr.args.len() == 1 && matches!(expr.args[0], Expr::Wildcard { .. }) + || expr.args.is_empty()) + { + let RawAggregateExpr { + func, + args: _, + distinct, + filter, + order_by, + null_treatment, + } = expr; + return Ok(PlannerResult::Planned(Expr::AggregateFunction( + AggregateFunction::new_udf( + func, + vec![lit(COUNT_STAR_EXPANSION)], + distinct, + filter, + order_by, + null_treatment, + ), + ))); + } + + Ok(PlannerResult::Original(expr)) + } +} diff --git a/datafusion/functions-nested/Cargo.toml b/datafusion/functions-nested/Cargo.toml index a63175b36e215..9a7b1f460ef54 100644 --- a/datafusion/functions-nested/Cargo.toml +++ b/datafusion/functions-nested/Cargo.toml @@ -28,16 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true -[features] - [lib] name = "datafusion_functions_nested" -path = "src/lib.rs" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] arrow = { workspace = true } @@ -55,8 +53,8 @@ log = { workspace = true } paste = "1.0.14" [dev-dependencies] -criterion = { version = "0.5", features = ["async_tokio"] } -rand = "0.8.5" +criterion = { workspace = true, features = ["async_tokio"] } +rand = { workspace = true } [[bench]] harness = false diff --git a/datafusion/functions-nested/src/lib.rs b/datafusion/functions-nested/src/lib.rs index c47e4a696a1d8..41ebb4366cff0 100644 --- a/datafusion/functions-nested/src/lib.rs +++ b/datafusion/functions-nested/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/functions-table/Cargo.toml b/datafusion/functions-table/Cargo.toml index f722d698f3d38..78d59257dd480 100644 --- a/datafusion/functions-table/Cargo.toml +++ b/datafusion/functions-table/Cargo.toml @@ -28,14 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_functions_table" -path = "src/lib.rs" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] arrow = { workspace = true } diff --git a/datafusion/functions-table/src/lib.rs b/datafusion/functions-table/src/lib.rs index f5436f7bf8a6d..4a31760e7c4d4 100644 --- a/datafusion/functions-table/src/lib.rs +++ b/datafusion/functions-table/src/lib.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + pub mod generate_series; use datafusion_catalog::TableFunction; diff --git a/datafusion/functions-window-common/Cargo.toml b/datafusion/functions-window-common/Cargo.toml index b5df212b7d2ad..466e7bc68b486 100644 --- a/datafusion/functions-window-common/Cargo.toml +++ b/datafusion/functions-window-common/Cargo.toml @@ -28,14 +28,14 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_functions_window_common" -path = "src/lib.rs" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] datafusion-common = { workspace = true } diff --git a/datafusion/functions-window-common/src/lib.rs b/datafusion/functions-window-common/src/lib.rs index da8d096da5621..6f2a1ac0f33ff 100644 --- a/datafusion/functions-window-common/src/lib.rs +++ b/datafusion/functions-window-common/src/lib.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + //! Common user-defined window functionality for [DataFusion] //! //! [DataFusion]: diff --git a/datafusion/functions-window/Cargo.toml b/datafusion/functions-window/Cargo.toml index fc1bc51bcc665..e0c17c579b196 100644 --- a/datafusion/functions-window/Cargo.toml +++ b/datafusion/functions-window/Cargo.toml @@ -28,14 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_functions_window" -path = "src/lib.rs" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] datafusion-common = { workspace = true } diff --git a/datafusion/functions-window/src/lib.rs b/datafusion/functions-window/src/lib.rs index 9f8e54a0423b8..718b0bf1587b2 100644 --- a/datafusion/functions-window/src/lib.rs +++ b/datafusion/functions-window/src/lib.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + //! Window Function packages for [DataFusion]. //! //! This crate contains a collection of various window function packages for DataFusion, @@ -39,6 +45,9 @@ pub mod nth_value; pub mod ntile; pub mod rank; pub mod row_number; + +pub mod planner; + mod utils; /// Fluent-style API for creating `Expr`s diff --git a/datafusion/functions-window/src/planner.rs b/datafusion/functions-window/src/planner.rs new file mode 100644 index 0000000000000..8f48ca8b18dce --- /dev/null +++ b/datafusion/functions-window/src/planner.rs @@ -0,0 +1,61 @@ +// 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. + +//! SQL planning extensions like [`WindowFunctionPlanner`] + +use datafusion_common::Result; +use datafusion_expr::{ + expr::WindowFunction, + lit, + planner::{ExprPlanner, PlannerResult, RawWindowExpr}, + utils::COUNT_STAR_EXPANSION, + Expr, ExprFunctionExt, +}; + +#[derive(Debug)] +pub struct WindowFunctionPlanner; + +impl ExprPlanner for WindowFunctionPlanner { + fn plan_window(&self, expr: RawWindowExpr) -> Result> { + if expr.func_def.name() == "count" + && (expr.args.len() == 1 && matches!(expr.args[0], Expr::Wildcard { .. }) + || expr.args.is_empty()) + { + let RawWindowExpr { + func_def, + args: _, + partition_by, + order_by, + window_frame, + null_treatment, + } = expr; + return Ok(PlannerResult::Planned( + Expr::WindowFunction(WindowFunction::new( + func_def, + vec![lit(COUNT_STAR_EXPANSION)], + )) + .partition_by(partition_by) + .order_by(order_by) + .window_frame(window_frame) + .null_treatment(null_treatment) + .build()?, + )); + } + + Ok(PlannerResult::Original(expr)) + } +} diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index c00997853bb39..b44127d6a1b7b 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -28,6 +28,9 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true @@ -60,9 +63,6 @@ unicode_expressions = ["unicode-segmentation"] [lib] name = "datafusion_functions" -path = "src/lib.rs" - -# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] arrow = { workspace = true } @@ -89,7 +89,7 @@ uuid = { version = "1.13", features = ["v4"], optional = true } [dev-dependencies] arrow = { workspace = true, features = ["test_utils"] } -criterion = "0.5" +criterion = { workspace = true } rand = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync"] } diff --git a/datafusion/functions/benches/helper.rs b/datafusion/functions/benches/helper.rs index c7c405bc46969..0dbb4b0027d42 100644 --- a/datafusion/functions/benches/helper.rs +++ b/datafusion/functions/benches/helper.rs @@ -35,7 +35,7 @@ pub fn gen_string_array( let rng_ref = &mut rng; let corpus = "DataFusionДатаФусион数据融合📊🔥"; // includes utf8 encoding with 1~4 bytes - let corpus_char_count = corpus.chars().count(); + let corpus = corpus.chars().collect::>(); let mut output_string_vec: Vec> = Vec::with_capacity(n_rows); for _ in 0..n_rows { @@ -46,8 +46,7 @@ pub fn gen_string_array( // Generate random UTF8 string let mut generated_string = String::with_capacity(str_len_chars); for _ in 0..str_len_chars { - let idx = rng_ref.gen_range(0..corpus_char_count); - let char = corpus.chars().nth(idx).unwrap(); + let char = corpus[rng_ref.gen_range(0..corpus.len())]; generated_string.push(char); } output_string_vec.push(Some(generated_string)); diff --git a/datafusion/functions/src/lib.rs b/datafusion/functions/src/lib.rs index ffb0ab5f51ecb..de2571779d423 100644 --- a/datafusion/functions/src/lib.rs +++ b/datafusion/functions/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/macros/Cargo.toml b/datafusion/macros/Cargo.toml index 8946640f20bb8..737d2ed72874e 100644 --- a/datafusion/macros/Cargo.toml +++ b/datafusion/macros/Cargo.toml @@ -27,6 +27,9 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true diff --git a/datafusion/macros/src/user_doc.rs b/datafusion/macros/src/user_doc.rs index 6ca90ed376c37..c6510c1564232 100644 --- a/datafusion/macros/src/user_doc.rs +++ b/datafusion/macros/src/user_doc.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + extern crate proc_macro; use datafusion_expr::scalar_doc_sections::doc_sections_const; use proc_macro::TokenStream; diff --git a/datafusion/optimizer/Cargo.toml b/datafusion/optimizer/Cargo.toml index 3f5ec9b0da030..3413b365f67de 100644 --- a/datafusion/optimizer/Cargo.toml +++ b/datafusion/optimizer/Cargo.toml @@ -28,12 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_optimizer" -path = "src/lib.rs" [features] recursive_protection = ["dep:recursive"] diff --git a/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs b/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs deleted file mode 100644 index f517761b1e330..0000000000000 --- a/datafusion/optimizer/src/analyzer/count_wildcard_rule.rs +++ /dev/null @@ -1,277 +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::analyzer::AnalyzerRule; - -use crate::utils::NamePreserver; -use datafusion_common::config::ConfigOptions; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::Result; -use datafusion_expr::expr::{AggregateFunction, AggregateFunctionParams, WindowFunction}; -use datafusion_expr::utils::COUNT_STAR_EXPANSION; -use datafusion_expr::{lit, Expr, LogicalPlan, WindowFunctionDefinition}; - -/// Rewrite `Count(Expr::Wildcard)` to `Count(Expr::Literal)`. -/// -/// Resolves issue: -#[derive(Default, Debug)] -pub struct CountWildcardRule {} - -impl CountWildcardRule { - pub fn new() -> Self { - Self {} - } -} - -impl AnalyzerRule for CountWildcardRule { - fn analyze(&self, plan: LogicalPlan, _: &ConfigOptions) -> Result { - plan.transform_down_with_subqueries(analyze_internal).data() - } - - fn name(&self) -> &str { - "count_wildcard_rule" - } -} - -fn is_wildcard(expr: &Expr) -> bool { - matches!(expr, Expr::Wildcard { .. }) -} - -fn is_count_star_aggregate(aggregate_function: &AggregateFunction) -> bool { - matches!(aggregate_function, - AggregateFunction { - func, - params: AggregateFunctionParams { args, .. }, - } if func.name() == "count" && (args.len() == 1 && is_wildcard(&args[0]) || args.is_empty())) -} - -fn is_count_star_window_aggregate(window_function: &WindowFunction) -> bool { - let args = &window_function.params.args; - matches!(window_function.fun, - WindowFunctionDefinition::AggregateUDF(ref udaf) - if udaf.name() == "count" && (args.len() == 1 && is_wildcard(&args[0]) || args.is_empty())) -} - -fn analyze_internal(plan: LogicalPlan) -> Result> { - let name_preserver = NamePreserver::new(&plan); - plan.map_expressions(|expr| { - let original_name = name_preserver.save(&expr); - let transformed_expr = expr.transform_up(|expr| match expr { - Expr::WindowFunction(mut window_function) - if is_count_star_window_aggregate(&window_function) => - { - window_function.params.args = vec![lit(COUNT_STAR_EXPANSION)]; - Ok(Transformed::yes(Expr::WindowFunction(window_function))) - } - Expr::AggregateFunction(mut aggregate_function) - if is_count_star_aggregate(&aggregate_function) => - { - aggregate_function.params.args = vec![lit(COUNT_STAR_EXPANSION)]; - Ok(Transformed::yes(Expr::AggregateFunction( - aggregate_function, - ))) - } - _ => Ok(Transformed::no(expr)), - })?; - Ok(transformed_expr.update_data(|data| original_name.restore(data))) - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::test::*; - use arrow::datatypes::DataType; - use datafusion_common::ScalarValue; - use datafusion_expr::expr::Sort; - use datafusion_expr::ExprFunctionExt; - use datafusion_expr::{ - col, exists, in_subquery, logical_plan::LogicalPlanBuilder, out_ref_col, - scalar_subquery, wildcard, WindowFrame, WindowFrameBound, WindowFrameUnits, - }; - use datafusion_functions_aggregate::count::count_udaf; - use datafusion_functions_aggregate::expr_fn::max; - use std::sync::Arc; - - use datafusion_functions_aggregate::expr_fn::{count, sum}; - - fn assert_plan_eq(plan: LogicalPlan, expected: &str) -> Result<()> { - assert_analyzed_plan_eq_display_indent( - Arc::new(CountWildcardRule::new()), - plan, - expected, - ) - } - - #[test] - fn test_count_wildcard_on_sort() -> Result<()> { - let table_scan = test_table_scan()?; - let plan = LogicalPlanBuilder::from(table_scan) - .aggregate(vec![col("b")], vec![count(wildcard())])? - .project(vec![count(wildcard())])? - .sort(vec![count(wildcard()).sort(true, false)])? - .build()?; - let expected = "Sort: count(*) ASC NULLS LAST [count(*):Int64]\ - \n Projection: count(*) [count(*):Int64]\ - \n Aggregate: groupBy=[[test.b]], aggr=[[count(Int64(1)) AS count(*)]] [b:UInt32, count(*):Int64]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_count_wildcard_on_where_in() -> Result<()> { - let table_scan_t1 = test_table_scan_with_name("t1")?; - let table_scan_t2 = test_table_scan_with_name("t2")?; - - let plan = LogicalPlanBuilder::from(table_scan_t1) - .filter(in_subquery( - col("a"), - Arc::new( - LogicalPlanBuilder::from(table_scan_t2) - .aggregate(Vec::::new(), vec![count(wildcard())])? - .project(vec![count(wildcard())])? - .build()?, - ), - ))? - .build()?; - - let expected = "Filter: t1.a IN () [a:UInt32, b:UInt32, c:UInt32]\ - \n Subquery: [count(*):Int64]\ - \n Projection: count(*) [count(*):Int64]\ - \n Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] [count(*):Int64]\ - \n TableScan: t2 [a:UInt32, b:UInt32, c:UInt32]\ - \n TableScan: t1 [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_count_wildcard_on_where_exists() -> Result<()> { - let table_scan_t1 = test_table_scan_with_name("t1")?; - let table_scan_t2 = test_table_scan_with_name("t2")?; - - let plan = LogicalPlanBuilder::from(table_scan_t1) - .filter(exists(Arc::new( - LogicalPlanBuilder::from(table_scan_t2) - .aggregate(Vec::::new(), vec![count(wildcard())])? - .project(vec![count(wildcard())])? - .build()?, - )))? - .build()?; - - let expected = "Filter: EXISTS () [a:UInt32, b:UInt32, c:UInt32]\ - \n Subquery: [count(*):Int64]\ - \n Projection: count(*) [count(*):Int64]\ - \n Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] [count(*):Int64]\ - \n TableScan: t2 [a:UInt32, b:UInt32, c:UInt32]\ - \n TableScan: t1 [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_count_wildcard_on_where_scalar_subquery() -> Result<()> { - let table_scan_t1 = test_table_scan_with_name("t1")?; - let table_scan_t2 = test_table_scan_with_name("t2")?; - - let plan = LogicalPlanBuilder::from(table_scan_t1) - .filter( - scalar_subquery(Arc::new( - LogicalPlanBuilder::from(table_scan_t2) - .filter(out_ref_col(DataType::UInt32, "t1.a").eq(col("t2.a")))? - .aggregate( - Vec::::new(), - vec![count(lit(COUNT_STAR_EXPANSION))], - )? - .project(vec![count(lit(COUNT_STAR_EXPANSION))])? - .build()?, - )) - .gt(lit(ScalarValue::UInt8(Some(0)))), - )? - .project(vec![col("t1.a"), col("t1.b")])? - .build()?; - - let expected = "Projection: t1.a, t1.b [a:UInt32, b:UInt32]\ - \n Filter: () > UInt8(0) [a:UInt32, b:UInt32, c:UInt32]\ - \n Subquery: [count(Int64(1)):Int64]\ - \n Projection: count(Int64(1)) [count(Int64(1)):Int64]\ - \n Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] [count(Int64(1)):Int64]\ - \n Filter: outer_ref(t1.a) = t2.a [a:UInt32, b:UInt32, c:UInt32]\ - \n TableScan: t2 [a:UInt32, b:UInt32, c:UInt32]\ - \n TableScan: t1 [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - #[test] - fn test_count_wildcard_on_window() -> Result<()> { - let table_scan = test_table_scan()?; - - let plan = LogicalPlanBuilder::from(table_scan) - .window(vec![Expr::WindowFunction(WindowFunction::new( - WindowFunctionDefinition::AggregateUDF(count_udaf()), - vec![wildcard()], - )) - .order_by(vec![Sort::new(col("a"), false, true)]) - .window_frame(WindowFrame::new_bounds( - WindowFrameUnits::Range, - WindowFrameBound::Preceding(ScalarValue::UInt32(Some(6))), - WindowFrameBound::Following(ScalarValue::UInt32(Some(2))), - )) - .build()?])? - .project(vec![count(wildcard())])? - .build()?; - - let expected = "Projection: count(Int64(1)) AS count(*) [count(*):Int64]\ - \n WindowAggr: windowExpr=[[count(Int64(1)) ORDER BY [test.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING AS count(*) ORDER BY [test.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING]] [a:UInt32, b:UInt32, c:UInt32, count(*) ORDER BY [test.a DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING:Int64]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_count_wildcard_on_aggregate() -> Result<()> { - let table_scan = test_table_scan()?; - let plan = LogicalPlanBuilder::from(table_scan) - .aggregate(Vec::::new(), vec![count(wildcard())])? - .project(vec![count(wildcard())])? - .build()?; - - let expected = "Projection: count(*) [count(*):Int64]\ - \n Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] [count(*):Int64]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } - - #[test] - fn test_count_wildcard_on_non_count_aggregate() -> Result<()> { - let table_scan = test_table_scan()?; - let res = LogicalPlanBuilder::from(table_scan) - .aggregate(Vec::::new(), vec![sum(wildcard())]); - assert!(res.is_err()); - Ok(()) - } - - #[test] - fn test_count_wildcard_on_nesting() -> Result<()> { - let table_scan = test_table_scan()?; - let plan = LogicalPlanBuilder::from(table_scan) - .aggregate(Vec::::new(), vec![max(count(wildcard()))])? - .project(vec![count(wildcard())])? - .build()?; - - let expected = "Projection: count(Int64(1)) AS count(*) [count(*):Int64]\ - \n Aggregate: groupBy=[[]], aggr=[[max(count(Int64(1))) AS max(count(*))]] [max(count(*)):Int64;N]\ - \n TableScan: test [a:UInt32, b:UInt32, c:UInt32]"; - assert_plan_eq(plan, expected) - } -} diff --git a/datafusion/optimizer/src/analyzer/mod.rs b/datafusion/optimizer/src/analyzer/mod.rs index 9d0ac6b54cf45..c506616d142e8 100644 --- a/datafusion/optimizer/src/analyzer/mod.rs +++ b/datafusion/optimizer/src/analyzer/mod.rs @@ -28,7 +28,6 @@ use datafusion_common::Result; use datafusion_expr::expr_rewriter::FunctionRewrite; use datafusion_expr::{InvariantLevel, LogicalPlan}; -use crate::analyzer::count_wildcard_rule::CountWildcardRule; use crate::analyzer::expand_wildcard_rule::ExpandWildcardRule; use crate::analyzer::inline_table_scan::InlineTableScan; use crate::analyzer::resolve_grouping_function::ResolveGroupingFunction; @@ -37,7 +36,6 @@ use crate::utils::log_plan; use self::function_rewrite::ApplyFunctionRewrites; -pub mod count_wildcard_rule; pub mod expand_wildcard_rule; pub mod function_rewrite; pub mod inline_table_scan; @@ -106,7 +104,6 @@ impl Analyzer { // [Expr::Wildcard] should be expanded before [TypeCoercion] Arc::new(ResolveGroupingFunction::new()), Arc::new(TypeCoercion::new()), - Arc::new(CountWildcardRule::new()), ]; Self::with_rules(rules) } diff --git a/datafusion/optimizer/src/lib.rs b/datafusion/optimizer/src/lib.rs index 614284e1b4774..61ca9b31cd29b 100644 --- a/datafusion/optimizer/src/lib.rs +++ b/datafusion/optimizer/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs index a33ecbc3a1fb3..b59acd72a26dd 100644 --- a/datafusion/optimizer/tests/optimizer_integration.rs +++ b/datafusion/optimizer/tests/optimizer_integration.rs @@ -23,11 +23,14 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef, TimeUnit}; use datafusion_common::config::ConfigOptions; use datafusion_common::{assert_contains, plan_err, Result, TableReference}; +use datafusion_expr::planner::ExprPlanner; use datafusion_expr::sqlparser::dialect::PostgreSqlDialect; use datafusion_expr::test::function_stub::sum_udaf; use datafusion_expr::{AggregateUDF, LogicalPlan, ScalarUDF, TableSource, WindowUDF}; use datafusion_functions_aggregate::average::avg_udaf; use datafusion_functions_aggregate::count::count_udaf; +use datafusion_functions_aggregate::planner::AggregateFunctionPlanner; +use datafusion_functions_window::planner::WindowFunctionPlanner; use datafusion_optimizer::analyzer::type_coercion::TypeCoercionRewriter; use datafusion_optimizer::analyzer::Analyzer; use datafusion_optimizer::optimizer::Optimizer; @@ -195,7 +198,7 @@ fn between_date32_plus_interval() -> Result<()> { WHERE col_date32 between '1998-03-18' AND cast('1998-03-18' as date) + INTERVAL '90 days'"; let plan = test_sql(sql)?; let expected = - "Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]]\ + "Aggregate: groupBy=[[]], aggr=[[count(*)]]\ \n Projection: \ \n Filter: test.col_date32 >= Date32(\"1998-03-18\") AND test.col_date32 <= Date32(\"1998-06-16\")\ \n TableScan: test projection=[col_date32]"; @@ -209,7 +212,7 @@ fn between_date64_plus_interval() -> Result<()> { WHERE col_date64 between '1998-03-18T00:00:00' AND cast('1998-03-18' as date) + INTERVAL '90 days'"; let plan = test_sql(sql)?; let expected = - "Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]]\ + "Aggregate: groupBy=[[]], aggr=[[count(*)]]\ \n Projection: \ \n Filter: test.col_date64 >= Date64(\"1998-03-18\") AND test.col_date64 <= Date64(\"1998-06-16\")\ \n TableScan: test projection=[col_date64]"; @@ -266,7 +269,7 @@ fn push_down_filter_groupby_expr_contains_alias() { let sql = "SELECT * FROM (SELECT (col_int32 + col_uint32) AS c, count(*) FROM test GROUP BY 1) where c > 3"; let plan = test_sql(sql).unwrap(); let expected = "Projection: test.col_int32 + test.col_uint32 AS c, count(*)\ - \n Aggregate: groupBy=[[test.col_int32 + CAST(test.col_uint32 AS Int32)]], aggr=[[count(Int64(1)) AS count(*)]]\ + \n Aggregate: groupBy=[[test.col_int32 + CAST(test.col_uint32 AS Int32)]], aggr=[[count(*)]]\ \n Filter: test.col_int32 + CAST(test.col_uint32 AS Int32) > Int32(3)\ \n TableScan: test projection=[col_int32, col_uint32]"; assert_eq!(expected, format!("{plan}")); @@ -311,7 +314,7 @@ fn eliminate_redundant_null_check_on_count() { let plan = test_sql(sql).unwrap(); let expected = "\ Projection: test.col_int32, count(*) AS c\ - \n Aggregate: groupBy=[[test.col_int32]], aggr=[[count(Int64(1)) AS count(*)]]\ + \n Aggregate: groupBy=[[test.col_int32]], aggr=[[count(*)]]\ \n TableScan: test projection=[col_int32]"; assert_eq!(expected, format!("{plan}")); } @@ -422,7 +425,12 @@ fn test_sql(sql: &str) -> Result { let context_provider = MyContextProvider::default() .with_udaf(sum_udaf()) .with_udaf(count_udaf()) - .with_udaf(avg_udaf()); + .with_udaf(avg_udaf()) + .with_expr_planners(vec![ + Arc::new(AggregateFunctionPlanner), + Arc::new(WindowFunctionPlanner), + ]); + let sql_to_rel = SqlToRel::new(&context_provider); let plan = sql_to_rel.sql_statement_to_plan(statement.clone())?; @@ -440,6 +448,7 @@ fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {} struct MyContextProvider { options: ConfigOptions, udafs: HashMap>, + expr_planners: Vec>, } impl MyContextProvider { @@ -448,6 +457,11 @@ impl MyContextProvider { self.udafs.insert(udaf.name().to_lowercase(), udaf); self } + + fn with_expr_planners(mut self, expr_planners: Vec>) -> Self { + self.expr_planners = expr_planners; + self + } } impl ContextProvider for MyContextProvider { @@ -516,6 +530,10 @@ impl ContextProvider for MyContextProvider { fn udwf_names(&self) -> Vec { Vec::new() } + + fn get_expr_planners(&self) -> &[Arc] { + &self.expr_planners + } } struct MyTableSource { diff --git a/datafusion/physical-expr-common/Cargo.toml b/datafusion/physical-expr-common/Cargo.toml index 14d6ca64d15e6..a5a12b5527b7d 100644 --- a/datafusion/physical-expr-common/Cargo.toml +++ b/datafusion/physical-expr-common/Cargo.toml @@ -28,12 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_physical_expr_common" -path = "src/lib.rs" [dependencies] ahash = { workspace = true } diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index 809c619e9845a..b37d9a7773eeb 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -384,7 +384,7 @@ where // value is "small" let payload = if value.len() <= SHORT_VALUE_LEN { - let inline = value.iter().fold(0usize, |acc, &x| acc << 8 | x as usize); + let inline = value.iter().fold(0usize, |acc, &x| (acc << 8) | x as usize); // is value is already present in the set? let entry = self.map.find_mut(hash, |header| { diff --git a/datafusion/physical-expr-common/src/lib.rs b/datafusion/physical-expr-common/src/lib.rs index a05f1c96306f2..440f044d88eb6 100644 --- a/datafusion/physical-expr-common/src/lib.rs +++ b/datafusion/physical-expr-common/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/physical-expr-common/src/physical_expr.rs b/datafusion/physical-expr-common/src/physical_expr.rs index b1b889136b35f..cc2ff2f24790f 100644 --- a/datafusion/physical-expr-common/src/physical_expr.rs +++ b/datafusion/physical-expr-common/src/physical_expr.rs @@ -26,10 +26,13 @@ use arrow::array::BooleanArray; use arrow::compute::filter_record_batch; use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; -use datafusion_common::{internal_err, not_impl_err, Result}; +use datafusion_common::{internal_err, not_impl_err, Result, ScalarValue}; use datafusion_expr_common::columnar_value::ColumnarValue; use datafusion_expr_common::interval_arithmetic::Interval; use datafusion_expr_common::sort_properties::ExprProperties; +use datafusion_expr_common::statistics::Distribution; + +use itertools::izip; /// Shared [`PhysicalExpr`]. pub type PhysicalExprRef = Arc; @@ -98,11 +101,16 @@ pub trait PhysicalExpr: Send + Sync + Display + Debug + DynEq + DynHash { /// Computes the output interval for the expression, given the input /// intervals. /// - /// # Arguments + /// # Parameters /// /// * `children` are the intervals for the children (inputs) of this /// expression. /// + /// # Returns + /// + /// A `Result` containing the output interval for the expression in + /// case of success, or an error object in case of failure. + /// /// # Example /// /// If the expression is `a + b`, and the input intervals are `a: [1, 2]` @@ -116,19 +124,20 @@ pub trait PhysicalExpr: Send + Sync + Display + Debug + DynEq + DynHash { /// /// This is used to propagate constraints down through an expression tree. /// - /// # Arguments + /// # Parameters /// /// * `interval` is the currently known interval for this expression. /// * `children` are the current intervals for the children of this expression. /// /// # Returns /// - /// A `Vec` of new intervals for the children, in order. + /// A `Result` containing a `Vec` of new intervals for the children (in order) + /// in case of success, or an error object in case of failure. /// /// If constraint propagation reveals an infeasibility for any child, returns - /// [`None`]. If none of the children intervals change as a result of propagation, - /// may return an empty vector instead of cloning `children`. This is the default - /// (and conservative) return value. + /// [`None`]. If none of the children intervals change as a result of + /// propagation, may return an empty vector instead of cloning `children`. + /// This is the default (and conservative) return value. /// /// # Example /// @@ -144,6 +153,111 @@ pub trait PhysicalExpr: Send + Sync + Display + Debug + DynEq + DynHash { Ok(Some(vec![])) } + /// Computes the output statistics for the expression, given the input + /// statistics. + /// + /// # Parameters + /// + /// * `children` are the statistics for the children (inputs) of this + /// expression. + /// + /// # Returns + /// + /// A `Result` containing the output statistics for the expression in + /// case of success, or an error object in case of failure. + /// + /// Expressions (should) implement this function and utilize the independence + /// assumption, match on children distribution types and compute the output + /// statistics accordingly. The default implementation simply creates an + /// unknown output distribution by combining input ranges. This logic loses + /// distribution information, but is a safe default. + fn evaluate_statistics(&self, children: &[&Distribution]) -> Result { + let children_ranges = children + .iter() + .map(|c| c.range()) + .collect::>>()?; + let children_ranges_refs = children_ranges.iter().collect::>(); + let output_interval = self.evaluate_bounds(children_ranges_refs.as_slice())?; + let dt = output_interval.data_type(); + if dt.eq(&DataType::Boolean) { + let p = if output_interval.eq(&Interval::CERTAINLY_TRUE) { + ScalarValue::new_one(&dt) + } else if output_interval.eq(&Interval::CERTAINLY_FALSE) { + ScalarValue::new_zero(&dt) + } else { + ScalarValue::try_from(&dt) + }?; + Distribution::new_bernoulli(p) + } else { + Distribution::new_from_interval(output_interval) + } + } + + /// Updates children statistics using the given parent statistic for this + /// expression. + /// + /// This is used to propagate statistics down through an expression tree. + /// + /// # Parameters + /// + /// * `parent` is the currently known statistics for this expression. + /// * `children` are the current statistics for the children of this expression. + /// + /// # Returns + /// + /// A `Result` containing a `Vec` of new statistics for the children (in order) + /// in case of success, or an error object in case of failure. + /// + /// If statistics propagation reveals an infeasibility for any child, returns + /// [`None`]. If none of the children statistics change as a result of + /// propagation, may return an empty vector instead of cloning `children`. + /// This is the default (and conservative) return value. + /// + /// Expressions (should) implement this function and apply Bayes rule to + /// reconcile and update parent/children statistics. This involves utilizing + /// the independence assumption, and matching on distribution types. The + /// default implementation simply creates an unknown distribution if it can + /// narrow the range by propagating ranges. This logic loses distribution + /// information, but is a safe default. + fn propagate_statistics( + &self, + parent: &Distribution, + children: &[&Distribution], + ) -> Result>> { + let children_ranges = children + .iter() + .map(|c| c.range()) + .collect::>>()?; + let children_ranges_refs = children_ranges.iter().collect::>(); + let parent_range = parent.range()?; + let Some(propagated_children) = + self.propagate_constraints(&parent_range, children_ranges_refs.as_slice())? + else { + return Ok(None); + }; + izip!(propagated_children.into_iter(), children_ranges, children) + .map(|(new_interval, old_interval, child)| { + if new_interval == old_interval { + // We weren't able to narrow the range, preserve the old statistics. + Ok((*child).clone()) + } else if new_interval.data_type().eq(&DataType::Boolean) { + let dt = old_interval.data_type(); + let p = if new_interval.eq(&Interval::CERTAINLY_TRUE) { + ScalarValue::new_one(&dt) + } else if new_interval.eq(&Interval::CERTAINLY_FALSE) { + ScalarValue::new_zero(&dt) + } else { + unreachable!("Given that we have a range reduction for a boolean interval, we should have certainty") + }?; + Distribution::new_bernoulli(p) + } else { + Distribution::new_from_interval(new_interval) + } + }) + .collect::>() + .map(Some) + } + /// Calculates the properties of this [`PhysicalExpr`] based on its /// children's properties (i.e. order and range), recursively aggregating /// the information from its children. In cases where the [`PhysicalExpr`] @@ -155,7 +269,7 @@ pub trait PhysicalExpr: Send + Sync + Display + Debug + DynEq + DynHash { } /// [`PhysicalExpr`] can't be constrained by [`Eq`] directly because it must remain object -/// safe. To ease implementation blanket implementation is provided for [`Eq`] types. +/// safe. To ease implementation, blanket implementation is provided for [`Eq`] types. pub trait DynEq { fn dyn_eq(&self, other: &dyn Any) -> bool; } diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs index b150d3dc9bd38..601b2a23d09da 100644 --- a/datafusion/physical-expr-common/src/sort_expr.rs +++ b/datafusion/physical-expr-common/src/sort_expr.rs @@ -172,13 +172,11 @@ impl PhysicalSortExpr { let nullable = self.expr.nullable(schema).unwrap_or(true); self.expr.eq(&requirement.expr) && if nullable { - requirement - .options - .map_or(true, |opts| self.options == opts) + requirement.options.is_none_or(|opts| self.options == opts) } else { requirement .options - .map_or(true, |opts| self.options.descending == opts.descending) + .is_none_or(|opts| self.options.descending == opts.descending) } } } @@ -293,7 +291,7 @@ impl PhysicalSortRequirement { self.expr.eq(&other.expr) && other .options - .map_or(true, |other_opts| self.options == Some(other_opts)) + .is_none_or(|other_opts| self.options == Some(other_opts)) } #[deprecated(since = "43.0.0", note = "use LexRequirement::from_lex_ordering")] diff --git a/datafusion/physical-expr/Cargo.toml b/datafusion/physical-expr/Cargo.toml index 38e8b44791ab8..72baa0db00a21 100644 --- a/datafusion/physical-expr/Cargo.toml +++ b/datafusion/physical-expr/Cargo.toml @@ -28,12 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_physical_expr" -path = "src/lib.rs" [dependencies] ahash = { workspace = true } @@ -53,7 +55,7 @@ petgraph = "0.7.1" [dev-dependencies] arrow = { workspace = true, features = ["test_utils"] } -criterion = "0.5" +criterion = { workspace = true } datafusion-functions = { workspace = true } rand = { workspace = true } rstest = { workspace = true } diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs index 96208cc5e32c9..838cb26807a9f 100755 --- a/datafusion/physical-expr/src/equivalence/properties.rs +++ b/datafusion/physical-expr/src/equivalence/properties.rs @@ -640,7 +640,7 @@ impl EquivalenceProperties { req.expr.eq(&existing.expr) && req .options - .map_or(true, |req_opts| req_opts == existing.options) + .is_none_or(|req_opts| req_opts == existing.options) }, ) }) diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index 052054bad6c1c..1f16c5471ed77 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/datafusion/physical-expr/src/expressions/binary.rs @@ -20,6 +20,7 @@ mod kernels; use std::hash::Hash; use std::{any::Any, sync::Arc}; +use crate::expressions::binary::kernels::concat_elements_utf8view; use crate::intervals::cp_solver::{propagate_arithmetic, propagate_comparison}; use crate::PhysicalExpr; @@ -36,10 +37,14 @@ use datafusion_common::{internal_err, Result, ScalarValue}; use datafusion_expr::binary::BinaryTypeCoercer; use datafusion_expr::interval_arithmetic::{apply_operator, Interval}; use datafusion_expr::sort_properties::ExprProperties; +use datafusion_expr::statistics::Distribution::{Bernoulli, Gaussian}; +use datafusion_expr::statistics::{ + combine_bernoullis, combine_gaussians, create_bernoulli_from_comparison, + new_generic_from_binary_op, Distribution, +}; use datafusion_expr::{ColumnarValue, Operator}; use datafusion_physical_expr_common::datum::{apply, apply_cmp, apply_cmp_for_nested}; -use crate::expressions::binary::kernels::concat_elements_utf8view; use kernels::{ bitwise_and_dyn, bitwise_and_dyn_scalar, bitwise_or_dyn, bitwise_or_dyn_scalar, bitwise_shift_left_dyn, bitwise_shift_left_dyn_scalar, bitwise_shift_right_dyn, @@ -486,6 +491,37 @@ impl PhysicalExpr for BinaryExpr { } } + fn evaluate_statistics(&self, children: &[&Distribution]) -> Result { + let (left, right) = (children[0], children[1]); + + if self.op.is_numerical_operators() { + // We might be able to construct the output statistics more accurately, + // without falling back to an unknown distribution, if we are dealing + // with Gaussian distributions and numerical operations. + if let (Gaussian(left), Gaussian(right)) = (left, right) { + if let Some(result) = combine_gaussians(&self.op, left, right)? { + return Ok(Gaussian(result)); + } + } + } else if self.op.is_logic_operator() { + // If we are dealing with logical operators, we expect (and can only + // operate on) Bernoulli distributions. + return if let (Bernoulli(left), Bernoulli(right)) = (left, right) { + combine_bernoullis(&self.op, left, right).map(Bernoulli) + } else { + internal_err!( + "Logical operators are only compatible with `Bernoulli` distributions" + ) + }; + } else if self.op.supports_propagation() { + // If we are handling comparison operators, we expect (and can only + // operate on) numeric distributions. + return create_bernoulli_from_comparison(&self.op, left, right); + } + // Fall back to an unknown distribution with only summary statistics: + new_generic_from_binary_op(&self.op, left, right) + } + /// For each operator, [`BinaryExpr`] has distinct rules. /// TODO: There may be rules specific to some data types and expression ranges. fn get_properties(&self, children: &[ExprProperties]) -> Result { @@ -732,6 +768,7 @@ pub fn similar_to( mod tests { use super::*; use crate::expressions::{col, lit, try_cast, Column, Literal}; + use datafusion_common::plan_datafusion_err; /// Performs a binary operation, applying any type coercion necessary @@ -4379,4 +4416,260 @@ mod tests { ) .unwrap(); } + + pub fn binary_expr( + left: Arc, + op: Operator, + right: Arc, + schema: &Schema, + ) -> Result { + Ok(binary_op(left, op, right, schema)? + .as_any() + .downcast_ref::() + .unwrap() + .clone()) + } + + /// Test for Uniform-Uniform, Unknown-Uniform, Uniform-Unknown and Unknown-Unknown evaluation. + #[test] + fn test_evaluate_statistics_combination_of_range_holders() -> Result<()> { + let schema = &Schema::new(vec![Field::new("a", DataType::Float64, false)]); + let a = Arc::new(Column::new("a", 0)) as _; + let b = lit(ScalarValue::from(12.0)); + + let left_interval = Interval::make(Some(0.0), Some(12.0))?; + let right_interval = Interval::make(Some(12.0), Some(36.0))?; + let (left_mean, right_mean) = (ScalarValue::from(6.0), ScalarValue::from(24.0)); + let (left_med, right_med) = (ScalarValue::from(6.0), ScalarValue::from(24.0)); + + for children in [ + vec![ + &Distribution::new_uniform(left_interval.clone())?, + &Distribution::new_uniform(right_interval.clone())?, + ], + vec![ + &Distribution::new_generic( + left_mean.clone(), + left_med.clone(), + ScalarValue::Float64(None), + left_interval.clone(), + )?, + &Distribution::new_uniform(right_interval.clone())?, + ], + vec![ + &Distribution::new_uniform(right_interval.clone())?, + &Distribution::new_generic( + right_mean.clone(), + right_med.clone(), + ScalarValue::Float64(None), + right_interval.clone(), + )?, + ], + vec![ + &Distribution::new_generic( + left_mean.clone(), + left_med.clone(), + ScalarValue::Float64(None), + left_interval.clone(), + )?, + &Distribution::new_generic( + right_mean.clone(), + right_med.clone(), + ScalarValue::Float64(None), + right_interval.clone(), + )?, + ], + ] { + let ops = vec![ + Operator::Plus, + Operator::Minus, + Operator::Multiply, + Operator::Divide, + ]; + + for op in ops { + let expr = binary_expr(Arc::clone(&a), op, Arc::clone(&b), schema)?; + assert_eq!( + expr.evaluate_statistics(&children)?, + new_generic_from_binary_op(&op, children[0], children[1])? + ); + } + } + Ok(()) + } + + #[test] + fn test_evaluate_statistics_bernoulli() -> Result<()> { + let schema = &Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ]); + let a = Arc::new(Column::new("a", 0)) as _; + let b = Arc::new(Column::new("b", 1)) as _; + let eq = Arc::new(binary_expr( + Arc::clone(&a), + Operator::Eq, + Arc::clone(&b), + schema, + )?); + let neq = Arc::new(binary_expr(a, Operator::NotEq, b, schema)?); + + let left_stat = &Distribution::new_uniform(Interval::make(Some(0), Some(7))?)?; + let right_stat = &Distribution::new_uniform(Interval::make(Some(4), Some(11))?)?; + + // Intervals: [0, 7], [4, 11]. + // The intersection is [4, 7], so the probability of equality is 4 / 64 = 1 / 16. + assert_eq!( + eq.evaluate_statistics(&[left_stat, right_stat])?, + Distribution::new_bernoulli(ScalarValue::from(1.0 / 16.0))? + ); + + // The probability of being distinct is 1 - 1 / 16 = 15 / 16. + assert_eq!( + neq.evaluate_statistics(&[left_stat, right_stat])?, + Distribution::new_bernoulli(ScalarValue::from(15.0 / 16.0))? + ); + + Ok(()) + } + + #[test] + fn test_propagate_statistics_combination_of_range_holders_arithmetic() -> Result<()> { + let schema = &Schema::new(vec![Field::new("a", DataType::Float64, false)]); + let a = Arc::new(Column::new("a", 0)) as _; + let b = lit(ScalarValue::from(12.0)); + + let left_interval = Interval::make(Some(0.0), Some(12.0))?; + let right_interval = Interval::make(Some(12.0), Some(36.0))?; + + let parent = Distribution::new_uniform(Interval::make(Some(-432.), Some(432.))?)?; + let children = vec![ + vec![ + Distribution::new_uniform(left_interval.clone())?, + Distribution::new_uniform(right_interval.clone())?, + ], + vec![ + Distribution::new_generic( + ScalarValue::from(6.), + ScalarValue::from(6.), + ScalarValue::Float64(None), + left_interval.clone(), + )?, + Distribution::new_uniform(right_interval.clone())?, + ], + vec![ + Distribution::new_uniform(left_interval.clone())?, + Distribution::new_generic( + ScalarValue::from(12.), + ScalarValue::from(12.), + ScalarValue::Float64(None), + right_interval.clone(), + )?, + ], + vec![ + Distribution::new_generic( + ScalarValue::from(6.), + ScalarValue::from(6.), + ScalarValue::Float64(None), + left_interval.clone(), + )?, + Distribution::new_generic( + ScalarValue::from(12.), + ScalarValue::from(12.), + ScalarValue::Float64(None), + right_interval.clone(), + )?, + ], + ]; + + let ops = vec![ + Operator::Plus, + Operator::Minus, + Operator::Multiply, + Operator::Divide, + ]; + + for child_view in children { + let child_refs = child_view.iter().collect::>(); + for op in &ops { + let expr = binary_expr(Arc::clone(&a), *op, Arc::clone(&b), schema)?; + assert_eq!( + expr.propagate_statistics(&parent, child_refs.as_slice())?, + Some(child_view.clone()) + ); + } + } + Ok(()) + } + + #[test] + fn test_propagate_statistics_combination_of_range_holders_comparison() -> Result<()> { + let schema = &Schema::new(vec![Field::new("a", DataType::Float64, false)]); + let a = Arc::new(Column::new("a", 0)) as _; + let b = lit(ScalarValue::from(12.0)); + + let left_interval = Interval::make(Some(0.0), Some(12.0))?; + let right_interval = Interval::make(Some(6.0), Some(18.0))?; + + let one = ScalarValue::from(1.0); + let parent = Distribution::new_bernoulli(one)?; + let children = vec![ + vec![ + Distribution::new_uniform(left_interval.clone())?, + Distribution::new_uniform(right_interval.clone())?, + ], + vec![ + Distribution::new_generic( + ScalarValue::from(6.), + ScalarValue::from(6.), + ScalarValue::Float64(None), + left_interval.clone(), + )?, + Distribution::new_uniform(right_interval.clone())?, + ], + vec![ + Distribution::new_uniform(left_interval.clone())?, + Distribution::new_generic( + ScalarValue::from(12.), + ScalarValue::from(12.), + ScalarValue::Float64(None), + right_interval.clone(), + )?, + ], + vec![ + Distribution::new_generic( + ScalarValue::from(6.), + ScalarValue::from(6.), + ScalarValue::Float64(None), + left_interval.clone(), + )?, + Distribution::new_generic( + ScalarValue::from(12.), + ScalarValue::from(12.), + ScalarValue::Float64(None), + right_interval.clone(), + )?, + ], + ]; + + let ops = vec![ + Operator::Eq, + Operator::Gt, + Operator::GtEq, + Operator::Lt, + Operator::LtEq, + ]; + + for child_view in children { + let child_refs = child_view.iter().collect::>(); + for op in &ops { + let expr = binary_expr(Arc::clone(&a), *op, Arc::clone(&b), schema)?; + assert!(expr + .propagate_statistics(&parent, child_refs.as_slice())? + .is_some()); + } + } + + Ok(()) + } } diff --git a/datafusion/physical-expr/src/expressions/negative.rs b/datafusion/physical-expr/src/expressions/negative.rs index dc863ccff511b..8795545274a2e 100644 --- a/datafusion/physical-expr/src/expressions/negative.rs +++ b/datafusion/physical-expr/src/expressions/negative.rs @@ -28,9 +28,12 @@ use arrow::{ datatypes::{DataType, Schema}, record_batch::RecordBatch, }; -use datafusion_common::{plan_err, Result}; +use datafusion_common::{internal_err, plan_err, Result}; use datafusion_expr::interval_arithmetic::Interval; use datafusion_expr::sort_properties::ExprProperties; +use datafusion_expr::statistics::Distribution::{ + self, Bernoulli, Exponential, Gaussian, Generic, Uniform, +}; use datafusion_expr::{ type_coercion::{is_interval, is_null, is_signed_numeric, is_timestamp}, ColumnarValue, @@ -89,14 +92,13 @@ impl PhysicalExpr for NegativeExpr { } fn evaluate(&self, batch: &RecordBatch) -> Result { - let arg = self.arg.evaluate(batch)?; - match arg { + match self.arg.evaluate(batch)? { ColumnarValue::Array(array) => { let result = neg_wrapping(array.as_ref())?; Ok(ColumnarValue::Array(result)) } ColumnarValue::Scalar(scalar) => { - Ok(ColumnarValue::Scalar((scalar.arithmetic_negate())?)) + Ok(ColumnarValue::Scalar(scalar.arithmetic_negate()?)) } } } @@ -116,10 +118,7 @@ impl PhysicalExpr for NegativeExpr { /// It replaces the upper and lower bounds after multiplying them with -1. /// Ex: `(a, b]` => `[-b, -a)` fn evaluate_bounds(&self, children: &[&Interval]) -> Result { - Interval::try_new( - children[0].upper().arithmetic_negate()?, - children[0].lower().arithmetic_negate()?, - ) + children[0].arithmetic_negate() } /// Returns a new [`Interval`] of a NegativeExpr that has the existing `interval` given that @@ -129,17 +128,37 @@ impl PhysicalExpr for NegativeExpr { interval: &Interval, children: &[&Interval], ) -> Result>> { - let child_interval = children[0]; - let negated_interval = Interval::try_new( - interval.upper().arithmetic_negate()?, - interval.lower().arithmetic_negate()?, - )?; + let negated_interval = interval.arithmetic_negate()?; - Ok(child_interval + Ok(children[0] .intersect(negated_interval)? .map(|result| vec![result])) } + fn evaluate_statistics(&self, children: &[&Distribution]) -> Result { + match children[0] { + Uniform(u) => Distribution::new_uniform(u.range().arithmetic_negate()?), + Exponential(e) => Distribution::new_exponential( + e.rate().clone(), + e.offset().arithmetic_negate()?, + !e.positive_tail(), + ), + Gaussian(g) => Distribution::new_gaussian( + g.mean().arithmetic_negate()?, + g.variance().clone(), + ), + Bernoulli(_) => { + internal_err!("NegativeExpr cannot operate on Boolean datatypes") + } + Generic(u) => Distribution::new_generic( + u.mean().arithmetic_negate()?, + u.median().arithmetic_negate()?, + u.variance().clone(), + u.range().arithmetic_negate()?, + ), + } + } + /// The ordering of a [`NegativeExpr`] is simply the reverse of its child. fn get_properties(&self, children: &[ExprProperties]) -> Result { Ok(ExprProperties { @@ -181,7 +200,7 @@ mod tests { use arrow::datatypes::DataType::{Float32, Float64, Int16, Int32, Int64, Int8}; use arrow::datatypes::*; use datafusion_common::cast::as_primitive_array; - use datafusion_common::DataFusionError; + use datafusion_common::{DataFusionError, ScalarValue}; use paste::paste; @@ -233,6 +252,67 @@ mod tests { Ok(()) } + #[test] + fn test_evaluate_statistics() -> Result<()> { + let negative_expr = NegativeExpr::new(Arc::new(Column::new("a", 0))); + + // Uniform + assert_eq!( + negative_expr.evaluate_statistics(&[&Distribution::new_uniform( + Interval::make(Some(-2.), Some(3.))? + )?])?, + Distribution::new_uniform(Interval::make(Some(-3.), Some(2.))?)? + ); + + // Bernoulli + assert!(negative_expr + .evaluate_statistics(&[&Distribution::new_bernoulli(ScalarValue::from( + 0.75 + ))?]) + .is_err()); + + // Exponential + assert_eq!( + negative_expr.evaluate_statistics(&[&Distribution::new_exponential( + ScalarValue::from(1.), + ScalarValue::from(1.), + true + )?])?, + Distribution::new_exponential( + ScalarValue::from(1.), + ScalarValue::from(-1.), + false + )? + ); + + // Gaussian + assert_eq!( + negative_expr.evaluate_statistics(&[&Distribution::new_gaussian( + ScalarValue::from(15), + ScalarValue::from(225), + )?])?, + Distribution::new_gaussian(ScalarValue::from(-15), ScalarValue::from(225),)? + ); + + // Unknown + assert_eq!( + negative_expr.evaluate_statistics(&[&Distribution::new_generic( + ScalarValue::from(15), + ScalarValue::from(15), + ScalarValue::from(10), + Interval::make(Some(10), Some(20))? + )?])?, + Distribution::new_generic( + ScalarValue::from(-15), + ScalarValue::from(-15), + ScalarValue::from(10), + Interval::make(Some(-20), Some(-10))? + )? + ); + + Ok(()) + } + #[test] fn test_propagate_constraints() -> Result<()> { let negative_expr = NegativeExpr::new(Arc::new(Column::new("a", 0))); @@ -249,6 +329,35 @@ mod tests { Ok(()) } + #[test] + fn test_propagate_statistics_range_holders() -> Result<()> { + let negative_expr = NegativeExpr::new(Arc::new(Column::new("a", 0))); + let original_child_interval = Interval::make(Some(-2), Some(3))?; + let after_propagation = Interval::make(Some(-2), Some(0))?; + + let parent = Distribution::new_uniform(Interval::make(Some(0), Some(4))?)?; + let children: Vec> = vec![ + vec![Distribution::new_uniform(original_child_interval.clone())?], + vec![Distribution::new_generic( + ScalarValue::from(0), + ScalarValue::from(0), + ScalarValue::Int32(None), + original_child_interval.clone(), + )?], + ]; + + for child_view in children { + let child_refs: Vec<_> = child_view.iter().collect(); + let actual = negative_expr.propagate_statistics(&parent, &child_refs)?; + let expected = Some(vec![Distribution::new_from_interval( + after_propagation.clone(), + )?]); + assert_eq!(actual, expected); + } + + Ok(()) + } + #[test] fn test_negation_valid_types() -> Result<()> { let negatable_types = [ diff --git a/datafusion/physical-expr/src/expressions/not.rs b/datafusion/physical-expr/src/expressions/not.rs index 440c4e9557bdf..ddf7c739b692a 100644 --- a/datafusion/physical-expr/src/expressions/not.rs +++ b/datafusion/physical-expr/src/expressions/not.rs @@ -23,10 +23,12 @@ use std::hash::Hash; use std::sync::Arc; use crate::PhysicalExpr; + use arrow::datatypes::{DataType, Schema}; use arrow::record_batch::RecordBatch; -use datafusion_common::{cast::as_boolean_array, Result, ScalarValue}; +use datafusion_common::{cast::as_boolean_array, internal_err, Result, ScalarValue}; use datafusion_expr::interval_arithmetic::Interval; +use datafusion_expr::statistics::Distribution::{self, Bernoulli}; use datafusion_expr::ColumnarValue; /// Not expression @@ -82,8 +84,7 @@ impl PhysicalExpr for NotExpr { } fn evaluate(&self, batch: &RecordBatch) -> Result { - let evaluate_arg = self.arg.evaluate(batch)?; - match evaluate_arg { + match self.arg.evaluate(batch)? { ColumnarValue::Array(array) => { let array = as_boolean_array(&array)?; Ok(ColumnarValue::Array(Arc::new( @@ -95,9 +96,7 @@ impl PhysicalExpr for NotExpr { return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None))); } let bool_value: bool = scalar.try_into()?; - Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some( - !bool_value, - )))) + Ok(ColumnarValue::Scalar(ScalarValue::from(!bool_value))) } } } @@ -112,9 +111,70 @@ impl PhysicalExpr for NotExpr { ) -> Result> { Ok(Arc::new(NotExpr::new(Arc::clone(&children[0])))) } + fn evaluate_bounds(&self, children: &[&Interval]) -> Result { children[0].not() } + + fn propagate_constraints( + &self, + interval: &Interval, + children: &[&Interval], + ) -> Result>> { + let complemented_interval = interval.not()?; + + Ok(children[0] + .intersect(complemented_interval)? + .map(|result| vec![result])) + } + + fn evaluate_statistics(&self, children: &[&Distribution]) -> Result { + match children[0] { + Bernoulli(b) => { + let p_value = b.p_value(); + if p_value.is_null() { + Ok(children[0].clone()) + } else { + let one = ScalarValue::new_one(&p_value.data_type())?; + Distribution::new_bernoulli(one.sub_checked(p_value)?) + } + } + _ => internal_err!("NotExpr can only operate on Boolean datatypes"), + } + } + + fn propagate_statistics( + &self, + parent: &Distribution, + children: &[&Distribution], + ) -> Result>> { + match (parent, children[0]) { + (Bernoulli(parent), Bernoulli(child)) => { + let parent_range = parent.range(); + let result = if parent_range == Interval::CERTAINLY_TRUE { + if child.range() == Interval::CERTAINLY_TRUE { + None + } else { + Some(vec![Distribution::new_bernoulli(ScalarValue::new_zero( + &child.data_type(), + )?)?]) + } + } else if parent_range == Interval::CERTAINLY_FALSE { + if child.range() == Interval::CERTAINLY_FALSE { + None + } else { + Some(vec![Distribution::new_bernoulli(ScalarValue::new_one( + &child.data_type(), + )?)?]) + } + } else { + Some(vec![]) + }; + Ok(result) + } + _ => internal_err!("NotExpr can only operate on Boolean datatypes"), + } + } } /// Creates a unary expression NOT @@ -124,10 +184,12 @@ pub fn not(arg: Arc) -> Result> { #[cfg(test)] mod tests { + use std::sync::LazyLock; + use super::*; - use crate::expressions::col; + use crate::expressions::{col, Column}; + use arrow::{array::BooleanArray, datatypes::*}; - use std::sync::LazyLock; #[test] fn neg_op() -> Result<()> { @@ -182,10 +244,81 @@ mod tests { expected_interval: Interval, ) -> Result<()> { let not_expr = not(col("a", &schema())?)?; + assert_eq!(not_expr.evaluate_bounds(&[&interval])?, expected_interval); + Ok(()) + } + + #[test] + fn test_evaluate_statistics() -> Result<()> { + let _schema = &Schema::new(vec![Field::new("a", DataType::Boolean, false)]); + let a = Arc::new(Column::new("a", 0)) as _; + let expr = not(a)?; + + // Uniform with non-boolean bounds + assert!(expr + .evaluate_statistics(&[&Distribution::new_uniform( + Interval::make_unbounded(&DataType::Float64)? + )?]) + .is_err()); + + // Exponential + assert!(expr + .evaluate_statistics(&[&Distribution::new_exponential( + ScalarValue::from(1.0), + ScalarValue::from(1.0), + true + )?]) + .is_err()); + + // Gaussian + assert!(expr + .evaluate_statistics(&[&Distribution::new_gaussian( + ScalarValue::from(1.0), + ScalarValue::from(1.0), + )?]) + .is_err()); + + // Bernoulli assert_eq!( - not_expr.evaluate_bounds(&[&interval]).unwrap(), - expected_interval + expr.evaluate_statistics(&[&Distribution::new_bernoulli( + ScalarValue::from(0.0), + )?])?, + Distribution::new_bernoulli(ScalarValue::from(1.))? ); + + assert_eq!( + expr.evaluate_statistics(&[&Distribution::new_bernoulli( + ScalarValue::from(1.0), + )?])?, + Distribution::new_bernoulli(ScalarValue::from(0.))? + ); + + assert_eq!( + expr.evaluate_statistics(&[&Distribution::new_bernoulli( + ScalarValue::from(0.25), + )?])?, + Distribution::new_bernoulli(ScalarValue::from(0.75))? + ); + + assert!(expr + .evaluate_statistics(&[&Distribution::new_generic( + ScalarValue::Null, + ScalarValue::Null, + ScalarValue::Null, + Interval::make_unbounded(&DataType::UInt8)? + )?]) + .is_err()); + + // Unknown with non-boolean interval as range + assert!(expr + .evaluate_statistics(&[&Distribution::new_generic( + ScalarValue::Null, + ScalarValue::Null, + ScalarValue::Null, + Interval::make_unbounded(&DataType::Float64)? + )?]) + .is_err()); + Ok(()) } diff --git a/datafusion/physical-expr/src/intervals/cp_solver.rs b/datafusion/physical-expr/src/intervals/cp_solver.rs index cb29109684fef..a53814c3ad2b8 100644 --- a/datafusion/physical-expr/src/intervals/cp_solver.rs +++ b/datafusion/physical-expr/src/intervals/cp_solver.rs @@ -15,7 +15,130 @@ // specific language governing permissions and limitations // under the License. -//! Constraint propagator/solver for custom PhysicalExpr graphs. +//! Constraint propagator/solver for custom [`PhysicalExpr`] graphs. +//! +//! The constraint propagator/solver in DataFusion uses interval arithmetic to +//! perform mathematical operations on intervals, which represent a range of +//! possible values rather than a single point value. This allows for the +//! propagation of ranges through mathematical operations, and can be used to +//! compute bounds for a complicated expression. The key idea is that by +//! breaking down a complicated expression into simpler terms, and then +//! combining the bounds for those simpler terms, one can obtain bounds for the +//! overall expression. +//! +//! This way of using interval arithmetic to compute bounds for a complex +//! expression by combining the bounds for the constituent terms within the +//! original expression allows us to reason about the range of possible values +//! of the expression. This information later can be used in range pruning of +//! the provably unnecessary parts of `RecordBatch`es. +//! +//! # Example +//! +//! For example, consider a mathematical expression such as `x^2 + y = 4` \[1\]. +//! Since this expression would be a binary tree in [`PhysicalExpr`] notation, +//! this type of an hierarchical computation is well-suited for a graph based +//! implementation. In such an implementation, an equation system `f(x) = 0` is +//! represented by a directed acyclic expression graph (DAEG). +//! +//! In order to use interval arithmetic to compute bounds for this expression, +//! one would first determine intervals that represent the possible values of +//! `x` and `y`` Let's say that the interval for `x` is `[1, 2]` and the interval +//! for `y` is `[-3, 1]`. In the chart below, you can see how the computation +//! takes place. +//! +//! # References +//! +//! 1. Kabak, Mehmet Ozan. Analog Circuit Start-Up Behavior Analysis: An Interval +//! Arithmetic Based Approach, Chapter 4. Stanford University, 2015. +//! 2. Moore, Ramon E. Interval analysis. Vol. 4. Englewood Cliffs: Prentice-Hall, 1966. +//! 3. F. Messine, "Deterministic global optimization using interval constraint +//! propagation techniques," RAIRO-Operations Research, vol. 38, no. 04, +//! pp. 277-293, 2004. +//! +//! # Illustration +//! +//! ## Computing bounds for an expression using interval arithmetic +//! +//! ```text +//! +-----+ +-----+ +//! +----| + |----+ +----| + |----+ +//! | | | | | | | | +//! | +-----+ | | +-----+ | +//! | | | | +//! +-----+ +-----+ +-----+ +-----+ +//! | 2 | | y | | 2 | [1, 4] | y | +//! |[.] | | | |[.] | | | +//! +-----+ +-----+ +-----+ +-----+ +//! | | +//! | | +//! +---+ +---+ +//! | x | [1, 2] | x | [1, 2] +//! +---+ +---+ +//! +//! (a) Bottom-up evaluation: Step 1 (b) Bottom up evaluation: Step 2 +//! +//! [1 - 3, 4 + 1] = [-2, 5] +//! +-----+ +-----+ +//! +----| + |----+ +----| + |----+ +//! | | | | | | | | +//! | +-----+ | | +-----+ | +//! | | | | +//! +-----+ +-----+ +-----+ +-----+ +//! | 2 |[1, 4] | y | | 2 |[1, 4] | y | +//! |[.] | | | |[.] | | | +//! +-----+ +-----+ +-----+ +-----+ +//! | [-3, 1] | [-3, 1] +//! | | +//! +---+ +---+ +//! | x | [1, 2] | x | [1, 2] +//! +---+ +---+ +//! +//! (c) Bottom-up evaluation: Step 3 (d) Bottom-up evaluation: Step 4 +//! ``` +//! +//! ## Top-down constraint propagation using inverse semantics +//! +//! ```text +//! [-2, 5] ∩ [4, 4] = [4, 4] [4, 4] +//! +-----+ +-----+ +//! +----| + |----+ +----| + |----+ +//! | | | | | | | | +//! | +-----+ | | +-----+ | +//! | | | | +//! +-----+ +-----+ +-----+ +-----+ +//! | 2 | [1, 4] | y | | 2 | [1, 4] | y | [0, 1]* +//! |[.] | | | |[.] | | | +//! +-----+ +-----+ +-----+ +-----+ +//! | [-3, 1] | +//! | | +//! +---+ +---+ +//! | x | [1, 2] | x | [1, 2] +//! +---+ +---+ +//! +//! (a) Top-down propagation: Step 1 (b) Top-down propagation: Step 2 +//! +//! [1 - 3, 4 + 1] = [-2, 5] +//! +-----+ +-----+ +//! +----| + |----+ +----| + |----+ +//! | | | | | | | | +//! | +-----+ | | +-----+ | +//! | | | | +//! +-----+ +-----+ +-----+ +-----+ +//! | 2 |[3, 4]** | y | | 2 |[3, 4] | y | +//! |[.] | | | |[.] | | | +//! +-----+ +-----+ +-----+ +-----+ +//! | [0, 1] | [-3, 1] +//! | | +//! +---+ +---+ +//! | x | [1, 2] | x | [sqrt(3), 2]*** +//! +---+ +---+ +//! +//! (c) Top-down propagation: Step 3 (d) Top-down propagation: Step 4 +//! +//! * [-3, 1] ∩ ([4, 4] - [1, 4]) = [0, 1] +//! ** [1, 4] ∩ ([4, 4] - [0, 1]) = [3, 4] +//! *** [1, 2] ∩ [sqrt(3), sqrt(4)] = [sqrt(3), 2] +//! ``` use std::collections::HashSet; use std::fmt::{Display, Formatter}; @@ -39,84 +162,6 @@ use petgraph::stable_graph::{DefaultIx, StableGraph}; use petgraph::visit::{Bfs, Dfs, DfsPostOrder, EdgeRef}; use petgraph::Outgoing; -// Interval arithmetic provides a way to perform mathematical operations on -// intervals, which represent a range of possible values rather than a single -// point value. This allows for the propagation of ranges through mathematical -// operations, and can be used to compute bounds for a complicated expression. -// The key idea is that by breaking down a complicated expression into simpler -// terms, and then combining the bounds for those simpler terms, one can -// obtain bounds for the overall expression. -// -// For example, consider a mathematical expression such as x^2 + y = 4. Since -// it would be a binary tree in [PhysicalExpr] notation, this type of an -// hierarchical computation is well-suited for a graph based implementation. -// In such an implementation, an equation system f(x) = 0 is represented by a -// directed acyclic expression graph (DAEG). -// -// In order to use interval arithmetic to compute bounds for this expression, -// one would first determine intervals that represent the possible values of x -// and y. Let's say that the interval for x is [1, 2] and the interval for y -// is [-3, 1]. In the chart below, you can see how the computation takes place. -// -// This way of using interval arithmetic to compute bounds for a complex -// expression by combining the bounds for the constituent terms within the -// original expression allows us to reason about the range of possible values -// of the expression. This information later can be used in range pruning of -// the provably unnecessary parts of `RecordBatch`es. -// -// References -// 1 - Kabak, Mehmet Ozan. Analog Circuit Start-Up Behavior Analysis: An Interval -// Arithmetic Based Approach, Chapter 4. Stanford University, 2015. -// 2 - Moore, Ramon E. Interval analysis. Vol. 4. Englewood Cliffs: Prentice-Hall, 1966. -// 3 - F. Messine, "Deterministic global optimization using interval constraint -// propagation techniques," RAIRO-Operations Research, vol. 38, no. 04, -// pp. 277{293, 2004. -// -// ``` text -// Computing bounds for an expression using interval arithmetic. Constraint propagation through a top-down evaluation of the expression -// graph using inverse semantics. -// -// [-2, 5] ∩ [4, 4] = [4, 4] [4, 4] -// +-----+ +-----+ +-----+ +-----+ -// +----| + |----+ +----| + |----+ +----| + |----+ +----| + |----+ -// | | | | | | | | | | | | | | | | -// | +-----+ | | +-----+ | | +-----+ | | +-----+ | -// | | | | | | | | -// +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ -// | 2 | | y | | 2 | [1, 4] | y | | 2 | [1, 4] | y | | 2 | [1, 4] | y | [0, 1]* -// |[.] | | | |[.] | | | |[.] | | | |[.] | | | -// +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ -// | | | [-3, 1] | -// | | | | -// +---+ +---+ +---+ +---+ -// | x | [1, 2] | x | [1, 2] | x | [1, 2] | x | [1, 2] -// +---+ +---+ +---+ +---+ -// -// (a) Bottom-up evaluation: Step1 (b) Bottom up evaluation: Step2 (a) Top-down propagation: Step1 (b) Top-down propagation: Step2 -// -// [1 - 3, 4 + 1] = [-2, 5] [1 - 3, 4 + 1] = [-2, 5] -// +-----+ +-----+ +-----+ +-----+ -// +----| + |----+ +----| + |----+ +----| + |----+ +----| + |----+ -// | | | | | | | | | | | | | | | | -// | +-----+ | | +-----+ | | +-----+ | | +-----+ | -// | | | | | | | | -// +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ -// | 2 |[1, 4] | y | | 2 |[1, 4] | y | | 2 |[3, 4]** | y | | 2 |[1, 4] | y | -// |[.] | | | |[.] | | | |[.] | | | |[.] | | | -// +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ +-----+ -// | [-3, 1] | [-3, 1] | [0, 1] | [-3, 1] -// | | | | -// +---+ +---+ +---+ +---+ -// | x | [1, 2] | x | [1, 2] | x | [1, 2] | x | [sqrt(3), 2]*** -// +---+ +---+ +---+ +---+ -// -// (c) Bottom-up evaluation: Step3 (d) Bottom-up evaluation: Step4 (c) Top-down propagation: Step3 (d) Top-down propagation: Step4 -// -// * [-3, 1] ∩ ([4, 4] - [1, 4]) = [0, 1] -// ** [1, 4] ∩ ([4, 4] - [0, 1]) = [3, 4] -// *** [1, 2] ∩ [sqrt(3), sqrt(4)] = [sqrt(3), 2] -// ``` - /// This object implements a directed acyclic expression graph (DAEG) that /// is used to compute ranges for expressions through interval arithmetic. #[derive(Clone, Debug)] @@ -125,18 +170,6 @@ pub struct ExprIntervalGraph { root: NodeIndex, } -impl ExprIntervalGraph { - /// Estimate size of bytes including `Self`. - pub fn size(&self) -> usize { - let node_memory_usage = self.graph.node_count() - * (size_of::() + size_of::()); - let edge_memory_usage = - self.graph.edge_count() * (size_of::() + size_of::() * 2); - - size_of_val(self) + node_memory_usage + edge_memory_usage - } -} - /// This object encapsulates all possible constraint propagation results. #[derive(PartialEq, Debug)] pub enum PropagationResult { @@ -153,6 +186,12 @@ pub struct ExprIntervalGraphNode { interval: Interval, } +impl PartialEq for ExprIntervalGraphNode { + fn eq(&self, other: &Self) -> bool { + self.expr.eq(&other.expr) + } +} + impl Display for ExprIntervalGraphNode { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.expr) @@ -160,7 +199,7 @@ impl Display for ExprIntervalGraphNode { } impl ExprIntervalGraphNode { - /// Constructs a new DAEG node with an [-∞, ∞] range. + /// Constructs a new DAEG node with an `[-∞, ∞]` range. pub fn new_unbounded(expr: Arc, dt: &DataType) -> Result { Interval::make_unbounded(dt) .map(|interval| ExprIntervalGraphNode { expr, interval }) @@ -178,7 +217,7 @@ impl ExprIntervalGraphNode { /// This function creates a DAEG node from DataFusion's [`ExprTreeNode`] /// object. Literals are created with definite, singleton intervals while - /// any other expression starts with an indefinite interval ([-∞, ∞]). + /// any other expression starts with an indefinite interval (`[-∞, ∞]`). pub fn make_node(node: &ExprTreeNode, schema: &Schema) -> Result { let expr = Arc::clone(&node.expr); if let Some(literal) = expr.as_any().downcast_ref::() { @@ -192,30 +231,24 @@ impl ExprIntervalGraphNode { } } -impl PartialEq for ExprIntervalGraphNode { - fn eq(&self, other: &Self) -> bool { - self.expr.eq(&other.expr) - } -} - /// This function refines intervals `left_child` and `right_child` by applying /// constraint propagation through `parent` via operation. The main idea is /// that we can shrink ranges of variables x and y using parent interval p. /// -/// Assuming that x,y and p has ranges [xL, xU], [yL, yU], and [pL, pU], we +/// Assuming that x,y and p has ranges `[xL, xU]`, `[yL, yU]`, and `[pL, pU]`, we /// apply the following operations: /// - For plus operation, specifically, we would first do -/// - [xL, xU] <- ([pL, pU] - [yL, yU]) ∩ [xL, xU], and then -/// - [yL, yU] <- ([pL, pU] - [xL, xU]) ∩ [yL, yU]. +/// - `[xL, xU]` <- (`[pL, pU]` - `[yL, yU]`) ∩ `[xL, xU]`, and then +/// - `[yL, yU]` <- (`[pL, pU]` - `[xL, xU]`) ∩ `[yL, yU]`. /// - For minus operation, specifically, we would first do -/// - [xL, xU] <- ([yL, yU] + [pL, pU]) ∩ [xL, xU], and then -/// - [yL, yU] <- ([xL, xU] - [pL, pU]) ∩ [yL, yU]. +/// - `[xL, xU]` <- (`[yL, yU]` + `[pL, pU]`) ∩ `[xL, xU]`, and then +/// - `[yL, yU]` <- (`[xL, xU]` - `[pL, pU]`) ∩ `[yL, yU]`. /// - For multiplication operation, specifically, we would first do -/// - [xL, xU] <- ([pL, pU] / [yL, yU]) ∩ [xL, xU], and then -/// - [yL, yU] <- ([pL, pU] / [xL, xU]) ∩ [yL, yU]. +/// - `[xL, xU]` <- (`[pL, pU]` / `[yL, yU]`) ∩ `[xL, xU]`, and then +/// - `[yL, yU]` <- (`[pL, pU]` / `[xL, xU]`) ∩ `[yL, yU]`. /// - For division operation, specifically, we would first do -/// - [xL, xU] <- ([yL, yU] * [pL, pU]) ∩ [xL, xU], and then -/// - [yL, yU] <- ([xL, xU] / [pL, pU]) ∩ [yL, yU]. +/// - `[xL, xU]` <- (`[yL, yU]` * `[pL, pU]`) ∩ `[xL, xU]`, and then +/// - `[yL, yU]` <- (`[xL, xU]` / `[pL, pU]`) ∩ `[yL, yU]`. pub fn propagate_arithmetic( op: &Operator, parent: &Interval, @@ -361,18 +394,30 @@ impl ExprIntervalGraph { self.graph.node_count() } + /// Estimate size of bytes including `Self`. + pub fn size(&self) -> usize { + let node_memory_usage = self.graph.node_count() + * (size_of::() + size_of::()); + let edge_memory_usage = + self.graph.edge_count() * (size_of::() + size_of::() * 2); + + size_of_val(self) + node_memory_usage + edge_memory_usage + } + // Sometimes, we do not want to calculate and/or propagate intervals all // way down to leaf expressions. For example, assume that we have a // `SymmetricHashJoin` which has a child with an output ordering like: // + // ```text // PhysicalSortExpr { // expr: BinaryExpr('a', +, 'b'), // sort_option: .. // } + // ``` // - // i.e. its output order comes from a clause like "ORDER BY a + b". In such - // a case, we must calculate the interval for the BinaryExpr('a', +, 'b') - // instead of the columns inside this BinaryExpr, because this interval + // i.e. its output order comes from a clause like `ORDER BY a + b`. In such + // a case, we must calculate the interval for the `BinaryExpr(a, +, b)` + // instead of the columns inside this `BinaryExpr`, because this interval // decides whether we prune or not. Therefore, children `PhysicalExpr`s of // this `BinaryExpr` may be pruned for performance. The figure below // explains this example visually. @@ -510,9 +555,6 @@ impl ExprIntervalGraph { /// Computes bounds for an expression using interval arithmetic via a /// bottom-up traversal. /// - /// # Arguments - /// * `leaf_bounds` - &[(usize, Interval)]. Provide NodeIndex, Interval tuples for leaf variables. - /// /// # Examples /// /// ``` @@ -570,7 +612,7 @@ impl ExprIntervalGraph { self.graph[node].expr.evaluate_bounds(&children_intervals)?; } } - Ok(&self.graph[self.root].interval) + Ok(self.graph[self.root].interval()) } /// Updates/shrinks bounds for leaf expressions using interval arithmetic @@ -579,8 +621,6 @@ impl ExprIntervalGraph { &mut self, given_range: Interval, ) -> Result { - let mut bfs = Bfs::new(&self.graph, self.root); - // Adjust the root node with the given range: if let Some(interval) = self.graph[self.root].interval.intersect(given_range)? { self.graph[self.root].interval = interval; @@ -588,6 +628,8 @@ impl ExprIntervalGraph { return Ok(PropagationResult::Infeasible); } + let mut bfs = Bfs::new(&self.graph, self.root); + while let Some(node) = bfs.next(&self.graph) { let neighbors = self.graph.neighbors_directed(node, Outgoing); let mut children = neighbors.collect::>(); diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index 11d6f54a7cc38..0a448fa6a2e99 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] @@ -35,6 +40,7 @@ pub mod udf { #[allow(deprecated)] pub use crate::scalar_function::create_physical_expr; } +pub mod statistics; pub mod utils; pub mod window; diff --git a/datafusion/physical-expr/src/statistics/mod.rs b/datafusion/physical-expr/src/statistics/mod.rs new file mode 100644 index 0000000000000..02897e0594578 --- /dev/null +++ b/datafusion/physical-expr/src/statistics/mod.rs @@ -0,0 +1,20 @@ +// 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. + +//! Statistics and constraint propagation library + +pub mod stats_solver; diff --git a/datafusion/physical-expr/src/statistics/stats_solver.rs b/datafusion/physical-expr/src/statistics/stats_solver.rs new file mode 100644 index 0000000000000..ec58076caf3b1 --- /dev/null +++ b/datafusion/physical-expr/src/statistics/stats_solver.rs @@ -0,0 +1,287 @@ +// 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 crate::expressions::Literal; +use crate::intervals::cp_solver::PropagationResult; +use crate::physical_expr::PhysicalExpr; +use crate::utils::{build_dag, ExprTreeNode}; + +use arrow::datatypes::{DataType, Schema}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::statistics::Distribution; +use datafusion_expr_common::interval_arithmetic::Interval; + +use petgraph::adj::DefaultIx; +use petgraph::prelude::Bfs; +use petgraph::stable_graph::{NodeIndex, StableGraph}; +use petgraph::visit::DfsPostOrder; +use petgraph::Outgoing; + +/// This object implements a directed acyclic expression graph (DAEG) that +/// is used to compute statistics/distributions for expressions hierarchically. +#[derive(Clone, Debug)] +pub struct ExprStatisticsGraph { + graph: StableGraph, + root: NodeIndex, +} + +/// This is a node in the DAEG; it encapsulates a reference to the actual +/// [`PhysicalExpr`] as well as its statistics/distribution. +#[derive(Clone, Debug)] +pub struct ExprStatisticsGraphNode { + expr: Arc, + dist: Distribution, +} + +impl ExprStatisticsGraphNode { + /// Constructs a new DAEG node based on the given interval with a + /// `Uniform` distribution. + fn new_uniform(expr: Arc, interval: Interval) -> Result { + Distribution::new_uniform(interval) + .map(|dist| ExprStatisticsGraphNode { expr, dist }) + } + + /// Constructs a new DAEG node with a `Bernoulli` distribution having an + /// unknown success probability. + fn new_bernoulli(expr: Arc) -> Result { + Distribution::new_bernoulli(ScalarValue::Float64(None)) + .map(|dist| ExprStatisticsGraphNode { expr, dist }) + } + + /// Constructs a new DAEG node with a `Generic` distribution having no + /// definite summary statistics. + fn new_generic(expr: Arc, dt: &DataType) -> Result { + let interval = Interval::make_unbounded(dt)?; + let dist = Distribution::new_from_interval(interval)?; + Ok(ExprStatisticsGraphNode { expr, dist }) + } + + /// Get the [`Distribution`] object representing the statistics of the + /// expression. + pub fn distribution(&self) -> &Distribution { + &self.dist + } + + /// This function creates a DAEG node from DataFusion's [`ExprTreeNode`] + /// object. Literals are created with `Uniform` distributions with a + /// definite, singleton interval. Expressions with a `Boolean` data type + /// result in a`Bernoulli` distribution with an unknown success probability. + /// Any other expression starts with an `Unknown` distribution with an + /// indefinite range (i.e. `[-∞, ∞]`). + pub fn make_node(node: &ExprTreeNode, schema: &Schema) -> Result { + let expr = Arc::clone(&node.expr); + if let Some(literal) = expr.as_any().downcast_ref::() { + let value = literal.value(); + Interval::try_new(value.clone(), value.clone()) + .and_then(|interval| Self::new_uniform(expr, interval)) + } else { + expr.data_type(schema).and_then(|dt| { + if dt.eq(&DataType::Boolean) { + Self::new_bernoulli(expr) + } else { + Self::new_generic(expr, &dt) + } + }) + } + } +} + +impl ExprStatisticsGraph { + pub fn try_new(expr: Arc, schema: &Schema) -> Result { + // Build the full graph: + let (root, graph) = build_dag(expr, &|node| { + ExprStatisticsGraphNode::make_node(node, schema) + })?; + Ok(Self { graph, root }) + } + + /// This function assigns given distributions to expressions in the DAEG. + /// The argument `assignments` associates indices of sought expressions + /// with their corresponding new distributions. + pub fn assign_statistics(&mut self, assignments: &[(usize, Distribution)]) { + for (index, stats) in assignments { + let node_index = NodeIndex::from(*index as DefaultIx); + self.graph[node_index].dist = stats.clone(); + } + } + + /// Computes statistics/distributions for an expression via a bottom-up + /// traversal. + pub fn evaluate_statistics(&mut self) -> Result<&Distribution> { + let mut dfs = DfsPostOrder::new(&self.graph, self.root); + while let Some(idx) = dfs.next(&self.graph) { + let neighbors = self.graph.neighbors_directed(idx, Outgoing); + let mut children_statistics = neighbors + .map(|child| self.graph[child].distribution()) + .collect::>(); + // Note that all distributions are assumed to be independent. + if !children_statistics.is_empty() { + // Reverse to align with `PhysicalExpr`'s children: + children_statistics.reverse(); + self.graph[idx].dist = self.graph[idx] + .expr + .evaluate_statistics(&children_statistics)?; + } + } + Ok(self.graph[self.root].distribution()) + } + + /// Runs a propagation mechanism in a top-down manner to update statistics + /// of leaf nodes. + pub fn propagate_statistics( + &mut self, + given_stats: Distribution, + ) -> Result { + // Adjust the root node with the given statistics: + let root_range = self.graph[self.root].dist.range()?; + let given_range = given_stats.range()?; + if let Some(interval) = root_range.intersect(&given_range)? { + if interval != root_range { + // If the given statistics enable us to obtain a more precise + // range for the root, update it: + let subset = root_range.contains(given_range)?; + self.graph[self.root].dist = if subset == Interval::CERTAINLY_TRUE { + // Given statistics is strictly more informative, use it as is: + given_stats + } else { + // Intersecting ranges gives us a more precise range: + Distribution::new_from_interval(interval)? + }; + } + } else { + return Ok(PropagationResult::Infeasible); + } + + let mut bfs = Bfs::new(&self.graph, self.root); + + while let Some(node) = bfs.next(&self.graph) { + let neighbors = self.graph.neighbors_directed(node, Outgoing); + let mut children = neighbors.collect::>(); + // If the current expression is a leaf, its statistics is now final. + // So, just continue with the propagation procedure: + if children.is_empty() { + continue; + } + // Reverse to align with `PhysicalExpr`'s children: + children.reverse(); + let children_stats = children + .iter() + .map(|child| self.graph[*child].distribution()) + .collect::>(); + let node_statistics = self.graph[node].distribution(); + let propagated_statistics = self.graph[node] + .expr + .propagate_statistics(node_statistics, &children_stats)?; + if let Some(propagated_stats) = propagated_statistics { + for (child_idx, stats) in children.into_iter().zip(propagated_stats) { + self.graph[child_idx].dist = stats; + } + } else { + // The constraint is infeasible, report: + return Ok(PropagationResult::Infeasible); + } + } + Ok(PropagationResult::Success) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::expressions::{binary, try_cast, Column}; + use crate::intervals::cp_solver::PropagationResult; + use crate::statistics::stats_solver::ExprStatisticsGraph; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{Result, ScalarValue}; + use datafusion_expr_common::interval_arithmetic::Interval; + use datafusion_expr_common::operator::Operator; + use datafusion_expr_common::statistics::Distribution; + use datafusion_expr_common::type_coercion::binary::BinaryTypeCoercer; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + + pub fn binary_expr( + left: Arc, + op: Operator, + right: Arc, + schema: &Schema, + ) -> Result> { + let left_type = left.data_type(schema)?; + let right_type = right.data_type(schema)?; + let binary_type_coercer = BinaryTypeCoercer::new(&left_type, &op, &right_type); + let (lhs, rhs) = binary_type_coercer.get_input_types()?; + + let left_expr = try_cast(left, schema, lhs)?; + let right_expr = try_cast(right, schema, rhs)?; + binary(left_expr, op, right_expr, schema) + } + + #[test] + fn test_stats_integration() -> Result<()> { + let schema = &Schema::new(vec![ + Field::new("a", DataType::Float64, false), + Field::new("b", DataType::Float64, false), + Field::new("c", DataType::Float64, false), + Field::new("d", DataType::Float64, false), + ]); + + let a = Arc::new(Column::new("a", 0)) as _; + let b = Arc::new(Column::new("b", 1)) as _; + let c = Arc::new(Column::new("c", 2)) as _; + let d = Arc::new(Column::new("d", 3)) as _; + + let left = binary_expr(a, Operator::Plus, b, schema)?; + let right = binary_expr(c, Operator::Minus, d, schema)?; + let expr = binary_expr(left, Operator::Eq, right, schema)?; + + let mut graph = ExprStatisticsGraph::try_new(expr, schema)?; + // 2, 5 and 6 are BinaryExpr + graph.assign_statistics(&[ + ( + 0usize, + Distribution::new_uniform(Interval::make(Some(0.), Some(1.))?)?, + ), + ( + 1usize, + Distribution::new_uniform(Interval::make(Some(0.), Some(2.))?)?, + ), + ( + 3usize, + Distribution::new_uniform(Interval::make(Some(1.), Some(3.))?)?, + ), + ( + 4usize, + Distribution::new_uniform(Interval::make(Some(1.), Some(5.))?)?, + ), + ]); + let ev_stats = graph.evaluate_statistics()?; + assert_eq!( + ev_stats, + &Distribution::new_bernoulli(ScalarValue::Float64(None))? + ); + + let one = ScalarValue::new_one(&DataType::Float64)?; + assert_eq!( + graph.propagate_statistics(Distribution::new_bernoulli(one)?)?, + PropagationResult::Success + ); + Ok(()) + } +} diff --git a/datafusion/physical-optimizer/Cargo.toml b/datafusion/physical-optimizer/Cargo.toml index c9c86e9c8d5c1..e8473e6556d1d 100644 --- a/datafusion/physical-optimizer/Cargo.toml +++ b/datafusion/physical-optimizer/Cargo.toml @@ -28,6 +28,9 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true diff --git a/datafusion/physical-optimizer/src/lib.rs b/datafusion/physical-optimizer/src/lib.rs index c2beab0320491..2613b95bbdc04 100644 --- a/datafusion/physical-optimizer/src/lib.rs +++ b/datafusion/physical-optimizer/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index f0afdaa2de3d9..4a10398e5a9ef 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -28,6 +28,9 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true @@ -36,7 +39,6 @@ force_hash_collisions = [] [lib] name = "datafusion_physical_plan" -path = "src/lib.rs" [dependencies] ahash = { workspace = true } @@ -63,7 +65,7 @@ pin-project-lite = "^0.2.7" tokio = { workspace = true } [dev-dependencies] -criterion = { version = "0.5", features = ["async_futures"] } +criterion = { workspace = true, features = ["async_futures"] } datafusion-functions-aggregate = { workspace = true } datafusion-functions-window = { workspace = true } rand = { workspace = true } @@ -75,10 +77,6 @@ tokio = { workspace = true, features = [ "parking_lot", ] } -[[bench]] -harness = false -name = "spm" - [[bench]] harness = false name = "partial_ordering" diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index 85b41da85742b..0947a2ff55391 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -1347,10 +1347,10 @@ mod tests { use crate::common::collect; use crate::execution_plan::Boundedness; use crate::expressions::col; - use crate::memory::MemorySourceConfig; use crate::metrics::MetricValue; use crate::test::assert_is_pending; use crate::test::exec::{assert_strong_count_converges_to_zero, BlockingExec}; + use crate::test::TestMemoryExec; use crate::RecordBatchStream; use arrow::array::{ @@ -2207,7 +2207,7 @@ mod tests { vec![test_last_value_agg_expr(&schema, sort_options)?] }; - let memory_exec = MemorySourceConfig::try_new_exec( + let memory_exec = TestMemoryExec::try_new_exec( &[ vec![partition1], vec![partition2], @@ -2442,11 +2442,8 @@ mod tests { }) .collect(); - let input = MemorySourceConfig::try_new_exec( - &[input_batches], - Arc::clone(&schema), - None, - )?; + let input = + TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; let aggregate_exec = Arc::new(AggregateExec::try_new( AggregateMode::Single, @@ -2557,7 +2554,7 @@ mod tests { .build() .map(Arc::new)?]; - let input = MemorySourceConfig::try_new_exec( + let input = TestMemoryExec::try_new_exec( &[vec![batch.clone()]], Arc::::clone(&batch.schema()), None, @@ -2627,7 +2624,7 @@ mod tests { ]; let input = - MemorySourceConfig::try_new_exec(&[input_data], Arc::clone(&schema), None)?; + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; let aggregate_exec = Arc::new(AggregateExec::try_new( AggregateMode::Partial, group_by, @@ -2714,7 +2711,7 @@ mod tests { ]; let input = - MemorySourceConfig::try_new_exec(&[input_data], Arc::clone(&schema), None)?; + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; let aggregate_exec = Arc::new(AggregateExec::try_new( AggregateMode::Partial, group_by, @@ -2829,7 +2826,7 @@ mod tests { create_record_batch(&schema, (vec![2, 3, 4, 4], vec![1.0, 2.0, 3.0, 4.0]))?, ]; let plan: Arc = - MemorySourceConfig::try_new_exec(&[batches], Arc::clone(&schema), None)?; + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; let grouping_set = PhysicalGroupBy::new( vec![(col("a", &schema)?, "a".to_string())], diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index 5168c3cc101f2..c4e738cb3ad15 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -20,10 +20,10 @@ use std::any::Any; use std::sync::Arc; -use super::{common, DisplayAs, PlanProperties, SendableRecordBatchStream, Statistics}; +use crate::memory::MemoryStream; +use crate::{common, DisplayAs, PlanProperties, SendableRecordBatchStream, Statistics}; use crate::{ execution_plan::{Boundedness, EmissionType}, - memory::MemoryStream, DisplayFormatType, ExecutionPlan, Partitioning, }; diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index 5866f0938e41f..a66873bc6576e 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -624,7 +624,6 @@ mod tests { use crate::expressions::*; use crate::test; use crate::test::exec::StatisticsExec; - use arrow::datatypes::{Field, Schema, UnionFields, UnionMode}; use datafusion_common::ScalarValue; diff --git a/datafusion/physical-plan/src/joins/hash_join.rs b/datafusion/physical-plan/src/joins/hash_join.rs index 2983478ada743..23ffd2a28d3c8 100644 --- a/datafusion/physical-plan/src/joins/hash_join.rs +++ b/datafusion/physical-plan/src/joins/hash_join.rs @@ -1638,7 +1638,7 @@ impl EmbeddedProjection for HashJoinExec { #[cfg(test)] mod tests { use super::*; - use crate::memory::MemorySourceConfig; + use crate::test::TestMemoryExec; use crate::{ common, expressions::Column, repartition::RepartitionExec, test::build_table_i32, test::exec::MockExec, @@ -1680,7 +1680,7 @@ mod tests { ) -> Arc { let batch = build_table_i32(a, b, c); let schema = batch.schema(); - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() } fn join( @@ -2083,7 +2083,7 @@ mod tests { build_table_i32(("a1", &vec![2]), ("b2", &vec![2]), ("c1", &vec![9])); let schema = batch1.schema(); let left = - MemorySourceConfig::try_new_exec(&[vec![batch1], vec![batch2]], schema, None) + TestMemoryExec::try_new_exec(&[vec![batch1], vec![batch2]], schema, None) .unwrap(); let right = build_table( @@ -2155,7 +2155,7 @@ mod tests { let schema = batch1.schema(); let left = - MemorySourceConfig::try_new_exec(&[vec![batch1], vec![batch2]], schema, None) + TestMemoryExec::try_new_exec(&[vec![batch1], vec![batch2]], schema, None) .unwrap(); let right = build_table( ("a2", &vec![20, 30, 10]), @@ -2209,7 +2209,7 @@ mod tests { build_table_i32(("a2", &vec![30]), ("b1", &vec![5]), ("c2", &vec![90])); let schema = batch1.schema(); let right = - MemorySourceConfig::try_new_exec(&[vec![batch1], vec![batch2]], schema, None) + TestMemoryExec::try_new_exec(&[vec![batch1], vec![batch2]], schema, None) .unwrap(); let on = vec![( @@ -2288,8 +2288,7 @@ mod tests { ) -> Arc { let batch = build_table_i32(a, b, c); let schema = batch.schema(); - MemorySourceConfig::try_new_exec(&[vec![batch.clone(), batch]], schema, None) - .unwrap() + TestMemoryExec::try_new_exec(&[vec![batch.clone(), batch]], schema, None).unwrap() } #[apply(batch_sizes)] @@ -2394,8 +2393,7 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right.schema()).unwrap()) as _, )]; let schema = right.schema(); - let right = - MemorySourceConfig::try_new_exec(&[vec![right]], schema, None).unwrap(); + let right = TestMemoryExec::try_new_exec(&[vec![right]], schema, None).unwrap(); let join = join(left, right, on, &JoinType::Left, false).unwrap(); let columns = columns(&join.schema()); @@ -2432,8 +2430,7 @@ mod tests { Arc::new(Column::new_with_schema("b2", &right.schema()).unwrap()) as _, )]; let schema = right.schema(); - let right = - MemorySourceConfig::try_new_exec(&[vec![right]], schema, None).unwrap(); + let right = TestMemoryExec::try_new_exec(&[vec![right]], schema, None).unwrap(); let join = join(left, right, on, &JoinType::Full, false).unwrap(); let columns = columns(&join.schema()); @@ -3738,13 +3735,12 @@ mod tests { let n: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); let batch = RecordBatch::try_new(Arc::clone(&schema), vec![dates, n])?; let left = - MemorySourceConfig::try_new_exec(&[vec![batch]], Arc::clone(&schema), None) + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None) .unwrap(); let dates: ArrayRef = Arc::new(Date32Array::from(vec![19108, 19108, 19109])); let n: ArrayRef = Arc::new(Int32Array::from(vec![4, 5, 6])); let batch = RecordBatch::try_new(Arc::clone(&schema), vec![dates, n])?; - let right = - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap(); + let right = TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap(); let on = vec![( Arc::new(Column::new_with_schema("date", &left.schema()).unwrap()) as _, Arc::new(Column::new_with_schema("date", &right.schema()).unwrap()) as _, @@ -4034,7 +4030,7 @@ mod tests { ("b1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), ("c1", &vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 0]), ); - let left = MemorySourceConfig::try_new_exec( + let left = TestMemoryExec::try_new_exec( &[vec![left_batch.clone()], vec![left_batch.clone()]], left_batch.schema(), None, @@ -4045,7 +4041,7 @@ mod tests { ("b2", &vec![12, 13]), ("c2", &vec![14, 15]), ); - let right = MemorySourceConfig::try_new_exec( + let right = TestMemoryExec::try_new_exec( &[vec![right_batch.clone()], vec![right_batch.clone()]], right_batch.schema(), None, @@ -4130,7 +4126,7 @@ mod tests { ) .unwrap(); let schema_ref = batch.schema(); - MemorySourceConfig::try_new_exec(&[vec![batch]], schema_ref, None).unwrap() + TestMemoryExec::try_new_exec(&[vec![batch]], schema_ref, None).unwrap() } #[tokio::test] diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index 6de6b3b4dff44..64dfc8219b644 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -1030,8 +1030,7 @@ impl EmbeddedProjection for NestedLoopJoinExec { #[cfg(test)] pub(crate) mod tests { use super::*; - use crate::memory::MemorySourceConfig; - use crate::source::DataSourceExec; + use crate::test::TestMemoryExec; use crate::{ common, expressions::Column, repartition::RepartitionExec, test::build_table_i32, }; @@ -1072,7 +1071,7 @@ pub(crate) mod tests { }; let mut source = - MemorySourceConfig::try_new(&[batches], Arc::clone(&schema), None).unwrap(); + TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None).unwrap(); if !sorted_column_names.is_empty() { let mut sort_info = LexOrdering::default(); for name in sorted_column_names { @@ -1089,7 +1088,7 @@ pub(crate) mod tests { source = source.try_with_sort_information(vec![sort_info]).unwrap(); } - Arc::new(DataSourceExec::new(Arc::new(source))) + Arc::new(TestMemoryExec::update_cache(Arc::new(source))) } fn build_left_table() -> Arc { diff --git a/datafusion/physical-plan/src/joins/sort_merge_join.rs b/datafusion/physical-plan/src/joins/sort_merge_join.rs index a3e835c64131c..6c933ca218078 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join.rs @@ -2547,7 +2547,7 @@ mod tests { use crate::joins::sort_merge_join::{get_corrected_filter_mask, JoinedRecordBatches}; use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; use crate::joins::SortMergeJoinExec; - use crate::memory::MemorySourceConfig; + use crate::test::TestMemoryExec; use crate::test::{build_table_i32, build_table_i32_two_cols}; use crate::{common, ExecutionPlan}; @@ -2558,12 +2558,12 @@ mod tests { ) -> Arc { let batch = build_table_i32(a, b, c); let schema = batch.schema(); - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() } fn build_table_from_batches(batches: Vec) -> Arc { let schema = batches.first().unwrap().schema(); - MemorySourceConfig::try_new_exec(&[batches], schema, None).unwrap() + TestMemoryExec::try_new_exec(&[batches], schema, None).unwrap() } fn build_date_table( @@ -2588,7 +2588,7 @@ mod tests { .unwrap(); let schema = batch.schema(); - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() } fn build_date64_table( @@ -2613,7 +2613,7 @@ mod tests { .unwrap(); let schema = batch.schema(); - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() } /// returns a table with 3 columns of i32 in memory @@ -2636,7 +2636,7 @@ mod tests { ], ) .unwrap(); - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() } pub fn build_table_two_cols( @@ -2645,7 +2645,7 @@ mod tests { ) -> Arc { let batch = build_table_i32_two_cols(a, b); let schema = batch.schema(); - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() } fn join( diff --git a/datafusion/physical-plan/src/joins/test_utils.rs b/datafusion/physical-plan/src/joins/test_utils.rs index 9932c647be0ab..e70007aa651f7 100644 --- a/datafusion/physical-plan/src/joins/test_utils.rs +++ b/datafusion/physical-plan/src/joins/test_utils.rs @@ -23,9 +23,8 @@ use crate::joins::utils::{JoinFilter, JoinOn}; use crate::joins::{ HashJoinExec, PartitionMode, StreamJoinPartitionMode, SymmetricHashJoinExec, }; -use crate::memory::MemorySourceConfig; use crate::repartition::RepartitionExec; -use crate::source::DataSourceExec; +use crate::test::TestMemoryExec; use crate::{common, ExecutionPlan, ExecutionPlanProperties, Partitioning}; use arrow::array::{ @@ -530,14 +529,14 @@ pub fn create_memory_table( right_sorted: Vec, ) -> Result<(Arc, Arc)> { let left_schema = left_partition[0].schema(); - let left = MemorySourceConfig::try_new(&[left_partition], left_schema, None)? + let left = TestMemoryExec::try_new(&[left_partition], left_schema, None)? .try_with_sort_information(left_sorted)?; let right_schema = right_partition[0].schema(); - let right = MemorySourceConfig::try_new(&[right_partition], right_schema, None)? + let right = TestMemoryExec::try_new(&[right_partition], right_schema, None)? .try_with_sort_information(right_sorted)?; Ok(( - Arc::new(DataSourceExec::new(Arc::new(left))), - Arc::new(DataSourceExec::new(Arc::new(right))), + Arc::new(TestMemoryExec::update_cache(Arc::new(left))), + Arc::new(TestMemoryExec::update_cache(Arc::new(right))), )) } diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 9210e3b0273c3..6ddaef1a2d280 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -15,8 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 - #![deny(clippy::clone_on_ref_ptr)] //! Traits for physical query plan, supporting parallel execution for partitioned relations. @@ -79,13 +83,11 @@ pub mod unnest; pub mod values; pub mod windows; pub mod work_table; - pub mod udaf { pub use datafusion_expr::StatisticsArgs; pub use datafusion_physical_expr::aggregate::AggregateFunctionExpr; } pub mod coalesce; -pub mod source; #[cfg(test)] pub mod test; diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index 0077804bdfc9a..fd338cc913538 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -22,696 +22,22 @@ use std::fmt; use std::sync::Arc; use std::task::{Context, Poll}; -use super::{ - common, ColumnarValue, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, - PhysicalExpr, PlanProperties, RecordBatchStream, SendableRecordBatchStream, - Statistics, -}; use crate::execution_plan::{Boundedness, EmissionType}; -use crate::projection::{ - all_alias_free_columns, new_projections_for_columns, ProjectionExec, +use crate::{ + DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, + RecordBatchStream, SendableRecordBatchStream, Statistics, }; -use crate::source::{DataSource, DataSourceExec}; -use arrow::array::{RecordBatch, RecordBatchOptions}; -use arrow::datatypes::{Schema, SchemaRef}; -use datafusion_common::{ - internal_err, plan_err, project_schema, Constraints, Result, ScalarValue, -}; +use arrow::array::RecordBatch; +use arrow::datatypes::SchemaRef; +use datafusion_common::{internal_err, Result}; use datafusion_execution::memory_pool::MemoryReservation; use datafusion_execution::TaskContext; -use datafusion_physical_expr::equivalence::ProjectionMapping; -use datafusion_physical_expr::expressions::Column; -use datafusion_physical_expr::utils::collect_columns; -use datafusion_physical_expr::{EquivalenceProperties, LexOrdering}; +use datafusion_physical_expr::EquivalenceProperties; use futures::Stream; use parking_lot::RwLock; -/// Execution plan for reading in-memory batches of data -#[derive(Clone)] -#[deprecated( - since = "46.0.0", - note = "use MemorySourceConfig and DataSourceExec instead" -)] -pub struct MemoryExec { - inner: DataSourceExec, - /// The partitions to query - partitions: Vec>, - /// Optional projection - projection: Option>, - // Sort information: one or more equivalent orderings - sort_information: Vec, - /// if partition sizes should be displayed - show_sizes: bool, -} - -#[allow(unused, deprecated)] -impl fmt::Debug for MemoryExec { - fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - self.inner.fmt_as(DisplayFormatType::Default, f) - } -} - -#[allow(unused, deprecated)] -impl DisplayAs for MemoryExec { - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - self.inner.fmt_as(t, f) - } -} - -#[allow(unused, deprecated)] -impl ExecutionPlan for MemoryExec { - fn name(&self) -> &'static str { - "MemoryExec" - } - - /// Return a reference to Any that can be used for downcasting - fn as_any(&self) -> &dyn Any { - self - } - - fn properties(&self) -> &PlanProperties { - self.inner.properties() - } - - fn children(&self) -> Vec<&Arc> { - // This is a leaf node and has no children - vec![] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> Result> { - // MemoryExec has no children - if children.is_empty() { - Ok(self) - } else { - internal_err!("Children cannot be replaced in {self:?}") - } - } - - fn execute( - &self, - partition: usize, - context: Arc, - ) -> Result { - self.inner.execute(partition, context) - } - - /// We recompute the statistics dynamically from the arrow metadata as it is pretty cheap to do so - fn statistics(&self) -> Result { - self.inner.statistics() - } - - fn try_swapping_with_projection( - &self, - projection: &ProjectionExec, - ) -> Result>> { - self.inner.try_swapping_with_projection(projection) - } -} - -#[allow(unused, deprecated)] -impl MemoryExec { - /// Create a new execution plan for reading in-memory record batches - /// The provided `schema` should not have the projection applied. - pub fn try_new( - partitions: &[Vec], - schema: SchemaRef, - projection: Option>, - ) -> Result { - let source = MemorySourceConfig::try_new(partitions, schema, projection.clone())?; - let data_source = DataSourceExec::new(Arc::new(source)); - Ok(Self { - inner: data_source, - partitions: partitions.to_vec(), - projection, - sort_information: vec![], - show_sizes: true, - }) - } - - /// Create a new execution plan from a list of constant values (`ValuesExec`) - pub fn try_new_as_values( - schema: SchemaRef, - data: Vec>>, - ) -> Result { - if data.is_empty() { - return plan_err!("Values list cannot be empty"); - } - - let n_row = data.len(); - let n_col = schema.fields().len(); - - // We have this single row batch as a placeholder to satisfy evaluation argument - // and generate a single output row - let placeholder_schema = Arc::new(Schema::empty()); - let placeholder_batch = RecordBatch::try_new_with_options( - Arc::clone(&placeholder_schema), - vec![], - &RecordBatchOptions::new().with_row_count(Some(1)), - )?; - - // Evaluate each column - let arrays = (0..n_col) - .map(|j| { - (0..n_row) - .map(|i| { - let expr = &data[i][j]; - let result = expr.evaluate(&placeholder_batch)?; - - match result { - ColumnarValue::Scalar(scalar) => Ok(scalar), - ColumnarValue::Array(array) if array.len() == 1 => { - ScalarValue::try_from_array(&array, 0) - } - ColumnarValue::Array(_) => { - plan_err!("Cannot have array values in a values list") - } - } - }) - .collect::>>() - .and_then(ScalarValue::iter_to_array) - }) - .collect::>>()?; - - let batch = RecordBatch::try_new_with_options( - Arc::clone(&schema), - arrays, - &RecordBatchOptions::new().with_row_count(Some(n_row)), - )?; - - let partitions = vec![batch]; - Self::try_new_from_batches(Arc::clone(&schema), partitions) - } - - /// Create a new plan using the provided schema and batches. - /// - /// Errors if any of the batches don't match the provided schema, or if no - /// batches are provided. - pub fn try_new_from_batches( - schema: SchemaRef, - batches: Vec, - ) -> Result { - if batches.is_empty() { - return plan_err!("Values list cannot be empty"); - } - - for batch in &batches { - let batch_schema = batch.schema(); - if batch_schema != schema { - return plan_err!( - "Batch has invalid schema. Expected: {}, got: {}", - schema, - batch_schema - ); - } - } - - let partitions = vec![batches]; - let source = MemorySourceConfig { - partitions: partitions.clone(), - schema: Arc::clone(&schema), - projected_schema: Arc::clone(&schema), - projection: None, - sort_information: vec![], - show_sizes: true, - fetch: None, - }; - let data_source = DataSourceExec::new(Arc::new(source)); - Ok(Self { - inner: data_source, - partitions, - projection: None, - sort_information: vec![], - show_sizes: true, - }) - } - - fn memory_source_config(&self) -> MemorySourceConfig { - self.inner - .source() - .as_any() - .downcast_ref::() - .unwrap() - .clone() - } - - pub fn with_constraints(mut self, constraints: Constraints) -> Self { - self.inner = self.inner.with_constraints(constraints); - self - } - - /// Set `show_sizes` to determine whether to display partition sizes - pub fn with_show_sizes(mut self, show_sizes: bool) -> Self { - let mut memory_source = self.memory_source_config(); - memory_source.show_sizes = show_sizes; - self.show_sizes = show_sizes; - self.inner = DataSourceExec::new(Arc::new(memory_source)); - self - } - - /// Ref to constraints - pub fn constraints(&self) -> &Constraints { - self.properties().equivalence_properties().constraints() - } - - /// Ref to partitions - pub fn partitions(&self) -> &[Vec] { - &self.partitions - } - - /// Ref to projection - pub fn projection(&self) -> &Option> { - &self.projection - } - - /// Show sizes - pub fn show_sizes(&self) -> bool { - self.show_sizes - } - - /// Ref to sort information - pub fn sort_information(&self) -> &[LexOrdering] { - &self.sort_information - } - - /// A memory table can be ordered by multiple expressions simultaneously. - /// [`EquivalenceProperties`] keeps track of expressions that describe the - /// global ordering of the schema. These columns are not necessarily same; e.g. - /// ```text - /// ┌-------┐ - /// | a | b | - /// |---|---| - /// | 1 | 9 | - /// | 2 | 8 | - /// | 3 | 7 | - /// | 5 | 5 | - /// └---┴---┘ - /// ``` - /// where both `a ASC` and `b DESC` can describe the table ordering. With - /// [`EquivalenceProperties`], we can keep track of these equivalences - /// and treat `a ASC` and `b DESC` as the same ordering requirement. - /// - /// Note that if there is an internal projection, that projection will be - /// also applied to the given `sort_information`. - pub fn try_with_sort_information( - mut self, - sort_information: Vec, - ) -> Result { - self.sort_information = sort_information.clone(); - let mut memory_source = self.memory_source_config(); - memory_source = memory_source.try_with_sort_information(sort_information)?; - self.inner = DataSourceExec::new(Arc::new(memory_source)); - Ok(self) - } - - /// Arc clone of ref to original schema - pub fn original_schema(&self) -> SchemaRef { - Arc::clone(&self.inner.schema()) - } - - /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. - fn compute_properties( - schema: SchemaRef, - orderings: &[LexOrdering], - constraints: Constraints, - partitions: &[Vec], - ) -> PlanProperties { - PlanProperties::new( - EquivalenceProperties::new_with_orderings(schema, orderings) - .with_constraints(constraints), - Partitioning::UnknownPartitioning(partitions.len()), - EmissionType::Incremental, - Boundedness::Bounded, - ) - } -} - -/// Data source configuration for reading in-memory batches of data -#[derive(Clone)] -pub struct MemorySourceConfig { - /// The partitions to query - partitions: Vec>, - /// Schema representing the data before projection - schema: SchemaRef, - /// Schema representing the data after the optional projection is applied - projected_schema: SchemaRef, - /// Optional projection - projection: Option>, - /// Sort information: one or more equivalent orderings - sort_information: Vec, - /// if partition sizes should be displayed - show_sizes: bool, - /// The maximum number of records to read from this plan. If `None`, - /// all records after filtering are returned. - fetch: Option, -} - -impl DataSource for MemorySourceConfig { - fn open( - &self, - partition: usize, - _context: Arc, - ) -> Result { - Ok(Box::pin( - MemoryStream::try_new( - self.partitions[partition].clone(), - Arc::clone(&self.projected_schema), - self.projection.clone(), - )? - .with_fetch(self.fetch), - )) - } - - fn as_any(&self) -> &dyn Any { - self - } - - fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { - match t { - DisplayFormatType::Default | DisplayFormatType::Verbose => { - let partition_sizes: Vec<_> = - self.partitions.iter().map(|b| b.len()).collect(); - - let output_ordering = self - .sort_information - .first() - .map(|output_ordering| { - format!(", output_ordering={}", output_ordering) - }) - .unwrap_or_default(); - - let eq_properties = self.eq_properties(); - let constraints = eq_properties.constraints(); - let constraints = if constraints.is_empty() { - String::new() - } else { - format!(", {}", constraints) - }; - - let limit = self - .fetch - .map_or(String::new(), |limit| format!(", fetch={}", limit)); - if self.show_sizes { - write!( - f, - "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", - partition_sizes.len(), - ) - } else { - write!( - f, - "partitions={}{limit}{output_ordering}{constraints}", - partition_sizes.len(), - ) - } - } - } - } - - fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(self.partitions.len()) - } - - fn eq_properties(&self) -> EquivalenceProperties { - EquivalenceProperties::new_with_orderings( - Arc::clone(&self.projected_schema), - self.sort_information.as_slice(), - ) - } - - fn statistics(&self) -> Result { - Ok(common::compute_record_batch_statistics( - &self.partitions, - &self.schema, - self.projection.clone(), - )) - } - - fn with_fetch(&self, limit: Option) -> Option> { - let source = self.clone(); - Some(Arc::new(source.with_limit(limit))) - } - - fn fetch(&self) -> Option { - self.fetch - } - - fn try_swapping_with_projection( - &self, - projection: &ProjectionExec, - ) -> Result>> { - // If there is any non-column or alias-carrier expression, Projection should not be removed. - // This process can be moved into MemoryExec, but it would be an overlap of their responsibility. - all_alias_free_columns(projection.expr()) - .then(|| { - let all_projections = (0..self.schema.fields().len()).collect(); - let new_projections = new_projections_for_columns( - projection, - self.projection().as_ref().unwrap_or(&all_projections), - ); - - MemorySourceConfig::try_new_exec( - self.partitions(), - self.original_schema(), - Some(new_projections), - ) - .map(|e| e as _) - }) - .transpose() - } -} - -impl MemorySourceConfig { - /// Create a new `MemorySourceConfig` for reading in-memory record batches - /// The provided `schema` should not have the projection applied. - pub fn try_new( - partitions: &[Vec], - schema: SchemaRef, - projection: Option>, - ) -> Result { - let projected_schema = project_schema(&schema, projection.as_ref())?; - Ok(Self { - partitions: partitions.to_vec(), - schema, - projected_schema, - projection, - sort_information: vec![], - show_sizes: true, - fetch: None, - }) - } - - /// Create a new `DataSourceExec` plan for reading in-memory record batches - /// The provided `schema` should not have the projection applied. - pub fn try_new_exec( - partitions: &[Vec], - schema: SchemaRef, - projection: Option>, - ) -> Result> { - let source = Self::try_new(partitions, schema, projection)?; - Ok(Arc::new(DataSourceExec::new(Arc::new(source)))) - } - - /// Create a new execution plan from a list of constant values (`ValuesExec`) - pub fn try_new_as_values( - schema: SchemaRef, - data: Vec>>, - ) -> Result> { - if data.is_empty() { - return plan_err!("Values list cannot be empty"); - } - - let n_row = data.len(); - let n_col = schema.fields().len(); - - // We have this single row batch as a placeholder to satisfy evaluation argument - // and generate a single output row - let placeholder_schema = Arc::new(Schema::empty()); - let placeholder_batch = RecordBatch::try_new_with_options( - Arc::clone(&placeholder_schema), - vec![], - &RecordBatchOptions::new().with_row_count(Some(1)), - )?; - - // Evaluate each column - let arrays = (0..n_col) - .map(|j| { - (0..n_row) - .map(|i| { - let expr = &data[i][j]; - let result = expr.evaluate(&placeholder_batch)?; - - match result { - ColumnarValue::Scalar(scalar) => Ok(scalar), - ColumnarValue::Array(array) if array.len() == 1 => { - ScalarValue::try_from_array(&array, 0) - } - ColumnarValue::Array(_) => { - plan_err!("Cannot have array values in a values list") - } - } - }) - .collect::>>() - .and_then(ScalarValue::iter_to_array) - }) - .collect::>>()?; - - let batch = RecordBatch::try_new_with_options( - Arc::clone(&schema), - arrays, - &RecordBatchOptions::new().with_row_count(Some(n_row)), - )?; - - let partitions = vec![batch]; - Self::try_new_from_batches(Arc::clone(&schema), partitions) - } - - /// Create a new plan using the provided schema and batches. - /// - /// Errors if any of the batches don't match the provided schema, or if no - /// batches are provided. - pub fn try_new_from_batches( - schema: SchemaRef, - batches: Vec, - ) -> Result> { - if batches.is_empty() { - return plan_err!("Values list cannot be empty"); - } - - for batch in &batches { - let batch_schema = batch.schema(); - if batch_schema != schema { - return plan_err!( - "Batch has invalid schema. Expected: {}, got: {}", - schema, - batch_schema - ); - } - } - - let partitions = vec![batches]; - let source = Self { - partitions, - schema: Arc::clone(&schema), - projected_schema: Arc::clone(&schema), - projection: None, - sort_information: vec![], - show_sizes: true, - fetch: None, - }; - Ok(Arc::new(DataSourceExec::new(Arc::new(source)))) - } - - /// Set the limit of the files - pub fn with_limit(mut self, limit: Option) -> Self { - self.fetch = limit; - self - } - - /// Set `show_sizes` to determine whether to display partition sizes - pub fn with_show_sizes(mut self, show_sizes: bool) -> Self { - self.show_sizes = show_sizes; - self - } - - /// Ref to partitions - pub fn partitions(&self) -> &[Vec] { - &self.partitions - } - - /// Ref to projection - pub fn projection(&self) -> &Option> { - &self.projection - } - - /// Show sizes - pub fn show_sizes(&self) -> bool { - self.show_sizes - } - - /// Ref to sort information - pub fn sort_information(&self) -> &[LexOrdering] { - &self.sort_information - } - - /// A memory table can be ordered by multiple expressions simultaneously. - /// [`EquivalenceProperties`] keeps track of expressions that describe the - /// global ordering of the schema. These columns are not necessarily same; e.g. - /// ```text - /// ┌-------┐ - /// | a | b | - /// |---|---| - /// | 1 | 9 | - /// | 2 | 8 | - /// | 3 | 7 | - /// | 5 | 5 | - /// └---┴---┘ - /// ``` - /// where both `a ASC` and `b DESC` can describe the table ordering. With - /// [`EquivalenceProperties`], we can keep track of these equivalences - /// and treat `a ASC` and `b DESC` as the same ordering requirement. - /// - /// Note that if there is an internal projection, that projection will be - /// also applied to the given `sort_information`. - pub fn try_with_sort_information( - mut self, - mut sort_information: Vec, - ) -> Result { - // All sort expressions must refer to the original schema - let fields = self.schema.fields(); - let ambiguous_column = sort_information - .iter() - .flat_map(|ordering| ordering.clone()) - .flat_map(|expr| collect_columns(&expr.expr)) - .find(|col| { - fields - .get(col.index()) - .map(|field| field.name() != col.name()) - .unwrap_or(true) - }); - if let Some(col) = ambiguous_column { - return internal_err!( - "Column {:?} is not found in the original schema of the MemorySourceConfig", - col - ); - } - - // If there is a projection on the source, we also need to project orderings - if let Some(projection) = &self.projection { - let base_eqp = EquivalenceProperties::new_with_orderings( - self.original_schema(), - &sort_information, - ); - let proj_exprs = projection - .iter() - .map(|idx| { - let base_schema = self.original_schema(); - let name = base_schema.field(*idx).name(); - (Arc::new(Column::new(name, *idx)) as _, name.to_string()) - }) - .collect::>(); - let projection_mapping = - ProjectionMapping::try_new(&proj_exprs, &self.original_schema())?; - sort_information = base_eqp - .project(&projection_mapping, Arc::clone(&self.projected_schema)) - .into_oeq_class() - .into_inner(); - } - - self.sort_information = sort_information; - Ok(self) - } - - /// Arc clone of ref to original schema - pub fn original_schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} - /// Iterator over batches pub struct MemoryStream { /// Vector of record batches @@ -746,13 +72,13 @@ impl MemoryStream { } /// Set the memory reservation for the data - pub(super) fn with_reservation(mut self, reservation: MemoryReservation) -> Self { + pub fn with_reservation(mut self, reservation: MemoryReservation) -> Self { self.reservation = Some(reservation); self } /// Set the number of rows to produce - pub(super) fn with_fetch(mut self, fetch: Option) -> Self { + pub fn with_fetch(mut self, fetch: Option) -> Self { self.fetch = fetch; self } @@ -962,62 +288,6 @@ impl RecordBatchStream for LazyMemoryStream { } } -#[cfg(test)] -mod memory_exec_tests { - use std::sync::Arc; - - use crate::memory::MemorySourceConfig; - use crate::source::DataSourceExec; - use crate::ExecutionPlan; - - use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_physical_expr::expressions::col; - use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_physical_expr_common::sort_expr::LexOrdering; - - #[test] - fn test_memory_order_eq() -> datafusion_common::Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - let sort1 = LexOrdering::new(vec![ - PhysicalSortExpr { - expr: col("a", &schema)?, - options: SortOptions::default(), - }, - PhysicalSortExpr { - expr: col("b", &schema)?, - options: SortOptions::default(), - }, - ]); - let sort2 = LexOrdering::new(vec![PhysicalSortExpr { - expr: col("c", &schema)?, - options: SortOptions::default(), - }]); - let mut expected_output_order = LexOrdering::default(); - expected_output_order.extend(sort1.clone()); - expected_output_order.extend(sort2.clone()); - - let sort_information = vec![sort1.clone(), sort2.clone()]; - let mem_exec = Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[vec![]], schema, None)? - .try_with_sort_information(sort_information)?, - ))); - - assert_eq!( - mem_exec.properties().output_ordering().unwrap(), - &expected_output_order - ); - let eq_properties = mem_exec.properties().equivalence_properties(); - assert!(eq_properties.oeq_class().contains(&sort1)); - assert!(eq_properties.oeq_class().contains(&sort2)); - Ok(()) - } -} - #[cfg(test)] mod lazy_memory_tests { use super::*; @@ -1137,129 +407,3 @@ mod lazy_memory_tests { Ok(()) } } - -#[cfg(test)] -mod tests { - use super::*; - use crate::expressions::lit; - use crate::test::{self, make_partition}; - - use arrow::datatypes::{DataType, Field}; - use datafusion_common::assert_batches_eq; - use datafusion_common::stats::{ColumnStatistics, Precision}; - use futures::StreamExt; - - #[tokio::test] - async fn exec_with_limit() -> Result<()> { - let task_ctx = Arc::new(TaskContext::default()); - let batch = make_partition(7); - let schema = batch.schema(); - let batches = vec![batch.clone(), batch]; - - let exec = MemorySourceConfig::try_new_from_batches(schema, batches).unwrap(); - assert_eq!(exec.fetch(), None); - - let exec = exec.with_fetch(Some(4)).unwrap(); - assert_eq!(exec.fetch(), Some(4)); - - let mut it = exec.execute(0, task_ctx)?; - let mut results = vec![]; - while let Some(batch) = it.next().await { - results.push(batch?); - } - - let expected = [ - "+---+", "| i |", "+---+", "| 0 |", "| 1 |", "| 2 |", "| 3 |", "+---+", - ]; - assert_batches_eq!(expected, &results); - Ok(()) - } - - #[tokio::test] - async fn values_empty_case() -> Result<()> { - let schema = test::aggr_test_schema(); - let empty = MemorySourceConfig::try_new_as_values(schema, vec![]); - assert!(empty.is_err()); - Ok(()) - } - - #[test] - fn new_exec_with_batches() { - let batch = make_partition(7); - let schema = batch.schema(); - let batches = vec![batch.clone(), batch]; - let _exec = MemorySourceConfig::try_new_from_batches(schema, batches).unwrap(); - } - - #[test] - fn new_exec_with_batches_empty() { - let batch = make_partition(7); - let schema = batch.schema(); - let _ = MemorySourceConfig::try_new_from_batches(schema, Vec::new()).unwrap_err(); - } - - #[test] - fn new_exec_with_batches_invalid_schema() { - let batch = make_partition(7); - let batches = vec![batch.clone(), batch]; - - let invalid_schema = Arc::new(Schema::new(vec![ - Field::new("col0", DataType::UInt32, false), - Field::new("col1", DataType::Utf8, false), - ])); - let _ = MemorySourceConfig::try_new_from_batches(invalid_schema, batches) - .unwrap_err(); - } - - // Test issue: https://github.com/apache/datafusion/issues/8763 - #[test] - fn new_exec_with_non_nullable_schema() { - let schema = Arc::new(Schema::new(vec![Field::new( - "col0", - DataType::UInt32, - false, - )])); - let _ = MemorySourceConfig::try_new_as_values( - Arc::clone(&schema), - vec![vec![lit(1u32)]], - ) - .unwrap(); - // Test that a null value is rejected - let _ = MemorySourceConfig::try_new_as_values( - schema, - vec![vec![lit(ScalarValue::UInt32(None))]], - ) - .unwrap_err(); - } - - #[test] - fn values_stats_with_nulls_only() -> Result<()> { - let data = vec![ - vec![lit(ScalarValue::Null)], - vec![lit(ScalarValue::Null)], - vec![lit(ScalarValue::Null)], - ]; - let rows = data.len(); - let values = MemorySourceConfig::try_new_as_values( - Arc::new(Schema::new(vec![Field::new("col0", DataType::Null, true)])), - data, - )?; - - assert_eq!( - values.statistics()?, - Statistics { - num_rows: Precision::Exact(rows), - total_byte_size: Precision::Exact(8), // not important - column_statistics: vec![ColumnStatistics { - null_count: Precision::Exact(rows), // there are only nulls - distinct_count: Precision::Absent, - max_value: Precision::Absent, - min_value: Precision::Absent, - sum_value: Precision::Absent, - },], - } - ); - - Ok(()) - } -} diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index 6a8f247ec0e68..6e31f601e1525 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -20,10 +20,10 @@ use std::any::Any; use std::sync::Arc; -use super::{common, DisplayAs, PlanProperties, SendableRecordBatchStream, Statistics}; use crate::execution_plan::{Boundedness, EmissionType}; -use crate::{memory::MemoryStream, DisplayFormatType, ExecutionPlan, Partitioning}; - +use crate::memory::MemoryStream; +use crate::{common, DisplayAs, PlanProperties, SendableRecordBatchStream, Statistics}; +use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; use arrow::array::{ArrayRef, NullArray}; use arrow::array::{RecordBatch, RecordBatchOptions}; use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; @@ -178,7 +178,8 @@ impl ExecutionPlan for PlaceholderRowExec { #[cfg(test)] mod tests { use super::*; - use crate::{test, with_new_children_if_necessary}; + use crate::test; + use crate::with_new_children_if_necessary; #[test] fn with_new_children() -> Result<()> { diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index bf7d2c7f275c9..05b78e4e1da42 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -21,12 +21,12 @@ use std::any::Any; use std::sync::Arc; use std::task::{Context, Poll}; -use super::{ +use super::work_table::{ReservedBatches, WorkTable, WorkTableExec}; +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::{ metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}, - work_table::{ReservedBatches, WorkTable, WorkTableExec}, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; -use crate::execution_plan::{Boundedness, EmissionType}; use crate::{DisplayAs, DisplayFormatType, ExecutionPlan}; use arrow::datatypes::SchemaRef; @@ -156,10 +156,10 @@ impl ExecutionPlan for RecursiveQueryExec { vec![false, false] } - fn required_input_distribution(&self) -> Vec { + fn required_input_distribution(&self) -> Vec { vec![ - datafusion_physical_expr::Distribution::SinglePartition, - datafusion_physical_expr::Distribution::SinglePartition, + crate::Distribution::SinglePartition, + crate::Distribution::SinglePartition, ] } diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 25668fa67d5b0..40e68cfcae837 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -1051,6 +1051,7 @@ mod tests { use std::collections::HashSet; use super::*; + use crate::test::TestMemoryExec; use crate::{ test::{ assert_is_pending, @@ -1059,7 +1060,7 @@ mod tests { ErrorExec, MockExec, }, }, - {collect, expressions::col, memory::MemorySourceConfig}, + {collect, expressions::col}, }; use arrow::array::{ArrayRef, StringArray, UInt32Array}; @@ -1164,11 +1165,8 @@ mod tests { ) -> Result>> { let task_ctx = Arc::new(TaskContext::default()); // create physical plan - let exec = MemorySourceConfig::try_new_exec( - &input_partitions, - Arc::clone(schema), - None, - )?; + let exec = + TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(schema), None)?; let exec = RepartitionExec::try_new(exec, partitioning)?; // execute and collect results @@ -1559,11 +1557,8 @@ mod tests { let task_ctx = Arc::new(task_ctx); // create physical plan - let exec = MemorySourceConfig::try_new_exec( - &input_partitions, - Arc::clone(&schema), - None, - )?; + let exec = + TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(&schema), None)?; let exec = RepartitionExec::try_new(exec, partitioning)?; // pull partitions @@ -1604,8 +1599,7 @@ mod test { use arrow::datatypes::{DataType, Field, Schema}; use super::*; - use crate::memory::MemorySourceConfig; - use crate::source::DataSourceExec; + use crate::test::TestMemoryExec; use crate::union::UnionExec; use datafusion_physical_expr::expressions::col; @@ -1711,15 +1705,15 @@ mod test { } fn memory_exec(schema: &SchemaRef) -> Arc { - MemorySourceConfig::try_new_exec(&[vec![]], Arc::clone(schema), None).unwrap() + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(schema), None).unwrap() } fn sorted_memory_exec( schema: &SchemaRef, sort_exprs: LexOrdering, ) -> Arc { - Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[vec![]], Arc::clone(schema), None) + Arc::new(TestMemoryExec::update_cache(Arc::new( + TestMemoryExec::try_new(&[vec![]], Arc::clone(schema), None) .unwrap() .try_with_sort_information(vec![sort_exprs]) .unwrap(), diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index eeef73c45fc48..dc03c012d9be4 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -466,11 +466,11 @@ mod tests { use crate::collect; use crate::expressions::col; use crate::expressions::PhysicalSortExpr; - use crate::memory::MemorySourceConfig; use crate::sorts::sort::SortExec; use crate::test; use crate::test::assert_is_pending; use crate::test::exec::{assert_strong_count_converges_to_zero, BlockingExec}; + use crate::test::TestMemoryExec; use super::*; @@ -696,7 +696,7 @@ mod tests { ); let schema = batch1.schema(); - MemorySourceConfig::try_new_exec( + TestMemoryExec::try_new_exec( &[vec![batch1, batch2, batch3, batch4]], Arc::clone(&schema), None, @@ -881,7 +881,7 @@ mod tests { let batch = RecordBatch::try_new(Arc::clone(&schema), vec![data])?; let input = - MemorySourceConfig::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?; let partial_sort_exec = Arc::new(PartialSortExec::new( LexOrdering::new(vec![PhysicalSortExpr { @@ -987,7 +987,7 @@ mod tests { options: option_desc, }, ]), - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None)?, + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None)?, 2, )); diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 649468260e560..30b5abcf88970 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -1142,10 +1142,10 @@ mod tests { use crate::collect; use crate::execution_plan::Boundedness; use crate::expressions::col; - use crate::memory::MemorySourceConfig; use crate::test; use crate::test::assert_is_pending; use crate::test::exec::{assert_strong_count_converges_to_zero, BlockingExec}; + use crate::test::TestMemoryExec; use arrow::array::*; use arrow::compute::SortOptions; @@ -1531,7 +1531,7 @@ mod tests { let batch = RecordBatch::try_new(Arc::clone(&schema), vec![data]).unwrap(); let input = - MemorySourceConfig::try_new_exec(&[vec![batch]], Arc::clone(&schema), None) + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None) .unwrap(); let sort_exec = Arc::new(SortExec::new( @@ -1602,7 +1602,7 @@ mod tests { }, }, ]), - MemorySourceConfig::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?, + TestMemoryExec::try_new_exec(&[vec![batch]], Arc::clone(&schema), None)?, )); assert_eq!(DataType::Int32, *sort_exec.schema().field(0).data_type()); @@ -1688,7 +1688,7 @@ mod tests { }, }, ]), - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None)?, + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None)?, )); assert_eq!(DataType::Float32, *sort_exec.schema().field(0).data_type()); diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 2cc55d60292a3..454a068551754 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -384,12 +384,12 @@ mod tests { use crate::coalesce_partitions::CoalescePartitionsExec; use crate::execution_plan::{Boundedness, EmissionType}; use crate::expressions::col; - use crate::memory::MemorySourceConfig; use crate::metrics::{MetricValue, Timestamp}; use crate::repartition::RepartitionExec; use crate::sorts::sort::SortExec; use crate::stream::RecordBatchReceiverStream; use crate::test::exec::{assert_strong_count_converges_to_zero, BlockingExec}; + use crate::test::TestMemoryExec; use crate::test::{self, assert_is_pending, make_partition}; use crate::{collect, common}; @@ -451,7 +451,7 @@ mod tests { ]); let repartition_exec = RepartitionExec::try_new( - MemorySourceConfig::try_new_exec(&[rbs], schema, None).unwrap(), + TestMemoryExec::try_new_exec(&[rbs], schema, None).unwrap(), Partitioning::RoundRobinBatch(2), )?; let coalesce_batches_exec = @@ -543,7 +543,7 @@ mod tests { let schema = batch.schema(); let sort = LexOrdering::default(); // no sort expressions - let exec = MemorySourceConfig::try_new_exec( + let exec = TestMemoryExec::try_new_exec( &[vec![batch.clone()], vec![batch]], schema, None, @@ -736,7 +736,7 @@ mod tests { options: Default::default(), }, ]); - let exec = MemorySourceConfig::try_new_exec(partitions, schema, None).unwrap(); + let exec = TestMemoryExec::try_new_exec(partitions, schema, None).unwrap(); let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); let collected = collect(merge, context).await.unwrap(); @@ -844,7 +844,7 @@ mod tests { let sorted = basic_sort(csv, sort, context).await; let split: Vec<_> = sizes.iter().map(|x| split_batch(&sorted, *x)).collect(); - Ok(MemorySourceConfig::try_new_exec(&split, sorted.schema(), None).unwrap()) + Ok(TestMemoryExec::try_new_exec(&split, sorted.schema(), None).unwrap()) } #[tokio::test] @@ -972,8 +972,8 @@ mod tests { }, }, ]); - let exec = MemorySourceConfig::try_new_exec(&[vec![b1], vec![b2]], schema, None) - .unwrap(); + let exec = + TestMemoryExec::try_new_exec(&[vec![b1], vec![b2]], schema, None).unwrap(); let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); let collected = collect(merge, task_ctx).await.unwrap(); @@ -1015,8 +1015,7 @@ mod tests { nulls_first: true, }, }]); - let exec = - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap(); + let exec = TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap(); let merge = Arc::new(SortPreservingMergeExec::new(sort, exec).with_fetch(Some(2))); @@ -1051,8 +1050,7 @@ mod tests { nulls_first: true, }, }]); - let exec = - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap(); + let exec = TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap(); let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); let collected = collect(merge, task_ctx).await.unwrap(); @@ -1161,8 +1159,8 @@ mod tests { expr: col("b", &schema).unwrap(), options: Default::default(), }]); - let exec = MemorySourceConfig::try_new_exec(&[vec![b1], vec![b2]], schema, None) - .unwrap(); + let exec = + TestMemoryExec::try_new_exec(&[vec![b1], vec![b2]], schema, None).unwrap(); let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); let collected = collect(Arc::clone(&merge) as Arc, task_ctx) @@ -1273,7 +1271,7 @@ mod tests { }, }]); - let exec = MemorySourceConfig::try_new_exec(&partitions, schema, None).unwrap(); + let exec = TestMemoryExec::try_new_exec(&partitions, schema, None).unwrap(); let merge = Arc::new(SortPreservingMergeExec::new(sort, exec)); let collected = collect(merge, task_ctx).await.unwrap(); diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index ad0e43503b2b3..7d0e3778452f3 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -17,27 +17,337 @@ //! Utilities for testing datafusion-physical-plan +use std::any::Any; use std::collections::HashMap; +use std::fmt; +use std::fmt::{Debug, Formatter}; use std::pin::Pin; use std::sync::Arc; +use std::task::Context; + +use crate::common; +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::memory::MemoryStream; +use crate::metrics::MetricsSet; +use crate::stream::RecordBatchStreamAdapter; +use crate::streaming::PartitionStream; +use crate::ExecutionPlan; +use crate::{DisplayAs, DisplayFormatType, PlanProperties}; use arrow::array::{Array, ArrayRef, Int32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use datafusion_common::{ + config::ConfigOptions, internal_err, project_schema, Result, Statistics, +}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; -use futures::{Future, FutureExt}; +use datafusion_physical_expr::{ + equivalence::ProjectionMapping, expressions::Column, utils::collect_columns, + EquivalenceProperties, LexOrdering, Partitioning, +}; -use crate::memory::MemorySourceConfig; -use crate::source::DataSourceExec; -use crate::stream::RecordBatchStreamAdapter; -use crate::streaming::PartitionStream; -use crate::ExecutionPlan; +use futures::{Future, FutureExt}; pub mod exec; +/// `TestMemoryExec` is a mock equivalent to [`MemorySourceConfig`] with [`ExecutionPlan`] implemented for testing. +/// i.e. It has some but not all the functionality of [`MemorySourceConfig`]. +/// This implements an in-memory DataSource rather than explicitly implementing a trait. +/// It is implemented in this manner to keep relevant unit tests in place +/// while avoiding circular dependencies between `datafusion-physical-plan` and `datafusion-datasource`. +/// +/// [`MemorySourceConfig`]: https://github.com/apache/datafusion/tree/main/datafusion/datasource/src/memory.rs +#[derive(Clone, Debug)] +pub struct TestMemoryExec { + /// The partitions to query + partitions: Vec>, + /// Schema representing the data before projection + schema: SchemaRef, + /// Schema representing the data after the optional projection is applied + projected_schema: SchemaRef, + /// Optional projection + projection: Option>, + /// Sort information: one or more equivalent orderings + sort_information: Vec, + /// if partition sizes should be displayed + show_sizes: bool, + /// The maximum number of records to read from this plan. If `None`, + /// all records after filtering are returned. + fetch: Option, + cache: PlanProperties, +} + +impl DisplayAs for TestMemoryExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + write!(f, "DataSourceExec: ")?; + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + let partition_sizes: Vec<_> = + self.partitions.iter().map(|b| b.len()).collect(); + + let output_ordering = self + .sort_information + .first() + .map(|output_ordering| { + format!(", output_ordering={}", output_ordering) + }) + .unwrap_or_default(); + + let eq_properties = self.eq_properties(); + let constraints = eq_properties.constraints(); + let constraints = if constraints.is_empty() { + String::new() + } else { + format!(", {}", constraints) + }; + + let limit = self + .fetch + .map_or(String::new(), |limit| format!(", fetch={}", limit)); + if self.show_sizes { + write!( + f, + "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}", + partition_sizes.len(), + ) + } else { + write!( + f, + "partitions={}{limit}{output_ordering}{constraints}", + partition_sizes.len(), + ) + } + } + } + } +} + +impl ExecutionPlan for TestMemoryExec { + fn name(&self) -> &'static str { + "DataSourceExec" + } + + fn as_any(&self) -> &dyn Any { + unimplemented!() + } + + fn properties(&self) -> &PlanProperties { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + Vec::new() + } + + fn with_new_children( + self: Arc, + _: Vec>, + ) -> Result> { + unimplemented!() + } + + fn repartitioned( + &self, + _target_partitions: usize, + _config: &ConfigOptions, + ) -> Result>> { + unimplemented!() + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.open(partition, context) + } + + fn metrics(&self) -> Option { + unimplemented!() + } + + fn statistics(&self) -> Result { + self.statistics() + } + + fn fetch(&self) -> Option { + self.fetch + } +} + +impl TestMemoryExec { + fn open( + &self, + partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin( + MemoryStream::try_new( + self.partitions[partition].clone(), + Arc::clone(&self.projected_schema), + self.projection.clone(), + )? + .with_fetch(self.fetch), + )) + } + + fn compute_properties(&self) -> PlanProperties { + PlanProperties::new( + self.eq_properties(), + self.output_partitioning(), + EmissionType::Incremental, + Boundedness::Bounded, + ) + } + + fn output_partitioning(&self) -> Partitioning { + Partitioning::UnknownPartitioning(self.partitions.len()) + } + + fn eq_properties(&self) -> EquivalenceProperties { + EquivalenceProperties::new_with_orderings( + Arc::clone(&self.projected_schema), + self.sort_information.as_slice(), + ) + } + + fn statistics(&self) -> Result { + Ok(common::compute_record_batch_statistics( + &self.partitions, + &self.schema, + self.projection.clone(), + )) + } + + pub fn try_new( + partitions: &[Vec], + schema: SchemaRef, + projection: Option>, + ) -> Result { + let projected_schema = project_schema(&schema, projection.as_ref())?; + Ok(Self { + partitions: partitions.to_vec(), + schema, + cache: PlanProperties::new( + EquivalenceProperties::new_with_orderings( + Arc::clone(&projected_schema), + vec![].as_slice(), + ), + Partitioning::UnknownPartitioning(partitions.len()), + EmissionType::Incremental, + Boundedness::Bounded, + ), + projected_schema, + projection, + sort_information: vec![], + show_sizes: true, + fetch: None, + }) + } + + /// Create a new `DataSourceExec` Equivalent plan for reading in-memory record batches + /// The provided `schema` should not have the projection applied. + pub fn try_new_exec( + partitions: &[Vec], + schema: SchemaRef, + projection: Option>, + ) -> Result> { + let mut source = Self::try_new(partitions, schema, projection)?; + let cache = source.compute_properties(); + source.cache = cache; + Ok(Arc::new(source)) + } + + // Equivalent of `DataSourceExec::new` + pub fn update_cache(source: Arc) -> TestMemoryExec { + let cache = source.compute_properties(); + let source = &*source; + let mut source = source.clone(); + source.cache = cache; + source + } + + /// Set the limit of the files + pub fn with_limit(mut self, limit: Option) -> Self { + self.fetch = limit; + self + } + + /// Ref to partitions + pub fn partitions(&self) -> &[Vec] { + &self.partitions + } + + /// Ref to projection + pub fn projection(&self) -> &Option> { + &self.projection + } + + /// Ref to sort information + pub fn sort_information(&self) -> &[LexOrdering] { + &self.sort_information + } + + /// refer to `try_with_sort_information` at MemorySourceConfig for more information. + /// https://github.com/apache/datafusion/tree/main/datafusion/datasource/src/memory.rs + pub fn try_with_sort_information( + mut self, + mut sort_information: Vec, + ) -> Result { + // All sort expressions must refer to the original schema + let fields = self.schema.fields(); + let ambiguous_column = sort_information + .iter() + .flat_map(|ordering| ordering.clone()) + .flat_map(|expr| collect_columns(&expr.expr)) + .find(|col| { + fields + .get(col.index()) + .map(|field| field.name() != col.name()) + .unwrap_or(true) + }); + if let Some(col) = ambiguous_column { + return internal_err!( + "Column {:?} is not found in the original schema of the TestMemoryExec", + col + ); + } + + // If there is a projection on the source, we also need to project orderings + if let Some(projection) = &self.projection { + let base_eqp = EquivalenceProperties::new_with_orderings( + self.original_schema(), + &sort_information, + ); + let proj_exprs = projection + .iter() + .map(|idx| { + let base_schema = self.original_schema(); + let name = base_schema.field(*idx).name(); + (Arc::new(Column::new(name, *idx)) as _, name.to_string()) + }) + .collect::>(); + let projection_mapping = + ProjectionMapping::try_new(&proj_exprs, &self.original_schema())?; + sort_information = base_eqp + .project(&projection_mapping, Arc::clone(&self.projected_schema)) + .into_oeq_class() + .into_inner(); + } + + self.sort_information = sort_information; + Ok(self) + } + + /// Arc clone of ref to original schema + pub fn original_schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + /// Asserts that given future is pending. pub fn assert_is_pending<'a, T>(fut: &mut Pin + Send + 'a>>) { let waker = futures::task::noop_waker(); - let mut cx = futures::task::Context::from_waker(&waker); + let mut cx = Context::from_waker(&waker); let poll = fut.poll_unpin(&mut cx); assert!(poll.is_pending()); @@ -117,7 +427,7 @@ pub fn build_table_scan_i32( ) -> Arc { let batch = build_table_i32(a, b, c); let schema = batch.schema(); - MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None).unwrap() + TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap() } /// Return a RecordBatch with a single Int32 array with values (0..sz) in a field named "i" @@ -157,26 +467,24 @@ pub fn scan_partitioned_utf8(partitions: usize) -> Arc { } /// Returns a `DataSourceExec` that scans `partitions` of 100 batches each -pub fn mem_exec(partitions: usize) -> DataSourceExec { +pub fn mem_exec(partitions: usize) -> TestMemoryExec { let data: Vec> = (0..partitions).map(|_| vec![make_partition(100)]).collect(); let schema = data[0][0].schema(); let projection = None; - DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&data, schema, projection).unwrap(), - )) + + TestMemoryExec::try_new(&data, schema, projection).unwrap() } -pub fn mem_exec_utf8(partitions: usize) -> DataSourceExec { +pub fn mem_exec_utf8(partitions: usize) -> TestMemoryExec { let data: Vec> = (0..partitions) .map(|_| vec![make_partition_utf8(100)]) .collect(); let schema = data[0][0].schema(); let projection = None; - DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&data, schema, projection).unwrap(), - )) + + TestMemoryExec::try_new(&data, schema, projection).unwrap() } // Construct a stream partition for test purposes diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index e1972d267b972..68d1803b7133e 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -640,10 +640,9 @@ fn stats_union(mut left: Statistics, right: Statistics) -> Statistics { mod tests { use super::*; use crate::collect; - use crate::memory::MemorySourceConfig; use crate::test; + use crate::test::TestMemoryExec; - use crate::source::DataSourceExec; use arrow::compute::SortOptions; use arrow::datatypes::DataType; use datafusion_common::ScalarValue; @@ -865,12 +864,12 @@ mod tests { .iter() .map(|ordering| convert_to_sort_exprs(ordering)) .collect::>(); - let child1 = Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[], Arc::clone(&schema), None)? + let child1 = Arc::new(TestMemoryExec::update_cache(Arc::new( + TestMemoryExec::try_new(&[], Arc::clone(&schema), None)? .try_with_sort_information(first_orderings)?, ))); - let child2 = Arc::new(DataSourceExec::new(Arc::new( - MemorySourceConfig::try_new(&[], Arc::clone(&schema), None)? + let child2 = Arc::new(TestMemoryExec::update_cache(Arc::new( + TestMemoryExec::try_new(&[], Arc::clone(&schema), None)? .try_with_sort_information(second_orderings)?, ))); diff --git a/datafusion/physical-plan/src/values.rs b/datafusion/physical-plan/src/values.rs index 6ab5cc84a21f5..b90c50510cb03 100644 --- a/datafusion/physical-plan/src/values.rs +++ b/datafusion/physical-plan/src/values.rs @@ -20,13 +20,12 @@ use std::any::Any; use std::sync::Arc; -use super::{common, DisplayAs, PlanProperties, SendableRecordBatchStream, Statistics}; use crate::execution_plan::{Boundedness, EmissionType}; +use crate::memory::MemoryStream; +use crate::{common, DisplayAs, PlanProperties, SendableRecordBatchStream, Statistics}; use crate::{ - memory::MemoryStream, ColumnarValue, DisplayFormatType, ExecutionPlan, Partitioning, - PhysicalExpr, + ColumnarValue, DisplayFormatType, ExecutionPlan, Partitioning, PhysicalExpr, }; - use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::{RecordBatch, RecordBatchOptions}; use datafusion_common::{internal_err, plan_err, Result, ScalarValue}; diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index 1e21d0757c414..c78c870ff3831 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -1190,9 +1190,9 @@ mod tests { use crate::common::collect; use crate::expressions::PhysicalSortExpr; - use crate::memory::MemorySourceConfig; use crate::projection::ProjectionExec; use crate::streaming::{PartitionStream, StreamingTableExec}; + use crate::test::TestMemoryExec; use crate::windows::{ create_udwf_window_expr, create_window_expr, BoundedWindowAggExec, InputOrderMode, }; @@ -1551,7 +1551,7 @@ mod tests { vec![Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3]))], )?; - let memory_exec = MemorySourceConfig::try_new_exec( + let memory_exec = TestMemoryExec::try_new_exec( &[vec![batch.clone(), batch.clone(), batch.clone()]], Arc::clone(&schema), None, diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs index a6ced527cbb2f..d3d29bfad7cec 100644 --- a/datafusion/physical-plan/src/work_table.rs +++ b/datafusion/physical-plan/src/work_table.rs @@ -20,12 +20,12 @@ use std::any::Any; use std::sync::{Arc, Mutex}; -use super::{ +use crate::execution_plan::{Boundedness, EmissionType}; +use crate::memory::MemoryStream; +use crate::{ metrics::{ExecutionPlanMetricsSet, MetricsSet}, SendableRecordBatchStream, Statistics, }; -use crate::execution_plan::{Boundedness, EmissionType}; -use crate::memory::MemoryStream; use crate::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; use arrow::datatypes::SchemaRef; diff --git a/datafusion/proto-common/Cargo.toml b/datafusion/proto-common/Cargo.toml index 3454b26ebcf6f..957cbc253616b 100644 --- a/datafusion/proto-common/Cargo.toml +++ b/datafusion/proto-common/Cargo.toml @@ -31,9 +31,11 @@ rust-version = { workspace = true } # Exclude proto files so crates.io consumers don't need protoc exclude = ["*.proto"] +[package.metadata.docs.rs] +all-features = true + [lib] name = "datafusion_proto_common" -path = "src/lib.rs" [features] default = [] diff --git a/datafusion/proto-common/src/lib.rs b/datafusion/proto-common/src/lib.rs index 9272a4e87960c..56cd42ee5067f 100644 --- a/datafusion/proto-common/src/lib.rs +++ b/datafusion/proto-common/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index fb5d414dcec44..00d4969182cf6 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -31,9 +31,11 @@ rust-version = { workspace = true } # Exclude proto files so crates.io consumers don't need protoc exclude = ["*.proto"] +[package.metadata.docs.rs] +all-features = true + [lib] name = "datafusion_proto" -path = "src/lib.rs" [features] default = ["parquet"] diff --git a/datafusion/proto/src/lib.rs b/datafusion/proto/src/lib.rs index 9f2973e3c958e..5d84be1cff55b 100644 --- a/datafusion/proto/src/lib.rs +++ b/datafusion/proto/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index a575a42d0b6c4..2c596255587b2 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -32,6 +32,7 @@ use datafusion::datasource::file_format::parquet::ParquetSink; #[cfg(feature = "parquet")] use datafusion::datasource::physical_plan::ParquetSource; use datafusion::datasource::physical_plan::{AvroSource, CsvSource, FileScanConfig}; +use datafusion::datasource::source::DataSourceExec; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::execution::FunctionRegistry; use datafusion::physical_expr::aggregate::AggregateFunctionExpr; @@ -57,7 +58,6 @@ use datafusion::physical_plan::projection::ProjectionExec; use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::sorts::sort::SortExec; use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; -use datafusion::physical_plan::source::DataSourceExec; use datafusion::physical_plan::union::{InterleaveExec, UnionExec}; use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec}; use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; diff --git a/datafusion/sql/Cargo.toml b/datafusion/sql/Cargo.toml index c4a404975d291..4435ee0f56cbc 100644 --- a/datafusion/sql/Cargo.toml +++ b/datafusion/sql/Cargo.toml @@ -28,12 +28,14 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_sql" -path = "src/lib.rs" [features] default = ["unicode_expressions", "unparser"] diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs index 1cf3dcb289a62..035749a78941d 100644 --- a/datafusion/sql/src/expr/function.rs +++ b/datafusion/sql/src/expr/function.rs @@ -23,7 +23,7 @@ use datafusion_common::{ DFSchema, Dependency, Result, }; use datafusion_expr::expr::{ScalarFunction, Unnest}; -use datafusion_expr::planner::PlannerResult; +use datafusion_expr::planner::{PlannerResult, RawAggregateExpr, RawWindowExpr}; use datafusion_expr::{ expr, qualified_wildcard, wildcard, Expr, ExprFunctionExt, ExprSchemable, WindowFrame, WindowFunctionDefinition, @@ -315,15 +315,38 @@ impl SqlToRel<'_, S> { }; if let Ok(fun) = self.find_window_func(&name) { - return Expr::WindowFunction(expr::WindowFunction::new( - fun, - self.function_args_to_expr(args, schema, planner_context)?, - )) - .partition_by(partition_by) - .order_by(order_by) - .window_frame(window_frame) - .null_treatment(null_treatment) - .build(); + let args = self.function_args_to_expr(args, schema, planner_context)?; + let mut window_expr = RawWindowExpr { + func_def: fun, + args, + partition_by, + order_by, + window_frame, + null_treatment, + }; + + for planner in self.context_provider.get_expr_planners().iter() { + match planner.plan_window(window_expr)? { + PlannerResult::Planned(expr) => return Ok(expr), + PlannerResult::Original(expr) => window_expr = expr, + } + } + + let RawWindowExpr { + func_def, + args, + partition_by, + order_by, + window_frame, + null_treatment, + } = window_expr; + + return Expr::WindowFunction(expr::WindowFunction::new(func_def, args)) + .partition_by(partition_by) + .order_by(order_by) + .window_frame(window_frame) + .null_treatment(null_treatment) + .build(); } } else { // User defined aggregate functions (UDAF) have precedence in case it has the same name as a scalar built-in function @@ -341,8 +364,33 @@ impl SqlToRel<'_, S> { .map(|e| self.sql_expr_to_logical_expr(*e, schema, planner_context)) .transpose()? .map(Box::new); + + let mut aggregate_expr = RawAggregateExpr { + func: fm, + args, + distinct, + filter, + order_by, + null_treatment, + }; + for planner in self.context_provider.get_expr_planners().iter() { + match planner.plan_aggregate(aggregate_expr)? { + PlannerResult::Planned(expr) => return Ok(expr), + PlannerResult::Original(expr) => aggregate_expr = expr, + } + } + + let RawAggregateExpr { + func, + args, + distinct, + filter, + order_by, + null_treatment, + } = aggregate_expr; + return Ok(Expr::AggregateFunction(expr::AggregateFunction::new_udf( - fm, + func, args, distinct, filter, diff --git a/datafusion/sql/src/expr/value.rs b/datafusion/sql/src/expr/value.rs index e81bfa0dc55f3..168348aee222e 100644 --- a/datafusion/sql/src/expr/value.rs +++ b/datafusion/sql/src/expr/value.rs @@ -304,7 +304,7 @@ fn try_decode_hex_literal(s: &str) -> Option> { for i in (start_idx..hex_bytes.len()).step_by(2) { let high = try_decode_hex_char(hex_bytes[i])?; let low = try_decode_hex_char(hex_bytes[i + 1])?; - decoded_bytes.push(high << 4 | low); + decoded_bytes.push((high << 4) | low); } Some(decoded_bytes) diff --git a/datafusion/sql/src/lib.rs b/datafusion/sql/src/lib.rs index 16a3d6d007cfe..d552efa8254ce 100644 --- a/datafusion/sql/src/lib.rs +++ b/datafusion/sql/src/lib.rs @@ -15,6 +15,11 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] // Make cheap clones clear: https://github.com/apache/datafusion/issues/11143 #![deny(clippy::clone_on_ref_ptr)] diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 74055d979145b..fbe6d6501c860 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -441,7 +441,7 @@ impl SqlToRel<'_, S> { plan }; - let constraints = Self::new_constraint_from_table_constraints( + let constraints = self.new_constraint_from_table_constraints( &all_constraints, plan.schema(), )?; @@ -465,7 +465,7 @@ impl SqlToRel<'_, S> { schema, }; let plan = LogicalPlan::EmptyRelation(plan); - let constraints = Self::new_constraint_from_table_constraints( + let constraints = self.new_constraint_from_table_constraints( &all_constraints, plan.schema(), )?; @@ -1434,7 +1434,7 @@ impl SqlToRel<'_, S> { let name = self.object_name_to_table_reference(name)?; let constraints = - Self::new_constraint_from_table_constraints(&all_constraints, &df_schema)?; + self.new_constraint_from_table_constraints(&all_constraints, &df_schema)?; Ok(LogicalPlan::Ddl(DdlStatement::CreateExternalTable( PlanCreateExternalTable { schema: df_schema, @@ -1454,8 +1454,34 @@ impl SqlToRel<'_, S> { ))) } + /// Get the indices of the constraint columns in the schema. + /// If any column is not found, return an error. + fn get_constraint_column_indices( + &self, + df_schema: &DFSchemaRef, + columns: &[Ident], + constraint_name: &str, + ) -> Result> { + let field_names = df_schema.field_names(); + columns + .iter() + .map(|ident| { + let column = self.ident_normalizer.normalize(ident.clone()); + field_names + .iter() + .position(|item| *item == column) + .ok_or_else(|| { + plan_datafusion_err!( + "Column for {constraint_name} not found in schema: {column}" + ) + }) + }) + .collect::>>() + } + /// Convert each [TableConstraint] to corresponding [Constraint] fn new_constraint_from_table_constraints( + &self, constraints: &[TableConstraint], df_schema: &DFSchemaRef, ) -> Result { @@ -1463,46 +1489,25 @@ impl SqlToRel<'_, S> { .iter() .map(|c: &TableConstraint| match c { TableConstraint::Unique { name, columns, .. } => { - let field_names = df_schema.field_names(); - // Get unique constraint indices in the schema: - let indices = columns - .iter() - .map(|u| { - let idx = field_names - .iter() - .position(|item| *item == u.value) - .ok_or_else(|| { - let name = name - .as_ref() - .map(|name| format!("with name '{name}' ")) - .unwrap_or("".to_string()); - DataFusionError::Execution( - format!("Column for unique constraint {}not found in schema: {}", name,u.value) - ) - })?; - Ok(idx) - }) - .collect::>>()?; + let constraint_name = match name { + Some(name) => &format!("unique constraint with name '{name}'"), + None => "unique constraint", + }; + // Get unique constraint indices in the schema + let indices = self.get_constraint_column_indices( + df_schema, + columns, + constraint_name, + )?; Ok(Constraint::Unique(indices)) } TableConstraint::PrimaryKey { columns, .. } => { - let field_names = df_schema.field_names(); - // Get primary key indices in the schema: - let indices = columns - .iter() - .map(|pk| { - let idx = field_names - .iter() - .position(|item| *item == pk.value) - .ok_or_else(|| { - DataFusionError::Execution(format!( - "Column for primary key not found in schema: {}", - pk.value - )) - })?; - Ok(idx) - }) - .collect::>>()?; + // Get primary key indices in the schema + let indices = self.get_constraint_column_indices( + df_schema, + columns, + "primary key", + )?; Ok(Constraint::PrimaryKey(indices)) } TableConstraint::ForeignKey { .. } => { diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 1df18302687ec..9c0d6316adb2b 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -1461,14 +1461,14 @@ fn select_simple_aggregate_with_groupby_and_column_is_in_aggregate_and_groupby() fn select_simple_aggregate_with_groupby_can_use_positions() { quick_test( "SELECT state, age AS b, count(1) FROM person GROUP BY 1, 2", - "Projection: person.state, person.age AS b, count(Int64(1))\ - \n Aggregate: groupBy=[[person.state, person.age]], aggr=[[count(Int64(1))]]\ + "Projection: person.state, person.age AS b, count(*)\ + \n Aggregate: groupBy=[[person.state, person.age]], aggr=[[count(*)]]\ \n TableScan: person", ); quick_test( "SELECT state, age AS b, count(1) FROM person GROUP BY 2, 1", - "Projection: person.state, person.age AS b, count(Int64(1))\ - \n Aggregate: groupBy=[[person.age, person.state]], aggr=[[count(Int64(1))]]\ + "Projection: person.state, person.age AS b, count(*)\ + \n Aggregate: groupBy=[[person.age, person.state]], aggr=[[count(*)]]\ \n TableScan: person", ); } @@ -1630,8 +1630,8 @@ fn test_wildcard() { #[test] fn select_count_one() { let sql = "SELECT count(1) FROM person"; - let expected = "Projection: count(Int64(1))\ - \n Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]]\ + let expected = "Projection: count(*)\ + \n Aggregate: groupBy=[[]], aggr=[[count(*)]]\ \n TableScan: person"; quick_test(sql, expected); } diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml index f1d37c7202d6d..bc1b283eda65f 100644 --- a/datafusion/sqllogictest/Cargo.toml +++ b/datafusion/sqllogictest/Cargo.toml @@ -27,12 +27,14 @@ repository = { workspace = true } rust-version = { workspace = true } version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true [lib] name = "datafusion_sqllogictest" -path = "src/lib.rs" [dependencies] arrow = { workspace = true } @@ -51,7 +53,7 @@ object_store = { workspace = true } postgres-protocol = { version = "0.6.7", optional = true } postgres-types = { version = "0.2.8", features = ["derive", "with-chrono-0_4"], optional = true } rust_decimal = { version = "1.36.0", features = ["tokio-pg"] } -sqllogictest = "0.27.1" +sqllogictest = "0.27.2" sqlparser = { workspace = true } tempfile = { workspace = true } testcontainers = { version = "0.23", features = ["default"], optional = true } diff --git a/datafusion/sqllogictest/src/lib.rs b/datafusion/sqllogictest/src/lib.rs index 0ea55782d34e6..ee20e70d14f41 100644 --- a/datafusion/sqllogictest/src/lib.rs +++ b/datafusion/sqllogictest/src/lib.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + //! DataFusion sqllogictest driver mod engines; diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 7caa81d64e5b7..f175973f92a15 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -6276,6 +6276,26 @@ physical_plan 05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c5], file_type=csv, has_header=true +# test count(null) case (null with type) + +statement count 0 +create table t(a int, b int) as values (1, 3), (2, 4), (3, 5); + +query I +select count(null::bigint) from t; +---- +0 + +query TT +explain select count(null::bigint) from t; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[count(Int64(NULL)) AS count(NULL)]] +02)--TableScan: t projection=[] +physical_plan +01)AggregateExec: mode=Single, gby=[], aggr=[count(NULL)] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + ####### # Group median test ####### diff --git a/datafusion/sqllogictest/test_files/avro.slt b/datafusion/sqllogictest/test_files/avro.slt index 80bf0bc2dd5a0..20179e0c5bdca 100644 --- a/datafusion/sqllogictest/test_files/avro.slt +++ b/datafusion/sqllogictest/test_files/avro.slt @@ -243,7 +243,7 @@ query TT EXPLAIN SELECT count(*) from alltypes_plain ---- logical_plan -01)Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +01)Aggregate: groupBy=[[]], aggr=[[count(*)]] 02)--TableScan: alltypes_plain projection=[] physical_plan 01)AggregateExec: mode=Final, gby=[], aggr=[count(*)] diff --git a/datafusion/sqllogictest/test_files/coalesce.slt b/datafusion/sqllogictest/test_files/coalesce.slt index 5f2d2f0d1da9d..e7cf31dc690b7 100644 --- a/datafusion/sqllogictest/test_files/coalesce.slt +++ b/datafusion/sqllogictest/test_files/coalesce.slt @@ -442,4 +442,4 @@ drop table test query T select coalesce(arrow_cast('', 'Utf8View'), arrow_cast('', 'Dictionary(UInt32, Utf8)')); ---- -(empty) \ No newline at end of file +(empty) diff --git a/datafusion/sqllogictest/test_files/copy.slt b/datafusion/sqllogictest/test_files/copy.slt index 7dd85b3ae2d8b..f39ff56ce4499 100644 --- a/datafusion/sqllogictest/test_files/copy.slt +++ b/datafusion/sqllogictest/test_files/copy.slt @@ -631,4 +631,3 @@ COPY source_table to '/tmp/table.parquet' (row_group_size 55 + 102); # Copy using execution.keep_partition_by_columns with an invalid value query error DataFusion error: Invalid or Unsupported Configuration: provided value for 'execution.keep_partition_by_columns' was not recognized: "invalid_value" COPY source_table to '/tmp/table.parquet' OPTIONS (execution.keep_partition_by_columns invalid_value); - diff --git a/datafusion/sqllogictest/test_files/count_star_rule.slt b/datafusion/sqllogictest/test_files/count_star_rule.slt index d660257b609d9..0efd9e99889f2 100644 --- a/datafusion/sqllogictest/test_files/count_star_rule.slt +++ b/datafusion/sqllogictest/test_files/count_star_rule.slt @@ -31,44 +31,44 @@ query TT EXPLAIN SELECT COUNT() FROM (SELECT 1 AS a, 2 AS b) AS t; ---- logical_plan -01)Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count()]] +01)Aggregate: groupBy=[[]], aggr=[[count(*)]] 02)--SubqueryAlias: t 03)----EmptyRelation physical_plan -01)ProjectionExec: expr=[1 as count()] +01)ProjectionExec: expr=[1 as count(*)] 02)--PlaceholderRowExec query TT EXPLAIN SELECT t1.a, COUNT() FROM t1 GROUP BY t1.a; ---- logical_plan -01)Aggregate: groupBy=[[t1.a]], aggr=[[count(Int64(1)) AS count()]] +01)Aggregate: groupBy=[[t1.a]], aggr=[[count(*)]] 02)--TableScan: t1 projection=[a] physical_plan -01)AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count()] +01)AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(*)] 02)--CoalesceBatchesExec: target_batch_size=8192 03)----RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -05)--------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count()] +05)--------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(*)] 06)----------DataSourceExec: partitions=1, partition_sizes=[1] query TT EXPLAIN SELECT t1.a, COUNT() AS cnt FROM t1 GROUP BY t1.a HAVING COUNT() > 0; ---- logical_plan -01)Projection: t1.a, count() AS cnt -02)--Filter: count() > Int64(0) -03)----Aggregate: groupBy=[[t1.a]], aggr=[[count(Int64(1)) AS count()]] +01)Projection: t1.a, count(*) AS cnt +02)--Filter: count(*) > Int64(0) +03)----Aggregate: groupBy=[[t1.a]], aggr=[[count(*)]] 04)------TableScan: t1 projection=[a] physical_plan -01)ProjectionExec: expr=[a@0 as a, count()@1 as cnt] +01)ProjectionExec: expr=[a@0 as a, count(*)@1 as cnt] 02)--CoalesceBatchesExec: target_batch_size=8192 -03)----FilterExec: count()@1 > 0 -04)------AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count()] +03)----FilterExec: count(*)@1 > 0 +04)------AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[count(*)] 05)--------CoalesceBatchesExec: target_batch_size=8192 06)----------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 07)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -08)--------------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count()] +08)--------------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(*)] 09)----------------DataSourceExec: partitions=1, partition_sizes=[1] query II @@ -80,12 +80,12 @@ query TT EXPLAIN SELECT a, COUNT() OVER (PARTITION BY a) AS count_a FROM t1; ---- logical_plan -01)Projection: t1.a, count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS count_a -02)--WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] +01)Projection: t1.a, count(*) PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS count_a +02)--WindowAggr: windowExpr=[[count(*) PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] 03)----TableScan: t1 projection=[a] physical_plan -01)ProjectionExec: expr=[a@0 as a, count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as count_a] -02)--WindowAggExec: wdw=[count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "count() PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +01)ProjectionExec: expr=[a@0 as a, count(*) PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as count_a] +02)--WindowAggExec: wdw=[count(*) PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "count(*) PARTITION BY [t1.a] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] 03)----SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false] 04)------DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/datafusion/sqllogictest/test_files/ddl.slt b/datafusion/sqllogictest/test_files/ddl.slt index aefc2672b5396..6f75a7d7f8fdb 100644 --- a/datafusion/sqllogictest/test_files/ddl.slt +++ b/datafusion/sqllogictest/test_files/ddl.slt @@ -827,4 +827,3 @@ drop table table_with_pk; statement ok set datafusion.catalog.information_schema = false; - diff --git a/datafusion/sqllogictest/test_files/errors.slt b/datafusion/sqllogictest/test_files/errors.slt index a35a4d6f28dc0..dc7a53adf889d 100644 --- a/datafusion/sqllogictest/test_files/errors.slt +++ b/datafusion/sqllogictest/test_files/errors.slt @@ -184,4 +184,4 @@ query error DataFusion error: Schema error: No field named ammp\. Did you mean ' select ammp from a; statement ok -drop table a; \ No newline at end of file +drop table a; diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 037565ce05f9b..0d5eab6cf56dc 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -178,7 +178,6 @@ logical_plan after inline_table_scan SAME TEXT AS ABOVE logical_plan after expand_wildcard_rule SAME TEXT AS ABOVE logical_plan after resolve_grouping_function SAME TEXT AS ABOVE logical_plan after type_coercion SAME TEXT AS ABOVE -logical_plan after count_wildcard_rule SAME TEXT AS ABOVE analyzed_logical_plan SAME TEXT AS ABOVE logical_plan after eliminate_nested_union SAME TEXT AS ABOVE logical_plan after simplify_expressions SAME TEXT AS ABOVE @@ -427,7 +426,7 @@ logical_plan 02)--TableScan: t1 projection=[a] 03)--SubqueryAlias: __correlated_sq_1 04)----Projection: -05)------Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +05)------Aggregate: groupBy=[[]], aggr=[[count(*)]] 06)--------TableScan: t2 projection=[] physical_plan 01)NestedLoopJoinExec: join_type=LeftSemi diff --git a/datafusion/sqllogictest/test_files/ident_normalization.slt b/datafusion/sqllogictest/test_files/ident_normalization.slt index 996093c3ad9ca..ac2f460ebc430 100644 --- a/datafusion/sqllogictest/test_files/ident_normalization.slt +++ b/datafusion/sqllogictest/test_files/ident_normalization.slt @@ -132,3 +132,44 @@ query T SELECT CONCAT('Hello', 'World') ---- HelloWorld + +# Restore default setting +statement ok +set datafusion.sql_parser.enable_ident_normalization = true; + +########## +## Constraint Column Name Normalization +########## + +# Test issue https://github.com/apache/datafusion/issues/14340 +statement ok +create table test_pk_constraint(COLUMN_NAME TEXT NOT NULL, constraint COLUMN_NAME_PK primary key (COLUMN_NAME)) + +statement ok +create table test_unique_constraint(cOlUmn_name TEXT NOT NULL, unique(COLUMN_NAME)) + +statement ok +drop table test_pk_constraint; + +statement ok +drop table test_unique_constraint; + +# Test with normalization disabled +statement ok +set datafusion.sql_parser.enable_ident_normalization = false; + +statement error Error during planning: Column for primary key not found in schema: COLUMN_NAME +create table test_pk_constraint(column_name TEXT NOT NULL, constraint COLUMN_NAME_PK primary key (COLUMN_NAME)) + +statement ok +create table test_pk_constraint(COLUMN_NAME TEXT NOT NULL, constraint COLUMN_NAME_PK primary key (COLUMN_NAME)) + +statement error Error during planning: Column for unique constraint not found in schema: COLUMN_NAME +create table test_unique_constraint(column_name TEXT NOT NULL, unique(COLUMN_NAME)) + +statement ok +drop table test_pk_constraint; + +# Restore default setting +statement ok +set datafusion.sql_parser.enable_ident_normalization = true; diff --git a/datafusion/sqllogictest/test_files/insert.slt b/datafusion/sqllogictest/test_files/insert.slt index ee76ee1c55115..32428fdef7654 100644 --- a/datafusion/sqllogictest/test_files/insert.slt +++ b/datafusion/sqllogictest/test_files/insert.slt @@ -61,7 +61,7 @@ logical_plan 02)--Projection: sum(aggregate_test_100.c4) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS field1, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS field2 03)----Sort: aggregate_test_100.c1 ASC NULLS LAST 04)------Projection: sum(aggregate_test_100.c4) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, aggregate_test_100.c1 -05)--------WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(Int64(1)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] +05)--------WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] 06)----------TableScan: aggregate_test_100 projection=[c1, c4, c9] physical_plan 01)DataSinkExec: sink=MemoryTable (partitions=1) @@ -122,7 +122,7 @@ FROM aggregate_test_100 logical_plan 01)Dml: op=[Insert Into] table=[table_without_values] 02)--Projection: sum(aggregate_test_100.c4) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS field1, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS field2 -03)----WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(Int64(1)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] +03)----WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] 04)------TableScan: aggregate_test_100 projection=[c1, c4, c9] physical_plan 01)DataSinkExec: sink=MemoryTable (partitions=1) @@ -172,7 +172,7 @@ logical_plan 02)--Projection: a1 AS a1, a2 AS a2 03)----Sort: aggregate_test_100.c1 ASC NULLS LAST 04)------Projection: sum(aggregate_test_100.c4) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS a1, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS a2, aggregate_test_100.c1 -05)--------WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(Int64(1)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] +05)--------WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] 06)----------TableScan: aggregate_test_100 projection=[c1, c4, c9] physical_plan 01)DataSinkExec: sink=MemoryTable (partitions=8) diff --git a/datafusion/sqllogictest/test_files/insert_to_external.slt b/datafusion/sqllogictest/test_files/insert_to_external.slt index ee1d67c5e26d2..752e8ce0e4ffe 100644 --- a/datafusion/sqllogictest/test_files/insert_to_external.slt +++ b/datafusion/sqllogictest/test_files/insert_to_external.slt @@ -352,7 +352,7 @@ logical_plan 02)--Projection: sum(aggregate_test_100.c4) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS field1, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS field2 03)----Sort: aggregate_test_100.c1 ASC NULLS LAST 04)------Projection: sum(aggregate_test_100.c4) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, aggregate_test_100.c1 -05)--------WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(Int64(1)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] +05)--------WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] 06)----------TableScan: aggregate_test_100 projection=[c1, c4, c9] physical_plan 01)DataSinkExec: sink=ParquetSink(file_groups=[]) @@ -414,7 +414,7 @@ FROM aggregate_test_100 logical_plan 01)Dml: op=[Insert Into] table=[table_without_values] 02)--Projection: sum(aggregate_test_100.c4) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS field1, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS field2 -03)----WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(Int64(1)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] +03)----WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] 04)------TableScan: aggregate_test_100 projection=[c1, c4, c9] physical_plan 01)DataSinkExec: sink=ParquetSink(file_groups=[]) diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index 5d311bc43293f..5b5368f6b0f4c 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -1396,7 +1396,7 @@ group by t1_id ---- logical_plan 01)Projection: count(*) -02)--Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(Int64(1)) AS count(*)]] +02)--Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(*)]] 03)----Projection: join_t1.t1_id 04)------Inner Join: join_t1.t1_id = join_t2.t2_id 05)--------TableScan: join_t1 projection=[t1_id] @@ -4442,7 +4442,7 @@ FROM my_catalog.my_schema.table_with_many_types AS l JOIN my_catalog.my_schema.table_with_many_types AS r ON l.binary_col = r.binary_col ---- logical_plan -01)Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +01)Aggregate: groupBy=[[]], aggr=[[count(*)]] 02)--Projection: 03)----Inner Join: l.binary_col = r.binary_col 04)------SubqueryAlias: l diff --git a/datafusion/sqllogictest/test_files/json.slt b/datafusion/sqllogictest/test_files/json.slt index dd310f7f2bf6b..466bba5566971 100644 --- a/datafusion/sqllogictest/test_files/json.slt +++ b/datafusion/sqllogictest/test_files/json.slt @@ -54,7 +54,7 @@ query TT EXPLAIN SELECT count(*) from json_test ---- logical_plan -01)Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +01)Aggregate: groupBy=[[]], aggr=[[count(*)]] 02)--TableScan: json_test projection=[] physical_plan 01)AggregateExec: mode=Final, gby=[], aggr=[count(*)] diff --git a/datafusion/sqllogictest/test_files/limit.slt b/datafusion/sqllogictest/test_files/limit.slt index 4e74b27b875f7..b4487be850ac3 100644 --- a/datafusion/sqllogictest/test_files/limit.slt +++ b/datafusion/sqllogictest/test_files/limit.slt @@ -307,7 +307,7 @@ query TT EXPLAIN SELECT COUNT(*) FROM (SELECT a FROM t1 LIMIT 3 OFFSET 11); ---- logical_plan -01)Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +01)Aggregate: groupBy=[[]], aggr=[[count(*)]] 02)--Limit: skip=11, fetch=3 03)----TableScan: t1 projection=[], fetch=14 physical_plan @@ -325,7 +325,7 @@ query TT EXPLAIN SELECT COUNT(*) FROM (SELECT a FROM t1 LIMIT 3 OFFSET 8); ---- logical_plan -01)Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +01)Aggregate: groupBy=[[]], aggr=[[count(*)]] 02)--Limit: skip=8, fetch=3 03)----TableScan: t1 projection=[], fetch=11 physical_plan @@ -343,7 +343,7 @@ query TT EXPLAIN SELECT COUNT(*) FROM (SELECT a FROM t1 OFFSET 8); ---- logical_plan -01)Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +01)Aggregate: groupBy=[[]], aggr=[[count(*)]] 02)--Limit: skip=8, fetch=None 03)----TableScan: t1 projection=[] physical_plan @@ -360,7 +360,7 @@ query TT EXPLAIN SELECT COUNT(*) FROM (SELECT a FROM t1 WHERE a > 3 LIMIT 3 OFFSET 6); ---- logical_plan -01)Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +01)Aggregate: groupBy=[[]], aggr=[[count(*)]] 02)--Projection: 03)----Limit: skip=6, fetch=3 04)------Filter: t1.a > Int32(3) diff --git a/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt b/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt index de6a153f58d98..8c87af75ed16a 100644 --- a/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt +++ b/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt @@ -48,8 +48,8 @@ FROM test_table t GROUP BY 1, 2, 3, 4 ---- logical_plan -01)Projection: t.c1, Int64(99999), t.c5 + t.c8, Utf8("test"), count(Int64(1)) -02)--Aggregate: groupBy=[[t.c1, t.c5 + t.c8]], aggr=[[count(Int64(1))]] +01)Projection: t.c1, Int64(99999), t.c5 + t.c8, Utf8("test"), count(*) +02)--Aggregate: groupBy=[[t.c1, t.c5 + t.c8]], aggr=[[count(*)]] 03)----SubqueryAlias: t 04)------TableScan: test_table projection=[c1, c5, c8] @@ -60,8 +60,8 @@ FROM test_table t group by 1, 2, 3 ---- logical_plan -01)Projection: Int64(123), Int64(456), Int64(789), count(Int64(1)), avg(t.c12) -02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1)), avg(t.c12)]] +01)Projection: Int64(123), Int64(456), Int64(789), count(*), avg(t.c12) +02)--Aggregate: groupBy=[[]], aggr=[[count(*), avg(t.c12)]] 03)----SubqueryAlias: t 04)------TableScan: test_table projection=[c12] @@ -72,8 +72,8 @@ FROM test_table t GROUP BY 1, 2 ---- logical_plan -01)Projection: Date32("2023-05-04") AS dt, Boolean(true) AS today_filter, count(Int64(1)) -02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +01)Projection: Date32("2023-05-04") AS dt, Boolean(true) AS today_filter, count(*) +02)--Aggregate: groupBy=[[]], aggr=[[count(*)]] 03)----SubqueryAlias: t 04)------TableScan: test_table projection=[] @@ -90,8 +90,8 @@ FROM test_table t GROUP BY 1 ---- logical_plan -01)Projection: Boolean(true) AS NOT date_part(Utf8("MONTH"),now()) BETWEEN Int64(50) AND Int64(60), count(Int64(1)) -02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +01)Projection: Boolean(true) AS NOT date_part(Utf8("MONTH"),now()) BETWEEN Int64(50) AND Int64(60), count(*) +02)--Aggregate: groupBy=[[]], aggr=[[count(*)]] 03)----SubqueryAlias: t 04)------TableScan: test_table projection=[] diff --git a/datafusion/sqllogictest/test_files/select.slt b/datafusion/sqllogictest/test_files/select.slt index e12bdca37e6fe..dcd373546d793 100644 --- a/datafusion/sqllogictest/test_files/select.slt +++ b/datafusion/sqllogictest/test_files/select.slt @@ -1541,7 +1541,7 @@ LIMIT 4) GROUP BY c2; ---- logical_plan -01)Aggregate: groupBy=[[aggregate_test_100.c2]], aggr=[[count(Int64(1)) AS count(*)]] +01)Aggregate: groupBy=[[aggregate_test_100.c2]], aggr=[[count(*)]] 02)--Projection: aggregate_test_100.c2 03)----Sort: aggregate_test_100.c1 ASC NULLS LAST, aggregate_test_100.c2 ASC NULLS LAST, fetch=4 04)------Projection: aggregate_test_100.c2, aggregate_test_100.c1 diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 264392fc1017c..c847f433f7fcc 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -555,7 +555,7 @@ logical_plan 03)----Subquery: 04)------Projection: count(*) 05)--------Filter: sum(outer_ref(t1.t1_int) + t2.t2_id) > Int64(0) -06)----------Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*), sum(CAST(outer_ref(t1.t1_int) + t2.t2_id AS Int64))]] +06)----------Aggregate: groupBy=[[]], aggr=[[count(*), sum(CAST(outer_ref(t1.t1_int) + t2.t2_id AS Int64))]] 07)------------Filter: outer_ref(t1.t1_name) = t2.t2_name 08)--------------TableScan: t2 09)----TableScan: t1 projection=[t1_id, t1_name, t1_int] @@ -738,7 +738,7 @@ explain select (select count(*) from t1) as b logical_plan 01)Projection: __scalar_sq_1.count(*) AS b 02)--SubqueryAlias: __scalar_sq_1 -03)----Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +03)----Aggregate: groupBy=[[]], aggr=[[count(*)]] 04)------TableScan: t1 projection=[] #simple_uncorrelated_scalar_subquery2 @@ -746,13 +746,13 @@ query TT explain select (select count(*) from t1) as b, (select count(1) from t2) ---- logical_plan -01)Projection: __scalar_sq_1.count(*) AS b, __scalar_sq_2.count(Int64(1)) AS count(Int64(1)) +01)Projection: __scalar_sq_1.count(*) AS b, __scalar_sq_2.count(*) AS count(*) 02)--Left Join: 03)----SubqueryAlias: __scalar_sq_1 -04)------Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +04)------Aggregate: groupBy=[[]], aggr=[[count(*)]] 05)--------TableScan: t1 projection=[] 06)----SubqueryAlias: __scalar_sq_2 -07)------Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +07)------Aggregate: groupBy=[[]], aggr=[[count(*)]] 08)--------TableScan: t2 projection=[] statement ok @@ -762,20 +762,20 @@ query TT explain select (select count(*) from t1) as b, (select count(1) from t2) ---- logical_plan -01)Projection: __scalar_sq_1.count(*) AS b, __scalar_sq_2.count(Int64(1)) AS count(Int64(1)) +01)Projection: __scalar_sq_1.count(*) AS b, __scalar_sq_2.count(*) AS count(*) 02)--Left Join: 03)----SubqueryAlias: __scalar_sq_1 -04)------Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +04)------Aggregate: groupBy=[[]], aggr=[[count(*)]] 05)--------TableScan: t1 projection=[] 06)----SubqueryAlias: __scalar_sq_2 -07)------Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +07)------Aggregate: groupBy=[[]], aggr=[[count(*)]] 08)--------TableScan: t2 projection=[] physical_plan -01)ProjectionExec: expr=[count(*)@0 as b, count(Int64(1))@1 as count(Int64(1))] +01)ProjectionExec: expr=[count(*)@0 as b, count(*)@1 as count(*)] 02)--NestedLoopJoinExec: join_type=Left 03)----ProjectionExec: expr=[4 as count(*)] 04)------PlaceholderRowExec -05)----ProjectionExec: expr=[4 as count(Int64(1))] +05)----ProjectionExec: expr=[4 as count(*)] 06)------PlaceholderRowExec statement ok @@ -796,7 +796,7 @@ logical_plan 03)----TableScan: t1 projection=[t1_id, t1_int] 04)----SubqueryAlias: __scalar_sq_1 05)------Projection: count(*), t2.t2_int, Boolean(true) AS __always_true -06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 07)----------TableScan: t2 projection=[t2_int] query II rowsort @@ -818,7 +818,7 @@ logical_plan 03)----TableScan: t1 projection=[t1_id, t1_int] 04)----SubqueryAlias: __scalar_sq_1 05)------Projection: count(*), t2.t2_int, Boolean(true) AS __always_true -06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 07)----------TableScan: t2 projection=[t2_int] query II rowsort @@ -839,7 +839,7 @@ logical_plan 03)----TableScan: t1 projection=[t1_id, t1_int] 04)----SubqueryAlias: __scalar_sq_1 05)------Projection: count(*) AS _cnt, t2.t2_int, Boolean(true) AS __always_true -06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 07)----------TableScan: t2 projection=[t2_int] query II rowsort @@ -860,7 +860,7 @@ logical_plan 03)----TableScan: t1 projection=[t1_id, t1_int] 04)----SubqueryAlias: __scalar_sq_1 05)------Projection: count(*) + Int64(2) AS _cnt, t2.t2_int, Boolean(true) AS __always_true -06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 07)----------TableScan: t2 projection=[t2_int] query II rowsort @@ -883,7 +883,7 @@ logical_plan 05)--------TableScan: t1 projection=[t1_id, t1_int] 06)--------SubqueryAlias: __scalar_sq_1 07)----------Projection: count(*), t2.t2_id, Boolean(true) AS __always_true -08)------------Aggregate: groupBy=[[t2.t2_id]], aggr=[[count(Int64(1)) AS count(*)]] +08)------------Aggregate: groupBy=[[t2.t2_id]], aggr=[[count(*)]] 09)--------------TableScan: t2 projection=[t2_id] query I rowsort @@ -905,7 +905,7 @@ logical_plan 04)----SubqueryAlias: __scalar_sq_1 05)------Projection: count(*) + Int64(2) AS cnt_plus_2, t2.t2_int 06)--------Filter: count(*) > Int64(1) -07)----------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +07)----------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 08)------------TableScan: t2 projection=[t2_int] query II rowsort @@ -927,7 +927,7 @@ logical_plan 03)----TableScan: t1 projection=[t1_id, t1_int] 04)----SubqueryAlias: __scalar_sq_1 05)------Projection: count(*) + Int64(2) AS cnt_plus_2, t2.t2_int, count(*), Boolean(true) AS __always_true -06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +06)--------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 07)----------TableScan: t2 projection=[t2_int] query II rowsort @@ -951,7 +951,7 @@ logical_plan 06)----------TableScan: t1 projection=[t1_int] 07)--------SubqueryAlias: __scalar_sq_1 08)----------Projection: count(*), t2.t2_int, Boolean(true) AS __always_true -09)------------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +09)------------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 10)--------------TableScan: t2 projection=[t2_int] query I rowsort @@ -972,7 +972,7 @@ logical_plan 05)--------TableScan: t1 projection=[t1_int] 06)--------SubqueryAlias: __scalar_sq_1 07)----------Projection: count(*) AS cnt, t2.t2_int, Boolean(true) AS __always_true -08)------------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +08)------------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 09)--------------TableScan: t2 projection=[t2_int] @@ -1002,7 +1002,7 @@ logical_plan 05)--------TableScan: t1 projection=[t1_int] 06)--------SubqueryAlias: __scalar_sq_1 07)----------Projection: count(*) + Int64(1) + Int64(1) AS cnt_plus_two, t2.t2_int, count(*), Boolean(true) AS __always_true -08)------------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +08)------------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 09)--------------TableScan: t2 projection=[t2_int] query I rowsort @@ -1031,7 +1031,7 @@ logical_plan 05)--------TableScan: t1 projection=[t1_int] 06)--------SubqueryAlias: __scalar_sq_1 07)----------Projection: CASE WHEN count(*) = Int64(1) THEN Int64(NULL) ELSE count(*) END AS cnt, t2.t2_int, Boolean(true) AS __always_true -08)------------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(Int64(1)) AS count(*)]] +08)------------Aggregate: groupBy=[[t2.t2_int]], aggr=[[count(*)]] 09)--------------TableScan: t2 projection=[t2_int] diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q1.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q1.slt.part index 2616b7b75b304..6a41ecb51bf43 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q1.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q1.slt.part @@ -42,7 +42,7 @@ explain select logical_plan 01)Sort: lineitem.l_returnflag ASC NULLS LAST, lineitem.l_linestatus ASC NULLS LAST 02)--Projection: lineitem.l_returnflag, lineitem.l_linestatus, sum(lineitem.l_quantity) AS sum_qty, sum(lineitem.l_extendedprice) AS sum_base_price, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount) AS sum_disc_price, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax) AS sum_charge, avg(lineitem.l_quantity) AS avg_qty, avg(lineitem.l_extendedprice) AS avg_price, avg(lineitem.l_discount) AS avg_disc, count(*) AS count_order -03)----Aggregate: groupBy=[[lineitem.l_returnflag, lineitem.l_linestatus]], aggr=[[sum(lineitem.l_quantity), sum(lineitem.l_extendedprice), sum(__common_expr_1) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount), sum(__common_expr_1 * (Decimal128(Some(1),20,0) + lineitem.l_tax)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax), avg(lineitem.l_quantity), avg(lineitem.l_extendedprice), avg(lineitem.l_discount), count(Int64(1)) AS count(*)]] +03)----Aggregate: groupBy=[[lineitem.l_returnflag, lineitem.l_linestatus]], aggr=[[sum(lineitem.l_quantity), sum(lineitem.l_extendedprice), sum(__common_expr_1) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount), sum(__common_expr_1 * (Decimal128(Some(1),20,0) + lineitem.l_tax)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax), avg(lineitem.l_quantity), avg(lineitem.l_extendedprice), avg(lineitem.l_discount), count(*)]] 04)------Projection: lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount) AS __common_expr_1, lineitem.l_quantity, lineitem.l_extendedprice, lineitem.l_discount, lineitem.l_tax, lineitem.l_returnflag, lineitem.l_linestatus 05)--------Filter: lineitem.l_shipdate <= Date32("1998-09-02") 06)----------TableScan: lineitem projection=[l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate], partial_filters=[lineitem.l_shipdate <= Date32("1998-09-02")] diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q13.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q13.slt.part index eb41445c3c138..68532733c6617 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q13.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q13.slt.part @@ -42,7 +42,7 @@ limit 10; logical_plan 01)Sort: custdist DESC NULLS FIRST, c_orders.c_count DESC NULLS FIRST, fetch=10 02)--Projection: c_orders.c_count, count(*) AS custdist -03)----Aggregate: groupBy=[[c_orders.c_count]], aggr=[[count(Int64(1)) AS count(*)]] +03)----Aggregate: groupBy=[[c_orders.c_count]], aggr=[[count(*)]] 04)------SubqueryAlias: c_orders 05)--------Projection: count(orders.o_orderkey) AS c_count 06)----------Aggregate: groupBy=[[customer.c_custkey]], aggr=[[count(orders.o_orderkey)]] diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q21.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q21.slt.part index 9e39732689da9..eb10f4c8d1958 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q21.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q21.slt.part @@ -60,7 +60,7 @@ order by logical_plan 01)Sort: numwait DESC NULLS FIRST, supplier.s_name ASC NULLS LAST 02)--Projection: supplier.s_name, count(*) AS numwait -03)----Aggregate: groupBy=[[supplier.s_name]], aggr=[[count(Int64(1)) AS count(*)]] +03)----Aggregate: groupBy=[[supplier.s_name]], aggr=[[count(*)]] 04)------Projection: supplier.s_name 05)--------LeftAnti Join: l1.l_orderkey = __correlated_sq_2.l_orderkey Filter: __correlated_sq_2.l_suppkey != l1.l_suppkey 06)----------LeftSemi Join: l1.l_orderkey = __correlated_sq_1.l_orderkey Filter: __correlated_sq_1.l_suppkey != l1.l_suppkey diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part index 9ad99361256c0..af8b7948c1cf9 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part @@ -58,7 +58,7 @@ order by logical_plan 01)Sort: custsale.cntrycode ASC NULLS LAST 02)--Projection: custsale.cntrycode, count(*) AS numcust, sum(custsale.c_acctbal) AS totacctbal -03)----Aggregate: groupBy=[[custsale.cntrycode]], aggr=[[count(Int64(1)) AS count(*), sum(custsale.c_acctbal)]] +03)----Aggregate: groupBy=[[custsale.cntrycode]], aggr=[[count(*), sum(custsale.c_acctbal)]] 04)------SubqueryAlias: custsale 05)--------Projection: substr(customer.c_phone, Int64(1), Int64(2)) AS cntrycode, customer.c_acctbal 06)----------Inner Join: Filter: CAST(customer.c_acctbal AS Decimal128(19, 6)) > __scalar_sq_2.avg(customer.c_acctbal) diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q4.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q4.slt.part index fb93850ab095c..766b21c22f24d 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q4.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q4.slt.part @@ -42,7 +42,7 @@ order by logical_plan 01)Sort: orders.o_orderpriority ASC NULLS LAST 02)--Projection: orders.o_orderpriority, count(*) AS order_count -03)----Aggregate: groupBy=[[orders.o_orderpriority]], aggr=[[count(Int64(1)) AS count(*)]] +03)----Aggregate: groupBy=[[orders.o_orderpriority]], aggr=[[count(*)]] 04)------Projection: orders.o_orderpriority 05)--------LeftSemi Join: orders.o_orderkey = __correlated_sq_1.l_orderkey 06)----------Projection: orders.o_orderkey, orders.o_orderpriority diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index dfac9c031074a..57207f00f7abd 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -449,7 +449,7 @@ SELECT count(*) FROM ( ---- logical_plan 01)Projection: count(*) -02)--Aggregate: groupBy=[[t1.name]], aggr=[[count(Int64(1)) AS count(*)]] +02)--Aggregate: groupBy=[[t1.name]], aggr=[[count(*)]] 03)----Union 04)------Aggregate: groupBy=[[t1.name]], aggr=[[]] 05)--------TableScan: t1 projection=[name] @@ -493,7 +493,7 @@ logical_plan 02)--Union 03)----Projection: count(*) AS cnt 04)------Limit: skip=0, fetch=3 -05)--------Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +05)--------Aggregate: groupBy=[[]], aggr=[[count(*)]] 06)----------SubqueryAlias: a 07)------------Projection: 08)--------------Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[]] @@ -651,7 +651,7 @@ select x, y from (select 1 as x , max(10) as y) b logical_plan 01)Union 02)--Projection: count(*) AS count, a.n -03)----Aggregate: groupBy=[[a.n]], aggr=[[count(Int64(1)) AS count(*)]] +03)----Aggregate: groupBy=[[a.n]], aggr=[[count(*)]] 04)------SubqueryAlias: a 05)--------Projection: Int64(5) AS n 06)----------EmptyRelation diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index ca4713e7d5160..6c00af879e762 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -1305,7 +1305,7 @@ EXPLAIN SELECT ---- logical_plan 01)Projection: sum(aggregate_test_100.c4) PARTITION BY [aggregate_test_100.c1, aggregate_test_100.c2] ORDER BY [aggregate_test_100.c2 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING, count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c2 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING -02)--WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c2 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING AS count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c2 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] +02)--WindowAggr: windowExpr=[[count(*) PARTITION BY [aggregate_test_100.c1] ORDER BY [aggregate_test_100.c2 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] 03)----Projection: aggregate_test_100.c1, aggregate_test_100.c2, sum(aggregate_test_100.c4) PARTITION BY [aggregate_test_100.c1, aggregate_test_100.c2] ORDER BY [aggregate_test_100.c2 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING 04)------WindowAggr: windowExpr=[[sum(CAST(aggregate_test_100.c4 AS Int64)) PARTITION BY [aggregate_test_100.c1, aggregate_test_100.c2] ORDER BY [aggregate_test_100.c2 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING]] 05)--------TableScan: aggregate_test_100 projection=[c1, c2, c4] @@ -1765,7 +1765,7 @@ EXPLAIN SELECT count(*) as global_count FROM ---- logical_plan 01)Projection: count(*) AS global_count -02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]] +02)--Aggregate: groupBy=[[]], aggr=[[count(*)]] 03)----SubqueryAlias: a 04)------Projection: 05)--------Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[]] @@ -2571,10 +2571,10 @@ logical_plan 01)Projection: sum1, sum2, sum3, min1, min2, min3, max1, max2, max3, cnt1, cnt2, sumr1, sumr2, sumr3, minr1, minr2, minr3, maxr1, maxr2, maxr3, cntr1, cntr2, sum4, cnt3 02)--Sort: annotated_data_finite.inc_col DESC NULLS FIRST, fetch=5 03)----Projection: sum(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING AS sum1, sum(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING AS sum2, sum(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING AS sum3, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING AS min1, min(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING AS min2, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING AS min3, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING AS max1, max(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING AS max2, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING AS max3, count(*) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 4 PRECEDING AND 8 FOLLOWING AS cnt1, count(*) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING AS cnt2, sum(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 1 PRECEDING AND 4 FOLLOWING AS sumr1, sum(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 1 PRECEDING AND 8 FOLLOWING AS sumr2, sum(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 5 FOLLOWING AS sumr3, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING AS minr1, min(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING AS minr2, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING AS minr3, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING AS maxr1, max(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING AS maxr2, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING AS maxr3, count(*) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING AS cntr1, count(*) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING AS cntr2, sum(annotated_data_finite.desc_col) ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING AS sum4, count(*) ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING AS cnt3, annotated_data_finite.inc_col -04)------WindowAggr: windowExpr=[[sum(__common_expr_1 AS annotated_data_finite.desc_col) ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING, count(Int64(1)) ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING AS count(*) ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING]] +04)------WindowAggr: windowExpr=[[sum(__common_expr_1 AS annotated_data_finite.desc_col) ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING, count(*) ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING]] 05)--------Projection: __common_expr_1, annotated_data_finite.inc_col, sum(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 1 PRECEDING AND 4 FOLLOWING, sum(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 1 PRECEDING AND 8 FOLLOWING, sum(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 5 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, count(*) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING, count(*) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING, sum(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, sum(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, sum(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, count(*) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 4 PRECEDING AND 8 FOLLOWING, count(*) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING -06)----------WindowAggr: windowExpr=[[sum(__common_expr_2 AS annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, sum(__common_expr_1 AS annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, sum(__common_expr_2 AS annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, count(Int64(1)) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 4 PRECEDING AND 8 FOLLOWING AS count(*) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 4 PRECEDING AND 8 FOLLOWING, count(Int64(1)) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING AS count(*) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING]] -07)------------WindowAggr: windowExpr=[[sum(__common_expr_2 AS annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 1 PRECEDING AND 4 FOLLOWING, sum(__common_expr_1 AS annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 1 PRECEDING AND 8 FOLLOWING, sum(__common_expr_1 AS annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 5 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, count(Int64(1)) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING AS count(*) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING, count(Int64(1)) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING AS count(*) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING]] +06)----------WindowAggr: windowExpr=[[sum(__common_expr_2 AS annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, sum(__common_expr_1 AS annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, sum(__common_expr_2 AS annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, count(*) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] RANGE BETWEEN 4 PRECEDING AND 8 FOLLOWING, count(*) ORDER BY [annotated_data_finite.ts ASC NULLS LAST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING]] +07)------------WindowAggr: windowExpr=[[sum(__common_expr_2 AS annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 1 PRECEDING AND 4 FOLLOWING, sum(__common_expr_1 AS annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 1 PRECEDING AND 8 FOLLOWING, sum(__common_expr_1 AS annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 5 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, min(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 10 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.desc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 5 PRECEDING AND 1 FOLLOWING, max(annotated_data_finite.inc_col) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND 10 FOLLOWING, count(*) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING, count(*) ORDER BY [annotated_data_finite.ts DESC NULLS FIRST] ROWS BETWEEN 8 PRECEDING AND 1 FOLLOWING]] 08)--------------Projection: CAST(annotated_data_finite.desc_col AS Int64) AS __common_expr_1, CAST(annotated_data_finite.inc_col AS Int64) AS __common_expr_2, annotated_data_finite.ts, annotated_data_finite.inc_col, annotated_data_finite.desc_col 09)----------------TableScan: annotated_data_finite projection=[ts, inc_col, desc_col] physical_plan @@ -4112,7 +4112,7 @@ EXPLAIN select count(*) over (partition by a order by a) from (select * from a w ---- logical_plan 01)Projection: count(*) PARTITION BY [a.a] ORDER BY [a.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW -02)--WindowAggr: windowExpr=[[count(Int64(1)) PARTITION BY [a.a] ORDER BY [a.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS count(*) PARTITION BY [a.a] ORDER BY [a.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +02)--WindowAggr: windowExpr=[[count(*) PARTITION BY [a.a] ORDER BY [a.a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] 03)----Filter: a.a = Int64(1) 04)------TableScan: a projection=[a] physical_plan diff --git a/datafusion/substrait/src/lib.rs b/datafusion/substrait/src/lib.rs index f33e86a2d20c0..a7493a48e4c57 100644 --- a/datafusion/substrait/src/lib.rs +++ b/datafusion/substrait/src/lib.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + //! Serialize / Deserialize DataFusion Plans to [Substrait.io] //! //! This crate provides support for serializing and deserializing both DataFusion diff --git a/datafusion/substrait/src/physical_plan/producer.rs b/datafusion/substrait/src/physical_plan/producer.rs index 3fc94a33442bf..e8c15731228cf 100644 --- a/datafusion/substrait/src/physical_plan/producer.rs +++ b/datafusion/substrait/src/physical_plan/producer.rs @@ -23,8 +23,8 @@ use crate::variation_const::{ }; use datafusion::arrow::datatypes::DataType; +use datafusion::datasource::source::DataSourceExec; use datafusion::error::{DataFusionError, Result}; -use datafusion::physical_plan::source::DataSourceExec; use datafusion::physical_plan::{displayable, ExecutionPlan}; use datafusion::datasource::physical_plan::{FileScanConfig, ParquetSource}; diff --git a/datafusion/substrait/tests/cases/consumer_integration.rs b/datafusion/substrait/tests/cases/consumer_integration.rs index 219f656bb471e..086c0858115ba 100644 --- a/datafusion/substrait/tests/cases/consumer_integration.rs +++ b/datafusion/substrait/tests/cases/consumer_integration.rs @@ -50,9 +50,9 @@ mod tests { let plan_str = tpch_plan_to_string(1).await?; assert_eq!( plan_str, - "Projection: LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS, sum(LINEITEM.L_QUANTITY) AS SUM_QTY, sum(LINEITEM.L_EXTENDEDPRICE) AS SUM_BASE_PRICE, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS SUM_DISC_PRICE, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT * Int32(1) + LINEITEM.L_TAX) AS SUM_CHARGE, avg(LINEITEM.L_QUANTITY) AS AVG_QTY, avg(LINEITEM.L_EXTENDEDPRICE) AS AVG_PRICE, avg(LINEITEM.L_DISCOUNT) AS AVG_DISC, count(Int64(1)) AS COUNT_ORDER\ + "Projection: LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS, sum(LINEITEM.L_QUANTITY) AS SUM_QTY, sum(LINEITEM.L_EXTENDEDPRICE) AS SUM_BASE_PRICE, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS SUM_DISC_PRICE, sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT * Int32(1) + LINEITEM.L_TAX) AS SUM_CHARGE, avg(LINEITEM.L_QUANTITY) AS AVG_QTY, avg(LINEITEM.L_EXTENDEDPRICE) AS AVG_PRICE, avg(LINEITEM.L_DISCOUNT) AS AVG_DISC, count(*) AS COUNT_ORDER\ \n Sort: LINEITEM.L_RETURNFLAG ASC NULLS LAST, LINEITEM.L_LINESTATUS ASC NULLS LAST\ - \n Aggregate: groupBy=[[LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS]], aggr=[[sum(LINEITEM.L_QUANTITY), sum(LINEITEM.L_EXTENDEDPRICE), sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT), sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT * Int32(1) + LINEITEM.L_TAX), avg(LINEITEM.L_QUANTITY), avg(LINEITEM.L_EXTENDEDPRICE), avg(LINEITEM.L_DISCOUNT), count(Int64(1))]]\ + \n Aggregate: groupBy=[[LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS]], aggr=[[sum(LINEITEM.L_QUANTITY), sum(LINEITEM.L_EXTENDEDPRICE), sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT), sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT * Int32(1) + LINEITEM.L_TAX), avg(LINEITEM.L_QUANTITY), avg(LINEITEM.L_EXTENDEDPRICE), avg(LINEITEM.L_DISCOUNT), count(*)]]\ \n Projection: LINEITEM.L_RETURNFLAG, LINEITEM.L_LINESTATUS, LINEITEM.L_QUANTITY, LINEITEM.L_EXTENDEDPRICE, LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT), LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) * (CAST(Int32(1) AS Decimal128(15, 2)) + LINEITEM.L_TAX), LINEITEM.L_DISCOUNT\ \n Filter: LINEITEM.L_SHIPDATE <= Date32(\"1998-12-01\") - IntervalDayTime(\"IntervalDayTime { days: 0, milliseconds: 10368000 }\")\ \n TableScan: LINEITEM" @@ -119,9 +119,9 @@ mod tests { let plan_str = tpch_plan_to_string(4).await?; assert_eq!( plan_str, - "Projection: ORDERS.O_ORDERPRIORITY, count(Int64(1)) AS ORDER_COUNT\ + "Projection: ORDERS.O_ORDERPRIORITY, count(*) AS ORDER_COUNT\ \n Sort: ORDERS.O_ORDERPRIORITY ASC NULLS LAST\ - \n Aggregate: groupBy=[[ORDERS.O_ORDERPRIORITY]], aggr=[[count(Int64(1))]]\ + \n Aggregate: groupBy=[[ORDERS.O_ORDERPRIORITY]], aggr=[[count(*)]]\ \n Projection: ORDERS.O_ORDERPRIORITY\ \n Filter: ORDERS.O_ORDERDATE >= CAST(Utf8(\"1993-07-01\") AS Date32) AND ORDERS.O_ORDERDATE < CAST(Utf8(\"1993-10-01\") AS Date32) AND EXISTS ()\ \n Subquery:\ @@ -269,10 +269,10 @@ mod tests { let plan_str = tpch_plan_to_string(13).await?; assert_eq!( plan_str, - "Projection: count(ORDERS.O_ORDERKEY) AS C_COUNT, count(Int64(1)) AS CUSTDIST\ - \n Sort: count(Int64(1)) DESC NULLS FIRST, count(ORDERS.O_ORDERKEY) DESC NULLS FIRST\ - \n Projection: count(ORDERS.O_ORDERKEY), count(Int64(1))\ - \n Aggregate: groupBy=[[count(ORDERS.O_ORDERKEY)]], aggr=[[count(Int64(1))]]\ + "Projection: count(ORDERS.O_ORDERKEY) AS C_COUNT, count(*) AS CUSTDIST\ + \n Sort: count(*) DESC NULLS FIRST, count(ORDERS.O_ORDERKEY) DESC NULLS FIRST\ + \n Projection: count(ORDERS.O_ORDERKEY), count(*)\ + \n Aggregate: groupBy=[[count(ORDERS.O_ORDERKEY)]], aggr=[[count(*)]]\ \n Projection: count(ORDERS.O_ORDERKEY)\ \n Aggregate: groupBy=[[CUSTOMER.C_CUSTKEY]], aggr=[[count(ORDERS.O_ORDERKEY)]]\ \n Projection: CUSTOMER.C_CUSTKEY, ORDERS.O_ORDERKEY\ @@ -410,10 +410,10 @@ mod tests { let plan_str = tpch_plan_to_string(21).await?; assert_eq!( plan_str, - "Projection: SUPPLIER.S_NAME, count(Int64(1)) AS NUMWAIT\ + "Projection: SUPPLIER.S_NAME, count(*) AS NUMWAIT\ \n Limit: skip=0, fetch=100\ - \n Sort: count(Int64(1)) DESC NULLS FIRST, SUPPLIER.S_NAME ASC NULLS LAST\ - \n Aggregate: groupBy=[[SUPPLIER.S_NAME]], aggr=[[count(Int64(1))]]\ + \n Sort: count(*) DESC NULLS FIRST, SUPPLIER.S_NAME ASC NULLS LAST\ + \n Aggregate: groupBy=[[SUPPLIER.S_NAME]], aggr=[[count(*)]]\ \n Projection: SUPPLIER.S_NAME\ \n Filter: SUPPLIER.S_SUPPKEY = LINEITEM.L_SUPPKEY AND ORDERS.O_ORDERKEY = LINEITEM.L_ORDERKEY AND ORDERS.O_ORDERSTATUS = Utf8(\"F\") AND LINEITEM.L_RECEIPTDATE > LINEITEM.L_COMMITDATE AND EXISTS () AND NOT EXISTS () AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_NAME = Utf8(\"SAUDI ARABIA\")\ \n Subquery:\ @@ -438,9 +438,9 @@ mod tests { let plan_str = tpch_plan_to_string(22).await?; assert_eq!( plan_str, - "Projection: substr(CUSTOMER.C_PHONE,Int32(1),Int32(2)) AS CNTRYCODE, count(Int64(1)) AS NUMCUST, sum(CUSTOMER.C_ACCTBAL) AS TOTACCTBAL\ + "Projection: substr(CUSTOMER.C_PHONE,Int32(1),Int32(2)) AS CNTRYCODE, count(*) AS NUMCUST, sum(CUSTOMER.C_ACCTBAL) AS TOTACCTBAL\ \n Sort: substr(CUSTOMER.C_PHONE,Int32(1),Int32(2)) ASC NULLS LAST\ - \n Aggregate: groupBy=[[substr(CUSTOMER.C_PHONE,Int32(1),Int32(2))]], aggr=[[count(Int64(1)), sum(CUSTOMER.C_ACCTBAL)]]\ + \n Aggregate: groupBy=[[substr(CUSTOMER.C_PHONE,Int32(1),Int32(2))]], aggr=[[count(*), sum(CUSTOMER.C_ACCTBAL)]]\ \n Projection: substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)), CUSTOMER.C_ACCTBAL\ \n Filter: (substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8(\"13\") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8(\"31\") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8(\"23\") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8(\"29\") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8(\"30\") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8(\"18\") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8(\"17\") AS Utf8)) AND CUSTOMER.C_ACCTBAL > () AND NOT EXISTS ()\ \n Subquery:\ diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 5fb357dfcd23d..68856117a38c5 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -687,7 +687,7 @@ async fn simple_intersect() -> Result<()> { // Substrait treats both count(*) and count(1) the same assert_expected_plan( "SELECT count(*) FROM (SELECT data.a FROM data INTERSECT SELECT data2.a FROM data2);", - "Aggregate: groupBy=[[]], aggr=[[count(Int64(1)) AS count(*)]]\ + "Aggregate: groupBy=[[]], aggr=[[count(*)]]\ \n Projection: \ \n LeftSemi Join: data.a = data2.a\ \n Aggregate: groupBy=[[data.a]], aggr=[[]]\ @@ -822,7 +822,7 @@ async fn simple_intersect_table_reuse() -> Result<()> { // Schema check works because we set aliases to what the Substrait consumer will generate. assert_expected_plan( "SELECT count(1) FROM (SELECT left.a FROM data AS left INTERSECT SELECT right.a FROM data AS right);", - "Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]]\ + "Aggregate: groupBy=[[]], aggr=[[count(*)]]\ \n Projection: \ \n LeftSemi Join: left.a = right.a\ \n SubqueryAlias: left\ diff --git a/datafusion/wasmtest/Cargo.toml b/datafusion/wasmtest/Cargo.toml index 7db051ad191f7..30d5bcaedcb73 100644 --- a/datafusion/wasmtest/Cargo.toml +++ b/datafusion/wasmtest/Cargo.toml @@ -27,6 +27,9 @@ license = { workspace = true } authors = { workspace = true } rust-version = { workspace = true } +[package.metadata.docs.rs] +all-features = true + [lints] workspace = true @@ -47,14 +50,12 @@ datafusion-common = { workspace = true, default-features = true } datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-optimizer = { workspace = true, default-features = true } -datafusion-physical-expr = { workspace = true, default-features = true } datafusion-physical-plan = { workspace = true } datafusion-sql = { workspace = true } # getrandom must be compiled with js feature getrandom = { version = "0.2.8", features = ["js"] } wasm-bindgen = "0.2.99" -wasm-bindgen-futures = "0.4.49" [dev-dependencies] tokio = { workspace = true } diff --git a/datafusion/wasmtest/src/lib.rs b/datafusion/wasmtest/src/lib.rs index c7c620d1be3a4..e2ba50beb6574 100644 --- a/datafusion/wasmtest/src/lib.rs +++ b/datafusion/wasmtest/src/lib.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_auto_cfg))] + extern crate wasm_bindgen; use datafusion_common::{DFSchema, ScalarValue}; diff --git a/rust-toolchain.toml b/rust-toolchain.toml index bd764d2010186..11f4fb798c376 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -19,5 +19,5 @@ # to compile this workspace and run CI jobs. [toolchain] -channel = "1.84.1" +channel = "1.85.0" components = ["rustfmt", "clippy"] diff --git a/test-utils/src/array_gen/string.rs b/test-utils/src/array_gen/string.rs index a405cb76b1bd2..ac659ae67bc0e 100644 --- a/test-utils/src/array_gen/string.rs +++ b/test-utils/src/array_gen/string.rs @@ -97,8 +97,7 @@ fn random_string(rng: &mut StdRng, max_len: usize) -> String { let len = rng.gen_range(1..=max_len); rng.sample_iter::(rand::distributions::Standard) .take(len) - .map(char::from) - .collect::() + .collect() } } }